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

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)
 }
 

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),

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) {