From ae196cc71fd106e6006bb843c9b185f2e88b171e Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Thu, 27 Aug 2026 14:31:23 +0000 Subject: [PATCH 1/2] Allow restarting shim >= 0.21.3 with running tasks Closes: https://github.com/dstackai/dstack/issues/4182 --- .../server/services/runner/client.py | 45 +++++++++---------- 1 file changed, 22 insertions(+), 23 deletions(-) diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index 695afbca7..c4b778a25 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -318,6 +318,10 @@ class ShimClient: # `/api/shutdown` _SHUTDOWN_MIN_SHIM_VERSION = (0, 20, 1) + # Whether it is safe to restart the shim while at least one task is still running + # (not terminated). Other task statuses are not restart-safe regardless of the shim version + _RESTART_SAFE_RUNNING_STATUS_MIN_SHIM_VERSION = (0, 21, 3) + _shim_version_string: str _shim_version_tuple: Optional["_Version"] _api_version: int @@ -358,36 +362,28 @@ def is_api_v2_supported(self) -> bool: return self._api_version == 2 def is_instance_health_supported(self) -> bool: - if not self._negotiated: - self._negotiate() + shim_version_tuple = self.get_version_tuple() return ( - self._shim_version_tuple is None - or self._shim_version_tuple >= self._INSTANCE_HEALTH_MIN_SHIM_VERSION + shim_version_tuple is None + or shim_version_tuple >= self._INSTANCE_HEALTH_MIN_SHIM_VERSION ) def is_instance_info_supported(self) -> bool: - if not self._negotiated: - self._negotiate() + shim_version_tuple = self.get_version_tuple() return ( - self._shim_version_tuple is None - or self._shim_version_tuple >= self._INSTANCE_INFO_MIN_SHIM_VERSION + shim_version_tuple is None + or shim_version_tuple >= self._INSTANCE_INFO_MIN_SHIM_VERSION ) def are_components_supported(self) -> bool: - if not self._negotiated: - self._negotiate() + shim_version_tuple = self.get_version_tuple() return ( - self._shim_version_tuple is None - or self._shim_version_tuple >= self._COMPONENTS_MIN_SHIM_VERSION + shim_version_tuple is None or shim_version_tuple >= self._COMPONENTS_MIN_SHIM_VERSION ) def is_shutdown_supported(self) -> bool: - if not self._negotiated: - self._negotiate() - return ( - self._shim_version_tuple is None - or self._shim_version_tuple >= self._SHUTDOWN_MIN_SHIM_VERSION - ) + shim_version_tuple = self.get_version_tuple() + return shim_version_tuple is None or shim_version_tuple >= self._SHUTDOWN_MIN_SHIM_VERSION @overload def healthcheck(self) -> Optional[HealthcheckResponse]: ... @@ -671,11 +667,14 @@ def _negotiate(self, healthcheck_response: Optional[requests.Response] = None) - self._negotiated = True def _get_restart_safe_task_statuses(self) -> list[TaskStatus]: - # TODO: Rework shim's DockerRunner.Run() so that it does not wait for container termination - # (this at least requires replacing .waitContainer() with periodic polling of container - # statuses and moving some cleanup defer calls to .Terminate() and/or .Remove()) and add - # TaskStatus.RUNNING to the list of restart-safe task statuses for supported shim versions. - return [TaskStatus.TERMINATED] + statuses = [TaskStatus.TERMINATED] + shim_version_tuple = self.get_version_tuple() + if ( + shim_version_tuple is None + or shim_version_tuple >= self._RESTART_SAFE_RUNNING_STATUS_MIN_SHIM_VERSION + ): + statuses.append(TaskStatus.RUNNING) + return statuses def _make_session_and_base_url( From 20d4682a47f3ae52e862fa0f52643424d452ad91 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Thu, 27 Aug 2026 14:39:53 +0000 Subject: [PATCH 2/2] Add a helper method --- .../server/services/runner/client.py | 30 ++++++------------- 1 file changed, 9 insertions(+), 21 deletions(-) diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index c4b778a25..4b51f0bb6 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -362,28 +362,16 @@ def is_api_v2_supported(self) -> bool: return self._api_version == 2 def is_instance_health_supported(self) -> bool: - shim_version_tuple = self.get_version_tuple() - return ( - shim_version_tuple is None - or shim_version_tuple >= self._INSTANCE_HEALTH_MIN_SHIM_VERSION - ) + return self._check_min_version(self._INSTANCE_HEALTH_MIN_SHIM_VERSION) def is_instance_info_supported(self) -> bool: - shim_version_tuple = self.get_version_tuple() - return ( - shim_version_tuple is None - or shim_version_tuple >= self._INSTANCE_INFO_MIN_SHIM_VERSION - ) + return self._check_min_version(self._INSTANCE_INFO_MIN_SHIM_VERSION) def are_components_supported(self) -> bool: - shim_version_tuple = self.get_version_tuple() - return ( - shim_version_tuple is None or shim_version_tuple >= self._COMPONENTS_MIN_SHIM_VERSION - ) + return self._check_min_version(self._COMPONENTS_MIN_SHIM_VERSION) def is_shutdown_supported(self) -> bool: - shim_version_tuple = self.get_version_tuple() - return shim_version_tuple is None or shim_version_tuple >= self._SHUTDOWN_MIN_SHIM_VERSION + return self._check_min_version(self._SHUTDOWN_MIN_SHIM_VERSION) @overload def healthcheck(self) -> Optional[HealthcheckResponse]: ... @@ -666,13 +654,13 @@ def _negotiate(self, healthcheck_response: Optional[requests.Response] = None) - self._api_version = api_version self._negotiated = True + def _check_min_version(self, min_version: "_Version") -> bool: + current_version = self.get_version_tuple() + return current_version is None or current_version >= min_version + def _get_restart_safe_task_statuses(self) -> list[TaskStatus]: statuses = [TaskStatus.TERMINATED] - shim_version_tuple = self.get_version_tuple() - if ( - shim_version_tuple is None - or shim_version_tuple >= self._RESTART_SAFE_RUNNING_STATUS_MIN_SHIM_VERSION - ): + if self._check_min_version(self._RESTART_SAFE_RUNNING_STATUS_MIN_SHIM_VERSION): statuses.append(TaskStatus.RUNNING) return statuses