From a6bb5a9cc0cdcb8f67e1ef740e0c85defe366c83 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Wed, 2 Sep 2026 15:37:19 +0000 Subject: [PATCH] Separate shim and runner API errors from connection errors `runner_ssh_tunnel` caught `DstackError` next to `requests.RequestException`, so an error reported by the peer's API was indistinguishable from a connectivity failure. `get_task()` on an unknown id answered 404 and surfaced as the same `False` the decorator returns when the SSH tunnel is down. `RunnerClient` made it worse by letting bare `requests.exceptions.HTTPError` escape, which is a `RequestException` subclass. Wrap runner HTTP errors as `RunnerHTTPError`, mirroring `ShimClient`, and share the wrapper implementation between the two clients. The tunnel now catches connection-level errors only, so anything escaping a client is either transport (`requests.RequestException`) or an answer from the peer (`ShimError`, `RunnerError`). `SSHError` is dropped from the catches as well: the tunnel is opened outside the guarded block, so it cannot be raised by the wrapped function. Each call site now decides what an API error means for the job instead of inheriting a silent `False`. Two of them change behavior: * `RunnerClient.healthcheck()` no longer masks HTTP errors. A peer answering with an error status is not a runner that has yet to start and is not expected to become one, so it counts as unreachable (`_RunnerAvailability.UNREACHABLE`) rather than unavailable. * `_process_running()` no longer reports a `LogStorageError` as a disconnect, which used to terminate the job as "Instance is unreachable" while the instance was healthy and only the log storage was down. `runner_timestamp` is not advanced, so the same logs and job state events are pulled again instead of being lost. The remaining sites keep their previous outcome and log the cause. `_get_gpu_driver()` and `_maybe_install_components()` gain catch-alls so optional metadata and opportunistic component installation cannot fail an instance check. Known gap, to be addressed separately: `pydantic.ValidationError` is a `ValueError`, so a response body that fails to validate still escapes. That is pre-existing -- neither `DstackError` nor `RequestException` covered it either. --- .../pipeline_tasks/instances/check.py | 25 ++- .../background/pipeline_tasks/jobs_running.py | 148 +++++++++++------- .../pipeline_tasks/jobs_terminating.py | 44 +++--- .../server/services/jobs/__init__.py | 3 +- .../server/services/runner/client.py | 137 +++++++++------- .../_internal/server/services/runner/ssh.py | 15 +- .../server/services/runner/test_client.py | 34 ++++ .../server/services/runner/test_ssh.py | 119 ++++++++++++++ 8 files changed, 378 insertions(+), 147 deletions(-) create mode 100644 src/tests/_internal/server/services/runner/test_ssh.py diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py index 793d0577d..5dab47cca 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py @@ -435,7 +435,10 @@ def _check_instance_inner( logger.warning("%s: error removing dangling tasks: %s", fmt(instance), exc) # There should be no shim API calls after this function call since it can request shim restart. - _maybe_install_components(instance, shim_client) + try: + _maybe_install_components(instance, shim_client) + except Exception as exc: + logger.warning("%s: error installing components: %s", fmt(instance), exc) return runner_client.healthcheck_response_to_instance_check( healthcheck_response, instance_health_response, @@ -453,29 +456,25 @@ def _get_gpu_driver( """ try: instance_info = shim_client.get_instance_info() - except requests.RequestException as exc: + return runner_client.instance_info_response_to_gpu_driver(instance_info) + except (requests.RequestException, runner_client.ShimError) as exc: logger.warning( "Instance %s: shim.get_instance_info(): request error: %s", instance_model.name, exc ) - return None - try: - return runner_client.instance_info_response_to_gpu_driver(instance_info) except ValueError as exc: logger.warning("Instance %s: unexpected instance info: %s", instance_model.name, exc) - return None + except Exception: + logger.exception( + "Instance %s: unexpected error retrieving the GPU driver", instance_model.name + ) + return None def _maybe_install_components( instance_model: InstanceModel, shim_client: runner_client.ShimClient, ) -> None: - try: - components = shim_client.get_components() - except requests.RequestException as exc: - logger.warning( - "Instance %s: shim.get_components(): request error: %s", instance_model.name, exc - ) - return + components = shim_client.get_components() if components is None: logger.debug("Instance %s: no components info", instance_model.name) return 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 2cfc889ff..159ff8426 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -826,21 +826,25 @@ async def _process_provisioning_status( if not server_settings.SSHPROXY_ENFORCED: ssh_user = job_provisioning_data.username user_ssh_key = get_or_error(context.run.run_spec.ssh_key_pub).strip() - success = await run_async( - _process_provisioning_with_shim, - server_ssh_private_keys, - job_provisioning_data, - None, - run=context.run, - job_model=context.job_model, - jrd=get_job_runtime_data(context.job_model), - jpd=job_provisioning_data, - volumes=startup_context.volumes, - registry_auth=context.job.job_spec.registry_auth, - public_keys=public_keys, - ssh_user=ssh_user, - ssh_key=user_ssh_key, - ) + try: + success = await run_async( + _process_provisioning_with_shim, + server_ssh_private_keys, + job_provisioning_data, + None, + run=context.run, + job_model=context.job_model, + jrd=get_job_runtime_data(context.job_model), + jpd=job_provisioning_data, + volumes=startup_context.volumes, + registry_auth=context.job.job_spec.registry_auth, + public_keys=public_keys, + 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) + success = False if success: _set_job_status(context.job_model, result, JobStatus.PULLING) result.job_update_map["skip_min_processing_interval"] = True @@ -924,14 +928,20 @@ async def _process_pulling_status( fmt(context.job_model), context.job_submission.age, ) - shim_state = await run_async( - _sync_shim_pulling_state, - server_ssh_private_keys, - job_provisioning_data, - None, - job_model=context.job_model, - jrd=_get_result_job_runtime_data(context.job_model, result), - ) + try: + shim_state = await run_async( + _sync_shim_pulling_state, + server_ssh_private_keys, + job_provisioning_data, + None, + job_model=context.job_model, + jrd=_get_result_job_runtime_data(context.job_model, result), + ) + except client.ShimHTTPError 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) + shim_state = False if shim_state is not False: if shim_state.job_runtime_data is not None: _set_job_runtime_data(result, shim_state.job_runtime_data) @@ -1495,6 +1505,8 @@ def _process_provisioning_with_shim( class _RunnerAvailability(enum.Enum): AVAILABLE = "available" UNAVAILABLE = "unavailable" + UNREACHABLE = "unreachable" + """Reached the runner port, but the peer is not a working runner.""" class _ShimPullingState(enum.Enum): @@ -1515,7 +1527,14 @@ class _SyncShimPullingStateResult: @runner_ssh_tunnel def _get_runner_availability(addresses: Mapping[int, client.LocalAddress]) -> _RunnerAvailability: runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) - if runner_client.healthcheck() is None: + 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) + return _RunnerAvailability.UNREACHABLE + if healthcheck_response is None: return _RunnerAvailability.UNAVAILABLE return _RunnerAvailability.AVAILABLE @@ -1630,31 +1649,37 @@ def _submit_job_to_runner( instance_env = None runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) - if runner_client.healthcheck() is None: - return _SubmitJobToRunnerResult(success=success_if_not_available) + try: + if runner_client.healthcheck() is None: + return _SubmitJobToRunnerResult(success=success_if_not_available) - runner_client.submit_job( - run=run, - job=job, - cluster_info=cluster_info, - # Do not send all the secrets since interpolation is already done by the server. - # TODO: Passing secrets may be necessary for filtering out secret values from logs. - secrets={}, - repo_credentials=repo_credentials, - instance_env=instance_env, - router_env=router_env, - ) - for archive_id, archive in file_archives: - logger.debug("%s: uploading file archive: %s", fmt(job_model), archive_id) - runner_client.upload_archive(archive_id, archive) - if code is None and not runner_client.is_code_upload_optional(): - # Old runner, we must call `/api/upload_code` to proceed - code = b"" - if code is not None: - logger.debug("%s: uploading code", fmt(job_model)) - runner_client.upload_code(code) - logger.debug("%s: starting job", fmt(job_model)) - job_info = runner_client.run_job() + runner_client.submit_job( + run=run, + job=job, + cluster_info=cluster_info, + # Do not send all the secrets since interpolation is already done by the server. + # TODO: Passing secrets may be necessary for filtering out secret values from logs. + secrets={}, + repo_credentials=repo_credentials, + instance_env=instance_env, + router_env=router_env, + ) + for archive_id, archive in file_archives: + logger.debug("%s: uploading file archive: %s", fmt(job_model), archive_id) + runner_client.upload_archive(archive_id, archive) + if code is None and not runner_client.is_code_upload_optional(): + # Old runner, we must call `/api/upload_code` to proceed + code = b"" + if code is not None: + logger.debug("%s: uploading code", fmt(job_model)) + 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) + return _SubmitJobToRunnerResult(success=False) if job_info is not None: if jrd is not None: jrd = jrd.model_copy( @@ -1680,14 +1705,25 @@ def _process_running( ) -> Union[_ProcessRunningResult, Literal[False]]: runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) timestamp = job_model.runner_timestamp or 0 - resp = runner_client.pull(timestamp) - logs_services.write_logs( - project=run_model.project, - run_name=run_model.run_name, - job_submission_id=job_model.id, - runner_logs=resp.runner_logs, - job_logs=resp.job_logs, - ) + try: + resp = runner_client.pull(timestamp) + except client.RunnerHTTPError as e: + logger.warning("%s: runner failed to serve the pull request: %s", fmt(job_model), e) + return False + try: + logs_services.write_logs( + project=run_model.project, + run_name=run_model.run_name, + job_submission_id=job_model.id, + runner_logs=resp.runner_logs, + job_logs=resp.job_logs, + ) + except logs_services.LogStorageError as e: + # The instance is reachable, the log storage is not, so this must not be reported as a + # disconnect. Nothing is updated: `runner_timestamp` is not advanced, so the same logs + # and job state events are pulled again next time instead of being lost. + logger.error("%s: failed to write logs: %s", fmt(job_model), e) + return _ProcessRunningResult() result = _ProcessRunningResult( job_update_map=_JobUpdateMap(runner_timestamp=resp.last_updated) ) 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 0a530a714..ace5d0c6f 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py @@ -830,25 +830,33 @@ async def _stop_container( def _shim_submit_stop(addresses: Mapping[int, client.LocalAddress], job_model: JobModel) -> bool: shim_client = client.ShimClient.from_address(addresses[DSTACK_SHIM_HTTP_PORT]) - resp = shim_client.healthcheck() - if resp is None: - logger.debug("%s: can't stop container, shim is not available yet", fmt(job_model)) + try: + resp = shim_client.healthcheck() + if resp is None: + logger.debug("%s: can't stop container, shim is not available yet", fmt(job_model)) + return False + + if shim_client.is_api_v2_supported(): + reason = ( + None + if job_model.termination_reason is None + else job_model.termination_reason.value + ) + shim_client.terminate_task( + task_id=job_model.id, + reason=reason, + message=job_model.termination_reason_message, + timeout=0, + ) + if not settings.SERVER_KEEP_SHIM_TASKS: + shim_client.remove_task(task_id=job_model.id) + else: + shim_client.stop(force=True) + except client.ShimHTTPError 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) return False - - if shim_client.is_api_v2_supported(): - reason = ( - None if job_model.termination_reason is None else job_model.termination_reason.value - ) - shim_client.terminate_task( - task_id=job_model.id, - reason=reason, - message=job_model.termination_reason_message, - timeout=0, - ) - if not settings.SERVER_KEEP_SHIM_TASKS: - shim_client.remove_task(task_id=job_model.id) - else: - shim_client.stop(force=True) return True diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index dafacc2ee..260ce93f7 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -387,7 +387,8 @@ def _stop_runner( runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) try: runner_client.stop() - except requests.RequestException: + 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 4b51f0bb6..51129998f 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -59,6 +59,68 @@ """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`. + + 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: + + try: + + except requests.exceptions.HTTPError as e: + raise RunnerHTTPError() from e + """ + + def __str__(self) -> str: + return self.message + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self.status_code})" + + @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 + + @property + def message(self) -> str: + cause = self._cause + if cause is None: + return "unknown_error" + return str(cause) + + @property + def _cause(self) -> Optional[requests.exceptions.HTTPError]: + cause = self.__cause__ + if isinstance(cause, requests.exceptions.HTTPError): + return cause + return None + + +class RunnerError(DstackError): + pass + + +class RunnerHTTPError(_HTTPErrorWrapper, RunnerError): + pass + + +class ShimError(DstackError): + pass + + +class ShimHTTPError(_HTTPErrorWrapper, ShimError): + pass + + +class ShimAPIVersionError(ShimError): + pass + + class RunnerClient: # `/api/upload_code` call is not required if there is no code _OPTIONAL_CODE_UPLOAD_MIN_VERSION = (0, 20, 17) @@ -99,6 +161,13 @@ def is_code_upload_optional(self) -> bool: return version_tuple is None or version_tuple >= self._OPTIONAL_CODE_UPLOAD_MIN_VERSION 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. + """ try: healthcheck_response = self._healthcheck() except requests.exceptions.RequestException: @@ -111,7 +180,7 @@ def get_metrics(self) -> Optional[MetricsResponse]: resp = self._session.get(self._url("/api/metrics"), timeout=REQUEST_TIMEOUT) if resp.status_code == 404: return None - resp.raise_for_status() + self._raise_for_status(resp) return validate_extra_ignore(MetricsResponse, resp.json()) def submit_job( @@ -161,7 +230,7 @@ def submit_job( headers={"Content-Type": "application/json"}, timeout=REQUEST_TIMEOUT, ) - resp.raise_for_status() + self._raise_for_status(resp) def upload_archive(self, id: uuid.UUID, file: Union[BinaryIO, bytes]): resp = self._session.post( @@ -169,17 +238,17 @@ def upload_archive(self, id: uuid.UUID, file: Union[BinaryIO, bytes]): files={"archive": (str(id), file)}, timeout=UPLOAD_CODE_REQUEST_TIMEOUT, ) - resp.raise_for_status() + self._raise_for_status(resp) def upload_code(self, file: Union[BinaryIO, bytes]): resp = self._session.post( self._url("/api/upload_code"), data=file, timeout=UPLOAD_CODE_REQUEST_TIMEOUT ) - resp.raise_for_status() + self._raise_for_status(resp) def run_job(self) -> Optional[JobInfoResponse]: resp = self._session.post(self._url("/api/run"), timeout=REQUEST_TIMEOUT) - resp.raise_for_status() + self._raise_for_status(resp) if not _is_json_response(resp): # Old runner or runner failed to get job info return None @@ -189,19 +258,25 @@ def pull(self, timestamp: int) -> PullResponse: resp = self._session.get( self._url("/api/pull"), params={"timestamp": timestamp}, timeout=REQUEST_TIMEOUT ) - resp.raise_for_status() + self._raise_for_status(resp) return validate_extra_ignore(PullResponse, resp.json()) def stop(self): resp = self._session.post(self._url("/api/stop"), timeout=REQUEST_TIMEOUT) - resp.raise_for_status() + self._raise_for_status(resp) def _url(self, path: str) -> str: return f"{self._base_url}/{path.lstrip('/')}" + def _raise_for_status(self, response: requests.Response) -> None: + try: + response.raise_for_status() + except requests.exceptions.HTTPError as e: + raise RunnerHTTPError() from e + def _healthcheck(self) -> HealthcheckResponse: resp = self._session.get(self._url("/api/healthcheck"), timeout=REQUEST_TIMEOUT) - resp.raise_for_status() + self._raise_for_status(resp) return validate_extra_ignore(HealthcheckResponse, resp.json()) def _negotiate(self, healthcheck_response: Optional[HealthcheckResponse] = None) -> None: @@ -214,52 +289,6 @@ def _negotiate(self, healthcheck_response: Optional[HealthcheckResponse] = None) self._negotiated = True -class ShimError(DstackError): - pass - - -class ShimHTTPError(ShimError): - """ - An HTTP error wrapper for `requests.exceptions.HTTPError`. Should be used as follows: - - try: - - except requests.exceptions.HTTPError as e: - raise ShimHTTPError() from e - """ - - def __str__(self) -> str: - return self.message - - def __repr__(self) -> str: - return f"{self.__class__.__name__}({self.status_code})" - - @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 - - @property - def message(self) -> str: - cause = self._cause - if cause is None: - return "unknown_error" - return str(cause) - - @property - def _cause(self) -> Optional[requests.exceptions.HTTPError]: - cause = self.__cause__ - if isinstance(cause, requests.exceptions.HTTPError): - return cause - return None - - -class ShimAPIVersionError(ShimError): - pass - - class ComponentList: _items: dict[ComponentName, ComponentInfo] diff --git a/src/dstack/_internal/server/services/runner/ssh.py b/src/dstack/_internal/server/services/runner/ssh.py index a7009df46..8b44303bd 100644 --- a/src/dstack/_internal/server/services/runner/ssh.py +++ b/src/dstack/_internal/server/services/runner/ssh.py @@ -5,7 +5,7 @@ import requests from typing_extensions import Concatenate, ParamSpec -from dstack._internal.core.errors import DstackError, SSHError +from dstack._internal.core.errors import SSHError from dstack._internal.core.models.runs import JobProvisioningData, JobRuntimeData from dstack._internal.server import settings from dstack._internal.server.services.runner.client import LocalAddress @@ -38,6 +38,10 @@ def runner_ssh_tunnel( There are no retries: a transient transport failure fails the call, and the callers must retry. In high-latency setups, tune `DSTACK_SERVER_SSH_CONNECT_TIMEOUT`. + + Only connection errors are converted to `False`. Errors reported by the peer's API + (`ShimError`, `RunnerError`) mean the shim or the runner was reached and answered, so they + propagate and the wrapped function or its caller must decide what they mean for the job. """ @functools.wraps(func) @@ -74,7 +78,7 @@ def wrapper( return False try: return func(conn.forwarded_paths(), *args, **kwargs) - except (DstackError, requests.RequestException): + except requests.RequestException: return False finally: conn.close() @@ -95,10 +99,11 @@ def wrapper( return False # couldn't establish at all try: return func(conn.forwarded_paths(), *args, **kwargs) - except (SSHError, requests.ConnectionError): + except requests.ConnectionError: instance_connection_pool.drop(conn.key) # dead ssh connection, re-open - except (DstackError, requests.RequestException): - return False # reached runner, app-level fail; don't re-open ssh connection + except requests.RequestException: + # Reached the peer, e.g. a read timeout — do not re-open the ssh connection + return False return False return wrapper diff --git a/src/tests/_internal/server/services/runner/test_client.py b/src/tests/_internal/server/services/runner/test_client.py index 0257dcd51..b2ead377a 100644 --- a/src/tests/_internal/server/services/runner/test_client.py +++ b/src/tests/_internal/server/services/runner/test_client.py @@ -4,6 +4,7 @@ from typing import Optional import pytest +import requests import requests_mock from gpuhunt import AcceleratorVendor @@ -32,6 +33,7 @@ ) from dstack._internal.server.services.runner.client import ( RunnerClient, + RunnerHTTPError, ShimClient, ShimHTTPError, _parse_version, @@ -140,6 +142,38 @@ 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") + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + with pytest.raises(RunnerHTTPError) 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 not isinstance(exc, requests.RequestException) + + def test_healthcheck_returns_none_on_connection_error(self, adapter: requests_mock.Adapter): + adapter.register_uri( + "GET", "/api/healthcheck", exc=requests.exceptions.ConnectionError("refused") + ) + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + assert client.healthcheck() is None + + def test_healthcheck_raises_on_http_error(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): + client.healthcheck() + + class TestShimClientNegotiate(BaseShimClientTest): @pytest.mark.parametrize( ["expected_shim_version", "expected_api_version"], diff --git a/src/tests/_internal/server/services/runner/test_ssh.py b/src/tests/_internal/server/services/runner/test_ssh.py new file mode 100644 index 000000000..97555a54a --- /dev/null +++ b/src/tests/_internal/server/services/runner/test_ssh.py @@ -0,0 +1,119 @@ +from collections.abc import Mapping +from pathlib import Path +from unittest.mock import Mock, patch + +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.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")} + + +class BaseRunnerSSHTunnelTest: + @pytest.fixture + def conn(self): + conn = Mock() + conn.forwarded_paths.return_value = FORWARDED_PATHS + return conn + + def call(self, func, dockerized: bool = False): + decorated = runner_ssh_tunnel(func) + return decorated( + ("private_key", None), get_job_provisioning_data(dockerized=dockerized), None + ) + + +class TestEphemeralConnection(BaseRunnerSSHTunnelTest): + """`dockerized=False` takes the branch that opens a fresh connection per call.""" + + @pytest.fixture(autouse=True) + def instance_connection(self, conn): + with patch( + "dstack._internal.server.services.runner.ssh.InstanceConnection", return_value=conn + ): + yield conn + + def test_returns_result(self, conn): + assert self.call(lambda addresses: "result") == "result" + assert conn.close.call_count == 1 + + def test_ssh_error_on_open_returns_false(self, conn): + conn.open.side_effect = SSHError("no route") + + assert self.call(lambda addresses: "result") is False + + @pytest.mark.parametrize( + "exc", + [ + requests.ConnectionError("refused"), + requests.ReadTimeout("too slow"), + requests.exceptions.ChunkedEncodingError("truncated"), + ], + ) + def test_connection_errors_return_false(self, conn, exc): + def func(addresses: Mapping[int, LocalAddress]): + raise exc + + assert self.call(func) is False + assert conn.close.call_count == 1 + + def test_api_errors_propagate(self, conn): + 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 + + with pytest.raises(ShimHTTPError): + self.call(func) + # the connection is still released + assert conn.close.call_count == 1 + + +class TestPooledConnection(BaseRunnerSSHTunnelTest): + """`dockerized=True` takes the branch that reuses pooled connections.""" + + @pytest.fixture(autouse=True) + def pool(self, conn): + with patch( + "dstack._internal.server.services.runner.ssh.instance_connection_pool" + ) as pool_mock: + pool_mock.get_or_open.return_value = conn + yield pool_mock + + def test_returns_result(self, pool): + assert self.call(lambda addresses: "result", dockerized=True) == "result" + assert pool.drop.call_count == 0 + + def test_connection_error_drops_and_retries_once(self, pool): + def func(addresses: Mapping[int, LocalAddress]): + raise requests.ConnectionError("refused") + + assert self.call(func, dockerized=True) is False + assert pool.get_or_open.call_count == 2 + assert pool.drop.call_count == 2 + + def test_other_connection_errors_do_not_retry(self, pool): + def func(addresses: Mapping[int, LocalAddress]): + raise requests.ReadTimeout("too slow") + + assert self.call(func, dockerized=True) is False + assert pool.get_or_open.call_count == 1 + assert pool.drop.call_count == 0 + + def test_api_errors_propagate(self, pool): + 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 + + with pytest.raises(ShimHTTPError): + self.call(func, dockerized=True) + assert pool.get_or_open.call_count == 1 + assert pool.drop.call_count == 0