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
+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