Skip to content

[shim] Commit task state changes via TaskStorage.Modify() - #4200

Merged
un-def merged 1 commit into
masterfrom
issue_4182_shim_task_storage_modify
Aug 24, 2026
Merged

[shim] Commit task state changes via TaskStorage.Modify()#4200
un-def merged 1 commit into
masterfrom
issue_4182_shim_task_storage_modify

Conversation

@un-def

@un-def un-def commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator

Groundwork for reworking DockerRunner.Run() so that it only starts a container, leaving container completion and cleanup to a periodic background job. Such a job takes task locks and commits task states concurrently with request handlers, which the current code is not prepared for.

  • TaskStorage.Update(), which overwrites the stored task with the caller's copy, is replaced with TaskStorage.Modify(), which applies a mutator to the stored task under the storage lock. The mutation is discarded if the mutator returns an error or if the resulting status transition is not allowed, that is, a partially applied mutation never reaches the storage. The transition is only checked if the status changes, therefore internal state can be committed without a status change, e.g., to publish container ports.
  • Task.Lock() no longer treats contention as a fatal error and waits for the lock instead. Task.TryLock() is added for callers that can retry later, such as the future background job. As long as there is at most one exclusive operation per task, contention means a bug, hence log.Fatal(); with a background job inspecting tasks every second, it is expected.
  • Run() commits the runner dir and the acquired GPUs as soon as they are acquired, rather than relying on the next status update to publish them, so that they can be cleaned up even if the local copy of the task is dropped.
  • DockerRunner.client is now a docker.APIClient interface, so that container states can be faked in tests.

The running -> running transition is no longer allowed, as same-status commits are not checked anymore. Docker tests now assert the final state committed for a successfully executed task, which was not covered before.

No behavior changes are expected, except that two concurrent exclusive operations on the same task no longer terminate the shim process.

Part-of: #4182

Groundwork for reworking DockerRunner.Run() so that it only starts a container,
leaving container completion and cleanup to a periodic background job. Such a job
takes task locks and commits task states concurrently with request handlers,
which the current code is not prepared for.

* `TaskStorage.Update()`, which overwrites the stored task with the caller's
  copy, is replaced with `TaskStorage.Modify()`, which applies a mutator to the
  stored task under the storage lock. The mutation is discarded if the mutator
  returns an error or if the resulting status transition is not allowed, that is,
  a partially applied mutation never reaches the storage. The transition is only
  checked if the status changes, therefore internal state can be committed
  without a status change, e.g., to publish container ports.
* `Task.Lock()` no longer treats contention as a fatal error and waits for the
  lock instead. `Task.TryLock()` is added for callers that can retry later, such
  as the future background job. As long as there is at most one exclusive
  operation per task, contention means a bug, hence `log.Fatal()`; with a
  background job inspecting tasks every second, it is expected.
* `Run()` commits the runner dir and the acquired GPUs as soon as they are
  acquired, rather than relying on the next status update to publish them, so
  that they can be cleaned up even if the local copy of the task is dropped.
* `DockerRunner.client` is now a `docker.APIClient` interface, so that container
  states can be faked in tests.

The `running -> running` transition is no longer allowed, as same-status commits
are not checked anymore. Docker tests now assert the final state committed for a
successfully executed task, which was not covered before.

No behavior changes are expected, except that two concurrent exclusive operations
on the same task no longer terminate the shim process.

Part-of: #4182
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@un-def
un-def merged commit 6d2312e into master Aug 24, 2026
27 checks passed
@un-def
un-def deleted the issue_4182_shim_task_storage_modify branch August 24, 2026 12:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant