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
This commit is contained in:
Justin Bradford
2025-10-27 16:49:38 +10:00
committed by GitHub
parent 0153e9f9e2
commit 3b4bcb7932
7 changed files with 653 additions and 1 deletions
+31
View File
@@ -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) {
+188
View File
@@ -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
+6
View File
@@ -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
}
+419
View File
@@ -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 {