Scope runner pull, metrics and stop to the job submission

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)

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
 }

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)
+	}
+}

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(

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
 }

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,

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,

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)

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))

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)

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-----