From e6d0dde98500041c2e4eaf51f777f37ce6d4fd50 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Wed, 26 Aug 2026 15:01:20 +0000 Subject: [PATCH 1/2] [shim] Persist task state to clean up tasks after a restart 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: https://github.com/dstackai/dstack/issues/4182 Co-Authored-By: Claude Opus 5 (1M context) --- runner/internal/shim/docker.go | 118 +++++++++++++---- runner/internal/shim/docker_test.go | 91 +++++++++++-- runner/internal/shim/models.go | 1 + runner/internal/shim/process.go | 2 +- runner/internal/shim/process_test.go | 7 +- runner/internal/shim/state.go | 189 +++++++++++++++++++++++++++ runner/internal/shim/state_test.go | 140 ++++++++++++++++++++ runner/internal/shim/task.go | 17 +-- runner/internal/shim/volumes.go | 10 +- 9 files changed, 517 insertions(+), 58 deletions(-) create mode 100644 runner/internal/shim/state.go create mode 100644 runner/internal/shim/state_test.go diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index 3db376fdc..1cbdaf4c1 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -17,6 +17,7 @@ import ( "sync" "time" + dockertypes "github.com/docker/docker/api/types" "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/filters" "github.com/docker/docker/api/types/image" @@ -156,6 +157,8 @@ type DockerRunner struct { gpuVendor gpu.GpuVendor gpuLock *GpuLock tasks TaskStorage + // stateMu serializes task state file updates, see saveTaskState() + stateMu sync.Mutex } func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*DockerRunner, error) { @@ -203,6 +206,9 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke if err := runner.restoreStateFromContainers(ctx); err != nil { return nil, fmt.Errorf("failed to restore state from containers: %w", err) } + // Must be called after the tasks are restored, as it uses them to tell the dirs of + // the live tasks from the orphaned ones + runner.sweepOrphanedTaskDirs(ctx) return runner, nil } @@ -283,29 +289,70 @@ func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error { break } } - if len(gpuIDs) > 0 { + state, restored := readRestoredTaskState(ctx, taskID, runnerDir) + config := state.Config + if !restored { + // Containers created by shim versions that did not write the state file have + // no config to restore. The volumes can still be recovered from the container + // mounts, unlike the host SSH keys, which are only known to the state file + config.Volumes = volumesFromMounts(containerShort.Mounts) + } + if state.CleanedUp { + // The resources of this task, its GPUs included, have already been released, + // therefore the task owns nothing + gpuIDs = nil + } else if len(gpuIDs) > 0 { // A GPU already locked by another restored task is not locked again and, // therefore, is not owned by this task -- otherwise, cleaning up this task // would release a GPU that the other one is still using gpuIDs = d.gpuLock.Lock(ctx, gpuIDs) log.Debug(ctx, "locked GPU(s) due to running task", "task", taskID, "gpus", gpuIDs) } - // Tasks are restored as running regardless of the container state, letting - // ProcessTasks() decide whether the container is still running and, if it is - // not, why it finished. This way, the termination reason of a task that - // finished while the shim was not running is not lost - task := NewTask(taskID, TaskStatusRunning, containerName, containerID, gpuIDs, ports, runnerDir) + // A task with a recorded termination reason has been terminated before the shim + // restarted, and its reason must not be overridden by the container exit code. + // Otherwise the task is restored as running regardless of the container state, + // letting ProcessTasks() decide whether the container is still running and, if + // it is not, why it finished + status := TaskStatusRunning + if state.TerminationReason != "" { + status = TaskStatusTerminated + } + task := NewTask(taskID, status) + task.TerminationReason = state.TerminationReason + task.TerminationMessage = state.TerminationMessage + task.config = config + task.containerName = containerName + task.containerID = containerID + task.gpuIDs = gpuIDs + task.ports = ports + task.runnerDir = runnerDir + task.cleanedUp = state.CleanedUp if !d.tasks.Add(task) { log.Error(ctx, "duplicate restored task", "task", taskID) // Nothing will release the GPUs of a task that is not stored d.gpuLock.Release(ctx, gpuIDs) continue } - log.Debug(ctx, "restored task", "task", taskID, "state", containerShort.State, "gpus", gpuIDs) + log.Debug( + ctx, "restored task", + "task", taskID, "status", status, "state", containerShort.State, "gpus", gpuIDs, + ) } return nil } +// volumesFromMounts recovers the volumes attached to a task from its container mounts. +// Only the volume names are recovered, which is all that is needed to unmount them +func volumesFromMounts(mounts []dockertypes.MountPoint) []VolumeInfo { + var volumes []VolumeInfo + for _, mount := range mounts { + if name, found := strings.CutPrefix(mount.Source, volumeMountPointDir+"/"); found { + volumes = append(volumes, VolumeInfo{Name: name}) + } + } + return volumes +} + func (d *DockerRunner) Resources(ctx context.Context) Resources { cpuCount := host.GetCpuCount(ctx) totalMemory, err := host.GetTotalMemory(ctx) @@ -375,7 +422,7 @@ func (d *DockerRunner) Submit(ctx context.Context, cfg TaskConfig) error { // 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 { +func (d *DockerRunner) commit(ctx context.Context, task *Task, mutate func(*Task)) error { updatedTask, err := d.tasks.Modify(task.ID, func(t *Task) error { mutate(t) return nil @@ -384,6 +431,7 @@ func (d *DockerRunner) commit(task *Task, mutate func(*Task)) error { return err } *task = updatedTask + d.saveTaskState(ctx, task.ID) return nil } @@ -395,7 +443,7 @@ func (d *DockerRunner) commitTerminated(ctx context.Context, task *Task) { // the last successful commit is the actual state, nothing to commit return } - if err := d.commit(task, func(t *Task) { + if err := d.commit(ctx, 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) @@ -439,7 +487,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { } } // Hand the task over to ProcessTasks() - if commitErr := d.commit(&task, func(t *Task) { + if commitErr := d.commit(ctx, &task, func(t *Task) { t.startInFlight = false if task.containerID != "" { t.containerID = task.containerID @@ -455,7 +503,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { } }() - if err := d.commit(&task, func(t *Task) { + if err := d.commit(ctx, &task, func(t *Task) { t.startInFlight = true t.SetStatusPreparing() }); err != nil { @@ -471,7 +519,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { log.Trace(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 { + if err := d.commit(ctx, &task, func(t *Task) { t.runnerDir = runnerDir }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } @@ -487,7 +535,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { } else { gpuIDs = []string{} } - if err := d.commit(&task, func(t *Task) { t.gpuIDs = gpuIDs }); err != nil { + if err := d.commit(ctx, &task, func(t *Task) { t.gpuIDs = gpuIDs }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } @@ -520,7 +568,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { log.Debug(ctx, "Pulling image") pullCtx, cancelPull := context.WithTimeout(ctx, ImagePullTimeout) defer cancelPull() - if err := d.commit(&task, func(t *Task) { t.SetStatusPulling(cancelPull) }); err != nil { + if err := d.commit(ctx, &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. @@ -534,7 +582,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { } log.Debug(ctx, "Creating container", "task", task.ID, "name", task.containerName) - if err := d.commit(&task, func(t *Task) { t.SetStatusCreating() }); err != nil { + if err := d.commit(ctx, &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 { @@ -574,7 +622,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { } // The container is running, the task is now processed in the background - if err := d.commit(&task, func(t *Task) { + if err := d.commit(ctx, &task, func(t *Task) { // startContainer sets the ports field, the retry above may have // replaced the container t.containerID = task.containerID @@ -608,16 +656,7 @@ func (d *DockerRunner) cleanupLocked(ctx context.Context, task *Task) { return } log.Debug(ctx, "releasing task resources", "task", task.ID) - cfg := task.config - if err := unmountVolumes(ctx, cfg); err != nil { - log.Error(ctx, "failed to unmount volumes", "task", task.ID, "err", err) - } - if len(cfg.HostSshKeys) > 0 { - ak := AuthorizedKeys{user: cfg.HostSshUser, lookup: user.Lookup} - if err := ak.RemovePublicKeys(cfg.HostSshKeys); err != nil { - log.Error(ctx, "failed to remove public keys", "task", task.ID, "err", err) - } - } + releaseTaskResources(ctx, task.config) if len(task.gpuIDs) > 0 { releasedGpuIDs := d.gpuLock.Release(ctx, task.gpuIDs) log.Debug(ctx, "released GPU(s)", "task", task.ID, "gpus", releasedGpuIDs) @@ -631,6 +670,22 @@ func (d *DockerRunner) cleanupLocked(ctx context.Context, task *Task) { }); err != nil && !errors.Is(err, ErrNotFound) { log.Error(ctx, "failed to commit cleaned up state", "task", task.ID, "err", err) } + d.saveTaskState(ctx, task.ID) +} + +// releaseTaskResources releases the host resources acquired for a task: volumes and +// host SSH keys. Unlike GPU locks, which are only kept in memory, these outlive the +// shim process, therefore they are released by task config and not by task +func releaseTaskResources(ctx context.Context, cfg TaskConfig) { + if err := unmountVolumes(ctx, cfg); err != nil { + log.Error(ctx, "failed to unmount volumes", "err", err) + } + if len(cfg.HostSshKeys) > 0 { + ak := AuthorizedKeys{user: cfg.HostSshUser, lookup: user.Lookup} + if err := ak.RemovePublicKeys(cfg.HostSshKeys); err != nil { + log.Error(ctx, "failed to remove public keys", "err", err) + } + } } // Terminate aborts running operations (pulling an image, running a container) and sets task status to terminated @@ -731,7 +786,10 @@ func (d *DockerRunner) remove(ctx context.Context, task *Task) (err error) { // Failed attempts to remove or rename runner dir are considered non-fatal if err := os.RemoveAll(task.runnerDir); err != nil { log.Error(ctx, "failed to remove runner directory", "dir", task.runnerDir, "err", err) - trashName := fmt.Sprintf(".trash-%s-%d", task.runnerDir, time.Now().UnixMicro()) + trashName := filepath.Join( + filepath.Dir(task.runnerDir), + fmt.Sprintf(".trash-%s-%d", filepath.Base(task.runnerDir), time.Now().UnixMicro()), + ) if err := os.Rename(task.runnerDir, trashName); err != nil { log.Error(ctx, "failed to rename runner directory", "dir", task.runnerDir, "err", err) } @@ -1317,8 +1375,12 @@ func (c *CLIArgs) DockerPorts() []int { return []int{c.Runner.HTTPPort, c.Runner.SSHPort} } +func (c *CLIArgs) RunnersDir() string { + return filepath.Join(c.Shim.HomeDir, "runners") +} + func (c *CLIArgs) MakeRunnerDir(name string) (string, error) { - runnerTemp := filepath.Join(c.Shim.HomeDir, "runners", name) + runnerTemp := filepath.Join(c.RunnersDir(), name) if err := os.MkdirAll(runnerTemp, 0o755); err != nil { return "", fmt.Errorf("create runner directory: %w", err) } diff --git a/runner/internal/shim/docker_test.go b/runner/internal/shim/docker_test.go index fd46d51bd..4a80703da 100644 --- a/runner/internal/shim/docker_test.go +++ b/runner/internal/shim/docker_test.go @@ -5,6 +5,8 @@ import ( "encoding/hex" "errors" "math/rand" + "os" + "path/filepath" "sync" "testing" "time" @@ -36,7 +38,7 @@ func TestDocker_SSHServer(t *testing.T) { params := &dockerParametersMock{ commands: []string{"/usr/sbin/sshd -V 2>&1 | grep OpenSSH"}, sshShellCommands: true, - runnerDir: t.TempDir(), + runnersDir: t.TempDir(), } timeout := 180 // seconds @@ -60,8 +62,8 @@ func TestDocker_ShmNoexecByDefault(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"mount | grep '/dev/shm .*size=65536k' | grep noexec"}, - runnerDir: t.TempDir(), + commands: []string{"mount | grep '/dev/shm .*size=65536k' | grep noexec"}, + runnersDir: t.TempDir(), } timeout := 180 // seconds @@ -85,8 +87,8 @@ func TestDocker_ShmExecIfSizeSpecified(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"mount | grep '/dev/shm .*size=1024k' | grep -v noexec"}, - runnerDir: t.TempDir(), + commands: []string{"mount | grep '/dev/shm .*size=1024k' | grep -v noexec"}, + runnersDir: t.TempDir(), } timeout := 180 // seconds @@ -111,8 +113,8 @@ func TestDocker_ContainerExitedWithError(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"echo failed for a reason", "exit 3"}, - runnerDir: t.TempDir(), + commands: []string{"echo failed for a reason", "exit 3"}, + runnersDir: t.TempDir(), } timeout := 180 // seconds @@ -142,8 +144,8 @@ func TestDocker_RestoredTaskIsTerminated(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"sleep 3", "exit 7"}, - runnerDir: t.TempDir(), + commands: []string{"sleep 3", "exit 7"}, + runnersDir: t.TempDir(), } timeout := 180 // seconds @@ -167,6 +169,54 @@ func TestDocker_RestoredTaskIsTerminated(t *testing.T) { assert.Equal(t, string(types.TerminationReasonContainerExitedWithError), taskInfo.TerminationReason) } +// TestDocker_RestoredTaskKeepsTerminationReason covers the task state file: why a task +// was terminated cannot be recovered from its container +func TestDocker_RestoredTaskKeepsTerminationReason(t *testing.T) { + if testing.Short() { + t.Skip() + } + + params := &dockerParametersMock{ + commands: []string{"sleep 60"}, + runnersDir: t.TempDir(), + } + + timeout := 180 // seconds + ctx, cancel := context.WithTimeout(t.Context(), time.Duration(timeout)*time.Second) + defer cancel() + + dockerRunner, err := NewDockerRunner(ctx, params) + require.NoError(t, err) + + taskConfig := createTaskConfig(t) + require.NoError(t, dockerRunner.Submit(ctx, taskConfig)) + require.NoError(t, dockerRunner.Start(ctx, taskConfig.ID)) + require.NoError(t, dockerRunner.Terminate( + ctx, taskConfig.ID, 0, string(types.TerminationReasonTerminatedByUser), "bye", + )) + + // The restarted shim restores the task from its container and its state file + restartedRunner, err := NewDockerRunner(ctx, params) + require.NoError(t, err) + defer cleanupTask(t, restartedRunner, taskConfig.ID) + + taskInfo := restartedRunner.TaskInfo(taskConfig.ID) + assert.Equal(t, TaskStatusTerminated, taskInfo.Status) + assert.Equal(t, string(types.TerminationReasonTerminatedByUser), taskInfo.TerminationReason) + assert.Equal(t, "bye", taskInfo.TerminationMessage) +} + +func TestVolumesFromMounts(t *testing.T) { + mounts := []dockertypes.MountPoint{ + {Source: "/root/.dstack/runners/task-0-0-1234abcd", Destination: consts.RunnerTempDir}, + {Source: "/mnt/disks/dstack-volumes/volume-1", Destination: "/volume"}, + {Source: "/home/dstack/data", Destination: "/instance-data"}, + } + + assert.Equal(t, []VolumeInfo{{Name: "volume-1"}}, volumesFromMounts(mounts)) + assert.Nil(t, volumesFromMounts(nil)) +} + func TestConfigureGpus_Nvidia(t *testing.T) { containerConfig := &container.Config{} hostConfig := &container.HostConfig{} @@ -212,7 +262,8 @@ func TestShouldRetryWithoutNvidiaDisplayCapability(t *testing.T) { type dockerParametersMock struct { commands []string sshShellCommands bool - runnerDir string + // runnersDir is the parent dir of all task runner dirs + runnersDir string } func (c *dockerParametersMock) DockerPrivileged() bool { @@ -240,12 +291,24 @@ func (c *dockerParametersMock) DockerPorts() []int { return []int{} } -func (c *dockerParametersMock) DockerMounts(string) ([]mount.Mount, error) { - return nil, nil +func (c *dockerParametersMock) DockerMounts(hostRunnerDir string) ([]mount.Mount, error) { + // The runner dir mount is how a restarted shim finds the dir of a restored task, + // and with it the task state file + return []mount.Mount{ + {Type: mount.TypeBind, Source: hostRunnerDir, Target: consts.RunnerTempDir}, + }, nil } -func (c *dockerParametersMock) MakeRunnerDir(string) (string, error) { - return c.runnerDir, nil +func (c *dockerParametersMock) RunnersDir() string { + return c.runnersDir +} + +func (c *dockerParametersMock) MakeRunnerDir(name string) (string, error) { + runnerDir := filepath.Join(c.runnersDir, name) + if err := os.MkdirAll(runnerDir, 0o755); err != nil { + return "", err + } + return runnerDir, nil } /* Utilities */ diff --git a/runner/internal/shim/models.go b/runner/internal/shim/models.go index 6cc093505..e527aa931 100644 --- a/runner/internal/shim/models.go +++ b/runner/internal/shim/models.go @@ -10,6 +10,7 @@ type DockerParameters interface { DockerShellCommands(authorizedKeys []string, runnerHttpAddress string) []string DockerMounts(string) ([]mount.Mount, error) DockerPorts() []int + RunnersDir() string MakeRunnerDir(name string) (string, error) DockerPJRTDevice() string } diff --git a/runner/internal/shim/process.go b/runner/internal/shim/process.go index 920caf658..16fe067eb 100644 --- a/runner/internal/shim/process.go +++ b/runner/internal/shim/process.go @@ -179,7 +179,7 @@ func (d *DockerRunner) finalizeLocked( ctx context.Context, task *Task, reason types.TerminationReason, message string, ) { d.cleanupLocked(ctx, task) - if err := d.commit(task, func(t *Task) { + if err := d.commit(ctx, task, func(t *Task) { t.SetStatusTerminated(string(reason), message) }); err != nil { log.Error(ctx, "failed to commit terminated status", "task", task.ID, "err", err) diff --git a/runner/internal/shim/process_test.go b/runner/internal/shim/process_test.go index f27a2cb8d..dd4b74727 100644 --- a/runner/internal/shim/process_test.go +++ b/runner/internal/shim/process_test.go @@ -214,7 +214,7 @@ func newTestRunner(t *testing.T, client docker.APIClient) *DockerRunner { t.Helper() return &DockerRunner{ client: client, - dockerParams: &dockerParametersMock{runnerDir: t.TempDir()}, + dockerParams: &dockerParametersMock{runnersDir: t.TempDir()}, gpuLock: newTestGpuLock(t), tasks: NewTaskStorage(), } @@ -230,7 +230,10 @@ func newTestGpuLock(t *testing.T, ids ...string) *GpuLock { } func newRunningTask(taskID string, containerID string) Task { - return NewTask(taskID, TaskStatusRunning, taskID+"-name", containerID, nil, nil, "") + task := NewTask(taskID, TaskStatusRunning) + task.containerName = taskID + "-name" + task.containerID = containerID + return task } func addTask(runner *DockerRunner, task Task) { diff --git a/runner/internal/shim/state.go b/runner/internal/shim/state.go new file mode 100644 index 000000000..8c41d39d0 --- /dev/null +++ b/runner/internal/shim/state.go @@ -0,0 +1,189 @@ +package shim + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/dstackai/dstack/runner/internal/common/log" +) + +// taskStateFileName is the name of the task state file in the task runner dir +const taskStateFileName = "task.json" + +// taskStateVersion is the version of the task state file format. It is only bumped on +// incompatible changes: new fields are ignored by older versions of the shim, and are +// zero when an older file is read by a newer version. +const taskStateVersion = 1 + +// taskState is the part of the task state that cannot be recovered from the task +// container, persisted so that the task can be finalized by another shim process. +// Everything that is recoverable -- the container ID, the GPUs, the ports -- is left +// out, keeping the container the source of truth for it. +type taskState struct { + Version int `json:"version"` + ID string `json:"id"` + // Config is the task config with the registry credentials stripped: they are only + // needed to pull the image, which is already pulled by the time the state is read. + // The rest is kept for cleanup -- volumes and host SSH keys in particular. + Config TaskConfig `json:"config"` + // TerminationReason is empty unless the task has been terminated, meaning that the + // outcome is to be determined from the container state + TerminationReason string `json:"termination_reason"` + TerminationMessage string `json:"termination_message"` + CleanedUp bool `json:"cleaned_up"` +} + +func newTaskState(task *Task) taskState { + config := task.config + config.RegistryUsername = "" + config.RegistryPassword = "" + return taskState{ + Version: taskStateVersion, + ID: task.ID, + Config: config, + TerminationReason: task.TerminationReason, + TerminationMessage: task.TerminationMessage, + CleanedUp: task.cleanedUp, + } +} + +// saveTaskState writes the state of the task to its runner dir. Tasks that have not +// acquired any resources yet, that is, tasks without a runner dir, are skipped. +// The stored task is snapshotted while holding a lock, so that a concurrent call +// cannot replace a newer state with an older one. +func (d *DockerRunner) saveTaskState(ctx context.Context, taskID string) { + d.stateMu.Lock() + defer d.stateMu.Unlock() + task, ok := d.tasks.Get(taskID) + if !ok || task.runnerDir == "" { + return + } + if err := writeTaskState(task.runnerDir, newTaskState(&task)); err != nil { + log.Error(ctx, "failed to save task state", "task", taskID, "err", err) + } +} + +func writeTaskState(dir string, state taskState) error { + data, err := json.Marshal(state) + if err != nil { + return fmt.Errorf("marshal task state: %w", err) + } + data = append(data, '\n') + path := filepath.Join(dir, taskStateFileName) + // Written to a temporary file and renamed, so that a state file is never partial + tempPath := path + ".tmp" + if err := writeFileSync(tempPath, data, 0o600); err != nil { + return fmt.Errorf("write task state: %w", err) + } + if err := os.Rename(tempPath, path); err != nil { + return fmt.Errorf("rename task state: %w", err) + } + return nil +} + +// readTaskState reads the state file from the task runner dir. The error wraps +// os.ErrNotExist if the dir has no state file, e.g., if the task was started by a shim +// version that did not write one. +func readTaskState(dir string) (taskState, error) { + var state taskState + path := filepath.Join(dir, taskStateFileName) + data, err := os.ReadFile(path) + if err != nil { + return state, fmt.Errorf("read task state: %w", err) + } + if err := json.Unmarshal(data, &state); err != nil { + return state, fmt.Errorf("unmarshal task state %s: %w", path, err) + } + if state.Version != taskStateVersion { + return state, fmt.Errorf("unsupported task state version %d in %s", state.Version, path) + } + return state, nil +} + +// readRestoredTaskState reads the state file of a task being restored from its +// container. It reports whether there was a state file to restore from; if there was +// not, the returned state is empty and the caller has to fall back to the container +func readRestoredTaskState(ctx context.Context, taskID string, runnerDir string) (taskState, bool) { + if runnerDir == "" { + return taskState{}, false + } + state, err := readTaskState(runnerDir) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + // A task started by a shim version that did not write the state file has none + log.Error(ctx, "failed to read task state", "task", taskID, "err", err) + } + return taskState{}, false + } + if state.ID != taskID { + log.Error(ctx, "task state belongs to another task", "task", taskID, "id", state.ID) + return taskState{}, false + } + return state, true +} + +// sweepOrphanedTaskDirs releases the resources of tasks that have a state file but no +// container, which normally means that the shim stopped running before the container +// was created, and removes their dirs. +// Dirs without a state file are left intact, as there is no way to tell whether they +// belong to a task, and so are the dirs of the tasks restored from containers. +func (d *DockerRunner) sweepOrphanedTaskDirs(ctx context.Context) { + runnersDir := d.dockerParams.RunnersDir() + entries, err := os.ReadDir(runnersDir) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + log.Error(ctx, "failed to list task dirs", "dir", runnersDir, "err", err) + } + return + } + for _, entry := range entries { + // Dot dirs are not task dirs, e.g., the .trash-* dirs left by Remove() + if !entry.IsDir() || strings.HasPrefix(entry.Name(), ".") { + continue + } + dir := filepath.Join(runnersDir, entry.Name()) + state, err := readTaskState(dir) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + log.Error(ctx, "failed to read task state", "dir", dir, "err", err) + } + continue + } + if _, ok := d.tasks.Get(state.ID); ok { + // the task has been restored from its container + continue + } + log.Warning(ctx, "cleaning up orphaned task dir", "task", state.ID, "dir", dir) + if !state.CleanedUp { + // GPU locks are in-memory, so there is nothing to release: a task without + // a container holds no GPUs after a restart + releaseTaskResources(ctx, state.Config) + } + if err := os.RemoveAll(dir); err != nil { + log.Error(ctx, "failed to remove orphaned task dir", "dir", dir, "err", err) + } + } +} + +// writeFileSync writes the file and flushes it to the disk, so that its content is not +// lost if the host restarts +func writeFileSync(path string, data []byte, perm os.FileMode) (err error) { + file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm) + if err != nil { + return err + } + defer func() { + if closeErr := file.Close(); err == nil { + err = closeErr + } + }() + if _, err := file.Write(data); err != nil { + return err + } + return file.Sync() +} diff --git a/runner/internal/shim/state_test.go b/runner/internal/shim/state_test.go new file mode 100644 index 000000000..a63d2b657 --- /dev/null +++ b/runner/internal/shim/state_test.go @@ -0,0 +1,140 @@ +package shim + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTaskState_RoundTrip(t *testing.T) { + dir := t.TempDir() + task := NewTaskFromConfig(TaskConfig{ + ID: "task-1", + Name: "task", + RegistryUsername: "user", + RegistryPassword: "secret", + HostSshUser: "root", + HostSshKeys: []string{"ssh-ed25519 AAAA"}, + Volumes: []VolumeInfo{{Name: "volume-1", Backend: "aws"}}, + }) + task.TerminationReason = "terminated_by_user" + task.TerminationMessage = "bye" + task.cleanedUp = true + + require.NoError(t, writeTaskState(dir, newTaskState(&task))) + + state, err := readTaskState(dir) + require.NoError(t, err) + assert.Equal(t, taskStateVersion, state.Version) + assert.Equal(t, "task-1", state.ID) + assert.Equal(t, []VolumeInfo{{Name: "volume-1", Backend: "aws"}}, state.Config.Volumes) + assert.Equal(t, "root", state.Config.HostSshUser) + assert.Equal(t, []string{"ssh-ed25519 AAAA"}, state.Config.HostSshKeys) + assert.Equal(t, "terminated_by_user", state.TerminationReason) + assert.Equal(t, "bye", state.TerminationMessage) + assert.True(t, state.CleanedUp) + + // The registry credentials are not persisted, and the task is not affected + assert.Empty(t, state.Config.RegistryUsername) + assert.Empty(t, state.Config.RegistryPassword) + assert.Equal(t, "secret", task.config.RegistryPassword) + data, err := os.ReadFile(filepath.Join(dir, taskStateFileName)) + require.NoError(t, err) + assert.NotContains(t, string(data), "secret") + + // The temporary file used to write the state atomically is not left behind + entries, err := os.ReadDir(dir) + require.NoError(t, err) + assert.Len(t, entries, 1) +} + +func TestReadTaskState_DoesNotExist(t *testing.T) { + _, err := readTaskState(t.TempDir()) + assert.ErrorIs(t, err, os.ErrNotExist) +} + +func TestReadTaskState_Corrupted(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, taskStateFileName), []byte("{"), 0o600)) + + _, err := readTaskState(dir) + assert.Error(t, err) + assert.NotErrorIs(t, err, os.ErrNotExist) +} + +func TestReadTaskState_UnsupportedVersion(t *testing.T) { + dir := t.TempDir() + state := []byte(`{"version": 999, "id": "task-1"}`) + require.NoError(t, os.WriteFile(filepath.Join(dir, taskStateFileName), state, 0o600)) + + _, err := readTaskState(dir) + assert.ErrorContains(t, err, "unsupported task state version") +} + +func TestSaveTaskState(t *testing.T) { + runner := newTestRunner(t, nil) + dir := t.TempDir() + task := NewTaskFromConfig(TaskConfig{ID: "task-1", Name: "task"}) + task.runnerDir = dir + addTask(runner, task) + + runner.saveTaskState(t.Context(), "task-1") + + state, err := readTaskState(dir) + require.NoError(t, err) + assert.Equal(t, "task-1", state.ID) +} + +// A task that has not acquired any resources yet has no dir to save the state to +func TestSaveTaskState_NoRunnerDir(t *testing.T) { + runner := newTestRunner(t, nil) + addTask(runner, NewTaskFromConfig(TaskConfig{ID: "task-1", Name: "task"})) + + runner.saveTaskState(t.Context(), "task-1") + + // The state file must not be written to the current working directory + assert.NoFileExists(t, taskStateFileName) +} + +func TestSweepOrphanedTaskDirs(t *testing.T) { + runner := newTestRunner(t, nil) + runnersDir := runner.dockerParams.RunnersDir() + + // A task that has no container: its resources are released and its dir is removed + orphanedDir := makeTaskDir(t, runnersDir, "orphaned-0-0-1234abcd", "orphaned-task") + // A task restored from its container: still in use + restoredDir := makeTaskDir(t, runnersDir, "restored-0-0-5678cdef", "restored-task") + addTask(runner, newRunningTask("restored-task", "container-1")) + // A dir without a state file may belong to anything, e.g., to a task started by a + // shim version that did not write the state file + legacyDir := filepath.Join(runnersDir, "legacy-0-0-90abef01") + require.NoError(t, os.MkdirAll(legacyDir, 0o755)) + // Dot dirs are not task dirs + trashDir := makeTaskDir(t, runnersDir, ".trash-orphaned-0-0-1234abcd-1", "trashed-task") + + runner.sweepOrphanedTaskDirs(t.Context()) + + assert.NoDirExists(t, orphanedDir) + assert.DirExists(t, restoredDir) + assert.DirExists(t, legacyDir) + assert.DirExists(t, trashDir) +} + +func TestSweepOrphanedTaskDirs_NoRunnersDir(t *testing.T) { + runner := newTestRunner(t, nil) + runner.dockerParams = &dockerParametersMock{runnersDir: filepath.Join(t.TempDir(), "missing")} + + runner.sweepOrphanedTaskDirs(t.Context()) +} + +func makeTaskDir(t *testing.T, runnersDir string, name string, taskID string) string { + t.Helper() + dir := filepath.Join(runnersDir, name) + require.NoError(t, os.MkdirAll(dir, 0o755)) + task := NewTaskFromConfig(TaskConfig{ID: taskID, Name: name}) + require.NoError(t, writeTaskState(dir, newTaskState(&task))) + return dir +} diff --git a/runner/internal/shim/task.go b/runner/internal/shim/task.go index 67ee0edc3..1a035b75d 100644 --- a/runner/internal/shim/task.go +++ b/runner/internal/shim/task.go @@ -139,17 +139,14 @@ func (t *Task) SetStatusTerminated(reason string, message string) { t.cancelPull = nil } -func NewTask(id string, status TaskStatus, containerName string, containerID string, gpuIDs []string, ports []PortMapping, runnerDir string) Task { +// NewTask returns a task with the given identity and status. The state fields are +// assigned by the caller, e.g., by restoreStateFromContainers() +func NewTask(id string, status TaskStatus) Task { return Task{ - ID: id, - Status: status, - containerName: containerName, - containerID: containerID, - runnerDir: runnerDir, - gpuIDs: gpuIDs, - ports: ports, - pullTracker: newPullTracker(), - mu: &sync.Mutex{}, + ID: id, + Status: status, + pullTracker: newPullTracker(), + mu: &sync.Mutex{}, } } diff --git a/runner/internal/shim/volumes.go b/runner/internal/shim/volumes.go index ed88acad1..f6d6096b0 100644 --- a/runner/internal/shim/volumes.go +++ b/runner/internal/shim/volumes.go @@ -7,6 +7,7 @@ import ( "fmt" "os" "os/exec" + "path/filepath" "strings" "time" @@ -90,10 +91,13 @@ func formatAndMountVolume(ctx context.Context, volume VolumeInfo) error { return nil } +// volumeMountPointDir is a dstack-specific dir for volumes, used to avoid clashes with +// host dirs. /mnt/disks is used since on some VM images other places may not be +// writable (e.g. GCP COS). +const volumeMountPointDir = "/mnt/disks/dstack-volumes" + func getVolumeMountPoint(volumeName string) string { - // Put volumes in dstack-specific dir to avoid clashes with host dirs. - // /mnt/disks is used since on some VM images other places may not be writable (e.g. GCP COS). - return fmt.Sprintf("/mnt/disks/dstack-volumes/%s", volumeName) + return filepath.Join(volumeMountPointDir, volumeName) } func prepareInstanceMountPoints(taskConfig TaskConfig) error { From 83c7ec3877d50f650ee4810a38c00f06683b398a Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Thu, 27 Aug 2026 08:03:46 +0000 Subject: [PATCH 2/2] Dummy commit to trigger Actions