Optimize db queries (dstackai#2928)

* Remove extra queries when processing running job

* Optimize process_instances queries

* Optimize process_runs queries

* Optimize process_fleets queries

* Optimize attribute loading

* Select for update with of=InstanceModel

* Remove extra selectin

* Select for update with of=VolumeModel

* Select with .load_only(InstanceModel.fleet_id)

* Select with .load_only(FleetModel.name)

Author: r4victor

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

@@ -7,7 +7,7 @@
 
 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
+from dstack._internal.server.models import FleetModel, InstanceModel, JobModel, RunModel
 from dstack._internal.server.services.fleets import (
     is_fleet_empty,
     is_fleet_in_use,
@@ -60,13 +60,14 @@ async def _process_next_fleet():
 async def _process_fleet(session: AsyncSession, fleet_model: FleetModel):
     logger.debug("Processing fleet %s", fleet_model.name)
     # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(
         select(FleetModel)
         .where(FleetModel.id == fleet_model.id)
-        .options(joinedload(FleetModel.project))
-        .options(joinedload(FleetModel.instances))
-        .options(joinedload(FleetModel.runs))
+        .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()

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

@@ -110,7 +110,6 @@ async def _process_connection(conn: GatewayConnection):
 async def _process_submitted_gateway(session: AsyncSession, gateway_model: GatewayModel):
     logger.info("%s: started gateway provisioning", fmt(gateway_model))
     # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(
         select(GatewayModel)
         .where(GatewayModel.id == gateway_model.id)
@@ -157,7 +156,6 @@ async def _process_provisioning_gateway(
     session: AsyncSession, gateway_model: GatewayModel
 ) -> None:
     # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(
         select(GatewayModel)
         .where(GatewayModel.id == gateway_model.id)

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

@@ -10,7 +10,7 @@
 from dstack._internal.core.models.profiles import parse_duration
 from dstack._internal.core.models.volumes import VolumeStatus
 from dstack._internal.server.db import get_db, get_session_ctx
-from dstack._internal.server.models import ProjectModel, VolumeModel
+from dstack._internal.server.models import ProjectModel, UserModel, VolumeModel
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.volumes import (
@@ -49,7 +49,7 @@ async def process_idle_volumes():
             select(VolumeModel)
             .where(VolumeModel.id.in_(volume_ids))
             .options(joinedload(VolumeModel.project).joinedload(ProjectModel.backends))
-            .options(joinedload(VolumeModel.user))
+            .options(joinedload(VolumeModel.user).load_only(UserModel.name))
             .options(joinedload(VolumeModel.attachments))
             .execution_options(populate_existing=True)
         )

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

@@ -9,7 +9,7 @@
 from pydantic import ValidationError
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, lazyload
+from sqlalchemy.orm import joinedload
 
 from dstack._internal import settings
 from dstack._internal.core.backends.base.compute import (
@@ -78,6 +78,7 @@
 from dstack._internal.server.models import (
     FleetModel,
     InstanceModel,
+    JobModel,
     PlacementGroupModel,
     ProjectModel,
 )
@@ -106,6 +107,7 @@
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
 from dstack._internal.utils.common import (
     get_current_datetime,
+    get_or_error,
     run_async,
 )
 from dstack._internal.utils.logging import get_logger
@@ -154,10 +156,11 @@ async def _process_next_instance():
                     InstanceModel.last_processed_at
                     < get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
-                .options(lazyload(InstanceModel.jobs))
+                .options(joinedload(InstanceModel.jobs).load_only(JobModel.id, JobModel.status))
+                .options(joinedload(InstanceModel.project).load_only(ProjectModel.ssh_private_key))
                 .order_by(InstanceModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True, key_share=True)
+                .with_for_update(skip_locked=True, key_share=True, of=InstanceModel)
             )
             instance = res.scalar()
             if instance is None:
@@ -171,23 +174,22 @@ async def _process_next_instance():
 
 
 async def _process_instance(session: AsyncSession, instance: InstanceModel):
-    # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
-    res = await session.execute(
-        select(InstanceModel)
-        .where(InstanceModel.id == instance.id)
-        .options(joinedload(InstanceModel.project).joinedload(ProjectModel.backends))
-        .options(joinedload(InstanceModel.jobs))
-        .options(joinedload(InstanceModel.fleet).joinedload(FleetModel.instances))
-        .execution_options(populate_existing=True)
-    )
-    instance = res.unique().scalar_one()
-    if (
-        instance.status == InstanceStatus.IDLE
-        and instance.termination_policy == TerminationPolicy.DESTROY_AFTER_IDLE
-        and not instance.jobs
+    if instance.status in (
+        InstanceStatus.PENDING,
+        InstanceStatus.TERMINATING,
     ):
-        await _mark_terminating_if_idle_duration_expired(instance)
+        # Refetch to load related attributes.
+        # Load related attributes only for statuses that always need them.
+        res = await session.execute(
+            select(InstanceModel)
+            .where(InstanceModel.id == instance.id)
+            .options(joinedload(InstanceModel.project).joinedload(ProjectModel.backends))
+            .options(joinedload(InstanceModel.jobs).load_only(JobModel.id, JobModel.status))
+            .options(joinedload(InstanceModel.fleet).joinedload(FleetModel.instances))
+            .execution_options(populate_existing=True)
+        )
+        instance = res.unique().scalar_one()
+
     if instance.status == InstanceStatus.PENDING:
         if instance.remote_connection_info is not None:
             await _add_remote(instance)
@@ -201,15 +203,23 @@ async def _process_instance(session: AsyncSession, instance: InstanceModel):
         InstanceStatus.IDLE,
         InstanceStatus.BUSY,
     ):
-        await _check_instance(instance)
+        idle_duration_expired = _check_and_mark_terminating_if_idle_duration_expired(instance)
+        if not idle_duration_expired:
+            await _check_instance(session, instance)
     elif instance.status == InstanceStatus.TERMINATING:
         await _terminate(instance)
 
     instance.last_processed_at = get_current_datetime()
     await session.commit()
 
 
-async def _mark_terminating_if_idle_duration_expired(instance: InstanceModel):
+def _check_and_mark_terminating_if_idle_duration_expired(instance: InstanceModel):
+    if not (
+        instance.status == InstanceStatus.IDLE
+        and instance.termination_policy == TerminationPolicy.DESTROY_AFTER_IDLE
+        and not instance.jobs
+    ):
+        return False
     idle_duration = _get_instance_idle_duration(instance)
     idle_seconds = instance.termination_idle_time
     delta = datetime.timedelta(seconds=idle_seconds)
@@ -225,6 +235,8 @@ async def _mark_terminating_if_idle_duration_expired(instance: InstanceModel):
                 "instance_status": instance.status.value,
             },
         )
