dstack 0.19.19__py3-none-any.whl → 0.19.21__py3-none-any.whl

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

@@ -9,13 +9,9 @@ from paramiko.ssh_exception import PasswordRequiredException
 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 import (
-    BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
-    BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT,
-)
 from dstack._internal.core.backends.base.compute import (
     ComputeWithCreateInstanceSupport,
     ComputeWithPlacementGroupSupport,
@@ -27,6 +23,10 @@ from dstack._internal.core.backends.base.compute import (
     get_shim_env,
     get_shim_pre_start_commands,
 )
+from dstack._internal.core.backends.features import (
+    BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
+    BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT,
+)
 from dstack._internal.core.backends.remote.provisioning import (
     detect_cpu_arch,
     get_host_info,
@@ -78,6 +78,7 @@ from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     FleetModel,
     InstanceModel,
+    JobModel,
     PlacementGroupModel,
     ProjectModel,
 )
@@ -104,7 +105,12 @@ from dstack._internal.server.services.placement import (
 from dstack._internal.server.services.runner import client as runner_client
 from dstack._internal.server.services.runner.client import HealthStatus
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
-from dstack._internal.utils.common import get_current_datetime, run_async
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.utils.common import (
+    get_current_datetime,
+    get_or_error,
+    run_async,
+)
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.network import get_ip_from_network, is_ip_among_addresses
 from dstack._internal.utils.ssh import (
@@ -131,6 +137,7 @@ async def process_instances(batch_size: int = 1):
     await asyncio.gather(*tasks)
 
 
+@sentry_utils.instrument_background_task
 async def _process_next_instance():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(InstanceModel.__tablename__)
     async with get_session_ctx() as session:
@@ -149,12 +156,13 @@ async def _process_next_instance():
                     ),
                     InstanceModel.id.not_in(lockset),
                     InstanceModel.last_processed_at
-                    < get_current_datetime().replace(tzinfo=None) - MIN_PROCESSING_INTERVAL,
+                    < 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:
@@ -168,23 +176,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)
@@ -198,7 +205,9 @@ 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)
 
@@ -206,7 +215,13 @@ async def _process_instance(session: AsyncSession, instance: InstanceModel):
     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)
@@ -222,6 +237,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:
@@ -461,7 +478,7 @@ def _deploy_instance(
 
 async def _create_instance(session: AsyncSession, instance: InstanceModel) -> None:
     if instance.last_retry_at is not None:
-        last_retry = instance.last_retry_at.replace(tzinfo=datetime.timezone.utc)
+        last_retry = instance.last_retry_at
         if get_current_datetime() < last_retry + timedelta(minutes=1):
             return
 
@@ -700,7 +717,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
@@ -719,12 +736,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.project_id)
+            .options(joinedload(ProjectModel.backends))
+        )
+        project = res.unique().scalar_one()
         await _wait_for_instance_provisioning_data(
-            project=instance.project,
+            project=project,
             instance=instance,
             job_provisioning_data=job_provisioning_data,
         )
@@ -801,7 +822,7 @@ async def _check_instance(instance: InstanceModel) -> None:
             instance.name,
             extra={"instance_name": instance.name},
         )
-        deadline = instance.termination_deadline.replace(tzinfo=datetime.timezone.utc)
+        deadline = instance.termination_deadline
         if get_current_datetime() > deadline:
             instance.status = InstanceStatus.TERMINATING
             instance.termination_reason = "Termination deadline"
@@ -956,18 +977,12 @@ async def _terminate(instance: InstanceModel) -> None:
 
 def _next_termination_retry_at(instance: InstanceModel) -> datetime.datetime:
     assert instance.last_termination_retry_at is not None
-    return (
-        instance.last_termination_retry_at.replace(tzinfo=datetime.timezone.utc)
-        + TERMINATION_RETRY_TIMEOUT
-    )
+    return instance.last_termination_retry_at + TERMINATION_RETRY_TIMEOUT
 
 
 def _get_termination_deadline(instance: InstanceModel) -> datetime.datetime:
     assert instance.first_termination_retry_at is not None
