introduce service.gpus

Signed-off-by: Nicolas De Loof <[email protected]>

Author: ndeloof

loader/loader_test.go

@@ -2224,6 +2224,32 @@ services:
 	assert.ErrorContains(t, err, `capabilities is required`)
 }
 
+func TestServiceGpus(t *testing.T) {
+	p, err := loadYAML(`
+name: service-gpus
+services:
+  test:
+    image: redis:alpine
+    gpus:
+      - driver: nvidia
+      - driver: 3dfx
+        device_ids: ["voodoo2"]
+        capabilities: ["directX"]
+`)
+	assert.NilError(t, err)
+	assert.DeepEqual(t, p.Services["test"].Gpus, []types.DeviceRequest{
+		{
+			Driver: "nvidia",
+			Count:  -1,
+		},
+		{
+			Capabilities: []string{"directX"},
+			Driver:       "3dfx",
+			IDs:          []string{"voodoo2"},
+		},
+	})
+}
+
 func TestServicePullPolicy(t *testing.T) {
 	actual, err := loadYAML(`
 name: service-pull-policy

schema/compose-spec.json

@@ -267,6 +267,7 @@
         },
         "external_links": {"type": "array", "items": {"type": "string"}, "uniqueItems": true},
         "extra_hosts": {"$ref": "#/definitions/extra_hosts"},
+        "gpus": {"$ref": "#/definitions/gpus"},
         "group_add": {
           "type": "array",
           "items": {
@@ -651,6 +652,23 @@
       }
     },
 
+    "gpus": {
+      "id": "#/definitions/gpus",
+      "type": "array",
+      "items": {
+        "type": "object",
+        "properties": {
+          "capabilities": {"$ref": "#/definitions/list_of_strings"},
+          "count": {"type": ["string", "integer"]},
+          "device_ids": {"$ref": "#/definitions/list_of_strings"},
+          "driver":{"type": "string"},
+          "options":{"$ref": "#/definitions/list_or_dict"}
+        },
+        "additionalProperties": false,
+        "patternProperties": {"^x-": {}}
+      }
+    },
+
     "include": {
       "id": "#/definitions/include",
       "oneOf": [

transform/canonical.go

@@ -29,7 +29,6 @@ func init() {
 	transformers["services.*.build.secrets.*"] = transformFileMount
 	transformers["services.*.build.additional_contexts"] = transformKeyValue
 	transformers["services.*.depends_on"] = transformDependsOn
-	transformers["services.*.deploy.resources.reservations.devices.*"] = transformDeviceRequest
 	transformers["services.*.env_file"] = transformEnvFile
 	transformers["services.*.extends"] = transformExtends
 	transformers["services.*.networks"] = transformServiceNetworks

transform/defaults.go

@@ -26,6 +26,8 @@ func init() {
 	defaultValues["services.*.build"] = defaultBuildContext
 	defaultValues["services.*.secrets.*"] = defaultSecretMount
 	defaultValues["services.*.ports.*"] = portDefaults
+	defaultValues["services.*.deploy.resources.reservations.devices.*"] = deviceRequestDefaults
+	defaultValues["services.*.gpus.*"] = deviceRequestDefaults
 }
 
 // SetDefaultValues transforms a compose model to set default values to missing attributes

transform/devices.go

@@ -22,19 +22,15 @@ import (
 	"github.com/compose-spec/compose-go/v2/tree"
 )
 
-func transformDeviceRequest(data any, p tree.Path, ignoreParseError bool) (any, error) {
-	switch v := data.(type) {
-	case map[string]any:
-		_, hasCount := v["count"]
-		_, hasIds := v["device_ids"]
-		if hasCount && hasIds {
-			return nil, fmt.Errorf(`%s: "count" and "device_ids" attributes are exclusive`, p)
-		}
-		if !hasCount && !hasIds {
-			v["count"] = "all"
-		}
-		return transformMapping(v, p, ignoreParseError)
-	default:
+func deviceRequestDefaults(data any, p tree.Path, _ bool) (any, error) {
+	v, ok := data.(map[string]any)
+	if !ok {
 		return data, fmt.Errorf("%s: invalid type %T for device request", p, v)
 	}
+	_, hasCount := v["count"]
+	_, hasIds := v["device_ids"]
+	if !hasCount && !hasIds {
+		v["count"] = "all"
+	}
+	return v, nil
 }