+        return True
+    return False
 
 
 async def _add_remote(instance: InstanceModel) -> None:
@@ -703,7 +715,7 @@ def _mark_terminated(instance: InstanceModel, termination_reason: str) -> None:
     )
 
 
-async def _check_instance(instance: InstanceModel) -> None:
+async def _check_instance(session: AsyncSession, instance: InstanceModel) -> None:
     if (
         instance.status == InstanceStatus.BUSY
         and instance.jobs
@@ -722,12 +734,16 @@ async def _check_instance(instance: InstanceModel) -> None:
         )
         return
 
-    job_provisioning_data = JobProvisioningData.__response__.parse_raw(
-        instance.job_provisioning_data
-    )
+    job_provisioning_data = get_or_error(get_instance_provisioning_data(instance))
     if job_provisioning_data.hostname is None:
+        res = await session.execute(
+            select(ProjectModel)
+            .where(ProjectModel.id == instance.id)
+            .options(joinedload(ProjectModel.backends))
+        )
+        project = res.scalar_one()
         await _wait_for_instance_provisioning_data(
-            project=instance.project,
+            project=project,
             instance=instance,
             job_provisioning_data=job_provisioning_data,
         )

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

@@ -9,7 +9,7 @@
 from dstack._internal.core.models.runs import JobStatus
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_session_ctx
-from dstack._internal.server.models import InstanceModel, JobMetricsPoint, JobModel
+from dstack._internal.server.models import InstanceModel, JobMetricsPoint, JobModel, ProjectModel
 from dstack._internal.server.schemas.runner import MetricsResponse
 from dstack._internal.server.services.instances import get_instance_ssh_private_keys
 from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_runtime_data
@@ -31,7 +31,11 @@ async def collect_metrics():
         res = await session.execute(
             select(JobModel)
             .where(JobModel.status.in_([JobStatus.RUNNING]))
-            .options(joinedload(JobModel.instance).joinedload(InstanceModel.project))
+            .options(
+                joinedload(JobModel.instance)
+                .joinedload(InstanceModel.project)
+                .load_only(ProjectModel.ssh_private_key)
+            )
             .order_by(JobModel.last_processed_at.asc())
             .limit(MAX_JOBS_FETCHED)
         )

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

@@ -9,7 +9,12 @@
 from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
 from dstack._internal.core.models.runs import JobStatus
 from dstack._internal.server.db import get_session_ctx
-from dstack._internal.server.models import InstanceModel, JobModel, JobPrometheusMetrics
+from dstack._internal.server.models import (
+    InstanceModel,
+    JobModel,
+    JobPrometheusMetrics,
+    ProjectModel,
+)
 from dstack._internal.server.services.instances import get_instance_ssh_private_keys
 from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_runtime_data
 from dstack._internal.server.services.runner import client
@@ -43,7 +48,11 @@ async def collect_prometheus_metrics():
                     JobPrometheusMetrics.collected_at < cutoff,
                 ),
             )
-            .options(joinedload(JobModel.instance).joinedload(InstanceModel.project))
+            .options(
+                joinedload(JobModel.instance)
+                .joinedload(InstanceModel.project)
+                .load_only(ProjectModel.ssh_private_key)
+            )
             .order_by(JobModel.last_processed_at.asc())
             .limit(MAX_JOBS_FETCHED)
         )