Separate shim and runner API errors from connection errors

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

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

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
 
 

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