From 3b4bcb7932f65bb651160e76fa065a95666c0337 Mon Sep 17 00:00:00 2001 From: Justin Bradford Date: Sun, 26 Oct 2025 23:49:38 -0700 Subject: [PATCH] feat: support service.gpus and reservations.devices (#156) * feat: support compose.yaml `service.gpus` and `service.deploy.resources.reservations.devices` DeviceRequests to Docker container * Fix lint error * Adjust test comments to work with linter * Rename DeviceRequests to DeviceReservations to be more consistent with compose nomenclature * Device reservation changes are immutable and should trigger a container recreate rather than update --- internal/machine/docker/server.go | 1 + pkg/api/resources.go | 6 + pkg/client/compose/service.go | 31 ++ pkg/client/compose/service_test.go | 188 ++++++++ pkg/client/deploy/container.go | 6 + pkg/client/deploy/container_test.go | 419 ++++++++++++++++++ .../3-concepts/5-compose/1-support-matrix.md | 3 +- 7 files changed, 653 insertions(+), 1 deletion(-) diff --git a/internal/machine/docker/server.go b/internal/machine/docker/server.go index fa423c2a..b5f3807f 100644 --- a/internal/machine/docker/server.go +++ b/internal/machine/docker/server.go @@ -611,6 +611,7 @@ func (s *Server) CreateServiceContainer( NanoCPUs: spec.Container.Resources.CPU, Memory: spec.Container.Resources.Memory, MemoryReservation: spec.Container.Resources.MemoryReservation, + DeviceRequests: spec.Container.Resources.DeviceReservations, }, // Restart service containers if they exit or a machine restarts unless they are explicitly stopped. // For one-off containers and batch jobs we plan to use a different service type/mode. diff --git a/pkg/api/resources.go b/pkg/api/resources.go index 97b4e460..8aece7d3 100644 --- a/pkg/api/resources.go +++ b/pkg/api/resources.go @@ -1,5 +1,9 @@ package api +import ( + "github.com/docker/docker/api/types/container" +) + const ( MilliCore = 1_000_000 Core = 1000 * MilliCore @@ -13,4 +17,6 @@ type ContainerResources struct { // MemoryReservation is the minimum amount of memory (in bytes) the container needs to run efficiently. // TODO: implement a placement constraint that checks available memory on machines. MemoryReservation int64 + // Device reservations/requests for access to things like GPUs + DeviceReservations []container.DeviceRequest } diff --git a/pkg/client/compose/service.go b/pkg/client/compose/service.go index 681a6363..953e1a93 100644 --- a/pkg/client/compose/service.go +++ b/pkg/client/compose/service.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/compose-spec/compose-go/v2/types" + "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/mount" "github.com/opencontainers/go-digest" "github.com/psviderski/uncloud/pkg/api" @@ -124,6 +125,9 @@ func resourcesFromCompose(service types.ServiceConfig) api.ContainerResources { MemoryReservation: int64(service.MemReservation), } + // Convert GPU device requests from compose format, appending "gpu" capability. + resources.DeviceReservations = append(resources.DeviceReservations, deviceReservationsFromCompose(service.Gpus, "gpu")...) + // Map resources from deploy section if specified. if service.Deploy != nil { if service.Deploy.Resources.Limits != nil { @@ -139,12 +143,39 @@ func resourcesFromCompose(service types.ServiceConfig) api.ContainerResources { if service.Deploy.Resources.Reservations.MemoryBytes > 0 { resources.MemoryReservation = int64(service.Deploy.Resources.Reservations.MemoryBytes) } + // Handle arbitrary device reservations (same structure as Gpus above). + resources.DeviceReservations = append(resources.DeviceReservations, deviceReservationsFromCompose(service.Deploy.Resources.Reservations.Devices)...) } } return resources } +// Converts compose-go DeviceRequest format to Docker API DeviceRequest format. +// Additional capabilities can be appended via extraCapabilities (e.g., "gpu" for service.Gpus). +func deviceReservationsFromCompose(devices []types.DeviceRequest, extraCapabilities ...string) []container.DeviceRequest { + if devices == nil { + return nil + } + + requests := make([]container.DeviceRequest, 0, len(devices)) + for _, deviceRequest := range devices { + // Docker expects an OR'd list of AND'd capabilities (e.g. [][]string), + // but compose-go provides a single AND'd list (e.g. []string). + capabilities := [][]string{append(deviceRequest.Capabilities, extraCapabilities...)} + + spec := container.DeviceRequest{ + Driver: deviceRequest.Driver, + Count: int(deviceRequest.Count), + DeviceIDs: deviceRequest.IDs, + Capabilities: capabilities, + Options: deviceRequest.Options, + } + requests = append(requests, spec) + } + return requests +} + func volumeSpecsFromCompose( volumes types.Volumes, serviceVolumes []types.ServiceVolumeConfig, ) ([]api.VolumeSpec, []api.VolumeMount, error) { diff --git a/pkg/client/compose/service_test.go b/pkg/client/compose/service_test.go index b1267621..b41951e4 100644 --- a/pkg/client/compose/service_test.go +++ b/pkg/client/compose/service_test.go @@ -10,6 +10,7 @@ import ( "github.com/compose-spec/compose-go/v2/loader" "github.com/compose-spec/compose-go/v2/types" + "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/mount" "github.com/docker/go-units" "github.com/google/go-cmp/cmp" @@ -434,6 +435,193 @@ services: } } +func TestServiceSpecFromCompose_GPUs(t *testing.T) { + tests := []struct { + name string + composeYAML string + expectedDeviceReqs []container.DeviceRequest + }{ + { + name: "gpus_all_shorthand", + composeYAML: ` +services: + ai: + image: nvidia/cuda + gpus: all +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + Count: -1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + { + name: "gpus_device_ids", + composeYAML: ` +services: + ai: + image: nvidia/cuda + gpus: + - device_ids: ['0', '1'] + capabilities: [compute] +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + DeviceIDs: []string{"0", "1"}, + Capabilities: [][]string{{"compute", "gpu"}}, + }, + }, + }, + { + name: "gpus_count", + composeYAML: ` +services: + ai: + image: nvidia/cuda + gpus: + - count: 2 + capabilities: [utility] +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + Count: 2, + Capabilities: [][]string{{"utility", "gpu"}}, + }, + }, + }, + { + name: "gpus_driver_and_options", + composeYAML: ` +services: + ai: + image: nvidia/cuda + gpus: + - driver: nvidia + count: 1 + capabilities: [compute, utility] + options: + key1: value1 + key2: value2 +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 1, + Capabilities: [][]string{{"compute", "utility", "gpu"}}, + Options: map[string]string{ + "key1": "value1", + "key2": "value2", + }, + }, + }, + }, + { + name: "deploy_resources_reservations_devices", + composeYAML: ` +services: + ai: + image: nvidia/cuda + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: [gpu] +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + { + name: "gpus_with_existing_gpu_capability", + composeYAML: ` +services: + ai: + image: nvidia/cuda + gpus: + - capabilities: [gpu, compute] +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + // defaults to "all" when count not specified + Count: -1, + // gpu appended even if already present (matching Docker Compose behavior) + Capabilities: [][]string{{"gpu", "compute", "gpu"}}, + }, + }, + }, + { + name: "multiple_device_reservations", + composeYAML: ` +services: + ai: + image: nvidia/cuda + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 2 + capabilities: [gpu, compute] + - capabilities: [tpu] + count: 1 +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + Capabilities: [][]string{{"gpu", "compute"}}, + }, + { + Count: 1, + Capabilities: [][]string{{"tpu"}}, + }, + }, + }, + { + name: "tpu_reservation", + composeYAML: ` +services: + ai: + image: tensorflow/tensorflow:latest + deploy: + resources: + reservations: + devices: + - capabilities: [tpu] + count: 4 + driver: google +`, + expectedDeviceReqs: []container.DeviceRequest{ + { + Driver: "google", + Count: 4, + Capabilities: [][]string{{"tpu"}}, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + project, err := loadProjectFromContent(t, tt.composeYAML) + require.NoError(t, err) + + spec, err := ServiceSpecFromCompose(project, "ai") + require.NoError(t, err) + + assert.Equal(t, tt.expectedDeviceReqs, spec.Container.Resources.DeviceReservations, + "DeviceReservations should match expected") + }) + } +} + func TestServiceSpecFromCompose_XMachinesPlacement(t *testing.T) { tests := []struct { name string diff --git a/pkg/client/deploy/container.go b/pkg/client/deploy/container.go index 739e305a..361740d1 100644 --- a/pkg/client/deploy/container.go +++ b/pkg/client/deploy/container.go @@ -83,11 +83,17 @@ func EvalContainerSpecChange(current api.ServiceSpec, new api.ServiceSpec) Conta } } + // Device reservations are immutable, so we'll need to recreate if any have changed + if !reflect.DeepEqual(current.Container.Resources.DeviceReservations, newResources.DeviceReservations) { + return ContainerNeedsRecreate + } + // Check if any mutable properties changed. if !current.Caddy.Equals(new.Caddy) { return ContainerNeedsRecreate } + // Remaining resources are mutable. if !reflect.DeepEqual(current.Container.Resources, newResources) { return ContainerNeedsUpdate } diff --git a/pkg/client/deploy/container_test.go b/pkg/client/deploy/container_test.go index 705869e8..3283d4e9 100644 --- a/pkg/client/deploy/container_test.go +++ b/pkg/client/deploy/container_test.go @@ -3,6 +3,7 @@ package deploy import ( "testing" + "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/mount" "github.com/psviderski/uncloud/pkg/api" "github.com/stretchr/testify/assert" @@ -1350,6 +1351,398 @@ func TestEvalContainerSpecChange_Volumes(t *testing.T) { } } +func TestEvalContainerSpecChange_DeviceReservations(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + current api.ContainerResources + new api.ContainerResources + want ContainerSpecStatus + }{ + { + name: "empty", + current: api.ContainerResources{}, + new: api.ContainerResources{}, + want: ContainerUpToDate, + }, + { + name: "both nil", + current: api.ContainerResources{DeviceReservations: nil}, + new: api.ContainerResources{DeviceReservations: nil}, + want: ContainerUpToDate, + }, + { + name: "both empty", + current: api.ContainerResources{DeviceReservations: []container.DeviceRequest{}}, + new: api.ContainerResources{DeviceReservations: []container.DeviceRequest{}}, + want: ContainerUpToDate, + }, + { + name: "set GPU request", + current: api.ContainerResources{}, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: -1, // all + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "unset GPU request", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: -1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + new: api.ContainerResources{}, + want: ContainerNeedsRecreate, + }, + { + name: "identical GPU request", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + Capabilities: [][]string{{"gpu", "compute"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + Capabilities: [][]string{{"gpu", "compute"}}, + }, + }, + }, + want: ContainerUpToDate, + }, + { + name: "change GPU count", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 2, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "change GPU driver", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "amd", + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "change device IDs", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + DeviceIDs: []string{"0"}, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + DeviceIDs: []string{"0", "1"}, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "change capabilities", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu", "compute"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "set options", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + Options: map[string]string{ + "key": "value", + }, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "change options", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + Options: map[string]string{ + "key": "value1", + }, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + Options: map[string]string{ + "key": "value2", + }, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "unset options", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + Options: map[string]string{ + "key": "value", + }, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "add device request", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + { + Count: 1, + Capabilities: [][]string{{"tpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "remove device request", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + { + Count: 1, + Capabilities: [][]string{{"tpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "multiple identical device requests", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + Capabilities: [][]string{{"gpu", "compute"}}, + }, + { + Count: 1, + Capabilities: [][]string{{"tpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + Capabilities: [][]string{{"gpu", "compute"}}, + }, + { + Count: 1, + Capabilities: [][]string{{"tpu"}}, + }, + }, + }, + want: ContainerUpToDate, + }, + { + name: "reordered device requests", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + Capabilities: [][]string{{"gpu"}}, + }, + { + Count: 1, + Capabilities: [][]string{{"tpu"}}, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"tpu"}}, + }, + { + Driver: "nvidia", + Count: 2, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + want: ContainerNeedsRecreate, + }, + { + name: "complex GPU configuration identical", + current: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + DeviceIDs: []string{"0", "1"}, + Capabilities: [][]string{{"gpu", "compute", "utility"}}, + Options: map[string]string{ + "runtime": "nvidia", + "compute": "exclusive", + }, + }, + }, + }, + new: api.ContainerResources{ + DeviceReservations: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: 2, + DeviceIDs: []string{"0", "1"}, + Capabilities: [][]string{{"gpu", "compute", "utility"}}, + Options: map[string]string{ + "runtime": "nvidia", + "compute": "exclusive", + }, + }, + }, + }, + want: ContainerUpToDate, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + currentSpec := api.ServiceSpec{ + Container: api.ContainerSpec{ + Image: "nvidia/cuda:latest", + Resources: tt.current, + }, + } + newSpec := api.ServiceSpec{ + Container: api.ContainerSpec{ + Image: "nvidia/cuda:latest", + Resources: tt.new, + }, + } + + result := EvalContainerSpecChange(currentSpec, newSpec) + assert.Equal(t, tt.want, result) + }) + } +} + func TestEvalContainerSpecChange_Mixed(t *testing.T) { t.Parallel() @@ -1399,6 +1792,32 @@ func TestEvalContainerSpecChange_Mixed(t *testing.T) { }, want: ContainerNeedsRecreate, }, + { + name: "mutable memory change with immutable device reservation change", + current: api.ServiceSpec{ + Container: api.ContainerSpec{ + Image: "nvidia/cuda:latest", + Resources: api.ContainerResources{ + Memory: 100 * 1024 * 1024, + }, + }, + }, + new: api.ServiceSpec{ + Container: api.ContainerSpec{ + Image: "nvidia/cuda:latest", + Resources: api.ContainerResources{ + Memory: 200 * 1024 * 1024, + DeviceReservations: []container.DeviceRequest{ + { + Count: 1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + }, + }, + }, + want: ContainerNeedsRecreate, + }, } for _, tt := range tests { diff --git a/website/docs/3-concepts/5-compose/1-support-matrix.md b/website/docs/3-concepts/5-compose/1-support-matrix.md index 9d0bf551..47d1c7bd 100644 --- a/website/docs/3-concepts/5-compose/1-support-matrix.md +++ b/website/docs/3-concepts/5-compose/1-support-matrix.md @@ -16,6 +16,7 @@ The following table shows the support status for main Compose features: | `entrypoint` | ✅ Supported | Override container entrypoint | | `env_file` | ✅ Supported | Environment file | | `environment` | ✅ Supported | Environment variables | +| `gpus` | ✅ Supported | GPU device access | | `image` | ✅ Supported | Container image specification | | `init` | ✅ Supported | Run init process in container | | `labels` | ❌ Not supported | | @@ -39,7 +40,7 @@ The following table shows the support status for main Compose features: | `mode` | ✅ Supported | Either `global` or `replicated` | | `placement` | ❌ Not supported | Use `x-machines` extension | | `replicas` | ✅ Supported | Number of container replicas | -| `resources` | ⚠️ Limited | CPU and memory limits only | +| `resources` | ⚠️ Limited | CPU, memory limits and device reservations | | `restart_policy` | ❌ Not supported | Defaults to `unless-stopped` | | **Volumes** | | | | Named volumes | ✅ Supported | Docker volumes |