dstack 0.19.20__py3-none-any.whl → 0.19.22__py3-none-any.whl

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

@@ -1,21 +1,18 @@
 import asyncio
 import datetime
+import logging
 from datetime import timedelta
-from typing import Any, Dict, List, Optional, Tuple, cast
+from typing import Any, Dict, List, Optional, cast
 
 import requests
 from paramiko.pkey import PKey
 from paramiko.ssh_exception import PasswordRequiredException
 from pydantic import ValidationError
-from sqlalchemy import select
+from sqlalchemy import delete, func, 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,17 +24,9 @@ from dstack._internal.core.backends.base.compute import (
     get_shim_env,
     get_shim_pre_start_commands,
 )
-from dstack._internal.core.backends.remote.provisioning import (
-    detect_cpu_arch,
-    get_host_info,
-    get_paramiko_connection,
-    get_shim_healthcheck,
-    host_info_to_instance_type,
-    remove_dstack_runner_if_exists,
-    remove_host_info_if_exists,
-    run_pre_start_commands,
-    run_shim_as_systemd_service,
-    upload_envs,
+from dstack._internal.core.backends.features import (
+    BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
+    BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT,
 )
 from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
 
@@ -77,11 +66,14 @@ from dstack._internal.server.background.tasks.common import get_provisioning_tim
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     FleetModel,
+    InstanceHealthCheckModel,
     InstanceModel,
+    JobModel,
     PlacementGroupModel,
     ProjectModel,
 )
