[shim] Rework Run() into Start() and complete tasks in the background - #4203
Merged
Merged
Conversation
Previously, DockerRunner.Run() blocked in waitContainer() until the container exited, and only then set the final task status and ran its cleanup defers. As a result, nothing completed a task if the shim stopped running in the meantime: the volumes stayed mounted, the host SSH keys stayed in authorized_keys, and the task was never terminated with a proper reason. This is why the server only restarts the shim when all tasks are terminated, which blocks shim self-upgrade on instances running long-lived jobs. * `Run()` is renamed to `Start()` and returns as soon as the container is started. * The new `DockerRunner.ProcessTasks()` brings task states in line with their containers: it terminates tasks whose containers are not running anymore, reporting the exit code and the last log lines the same way as before, and releases the resources of terminated tasks. `ShimServer` calls it every second while it serves requests, starting before the first request is accepted, so that containers that exited while the shim was not running are processed first. * Each call makes one container list request for all tasks. A container is inspected only once it is not running anymore, since the list response reports the exit code as a part of a human-readable status string only. * Releasing the task resources is now a single idempotent function shared by `Start()`, `Terminate()`, `Remove()`, and `ProcessTasks()`, instead of the deferred calls of `Run()`. `Start()` owns the task until it returns (`Task.startInFlight`), so that the background job does not release the resources it is about to use. * Tasks are now restored from containers as running regardless of the container state, letting `ProcessTasks()` decide why a container finished. Previously, a container that exited while the shim was not running produced a task terminated with an empty termination reason. * A container that has never been started, which means the shim stopped running between creating and starting it, is now reported as `executor_error` instead of leaving the task running forever. The HTTP API is unchanged, as is the task status machine. Notes: - Tasks restored from containers have no task config, therefore their volumes and host SSH keys are still not cleaned up. Persisting the task state, which is needed for that, is a separate change (planned). - `Start()` still blocks while pulling the image, so tasks are only restart-safe in the running status. Adding it to the restart-safe statuses on the server is a separate change (not planned). Part-of: #4182 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
un-def
added a commit
that referenced
this pull request
Aug 26, 2026
TestDocker_SSHServer became flaky after #4203, which stopped Run() from waiting for the container: the tests now wait for the completion themselves, and the wait was capped at 60 seconds, while the container installs openssh-server over the network, which was measured to take from 24 to 52 seconds. #4203 also added three more Docker tests, all running in parallel with each other. A timed out wait also leaves the container behind, as the deferred Remove() call rejects a task that is still running. The runner dir of such a container is deleted along with the test temporary dir, so the following test runs, which adopt every task container on the host, keep failing to save the state of a task whose dir is gone, reporting errors unrelated to the tests being run. That is, one flaky run poisons the runs that follow it. * The wait timeout is raised to 150 seconds, which is still below the timeout of the tests themselves. * The deferred cleanup call now terminates the task before removing it, so that a failed test does not leave its container behind. * Docker tests no longer run in parallel: they share the Docker daemon, and a DockerRunner adopts every task container on the host, including the containers of the other tests. Two runners processing the same task write over each other's task state files, and one of them may stop a container that the other still considers running. * The task info is no longer read inside the require.Eventually() condition, which is called in a goroutine that may outlive a timed out call, racing with the assertions that follow it. Note that the Docker tests now take about 64 instead of 26 seconds, almost entirely due to TestDocker_SSHServer, which no longer overlaps with the others. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
un-def
added a commit
that referenced
this pull request
Aug 27, 2026
Since #4203, tasks outlive the shim process, but a restarted shim cannot finish them properly. Restoring a task from its container recovers the container ID, the GPUs, and the ports, but not the task config, so nothing tells the shim which volumes to unmount and which host SSH keys to remove. The reason a task was terminated for is lost as well, as it is not a property of the container. Each task now has a `task.json` file in its runner dir, holding what the container cannot tell: the task config, the termination reason and message, and whether the task resources are already released. It is written atomically and flushed to the disk on every task state change, starting right after the runner dir is created, that is, before any resource is acquired -- releasing a resource that was never acquired is a no-op, while the opposite order would leak. * Restored tasks get their config back, therefore their volumes are unmounted and their host SSH keys are removed once they finish. The volumes of a container created by an earlier shim version are still recovered from the container mounts; its host SSH keys are not recoverable. * A task with a recorded termination reason is restored as terminated, so that the reason reported by the server, e.g., `terminated_by_user`, is not replaced with the exit code of the container it stopped. * GPUs are locked back only for the restored tasks whose resources are not released yet, so that a task cleaned up before the restart does not hold them forever. * Task dirs that have a state file but no container are cleaned up on start. Normally, such a dir is left behind when the shim stops running before the container is created, e.g., while pulling the image. Dirs without a state file are left intact, as there is no way to tell whether they belong to a task. * The registry credentials are not persisted: the image is already pulled by the time the state is read back. Also fixes the runner dir rename fallback in remove(), which never worked: the trash name was built from the absolute path of the dir, so the rename target was a relative path under a `.trash-` dir that does not exist. Part-of: #4182 Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Previously, DockerRunner.Run() blocked in waitContainer() until the container exited, and only then set the final task status and ran its cleanup defers. As a result, nothing completed a task if the shim stopped running in the meantime: the volumes stayed mounted, the host SSH keys stayed in authorized_keys, and the task was never terminated with a proper reason. This is why the server only restarts the shim when all tasks are terminated, which blocks shim self-upgrade on instances running long-lived jobs.
Run()is renamed toStart()and returns as soon as the container is started.DockerRunner.ProcessTasks()brings task states in line with their containers: it terminates tasks whose containers are not running anymore, reporting the exit code and the last log lines the same way as before, and releases the resources of terminated tasks.ShimServercalls it every second while it serves requests, starting before the first request is accepted, so that containers that exited while the shim was not running are processed first.Start(),Terminate(),Remove(), andProcessTasks(), instead of the deferred calls ofRun().Start()owns the task until it returns (Task.startInFlight), so that the background job does not release the resources it is about to use.ProcessTasks()decide why a container finished. Previously, a container that exited while the shim was not running produced a task terminated with an empty termination reason.executor_errorinstead of leaving the task running forever.The HTTP API is unchanged, as is the task status machine.
Notes:
Start()still blocks while pulling the image, so tasks are only restart-safe in the running status. Adding it to the restart-safe statuses on the server is a separate change (not planned).Part-of: #4182