Make shim tasks survive a shim restart (rework DockerRunner.Run()) #4182

Problem

dstack-shim cannot survive a restart while any task is running, which has two consequences.

1. Shim self-upgrade is blocked on busy instances. The server only restarts the shim when every task is in a restart-safe status, and today that list is [terminated] only (ShimClient._get_restart_safe_task_statuses, is_safe_to_restart, used by
_maybe_install_shim). An instance running a long-lived service or a multi-day training job therefore keeps its old shim binary indefinitely, even after the new one has been downloaded and verified.

2. If the shim does go away mid-task, the task is never finalized. The cause is DockerRunner.Run(): it blocks in waitContainer() until the container exits and only then sets the final task status and runs its cleanup defers. So when the shim process is replaced, the service is restarted, or the host reboots:

• the defers never run — volumes stay mounted (unmountVolumes) and the run's host SSH keys stay in authorized_keys (AuthorizedKeys.RemovePublicKeys);
• after the restart, restoreStateFromContainers rebuilds the task as running, but nothing waits for the container anymore, so the task never moves to terminated on its own and its exit code and last log lines are lost — the server never sees done_by_runner / container_exited_with_error;
• GPUs stay locked for such a task until the server explicitly terminates it.

A graceful POST /api/shutdown doesn't help either: Run() goroutines are spawned detached from the server's background-job group, so shutdown neither waits for them nor cancels them.

Solution

Make Run() a start-only operation and move completion detection and cleanup into a periodic, restartable step that reads its inputs from the Docker daemon and from disk rather than from a live goroutine's stack.

• Run() → Start(): prepare, pull, create, start, set the task running, return. It stays blocking for the image pull, so pulling/creating remain restart-unsafe; only running becomes restart-safe.
• New DockerRunner.ProcessTasks(), called once per second by a ShimServer background job (started in Serve(), cancelled and awaited in Shutdown()). Per tick it makes a single labeled ContainerList call and inspects a container only when it is no longer running, then sets the task status (same done_by_runner / container_exited_with_error mapping and last-log message as today) and releases the task's resources.
• Resource release becomes one idempotent function shared by ProcessTasks, Terminate, Remove, and Start's failure path, instead of per-step defers.
• Persist the cleanup inputs that cannot be recovered from the daemon (host SSH keys, volume list, termination outcome) to a small per-task file under the task's runner dir, and rehydrate them in restoreStateFromContainers, so cleanup completes even for a task that finished while the shim was down. Sweep orphaned task dirs on startup.
• Finally, add running to the server's restart-safe task statuses, gated on the shim version that includes the above.

Implementation outline:

• TaskStorage.Modify(id, fn) replacing Update; Task.Lock made blocking plus a new TryLock (a periodic background job makes lock contention normal, and the current Lock treats it as fatal)
• Run → Start, ProcessTasks, shared cleanup, 1s ticker
• Per-task state file, restore rehydration, orphan-dir sweep
• Server: running as a restart-safe status for supported shim versions

This also fixes, as a side effect, a task restored after a restart reporting no termination reason, and a container left running when a task is terminated in the window between container creation and the running commit.

Workaround

Wait until the instance is idle. The server already refuses to restart a shim with non-terminated tasks, so the upgrade happens on its own once all jobs on the instance have finished; terminating the run (or the fleet) forces it earlier.

Would you like to help us implement this feature by sending a PR?

Yes