Optimize db queries (part 2) (#2929)

* Process fleets using one session

* Select only ids in locking select

* Do not load InstanceModel.jobs by default

* Add indexes on status columns

* Do not load all projects attrs for projects filtering

* Fix comments

* Do not load ProjectModel.default_gateway by default

* Add TODO

Author: r4victor

src/dstack/_internal/server/background/__init__.py

@@ -79,6 +79,11 @@ def start_background_tasks() -> AsyncIOScheduler:
         process_idle_volumes, IntervalTrigger(seconds=60, jitter=10), max_instances=1
     )
     _scheduler.add_job(process_placement_groups, IntervalTrigger(seconds=30, jitter=5))
+    _scheduler.add_job(
+        process_fleets,
+        IntervalTrigger(seconds=10, jitter=2),
+        max_instances=1,
+    )
     for replica in range(settings.SERVER_BACKGROUND_PROCESSING_FACTOR):
         # Add multiple copies of tasks if requested.
         # max_instances=1 for additional copies to avoid running too many tasks.
@@ -113,11 +118,5 @@ def start_background_tasks() -> AsyncIOScheduler:
             kwargs={"batch_size": 5},
             max_instances=2 if replica == 0 else 1,
         )
-        _scheduler.add_job(
-            process_fleets,
-            IntervalTrigger(seconds=10, jitter=2),
-            kwargs={"batch_size": 5},
-            max_instances=2 if replica == 0 else 1,
-        )
     _scheduler.start()
     return _scheduler

src/dstack/_internal/server/background/tasks/process_fleets.py

@@ -1,36 +1,35 @@
-import asyncio
 from datetime import timedelta
+from typing import List
 
-from sqlalchemy import select
+from sqlalchemy import select, update
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, load_only
 
 from dstack._internal.core.models.fleets import FleetStatus
 from dstack._internal.server.db import get_db, get_session_ctx
-from dstack._internal.server.models import FleetModel, InstanceModel, JobModel, RunModel
+from dstack._internal.server.models import (
+    FleetModel,
+    InstanceModel,
+    JobModel,
+    PlacementGroupModel,
+    RunModel,
+)
 from dstack._internal.server.services.fleets import (
     is_fleet_empty,
     is_fleet_in_use,
 )
 from dstack._internal.server.services.locking import get_locker
-from dstack._internal.server.services.placement import schedule_fleet_placement_groups_deletion
 from dstack._internal.utils.common import get_current_datetime
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
 
 
+BATCH_SIZE = 10
 MIN_PROCESSING_INTERVAL = timedelta(seconds=30)
 
 
-async def process_fleets(batch_size: int = 1):
-    tasks = []
-    for _ in range(batch_size):
-        tasks.append(_process_next_fleet())
-    await asyncio.gather(*tasks)
-
-
-async def _process_next_fleet():
+async def process_fleets():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(FleetModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
@@ -42,50 +41,62 @@ async def _process_next_fleet():
                     FleetModel.last_processed_at
                     < get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
+                .options(load_only(FleetModel.id))
                 .order_by(FleetModel.last_processed_at.asc())
-                .limit(1)
+                .limit(BATCH_SIZE)
                 .with_for_update(skip_locked=True, key_share=True)
             )
-            fleet_model = res.scalar()
-            if fleet_model is None:
-                return
-            lockset.add(fleet_model.id)
+            fleet_models = list(res.scalars().all())
+            fleet_ids = [fm.id for fm in fleet_models]
+            for fleet_id in fleet_ids:
+                lockset.add(fleet_id)
         try:
-            fleet_model_id = fleet_model.id
-            await _process_fleet(session=session, fleet_model=fleet_model)
+            await _process_fleets(session=session, fleet_models=fleet_models)
         finally:
-            lockset.difference_update([fleet_model_id])
+            lockset.difference_update(fleet_ids)
 
 
-async def _process_fleet(session: AsyncSession, fleet_model: FleetModel):
-    logger.debug("Processing fleet %s", fleet_model.name)
+async def _process_fleets(session: AsyncSession, fleet_models: List[FleetModel]):
+    fleet_ids = [fm.id for fm in fleet_models]
     # Refetch to load related attributes.
     res = await session.execute(
         select(FleetModel)
-        .where(FleetModel.id == fleet_model.id)
+        .where(FleetModel.id.in_(fleet_ids))
         .options(joinedload(FleetModel.instances).load_only(InstanceModel.deleted))
         .options(
             joinedload(FleetModel.instances).joinedload(InstanceModel.jobs).load_only(JobModel.id)
         )
         .options(joinedload(FleetModel.runs).load_only(RunModel.status))
         .execution_options(populate_existing=True)
     )
-    fleet_model = res.unique().scalar_one()
-    await _autodelete_fleet(session=session, fleet_model=fleet_model)
+    fleet_models = list(res.unique().scalars().all())
 
+    deleted_fleets_ids = []
+    now = get_current_datetime()
+    for fleet_model in fleet_models:
+        deleted = _autodelete_fleet(fleet_model)
+        if deleted:
+            deleted_fleets_ids.append(fleet_model.id)
+        fleet_model.last_processed_at = now
 
-async def _autodelete_fleet(session: AsyncSession, fleet_model: FleetModel):
+    await session.execute(
+        update(PlacementGroupModel)
+        .where(
+            PlacementGroupModel.fleet_id.in_(deleted_fleets_ids),
+        )
+        .values(fleet_deleted=True)
+    )
+    await session.commit()
+
+
+def _autodelete_fleet(fleet_model: FleetModel) -> bool:
     # Currently all empty fleets are autodeleted.
     # TODO: If fleets with `nodes: 0..` are supported, their deletion should be skipped.
     if is_fleet_in_use(fleet_model) or not is_fleet_empty(fleet_model):
-        fleet_model.last_processed_at = get_current_datetime()
-        await session.commit()
-        return
+        return False
 
     logger.info("Automatic cleanup of an empty fleet %s", fleet_model.name)
     fleet_model.status = FleetStatus.TERMINATED
     fleet_model.deleted = True
-    fleet_model.last_processed_at = get_current_datetime()
-    await schedule_fleet_placement_groups_deletion(session=session, fleet_id=fleet_model.id)
-    await session.commit()
     logger.info("Fleet %s deleted", fleet_model.name)
+    return True

src/dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -7,7 +7,7 @@
 
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, load_only
 
 from dstack._internal import settings
 from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT, DSTACK_SHIM_HTTP_PORT
@@ -110,6 +110,7 @@ async def _process_next_running_job():
                     JobModel.last_processed_at
                     < common_utils.get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
+                .options(load_only(JobModel.id))
                 .order_by(JobModel.last_processed_at.asc())
                 .limit(1)
                 .with_for_update(

src/dstack/_internal/server/background/tasks/process_runs.py

@@ -4,7 +4,7 @@
 
 from sqlalchemy import and_, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, selectinload
+from sqlalchemy.orm import joinedload, load_only, selectinload
 
 import dstack._internal.server.services.services.autoscalers as autoscalers
 from dstack._internal.core.errors import ServerError
@@ -102,6 +102,7 @@ async def _process_next_run():
                     ),
                 )
                 .options(joinedload(RunModel.jobs).load_only(JobModel.id))
+                .options(load_only(RunModel.id))
                 .order_by(RunModel.last_processed_at.asc())
                 .limit(1)
                 .with_for_update(skip_locked=True, key_share=True, of=RunModel)
@@ -134,7 +135,6 @@ async def _process_next_run():
 
 
 async def _process_run(session: AsyncSession, run_model: RunModel):
-    logger.debug("%s: processing run", fmt(run_model))
     # Refetch to load related attributes.
     res = await session.execute(
         select(RunModel)
@@ -150,6 +150,7 @@ async def _process_run(session: AsyncSession, run_model: RunModel):
         .execution_options(populate_existing=True)
     )
     run_model = res.unique().scalar_one()
