mirror of
https://github.com/psviderski/uncloud.git
synced 2026-08-27 19:43:34 +00:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user