From cd8aec24e0454a17fe2f1fc83467f2e8fd4bffe0 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Thu, 27 Aug 2026 13:28:00 +0000 Subject: [PATCH] [shim] Keep shim task files out of the container `~/.dstack/runners/` is mounted into the task container as `/tmp/runner`. It used to hold runner's files only, but shim now keeps its own files there as well -- the image pull log since #2903 and the task state file since #4220 -- sharing them with the runner and the user workload, which may corrupt or delete them. Nothing sensitive is stored there today, but the approach is unsafe: nothing stops a contributor from putting a secret into a file the container can read. That dir is now the task dir, private to shim, and only its new `runner` subdir is mounted into the container: ~/.dstack/runners// 0700, shim only task.json pull.log runner/ 0755, mounted as /tmp/runner * `runnerDir`/`runnersDir` are renamed to `taskDir`/`tasksDir` throughout, including the `DockerParameters` methods, to signal that the dir is managed by shim rather than by runner. The `runners` path itself is kept: an upgraded shim must find the dirs of the tasks created by the previous version. * The dir of a restored task is no longer derived from the container mounts, which now point at the `runner` subdir, but from the task ID in its state file. The tasks dir is scanned once on start, and the result is shared by the restore and the orphan sweep, which used to scan the dir a second time. * A task started by a shim version that did not write state files cannot be found by its state file, so its dir, named after the container, is looked up by name and, as before, removed along with the task. Co-Authored-By: Claude Opus 5 (1M context) --- runner/internal/shim/docker.go | 98 ++++++++++------- runner/internal/shim/docker_test.go | 152 +++++++++++++++++++++------ runner/internal/shim/models.go | 6 +- runner/internal/shim/process_test.go | 2 +- runner/internal/shim/state.go | 110 +++++++++++-------- runner/internal/shim/state_test.go | 68 +++++++----- runner/internal/shim/task.go | 5 +- 7 files changed, 300 insertions(+), 141 deletions(-) diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index 1cbdaf4c1..818d98154 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -203,12 +203,15 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke tasks: NewTaskStorage(), } - if err := runner.restoreStateFromContainers(ctx); err != nil { + // The task dirs are scanned once: the tasks whose dirs are claimed by a container + // are restored, the dirs of the rest are orphaned and swept + storedTasks := scanTaskDirs(ctx, dockerParams.TasksDir()) + if err := runner.restoreStateFromContainers(ctx, storedTasks); 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) + runner.sweepOrphanedTaskDirs(ctx, storedTasks) return runner, nil } @@ -219,8 +222,11 @@ func taskContainerFilters() filters.Args { } // restoreStateFromContainers regenerates TaskStorage and GpuLock inspecting containers +// and the task state files scanned by scanTaskDirs() // Used to restore shim state on restarts -func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error { +func (d *DockerRunner) restoreStateFromContainers( + ctx context.Context, storedTasks map[string]storedTask, +) error { listOptions := container.ListOptions{All: true, Filters: taskContainerFilters()} containers, err := d.client.ContainerList(ctx, listOptions) if err != nil { @@ -282,20 +288,17 @@ func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error { } ports = extractPorts(ctx, containerFull.NetworkSettings.Ports) } - var runnerDir string - for _, mount := range containerShort.Mounts { - if mount.Destination == consts.RunnerTempDir { - runnerDir = mount.Source - break - } - } - state, restored := readRestoredTaskState(ctx, taskID, runnerDir) + storedTask, restored := storedTasks[taskID] + state := storedTask.state config := state.Config + taskDir := storedTask.dir 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) + // Such a task still has a dir that must be removed along with the task + taskDir = d.findLegacyTaskDir(containerName) } if state.CleanedUp { // The resources of this task, its GPUs included, have already been released, @@ -325,7 +328,7 @@ func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error { task.containerID = containerID task.gpuIDs = gpuIDs task.ports = ports - task.runnerDir = runnerDir + task.taskDir = taskDir task.cleanedUp = state.CleanedUp if !d.tasks.Add(task) { log.Error(ctx, "duplicate restored task", "task", taskID) @@ -512,14 +515,14 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { cfg := task.config - runnerDir, err := d.dockerParams.MakeRunnerDir(task.containerName) + taskDir, err := d.dockerParams.MakeTaskDir(task.containerName) if err != nil { - return fmt.Errorf("make runner dir: %w", err) + return fmt.Errorf("make task dir: %w", err) } - log.Trace(ctx, "runner dir", "task", task.ID, "path", runnerDir) + log.Trace(ctx, "task dir", "task", task.ID, "path", taskDir) // 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(ctx, &task, func(t *Task) { t.runnerDir = runnerDir }); err != nil { + if err := d.commit(ctx, &task, func(t *Task) { t.taskDir = taskDir }); err != nil { return fmt.Errorf("%w: failed to update task %s: %w", ErrInternal, task.ID, err) } @@ -571,9 +574,7 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { 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. - // Maybe we should rename it to "task dir" (including the `/root/.dstack/runners` dir on the host). - pullLogPath := filepath.Join(runnerDir, "pull.log") + pullLogPath := filepath.Join(taskDir, "pull.log") if err = pullImage(pullCtx, d.client, cfg, pullLogPath, task.pullTracker); err != nil { errMessage := fmt.Sprintf("pullImage error: %s", err.Error()) log.Error(ctx, errMessage) @@ -782,16 +783,16 @@ func (d *DockerRunner) remove(ctx context.Context, task *Task) (err error) { // but the task may be removed before that happens d.cleanupLocked(ctx, task) // Normally, it should not be empty - if task.runnerDir != "" { - // 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) + if task.taskDir != "" { + // Failed attempts to remove or rename task dir are considered non-fatal + if err := os.RemoveAll(task.taskDir); err != nil { + log.Error(ctx, "failed to remove task directory", "dir", task.taskDir, "err", err) trashName := filepath.Join( - filepath.Dir(task.runnerDir), - fmt.Sprintf(".trash-%s-%d", filepath.Base(task.runnerDir), time.Now().UnixMicro()), + filepath.Dir(task.taskDir), + fmt.Sprintf(".trash-%s-%d", filepath.Base(task.taskDir), 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) + if err := os.Rename(task.taskDir, trashName); err != nil { + log.Error(ctx, "failed to rename task directory", "dir", task.taskDir, "err", err) } } } @@ -913,7 +914,7 @@ func (d *DockerRunner) createContainer( task *Task, options createContainerOptions, ) error { - mounts, err := d.dockerParams.DockerMounts(task.runnerDir) + mounts, err := d.dockerParams.DockerMounts(task.taskDir) if err != nil { return fmt.Errorf("get docker mounts: %w", err) } @@ -1356,11 +1357,11 @@ func (c *CLIArgs) DockerShellCommands(authorizedKeys []string, runnerHttpAddress return append(commands, strings.Join(runnerCommand, " ")) } -func (c *CLIArgs) DockerMounts(hostRunnerDir string) ([]mount.Mount, error) { +func (c *CLIArgs) DockerMounts(hostTaskDir string) ([]mount.Mount, error) { return []mount.Mount{ { Type: mount.TypeBind, - Source: hostRunnerDir, + Source: taskRunnerDir(hostTaskDir), Target: consts.RunnerTempDir, }, { @@ -1375,14 +1376,39 @@ 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") +// tasksDirName is the name of the dir inside shim's home dir that holds the dirs of +// the tasks. A task dir holds shim's own files, such as the task state file and the +// image pull log, and the runner dir, the only part of it mounted into the container. +// Historically, the whole task dir was mounted into the container and held runner's +// files only, hence the name, which is kept for backward compatibility: an upgraded +// shim must find the dirs of the tasks created by the previous version. +const tasksDirName = "runners" + +// taskRunnerDirName is the name of the dir inside a task dir that is mounted into the +// container as consts.RunnerTempDir. Only the files in this dir are shared with the +// container, the rest of the task dir is private to shim. +const taskRunnerDirName = "runner" + +func (c *CLIArgs) TasksDir() string { + return filepath.Join(c.Shim.HomeDir, tasksDirName) } -func (c *CLIArgs) MakeRunnerDir(name string) (string, error) { - runnerTemp := filepath.Join(c.RunnersDir(), name) - if err := os.MkdirAll(runnerTemp, 0o755); err != nil { +// MakeTaskDir creates the dir of the task, including the runner dir inside it, +// and returns the path to the task dir +func (c *CLIArgs) MakeTaskDir(name string) (string, error) { + taskDir := filepath.Join(c.TasksDir(), name) + // Only shim needs access to the task dir itself, unlike the runner dir, which is + // written by the container + if err := os.MkdirAll(taskDir, 0o700); err != nil { + return "", fmt.Errorf("create task directory: %w", err) + } + if err := os.MkdirAll(taskRunnerDir(taskDir), 0o755); err != nil { return "", fmt.Errorf("create runner directory: %w", err) } - return runnerTemp, nil + return taskDir, nil +} + +// taskRunnerDir returns the path to the runner dir inside the given task dir +func taskRunnerDir(taskDir string) string { + return filepath.Join(taskDir, taskRunnerDirName) } diff --git a/runner/internal/shim/docker_test.go b/runner/internal/shim/docker_test.go index 4a80703da..25d26c50b 100644 --- a/runner/internal/shim/docker_test.go +++ b/runner/internal/shim/docker_test.go @@ -38,7 +38,7 @@ func TestDocker_SSHServer(t *testing.T) { params := &dockerParametersMock{ commands: []string{"/usr/sbin/sshd -V 2>&1 | grep OpenSSH"}, sshShellCommands: true, - runnersDir: t.TempDir(), + tasksDir: t.TempDir(), } timeout := 180 // seconds @@ -62,8 +62,8 @@ func TestDocker_ShmNoexecByDefault(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"mount | grep '/dev/shm .*size=65536k' | grep noexec"}, - runnersDir: t.TempDir(), + commands: []string{"mount | grep '/dev/shm .*size=65536k' | grep noexec"}, + tasksDir: t.TempDir(), } timeout := 180 // seconds @@ -87,8 +87,8 @@ func TestDocker_ShmExecIfSizeSpecified(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"mount | grep '/dev/shm .*size=1024k' | grep -v noexec"}, - runnersDir: t.TempDir(), + commands: []string{"mount | grep '/dev/shm .*size=1024k' | grep -v noexec"}, + tasksDir: t.TempDir(), } timeout := 180 // seconds @@ -113,8 +113,8 @@ func TestDocker_ContainerExitedWithError(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"echo failed for a reason", "exit 3"}, - runnersDir: t.TempDir(), + commands: []string{"echo failed for a reason", "exit 3"}, + tasksDir: t.TempDir(), } timeout := 180 // seconds @@ -144,8 +144,8 @@ func TestDocker_RestoredTaskIsTerminated(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"sleep 3", "exit 7"}, - runnersDir: t.TempDir(), + commands: []string{"sleep 3", "exit 7"}, + tasksDir: t.TempDir(), } timeout := 180 // seconds @@ -177,8 +177,8 @@ func TestDocker_RestoredTaskKeepsTerminationReason(t *testing.T) { } params := &dockerParametersMock{ - commands: []string{"sleep 60"}, - runnersDir: t.TempDir(), + commands: []string{"sleep 60"}, + tasksDir: t.TempDir(), } timeout := 180 // seconds @@ -208,7 +208,7 @@ func TestDocker_RestoredTaskKeepsTerminationReason(t *testing.T) { func TestVolumesFromMounts(t *testing.T) { mounts := []dockertypes.MountPoint{ - {Source: "/root/.dstack/runners/task-0-0-1234abcd", Destination: consts.RunnerTempDir}, + {Source: "/root/.dstack/runners/task-0-0-1234abcd/runner", Destination: consts.RunnerTempDir}, {Source: "/mnt/disks/dstack-volumes/volume-1", Destination: "/volume"}, {Source: "/home/dstack/data", Destination: "/instance-data"}, } @@ -257,13 +257,37 @@ func TestShouldRetryWithoutNvidiaDisplayCapability(t *testing.T) { assert.False(t, shouldRetryWithoutNvidiaDisplayCapability(gpu.GpuVendorNvidia, nil)) } +func TestMakeTaskDir(t *testing.T) { + args := CLIArgs{} + args.Shim.HomeDir = t.TempDir() + + taskDir, err := args.MakeTaskDir("task-0-0-1234abcd") + require.NoError(t, err) + + assert.Equal(t, filepath.Join(args.Shim.HomeDir, "runners", "task-0-0-1234abcd"), taskDir) + info, err := os.Stat(taskDir) + require.NoError(t, err) + // Shim's files in the task dir, e.g. the task state file, are not world-readable + assert.Zero(t, info.Mode().Perm()&0o077, "task dir mode: %s", info.Mode()) + // Only the runner dir inside the task dir is shared with the container + runnerDir := filepath.Join(taskDir, "runner") + assert.DirExists(t, runnerDir) + mounts, err := args.DockerMounts(taskDir) + require.NoError(t, err) + assert.Contains(t, mounts, mount.Mount{ + Type: mount.TypeBind, + Source: runnerDir, + Target: consts.RunnerTempDir, + }) +} + /* Mocks */ type dockerParametersMock struct { commands []string sshShellCommands bool - // runnersDir is the parent dir of all task runner dirs - runnersDir string + // tasksDir is the parent dir of all task dirs + tasksDir string } func (c *dockerParametersMock) DockerPrivileged() bool { @@ -291,24 +315,22 @@ func (c *dockerParametersMock) DockerPorts() []int { return []int{} } -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 +func (c *dockerParametersMock) DockerMounts(hostTaskDir string) ([]mount.Mount, error) { return []mount.Mount{ - {Type: mount.TypeBind, Source: hostRunnerDir, Target: consts.RunnerTempDir}, + {Type: mount.TypeBind, Source: taskRunnerDir(hostTaskDir), Target: consts.RunnerTempDir}, }, nil } -func (c *dockerParametersMock) RunnersDir() string { - return c.runnersDir +func (c *dockerParametersMock) TasksDir() string { + return c.tasksDir } -func (c *dockerParametersMock) MakeRunnerDir(name string) (string, error) { - runnerDir := filepath.Join(c.runnersDir, name) - if err := os.MkdirAll(runnerDir, 0o755); err != nil { +func (c *dockerParametersMock) MakeTaskDir(name string) (string, error) { + taskDir := filepath.Join(c.tasksDir, name) + if err := os.MkdirAll(taskRunnerDir(taskDir), 0o755); err != nil { return "", err } - return runnerDir, nil + return taskDir, nil } /* Utilities */ @@ -544,7 +566,7 @@ func nvidiaContainer(containerID, taskID string, gpuIDs []string) (dockertypes.C State: containerStateRunning, Labels: map[string]string{LabelKeyIsTask: LabelValueTrue, LabelKeyTaskID: taskID}, Mounts: []dockertypes.MountPoint{ - {Destination: consts.RunnerTempDir, Source: "/root/.dstack/runners/" + containerID}, + {Destination: consts.RunnerTempDir, Source: "/root/.dstack/runners/" + containerID + "/runner"}, }, } inspection := dockertypes.ContainerJSON{ @@ -570,13 +592,81 @@ func newRestoreRunner(t *testing.T, gpuIDs []string, containers ...dockertypes.C gpuLock, err := NewGpuLock(gpus) require.NoError(t, err) return &DockerRunner{ - client: &restoreClientMock{containers: containers, inspect: map[string]dockertypes.ContainerJSON{}}, - gpuVendor: gpu.GpuVendorNvidia, - gpuLock: gpuLock, - tasks: NewTaskStorage(), + client: &restoreClientMock{containers: containers, inspect: map[string]dockertypes.ContainerJSON{}}, + dockerParams: &dockerParametersMock{tasksDir: t.TempDir()}, + gpuVendor: gpu.GpuVendorNvidia, + gpuLock: gpuLock, + tasks: NewTaskStorage(), } } +// TestRestoreState_FromStateFile checks that the task dir and the part of the task +// state that cannot be recovered from the container are restored from the state file +func TestRestoreState_FromStateFile(t *testing.T) { + containerShort, containerFull := nvidiaContainer("container-1", "task-1", nil) + runner := newRestoreRunner(t, nil, containerShort) + runner.client.(*restoreClientMock).inspect["container-1"] = containerFull + tasksDir := runner.dockerParams.TasksDir() + taskDir := filepath.Join(tasksDir, "container-1") + require.NoError(t, os.MkdirAll(taskRunnerDir(taskDir), 0o755)) + task := NewTaskFromConfig(TaskConfig{ + ID: "task-1", + Name: "task", + HostSshUser: "root", + HostSshKeys: []string{"ssh-ed25519 AAAA"}, + }) + task.SetStatusTerminated(string(types.TerminationReasonTerminatedByUser), "bye") + require.NoError(t, writeTaskState(taskDir, newTaskState(&task))) + + stored := scanTaskDirs(t.Context(), tasksDir) + require.NoError(t, runner.restoreStateFromContainers(t.Context(), stored)) + + restored, ok := runner.tasks.Get("task-1") + require.True(t, ok) + assert.Equal(t, taskDir, restored.taskDir) + assert.Equal(t, TaskStatusTerminated, restored.Status) + assert.Equal(t, string(types.TerminationReasonTerminatedByUser), restored.TerminationReason) + assert.Equal(t, "bye", restored.TerminationMessage) + // The host SSH keys are only known to the state file + assert.Equal(t, []string{"ssh-ed25519 AAAA"}, restored.config.HostSshKeys) +} + +// TestRestoreState_LegacyTaskDir checks that a task started by a shim version that did +// not write state files is restored with its dir, named after the container, so that the +// dir is still removed along with the task +func TestRestoreState_LegacyTaskDir(t *testing.T) { + containerShort, containerFull := nvidiaContainer("container-1", "task-1", nil) + containerShort.Mounts = append(containerShort.Mounts, dockertypes.MountPoint{ + Source: volumeMountPointDir + "/volume-1", Destination: "/volume", + }) + runner := newRestoreRunner(t, nil, containerShort) + runner.client.(*restoreClientMock).inspect["container-1"] = containerFull + legacyDir := filepath.Join(runner.dockerParams.TasksDir(), "container-1") + require.NoError(t, os.MkdirAll(legacyDir, 0o755)) + + require.NoError(t, runner.restoreStateFromContainers(t.Context(), nil)) + + restored, ok := runner.tasks.Get("task-1") + require.True(t, ok) + assert.Equal(t, legacyDir, restored.taskDir) + assert.Equal(t, TaskStatusRunning, restored.Status) + // Without a state file, the volumes are recovered from the container mounts + assert.Equal(t, []VolumeInfo{{Name: "volume-1"}}, restored.config.Volumes) +} + +// TestRestoreState_NoTaskDir checks that a task whose dir is gone is still restored +func TestRestoreState_NoTaskDir(t *testing.T) { + containerShort, containerFull := nvidiaContainer("container-1", "task-1", nil) + runner := newRestoreRunner(t, nil, containerShort) + runner.client.(*restoreClientMock).inspect["container-1"] = containerFull + + require.NoError(t, runner.restoreStateFromContainers(t.Context(), nil)) + + restored, ok := runner.tasks.Get("task-1") + require.True(t, ok) + assert.Empty(t, restored.taskDir) +} + // TestRestoreState_GpuLockedByAnotherTaskIsNotOwned checks that a task restored with // a GPU already locked by another restored task does not claim that GPU, so that // releasing the task's resources does not release a GPU still in use @@ -589,7 +679,7 @@ func TestRestoreState_GpuLockedByAnotherTaskIsNotOwned(t *testing.T) { mock.inspect["container-1"] = firstFull mock.inspect["container-2"] = secondFull - require.NoError(t, runner.restoreStateFromContainers(t.Context())) + require.NoError(t, runner.restoreStateFromContainers(t.Context(), nil)) firstTask, ok := runner.tasks.Get("task-1") require.True(t, ok) @@ -615,7 +705,7 @@ func TestRestoreState_DuplicateTaskReleasesGpus(t *testing.T) { mock.inspect["container-1"] = firstFull mock.inspect["container-2"] = secondFull - require.NoError(t, runner.restoreStateFromContainers(t.Context())) + require.NoError(t, runner.restoreStateFromContainers(t.Context(), nil)) assert.Len(t, runner.tasks.List(), 1) assert.True(t, runner.gpuLock.lock["GPU-beef"], "GPU-beef") diff --git a/runner/internal/shim/models.go b/runner/internal/shim/models.go index e527aa931..6166c7c7f 100644 --- a/runner/internal/shim/models.go +++ b/runner/internal/shim/models.go @@ -8,10 +8,10 @@ type DockerParameters interface { DockerPassEnv() []string DockerPrivileged() bool DockerShellCommands(authorizedKeys []string, runnerHttpAddress string) []string - DockerMounts(string) ([]mount.Mount, error) + DockerMounts(taskDir string) ([]mount.Mount, error) DockerPorts() []int - RunnersDir() string - MakeRunnerDir(name string) (string, error) + TasksDir() string + MakeTaskDir(name string) (string, error) DockerPJRTDevice() string } diff --git a/runner/internal/shim/process_test.go b/runner/internal/shim/process_test.go index dd4b74727..e88e3d20b 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{runnersDir: t.TempDir()}, + dockerParams: &dockerParametersMock{tasksDir: t.TempDir()}, gpuLock: newTestGpuLock(t), tasks: NewTaskStorage(), } diff --git a/runner/internal/shim/state.go b/runner/internal/shim/state.go index 8c41d39d0..7d7243f5a 100644 --- a/runner/internal/shim/state.go +++ b/runner/internal/shim/state.go @@ -12,7 +12,7 @@ import ( "github.com/dstackai/dstack/runner/internal/common/log" ) -// taskStateFileName is the name of the task state file in the task runner dir +// taskStateFileName is the name of the task state file in the task dir const taskStateFileName = "task.json" // taskStateVersion is the version of the task state file format. It is only bumped on @@ -52,18 +52,18 @@ func newTaskState(task *Task) taskState { } } -// 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. +// saveTaskState writes the state of the task to its task dir. Tasks that have not +// acquired any resources yet, that is, tasks without a task 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 == "" { + if !ok || task.taskDir == "" { return } - if err := writeTaskState(task.runnerDir, newTaskState(&task)); err != nil { + if err := writeTaskState(task.taskDir, newTaskState(&task)); err != nil { log.Error(ctx, "failed to save task state", "task", taskID, "err", err) } } @@ -86,9 +86,9 @@ func writeTaskState(dir string, state taskState) error { 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. +// readTaskState reads the state file from the task 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) @@ -105,48 +105,35 @@ func readTaskState(dir string) (taskState, error) { 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 +// storedTask is a task state file found in the tasks dir, along with the task dir it +// was read from +type storedTask struct { + dir string + state taskState } -// 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) +// scanTaskDirs reads the state files of all the task dirs, returning them by task ID. +// This is how the dirs of the tasks started by a previous shim process are found, both +// the dirs of the tasks restored from their containers and the dirs of the orphaned +// tasks swept by sweepOrphanedTaskDirs(). +// Dirs without a state file are skipped, as there is no way to tell whether they belong +// to a task; the dirs of the tasks started by a shim version that did not write state +// files are found by findLegacyTaskDir() instead. +func scanTaskDirs(ctx context.Context, tasksDir string) map[string]storedTask { + stored := make(map[string]storedTask) + entries, err := os.ReadDir(tasksDir) if err != nil { if !errors.Is(err, os.ErrNotExist) { - log.Error(ctx, "failed to list task dirs", "dir", runnersDir, "err", err) + log.Error(ctx, "failed to list task dirs", "dir", tasksDir, "err", err) } - return + return stored } 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()) + dir := filepath.Join(tasksDir, entry.Name()) state, err := readTaskState(dir) if err != nil { if !errors.Is(err, os.ErrNotExist) { @@ -154,18 +141,51 @@ func (d *DockerRunner) sweepOrphanedTaskDirs(ctx context.Context) { } continue } - if _, ok := d.tasks.Get(state.ID); ok { + if other, ok := stored[state.ID]; ok { + log.Error( + ctx, "duplicate task state, ignoring", + "task", state.ID, "dir", dir, "used", other.dir, + ) + continue + } + stored[state.ID] = storedTask{dir: dir, state: state} + } + return stored +} + +// findLegacyTaskDir returns the dir of a task started by a shim version that did not +// write state files, so that the dir is still removed along with the task. Such a dir +// cannot be found by scanTaskDirs() and is identified by its name, which is the name of +// the task container. Returns an empty string if there is no such dir. +func (d *DockerRunner) findLegacyTaskDir(containerName string) string { + if containerName == "" { + return "" + } + dir := filepath.Join(d.dockerParams.TasksDir(), containerName) + if info, err := os.Stat(dir); err != nil || !info.IsDir() { + return "" + } + return dir +} + +// sweepOrphanedTaskDirs releases the resources of the scanned tasks that have no +// container, which normally means that the shim stopped running before the container +// was created, and removes their dirs. The dirs of the tasks restored from their +// containers are left intact. +func (d *DockerRunner) sweepOrphanedTaskDirs(ctx context.Context, storedTasks map[string]storedTask) { + for taskID, stored := range storedTasks { + if _, ok := d.tasks.Get(taskID); 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 { + log.Warning(ctx, "cleaning up orphaned task dir", "task", taskID, "dir", stored.dir) + if !stored.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) + releaseTaskResources(ctx, stored.state.Config) } - if err := os.RemoveAll(dir); err != nil { - log.Error(ctx, "failed to remove orphaned task dir", "dir", dir, "err", err) + if err := os.RemoveAll(stored.dir); err != nil { + log.Error(ctx, "failed to remove orphaned task dir", "dir", stored.dir, "err", err) } } } diff --git a/runner/internal/shim/state_test.go b/runner/internal/shim/state_test.go index a63d2b657..635067413 100644 --- a/runner/internal/shim/state_test.go +++ b/runner/internal/shim/state_test.go @@ -78,7 +78,7 @@ func TestSaveTaskState(t *testing.T) { runner := newTestRunner(t, nil) dir := t.TempDir() task := NewTaskFromConfig(TaskConfig{ID: "task-1", Name: "task"}) - task.runnerDir = dir + task.taskDir = dir addTask(runner, task) runner.saveTaskState(t.Context(), "task-1") @@ -89,7 +89,7 @@ func TestSaveTaskState(t *testing.T) { } // A task that has not acquired any resources yet has no dir to save the state to -func TestSaveTaskState_NoRunnerDir(t *testing.T) { +func TestSaveTaskState_NoTaskDir(t *testing.T) { runner := newTestRunner(t, nil) addTask(runner, NewTaskFromConfig(TaskConfig{ID: "task-1", Name: "task"})) @@ -99,41 +99,61 @@ func TestSaveTaskState_NoRunnerDir(t *testing.T) { assert.NoFileExists(t, taskStateFileName) } +func TestScanTaskDirs(t *testing.T) { + tasksDir := t.TempDir() + dir := makeTaskDir(t, tasksDir, "task-0-0-1234abcd", "task-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 + require.NoError(t, os.MkdirAll(filepath.Join(tasksDir, "legacy-0-0-90abef01"), 0o755)) + // Dot dirs are not task dirs + makeTaskDir(t, tasksDir, ".trash-task-0-0-1234abcd-1", "trashed-task") + + stored := scanTaskDirs(t.Context(), tasksDir) + + require.Len(t, stored, 1) + require.Contains(t, stored, "task-1") + assert.Equal(t, dir, stored["task-1"].dir) + assert.Equal(t, "task-1", stored["task-1"].state.ID) +} + +func TestScanTaskDirs_NoTasksDir(t *testing.T) { + stored := scanTaskDirs(t.Context(), filepath.Join(t.TempDir(), "missing")) + + assert.Empty(t, stored) +} + +func TestFindLegacyTaskDir(t *testing.T) { + runner := newTestRunner(t, nil) + tasksDir := runner.dockerParams.TasksDir() + dir := filepath.Join(tasksDir, "legacy-0-0-90abef01") + require.NoError(t, os.MkdirAll(dir, 0o755)) + + assert.Equal(t, dir, runner.findLegacyTaskDir("legacy-0-0-90abef01")) + assert.Empty(t, runner.findLegacyTaskDir("unknown-0-0-1234abcd")) + // An empty container name must not resolve to the tasks dir itself + assert.Empty(t, runner.findLegacyTaskDir("")) +} + func TestSweepOrphanedTaskDirs(t *testing.T) { runner := newTestRunner(t, nil) - runnersDir := runner.dockerParams.RunnersDir() + tasksDir := runner.dockerParams.TasksDir() // 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") + orphanedDir := makeTaskDir(t, tasksDir, "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") + restoredDir := makeTaskDir(t, tasksDir, "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()) + runner.sweepOrphanedTaskDirs(t.Context(), scanTaskDirs(t.Context(), tasksDir)) 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 { +func makeTaskDir(t *testing.T, tasksDir string, name string, taskID string) string { t.Helper() - dir := filepath.Join(runnersDir, name) - require.NoError(t, os.MkdirAll(dir, 0o755)) + dir := filepath.Join(tasksDir, name) + require.NoError(t, os.MkdirAll(taskRunnerDir(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 1a035b75d..af98841ce 100644 --- a/runner/internal/shim/task.go +++ b/runner/internal/shim/task.go @@ -40,7 +40,10 @@ type Task struct { cancelPull context.CancelFunc gpuIDs []string ports []PortMapping - runnerDir string // path on host mapped to consts.RunnerDir in container + // taskDir is the path on the host to the dir holding shim's files related to the + // task, e.g., the task state file. Its runner subdir, and only it, is mounted into + // the container as consts.RunnerTempDir + taskDir string // startInFlight is true while Start() is working on the task. Start() owns the // task resources until it returns, therefore ProcessTasks() skips such tasks. // Tasks restored from containers are never in flight.