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

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
 }

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) {

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

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
 	}

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 {