Only register service replicas after probes pass

To avoid service disruptions during rolling
deployments, only register a newly started replica
to receive requests after all its probes pass, not
immediately after it becomes `running`.

Author: jvstme

src/dstack/_internal/proxy/gateway/services/registry.py

@@ -152,6 +152,8 @@ async def register_replica(
             )
 
         if old_service.find_replica(replica_id) is not None:
+            # NOTE: as of 0.19.25, the dstack server relies on the exact text of this error.
+            # See dstack._internal.server.services.services.register_replica
             raise ProxyError(f"Replica {replica_id} already exists in service {old_service.fmt()}")
 
         service = old_service.with_replicas(old_service.replicas + (replica,))

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

@@ -32,6 +32,7 @@
     JobSpec,
     JobStatus,
     JobTerminationReason,
+    ProbeSpec,
     Run,
     RunSpec,
     RunStatus,
@@ -70,6 +71,7 @@
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
 from dstack._internal.server.services.runs import (
+    is_job_ready,
     run_model_to_run,
 )
 from dstack._internal.server.services.secrets import get_project_secrets_mapping
@@ -140,6 +142,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
         select(JobModel)
         .where(JobModel.id == job_model.id)
         .options(joinedload(JobModel.instance).joinedload(InstanceModel.project))
+        .options(joinedload(JobModel.probes).load_only(ProbeModel.success_streak))
         .execution_options(populate_existing=True)
     )
     job_model = res.unique().scalar_one()
@@ -382,52 +385,21 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                         job_submission.age,
                     )
 
-    if (
-        initial_status != job_model.status
-        and job_model.status == JobStatus.RUNNING
-        and job_model.job_num == 0  # gateway connects only to the first node
-        and run.run_spec.configuration.type == "service"
-    ):
-        ssh_head_proxy: Optional[SSHConnectionParams] = None
-        ssh_head_proxy_private_key: Optional[str] = None
-        instance = common_utils.get_or_error(job_model.instance)
-        if instance.remote_connection_info is not None:
-            rci = RemoteConnectionInfo.__response__.parse_raw(instance.remote_connection_info)
-            if rci.ssh_proxy is not None:
-                ssh_head_proxy = rci.ssh_proxy
-                ssh_head_proxy_keys = common_utils.get_or_error(rci.ssh_proxy_keys)
-                ssh_head_proxy_private_key = ssh_head_proxy_keys[0].private
-        try:
-            await services.register_replica(
-                session,
-                run_model.gateway_id,
-                run,
-                job_model,
-                ssh_head_proxy,
-                ssh_head_proxy_private_key,
-            )
-        except GatewayError as e:
-            logger.warning(
-                "%s: failed to register service replica: %s, age=%s",
-                fmt(job_model),
-                e,
-                job_submission.age,
-            )
-            job_model.status = JobStatus.TERMINATING
-            job_model.termination_reason = JobTerminationReason.GATEWAY_ERROR
-        else:
-            for probe_num in range(len(job.job_spec.probes)):
-                session.add(
-                    ProbeModel(
-                        name=f"{job_model.job_name}-{probe_num}",
-                        job=job_model,
-                        probe_num=probe_num,
-                        due=common_utils.get_current_datetime(),
-                        success_streak=0,
-                        active=True,
-                    )
+    if initial_status != job_model.status and job_model.status == JobStatus.RUNNING:
+        job_model.probes = []
+        for probe_num in range(len(job.job_spec.probes)):
+            job_model.probes.append(
+                ProbeModel(
+                    name=f"{job_model.job_name}-{probe_num}",
+                    probe_num=probe_num,
+                    due=common_utils.get_current_datetime(),
+                    success_streak=0,
+                    active=True,
                 )
+            )
 
+    if job_model.status == JobStatus.RUNNING:
+        await _maybe_register_replica(session, run_model, run, job_model, job.job_spec.probes)
     if job_model.status == JobStatus.RUNNING:
         await _check_gpu_utilization(session, job_model, job)
 
@@ -822,6 +794,55 @@ def _should_terminate_job_due_to_disconnect(job_model: JobModel) -> bool:
     )
 
 
+async def _maybe_register_replica(
+    session: AsyncSession,
+    run_model: RunModel,
+    run: Run,
+    job_model: JobModel,
+    probe_specs: Iterable[ProbeSpec],
+) -> None:
+    """
+    Register the replica represented by this job to receive service requests if it is ready.
+    """
+
+    if (
+        run.run_spec.configuration.type != "service"
+        or job_model.registered
+        or job_model.job_num != 0  # only the first job in the replica receives service requests
+        or not is_job_ready(job_model.probes, probe_specs)
+    ):
+        return
+
+    ssh_head_proxy: Optional[SSHConnectionParams] = None
+    ssh_head_proxy_private_key: Optional[str] = None
+    instance = common_utils.get_or_error(job_model.instance)
+    if instance.remote_connection_info is not None:
+        rci: RemoteConnectionInfo = RemoteConnectionInfo.__response__.parse_raw(
+            instance.remote_connection_info
+        )
+        if rci.ssh_proxy is not None:
+            ssh_head_proxy = rci.ssh_proxy
+            ssh_head_proxy_keys = common_utils.get_or_error(rci.ssh_proxy_keys)
+            ssh_head_proxy_private_key = ssh_head_proxy_keys[0].private
+    try:
+        await services.register_replica(
+            session,
+            run_model.gateway_id,
+            run,
+            job_model,
+            ssh_head_proxy,
+            ssh_head_proxy_private_key,
+        )
+    except GatewayError as e:
+        logger.warning(
+            "%s: failed to register service replica: %s",
+            fmt(job_model),
+            e,
+        )
+        job_model.status = JobStatus.TERMINATING
+        job_model.termination_reason = JobTerminationReason.GATEWAY_ERROR
+
+
 async def _check_gpu_utilization(session: AsyncSession, job_model: JobModel, job: Job) -> None:
     policy = job.job_spec.utilization_policy
     if policy is None:

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

@@ -23,7 +23,6 @@
 from dstack._internal.server.models import (
     InstanceModel,
     JobModel,
-    ProbeModel,
     ProjectModel,
     RunModel,
     UserModel,
@@ -37,7 +36,7 @@
 from dstack._internal.server.services.prometheus.client_metrics import run_metrics
 from dstack._internal.server.services.runs import (
     fmt,
-    is_replica_ready,
+    is_replica_registered,
     process_terminating_run,
     retry_run_replica_jobs,
     run_model_to_run,
@@ -151,11 +150,6 @@ async def _process_run(session: AsyncSession, run_model: RunModel):
             .joinedload(JobModel.instance)
             .load_only(InstanceModel.fleet_id)
         )
-        .options(
-            selectinload(RunModel.jobs)
-            .joinedload(JobModel.probes)
-            .load_only(ProbeModel.success_streak)
-        )
         .execution_options(populate_existing=True)
     )
     run_model = res.unique().scalar_one()
@@ -465,6 +459,9 @@ async def _handle_run_replicas(
         run_spec=run_spec,
     )
     if _has_out_of_date_replicas(run_model):
+        assert run_spec.configuration.type == "service", (
+            "Rolling deployment is only supported for services"
+        )
         non_terminated_replica_count = len(
             {j.replica_num for j in run_model.jobs if not j.status.is_finished()}
         )
@@ -479,22 +476,24 @@ async def _handle_run_replicas(
             )
 
         replicas_to_stop_count = 0
-        # stop any out-of-date replicas that are not ready
+        # stop any out-of-date replicas that are not registered
         replicas_to_stop_count += sum(
             any(j.deployment_num < run_model.deployment_num for j in jobs)
             and any(
                 j.status not in [JobStatus.TERMINATING] + JobStatus.finished_statuses()
                 for j in jobs
             )
-            and not is_replica_ready(jobs)
+            and not is_replica_registered(jobs)
+            for _, jobs in group_jobs_by_replica_latest(run_model.jobs)
+        )
+        # stop excessive registered out-of-date replicas, except those that are already `terminating`
+        non_terminating_registered_replicas_count = sum(
+            is_replica_registered(jobs) and all(j.status != JobStatus.TERMINATING for j in jobs)
             for _, jobs in group_jobs_by_replica_latest(run_model.jobs)
         )
