From 91fbe2fbea6b141fd98440a3f605ee988f051481 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Fri, 14 Aug 2026 13:38:22 +0000 Subject: [PATCH] Don't process running job if other jobs in replica are not provisioned Fixes: https://github.com/dstackai/dstack/issues/4146 --- .../background/pipeline_tasks/jobs_running.py | 26 +++- .../pipeline_tasks/test_running_jobs.py | 111 ++++++++++++++++++ 2 files changed, 132 insertions(+), 5 deletions(-) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py index 77bf8576f..37fc51ab8 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -507,16 +507,32 @@ async def _prepare_startup_context( ) -> Optional[_StartupContext]: job_provisioning_data = get_or_error(context.job_provisioning_data) + # `_get_cluster_info` below requires every job in the replica to have provisioning + # data, so gate on that rather than on job statuses. A `SUBMITTED` job is simply not + # provisioned yet, while a job that has no provisioning data and is no longer + # `SUBMITTED` reached a terminal state without ever provisioning (e.g. due to no + # capacity) and never will. Deferring in the latter case is what allows recovery: + # the run pipeline retries or fails the replica, but it can only lock the run's jobs + # once this job is unlocked, which does not happen if we raise here. for other_job in context.run.jobs: - if ( - other_job.job_spec.replica_num == context.job.job_spec.replica_num - and other_job.job_submissions[-1].status == JobStatus.SUBMITTED - ): + if other_job.job_spec.replica_num != context.job.job_spec.replica_num: + continue + other_job_submission = other_job.job_submissions[-1] + if other_job_submission.job_provisioning_data is not None: + continue + if other_job_submission.status == JobStatus.SUBMITTED: logger.debug( "%s: waiting for all jobs in the replica to be provisioned", fmt(context.job_model), ) - return None + else: + logger.debug( + "%s: job %s in the replica has no provisioning data, waiting for the run" + " to be retried or terminated", + fmt(context.job_model), + other_job.job_spec.job_name, + ) + return None # If this run has a router replica group and this job is a worker, gate # startup on the router replica's state. The helper returns None for the diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 19dfd4219..8cea089da 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -3088,6 +3088,117 @@ async def _fake_session_ctx(): assert out.router_env == router_env +@pytest.mark.asyncio +class TestPrepareStartupContextReplicaGate: + """Tests for the https://github.com/dstackai/dstack/issues/4146 fix""" + + def _make_job( + self, + *, + replica_num: int, + status: JobStatus, + job_provisioning_data: Optional[MagicMock], + job_name: str = "run-0-0", + ) -> MagicMock: + job = MagicMock() + job.job_spec.replica_num = replica_num + job.job_spec.job_name = job_name + job_submission = MagicMock() + job_submission.status = status + job_submission.job_provisioning_data = job_provisioning_data + job.job_submissions = [job_submission] + return job + + def _make_context(self, other_jobs: list[MagicMock]) -> _ProcessContext: + job = self._make_job( + replica_num=0, + status=JobStatus.PROVISIONING, + job_provisioning_data=MagicMock(), + ) + run = MagicMock() + run.jobs = [job] + other_jobs + return _ProcessContext( + job_model=MagicMock(), + run_model=MagicMock(), + run=run, + job=job, + job_submission=MagicMock(job_runtime_data=None), + job_provisioning_data=MagicMock(), + instance_access_revoked=False, + ) + + async def _run_gate(self, context: _ProcessContext, result: _ProcessResult): + # `get_router_env_for_job` is the first thing after the gate. Returning `FAILED` + # makes passing the gate observable without mocking the whole startup path. + with patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.get_router_env_for_job", + return_value=RouterEnvStatus.FAILED, + ): + return await _prepare_startup_context(context=context, result=result) + + def _assert_gate_not_passed(self, result: _ProcessResult) -> None: + assert result.job_update_map == {} + + def _assert_gate_passed(self, result: _ProcessResult) -> None: + assert "Router replica is in a terminal state" in ( + result.job_update_map.get("termination_reason_message") or "" + ) + + async def test_defers_on_submitted_job_in_replica(self): + context = self._make_context( + [ + self._make_job( + replica_num=0, + status=JobStatus.SUBMITTED, + job_provisioning_data=None, + ) + ] + ) + result = _ProcessResult() + assert await self._run_gate(context, result) is None + self._assert_gate_not_passed(result) + + @pytest.mark.parametrize("status", [JobStatus.FAILED, JobStatus.TERMINATED, JobStatus.ABORTED]) + async def test_defers_on_job_in_replica_terminated_before_provisioning(self, status): + context = self._make_context( + [self._make_job(replica_num=0, status=status, job_provisioning_data=None)] + ) + result = _ProcessResult() + assert await self._run_gate(context, result) is None + # The job is not terminated here: the run pipeline decides whether the replica + # is retried or failed, and it can only do so once this job is unlocked. + self._assert_gate_not_passed(result) + + async def test_does_not_defer_on_done_job_in_replica(self): + # A job that already finished its work does not block its slower siblings. + context = self._make_context( + [ + self._make_job( + replica_num=0, + status=JobStatus.DONE, + job_provisioning_data=MagicMock(), + ) + ] + ) + result = _ProcessResult() + assert await self._run_gate(context, result) is None + self._assert_gate_passed(result) + + async def test_does_not_defer_on_job_in_other_replica(self): + context = self._make_context( + [ + self._make_job( + replica_num=1, + status=JobStatus.FAILED, + job_provisioning_data=None, + ) + ] + ) + result = _ProcessResult() + assert await self._run_gate(context, result) is None + self._assert_gate_passed(result) + + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestFetchRunModelDynamoBranch: