From a1337fdc3604dde090394b091b9e33a7021bcb71 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Tue, 20 Oct 2026 08:27:20 +0000 Subject: [PATCH] Wait for other jobs only before submitting to the runner `_prepare_startup_context` ran on every PROVISIONING and PULLING tick and returned early whenever the runner submission had to wait for other jobs: replica provisioning, the router IP, or referenced group IPs. The early return also skipped the status handlers' own work: copying the instance provisioning data into the job, failing the job when its instance is terminated, the provisioning timeout, and shim polling. With replica IP refs on newly provisioned instances, jobs waited for IPs that only the skipped copy would set, so the run stayed in provisioning forever. * Replace `_prepare_startup_context` with `_prepare_runner_submission`, called right before the runner submission once the runner is up. * Load volumes and interpolate secrets in the shim branch, which needs the interpolated `registry_auth`. The shim no longer waits for other jobs, so the image pull starts earlier. * Terminate workers whose router failed or did not get an IP in time in `_terminate_if_router_unavailable`, before both statuses as before. * Wait for every job in the replica to have a hostname, not only provisioning data, since shim submission no longer ensures it. * Drop the unused `secrets` parameter of `_submit_job_to_runner`. Fixes: https://github.com/dstackai/dstack/issues/4331 Co-Authored-By: Claude Opus 5.5 (1M context) --- .../background/pipeline_tasks/jobs_running.py | 380 ++++++++++-------- src/dstack/_internal/server/testing/common.py | 2 +- .../pipeline_tasks/test_running_jobs.py | 326 ++++++++++----- 3 files changed, 442 insertions(+), 266 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 35d74edc21..342cdbc301 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -404,20 +404,6 @@ class _ProcessResult: new_probe_models: list[ProbeModel] = field(default_factory=list) -@dataclass -class _StartupContext: - cluster_info: ClusterInfo - volumes: list[Volume] - secrets: dict[str, str] - repo_creds: Optional[RemoteRepoCreds] - router_env: Optional[Dict[str, str]] = None - """Dynamo-specific env (e.g. DSTACK_ROUTER_INTERNAL_IP) computed from the - router replica's state. Passed through to RunnerClient.submit_job, which - merges it into a deep-copied job_spec.env so the shared job_spec is not - mutated. None for SMG services, non-router runs, and the router - replica itself.""" - - async def _load_process_context(item: JobRunningPipelineItem) -> Optional[_ProcessContext]: async with get_session_ctx() as session: job_model = await _refetch_locked_job_model(session=session, item=item) @@ -436,8 +422,9 @@ async def _load_process_context(item: JobRunningPipelineItem) -> Optional[_Proce else: # PROVISIONING/PULLING jobs need same-replica siblings for cluster # coordination, plus — when the run has a router replica group — - # the router replica's job (cross-replica) so the env-injection - # gate in _prepare_startup_context can read its status / IP. + # the router replica's job (cross-replica) so the router checks in + # _terminate_if_router_unavailable and _prepare_runner_submission + # can read its status / IP. # _fetch_run_model handles both: same-replica jobs always, plus # all non-terminated jobs when one exists. run_spec = validate_json_extra_ignore(RunSpec, job_model.run.run_spec) @@ -490,19 +477,13 @@ async def _process_running_job(context: _ProcessContext) -> _ProcessResult: return result if context.job_model.status == JobStatus.PROVISIONING: - startup_context = await _prepare_startup_context(context=context, result=result) - if startup_context is None: + if _terminate_if_router_unavailable(context=context, result=result): return result - await _process_provisioning_status( - context=context, startup_context=startup_context, result=result - ) + await _process_provisioning_status(context=context, result=result) elif context.job_model.status == JobStatus.PULLING: - startup_context = await _prepare_startup_context(context=context, result=result) - if startup_context is None: + if _terminate_if_router_unavailable(context=context, result=result): return result - await _process_pulling_status( - context=context, startup_context=startup_context, result=result - ) + await _process_pulling_status(context=context, result=result) elif context.job_model.status == JobStatus.RUNNING: if _server_access_enabled(context): await job_server_connections_pool.ensure( @@ -528,49 +509,21 @@ async def _process_running_job(context: _ProcessContext) -> _ProcessResult: return result -async def _prepare_startup_context( - context: _ProcessContext, - result: _ProcessResult, -) -> 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: - 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), - ) - 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 +def _terminate_if_router_unavailable(context: _ProcessContext, result: _ProcessResult) -> bool: + """ + Terminates a worker job whose router replica is in a terminal state or did not acquire an + internal IP in time. Waiting for the router IP is up to `_prepare_runner_submission`. - # 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 - # router itself and for runs without a router group, so this whole block - # is a no-op in those cases. - router_env_outcome = get_router_env_for_job( + Returns `True` if the job was terminated. + """ + # `get_router_env_for_job` returns None for the router itself and for runs without + # a router group, so this check is a no-op in those cases. + router_env = get_router_env_for_job( run_model=context.run_model, run_spec=context.run.run_spec, job_model=context.job_model, ) - if router_env_outcome is RouterEnvStatus.FAILED: + if router_env is RouterEnvStatus.FAILED: # Router has reached a terminal state — the worker cannot recover by # waiting. Terminate it now with a clear reason instead of letting it # idle until the run-level reconciler tears the whole run down. @@ -583,100 +536,24 @@ async def _prepare_startup_context( "without a running router." ), ) - return None - if router_env_outcome is RouterEnvStatus.NOT_PROVISIONED: - # Router is alive but its internal_ip is not yet known. Defer this - # worker — the next pipeline tick will re-check. Bound the wait so a - # router that is genuinely stuck can't burn worker instance-hours - # forever; see ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS. - waited_seconds = (get_current_datetime() - context.job_model.submitted_at).total_seconds() - if waited_seconds > ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS: - _terminate_job( - job_model=context.job_model, - job_update_map=result.job_update_map, - termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, - termination_reason_message=( - f"Router replica did not acquire an internal IP within " - f"{ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS}s; terminating worker." - ), - ) - return None - logger.debug( - "%s: waiting for router replica to be provisioned", - fmt(context.job_model), - ) - return None - # Past the enum branches, router_env_outcome is either None or a Dict. - # We don't mutate job_spec.env here — RunnerClient.submit_job merges it - # into a deep-copied spec, mirroring how instance_env is handled. - router_env: Optional[Dict[str, str]] = ( - router_env_outcome if isinstance(router_env_outcome, dict) else None - ) - - cluster_info = _get_cluster_info( - jobs=context.run.jobs, - replica_num=context.job.job_spec.replica_num, - job_provisioning_data=job_provisioning_data, - job_runtime_data=context.job_submission.job_runtime_data, - ) - - async with get_session_ctx() as session: - volumes = await get_job_attached_volumes( - session=session, - project=context.project, - run_spec=context.run.run_spec, - job_num=context.job.job_spec.job_num, - job_provisioning_data=job_provisioning_data, - ) - repo_creds_model = await get_repo_creds( - session=session, - repo=context.repo_model, - user=context.run_model.user, - ) - secrets = await get_project_secrets_mapping(session=session, project=context.project) - - repo_creds = repo_model_to_repo_head_with_creds( - context.repo_model, - repo_creds_model, - ).repo_creds - - try: - interpolate_job_spec_secrets(context.job.job_spec, secrets) - except InterpolatorError as e: - _terminate_job( - job_model=context.job_model, - job_update_map=result.job_update_map, - termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, - termination_reason_message=f"Secrets interpolation error: {e.args[0]}", - ) - return None - - commands = context.job.job_spec.commands - try: - for c in commands: - validate_groups_refs(c) - if not _substitute_groups_ip_refs(commands, context): - logger.debug( - "%s: waiting for referenced group IPs", - fmt(context.job_model), - ) - return None - except InterpolatorError as e: - _terminate_job( - job_model=context.job_model, - job_update_map=result.job_update_map, - termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, - termination_reason_message=f"Groups IP interpolation error: {e.args[0]}", - ) - return None - - return _StartupContext( - cluster_info=cluster_info, - volumes=volumes, - secrets=secrets, - repo_creds=repo_creds, - router_env=router_env, + return True + if router_env is not RouterEnvStatus.NOT_PROVISIONED: + return False + # Bound the wait for the router IP so a router that is genuinely stuck can't burn + # worker instance-hours forever; see ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS. + waited_seconds = (get_current_datetime() - context.job_model.submitted_at).total_seconds() + if waited_seconds <= ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS: + return False + _terminate_job( + job_model=context.job_model, + job_update_map=result.job_update_map, + termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, + termination_reason_message=( + f"Router replica did not acquire an internal IP within " + f"{ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS}s; terminating worker." + ), ) + return True async def _refetch_locked_job_model( @@ -808,7 +685,6 @@ async def _is_instance_access_revoked(session: AsyncSession, job_model: JobModel async def _process_provisioning_status( context: _ProcessContext, - startup_context: _StartupContext, result: _ProcessResult, ) -> None: job_provisioning_data = get_or_error(context.job_provisioning_data) @@ -826,6 +702,18 @@ async def _process_provisioning_status( fmt(context.job_model), context.job_submission.age, ) + async with get_session_ctx() as session: + volumes = await get_job_attached_volumes( + session=session, + project=context.project, + run_spec=context.run.run_spec, + job_num=context.job.job_spec.job_num, + job_provisioning_data=job_provisioning_data, + ) + secrets = await get_project_secrets_mapping(session=session, project=context.project) + # The shim pulls the image with the interpolated `registry_auth` + if not _interpolate_secrets_or_terminate(context, result, secrets): + return extra_authorized_keys = get_extra_authorized_keys(context.run.run_spec) public_keys = [context.project.ssh_public_key.strip(), *extra_authorized_keys] # Host access, unlike container access, is all or nothing -- the user key is added to @@ -845,7 +733,7 @@ async def _process_provisioning_status( job_model=context.job_model, jrd=get_job_runtime_data(context.job_model), jpd=job_provisioning_data, - volumes=startup_context.volumes, + volumes=volumes, registry_auth=context.job.job_spec.registry_auth, public_keys=public_keys, ssh_user=ssh_user, @@ -877,6 +765,11 @@ async def _process_provisioning_status( job_provisioning_data, None, ): + runner_submission = await _prepare_runner_submission(context, result) + if runner_submission is None: + # The runner is up, so the provisioning timeout below does not apply + # while waiting for other jobs. + return if not await _ensure_job_server_connection(context, result): return file_archives = await _get_job_file_archives( @@ -897,12 +790,11 @@ async def _process_provisioning_status( job_model=context.job_model, job=context.job, jrd=get_job_runtime_data(context.job_model), - cluster_info=startup_context.cluster_info, + cluster_info=runner_submission.cluster_info, code=code, file_archives=file_archives, - secrets=startup_context.secrets, - repo_credentials=startup_context.repo_creds, - router_env=startup_context.router_env, + repo_credentials=runner_submission.repo_creds, + router_env=runner_submission.router_env, success_if_not_available=False, ) _apply_submit_job_to_runner_result( @@ -936,7 +828,6 @@ async def _process_provisioning_status( async def _process_pulling_status( context: _ProcessContext, - startup_context: _StartupContext, result: _ProcessResult, ) -> None: job_provisioning_data = get_or_error(context.job_provisioning_data) @@ -995,6 +886,10 @@ async def _process_pulling_status( _reset_disconnected_at(context.job_model, result) return + runner_submission = await _prepare_runner_submission(context, result) + if runner_submission is None: + _reset_disconnected_at(context.job_model, result) + return if not await _ensure_job_server_connection(context, result): return file_archives = await _get_job_file_archives( @@ -1015,12 +910,11 @@ async def _process_pulling_status( job_model=context.job_model, job=context.job, jrd=job_runtime_data, - cluster_info=startup_context.cluster_info, + cluster_info=runner_submission.cluster_info, code=code, file_archives=file_archives, - secrets=startup_context.secrets, - repo_credentials=startup_context.repo_creds, - router_env=startup_context.router_env, + repo_credentials=runner_submission.repo_creds, + router_env=runner_submission.router_env, success_if_not_available=True, ) _apply_submit_job_to_runner_result( @@ -1426,6 +1320,155 @@ def _should_wait_for_other_nodes(run: Run, job: Job, job_model: JobModel) -> boo return False +def _interpolate_secrets_or_terminate( + context: _ProcessContext, + result: _ProcessResult, + secrets: Mapping[str, str], +) -> bool: + """ + Interpolates secrets into `context.job.job_spec` in place. + + Returns `False` if the job was terminated due to an interpolation error. + """ + try: + interpolate_job_spec_secrets(context.job.job_spec, secrets) + except InterpolatorError as e: + _terminate_job( + job_model=context.job_model, + job_update_map=result.job_update_map, + termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, + termination_reason_message=f"Secrets interpolation error: {e.args[0]}", + ) + return False + return True + + +@dataclass +class _RunnerSubmissionContext: + cluster_info: ClusterInfo + repo_creds: Optional[RemoteRepoCreds] + router_env: Optional[Dict[str, str]] = None + """Dynamo-specific env (e.g. DSTACK_ROUTER_INTERNAL_IP) computed from the + router replica's state. Passed through to RunnerClient.submit_job, which + merges it into a deep-copied job_spec.env so the shared job_spec is not + mutated. None for SMG services, non-router runs, and the router + replica itself.""" + + +async def _prepare_runner_submission( + context: _ProcessContext, + result: _ProcessResult, +) -> Optional[_RunnerSubmissionContext]: + """ + Waits for the other jobs that the runner submission depends on, collects the submission + data, and interpolates secrets and groups IP refs into `context.job.job_spec` in place. + + Returns `None` if the submission has to wait or the job was terminated. + """ + if _should_wait_for_replica_provisioning(context): + return None + + router_env = get_router_env_for_job( + run_model=context.run_model, + run_spec=context.run.run_spec, + job_model=context.job_model, + ) + if isinstance(router_env, RouterEnvStatus): + # `FAILED` and a timed-out wait are handled by `_terminate_if_router_unavailable` + logger.debug( + "%s: waiting for router replica to be provisioned", + fmt(context.job_model), + ) + return None + + cluster_info = _get_cluster_info( + jobs=context.run.jobs, + replica_num=context.job.job_spec.replica_num, + job_provisioning_data=get_or_error(context.job_provisioning_data), + job_runtime_data=context.job_submission.job_runtime_data, + ) + + async with get_session_ctx() as session: + repo_creds_model = await get_repo_creds( + session=session, + repo=context.repo_model, + user=context.run_model.user, + ) + secrets = await get_project_secrets_mapping(session=session, project=context.project) + + repo_creds = repo_model_to_repo_head_with_creds( + context.repo_model, + repo_creds_model, + ).repo_creds + + if not _interpolate_secrets_or_terminate(context, result, secrets): + return None + + commands = context.job.job_spec.commands + try: + for c in commands: + validate_groups_refs(c) + if not _substitute_groups_ip_refs(commands, context): + logger.debug( + "%s: waiting for referenced group IPs", + fmt(context.job_model), + ) + return None + except InterpolatorError as e: + _terminate_job( + job_model=context.job_model, + job_update_map=result.job_update_map, + termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, + termination_reason_message=f"Groups IP interpolation error: {e.args[0]}", + ) + return None + + return _RunnerSubmissionContext( + cluster_info=cluster_info, + repo_creds=repo_creds, + router_env=router_env, + ) + + +def _should_wait_for_replica_provisioning(context: _ProcessContext) -> bool: + """ + Whether any job in the replica has no provisioning data or no hostname yet. + `_get_cluster_info` needs the IPs of all of them. + """ + # 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 `_get_cluster_info` raises. + for other_job in context.run.jobs: + if other_job.job_spec.replica_num != context.job.job_spec.replica_num: + continue + other_job_submission = other_job.job_submissions[-1] + other_jpd = other_job_submission.job_provisioning_data + if other_jpd is not None and other_jpd.hostname is not None: + continue + if other_jpd is not None: + logger.debug( + "%s: waiting for job %s in the replica to have IP assigned", + fmt(context.job_model), + other_job.job_spec.job_name, + ) + elif other_job_submission.status == JobStatus.SUBMITTED: + logger.debug( + "%s: waiting for all jobs in the replica to be provisioned", + fmt(context.job_model), + ) + 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 True + return False + + @runner_ssh_tunnel def _process_provisioning_with_shim( addresses: Mapping[int, client.LocalAddress], @@ -1653,7 +1696,6 @@ def _submit_job_to_runner( cluster_info: ClusterInfo, code: Optional[bytes], file_archives: Iterable[tuple[uuid.UUID, bytes]], - secrets: Dict[str, str], repo_credentials: Optional[RemoteRepoCreds], router_env: Optional[Dict[str, str]], success_if_not_available: bool, diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index f6dc631812..f4a2acaba5 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -510,7 +510,7 @@ def get_job_provisioning_data( cpu_count: int = 1, memory_gib: float = 0.5, spot: bool = False, - hostname: str = "127.0.0.4", + hostname: Optional[str] = "127.0.0.4", internal_ip: Optional[str] = "127.0.0.4", price: float = 10.5, instance_type: Optional[InstanceType] = None, 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 08940cc544..5e85487313 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 @@ -31,6 +31,7 @@ ClusterInfo, ImagePullProgress, Job, + JobProvisioningData, JobRuntimeData, JobSpec, JobStatus, @@ -54,11 +55,13 @@ _build_replica_groups_ip_view, _fetch_run_model, _get_cluster_info, - _prepare_startup_context, + _prepare_runner_submission, _ProcessContext, _ProcessResult, _referenced_ips_ready, + _should_wait_for_replica_provisioning, _SubmitJobToRunnerResult, + _terminate_if_router_unavailable, ) from dstack._internal.server.background.pipeline_tasks.runs import RunPipeline from dstack._internal.server.models import JobModel, ProbeModel, ServiceReplicaRegistrationModel @@ -1501,6 +1504,135 @@ async def test_master_job_waits_for_workers( await session.refresh(master_job) assert master_job.status == JobStatus.RUNNING + async def test_provisioning_job_waiting_for_replica_ip_picks_up_instance_ip( + self, test_db, session: AsyncSession, worker: JobRunningWorker + ): + """Regression test for https://github.com/dstackai/dstack/issues/4331""" + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=get_run_spec( + repo_id=repo.name, configuration=_replica_ip_refs_service_configuration() + ), + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + job_provisioning_data=get_job_provisioning_data( + dockerized=True, hostname="1.2.3.4", internal_ip="10.0.0.1" + ), + ) + # The job references its own IP, which is known only after the job picks up + # the provisioning data of its newly provisioned instance + job = await create_job( + session=session, + run=run, + status=JobStatus.PROVISIONING, + submitted_at=get_current_datetime(), + job_provisioning_data=get_job_provisioning_data( + dockerized=True, hostname=None, internal_ip=None + ), + instance=instance, + instance_assigned=True, + replica_group_name="workers", + ) + + await _process_job(session, worker, job) + + await session.refresh(job) + assert job.status == JobStatus.PROVISIONING + job_provisioning_data = validate_json_extra_ignore( + JobProvisioningData, get_or_error(job.job_provisioning_data) + ) + assert job_provisioning_data.hostname == "1.2.3.4" + assert job_provisioning_data.internal_ip == "10.0.0.1" + + async def test_waiting_for_replica_ip_blocks_only_runner_submission( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ssh_tunnel_mock: Mock, + shim_client_mock: Mock, + runner_client_mock: Mock, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=get_run_spec( + repo_id=repo.name, configuration=_replica_ip_refs_service_configuration() + ), + ) + # Replica 0 is referenced by all replicas and its instance is not running yet + head_job = await create_job( + session=session, + run=run, + status=JobStatus.PROVISIONING, + job_provisioning_data=get_job_provisioning_data( + dockerized=True, hostname=None, internal_ip=None + ), + replica_num=0, + replica_group_name="workers", + ) + instance = await create_instance( + session=session, project=project, status=InstanceStatus.BUSY + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.PROVISIONING, + submitted_at=get_current_datetime(), + job_provisioning_data=get_job_provisioning_data( + dockerized=True, internal_ip="10.0.0.2" + ), + instance=instance, + instance_assigned=True, + replica_num=1, + replica_group_name="workers", + ) + + # The shim does not need the referenced IP, so the image is pulled right away + await _process_job(session, worker, job) + + shim_client_mock.submit_task.assert_called_once() + await session.refresh(job) + assert job.status == JobStatus.PULLING + + # The runner submission waits for the referenced IP, while the shim is still polled + shim_client_mock.get_task.return_value.status = TaskStatus.RUNNING + await _process_job(session, worker, job) + + shim_client_mock.get_task.assert_called_once() + runner_client_mock.submit_job.assert_not_called() + await session.refresh(job) + assert job.status == JobStatus.PULLING + + head_job.job_provisioning_data = get_job_provisioning_data( + dockerized=True, internal_ip="10.0.0.1" + ).model_dump_json() + await session.commit() + runner_client_mock.run_job.return_value = JobInfoResponse( + working_dir="/dstack/run", username="dstack" + ) + await _process_job(session, worker, job) + + runner_client_mock.submit_job.assert_called_once() + submitted_job = runner_client_mock.submit_job.call_args.kwargs["job"] + assert submitted_job.job_spec.commands == ["/bin/sh", "-i", "-c", "echo head=10.0.0.1"] + await session.refresh(job) + assert job.status == JobStatus.RUNNING + async def test_apply_skips_when_lock_token_changes_after_processing( self, test_db, session: AsyncSession, worker: JobRunningWorker ): @@ -3157,28 +3289,45 @@ def _router_service_configuration( ) -@pytest.mark.asyncio -class TestPrepareStartupContextRouterEnv: - def _make_context(self, *, submitted_at: datetime) -> _ProcessContext: - job_model = MagicMock() - job_model.submitted_at = submitted_at - job = MagicMock() - job.job_spec.replica_num = 0 - run = MagicMock() - run.jobs = [] - run.run_spec = MagicMock() - return _ProcessContext( - job_model=job_model, - run_model=MagicMock(), - run=run, - job=job, - job_submission=MagicMock(job_runtime_data=None), - job_provisioning_data=MagicMock(), - instance_access_revoked=False, - ) +def _replica_ip_refs_service_configuration() -> ServiceConfiguration: + return ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "groups": [ + { + "name": "workers", + "replicas": 2, + "commands": ["echo head=${{ groups[0].replicas[0].IP_ADDRESS }}"], + } + ], + } + ) - async def test_router_failed_terminates_worker(self): - context = self._make_context( + +def _make_router_worker_context(*, submitted_at: datetime) -> _ProcessContext: + job_model = MagicMock() + job_model.submitted_at = submitted_at + job = MagicMock() + job.job_spec.replica_num = 0 + run = MagicMock() + run.jobs = [] + run.run_spec = MagicMock() + return _ProcessContext( + job_model=job_model, + run_model=MagicMock(), + run=run, + job=job, + job_submission=MagicMock(job_runtime_data=None), + job_provisioning_data=MagicMock(), + instance_access_revoked=False, + ) + + +class TestTerminateIfRouterUnavailable: + def test_router_failed_terminates_worker(self): + context = _make_router_worker_context( submitted_at=datetime(2023, 1, 1, 12, 0, 0, tzinfo=timezone.utc), ) result = _ProcessResult() @@ -3186,8 +3335,7 @@ async def test_router_failed_terminates_worker(self): "dstack._internal.server.background.pipeline_tasks.jobs_running.get_router_env_for_job", return_value=RouterEnvStatus.FAILED, ): - out = await _prepare_startup_context(context=context, result=result) - assert out is None + assert _terminate_if_router_unavailable(context=context, result=result) assert ( result.job_update_map.get("termination_reason") == JobTerminationReason.TERMINATED_BY_SERVER @@ -3197,8 +3345,8 @@ async def test_router_failed_terminates_worker(self): ) @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) - async def test_router_not_provisioned_within_timeout_defers(self): - context = self._make_context( + def test_router_not_provisioned_within_timeout_does_not_terminate(self): + context = _make_router_worker_context( submitted_at=datetime(2023, 1, 1, 11, 45, 0, tzinfo=timezone.utc), ) result = _ProcessResult() @@ -3206,13 +3354,12 @@ async def test_router_not_provisioned_within_timeout_defers(self): "dstack._internal.server.background.pipeline_tasks.jobs_running.get_router_env_for_job", return_value=RouterEnvStatus.NOT_PROVISIONED, ): - out = await _prepare_startup_context(context=context, result=result) - assert out is None + assert not _terminate_if_router_unavailable(context=context, result=result) assert result.job_update_map == {} @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) - async def test_router_not_provisioned_past_timeout_terminates(self): - context = self._make_context( + def test_router_not_provisioned_past_timeout_terminates(self): + context = _make_router_worker_context( submitted_at=datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc), ) result = _ProcessResult() @@ -3220,8 +3367,7 @@ async def test_router_not_provisioned_past_timeout_terminates(self): "dstack._internal.server.background.pipeline_tasks.jobs_running.get_router_env_for_job", return_value=RouterEnvStatus.NOT_PROVISIONED, ): - out = await _prepare_startup_context(context=context, result=result) - assert out is None + assert _terminate_if_router_unavailable(context=context, result=result) assert ( result.job_update_map.get("termination_reason") == JobTerminationReason.TERMINATED_BY_SERVER @@ -3230,8 +3376,25 @@ async def test_router_not_provisioned_past_timeout_terminates(self): assert str(ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS) in msg assert "internal IP" in msg - async def test_router_env_dict_populates_startup_context(self): - context = self._make_context( + +@pytest.mark.asyncio +class TestPrepareRunnerSubmissionRouterEnv: + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) + async def test_router_not_provisioned_waits(self): + context = _make_router_worker_context( + submitted_at=datetime(2023, 1, 1, 11, 45, 0, tzinfo=timezone.utc), + ) + result = _ProcessResult() + with patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.get_router_env_for_job", + return_value=RouterEnvStatus.NOT_PROVISIONED, + ): + out = await _prepare_runner_submission(context=context, result=result) + assert out is None + assert result.job_update_map == {} + + async def test_router_env_dict_populates_submission_context(self): + context = _make_router_worker_context( submitted_at=datetime(2023, 1, 1, 12, 0, 0, tzinfo=timezone.utc), ) result = _ProcessResult() @@ -3250,11 +3413,6 @@ async def _fake_session_ctx(): "dstack._internal.server.background.pipeline_tasks.jobs_running.get_session_ctx", _fake_session_ctx, ), - patch( - "dstack._internal.server.background.pipeline_tasks.jobs_running.get_job_attached_volumes", - new_callable=AsyncMock, - return_value=[], - ), patch( "dstack._internal.server.background.pipeline_tasks.jobs_running.get_repo_creds", new_callable=AsyncMock, @@ -3279,13 +3437,13 @@ async def _fake_session_ctx(): ), ), ): - out = await _prepare_startup_context(context=context, result=result) + out = await _prepare_runner_submission(context=context, result=result) assert out is not None assert out.router_env == router_env @pytest.mark.asyncio -class TestPrepareStartupContextClusterWait: +class TestPrepareRunnerSubmissionClusterWait: def _make_context(self) -> _ProcessContext: job_model = MagicMock() job_model.submitted_at = datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc) @@ -3314,13 +3472,13 @@ def _make_context(self) -> _ProcessContext: async def test_submitted_sibling_defers(self): context = self._make_context() result = _ProcessResult() - out = await _prepare_startup_context(context=context, result=result) + out = await _prepare_runner_submission(context=context, result=result) assert out is None assert result.job_update_map == {} @pytest.mark.asyncio -class TestPrepareStartupContextGroupsIpWait: +class TestPrepareRunnerSubmissionGroupsIpWait: def _make_context(self) -> _ProcessContext: job_model = MagicMock() job_model.submitted_at = datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc) @@ -3363,11 +3521,6 @@ async def _fake_session_ctx(): "dstack._internal.server.background.pipeline_tasks.jobs_running.get_session_ctx", _fake_session_ctx, ), - patch( - "dstack._internal.server.background.pipeline_tasks.jobs_running.get_job_attached_volumes", - new_callable=AsyncMock, - return_value=[], - ), patch( "dstack._internal.server.background.pipeline_tasks.jobs_running.get_repo_creds", new_callable=AsyncMock, @@ -3400,14 +3553,13 @@ async def test_groups_ip_not_ready_defers(self): with ExitStack() as stack: for p in self._patches(): stack.enter_context(p) - out = await _prepare_startup_context(context=context, result=result) + out = await _prepare_runner_submission(context=context, result=result) assert out is None assert result.job_update_map == {} -@pytest.mark.asyncio -class TestPrepareStartupContextReplicaGate: - """Tests for the https://github.com/dstackai/dstack/issues/4146 fix""" +class TestShouldWaitForReplicaProvisioning: + """Includes tests for the https://github.com/dstackai/dstack/issues/4146 fix""" def _make_job( self, @@ -3444,24 +3596,7 @@ def _make_context(self, other_jobs: list[MagicMock]) -> _ProcessContext: 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): + def test_waits_for_submitted_job_in_replica(self): context = self._make_context( [ self._make_job( @@ -3471,22 +3606,30 @@ async def test_defers_on_submitted_job_in_replica(self): ) ] ) - result = _ProcessResult() - assert await self._run_gate(context, result) is None - self._assert_gate_not_passed(result) + assert _should_wait_for_replica_provisioning(context) + + def test_waits_for_job_in_replica_without_hostname(self): + context = self._make_context( + [ + self._make_job( + replica_num=0, + status=JobStatus.PROVISIONING, + job_provisioning_data=MagicMock(hostname=None), + ) + ] + ) + assert _should_wait_for_replica_provisioning(context) @pytest.mark.parametrize("status", [JobStatus.FAILED, JobStatus.TERMINATED, JobStatus.ABORTED]) - async def test_defers_on_job_in_replica_terminated_before_provisioning(self, status): + def test_waits_for_job_in_replica_terminated_before_provisioning(self, status): + # The job is not terminated: the run pipeline decides whether the replica + # is retried or failed, and it can only do so once this job is unlocked. 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) + assert _should_wait_for_replica_provisioning(context) - async def test_does_not_defer_on_done_job_in_replica(self): + def test_does_not_wait_for_done_job_in_replica(self): # A job that already finished its work does not block its slower siblings. context = self._make_context( [ @@ -3497,11 +3640,9 @@ async def test_does_not_defer_on_done_job_in_replica(self): ) ] ) - result = _ProcessResult() - assert await self._run_gate(context, result) is None - self._assert_gate_passed(result) + assert not _should_wait_for_replica_provisioning(context) - async def test_does_not_defer_on_job_in_other_replica(self): + def test_does_not_wait_for_job_in_other_replica(self): context = self._make_context( [ self._make_job( @@ -3511,13 +3652,11 @@ async def test_does_not_defer_on_job_in_other_replica(self): ) ] ) - result = _ProcessResult() - assert await self._run_gate(context, result) is None - self._assert_gate_passed(result) + assert not _should_wait_for_replica_provisioning(context) @pytest.mark.asyncio -class TestPrepareStartupContextReplicaIpWait: +class TestPrepareRunnerSubmissionReplicaIpWait: def _service_configuration(self) -> ServiceConfiguration: return ServiceConfiguration.model_validate( { @@ -3574,11 +3713,6 @@ async def _fake_session_ctx(): "dstack._internal.server.background.pipeline_tasks.jobs_running.get_session_ctx", _fake_session_ctx, ), - patch( - "dstack._internal.server.background.pipeline_tasks.jobs_running.get_job_attached_volumes", - new_callable=AsyncMock, - return_value=[], - ), patch( "dstack._internal.server.background.pipeline_tasks.jobs_running.get_repo_creds", new_callable=AsyncMock, @@ -3611,7 +3745,7 @@ async def test_replica_ip_not_ready_defers(self): with ExitStack() as stack: for p in self._patches(): stack.enter_context(p) - out = await _prepare_startup_context(context=context, result=result) + out = await _prepare_runner_submission(context=context, result=result) assert out is None assert result.job_update_map == {} assert context.job.job_spec.commands == ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"] @@ -3623,7 +3757,7 @@ async def test_replica_ip_ready_substitutes(self): with ExitStack() as stack: for p in self._patches(): stack.enter_context(p) - out = await _prepare_startup_context(context=context, result=result) + out = await _prepare_runner_submission(context=context, result=result) assert out is not None assert context.job.job_spec.commands == ["echo 10.0.0.5"] @@ -3637,7 +3771,7 @@ async def test_nodes_ref_in_service_terminates(self): with ExitStack() as stack: for p in self._patches(): stack.enter_context(p) - out = await _prepare_startup_context(context=context, result=result) + out = await _prepare_runner_submission(context=context, result=result) assert out is None assert result.job_update_map.get("status") == JobStatus.TERMINATING assert "Illegal reference name" in (