From 81b9adba200d1c7c28bdeb47f2e3d90e876c93d2 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Wed, 28 Oct 2026 10:01:56 +0000 Subject: [PATCH] Scope runner pull, metrics and stop to the job submission The server could not tell the runner it submitted a job to from another one at the same address. A container restarted by the backend (e.g., Kubernetes' default `restartPolicy: Always`) answers /api/pull with no job states, which the server took for a healthy running job, while the new runner waited for a submission that never came, exited, and was restarted again. If the address went to another job's container, the server would collect that job's logs, states and metrics, and could stop it. The server now sends `job_submission_id` with /api/pull, /api/metrics and /api/stop, and the runner answers 409 if it runs another job submission or none. The server then terminates the job, collects no logs or metrics from that runner, and does not stop it. Old servers do not send the ID, and old runners ignore it, so both behave as before. Part-of: https://github.com/dstackai/dstack/issues/4420 Co-Authored-By: Claude Opus 5.5 (1M context) --- runner/internal/runner/api/http.go | 31 +++++++ runner/internal/runner/api/server.go | 6 +- runner/internal/runner/api/server_test.go | 54 +++++++++++- runner/internal/runner/executor/base.go | 2 + runner/internal/runner/executor/query.go | 4 + .../background/pipeline_tasks/jobs_running.py | 17 +++- .../pipeline_tasks/jobs_terminating.py | 7 +- .../background/scheduled_tasks/metrics.py | 10 ++- .../server/services/jobs/__init__.py | 5 +- .../server/services/runner/client.py | 42 ++++++++-- src/dstack/_internal/server/testing/common.py | 14 ++++ .../pipeline_tasks/test_running_jobs.py | 48 +++++++++++ .../pipeline_tasks/test_terminating_jobs.py | 83 ++++++++++++++++++- .../scheduled_tasks/test_metrics.py | 39 ++++++++- .../server/services/runner/test_client.py | 47 +++++++++-- 15 files changed, 380 insertions(+), 29 deletions(-) diff --git a/runner/internal/runner/api/http.go b/runner/internal/runner/api/http.go index 6efa907b82..e099deb2bd 100644 --- a/runner/internal/runner/api/http.go +++ b/runner/internal/runner/api/http.go @@ -209,6 +209,37 @@ func (s *Server) stopPostHandler(w http.ResponseWriter, r *http.Request) (interf return nil, nil } +// jobSubmissionScoped refuses a request meant for another job submission than the runner's own, +// or for any if the runner has none. The backend may have restarted the container, starting the +// runner anew, or given the job's address to another container. Either way, the runner the job +// was submitted to is gone, and the server must not take this one for it. Old servers do not +// send the ID, and their requests are served as before. +// +// The lock is not held through the handler, so the job submission may change in between. It is +// only set by /api/submit, which is refused once the job is run, and before that the scoped +// endpoints have no job to act on. +func (s *Server) jobSubmissionScoped( + handler func(http.ResponseWriter, *http.Request) (interface{}, error), +) func(http.ResponseWriter, *http.Request) (interface{}, error) { + return func(w http.ResponseWriter, r *http.Request) (interface{}, error) { + query := r.URL.Query() + if !query.Has("job_submission_id") { + return handler(w, r) + } + s.executor.RLock() + actual := s.executor.GetJobSubmissionId() + s.executor.RUnlock() + if query.Get("job_submission_id") != actual { + msg := "no job submitted" + if actual != "" { + msg = fmt.Sprintf("another job submitted: %s", actual) + } + return nil, &api.Error{Status: http.StatusConflict, Msg: msg} + } + return handler(w, r) + } +} + func isMaxBytesError(err error) bool { var maxBytesError *http.MaxBytesError return errors.As(err, &maxBytesError) diff --git a/runner/internal/runner/api/server.go b/runner/internal/runner/api/server.go index 979fe115bc..d0ea0922c4 100644 --- a/runner/internal/runner/api/server.go +++ b/runner/internal/runner/api/server.go @@ -64,13 +64,13 @@ func NewServer(ctx context.Context, address string, version string, ex executor. version: version, } r.AddHandler("GET", "/api/healthcheck", s.healthcheckGetHandler) - r.AddHandler("GET", "/api/metrics", s.metricsGetHandler) + r.AddHandler("GET", "/api/metrics", s.jobSubmissionScoped(s.metricsGetHandler)) r.AddHandler("POST", "/api/submit", s.submitPostHandler) r.AddHandler("POST", "/api/upload_archive", s.uploadArchivePostHandler) r.AddHandler("POST", "/api/upload_code", s.uploadCodePostHandler) r.AddHandler("POST", "/api/run", s.runPostHandler) - r.AddHandler("GET", "/api/pull", s.pullGetHandler) - r.AddHandler("POST", "/api/stop", s.stopPostHandler) + r.AddHandler("GET", "/api/pull", s.jobSubmissionScoped(s.pullGetHandler)) + r.AddHandler("POST", "/api/stop", s.jobSubmissionScoped(s.stopPostHandler)) r.AddHandler("GET", "/logs_ws", s.logsWsGetHandler) return s, nil } diff --git a/runner/internal/runner/api/server_test.go b/runner/internal/runner/api/server_test.go index 674cd7b10c..3da09672fc 100644 --- a/runner/internal/runner/api/server_test.go +++ b/runner/internal/runner/api/server_test.go @@ -20,8 +20,9 @@ import ( // fakeExecutor implements executor.Executor with the minimum needed to drive the handlers. type fakeExecutor struct { - mu sync.RWMutex - state string + mu sync.RWMutex + state string + jobSubmissionId string } func (e *fakeExecutor) SetJob(schemas.SubmitBody) {} @@ -37,6 +38,7 @@ func (e *fakeExecutor) GetJobWsLogsHistory() []schemas.LogEvent { return nil } func (e *fakeExecutor) GetRunnerState() string { return e.state } func (e *fakeExecutor) SetRunnerState(state string) { e.state = state } +func (e *fakeExecutor) GetJobSubmissionId() string { return e.jobSubmissionId } func (e *fakeExecutor) SetJobState(context.Context, schemas.JobState) {} func (e *fakeExecutor) SetJobStateWithTerminationReason( @@ -199,3 +201,51 @@ func TestStop_DoesNotMarkLogsFinished(t *testing.T) { t.Fatal("the first pull after the executor finished must be the final one") } } + +// /api/stop has a side effect, so the runner itself must refuse a request meant for another job: +// the backend may have restarted the container, or given the job's address to another one. +func TestStop_JobSubmissionScoped(t *testing.T) { + cases := []struct { + name string + query string + status int + stopped bool + }{ + {"same job", "?job_submission_id=job-1", http.StatusOK, true}, + {"another job", "?job_submission_id=job-2", http.StatusConflict, false}, + // Old servers do not send the ID + {"not sent", "", http.StatusOK, true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + s := newTestServer(t, executor.ServeLogs) + s.executor.(*fakeExecutor).jobSubmissionId = "job-1" + stopped := false + s.cancelRun = func() { stopped = true } + + rec := httptest.NewRecorder() + s.srv.Handler.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/api/stop"+tc.query, nil)) + + require.Equal(t, tc.status, rec.Code) + require.Equal(t, tc.stopped, stopped) + }) + } +} + +// A runner restarted along with its container has no job, and must not pass for the runner the +// job was submitted to. +func TestJobSubmissionScoped_NoJobSubmitted(t *testing.T) { + s := newTestServer(t, executor.WaitSubmit) + for _, endpoint := range []struct{ method, path string }{ + {http.MethodGet, "/api/pull"}, + {http.MethodGet, "/api/metrics"}, + {http.MethodPost, "/api/stop"}, + } { + rec := httptest.NewRecorder() + s.srv.Handler.ServeHTTP( + rec, httptest.NewRequest(endpoint.method, endpoint.path+"?job_submission_id=job-1", nil), + ) + require.Equal(t, http.StatusConflict, rec.Code, endpoint.path) + require.Equal(t, "no job submitted\n", rec.Body.String(), endpoint.path) + } +} diff --git a/runner/internal/runner/executor/base.go b/runner/internal/runner/executor/base.go index 9f5af2d0a3..ababb6c650 100644 --- a/runner/internal/runner/executor/base.go +++ b/runner/internal/runner/executor/base.go @@ -29,6 +29,8 @@ type Executor interface { GetRunnerState() string SetRunnerState(state string) + // Empty until a job is submitted + GetJobSubmissionId() string SetJobState(ctx context.Context, state schemas.JobState) SetJobStateWithTerminationReason( diff --git a/runner/internal/runner/executor/query.go b/runner/internal/runner/executor/query.go index f3acbf20ac..6264d89636 100644 --- a/runner/internal/runner/executor/query.go +++ b/runner/internal/runner/executor/query.go @@ -23,6 +23,10 @@ func (ex *RunExecutor) GetRunnerState() string { return ex.state } +func (ex *RunExecutor) GetJobSubmissionId() string { + return ex.jobSubmission.Id +} + type OrderedEvent interface { GetTimestamp() int64 } 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 342cdbc301..01d6628af1 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -1769,7 +1769,22 @@ def _process_running( ) -> _ProcessRunningResult: runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) timestamp = job_model.runner_timestamp or 0 - resp = runner_client.pull(timestamp) + try: + resp = runner_client.pull(job_model.id, timestamp) + except client.RunnerJobSubmissionMismatchError as e: + # The runner the job was submitted to is gone, and nothing will ever give the job to + # this one. Unlike an unreachable instance, this is not expected to resolve on its own. + logger.warning("%s: runner no longer runs the job: %s", fmt(job_model), e) + result = _ProcessRunningResult() + _terminate_job( + job_model=job_model, + job_update_map=result.job_update_map, + termination_reason=JobTerminationReason.INSTANCE_UNREACHABLE, + termination_reason_message=( + "Runner no longer runs the job, the container may have been restarted" + ), + ) + return result try: logs_services.write_logs( project=run_model.project, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py index 4b2ddf6e22..269e2e10c5 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py @@ -799,7 +799,12 @@ def _pull_job_logs( job_update_map: _JobUpdateMap, ) -> bool: runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) - resp = runner_client.pull(job_model.runner_timestamp or 0) + try: + resp = runner_client.pull(job_model.id, job_model.runner_timestamp or 0) + except client.RunnerJobSubmissionMismatchError: + # Not a failure to report: the job's logs went away with the runner it was submitted to + logger.debug("%s: runner no longer runs the job, no logs to collect", fmt(job_model)) + return True try: logs_services.write_logs( project=run_model.project, diff --git a/src/dstack/_internal/server/background/scheduled_tasks/metrics.py b/src/dstack/_internal/server/background/scheduled_tasks/metrics.py index 3d6695352c..103e7795e8 100644 --- a/src/dstack/_internal/server/background/scheduled_tasks/metrics.py +++ b/src/dstack/_internal/server/background/scheduled_tasks/metrics.py @@ -1,5 +1,6 @@ import asyncio import json +import uuid from collections.abc import Mapping from typing import List, Optional @@ -123,7 +124,13 @@ async def _collect_job_metrics(job_model: JobModel) -> Optional[JobMetricsPoint] ssh_private_keys, jpd, jrd, + job_model.id, ) + except client.RunnerJobSubmissionMismatchError: + # The metrics would be of another container. The job itself is terminated by + # `JobRunningPipeline` on its next pull from the runner. + logger.debug("Runner no longer runs job %s, skipping metrics", job_model.job_name) + return None except client.PeerConnectionError as e: # The job may already be terminated when collecting metrics - that's ok. logger.warning("Failed to connect to job %s to collect metrics: %s", job_model.job_name, e) @@ -160,6 +167,7 @@ async def _collect_job_metrics(job_model: JobModel) -> Optional[JobMetricsPoint] @runner_ssh_tunnel def _pull_runner_metrics( addresses: Mapping[int, client.LocalAddress], + job_submission_id: uuid.UUID, ) -> Optional[MetricsResponse]: runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) - return runner_client.get_metrics() + return runner_client.get_metrics(job_submission_id) diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index 0326be9aca..158358d2ae 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -380,7 +380,10 @@ def _stop_runner( logger.debug("%s: stopping runner", fmt(job_model)) runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) try: - runner_client.stop() + runner_client.stop(job_model.id) + except client.RunnerJobSubmissionMismatchError: + # The runner the job was submitted to is gone, and this one runs no job of ours to stop + logger.debug("%s: runner no longer runs the job, nothing to stop", fmt(job_model)) except (requests.RequestException, client.RunnerError): # Stopping the runner is best-effort: the job is being terminated either way. logger.exception("%s: failed to stop runner gracefully", fmt(job_model)) diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index 3407c7d8e3..f95bcf734a 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -170,6 +170,14 @@ class RunnerResponseBodyError(_ResponseBodyError, RunnerResponseError): pass +class RunnerJobSubmissionMismatchError(RunnerResponseStatusError): + """ + The runner refused a request meant for another job submission than its own, or it has none. + The backend may have restarted the container, starting the runner anew, or given the job's + address to another container. Either way, the runner the job was submitted to is gone. + """ + + class ShimError(DstackError): pass @@ -246,11 +254,15 @@ def healthcheck(self) -> Optional[HealthcheckResponse]: self._negotiate(healthcheck_response) return healthcheck_response - def get_metrics(self) -> Optional[MetricsResponse]: - resp = self._session.get(self._url("/api/metrics"), timeout=REQUEST_TIMEOUT) + def get_metrics(self, job_submission_id: uuid.UUID) -> Optional[MetricsResponse]: + resp = self._session.get( + self._url("/api/metrics"), + params={"job_submission_id": str(job_submission_id)}, + timeout=REQUEST_TIMEOUT, + ) if resp.status_code == 404: return None - self._raise_for_status(resp) + self._raise_for_job_submission_status(resp) return self._response(MetricsResponse, resp) def submit_job( @@ -324,16 +336,22 @@ def run_job(self) -> Optional[JobInfoResponse]: return None return self._response(JobInfoResponse, resp) - def pull(self, timestamp: int) -> PullResponse: + def pull(self, job_submission_id: uuid.UUID, timestamp: int) -> PullResponse: resp = self._session.get( - self._url("/api/pull"), params={"timestamp": timestamp}, timeout=REQUEST_TIMEOUT + self._url("/api/pull"), + params={"job_submission_id": str(job_submission_id), "timestamp": timestamp}, + timeout=REQUEST_TIMEOUT, ) - self._raise_for_status(resp) + self._raise_for_job_submission_status(resp) return self._response(PullResponse, resp) - def stop(self): - resp = self._session.post(self._url("/api/stop"), timeout=REQUEST_TIMEOUT) - self._raise_for_status(resp) + def stop(self, job_submission_id: uuid.UUID): + resp = self._session.post( + self._url("/api/stop"), + params={"job_submission_id": str(job_submission_id)}, + timeout=REQUEST_TIMEOUT, + ) + self._raise_for_job_submission_status(resp) def _url(self, path: str) -> str: return f"{self._base_url}/{path.lstrip('/')}" @@ -348,6 +366,12 @@ def _raise_for_status(self, response: requests.Response) -> None: if response.status_code >= HTTPStatus.BAD_REQUEST: raise RunnerResponseStatusError(response) + def _raise_for_job_submission_status(self, response: requests.Response) -> None: + """`_raise_for_status` for the endpoints that take `job_submission_id`.""" + if response.status_code == HTTPStatus.CONFLICT: + raise RunnerJobSubmissionMismatchError(response) + self._raise_for_status(response) + def _healthcheck(self) -> HealthcheckResponse: resp = self._session.get(self._url("/api/healthcheck"), timeout=REQUEST_TIMEOUT) self._raise_for_status(resp) diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index f4a2acaba5..29f2eac87b 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -3,10 +3,12 @@ from collections.abc import Callable from contextlib import contextmanager from datetime import datetime, timezone +from http import HTTPStatus from typing import Any, Dict, List, Literal, Optional, Union from uuid import UUID import gpuhunt +import requests from sqlalchemy import delete, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload @@ -141,6 +143,7 @@ get_default_permissions, set_default_permissions, ) +from dstack._internal.server.services.runner.client import RunnerJobSubmissionMismatchError from dstack._internal.server.services.users import get_token_hash @@ -1365,6 +1368,17 @@ async def clear_events(session: AsyncSession) -> None: await session.execute(delete(EventModel)) +def get_runner_job_submission_mismatch_error( + path: str = "/api/pull", +) -> RunnerJobSubmissionMismatchError: + """The error `RunnerClient` raises when the runner has no job, e.g., after a restart.""" + response = requests.Response() + response.status_code = HTTPStatus.CONFLICT + response._content = b"no job submitted" + response.request = requests.Request(method="GET", url=f"http://localhost{path}").prepare() + return RunnerJobSubmissionMismatchError(response) + + def get_private_key_string() -> str: return """ -----BEGIN RSA PRIVATE KEY----- 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 5e85487313..2de105396f 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 @@ -98,6 +98,7 @@ get_job_provisioning_data, get_job_runtime_data, get_run_spec, + get_runner_job_submission_mismatch_error, get_volume_configuration, list_events, ) @@ -1732,6 +1733,7 @@ async def test_updates_running_job( ) await _process_job(session, worker, job) ssh_tunnel_cls.assert_called_once() + runner_client_mock.pull.assert_called_once_with(job.id, 0) await session.refresh(job) assert job.status == JobStatus.RUNNING @@ -1762,6 +1764,52 @@ async def test_updates_running_job( assert job.exit_status == 0 assert job.runner_timestamp == 2 + async def test_terminates_job_if_runner_no_longer_runs_it( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ): + """ + A runner restarted along with the container is reachable but will never be given the + job, so the job is terminated right away rather than after the disconnect timeout. + """ + 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) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(dockerized=False), + instance=instance, + instance_assigned=True, + ) + job.runner_timestamp = 1 + await session.commit() + + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel"), + patch( + "dstack._internal.server.services.runner.client.RunnerClient.from_address" + ) as runner_client_cls, + ): + runner_client_mock = runner_client_cls.return_value + runner_client_mock.pull.side_effect = get_runner_job_submission_mismatch_error() + await _process_job(session, worker, job) + + await session.refresh(job) + assert job.status == JobStatus.TERMINATING + assert job.termination_reason == JobTerminationReason.INSTANCE_UNREACHABLE + assert job.disconnected_at is None + assert job.runner_timestamp == 1 + async def test_running_job_disconnect_retries_then_terminates( self, test_db, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py index 80e062246d..978a2720a1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py @@ -42,6 +42,7 @@ get_job_provisioning_data, get_job_runtime_data, get_run_spec, + get_runner_job_submission_mismatch_error, get_volume_configuration, get_volume_provisioning_data, list_events, @@ -322,6 +323,40 @@ async def test_closes_requested_server_access( remove_server_connection_mock.assert_awaited_once_with(job.id) + async def test_asks_runner_to_stop_job_submission( + self, + test_db, + session: AsyncSession, + worker: JobTerminatingWorker, + ssh_tunnel_mock: Mock, + runner_client_mock: Mock, + ): + """The runner refuses a stop meant for another job submission than its own.""" + project = await create_project(session=session) + user = await create_user(session=session) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run(session=session, project=project, repo=repo, user=user) + job = await create_job( + session=session, + run=run, + status=JobStatus.TERMINATING, + termination_reason=JobTerminationReason.TERMINATED_BY_USER, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + ) + job.graceful_termination_attempts = 0 + _lock_job(job) + await session.commit() + + await worker.process(_job_to_pipeline_item(job)) + + runner_client_mock.stop.assert_called_once_with(job.id) + async def test_stops_job_gracefully_before_terminating_container( self, test_db, session: AsyncSession, worker: JobTerminatingWorker ): @@ -528,7 +563,7 @@ async def test_collects_last_logs_before_terminating_container( ): await worker.process(_job_to_pipeline_item(job)) - runner_client_mock.pull.assert_called_once_with(1) + runner_client_mock.pull.assert_called_once_with(job.id, 1) write_logs.assert_called_once() assert write_logs.call_args.kwargs["job_logs"][0].message == b"the tail" assert stop_container.await_count == int(container_stopped) @@ -537,6 +572,52 @@ async def test_collects_last_logs_before_terminating_container( # Advanced, so the next pass does not collect the same logs again assert job.runner_timestamp == 3 + async def test_does_not_wait_for_logs_if_runner_no_longer_runs_job( + self, + test_db, + session: AsyncSession, + worker: JobTerminatingWorker, + ssh_tunnel_mock: Mock, + runner_client_mock: Mock, + ): + """The job's logs went away with the runner it was submitted to, nothing to wait for.""" + project = await create_project(session=session) + user = await create_user(session=session) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run(session=session, project=project, repo=repo, user=user) + job = await create_job( + session=session, + run=run, + status=JobStatus.TERMINATING, + termination_reason=JobTerminationReason.INSTANCE_UNREACHABLE, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + running_at=get_current_datetime() - timedelta(minutes=1), + ) + job.graceful_termination_attempts = 1 + job.remove_at = get_current_datetime() + timedelta(seconds=30) + job.runner_timestamp = 1 + _lock_job(job) + await session.commit() + + runner_client_mock.pull.side_effect = get_runner_job_submission_mismatch_error() + with patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating._stop_container", + new=AsyncMock(return_value=True), + ) as stop_container: + await worker.process(_job_to_pipeline_item(job)) + + runner_client_mock.pull.assert_called_once_with(job.id, 1) + stop_container.assert_awaited_once() + + await session.refresh(job) + assert job.runner_timestamp == 1 + @pytest.mark.parametrize( "error", [ diff --git a/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py b/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py index 408f85ba61..59ac8a9c8c 100644 --- a/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py +++ b/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py @@ -26,6 +26,7 @@ create_run, create_user, get_job_provisioning_data, + get_runner_job_submission_mismatch_error, ) pytestmark = pytest.mark.usefixtures("image_config_mock") @@ -84,11 +85,47 @@ async def test_collects_metrics(self, test_db, session: AsyncSession): ) await collect_metrics() SSHTunnelMock.assert_called_once() - runner_client_mock.get_metrics.assert_called_once() + runner_client_mock.get_metrics.assert_called_once_with(job.id) res = await session.execute(select(JobMetricsPoint)) metrics_point = res.scalar_one() assert metrics_point.job_id == job.id + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_skips_metrics_if_runner_no_longer_runs_job( + self, test_db, session: AsyncSession + ): + """The metrics would be of another container, e.g., one restarted in place of the job's.""" + user = await create_user(session=session, global_role=GlobalRole.USER) + project = await create_project(session=session, owner=user) + repo = await create_repo(session=session, project_id=project.id) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + run = await create_run(session=session, project=project, repo=repo, user=user) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(), + instance_assigned=True, + instance=instance, + ) + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel"), + patch( + "dstack._internal.server.services.runner.client.RunnerClient.from_address" + ) as RunnerClientMock, + ): + RunnerClientMock.return_value.get_metrics.side_effect = ( + get_runner_job_submission_mismatch_error(path="/api/metrics") + ) + await collect_metrics() + res = await session.execute(select(JobMetricsPoint)) + assert res.scalars().all() == [] + class TestDeleteMetrics: @pytest.mark.asyncio diff --git a/src/tests/_internal/server/services/runner/test_client.py b/src/tests/_internal/server/services/runner/test_client.py index dddd946c2f..9a0f4d81b0 100644 --- a/src/tests/_internal/server/services/runner/test_client.py +++ b/src/tests/_internal/server/services/runner/test_client.py @@ -1,5 +1,5 @@ import uuid -from collections.abc import Generator +from collections.abc import Callable, Generator from datetime import datetime, timezone from typing import Optional @@ -33,6 +33,7 @@ ) from dstack._internal.server.services.runner.client import ( RunnerClient, + RunnerJobSubmissionMismatchError, RunnerResponseBodyError, RunnerResponseStatusError, ShimClient, @@ -147,38 +148,38 @@ def test_preserves_explicit_project_for_server_access(self, adapter: requests_mo class TestRunnerClientResponseErrors(BaseShimClientTest): def test_status_error_reports_endpoint_status_and_body(self, adapter: requests_mock.Adapter): adapter.register_uri( - "POST", "/api/stop", status_code=502, reason="Bad Gateway", text="upstream is down" + "POST", "/api/run", status_code=502, reason="Bad Gateway", text="upstream is down" ) client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) with pytest.raises(RunnerResponseStatusError) as excinfo: - client.stop() + client.run_job() exc = excinfo.value assert exc.status_code == 502 assert exc.response.status_code == 502 # The status line reason is dropped: Go derives it from the status code alone, # while the body carries the message the handler actually wrote. - assert str(exc) == "POST /api/stop: 502: upstream is down" + assert str(exc) == "POST /api/run: 502: upstream is down" assert repr(exc) == "RunnerResponseStatusError(502)" # API errors must not be confused with connection errors assert not isinstance(exc, requests.RequestException) def test_status_error_without_body(self, adapter: requests_mock.Adapter): - adapter.register_uri("POST", "/api/stop", status_code=500, text="") + adapter.register_uri("POST", "/api/run", status_code=500, text="") client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) with pytest.raises(RunnerResponseStatusError) as excinfo: - client.stop() + client.run_job() - assert str(excinfo.value) == "POST /api/stop: 500: " + assert str(excinfo.value) == "POST /api/run: 500: " def test_status_error_truncates_long_body(self, adapter: requests_mock.Adapter): - adapter.register_uri("POST", "/api/stop", status_code=500, text="x" * 4096) + adapter.register_uri("POST", "/api/run", status_code=500, text="x" * 4096) client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) with pytest.raises(RunnerResponseStatusError) as excinfo: - client.stop() + client.run_job() message = str(excinfo.value) assert message.endswith("x" * 512 + "...") @@ -225,6 +226,34 @@ def test_healthcheck_raises_on_error_status(self, adapter: requests_mock.Adapter client.healthcheck() +class TestRunnerClientJobSubmissionScoped(BaseShimClientTest): + @pytest.mark.parametrize( + ("method", "path", "call"), + [ + pytest.param("GET", "/api/pull", lambda c, id_: c.pull(id_, 0), id="pull"), + pytest.param("GET", "/api/metrics", lambda c, id_: c.get_metrics(id_), id="metrics"), + pytest.param("POST", "/api/stop", lambda c, id_: c.stop(id_), id="stop"), + ], + ) + def test_raises_mismatch_error_on_conflict( + self, + adapter: requests_mock.Adapter, + method: str, + path: str, + call: Callable[[RunnerClient, uuid.UUID], object], + ): + adapter.register_uri(method, path, status_code=409, text="no job submitted") + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + job_submission_id = uuid.uuid4() + + with pytest.raises(RunnerJobSubmissionMismatchError) as excinfo: + call(client, job_submission_id) + + assert adapter.last_request is not None + assert adapter.last_request.qs["job_submission_id"] == [str(job_submission_id)] + assert str(excinfo.value).endswith(": 409: no job submitted") + + class TestShimClientNegotiate(BaseShimClientTest): @pytest.mark.parametrize( ["expected_shim_version", "expected_api_version"],