-    return (
-        instance.first_termination_retry_at.replace(tzinfo=datetime.timezone.utc)
-        + TERMINATION_RETRY_MAX_DURATION
-    )
+    return instance.first_termination_retry_at + TERMINATION_RETRY_MAX_DURATION
 
 
 def _need_to_wait_fleet_provisioning(instance: InstanceModel) -> bool:
@@ -1102,27 +1117,26 @@ async def _create_placement_group(
 
 
 def _get_instance_idle_duration(instance: InstanceModel) -> datetime.timedelta:
-    last_time = instance.created_at.replace(tzinfo=datetime.timezone.utc)
+    last_time = instance.created_at
     if instance.last_job_processed_at is not None:
-        last_time = instance.last_job_processed_at.replace(tzinfo=datetime.timezone.utc)
+        last_time = instance.last_job_processed_at
     return get_current_datetime() - last_time
 
 
 def _get_retry_duration_deadline(instance: InstanceModel, retry: Retry) -> datetime.datetime:
-    return instance.created_at.replace(tzinfo=datetime.timezone.utc) + timedelta(
-        seconds=retry.duration
-    )
+    return instance.created_at + timedelta(seconds=retry.duration)
 
 
 def _get_provisioning_deadline(
     instance: InstanceModel,
     job_provisioning_data: JobProvisioningData,
 ) -> datetime.datetime:
+    assert instance.started_at is not None
     timeout_interval = get_provisioning_timeout(
         backend_type=job_provisioning_data.get_base_backend(),
         instance_type_name=job_provisioning_data.instance_type.name,
     )
-    return instance.started_at.replace(tzinfo=datetime.timezone.utc) + timeout_interval
+    return instance.started_at + timeout_interval
 
 
 def _ssh_keys_to_pkeys(ssh_keys: list[SSHKey]) -> list[PKey]:

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

@@ -9,12 +9,13 @@ from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT
 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
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
+from dstack._internal.server.utils import sentry_utils
 from dstack._internal.utils.common import batched, get_current_datetime, get_or_error, run_async
 from dstack._internal.utils.logging import get_logger
 
@@ -26,12 +27,17 @@ BATCH_SIZE = 10
 MIN_COLLECT_INTERVAL_SECONDS = 9
 
 
+@sentry_utils.instrument_background_task
 async def collect_metrics():
     async with get_session_ctx() as session:
         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)
         )
@@ -41,6 +47,7 @@ async def collect_metrics():
         await _collect_jobs_metrics(batch)
 
 
+@sentry_utils.instrument_background_task
 async def delete_metrics():
     now_timestamp_micro = int(get_current_datetime().timestamp() * 1_000_000)
     running_timestamp_micro_cutoff = (

dstack/_internal/server/background/tasks/process_placement_groups.py

@@ -12,12 +12,14 @@ from dstack._internal.server.models import PlacementGroupModel, ProjectModel
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.placement import placement_group_model_to_placement_group
+from dstack._internal.server.utils import sentry_utils
 from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
 
 
+@sentry_utils.instrument_background_task
 async def process_placement_groups():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(
         PlacementGroupModel.__tablename__

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

@@ -9,11 +9,17 @@ from sqlalchemy.orm import joinedload
 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
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
+from dstack._internal.server.utils import sentry_utils
 from dstack._internal.server.utils.common import gather_map_async
 from dstack._internal.utils.common import batched, get_current_datetime, get_or_error, run_async
 from dstack._internal.utils.logging import get_logger
@@ -29,6 +35,7 @@ MIN_COLLECT_INTERVAL_SECONDS = 9
 METRICS_TTL_SECONDS = 600
 
 
+@sentry_utils.instrument_background_task
 async def collect_prometheus_metrics():
     now = get_current_datetime()
     cutoff = now - timedelta(seconds=MIN_COLLECT_INTERVAL_SECONDS)
@@ -43,7 +50,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)
         )
@@ -52,6 +63,7 @@ async def collect_prometheus_metrics():
         await _collect_jobs_metrics(batch, now)
 
 
+@sentry_utils.instrument_background_task
 async def delete_prometheus_metrics():
     now = get_current_datetime()
     cutoff = now - timedelta(seconds=METRICS_TTL_SECONDS)

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

@@ -2,12 +2,12 @@ import asyncio
 import re
 import uuid
 from collections.abc import Iterable
