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_jobs_running.py b/src/tests/_internal/server/background/pipeline_tasks/test_jobs_running.py index 5e85487313..2de105396f 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_jobs_running.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_jobs_running.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_jobs_terminating.py b/src/tests/_internal/server/background/pipeline_tasks/test_jobs_terminating.py index 80e062246d..978a2720a1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_jobs_terminating.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_jobs_terminating.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"],