From a0ad2142b70f5d68299e3f428223b5c81ce72b8d Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Fri, 4 Sep 2026 13:06:10 +0000 Subject: [PATCH] Split shim and runner response errors into status and body Follow-up to #4251, which separated API errors from connection errors but left three loose ends. `ShimHTTPError`/`RunnerHTTPError` were misnamed. We raise them only for status codes, while "HTTP error" suggests anything about the protocol, including transport. Rename them under a parent that says what the family means -- the peer answered, the answer is unusable, and repeating the request is not expected to help: ShimError |-- ShimAPIVersionError # our bug, stays loud `-- ShimResponseError |-- ShimResponseStatusError # 4xx/5xx as API error codes `-- ShimResponseBodyError # the body cannot be read Call sites catch the `*ResponseError` parent: none of them cares which leaf it was. Build the errors from the `Response` instead of wrapping the one from `raise_for_status()`. Its message carried a reason phrase that Go derives from the status code alone, a client/server split that 4xx vs 5xx already says, and a URL whose authority is always localhost or a percent-encoded socket path -- but not the body, which is where shim and runner put the actual message. Before: 404 Client Error: Not Found for url: http+unix://%2Ftmp%2F.../api/tasks/abc After: GET /api/tasks/abc: 404: Task not found Parse response bodies with pydantic instead of `Response.json()`. Besides saving a decode and an intermediate dict, this fixes a misclassification: `Response.json()` raises `requests.JSONDecodeError`, which is a `RequestException`, so a peer sending garbage was reported as a transport failure and retried. Both malformed JSON and a schema mismatch now raise `ValidationError`, wrapped as `*ResponseBodyError`. Nothing inside a client method raises `RequestException` any more except genuine transport, which is what `runner_ssh_tunnel` already documents. This also closes the `pydantic.ValidationError` leak noted in #4251: an unparsable response body no longer escapes the pipeline tasks. --- .../background/pipeline_tasks/jobs_running.py | 27 +-- .../pipeline_tasks/jobs_terminating.py | 4 +- .../server/services/runner/client.py | 157 +++++++++++++----- .../server/services/runner/test_client.py | 99 ++++++++--- .../server/services/runner/test_ssh.py | 50 ++++-- 5 files changed, 245 insertions(+), 92 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 159ff8426..c14fe6666 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -842,8 +842,10 @@ async def _process_provisioning_status( ssh_user=ssh_user, ssh_key=user_ssh_key, ) - except client.ShimHTTPError as e: - logger.warning("%s: shim refused the task submission: %s", fmt(context.job_model), e) + except client.ShimResponseError as e: + logger.warning( + "%s: shim did not accept the task submission: %s", fmt(context.job_model), e + ) success = False if success: _set_job_status(context.job_model, result, JobStatus.PULLING) @@ -937,7 +939,7 @@ async def _process_pulling_status( job_model=context.job_model, jrd=_get_result_job_runtime_data(context.job_model, result), ) - except client.ShimHTTPError as e: + except client.ShimResponseError as e: # Same outcome as a connection error, `_handle_instance_unreachable()` below, # but the cause is now logged instead of being silently indistinguishable. logger.warning("%s: shim failed to report the task state: %s", fmt(context.job_model), e) @@ -1529,10 +1531,11 @@ def _get_runner_availability(addresses: Mapping[int, client.LocalAddress]) -> _R runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) try: healthcheck_response = runner_client.healthcheck() - except client.RunnerHTTPError as e: - # Unlike a runner that has not started yet, a peer answering with an error status - # is not expected to become a working runner, so this counts as unreachable. - logger.warning("Runner healthcheck returned an error status: %s", e) + except client.RunnerResponseError as e: + # Unlike a runner that has not started yet, a peer that answers with an error status + # or an unreadable body is not expected to become a working runner, so this counts + # as unreachable. + logger.warning("Runner healthcheck failed: %s", e) return _RunnerAvailability.UNREACHABLE if healthcheck_response is None: return _RunnerAvailability.UNAVAILABLE @@ -1675,10 +1678,10 @@ def _submit_job_to_runner( runner_client.upload_code(code) logger.debug("%s: starting job", fmt(job_model)) job_info = runner_client.run_job() - except client.RunnerHTTPError as e: - # The runner answered with an error status, so retrying the same submission - # is not expected to help. - logger.warning("%s: runner refused the job submission: %s", fmt(job_model), e) + except client.RunnerResponseError as e: + # The runner answered, but unusably, so retrying the same submission is not + # expected to help. + logger.warning("%s: runner did not accept the job submission: %s", fmt(job_model), e) return _SubmitJobToRunnerResult(success=False) if job_info is not None: if jrd is not None: @@ -1707,7 +1710,7 @@ def _process_running( timestamp = job_model.runner_timestamp or 0 try: resp = runner_client.pull(timestamp) - except client.RunnerHTTPError as e: + except client.RunnerResponseError as e: logger.warning("%s: runner failed to serve the pull request: %s", fmt(job_model), e) return False try: 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 ace5d0c6f..c1d4e475d 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py @@ -852,10 +852,10 @@ def _shim_submit_stop(addresses: Mapping[int, client.LocalAddress], job_model: J shim_client.remove_task(task_id=job_model.id) else: shim_client.stop(force=True) - except client.ShimHTTPError as e: + except client.ShimResponseError as e: # The job is being terminated either way; dangling tasks are cleared later # by `remove_dangling_tasks_from_instance()` on instance checks. - logger.warning("%s: shim refused to stop the container: %s", fmt(job_model), e) + logger.warning("%s: shim failed to stop the container: %s", fmt(job_model), e) return False return True diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index 51129998f..6171fd788 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -6,6 +6,7 @@ from typing import BinaryIO, Dict, List, Literal, Optional, TypeVar, Union, overload import packaging.version +import pydantic import requests import requests.exceptions import requests_unixsocket @@ -13,7 +14,11 @@ from dstack._internal.core.consts import DSTACK_PROJECT_ENV from dstack._internal.core.errors import DstackError -from dstack._internal.core.models.common import CoreModel, NetworkMode, validate_extra_ignore +from dstack._internal.core.models.common import ( + CoreModel, + NetworkMode, + validate_json_extra_ignore, +) from dstack._internal.core.models.envs import Env from dstack._internal.core.models.instances import GpuDriverInfo from dstack._internal.core.models.repos.remote import RemoteRepoCreds @@ -59,65 +64,121 @@ """A local TCP port or a Unix domain socket path the client connects to.""" -class _HTTPErrorWrapper(DstackError): - """ - A base class for wrappers of `requests.exceptions.HTTPError`. +_M = TypeVar("_M", bound=CoreModel) +"""A response model parsed from a peer's response body.""" - Wrapping keeps API-level errors (the peer answered with a non-2xx status) distinct from - connection-level errors, which stay `requests.RequestException`. Subclasses are raised - as follows, so that `status_code` and `message` can read the original error: +_MAX_ERROR_BODY_BYTES = 512 +"""How much of an unusable response body is kept in the error message.""" - try: - - except requests.exceptions.HTTPError as e: - raise RunnerHTTPError() from e + +class _ResponseError(DstackError): + """ + A base implementation for the `*ResponseError` families: the peer was reached and answered, + but the answer cannot be used. Unlike `requests.RequestException`, which means the request + did not get through, repeating the same request is not expected to help. """ + def __init__(self, response: requests.Response) -> None: + super().__init__() + self.response = response + def __str__(self) -> str: return self.message def __repr__(self) -> str: - return f"{self.__class__.__name__}({self.status_code})" + return f"{self.__class__.__name__}({self.response.status_code})" + + @property + def message(self) -> str: + raise NotImplementedError + + @property + def _request_line(self) -> str: + request = self.response.request + # The full URL is noise: the authority is always localhost or a percent-encoded + # socket path, since the client always talks to a locally forwarded port. + return f"{request.method} {request.path_url}" + + @property + def _body(self) -> str: + content = self.response.content + text = content[:_MAX_ERROR_BODY_BYTES].decode("utf-8", "replace").strip() + if not text: + return "" + if len(content) > _MAX_ERROR_BODY_BYTES: + text += "..." + return text + + +class _ResponseStatusError(_ResponseError): + """ + The peer answered with an error status. Both shim and runner use HTTP status codes as API + error codes and put the message in the body, so the body is worth more than the status line + reason, which Go generates from the status code alone. + """ @property def status_code(self) -> int: - cause = self._cause - if cause is not None and cause.response is not None: - return cause.response.status_code - return 0 + return self.response.status_code @property def message(self) -> str: - cause = self._cause - if cause is None: - return "unknown_error" - return str(cause) + return f"{self._request_line}: {self.status_code}: {self._body}" + + +class _ResponseBodyError(_ResponseError): + """ + The peer answered with a body we cannot read: malformed JSON, or a payload that does not + match the expected schema, e.g. because shim or runner is too old or too new. + """ + + def __init__(self, response: requests.Response, error: pydantic.ValidationError) -> None: + super().__init__(response) + self.error = error @property - def _cause(self) -> Optional[requests.exceptions.HTTPError]: - cause = self.__cause__ - if isinstance(cause, requests.exceptions.HTTPError): - return cause - return None + def message(self) -> str: + errors = self.error.errors() + detail = f"{len(errors)} validation error(s)" + if errors: + location = ".".join(str(item) for item in errors[0]["loc"]) or "" + detail = f"{detail}, first at {location}: {errors[0]['msg']}" + return f"{self._request_line}: {detail}; body: {self._body}" class RunnerError(DstackError): pass -class RunnerHTTPError(_HTTPErrorWrapper, RunnerError): +class RunnerResponseError(RunnerError): pass -class ShimError(DstackError): +class RunnerResponseStatusError(_ResponseStatusError, RunnerResponseError): + pass + + +class RunnerResponseBodyError(_ResponseBodyError, RunnerResponseError): pass -class ShimHTTPError(_HTTPErrorWrapper, ShimError): +class ShimError(DstackError): pass class ShimAPIVersionError(ShimError): + """Raised when a caller uses an API the peer does not support. Signals a server-side bug.""" + + +class ShimResponseError(ShimError): + pass + + +class ShimResponseStatusError(_ResponseStatusError, ShimResponseError): + pass + + +class ShimResponseBodyError(_ResponseBodyError, ShimResponseError): pass @@ -164,9 +225,10 @@ def healthcheck(self) -> Optional[HealthcheckResponse]: """ Returns the healthcheck response, or `None` if the runner cannot be reached. - Only connection errors mean "not up yet". A non-2xx status means something is - listening that is not a working runner, which is not expected to resolve on its own, - so `RunnerHTTPError` propagates instead of being reported as unavailable. + Only connection errors mean "not up yet". An error status or a body we cannot read + means something is listening that is not a working runner, which is not expected to + resolve on its own, so `RunnerResponseError` propagates instead of being reported as + unavailable. """ try: healthcheck_response = self._healthcheck() @@ -181,7 +243,7 @@ def get_metrics(self) -> Optional[MetricsResponse]: if resp.status_code == 404: return None self._raise_for_status(resp) - return validate_extra_ignore(MetricsResponse, resp.json()) + return self._response(MetricsResponse, resp) def submit_job( self, @@ -252,14 +314,14 @@ def run_job(self) -> Optional[JobInfoResponse]: if not _is_json_response(resp): # Old runner or runner failed to get job info return None - return validate_extra_ignore(JobInfoResponse, resp.json()) + return self._response(JobInfoResponse, resp) def pull(self, timestamp: int) -> PullResponse: resp = self._session.get( self._url("/api/pull"), params={"timestamp": timestamp}, timeout=REQUEST_TIMEOUT ) self._raise_for_status(resp) - return validate_extra_ignore(PullResponse, resp.json()) + return self._response(PullResponse, resp) def stop(self): resp = self._session.post(self._url("/api/stop"), timeout=REQUEST_TIMEOUT) @@ -268,16 +330,20 @@ def stop(self): def _url(self, path: str) -> str: return f"{self._base_url}/{path.lstrip('/')}" - def _raise_for_status(self, response: requests.Response) -> None: + def _response(self, model_cls: type[_M], response: requests.Response) -> _M: try: - response.raise_for_status() - except requests.exceptions.HTTPError as e: - raise RunnerHTTPError() from e + return validate_json_extra_ignore(model_cls, response.content) + except pydantic.ValidationError as e: + raise RunnerResponseBodyError(response, e) from e + + def _raise_for_status(self, response: requests.Response) -> None: + if response.status_code >= HTTPStatus.BAD_REQUEST: + raise RunnerResponseStatusError(response) def _healthcheck(self) -> HealthcheckResponse: resp = self._session.get(self._url("/api/healthcheck"), timeout=REQUEST_TIMEOUT) self._raise_for_status(resp) - return validate_extra_ignore(HealthcheckResponse, resp.json()) + return self._response(HealthcheckResponse, resp) def _negotiate(self, healthcheck_response: Optional[HealthcheckResponse] = None) -> None: if healthcheck_response is None: @@ -658,16 +724,15 @@ def _request( self._raise_for_status(resp) return resp - _M = TypeVar("_M", bound=CoreModel) - def _response(self, model_cls: type[_M], response: requests.Response) -> _M: - return validate_extra_ignore(model_cls, response.json()) + try: + return validate_json_extra_ignore(model_cls, response.content) + except pydantic.ValidationError as e: + raise ShimResponseBodyError(response, e) from e def _raise_for_status(self, response: requests.Response) -> None: - try: - response.raise_for_status() - except requests.exceptions.HTTPError as e: - raise ShimHTTPError() from e + if response.status_code >= HTTPStatus.BAD_REQUEST: + raise ShimResponseStatusError(response) def _negotiate(self, healthcheck_response: Optional[requests.Response] = None) -> None: if healthcheck_response is None: diff --git a/src/tests/_internal/server/services/runner/test_client.py b/src/tests/_internal/server/services/runner/test_client.py index b2ead377a..ad253fbb2 100644 --- a/src/tests/_internal/server/services/runner/test_client.py +++ b/src/tests/_internal/server/services/runner/test_client.py @@ -33,9 +33,11 @@ ) from dstack._internal.server.services.runner.client import ( RunnerClient, - RunnerHTTPError, + RunnerResponseBodyError, + RunnerResponseStatusError, ShimClient, - ShimHTTPError, + ShimResponseBodyError, + ShimResponseStatusError, _parse_version, healthcheck_response_to_instance_check, instance_info_response_to_gpu_driver, @@ -142,22 +144,71 @@ def test_preserves_explicit_project_for_server_access(self, adapter: requests_mo assert adapter.last_request.json()["job_spec"]["env"]["DSTACK_PROJECT"] == "other" -class TestRunnerClientRaiseForStatus(BaseShimClientTest): - def test_wraps_http_error(self, adapter: requests_mock.Adapter): - adapter.register_uri("POST", "/api/stop", status_code=502, reason="Bad Gateway") +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" + ) client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) - with pytest.raises(RunnerHTTPError) as excinfo: + with pytest.raises(RunnerResponseStatusError) as excinfo: client.stop() exc = excinfo.value assert exc.status_code == 502 - assert exc.message.startswith("502 Server Error: Bad Gateway") - assert str(exc).startswith("502 Server Error: Bad Gateway") - assert repr(exc) == "RunnerHTTPError(502)" - # API-level errors must not be confused with connection errors + 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 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="") + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + with pytest.raises(RunnerResponseStatusError) as excinfo: + client.stop() + + assert str(excinfo.value) == "POST /api/stop: 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) + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + with pytest.raises(RunnerResponseStatusError) as excinfo: + client.stop() + + message = str(excinfo.value) + assert message.endswith("x" * 512 + "...") + assert len(message) < 600 + + def test_body_error_on_malformed_json(self, adapter: requests_mock.Adapter): + adapter.register_uri("GET", "/api/healthcheck", text="not json") + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + with pytest.raises(RunnerResponseBodyError) as excinfo: + client.healthcheck() + + exc = excinfo.value + assert str(exc).startswith("GET /api/healthcheck: 1 validation error(s)") + assert str(exc).endswith("body: not json") + # Parsing the body with pydantic rather than `Response.json()` keeps this out of the + # `requests.RequestException` hierarchy, where the tunnel would read it as transport. + assert not isinstance(exc, requests.RequestException) + + def test_body_error_on_schema_mismatch(self, adapter: requests_mock.Adapter): + adapter.register_uri("GET", "/api/healthcheck", json={"service": "dstack-runner"}) + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + with pytest.raises(RunnerResponseBodyError) as excinfo: + client.healthcheck() + + exc = excinfo.value + assert "first at version: Field required" in str(exc) + assert exc.error.error_count() == 1 + def test_healthcheck_returns_none_on_connection_error(self, adapter: requests_mock.Adapter): adapter.register_uri( "GET", "/api/healthcheck", exc=requests.exceptions.ConnectionError("refused") @@ -166,11 +217,11 @@ def test_healthcheck_returns_none_on_connection_error(self, adapter: requests_mo assert client.healthcheck() is None - def test_healthcheck_raises_on_http_error(self, adapter: requests_mock.Adapter): + def test_healthcheck_raises_on_error_status(self, adapter: requests_mock.Adapter): adapter.register_uri("GET", "/api/healthcheck", status_code=500) client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) - with pytest.raises(RunnerHTTPError): + with pytest.raises(RunnerResponseStatusError): client.healthcheck() @@ -214,19 +265,29 @@ def test( self.assert_request(adapter, 0, "GET", "/api/healthcheck") -class TestShimClientRaiseForStatus(BaseShimClientTest): - def test(self, client: ShimClient, adapter: requests_mock.Adapter): - adapter.register_uri("GET", "/test/path", status_code=502, reason="Bad Gateway") +class TestShimClientResponseErrors(BaseShimClientTest): + def test_status_error(self, client: ShimClient, adapter: requests_mock.Adapter): + adapter.register_uri( + "GET", "/test/path", status_code=502, reason="Bad Gateway", text="Task not found" + ) response = client._request("GET", "/test/path") - with pytest.raises(ShimHTTPError) as excinfo: + with pytest.raises(ShimResponseStatusError) as excinfo: client._raise_for_status(response) exc = excinfo.value assert exc.status_code == 502 - assert exc.message.startswith("502 Server Error: Bad Gateway") - assert str(exc).startswith("502 Server Error: Bad Gateway") - assert repr(exc) == "ShimHTTPError(502)" + assert str(exc) == "GET /test/path: 502: Task not found" + assert repr(exc) == "ShimResponseStatusError(502)" + + def test_body_error(self, client: ShimClient, adapter: requests_mock.Adapter): + adapter.register_uri("GET", "/test/path", json={"service": "dstack-shim"}) + response = client._request("GET", "/test/path") + + with pytest.raises(ShimResponseBodyError) as excinfo: + client._response(HealthcheckResponse, response) + + assert "first at version: Field required" in str(excinfo.value) @pytest.mark.shim_version("0.18.30") diff --git a/src/tests/_internal/server/services/runner/test_ssh.py b/src/tests/_internal/server/services/runner/test_ssh.py index 97555a54a..9dd92218e 100644 --- a/src/tests/_internal/server/services/runner/test_ssh.py +++ b/src/tests/_internal/server/services/runner/test_ssh.py @@ -2,18 +2,46 @@ from pathlib import Path from unittest.mock import Mock, patch +import pydantic import pytest import requests from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT from dstack._internal.core.errors import SSHError -from dstack._internal.server.services.runner.client import LocalAddress, ShimHTTPError +from dstack._internal.server.schemas.runner import HealthcheckResponse +from dstack._internal.server.services.runner.client import ( + LocalAddress, + ShimResponseBodyError, + ShimResponseError, + ShimResponseStatusError, +) from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel from dstack._internal.server.testing.common import get_job_provisioning_data FORWARDED_PATHS = {DSTACK_SHIM_HTTP_PORT: Path("/tmp/shim.sock")} +def _shim_response(status_code: int, content: bytes) -> requests.Response: + response = requests.Response() + response.status_code = status_code + response._content = content + response.request = requests.Request( + method="GET", url="http://localhost/api/tasks/id" + ).prepare() + return response + + +def _make_error(error_cls: type[ShimResponseError]) -> ShimResponseError: + """Builds either leaf of the `ShimResponseError` family without going through a client.""" + if error_cls is ShimResponseStatusError: + return ShimResponseStatusError(_shim_response(404, b"Task not found")) + try: + HealthcheckResponse.model_validate({}) + except pydantic.ValidationError as error: + return ShimResponseBodyError(_shim_response(200, b"{}"), error) + raise AssertionError("expected a validation error") + + class BaseRunnerSSHTunnelTest: @pytest.fixture def conn(self): @@ -62,14 +90,12 @@ def func(addresses: Mapping[int, LocalAddress]): assert self.call(func) is False assert conn.close.call_count == 1 - def test_api_errors_propagate(self, conn): + @pytest.mark.parametrize("error_cls", [ShimResponseStatusError, ShimResponseBodyError]) + def test_api_errors_propagate(self, conn, error_cls): def func(addresses: Mapping[int, LocalAddress]): - try: - raise requests.exceptions.HTTPError("404 Client Error: Not Found") - except requests.exceptions.HTTPError as e: - raise ShimHTTPError() from e + raise _make_error(error_cls) - with pytest.raises(ShimHTTPError): + with pytest.raises(ShimResponseError): self.call(func) # the connection is still released assert conn.close.call_count == 1 @@ -106,14 +132,12 @@ def func(addresses: Mapping[int, LocalAddress]): assert pool.get_or_open.call_count == 1 assert pool.drop.call_count == 0 - def test_api_errors_propagate(self, pool): + @pytest.mark.parametrize("error_cls", [ShimResponseStatusError, ShimResponseBodyError]) + def test_api_errors_propagate(self, pool, error_cls): def func(addresses: Mapping[int, LocalAddress]): - try: - raise requests.exceptions.HTTPError("404 Client Error: Not Found") - except requests.exceptions.HTTPError as e: - raise ShimHTTPError() from e + raise _make_error(error_cls) - with pytest.raises(ShimHTTPError): + with pytest.raises(ShimResponseError): self.call(func, dockerized=True) assert pool.get_or_open.call_count == 1 assert pool.drop.call_count == 0