-from datetime import timedelta, timezone
+from datetime import timedelta
 from typing import Dict, List, Optional
 
 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
@@ -73,6 +73,7 @@ from dstack._internal.server.services.runs import (
 )
 from dstack._internal.server.services.secrets import get_project_secrets_mapping
 from dstack._internal.server.services.storage import get_default_storage
+from dstack._internal.server.utils import sentry_utils
 from dstack._internal.utils import common as common_utils
 from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator
 from dstack._internal.utils.logging import get_logger
@@ -94,6 +95,7 @@ async def process_running_jobs(batch_size: int = 1):
     await asyncio.gather(*tasks)
 
 
+@sentry_utils.instrument_background_task
 async def _process_next_running_job():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
     async with get_session_ctx() as session:
@@ -108,9 +110,9 @@ async def _process_next_running_job():
                     RunModel.status.not_in([RunStatus.TERMINATING]),
                     JobModel.id.not_in(lockset),
                     JobModel.last_processed_at
-                    < common_utils.get_current_datetime().replace(tzinfo=None)
-                    - MIN_PROCESSING_INTERVAL,
+                    < 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(
@@ -133,7 +135,6 @@ async def _process_next_running_job():
 
 async def _process_running_job(session: AsyncSession, job_model: JobModel):
     # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(
         select(JobModel)
         .where(JobModel.id == job_model.id)
@@ -144,7 +145,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))
@@ -160,143 +161,147 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
         job_model.status = JobStatus.TERMINATING
         job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
         job_model.last_processed_at = common_utils.get_current_datetime()
+        await session.commit()
         return
 
     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
-        ):
+    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,
+            job_provisioning_data=job_provisioning_data,
+            job_runtime_data=job_submission.job_runtime_data,
+        )
+
+        volumes = await get_job_attached_volumes(
+            session=session,
+            project=project,
+            run_spec=run.run_spec,
+            job_num=job.job_spec.job_num,
+            job_provisioning_data=job_provisioning_data,
+        )
+
+        repo_creds_model = await get_repo_creds(
+            session=session, repo=repo_model, user=run_model.user
+        )
+        repo_creds = repo_model_to_repo_head_with_creds(repo_model, repo_creds_model).repo_creds
+
+        secrets = await get_project_secrets_mapping(session=session, project=project)
+        try:
+            _interpolate_secrets(secrets, job.job_spec)
+        except InterpolatorError as e:
+            logger.info("%s: terminating due to secrets interpolation error", fmt(job_model))
+            job_model.status = JobStatus.TERMINATING
+            job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+            job_model.termination_reason_message = e.args[0]
             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,
-        job_provisioning_data=job_provisioning_data,
-        job_runtime_data=job_submission.job_runtime_data,
-    )
-
-    volumes = await get_job_attached_volumes(
-        session=session,
-        project=project,
-        run_spec=run.run_spec,
-        job_num=job.job_spec.job_num,
-        job_provisioning_data=job_provisioning_data,
-    )
-
     server_ssh_private_keys = get_instance_ssh_private_keys(
         common_utils.get_or_error(job_model.instance)
     )
 
-    secrets = await get_project_secrets_mapping(session=session, project=project)
-
-    try:
-        _interpolate_secrets(secrets, job.job_spec)
-    except InterpolatorError as e:
-        logger.info("%s: terminating due to secrets interpolation error", fmt(job_model))
-        job_model.status = JobStatus.TERMINATING
-        job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
-        job_model.termination_reason_message = e.args[0]
-        job_model.last_processed_at = common_utils.get_current_datetime()
-        return
-
-    repo_creds_model = await get_repo_creds(session=session, repo=repo_model, user=run_model.user)
-    repo_creds = repo_model_to_repo_head_with_creds(repo_model, repo_creds_model).repo_creds
-
-    initial_status = job_model.status
     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:
@@ -801,7 +806,7 @@ def _should_terminate_job_due_to_disconnect(job_model: JobModel) -> bool:
         return False
     return (
         common_utils.get_current_datetime()
-        > job_model.disconnected_at.replace(tzinfo=timezone.utc) + JOB_DISCONNECTED_RETRY_TIMEOUT
+        > job_model.disconnected_at + JOB_DISCONNECTED_RETRY_TIMEOUT
     )