Optimize process_instances queries

Author: r4victor

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,7 +156,8 @@ 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)
@@ -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_running_jobs.py

@@ -143,7 +143,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
     res = await session.execute(
         select(RunModel)
         .where(RunModel.id == job_model.run_id)
-        .options(joinedload(RunModel.project).joinedload(ProjectModel.backends))
+        .options(joinedload(RunModel.project))
         .options(joinedload(RunModel.user))
         .options(joinedload(RunModel.repo))
         .options(joinedload(RunModel.jobs))
@@ -163,22 +163,18 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
 
     job = find_job(run.jobs, job_model.replica_num, job_model.job_num)
 
-    # Wait until all other jobs in the replica are provisioned
-    for other_job in run.jobs:
-        if (
-            other_job.job_spec.replica_num == job.job_spec.replica_num
-            and other_job.job_submissions[-1].status == JobStatus.SUBMITTED
-        ):
-            job_model.last_processed_at = common_utils.get_current_datetime()
-            await session.commit()
-            return
-
-    server_ssh_private_keys = get_instance_ssh_private_keys(
-        common_utils.get_or_error(job_model.instance)
-    )
-
     initial_status = job_model.status
     if initial_status in [JobStatus.PROVISIONING, JobStatus.PULLING]:
+        # Wait until all other jobs in the replica are provisioned
+        for other_job in run.jobs:
+            if (
+                other_job.job_spec.replica_num == job.job_spec.replica_num
+                and other_job.job_submissions[-1].status == JobStatus.SUBMITTED
+            ):
+                job_model.last_processed_at = common_utils.get_current_datetime()
+                await session.commit()
+                return
+
         cluster_info = _get_cluster_info(
             jobs=run.jobs,
             replica_num=job.job_spec.replica_num,
@@ -210,94 +206,98 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
             job_model.last_processed_at = common_utils.get_current_datetime()
             return
 
+    server_ssh_private_keys = get_instance_ssh_private_keys(
+        common_utils.get_or_error(job_model.instance)
+    )
+
     if initial_status == JobStatus.PROVISIONING:
         if job_provisioning_data.hostname is None:
             await _wait_for_instance_provisioning_data(job_model=job_model)
+            job_model.last_processed_at = common_utils.get_current_datetime()
+            await session.commit()
+            return
+        if _should_wait_for_other_nodes(run, job, job_model):
+            job_model.last_processed_at = common_utils.get_current_datetime()
+            await session.commit()
+            return
+
+        # fails are acceptable until timeout is exceeded
+        if job_provisioning_data.dockerized:
+            logger.debug(
+                "%s: process provisioning job with shim, age=%s",
+                fmt(job_model),
+                job_submission.age,
+            )
+            ssh_user = job_provisioning_data.username
+            user_ssh_key = run.run_spec.ssh_key_pub.strip()
+            public_keys = [project.ssh_public_key.strip(), user_ssh_key]
+            if job_provisioning_data.backend == BackendType.LOCAL:
+                # No need to update ~/.ssh/authorized_keys when running shim locally
+                user_ssh_key = ""
+            success = await common_utils.run_async(
+                _process_provisioning_with_shim,
+                server_ssh_private_keys,
+                job_provisioning_data,
+                None,
+                run,
+                job_model,
+                job_provisioning_data,
+                volumes,
+                job.job_spec.registry_auth,
+                public_keys,
+                ssh_user,
+                user_ssh_key,
+            )
         else:
-            if _should_wait_for_other_nodes(run, job, job_model):
-                job_model.last_processed_at = common_utils.get_current_datetime()
-                await session.commit()
-                return
+            logger.debug(
+                "%s: process provisioning job without shim, age=%s",
+                fmt(job_model),
+                job_submission.age,
+            )
+            # FIXME: downloading file archives and code here is a waste of time if
+            # the runner is not ready yet
+            file_archives = await _get_job_file_archives(
+                session=session,
+                archive_mappings=job.job_spec.file_archives,
+                user=run_model.user,
+            )
+            code = await _get_job_code(
+                session=session,
+                project=project,
+                repo=repo_model,
+                code_hash=_get_repo_code_hash(run, job),
+            )
 
-            # fails are acceptable until timeout is exceeded
-            if job_provisioning_data.dockerized:
-                logger.debug(
-                    "%s: process provisioning job with shim, age=%s",
-                    fmt(job_model),
-                    job_submission.age,
-                )
-                ssh_user = job_provisioning_data.username
-                user_ssh_key = run.run_spec.ssh_key_pub.strip()
-                public_keys = [project.ssh_public_key.strip(), user_ssh_key]
-                if job_provisioning_data.backend == BackendType.LOCAL:
-                    # No need to update ~/.ssh/authorized_keys when running shim locally
-                    user_ssh_key = ""
-                success = await common_utils.run_async(
-                    _process_provisioning_with_shim,
-                    server_ssh_private_keys,
-                    job_provisioning_data,
-                    None,
-                    run,
-                    job_model,
-                    job_provisioning_data,
-                    volumes,
-                    job.job_spec.registry_auth,
-                    public_keys,
-                    ssh_user,
-                    user_ssh_key,
-                )
-            else:
-                logger.debug(
-                    "%s: process provisioning job without shim, age=%s",
+            success = await common_utils.run_async(
+                _submit_job_to_runner,
+                server_ssh_private_keys,
+                job_provisioning_data,
+                None,
+                run,
+                job_model,
+                job,
+                cluster_info,
+                code,
+                file_archives,
+                secrets,
+                repo_creds,
+                success_if_not_available=False,
+            )
+
+        if not success:
+            # check timeout
+            if job_submission.age > get_provisioning_timeout(
+                backend_type=job_provisioning_data.get_base_backend(),
+                instance_type_name=job_provisioning_data.instance_type.name,
+            ):
+                logger.warning(
+                    "%s: failed because runner has not become available in time, age=%s",
                     fmt(job_model),
                     job_submission.age,
                 )
-                # FIXME: downloading file archives and code here is a waste of time if
-                # the runner is not ready yet
-                file_archives = await _get_job_file_archives(
-                    session=session,
-                    archive_mappings=job.job_spec.file_archives,
-                    user=run_model.user,
-                )
-                code = await _get_job_code(
-                    session=session,
-                    project=project,
-                    repo=repo_model,
-                    code_hash=_get_repo_code_hash(run, job),
-                )
-
-                success = await common_utils.run_async(
-                    _submit_job_to_runner,
-                    server_ssh_private_keys,
-                    job_provisioning_data,
-                    None,
-                    run,
-                    job_model,
-                    job,
-                    cluster_info,
-                    code,
-                    file_archives,
-                    secrets,
-                    repo_creds,
-                    success_if_not_available=False,
-                )
-
-            if not success:
-                # check timeout
-                if job_submission.age > get_provisioning_timeout(
-                    backend_type=job_provisioning_data.get_base_backend(),
-                    instance_type_name=job_provisioning_data.instance_type.name,
-                ):
-                    logger.warning(
-                        "%s: failed because runner has not become available in time, age=%s",
-                        fmt(job_model),
-                        job_submission.age,
-                    )
-                    job_model.status = JobStatus.TERMINATING
-                    job_model.termination_reason = (
-                        JobTerminationReason.WAITING_RUNNER_LIMIT_EXCEEDED
-                    )
-                    # instance will be emptied by process_terminating_jobs
+                job_model.status = JobStatus.TERMINATING
+                job_model.termination_reason = JobTerminationReason.WAITING_RUNNER_LIMIT_EXCEEDED
+                # instance will be emptied by process_terminating_jobs
 
     else:  # fails are not acceptable
         if initial_status == JobStatus.PULLING:

src/tests/_internal/server/background/tasks/test_process_instances.py

@@ -224,7 +224,7 @@ async def test_check_shim_stop_termination_deadline(self, test_db, session: Asyn
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_check_shim_terminate_instance_by_dedaline(self, test_db, session: AsyncSession):
+    async def test_check_shim_terminate_instance_by_deadline(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         instance = await create_instance(
             session=session,
@@ -306,6 +306,7 @@ async def test_check_shim_process_ureachable_state(
         ) as healthcheck:
             healthcheck.return_value = HealthStatus(healthy=True, reason="OK")
             await process_instances()
+            healthcheck.assert_called()
 
         await session.refresh(instance)
 
@@ -329,7 +330,7 @@ async def test_terminate_by_idle_timeout(self, test_db, session: AsyncSession):
         await process_instances()
         await session.refresh(instance)
         assert instance is not None
-        assert instance.status == InstanceStatus.TERMINATED
+        assert instance.status == InstanceStatus.TERMINATING
         assert instance.termination_reason == "Idle timeout"