diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index 7fbda1458..1864b5f52 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 3f807dd37..7c634fa60 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 ea3ad7c96..7e04e93c0 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 37ea6d754..3cf1d194d 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}