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 5dab47cca..8185621f8 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py @@ -264,19 +264,18 @@ async def _run_instance_check( check_instance_info: bool, ) -> InstanceCheck: ssh_private_keys = get_instance_ssh_private_keys(instance_model) - instance_check = await run_async( - _check_instance_inner, - ssh_private_keys, - job_provisioning_data, - None, - instance=instance_model, - check_instance_health=check_instance_health, - check_instance_info=check_instance_info, - ) - # May return False if fails to establish ssh connection. - if instance_check is False: - return InstanceCheck(reachable=False, message="SSH or tunnel error") - return instance_check + try: + return await run_async( + _check_instance_inner, + ssh_private_keys, + job_provisioning_data, + None, + instance=instance_model, + check_instance_health=check_instance_health, + check_instance_info=check_instance_info, + ) + except runner_client.PeerConnectionError as e: + return InstanceCheck(reachable=False, message=f"SSH or tunnel error: {e}") def _get_health_status_for_instance_check( 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 c14fe6666..511823fff 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -4,7 +4,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field from datetime import datetime, timedelta -from typing import Dict, Iterable, Literal, Optional, Sequence, Union +from typing import Dict, Iterable, Optional, Sequence from sqlalchemy import and_, exists, false, func, or_, select, true, update from sqlalchemy.ext.asyncio import AsyncSession @@ -842,6 +842,10 @@ async def _process_provisioning_status( ssh_user=ssh_user, ssh_key=user_ssh_key, ) + except client.PeerConnectionError as e: + # Expected while the instance is still booting + logger.debug("%s: shim is unreachable: %s", fmt(context.job_model), e) + success = False except client.ShimResponseError as e: logger.warning( "%s: shim did not accept the task submission: %s", fmt(context.job_model), e @@ -857,49 +861,53 @@ async def _process_provisioning_status( fmt(context.job_model), context.job_submission.age, ) - runner_availability = await run_async( - _get_runner_availability, - server_ssh_private_keys, - job_provisioning_data, - None, - ) - if runner_availability == _RunnerAvailability.AVAILABLE: - if not await _ensure_job_server_connection(context, result): - return - file_archives = await _get_job_file_archives( - archive_mappings=context.job.job_spec.file_archives, - user=context.run_model.user, - ) - code = await _get_job_code( - project=context.project, - repo=context.repo_model, - code_hash=_get_repo_code_hash(context.run, context.job), - ) - submit_result = await run_async( - _submit_job_to_runner, + try: + if await run_async( + _is_runner_available, server_ssh_private_keys, job_provisioning_data, None, - run=context.run, - job_model=context.job_model, - job=context.job, - jrd=get_job_runtime_data(context.job_model), - cluster_info=startup_context.cluster_info, - code=code, - file_archives=file_archives, - secrets=startup_context.secrets, - repo_credentials=startup_context.repo_creds, - router_env=startup_context.router_env, - success_if_not_available=False, - ) - if submit_result is not False: + ): + if not await _ensure_job_server_connection(context, result): + return + file_archives = await _get_job_file_archives( + archive_mappings=context.job.job_spec.file_archives, + user=context.run_model.user, + ) + code = await _get_job_code( + project=context.project, + repo=context.repo_model, + code_hash=_get_repo_code_hash(context.run, context.job), + ) + submit_result = await run_async( + _submit_job_to_runner, + server_ssh_private_keys, + job_provisioning_data, + None, + run=context.run, + job_model=context.job_model, + job=context.job, + jrd=get_job_runtime_data(context.job_model), + cluster_info=startup_context.cluster_info, + code=code, + file_archives=file_archives, + secrets=startup_context.secrets, + repo_credentials=startup_context.repo_creds, + router_env=startup_context.router_env, + success_if_not_available=False, + ) _apply_submit_job_to_runner_result( job_model=context.job_model, result=result, submit_result=submit_result, ) - if submit_result is not False and submit_result.success: - return + if submit_result.success: + return + except client.PeerConnectionError as e: + # Expected while the instance is still booting + logger.debug("%s: runner is unreachable: %s", fmt(context.job_model), e) + except client.RunnerResponseError as e: + logger.warning("%s: runner healthcheck failed: %s", fmt(context.job_model), e) provisioning_timeout = get_provisioning_timeout( backend_type=job_provisioning_data.get_base_backend(), @@ -939,12 +947,7 @@ async def _process_pulling_status( job_model=context.job_model, jrd=_get_result_job_runtime_data(context.job_model, result), ) - 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) - 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) @@ -974,56 +977,59 @@ async def _process_pulling_status( # _ShimPullingState.READY job_runtime_data = _get_result_job_runtime_data(context.job_model, result) - runner_availability = await run_async( - _get_runner_availability, + if not await run_async( + _is_runner_available, server_ssh_private_keys, job_provisioning_data, job_runtime_data, - ) - if runner_availability == _RunnerAvailability.UNAVAILABLE: + ): _reset_disconnected_at(context.job_model, result) return - if runner_availability == _RunnerAvailability.AVAILABLE: - if not await _ensure_job_server_connection(context, result): - return - file_archives = await _get_job_file_archives( - archive_mappings=context.job.job_spec.file_archives, - user=context.run_model.user, - ) - code = await _get_job_code( - project=context.project, - repo=context.repo_model, - code_hash=_get_repo_code_hash(context.run, context.job), - ) - submit_result = await run_async( - _submit_job_to_runner, - server_ssh_private_keys, - job_provisioning_data, - job_runtime_data, - run=context.run, - job_model=context.job_model, - job=context.job, - jrd=job_runtime_data, - cluster_info=startup_context.cluster_info, - code=code, - file_archives=file_archives, - secrets=startup_context.secrets, - repo_credentials=startup_context.repo_creds, - router_env=startup_context.router_env, - success_if_not_available=True, - ) - if submit_result is not False: - _apply_submit_job_to_runner_result( - job_model=context.job_model, - result=result, - submit_result=submit_result, - ) - if submit_result is not False and submit_result.success: - _reset_disconnected_at(context.job_model, result) - return - - # SSH tunnel failed or READY but runner submit failed — treat as disconnect + if not await _ensure_job_server_connection(context, result): + return + file_archives = await _get_job_file_archives( + archive_mappings=context.job.job_spec.file_archives, + user=context.run_model.user, + ) + code = await _get_job_code( + project=context.project, + repo=context.repo_model, + code_hash=_get_repo_code_hash(context.run, context.job), + ) + submit_result = await run_async( + _submit_job_to_runner, + server_ssh_private_keys, + job_provisioning_data, + job_runtime_data, + run=context.run, + job_model=context.job_model, + job=context.job, + jrd=job_runtime_data, + cluster_info=startup_context.cluster_info, + code=code, + file_archives=file_archives, + secrets=startup_context.secrets, + repo_credentials=startup_context.repo_creds, + router_env=startup_context.router_env, + success_if_not_available=True, + ) + _apply_submit_job_to_runner_result( + job_model=context.job_model, + result=result, + submit_result=submit_result, + ) + if submit_result.success: + _reset_disconnected_at(context.job_model, result) + return + except client.PeerConnectionError as e: + logger.debug("%s: instance is unreachable: %s", fmt(context.job_model), e) + except (client.ShimResponseError, client.RunnerResponseError) as e: + # Same outcome as a connection error, but the cause is logged instead of being + # silently indistinguishable. + logger.warning("%s: shim or runner answered unusably: %s", fmt(context.job_model), e) + + # The peer could not be reached, or it is READY but the runner submit failed _handle_instance_unreachable(context, result, job_provisioning_data) @@ -1039,20 +1045,30 @@ async def _process_running_status( fmt(context.job_model), context.job_submission.age, ) - process_running_result = await run_async( - _process_running, - server_ssh_private_keys, - job_provisioning_data, - context.job_submission.job_runtime_data, - run_model=context.run_model, - job_model=context.job_model, - ) - if process_running_result is not False: - result.job_update_map.update(process_running_result.job_update_map) - _reset_disconnected_at(context.job_model, result) + try: + process_running_result = await run_async( + _process_running, + server_ssh_private_keys, + job_provisioning_data, + context.job_submission.job_runtime_data, + run_model=context.run_model, + job_model=context.job_model, + ) + except client.PeerConnectionError as e: + logger.debug("%s: instance is unreachable: %s", fmt(context.job_model), e) + _handle_instance_unreachable(context, result, job_provisioning_data) + return + except client.RunnerResponseError as e: + # Same outcome as a connection error, but the cause is logged instead of being + # silently indistinguishable. + logger.warning( + "%s: runner failed to serve the pull request: %s", fmt(context.job_model), e + ) + _handle_instance_unreachable(context, result, job_provisioning_data) return - _handle_instance_unreachable(context, result, job_provisioning_data) + result.job_update_map.update(process_running_result.job_update_map) + _reset_disconnected_at(context.job_model, result) async def _ensure_job_server_connection( @@ -1504,13 +1520,6 @@ def _process_provisioning_with_shim( return True -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): WAITING = "waiting" READY = "ready" @@ -1527,19 +1536,16 @@ class _SyncShimPullingStateResult: @runner_ssh_tunnel -def _get_runner_availability(addresses: Mapping[int, client.LocalAddress]) -> _RunnerAvailability: +def _is_runner_available(addresses: Mapping[int, client.LocalAddress]) -> bool: + """ + Whether the runner has started and is ready to accept a job. + + A peer that answers the healthcheck with an error status or an unreadable body is not + expected to become a working runner, so `RunnerResponseError` propagates and the callers + treat it as an unreachable instance, unlike a runner that has not started yet. + """ runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) - try: - healthcheck_response = runner_client.healthcheck() - 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 - return _RunnerAvailability.AVAILABLE + return runner_client.healthcheck() is not None @runner_ssh_tunnel @@ -1547,7 +1553,7 @@ def _sync_shim_pulling_state( addresses: Mapping[int, client.LocalAddress], job_model: JobModel, jrd: Optional[JobRuntimeData] = None, -) -> Union[_SyncShimPullingStateResult, Literal[False]]: +) -> _SyncShimPullingStateResult: shim_client = client.ShimClient.from_address(addresses[DSTACK_SHIM_HTTP_PORT]) image_pull_progress: Optional[ImagePullProgress] = None if shim_client.is_api_v2_supported(): @@ -1638,7 +1644,7 @@ def _submit_job_to_runner( repo_credentials: Optional[RemoteRepoCreds], router_env: Optional[Dict[str, str]], success_if_not_available: bool, -) -> Union[_SubmitJobToRunnerResult, Literal[False]]: +) -> _SubmitJobToRunnerResult: logger.debug("%s: submitting job spec", fmt(job_model)) logger.debug( "%s: repo clone URL is %s", @@ -1705,14 +1711,10 @@ def _process_running( addresses: Mapping[int, client.LocalAddress], run_model: RunModel, job_model: JobModel, -) -> Union[_ProcessRunningResult, Literal[False]]: +) -> _ProcessRunningResult: runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) timestamp = job_model.runner_timestamp or 0 - try: - resp = runner_client.pull(timestamp) - except client.RunnerResponseError as e: - logger.warning("%s: runner failed to serve the pull request: %s", fmt(job_model), e) - return False + resp = runner_client.pull(timestamp) try: logs_services.write_logs( project=run_model.project, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py index c1d4e475d..c5cd9c7b6 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py @@ -816,13 +816,17 @@ async def _stop_container( ssh_private_keys: tuple[str, Optional[str]], ) -> bool: if job_provisioning_data.dockerized: - return await common.run_async( - _shim_submit_stop, - ssh_private_keys, - job_provisioning_data, - None, - job_model, - ) + try: + return await common.run_async( + _shim_submit_stop, + ssh_private_keys, + job_provisioning_data, + None, + job_model, + ) + except client.PeerConnectionError as e: + logger.debug("%s: can't stop container, shim is unreachable: %s", fmt(job_model), e) + return False return True diff --git a/src/dstack/_internal/server/background/scheduled_tasks/metrics.py b/src/dstack/_internal/server/background/scheduled_tasks/metrics.py index 10f87aa08..a083d5956 100644 --- a/src/dstack/_internal/server/background/scheduled_tasks/metrics.py +++ b/src/dstack/_internal/server/background/scheduled_tasks/metrics.py @@ -131,15 +131,14 @@ async def _collect_job_metrics(job_model: JobModel) -> Optional[JobMetricsPoint] jpd, jrd, ) + 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) + return None except Exception: logger.exception("Failed to collect job %s metrics", job_model.job_name) return None - if isinstance(res, bool): - # The job may already be terminated when collecting metrics - that's ok. - logger.warning("Failed to connect to job %s to collect metrics", job_model.job_name) - return None - if res is None: logger.debug( ( diff --git a/src/dstack/_internal/server/background/scheduled_tasks/prometheus_metrics.py b/src/dstack/_internal/server/background/scheduled_tasks/prometheus_metrics.py index 1b5d21c79..fcd908843 100644 --- a/src/dstack/_internal/server/background/scheduled_tasks/prometheus_metrics.py +++ b/src/dstack/_internal/server/background/scheduled_tasks/prometheus_metrics.py @@ -128,15 +128,16 @@ async def _collect_job_metrics(job_model: JobModel) -> Optional[str]: jrd, job_model.id, ) - except Exception: - logger.exception("Failed to collect job %s Prometheus metrics", job_model.job_name) - return None - - if isinstance(res, bool): + except client.PeerConnectionError as e: logger.warning( - "Failed to connect to job %s to collect Prometheus metrics", job_model.job_name + "Failed to connect to job %s to collect Prometheus metrics: %s", + job_model.job_name, + e, ) return None + except Exception: + logger.exception("Failed to collect job %s Prometheus metrics", job_model.job_name) + return None if res is None: # Either not supported by shim or exporter is not available diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index 260ce93f7..4e3d26a5b 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -14,7 +14,6 @@ from dstack._internal.core.errors import ( ResourceNotExistsError, ServerClientError, - SSHError, ) from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import validate_json_extra_ignore @@ -374,8 +373,8 @@ async def stop_runner(job_model: JobModel, instance_model: InstanceModel): jrd = get_job_runtime_data(job_model) try: await run_async(_stop_runner, ssh_private_keys, jpd, jrd, job_model) - except SSHError: - logger.debug("%s: failed to stop runner", fmt(job_model)) + except client.PeerConnectionError as e: + logger.debug("%s: failed to stop runner: %s", fmt(job_model), e) @runner_ssh_tunnel diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index 6171fd788..8ade87d87 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -71,10 +71,18 @@ """How much of an unusable response body is kept in the error message.""" +class PeerConnectionError(DstackError): + """ + The shim or the runner could not be reached: the SSH tunnel could not be established or + broke, or the request did not get through. The counterpart of the `*ResponseError` + families: nothing is known about the peer's state, and retrying may help. + """ + + 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 + but the answer cannot be used. Unlike `PeerConnectionError`, which means the request did not get through, repeating the same request is not expected to help. """ diff --git a/src/dstack/_internal/server/services/runner/ssh.py b/src/dstack/_internal/server/services/runner/ssh.py index 8b44303bd..7823e0901 100644 --- a/src/dstack/_internal/server/services/runner/ssh.py +++ b/src/dstack/_internal/server/services/runner/ssh.py @@ -1,6 +1,6 @@ import functools from collections.abc import Mapping -from typing import Callable, Literal, Optional, TypeVar, Union +from typing import Callable, Optional, TypeVar import requests from typing_extensions import Concatenate, ParamSpec @@ -8,7 +8,7 @@ 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 +from dstack._internal.server.services.runner.client import LocalAddress, PeerConnectionError from dstack._internal.server.services.runner.pool import ( InstanceConnection, PrivateKeyOrPair, @@ -23,7 +23,7 @@ def runner_ssh_tunnel( func: Callable[Concatenate[Mapping[int, LocalAddress], P], R], ) -> Callable[ Concatenate[PrivateKeyOrPair, JobProvisioningData, Optional[JobRuntimeData], P], - Union[Literal[False], R], + R, ]: """ A decorator that opens an SSH tunnel to the runner instance for port forwarding. @@ -39,9 +39,11 @@ 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. + Raises: + PeerConnectionError: the peer could not be reached. Errors reported by the peer's API + (`ShimError`, `RunnerError`) mean the shim or the runner was reached and answered, + so they propagate as they are, and the wrapped function or its caller must decide + what they mean for the job. """ @functools.wraps(func) @@ -51,15 +53,11 @@ def wrapper( job_runtime_data: Optional[JobRuntimeData], *args: P.args, **kwargs: P.kwargs, - ) -> Union[Literal[False], R]: - """ - Returns: - is successful - """ + ) -> R: if job_provisioning_data.hostname is None or job_provisioning_data.ssh_port is None: - # The callers may try to establish tunnels even if hostname/ssh_port is missing - # and rely on `False` being returned in this case. - return False + # The callers may try to establish tunnels even before the instance is fully + # provisioned, and rely on this being reported as an unreachable peer. + raise PeerConnectionError("the instance hostname or SSH port is not known yet") if not settings.SERVER_SSH_POOL_ENABLED or not job_provisioning_data.dockerized: # Connections from dstack-server to runner's sshd are expected to be short @@ -74,12 +72,12 @@ def wrapper( ephemeral=True, ) conn.open() - except SSHError: - return False + except SSHError as e: + raise PeerConnectionError(f"failed to open an SSH connection: {e}") from e try: return func(conn.forwarded_paths(), *args, **kwargs) - except requests.RequestException: - return False + except requests.RequestException as e: + raise PeerConnectionError(f"the request did not get through: {e}") from e finally: conn.close() @@ -89,6 +87,7 @@ def wrapper( # b) stale control socket file left by killed master. # (Because we cannot rely solely on connection errors from `func` – it may swallow the errors.) # but we still want a fast retry in case master dies mid-request. + error: Optional[requests.ConnectionError] = None for _ in range(2): conn = instance_connection_pool.get_or_open( ssh_private_key=ssh_private_key, @@ -96,14 +95,15 @@ def wrapper( jrd=job_runtime_data, ) if conn is None: - return False # couldn't establish at all + raise PeerConnectionError("failed to open an SSH connection") try: return func(conn.forwarded_paths(), *args, **kwargs) - except requests.ConnectionError: + except requests.ConnectionError as e: instance_connection_pool.drop(conn.key) # dead ssh connection, re-open - except requests.RequestException: + error = e + except requests.RequestException as e: # Reached the peer, e.g. a read timeout — do not re-open the ssh connection - return False - return False + raise PeerConnectionError(f"the request did not get through: {e}") from e + raise PeerConnectionError(f"the request did not get through: {error}") from error return wrapper diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 413691d42..88a2245a1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -57,7 +57,6 @@ _ProcessContext, _ProcessResult, _referenced_ips_ready, - _RunnerAvailability, _SubmitJobToRunnerResult, ) from dstack._internal.server.background.pipeline_tasks.runs import RunPipeline @@ -1202,10 +1201,10 @@ async def test_pulling_shim_uses_runtime_port_mapping_for_runner_calls( ] expected_ports = {10022: 32771, 10999: 32772} - def assert_runner_availability(_, __, job_runtime_data): + def assert_runner_available(_, __, job_runtime_data): assert job_runtime_data is not None assert job_runtime_data.ports == expected_ports - return _RunnerAvailability.AVAILABLE + return True def assert_submit_job_to_runner(_, __, job_runtime_data, **kwargs): assert job_runtime_data is not None @@ -1214,9 +1213,9 @@ def assert_submit_job_to_runner(_, __, job_runtime_data, **kwargs): with ( patch( - "dstack._internal.server.background.pipeline_tasks.jobs_running._get_runner_availability", - side_effect=assert_runner_availability, - ) as get_runner_availability_mock, + "dstack._internal.server.background.pipeline_tasks.jobs_running._is_runner_available", + side_effect=assert_runner_available, + ) as is_runner_available_mock, patch( "dstack._internal.server.background.pipeline_tasks.jobs_running._submit_job_to_runner", side_effect=assert_submit_job_to_runner, @@ -1234,7 +1233,7 @@ def assert_submit_job_to_runner(_, __, job_runtime_data, **kwargs): ): await _process_job(session, worker, job) ssh_tunnel_mock.assert_called_once() - get_runner_availability_mock.assert_called_once() + is_runner_available_mock.assert_called_once() submit_job_to_runner_mock.assert_called_once() await session.refresh(job) @@ -1485,8 +1484,8 @@ async def invalidate_lock(*args, **kwargs): with ( patch( - "dstack._internal.server.background.pipeline_tasks.jobs_running._get_runner_availability", - return_value=_RunnerAvailability.AVAILABLE, + "dstack._internal.server.background.pipeline_tasks.jobs_running._is_runner_available", + return_value=True, ), patch( "dstack._internal.server.background.pipeline_tasks.jobs_running._get_job_file_archives", diff --git a/src/tests/_internal/server/services/runner/test_ssh.py b/src/tests/_internal/server/services/runner/test_ssh.py index 9dd92218e..5586e7243 100644 --- a/src/tests/_internal/server/services/runner/test_ssh.py +++ b/src/tests/_internal/server/services/runner/test_ssh.py @@ -11,6 +11,7 @@ from dstack._internal.server.schemas.runner import HealthcheckResponse from dstack._internal.server.services.runner.client import ( LocalAddress, + PeerConnectionError, ShimResponseBodyError, ShimResponseError, ShimResponseStatusError, @@ -49,11 +50,17 @@ def conn(self): conn.forwarded_paths.return_value = FORWARDED_PATHS return conn - def call(self, func, dockerized: bool = False): + def call(self, func, dockerized: bool = False, jpd=None): decorated = runner_ssh_tunnel(func) - return decorated( - ("private_key", None), get_job_provisioning_data(dockerized=dockerized), None - ) + if jpd is None: + jpd = get_job_provisioning_data(dockerized=dockerized) + return decorated(("private_key", None), jpd, None) + + def test_missing_connection_details_raise(self): + jpd = get_job_provisioning_data().model_copy(update={"hostname": None}) + + with pytest.raises(PeerConnectionError, match="hostname or SSH port is not known"): + self.call(lambda addresses: "result", jpd=jpd) class TestEphemeralConnection(BaseRunnerSSHTunnelTest): @@ -70,10 +77,12 @@ 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): + def test_ssh_error_on_open_raises(self, conn): conn.open.side_effect = SSHError("no route") - assert self.call(lambda addresses: "result") is False + with pytest.raises(PeerConnectionError, match="failed to open an SSH connection") as exc: + self.call(lambda addresses: "result") + assert isinstance(exc.value.__cause__, SSHError) @pytest.mark.parametrize( "exc", @@ -83,11 +92,14 @@ def test_ssh_error_on_open_returns_false(self, conn): requests.exceptions.ChunkedEncodingError("truncated"), ], ) - def test_connection_errors_return_false(self, conn, exc): + def test_connection_errors_raise(self, conn, exc): def func(addresses: Mapping[int, LocalAddress]): raise exc - assert self.call(func) is False + with pytest.raises(PeerConnectionError, match="did not get through") as raised: + self.call(func) + assert raised.value.__cause__ is exc + # the connection is still released assert conn.close.call_count == 1 @pytest.mark.parametrize("error_cls", [ShimResponseStatusError, ShimResponseBodyError]) @@ -117,10 +129,14 @@ def test_returns_result(self, pool): assert pool.drop.call_count == 0 def test_connection_error_drops_and_retries_once(self, pool): + error = requests.ConnectionError("refused") + def func(addresses: Mapping[int, LocalAddress]): - raise requests.ConnectionError("refused") + raise error - assert self.call(func, dockerized=True) is False + with pytest.raises(PeerConnectionError, match="did not get through") as raised: + self.call(func, dockerized=True) + assert raised.value.__cause__ is error assert pool.get_or_open.call_count == 2 assert pool.drop.call_count == 2 @@ -128,10 +144,17 @@ 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 + with pytest.raises(PeerConnectionError, match="did not get through"): + self.call(func, dockerized=True) assert pool.get_or_open.call_count == 1 assert pool.drop.call_count == 0 + def test_unopenable_connection_raises(self, pool): + pool.get_or_open.return_value = None + + with pytest.raises(PeerConnectionError, match="failed to open an SSH connection"): + self.call(lambda addresses: "result", dockerized=True) + @pytest.mark.parametrize("error_cls", [ShimResponseStatusError, ShimResponseBodyError]) def test_api_errors_propagate(self, pool, error_cls): def func(addresses: Mapping[int, LocalAddress]):