+    logger.debug("%s: processing run", fmt(run_model))
     try:
         if run_model.status == RunStatus.PENDING:
             await _process_pending_run(session, run_model)

src/dstack/_internal/server/background/tasks/process_submitted_jobs.py

@@ -5,7 +5,7 @@
 
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, lazyload, selectinload
+from sqlalchemy.orm import joinedload, load_only, selectinload
 
 from dstack._internal.core.backends.base.backend import Backend
 from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport
@@ -120,6 +120,7 @@ async def _process_next_submitted_job():
                     JobModel.status == JobStatus.SUBMITTED,
                     JobModel.id.not_in(lockset),
                 )
+                .options(load_only(JobModel.id))
                 # Jobs are process in FIFO sorted by priority globally,
                 # thus runs from different projects can "overtake" each other by using higher priorities.
                 # That's not a big problem as long as projects do not compete for the same compute resources.
@@ -152,7 +153,6 @@ async def _process_next_submitted_job():
 
 
 async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
-    logger.debug("%s: provisioning has started", fmt(job_model))
     # Refetch to load related attributes.
     res = await session.execute(
         select(JobModel).where(JobModel.id == job_model.id).options(joinedload(JobModel.instance))
@@ -166,6 +166,8 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         .options(joinedload(RunModel.fleet).joinedload(FleetModel.instances))
     )
     run_model = res.unique().scalar_one()
+    logger.debug("%s: provisioning has started", fmt(job_model))
+
     project = run_model.project
     run = run_model_to_run(run_model)
     run_spec = run.run_spec
@@ -227,7 +229,6 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
                 InstanceModel.deleted == False,
                 InstanceModel.total_blocks > InstanceModel.busy_blocks,
             )
-            .options(lazyload(InstanceModel.jobs))
             .order_by(InstanceModel.id)  # take locks in order
             .with_for_update(key_share=True)
         )

src/dstack/_internal/server/background/tasks/process_terminating_jobs.py

@@ -2,7 +2,7 @@
 
 from sqlalchemy import or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, lazyload
+from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.models.runs import JobStatus
 from dstack._internal.server.db import get_db, get_session_ctx
@@ -65,7 +65,6 @@ async def _process_next_terminating_job():
                         InstanceModel.id == job_model.used_instance_id,
                         InstanceModel.id.not_in(instance_lockset),
                     )
-                    .options(lazyload(InstanceModel.jobs))
                     .with_for_update(skip_locked=True, key_share=True)
                 )
                 instance_model = res.scalar()
@@ -94,6 +93,7 @@ async def _process_job(session: AsyncSession, job_model: JobModel):
         .options(
             joinedload(InstanceModel.project).joinedload(ProjectModel.backends),
             joinedload(InstanceModel.volume_attachments).joinedload(VolumeAttachmentModel.volume),
+            joinedload(InstanceModel.jobs).load_only(JobModel.id),
         )
     )
     instance_model = res.unique().scalar()

src/dstack/_internal/server/migrations/versions/50dd7ea98639_index_status_columns.py

@@ -0,0 +1,55 @@
+"""Index status columns
+
+Revision ID: 50dd7ea98639
+Revises: ec02a26a256c
+Create Date: 2025-07-25 10:36:25.127923
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "50dd7ea98639"
+down_revision = "ec02a26a256c"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("runs", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_runs_status"), ["status"], unique=False)
+
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_jobs_status"), ["status"], unique=False)
+
+    with op.batch_alter_table("fleets", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_fleets_status"), ["status"], unique=False)
+
+    with op.batch_alter_table("instances", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_instances_status"), ["status"], unique=False)
+
+    with op.batch_alter_table("volumes", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_volumes_status"), ["status"], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("runs", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_runs_status"))
+
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_jobs_status"))
+
+    with op.batch_alter_table("fleets", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_fleets_status"))
+
+    with op.batch_alter_table("instances", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_instances_status"))
+
+    with op.batch_alter_table("volumes", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_volumes_status"))
+
+    # ### end Alembic commands ###