-        ready_replica_count = sum(
-            is_replica_ready(jobs) for _, jobs in group_jobs_by_replica_latest(run_model.jobs)
+        replicas_to_stop_count += max(
+            0, non_terminating_registered_replicas_count - run_model.desired_replica_count
         )
-        if ready_replica_count > run_model.desired_replica_count:
-            # stop excessive ready out-of-date replicas
-            replicas_to_stop_count += ready_replica_count - run_model.desired_replica_count
         if replicas_to_stop_count:
             await scale_run_replicas(
                 session,

src/dstack/_internal/server/migrations/versions/3d7f6c2ec000_add_jobmodel_registered.py

@@ -0,0 +1,28 @@
+"""Add JobModel.registered
+
+Revision ID: 3d7f6c2ec000
+Revises: 74a1f55209bd
+Create Date: 2025-08-11 13:23:39.530103
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "3d7f6c2ec000"
+down_revision = "74a1f55209bd"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("registered", sa.Boolean(), server_default=sa.false(), nullable=False)
+        )
+
+
+def downgrade() -> None:
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.drop_column("registered")

src/dstack/_internal/server/models.py

@@ -430,6 +430,9 @@ class JobModel(BaseModel):
     probes: Mapped[list["ProbeModel"]] = relationship(
         back_populates="job", order_by="ProbeModel.probe_num"
     )
+    # Whether the replica is registered to receive service requests.
+    # Always `False` for non-service runs.
+    registered: Mapped[bool] = mapped_column(Boolean, server_default=false())
 
 
 class GatewayModel(BaseModel):

src/dstack/_internal/server/services/probes.py

@@ -1,6 +1,10 @@
-from dstack._internal.core.models.runs import Probe
+from dstack._internal.core.models.runs import Probe, ProbeSpec
 from dstack._internal.server.models import ProbeModel
 
 
 def probe_model_to_probe(probe_model: ProbeModel) -> Probe:
     return Probe(success_streak=probe_model.success_streak)
+
+
+def is_probe_ready(probe: ProbeModel, spec: ProbeSpec) -> bool:
+    return probe.success_streak >= spec.ready_after

src/dstack/_internal/server/services/proxy/repo.py

@@ -54,6 +54,7 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
                 RunModel.gateway_id.is_(None),
                 JobModel.run_name == run_name,
                 JobModel.status == JobStatus.RUNNING,
+                JobModel.registered == True,
                 JobModel.job_num == 0,
             )
             .options(

src/dstack/_internal/server/services/runs.py

@@ -41,6 +41,7 @@
     JobStatus,
     JobSubmission,
     JobTerminationReason,
+    ProbeSpec,
     Run,
     RunPlan,
     RunSpec,
@@ -58,6 +59,7 @@
 from dstack._internal.server.db import get_db
 from dstack._internal.server.models import (
     JobModel,
+    ProbeModel,
     ProjectModel,
     RepoModel,
     RunModel,
@@ -86,6 +88,7 @@
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import get_offers_by_requirements
 from dstack._internal.server.services.plugins import apply_plugin_policies
+from dstack._internal.server.services.probes import is_probe_ready
 from dstack._internal.server.services.projects import list_user_project_models
 from dstack._internal.server.services.resources import set_resources_defaults
 from dstack._internal.server.services.secrets import get_project_secrets_mapping
@@ -1185,8 +1188,8 @@ async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replica
         elif {JobStatus.PROVISIONING, JobStatus.PULLING} & statuses:
             # if there are any provisioning or pulling jobs, the replica is active and has the importance of 1
             active_replicas.append((1, is_out_of_date, replica_num, replica_jobs))
-        elif not is_replica_ready(replica_jobs):
-            # all jobs are running, but probes are failing, the replica is active and has the importance of 2
+        elif not is_replica_registered(replica_jobs):
+            # all jobs are running, but not receiving traffic, the replica is active and has the importance of 2
             active_replicas.append((2, is_out_of_date, replica_num, replica_jobs))
         else:
             # all jobs are running and ready, the replica is active and has the importance of 3
@@ -1273,15 +1276,13 @@ async def retry_run_replica_jobs(
         session.add(new_job_model)
 
 
-def is_replica_ready(jobs: Iterable[JobModel]) -> bool:
-    if not all(job.status == JobStatus.RUNNING for job in jobs):
-        return False
-    for job in jobs:
-        job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
-        for probe_spec, probe in zip(job_spec.probes, job.probes):
-            if probe.success_streak < probe_spec.ready_after:
-                return False
-    return True
+def is_job_ready(probes: Iterable[ProbeModel], probe_specs: Iterable[ProbeSpec]) -> bool:
+    return all(is_probe_ready(probe, probe_spec) for probe, probe_spec in zip(probes, probe_specs))
+
+
+def is_replica_registered(jobs: list[JobModel]) -> bool:
+    # Only job_num=0 is supposed to receive service requests
+    return jobs[0].registered
 
 
 def _remove_job_spec_sensitive_info(spec: JobSpec):