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