Repository: compose-spec/compose-go
Pull request #732 (2025-01-23 13:00:35)

support gpus: all alias
Add syntaxic sugar `gpus: all` to select all available GPUS, aligned with `docker run --gpus=all`


Patch:
diff --git a/schema/compose-spec.json b/schema/compose-spec.json
index a228d19..bfaa72d 100644
--- a/schema/compose-spec.json
+++ b/schema/compose-spec.json
@@ -657,19 +657,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-": {}}
-      }
+      "oneOf": [
+        {"type": "string", "enum": ["all"]},
+        {"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": {
diff --git a/transform/canonical.go b/transform/canonical.go
index 2ace976..d37eb1e 100644
--- a/transform/canonical.go
+++ b/transform/canonical.go
@@ -32,6 +32,7 @@ func init() {
 	transformers["services.*.env_file"] = transformEnvFile
 	transformers["services.*.label_file"] = transformStringOrList
 	transformers["services.*.extends"] = transformExtends
+	transformers["services.*.gpus"] = transformGpus
 	transformers["services.*.networks"] = transformServiceNetworks
 	transformers["services.*.volumes.*"] = transformVolumeMount
 	transformers["services.*.dns"] = transformStringOrList
diff --git a/transform/gpus.go b/transform/gpus.go
new file mode 100644
index 0000000..3c91b63
--- /dev/null
+++ b/transform/gpus.go
@@ -0,0 +1,38 @@
+/*
+   Copyright 2020 The Compose Specification Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package transform
+
+import (
+	"fmt"
+
+	"github.com/compose-spec/compose-go/v2/tree"
+)
+
+func transformGpus(data any, p tree.Path, ignoreParseError bool) (any, error) {
+	switch v := data.(type) {
+	case []any:
+		return transformSequence(v, p, ignoreParseError)
+	case string:
+		return []any{
+			map[string]any{
+				"count": "all",
+			},
+		}, nil
+	default:
+		return data, fmt.Errorf("%s: invalid type %T for gpus", p, v)
+	}
+}


Test patch:
diff --git a/loader/loader_test.go b/loader/loader_test.go
index fbe9387..234107f 100644
--- a/loader/loader_test.go
+++ b/loader/loader_test.go
@@ -3738,3 +3738,15 @@ services:
 	assert.NilError(t, err)
 	assert.Equal(t, len(p.Services["test"].DNS), 0)
 }
+
+func TestAllGPUS(t *testing.T) {
+	p, err := loadYAML(`
+name: load-all-gpus
+services:
+  test:
+    gpus: all
+`)
+	assert.NilError(t, err)
+	assert.Equal(t, len(p.Services["test"].Gpus), 1)
+	assert.Equal(t, p.Services["test"].Gpus[0].Count, types.DeviceCount(-1))
+}