From 477578c6a350152f8667f591d503dde7427f9649 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Mon, 24 Aug 2026 12:03:30 +0000 Subject: [PATCH] [shim] Commit task state changes via TaskStorage.Modify() 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: https://github.com/dstackai/dstack/issues/4182 Co-Authored-By: Claude Opus 5 (1M context) --- runner/internal/shim/docker.go | 89 ++++++++++++------- runner/internal/shim/docker_test.go | 14 +++ runner/internal/shim/task.go | 71 ++++++++++------ runner/internal/shim/task_test.go | 127 ++++++++++++++++++++++++---- 4 files changed, 231 insertions(+), 70 deletions(-) diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index 7fbda1458d..1864b5f52c 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -148,7 +148,7 @@ func (t *PullTracker) Progress() *ImagePullProgress { } type DockerRunner struct { - client *docker.Client + client docker.APIClient dockerParams DockerParameters dockerInfo dockersystem.Info baseEnv []string @@ -366,6 +366,37 @@ func (d *DockerRunner) Submit(ctx context.Context, cfg TaskConfig) error { return nil } +// commit applies mutate to the stored task and, on success, updates the local copy +// of the task accordingly. mutate is called before the local copy is updated, so it +// can read the local copy to publish fields set by the caller, e.g., containerID. +// The local copy is left intact if the update is rejected. +func (d *DockerRunner) commit(task *Task, mutate func(*Task)) error { + updatedTask, err := d.tasks.Modify(task.ID, func(t *Task) error { + mutate(t) + return nil + }) + if err != nil { + return err + } + *task = updatedTask + return nil +} + +// commitTerminated commits the terminated status set on the local copy of the task. +// Used by operations that set the status on their failure paths, where committing it +// at every such path would be too verbose. +func (d *DockerRunner) commitTerminated(ctx context.Context, task *Task) { + if task.Status != TaskStatusTerminated { + // the last successful commit is the actual state, nothing to commit + return + } + if err := d.commit(task, func(t *Task) { + t.SetStatusTerminated(task.TerminationReason, task.TerminationMessage) + }); err != nil && !errors.Is(err, ErrNotFound) { + log.Error(ctx, "failed to commit terminated status", "task", task.ID, "err", err) + } +} + func (d *DockerRunner) Run(ctx context.Context, taskID string) error { task, ok := d.tasks.Get(taskID) if !ok { @@ -377,17 +408,9 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error { return fmt.Errorf("%w: cannot run task %s with %s status", ErrRequest, task.ID, task.Status) } - defer func() { - if err := d.tasks.Update(task); err != nil { - if currentTask, ok := d.tasks.Get(task.ID); ok && currentTask.Status != task.Status { - // ignore error if task is gone or status has not changed, e.g., terminated -> terminated - log.Error(ctx, "failed to update", "task", task.ID, "err", err) - } - } - }() + defer func() { d.commitTerminated(ctx, &task) }() - task.SetStatusPreparing() - if err := d.tasks.Update(task); err != nil { + if err := d.commit(&task, func(t *Task) { t.SetStatusPreparing() }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } @@ -398,25 +421,32 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error { if err != nil { return fmt.Errorf("make runner dir: %w", err) } - task.runnerDir = runnerDir log.Debug(ctx, "runner dir", "task", task.ID, "path", runnerDir) + // Resources are committed as soon as they are acquired, so that they are not + // lost if the task is updated by another goroutine, e.g., terminated by the server + if err := d.commit(&task, func(t *Task) { t.runnerDir = runnerDir }); err != nil { + return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) + } + var gpuIDs []string if cfg.GPU != 0 { - gpuIDs, err := d.gpuLock.Acquire(ctx, cfg.GPU) + gpuIDs, err = d.gpuLock.Acquire(ctx, cfg.GPU) if err != nil { log.Error(ctx, err.Error()) task.SetStatusTerminated(string(types.TerminationReasonExecutorError), err.Error()) return fmt.Errorf("acquire GPU: %w", err) } - task.gpuIDs = gpuIDs log.Debug(ctx, "acquired GPU(s)", "task", task.ID, "gpus", gpuIDs) defer func() { - releasedGpuIDs := d.gpuLock.Release(ctx, task.gpuIDs) + releasedGpuIDs := d.gpuLock.Release(ctx, gpuIDs) log.Debug(ctx, "released GPU(s)", "task", task.ID, "gpus", releasedGpuIDs) }() } else { - task.gpuIDs = []string{} + gpuIDs = []string{} + } + if err := d.commit(&task, func(t *Task) { t.gpuIDs = gpuIDs }); err != nil { + return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } if len(cfg.HostSshKeys) > 0 { @@ -458,8 +488,7 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error { log.Debug(ctx, "Pulling image") pullCtx, cancelPull := context.WithTimeout(ctx, ImagePullTimeout) defer cancelPull() - task.SetStatusPulling(cancelPull) - if err := d.tasks.Update(task); err != nil { + if err := d.commit(&task, func(t *Task) { t.SetStatusPulling(cancelPull) }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } // Although it's called "runner dir", we also use it for shim task-related data. @@ -473,8 +502,7 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error { } log.Debug(ctx, "Creating container", "task", task.ID, "name", task.containerName) - task.SetStatusCreating() - if err := d.tasks.Update(task); err != nil { + if err := d.commit(&task, func(t *Task) { t.SetStatusCreating() }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } if err := d.createContainer(ctx, &task, createContainerOptions{}); err != nil { @@ -485,8 +513,11 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error { } log.Debug(ctx, "Running container", "task", task.ID, "name", task.containerName) - task.SetStatusRunning() - if err := d.tasks.Update(task); err != nil { + if err := d.commit(&task, func(t *Task) { + // createContainer sets the containerID field + t.containerID = task.containerID + t.SetStatusRunning() + }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } err = d.startContainer(ctx, &task) @@ -506,8 +537,12 @@ func (d *DockerRunner) Run(ctx context.Context, taskID string) error { } } if err == nil { - // startContainer sets `ports` field, committing update - if err := d.tasks.Update(task); err != nil { + if err := d.commit(&task, func(t *Task) { + // startContainer sets the ports field, the retry above may have + // replaced the container + t.containerID = task.containerID + t.ports = task.ports + }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } err = d.waitContainer(ctx, &task) @@ -541,11 +576,7 @@ func (d *DockerRunner) Terminate(ctx context.Context, taskID string, timeout uin } task.Lock(ctx) defer func() { task.Release(ctx) }() - defer func() { - if err := d.tasks.Update(task); err != nil { - log.Error(ctx, "failed to update task", "task", task.ID, "err", err) - } - }() + defer func() { d.commitTerminated(ctx, &task) }() return d.terminate(ctx, &task, timeout, reason, message) } diff --git a/runner/internal/shim/docker_test.go b/runner/internal/shim/docker_test.go index 3f807dd374..7c634fa604 100644 --- a/runner/internal/shim/docker_test.go +++ b/runner/internal/shim/docker_test.go @@ -12,6 +12,7 @@ import ( "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/mount" "github.com/dstackai/dstack/runner/internal/common/gpu" + "github.com/dstackai/dstack/runner/internal/common/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -42,6 +43,7 @@ func TestDocker_SSHServer(t *testing.T) { assert.NoError(t, dockerRunner.Submit(ctx, taskConfig)) assert.NoError(t, dockerRunner.Run(ctx, taskConfig.ID)) + assertTaskDone(t, dockerRunner, taskConfig.ID) } func TestDocker_ShmNoexecByDefault(t *testing.T) { @@ -67,6 +69,7 @@ func TestDocker_ShmNoexecByDefault(t *testing.T) { assert.NoError(t, dockerRunner.Submit(ctx, taskConfig)) assert.NoError(t, dockerRunner.Run(ctx, taskConfig.ID)) + assertTaskDone(t, dockerRunner, taskConfig.ID) } func TestDocker_ShmExecIfSizeSpecified(t *testing.T) { @@ -93,6 +96,7 @@ func TestDocker_ShmExecIfSizeSpecified(t *testing.T) { assert.NoError(t, dockerRunner.Submit(ctx, taskConfig)) assert.NoError(t, dockerRunner.Run(ctx, taskConfig.ID)) + assertTaskDone(t, dockerRunner, taskConfig.ID) } func TestConfigureGpus_Nvidia(t *testing.T) { @@ -193,6 +197,16 @@ func generateID(t *testing.T) string { return hex.EncodeToString(b)[:idLen] } +// assertTaskDone asserts that the final state of a successfully executed task +// is committed to the storage +func assertTaskDone(t *testing.T, runner *DockerRunner, taskID string) { + t.Helper() + taskInfo := runner.TaskInfo(taskID) + assert.Equal(t, TaskStatusTerminated, taskInfo.Status) + assert.Equal(t, string(types.TerminationReasonDoneByRunner), taskInfo.TerminationReason) + assert.Empty(t, taskInfo.TerminationMessage) +} + func createTaskConfig(t *testing.T) TaskConfig { return TaskConfig{ ID: generateID(t), diff --git a/runner/internal/shim/task.go b/runner/internal/shim/task.go index ea3ad7c960..7e04e93c05 100644 --- a/runner/internal/shim/task.go +++ b/runner/internal/shim/task.go @@ -48,12 +48,23 @@ type Task struct { } // Lock is used for exclusive operations, e.g, stopping a container, -// removing task data, etc. +// removing task data, etc. It blocks until the lock is acquired, since +// contention is expected, e.g., the server may terminate a task while it is +// being processed in the background. func (t *Task) Lock(ctx context.Context) { + t.mu.Lock() + log.Debug(ctx, "locked", "task", t.ID) +} + +// TryLock is a non-blocking version of Lock. It reports whether the lock has +// been acquired, so that the caller can retry later instead of waiting. +func (t *Task) TryLock(ctx context.Context) bool { if !t.mu.TryLock() { - log.Fatal(ctx, "already locked!", "task", t.ID) + log.Debug(ctx, "already locked", "task", t.ID) + return false } log.Debug(ctx, "locked", "task", t.ID) + return true } // Release should be called Unlock, but this name triggers govet copylocks check, @@ -66,13 +77,13 @@ func (t *Task) Release(ctx context.Context) { func (t *Task) IsTransitionAllowed(toStatus TaskStatus) bool { // same-state transitions are not allowed unless stated otherwise, meaning that - // task.Update(); task.Update() is not allowed is most cases. - // This is mainly done to avoid erroneous/concurrent updates, though this limits - // our ability to commit internal state more often. - // If this becomes a problem, consider allowing sameState->sameState transitions in general. + // two consecutive updates to the same status are not allowed in most cases. + // This is mainly done to avoid erroneous/concurrent updates. + // Note that TaskStorage.Modify() checks the transition only if the status changes, + // therefore committing internal state without changing the status is always allowed. switch toStatus { case TaskStatusPending: - // initial status, task should be Add()ed with it, not Update()d + // initial status, task should be Add()ed with it, not Modify()ed return false case TaskStatusPreparing: return t.Status == TaskStatusPending @@ -81,12 +92,11 @@ func (t *Task) IsTransitionAllowed(toStatus TaskStatus) bool { case TaskStatusCreating: return t.Status == TaskStatusPulling case TaskStatusRunning: - // allow running->running transition to update internal state, e.g., ports - return t.Status == TaskStatusCreating || t.Status == TaskStatusRunning + return t.Status == TaskStatusCreating case TaskStatusTerminated: // terminated -> terminated is also allowed since server _always_ tries to // terminate the task, even if it is already terminated, but this is a special case, - // see TaskStorage.Update() for details + // see TaskStorage.Modify() for details return true } return false @@ -152,7 +162,7 @@ type TaskStorage struct { mu sync.RWMutex } -// Get a _copy_ of all tasks. To "commit" changes, use Update() +// Get a _copy_ of all tasks. To "commit" changes, use Modify() func (ts *TaskStorage) List() []Task { ts.mu.RLock() defer ts.mu.RUnlock() @@ -163,7 +173,7 @@ func (ts *TaskStorage) List() []Task { return tasks } -// Get a _copy_ of the task. To "commit" changes, use Update() +// Get a _copy_ of the task. To "commit" changes, use Modify() func (ts *TaskStorage) Get(id string) (Task, bool) { ts.mu.RLock() defer ts.mu.RUnlock() @@ -182,29 +192,38 @@ func (ts *TaskStorage) Add(task Task) bool { return true } -// Update the _existing_ task. If the task is not in the storage, do nothing and return false -// If the current status is terminated, do nothing and return false -func (ts *TaskStorage) Update(task Task) error { +// Modify applies fn to a _copy_ of the _existing_ task and commits the copy, +// returning it on success. If the task is not in the storage, do nothing and +// return ErrNotFound. +// If fn returns an error, or if the resulting status transition is not allowed, +// the copy is discarded, that is, a partially applied fn never reaches the storage. +// The transition is checked only if fn changes the status, therefore fn is free to +// update the internal state of the task without changing its status. +// fn is called with the storage lock held, so it must be fast and must not block, +// in particular, it must not call the Docker API or touch the file system. +func (ts *TaskStorage) Modify(id string, fn func(*Task) error) (Task, error) { ts.mu.Lock() defer ts.mu.Unlock() - currentTask, ok := ts.tasks[task.ID] + currentTask, ok := ts.tasks[id] if !ok { - return ErrNotFound + return Task{}, ErrNotFound + } + task := currentTask + if err := fn(&task); err != nil { + return Task{}, err } - if !currentTask.IsTransitionAllowed(task.Status) { - return fmt.Errorf("%w: %s -> %s transition not allowed", ErrRequest, currentTask.Status, task.Status) + if task.Status != currentTask.Status && !currentTask.IsTransitionAllowed(task.Status) { + return Task{}, fmt.Errorf("%w: %s -> %s transition not allowed", ErrRequest, currentTask.Status, task.Status) } - if currentTask.Status == TaskStatusTerminated { + if currentTask.Status == TaskStatusTerminated && currentTask.TerminationReason != "" { // We ignore reason/message fields if they are already set to avoid // overriding these fields by the server, which _always_ tries to terminate the task, // even if it is not running - if currentTask.TerminationReason != "" { - task.TerminationReason = currentTask.TerminationReason - task.TerminationMessage = currentTask.TerminationMessage - } + task.TerminationReason = currentTask.TerminationReason + task.TerminationMessage = currentTask.TerminationMessage } - ts.tasks[task.ID] = task - return nil + ts.tasks[id] = task + return task, nil } func (ts *TaskStorage) Delete(id string) { diff --git a/runner/internal/shim/task_test.go b/runner/internal/shim/task_test.go index 37ea6d7542..3cf1d194d7 100644 --- a/runner/internal/shim/task_test.go +++ b/runner/internal/shim/task_test.go @@ -1,8 +1,11 @@ package shim import ( + "errors" "fmt" + "sync" "testing" + "time" "github.com/stretchr/testify/assert" ) @@ -43,37 +46,95 @@ func TestTaskStorage_Add_AlreadyExists(t *testing.T) { assert.Equal(t, storedTask, storage.tasks["1"]) } -func TestTaskStorage_Update_OK(t *testing.T) { +func TestTaskStorage_Modify_OK(t *testing.T) { storage := NewTaskStorage() - storedTask := Task{ID: "1", Status: TaskStatusRunning} - storage.tasks["1"] = storedTask - updatedTask := Task{ID: "1", Status: TaskStatusTerminated} + storage.tasks["1"] = Task{ID: "1", Status: TaskStatusRunning} + + task, err := storage.Modify("1", func(t *Task) error { + t.SetStatusTerminated("container_exited_with_error", "oom") + return nil + }) + assert.NoError(t, err) + assert.Equal(t, TaskStatusTerminated, task.Status) + assert.Equal(t, "container_exited_with_error", task.TerminationReason) + assert.Equal(t, task, storage.tasks["1"]) +} - err := storage.Update(updatedTask) - assert.Nil(t, err) - assert.Equal(t, updatedTask, storage.tasks["1"]) +// The transition is not checked unless the status changes, so that the internal +// state of a task can be committed at any time +func TestTaskStorage_Modify_NoStatusChange(t *testing.T) { + storage := NewTaskStorage() + storage.tasks["1"] = Task{ID: "1", Status: TaskStatusRunning} + ports := []PortMapping{{Host: 30000, Container: 10999}} + + task, err := storage.Modify("1", func(t *Task) error { + t.ports = ports + return nil + }) + assert.NoError(t, err) + assert.Equal(t, TaskStatusRunning, task.Status) + assert.Equal(t, ports, storage.tasks["1"].ports) } -func TestTaskStorage_Update_DoesNotExist(t *testing.T) { +func TestTaskStorage_Modify_DoesNotExist(t *testing.T) { storage := NewTaskStorage() - err := storage.Update(Task{ID: "1", Status: TaskStatusPending}) + _, err := storage.Modify("1", func(t *Task) error { + t.SetStatusPreparing() + return nil + }) assert.ErrorIs(t, err, ErrNotFound) - assert.Equal(t, 0, len(storage.tasks)) + assert.Empty(t, storage.tasks) } -func TestTaskStorage_Update_TransitionNotAllowed(t *testing.T) { +func TestTaskStorage_Modify_TransitionNotAllowed(t *testing.T) { storage := NewTaskStorage() storedTask := Task{ID: "1", Status: TaskStatusPending} storage.tasks["1"] = storedTask - updatedTask := Task{ID: "1", Status: TaskStatusRunning} - err := storage.Update(updatedTask) + _, err := storage.Modify("1", func(t *Task) error { + t.SetStatusRunning() + return nil + }) assert.ErrorIs(t, err, ErrRequest) - assert.ErrorContains(t, err, fmt.Sprintf("%s -> %s", storedTask.Status, updatedTask.Status)) + assert.ErrorContains(t, err, fmt.Sprintf("%s -> %s", TaskStatusPending, TaskStatusRunning)) assert.Equal(t, storedTask, storage.tasks["1"]) } +// A partially applied function must not reach the storage +func TestTaskStorage_Modify_Error(t *testing.T) { + storage := NewTaskStorage() + storedTask := Task{ID: "1", Status: TaskStatusPulling} + storage.tasks["1"] = storedTask + errFailed := errors.New("failed") + + _, err := storage.Modify("1", func(t *Task) error { + t.SetStatusCreating() + return errFailed + }) + assert.ErrorIs(t, err, errFailed) + assert.Equal(t, storedTask, storage.tasks["1"]) +} + +func TestTaskStorage_Modify_TerminationReasonNotOverridden(t *testing.T) { + storage := NewTaskStorage() + storage.tasks["1"] = Task{ + ID: "1", + Status: TaskStatusTerminated, + TerminationReason: "container_exited_with_error", + TerminationMessage: "oom", + } + + task, err := storage.Modify("1", func(t *Task) error { + t.SetStatusTerminated("terminated_by_server", "") + return nil + }) + assert.NoError(t, err) + assert.Equal(t, "container_exited_with_error", task.TerminationReason) + assert.Equal(t, "oom", task.TerminationMessage) + assert.Equal(t, "container_exited_with_error", storage.tasks["1"].TerminationReason) +} + func TestTaskStorage_Delete(t *testing.T) { storage := NewTaskStorage() storage.tasks["1"] = Task{ID: "1", Status: TaskStatusRunning} @@ -85,6 +146,42 @@ func TestTaskStorage_Delete(t *testing.T) { assert.Equal(t, 0, len(storage.tasks)) } +func TestTask_TryLock(t *testing.T) { + ctx := t.Context() + task := Task{ID: "1", mu: &sync.Mutex{}} + + assert.True(t, task.TryLock(ctx)) + assert.False(t, task.TryLock(ctx)) + + task.Release(ctx) + assert.True(t, task.TryLock(ctx)) +} + +func TestTask_Lock_WaitsForRelease(t *testing.T) { + ctx := t.Context() + task := Task{ID: "1", mu: &sync.Mutex{}} + task.Lock(ctx) + + locked := make(chan struct{}) + go func() { + task.Lock(ctx) + close(locked) + }() + + select { + case <-locked: + t.Fatal("Lock did not wait for Release") + case <-time.After(50 * time.Millisecond): + } + + task.Release(ctx) + select { + case <-locked: + case <-time.After(5 * time.Second): + t.Fatal("Lock did not acquire the lock after Release") + } +} + func TestTask_IsTransitionAllowed_true(t *testing.T) { testCases := []struct { oldStatus, newStatus TaskStatus @@ -97,7 +194,6 @@ func TestTask_IsTransitionAllowed_true(t *testing.T) { {TaskStatusPulling, TaskStatusTerminated}, {TaskStatusCreating, TaskStatusRunning}, {TaskStatusCreating, TaskStatusTerminated}, - {TaskStatusRunning, TaskStatusRunning}, {TaskStatusRunning, TaskStatusTerminated}, {TaskStatusTerminated, TaskStatusTerminated}, } @@ -115,6 +211,7 @@ func TestTask_IsTransitionAllowed_false(t *testing.T) { {TaskStatusPending, TaskStatusPending}, {TaskStatusPending, TaskStatusRunning}, {TaskStatusPulling, TaskStatusPending}, + {TaskStatusRunning, TaskStatusRunning}, } for _, tc := range testCases { task := Task{ID: "1", Status: tc.oldStatus}