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.