diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index af6d4b07d..3db376fdc 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -283,6 +283,13 @@ func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error { break } } + 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 @@ -290,13 +297,11 @@ func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error { task := NewTask(taskID, TaskStatusRunning, containerName, containerID, gpuIDs, ports, runnerDir) if !d.tasks.Add(task) { log.Error(ctx, "duplicate restored task", "task", taskID) - } else { - log.Debug(ctx, "restored task", "task", taskID, "state", containerShort.State, "gpus", gpuIDs) - } - if len(gpuIDs) > 0 { - lockedGpuIDs := d.gpuLock.Lock(ctx, gpuIDs) - log.Debug(ctx, "locked GPU(s) due to running task", "task", taskID, "gpus", lockedGpuIDs) + // 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) } return nil } diff --git a/runner/internal/shim/docker_test.go b/runner/internal/shim/docker_test.go index b629a0d18..fd46d51bd 100644 --- a/runner/internal/shim/docker_test.go +++ b/runner/internal/shim/docker_test.go @@ -9,10 +9,14 @@ import ( "testing" "time" + dockertypes "github.com/docker/docker/api/types" "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/mount" + docker "github.com/docker/docker/client" + "github.com/dstackai/dstack/runner/internal/common/consts" "github.com/dstackai/dstack/runner/internal/common/gpu" "github.com/dstackai/dstack/runner/internal/common/types" + "github.com/dstackai/dstack/runner/internal/shim/host" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -443,3 +447,114 @@ func TestPullTracker_NonBytesExtractingUnit(t *testing.T) { assert.Equal(t, uint64(200), p.TotalBytes) assert.True(t, p.IsTotalBytesFinal) } + +/* restoreStateFromContainers */ + +// restoreClientMock implements the only two docker.APIClient methods used by +// restoreStateFromContainers, letting it run without a Docker daemon. +// Any other method panics with a nil pointer dereference, which is intentional: +// this mock must be updated if restoreStateFromContainers starts using them +type restoreClientMock struct { + docker.APIClient + containers []dockertypes.Container + inspect map[string]dockertypes.ContainerJSON +} + +func (c *restoreClientMock) ContainerList( + context.Context, container.ListOptions, +) ([]dockertypes.Container, error) { + return c.containers, nil +} + +func (c *restoreClientMock) ContainerInspect( + _ context.Context, containerID string, +) (dockertypes.ContainerJSON, error) { + return c.inspect[containerID], nil +} + +// nvidiaContainer returns a task container summary and its inspection response +// with the given GPU IDs requested +func nvidiaContainer(containerID, taskID string, gpuIDs []string) (dockertypes.Container, dockertypes.ContainerJSON) { + summary := dockertypes.Container{ + ID: containerID, + Names: []string{"/" + containerID}, + State: containerStateRunning, + Labels: map[string]string{LabelKeyIsTask: LabelValueTrue, LabelKeyTaskID: taskID}, + Mounts: []dockertypes.MountPoint{ + {Destination: consts.RunnerTempDir, Source: "/root/.dstack/runners/" + containerID}, + }, + } + inspection := dockertypes.ContainerJSON{ + ContainerJSONBase: &dockertypes.ContainerJSONBase{ + HostConfig: &container.HostConfig{ + Resources: container.Resources{ + DeviceRequests: []container.DeviceRequest{{DeviceIDs: gpuIDs}}, + }, + }, + }, + Config: &container.Config{}, + NetworkSettings: &dockertypes.NetworkSettings{}, + } + return summary, inspection +} + +func newRestoreRunner(t *testing.T, gpuIDs []string, containers ...dockertypes.Container) *DockerRunner { + t.Helper() + gpus := make([]host.GpuInfo, 0, len(gpuIDs)) + for _, id := range gpuIDs { + gpus = append(gpus, host.GpuInfo{Vendor: gpu.GpuVendorNvidia, ID: id}) + } + 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(), + } +} + +// 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 +func TestRestoreState_GpuLockedByAnotherTaskIsNotOwned(t *testing.T) { + firstShort, firstFull := nvidiaContainer("container-1", "task-1", []string{"GPU-beef"}) + // Should not happen, but if it does, the second task must not take over the GPU + secondShort, secondFull := nvidiaContainer("container-2", "task-2", []string{"GPU-beef", "GPU-f00d"}) + runner := newRestoreRunner(t, []string{"GPU-beef", "GPU-f00d"}, firstShort, secondShort) + mock := runner.client.(*restoreClientMock) + mock.inspect["container-1"] = firstFull + mock.inspect["container-2"] = secondFull + + require.NoError(t, runner.restoreStateFromContainers(t.Context())) + + firstTask, ok := runner.tasks.Get("task-1") + require.True(t, ok) + assert.Equal(t, []string{"GPU-beef"}, firstTask.gpuIDs) + secondTask, ok := runner.tasks.Get("task-2") + require.True(t, ok) + assert.Equal(t, []string{"GPU-f00d"}, secondTask.gpuIDs, "GPU-beef is owned by task-1") + + // Releasing the second task must not free the GPU owned by the first one + runner.gpuLock.Release(t.Context(), secondTask.gpuIDs) + assert.True(t, runner.gpuLock.lock["GPU-beef"], "GPU-beef") + assert.False(t, runner.gpuLock.lock["GPU-f00d"], "GPU-f00d") +} + +// TestRestoreState_DuplicateTaskReleasesGpus checks that the GPUs locked for a task +// that cannot be stored are released, as nothing else will release them +func TestRestoreState_DuplicateTaskReleasesGpus(t *testing.T) { + firstShort, firstFull := nvidiaContainer("container-1", "task-1", []string{"GPU-beef"}) + // Two containers reporting the same task ID: the second one is not stored + secondShort, secondFull := nvidiaContainer("container-2", "task-1", []string{"GPU-f00d"}) + runner := newRestoreRunner(t, []string{"GPU-beef", "GPU-f00d"}, firstShort, secondShort) + mock := runner.client.(*restoreClientMock) + mock.inspect["container-1"] = firstFull + mock.inspect["container-2"] = secondFull + + require.NoError(t, runner.restoreStateFromContainers(t.Context())) + + assert.Len(t, runner.tasks.List(), 1) + assert.True(t, runner.gpuLock.lock["GPU-beef"], "GPU-beef") + assert.False(t, runner.gpuLock.lock["GPU-f00d"], "GPU-f00d is not owned by any task") +}