Allow restarting shim >= 0.21.3 with running tasks (#4224)

Closes: #4182

Author: un-def

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,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]: ...
@@ -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(