-from dstack._internal.server.schemas.runner import HealthcheckResponse
+from dstack._internal.server.schemas.instances import InstanceCheck
+from dstack._internal.server.schemas.runner import HealthcheckResponse, InstanceHealthResponse
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services.fleets import (
     fleet_model_to_fleet,
@@ -102,9 +94,25 @@ from dstack._internal.server.services.placement import (
     schedule_fleet_placement_groups_deletion,
 )
 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.server.utils.provisioning import (
+    detect_cpu_arch,
+    get_host_info,
+    get_paramiko_connection,
+    get_shim_healthcheck,
+    host_info_to_instance_type,
+    remove_dstack_runner_if_exists,
+    remove_host_info_if_exists,
+    run_pre_start_commands,
+    run_shim_as_systemd_service,
+    upload_envs,
+)
+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 +139,18 @@ async def process_instances(batch_size: int = 1):
     await asyncio.gather(*tasks)
 
 
+@sentry_utils.instrument_background_task
+async def delete_instance_health_checks():
+    now = get_current_datetime()
+    cutoff = now - timedelta(seconds=server_settings.SERVER_INSTANCE_HEALTH_TTL_SECONDS)
+    async with get_session_ctx() as session:
+        await session.execute(
+            delete(InstanceHealthCheckModel).where(InstanceHealthCheckModel.collected_at < cutoff)
+        )
+        await session.commit()
+
+
+@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 +169,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 +189,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 +218,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 +228,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 +250,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:
@@ -398,10 +428,10 @@ async def _add_remote(instance: InstanceModel) -> None:
 
 def _deploy_instance(
     remote_details: RemoteConnectionInfo,
-    pkeys: List[PKey],
+    pkeys: list[PKey],
     ssh_proxy_pkeys: Optional[list[PKey]],
-    authorized_keys: List[str],
-) -> Tuple[HealthStatus, Dict[str, Any], GoArchType]:
+    authorized_keys: list[str],
+) -> tuple[InstanceCheck, dict[str, Any], GoArchType]:
     with get_paramiko_connection(
         remote_details.ssh_user,
         remote_details.host,
@@ -449,19 +479,19 @@ def _deploy_instance(
         host_info = get_host_info(client, dstack_working_dir)
         logger.debug("Received a host_info %s", host_info)
 
-        raw_health = get_shim_healthcheck(client)
+        healthcheck_out = get_shim_healthcheck(client)
         try:
-            health_response = HealthcheckResponse.__response__.parse_raw(raw_health)
+            healthcheck = HealthcheckResponse.__response__.parse_raw(healthcheck_out)
         except ValueError as e:
-            raise ProvisioningError("Cannot read HealthcheckResponse") from e
-        health = runner_client.health_response_to_health_status(health_response)
+            raise ProvisioningError(f"Cannot parse HealthcheckResponse: {e}") from e
+        instance_check = runner_client.healthcheck_response_to_instance_check(healthcheck)
 
-        return health, host_info, arch
+        return instance_check, host_info, arch
 
 
 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 +730,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 +749,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,
         )
@@ -737,29 +771,65 @@ async def _check_instance(instance: InstanceModel) -> None:
 
     ssh_private_keys = get_instance_ssh_private_keys(instance)
 
+    health_check_cutoff = get_current_datetime() - timedelta(
+        seconds=server_settings.SERVER_INSTANCE_HEALTH_MIN_COLLECT_INTERVAL_SECONDS
+    )
+    res = await session.execute(
+        select(func.count(1)).where(
+            InstanceHealthCheckModel.instance_id == instance.id,
+            InstanceHealthCheckModel.collected_at > health_check_cutoff,
+        )
+    )
+    check_instance_health = res.scalar_one() == 0
+
     # May return False if fails to establish ssh connection
-    health_status_response = await run_async(
-        _instance_healthcheck,
+    instance_check = await run_async(
+        _check_instance_inner,
         ssh_private_keys,
         job_provisioning_data,
         None,
+        check_instance_health=check_instance_health,
     )
-    if isinstance(health_status_response, bool) or health_status_response is None:
-        health_status = HealthStatus(healthy=False, reason="SSH or tunnel error")
-    else:
-        health_status = health_status_response
+    if instance_check is False:
+        instance_check = InstanceCheck(reachable=False, message="SSH or tunnel error")
 
-    logger.debug(
-        "Check instance %s status. shim health: %s",
+    if instance_check.reachable and check_instance_health:
+        health_status = instance_check.get_health_status()
+    else:
+        # Keep previous health status
+        health_status = instance.health
+
+    loglevel = logging.DEBUG
+    if not instance_check.reachable and instance.status.is_available():
+        loglevel = logging.WARNING
+    elif check_instance_health and not health_status.is_healthy():
+        loglevel = logging.WARNING
+    logger.log(
+        loglevel,
+        "Instance %s check: reachable=%s health_status=%s message=%r",
         instance.name,
-        health_status,
-        extra={"instance_name": instance.name, "shim_health": health_status},
+        instance_check.reachable,
+        health_status.name,
+        instance_check.message,
+        extra={"instance_name": instance.name, "health_status": health_status},
     )
 
-    if health_status.healthy:
+    if instance_check.has_health_checks():
+        # ensured by has_health_checks()
+        assert instance_check.health_response is not None
+        health_check_model = InstanceHealthCheckModel(
+            instance_id=instance.id,
+            collected_at=get_current_datetime(),
+            status=health_status,
+            response=instance_check.health_response.json(),
+        )
+        session.add(health_check_model)
+
+    instance.health = health_status
+    instance.unreachable = not instance_check.reachable
+
+    if instance_check.reachable:
         instance.termination_deadline = None
-        instance.health_status = None
-        instance.unreachable = False
 
         if instance.status == InstanceStatus.PROVISIONING:
             instance.status = InstanceStatus.IDLE if not instance.jobs else InstanceStatus.BUSY
@@ -777,9 +847,6 @@ async def _check_instance(instance: InstanceModel) -> None:
     if instance.termination_deadline is None:
         instance.termination_deadline = get_current_datetime() + TERMINATION_DEADLINE_OFFSET
 
-    instance.health_status = health_status.reason
-    instance.unreachable = True
-
     if instance.status == InstanceStatus.PROVISIONING and instance.started_at is not None:
         provisioning_deadline = _get_provisioning_deadline(
             instance=instance,
@@ -795,13 +862,8 @@ async def _check_instance(instance: InstanceModel) -> None:
                     "instance_status": InstanceStatus.TERMINATING.value,
                 },
             )
-    elif instance.status in (InstanceStatus.IDLE, InstanceStatus.BUSY):
-        logger.warning(
-            "Instance %s shim is not available",
-            instance.name,
-            extra={"instance_name": instance.name},
-        )
-        deadline = instance.termination_deadline.replace(tzinfo=datetime.timezone.utc)
+    elif instance.status.is_available():
+        deadline = instance.termination_deadline
         if get_current_datetime() > deadline:
             instance.status = InstanceStatus.TERMINATING
             instance.termination_reason = "Termination deadline"
@@ -871,20 +933,30 @@ async def _wait_for_instance_provisioning_data(
 
 
 @runner_ssh_tunnel(ports=[DSTACK_SHIM_HTTP_PORT], retries=1)
-def _instance_healthcheck(ports: Dict[int, int]) -> HealthStatus:
+def _check_instance_inner(
+    ports: Dict[int, int], *, check_instance_health: bool = False
+) -> InstanceCheck:
+    instance_health_response: Optional[InstanceHealthResponse] = None
     shim_client = runner_client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
+    method = shim_client.healthcheck
     try:
-        resp = shim_client.healthcheck(unmask_exeptions=True)
-        if resp is None:
-            return HealthStatus(healthy=False, reason="Unknown reason")
-        return runner_client.health_response_to_health_status(resp)
+        healthcheck_response = method(unmask_exceptions=True)
+        if check_instance_health:
+            method = shim_client.get_instance_health
+            instance_health_response = method()
     except requests.RequestException as e:
-        return HealthStatus(healthy=False, reason=f"Can't request shim: {e}")
+        template = "shim.%s(): request error: %s"
+        args = (method.__func__.__name__, e)
+        logger.debug(template, *args)
+        return InstanceCheck(reachable=False, message=template % args)
     except Exception as e:
-        logger.exception("Unknown exception from shim.healthcheck: %s", e)
-        return HealthStatus(
-            healthy=False, reason=f"Unknown exception ({e.__class__.__name__}): {e}"
-        )
+        template = "shim.%s(): unexpected exception %s: %s"
+        args = (method.__func__.__name__, e.__class__.__name__, e)
+        logger.exception(template, *args)
+        return InstanceCheck(reachable=False, message=template % args)
+    return runner_client.healthcheck_response_to_instance_check(
+        healthcheck_response, instance_health_response
+    )
 
 
 async def _terminate(instance: InstanceModel) -> None:
@@ -956,18 +1028,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 +1168,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_probes.py

@@ -0,0 +1,164 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from functools import partial
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import httpx
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from httpx import AsyncClient, AsyncHTTPTransport
+from sqlalchemy import select, update
+from sqlalchemy.orm import joinedload
+
+from dstack._internal.core.errors import SSHError
+from dstack._internal.core.models.runs import JobSpec, JobStatus, ProbeSpec
+from dstack._internal.core.services.ssh.tunnel import (
+    SSH_DEFAULT_OPTIONS,
+    IPSocket,
+    SocketPair,
+    UnixSocket,
+)
+from dstack._internal.server.db import get_db, get_session_ctx
+from dstack._internal.server.models import InstanceModel, JobModel, ProbeModel
+from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.logging import fmt
+from dstack._internal.server.services.ssh import container_ssh_tunnel
+from dstack._internal.utils.common import get_current_datetime, get_or_error
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+BATCH_SIZE = 100
+SSH_CONNECT_TIMEOUT = timedelta(seconds=10)
+PROCESSING_OVERHEAD_TIMEOUT = timedelta(minutes=1)
+PROBES_SCHEDULER = AsyncIOScheduler()
+
+
+async def process_probes():
+    probe_lock, probe_lockset = get_locker(get_db().dialect_name).get_lockset(
+        ProbeModel.__tablename__
+    )
+    async with get_session_ctx() as session:
+        async with probe_lock:
+            res = await session.execute(
+                select(ProbeModel.id)
+                .where(ProbeModel.id.not_in(probe_lockset))
+                .where(ProbeModel.active == True)
+                .where(ProbeModel.due <= get_current_datetime())
+                .order_by(ProbeModel.due.asc())
+                .limit(BATCH_SIZE)
+                .with_for_update(skip_locked=True, key_share=True)
+            )
+            probe_ids = res.unique().scalars().all()
+            probe_lockset.update(probe_ids)
+
+        try:
+            # Refetch to load all attributes.
+            # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
+            res = await session.execute(
+                select(ProbeModel)
+                .where(ProbeModel.id.in_(probe_ids))
+                .options(
+                    joinedload(ProbeModel.job)
+                    .joinedload(JobModel.instance)
+                    .joinedload(InstanceModel.project)
+                )
+                .options(joinedload(ProbeModel.job))
+                .execution_options(populate_existing=True)
+            )
+            probes = res.unique().scalars().all()
+            for probe in probes:
+                if probe.job.status != JobStatus.RUNNING:
+                    probe.active = False
+                else:
+                    job_spec: JobSpec = JobSpec.__response__.parse_raw(probe.job.job_spec_data)
+                    probe_spec = job_spec.probes[probe.probe_num]
+                    # Schedule the next probe execution in case this execution is interrupted
+                    probe.due = get_current_datetime() + _get_probe_async_processing_timeout(
+                        probe_spec
+                    )
+                    # Execute the probe asynchronously outside of the DB session
+                    PROBES_SCHEDULER.add_job(partial(_process_probe_async, probe, probe_spec))
+            await session.commit()
+        finally:
+            probe_lockset.difference_update(probe_ids)
+
+
+async def _process_probe_async(probe: ProbeModel, probe_spec: ProbeSpec) -> None:
+    start = get_current_datetime()
+    logger.debug("%s: processing probe", fmt(probe))
+    success = await _execute_probe(probe, probe_spec)
+
+    async with get_session_ctx() as session:
+        async with get_locker(get_db().dialect_name).lock_ctx(
+            ProbeModel.__tablename__, [probe.id]
+        ):
+            await session.execute(
+                update(ProbeModel)
+                .where(ProbeModel.id == probe.id)
+                .values(
+                    success_streak=0 if not success else ProbeModel.success_streak + 1,
+                    due=get_current_datetime() + timedelta(seconds=probe_spec.interval),
+                )
+            )
+    logger.debug(
+        "%s: probe processing took %ss",
+        fmt(probe),
+        (get_current_datetime() - start).total_seconds(),
+    )
+
+
+async def _execute_probe(probe: ProbeModel, probe_spec: ProbeSpec) -> bool:
+    """
+    Returns:
+        Whether probe execution was successful.
+    """
+
+    try:
+        async with _get_service_replica_client(probe.job) as client:
+            resp = await client.request(
+                method=probe_spec.method,
+                url="http://dstack" + probe_spec.url,
+                headers=[(h.name, h.value) for h in probe_spec.headers],
+                data=probe_spec.body,
+                timeout=probe_spec.timeout,
+                follow_redirects=False,
+            )
+            logger.debug("%s: probe status code: %s", fmt(probe), resp.status_code)
+            return resp.is_success
+    except (SSHError, httpx.RequestError) as e:
+        logger.debug("%s: probe failed: %r", fmt(probe), e)
+        return False
+
+
+def _get_probe_async_processing_timeout(probe_spec: ProbeSpec) -> timedelta:
+    return (
+        timedelta(seconds=probe_spec.timeout)
+        + SSH_CONNECT_TIMEOUT
+        + PROCESSING_OVERHEAD_TIMEOUT  # slow db queries and other unforeseen conditions
+    )
+
+
+@asynccontextmanager
+async def _get_service_replica_client(job: JobModel) -> AsyncGenerator[AsyncClient, None]:
+    options = {
+        **SSH_DEFAULT_OPTIONS,
+        "ConnectTimeout": str(int(SSH_CONNECT_TIMEOUT.total_seconds())),
+    }
+    job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
+    with TemporaryDirectory() as temp_dir:
+        app_socket_path = (Path(temp_dir) / "replica.sock").absolute()
+        async with container_ssh_tunnel(
+            job=job,
+            forwarded_sockets=[
+                SocketPair(
+                    remote=IPSocket("localhost", get_or_error(job_spec.service_port)),
+                    local=UnixSocket(app_socket_path),
+                ),
+            ],
+            options=options,
+        ):
+            async with AsyncClient(
+                transport=AsyncHTTPTransport(uds=str(app_socket_path))
+            ) as client:
+                yield client