Skip to content

Allow restarting shim >= 0.21.3 with running tasks #4224

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
un-def merged 2 commits into from
Aug 27, 2026
Merged

Allow restarting shim >= 0.21.3 with running tasks #4224

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ae196cc
Allow restarting shim >= 0.21.3 with running tasks
un-def Aug 27, 2026
20d4682
Add a helper method
un-def Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 16 additions & 29 deletions src/dstack/_internal/server/services/runner/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -358,36 +362,16 @@ 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()
return (
self._shim_version_tuple is None
or self._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:
if not self._negotiated:
self._negotiate()
return (
self._shim_version_tuple is None
or self._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:
if not self._negotiated:
self._negotiate()
return (
self._shim_version_tuple is None
or self._shim_version_tuple >= self._COMPONENTS_MIN_SHIM_VERSION
)
return self._check_min_version(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
)
return self._check_min_version(self._SHUTDOWN_MIN_SHIM_VERSION)

@overload
def healthcheck(self) -> Optional[HealthcheckResponse]: ...
Expand Down Expand Up @@ -670,12 +654,15 @@ 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]:
# 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]
if self._check_min_version(self._RESTART_SAFE_RUNNING_STATUS_MIN_SHIM_VERSION):
statuses.append(TaskStatus.RUNNING)
return statuses


def _make_session_and_base_url(
Expand Down
Loading