schema for matrix

See iterative/dvc#9725.

examples/valid/matrix.dvc.yaml

@@ -0,0 +1,18 @@
+vars:
+- datasets: [dataset1/, dataset2/]
+
+stages:
+  stage1:
+    matrix:
+      processor: [processor1, processor2]
+      dataset: ${datasets}
+      config:
+        - n_estimators: 150
+          max_depth: 20
+        - n_estimators: 120
+          max_depth: 30
+      labels:
+        - [label1, label2, label3]
+        - [labelX, labelY, labelZ]
+
+    cmd: python train.py ${item.processor} ${item.dataset} ${item.config} ${item.labels.0} ${item.labels.1} ${item.labels.2}

gen.py

@@ -295,10 +295,28 @@ class ForeachDo(BaseModel):
     do: Stage = Field(..., description=DO_DESC)
 
     class Config:
+        allow_mutation = False
+        extra = "forbid"
+
+
+MATRIX_DESC = """\
+Generate stages based on combination of variables.
+
+The variable can be a list of values, or a parametrized string referencing a \
+list."""
+
+
+class Matrix(Stage):
+    matrix: dict[str, list[Any] | ParametrizedString] = Field(
+        ..., description=MATRIX_DESC
+    )
+
+    class Config:
+        allow_mutation = False
         extra = "forbid"
 
 
-Definition = ForeachDo | Stage
+Definition = ForeachDo | Matrix | Stage
 
 
 VARS_DESC = """\

schema.json

@@ -20,6 +20,9 @@
           {
             "$ref": "#/definitions/ForeachDo"
           },
+          {
+            "$ref": "#/definitions/Matrix"
+          },
           {
             "$ref": "#/definitions/Stage"
           }
@@ -497,6 +500,130 @@
       ],
       "additionalProperties": false
     },
+    "Matrix": {
+      "title": "Matrix",
+      "description": "A named stage of a pipeline.",
+      "type": "object",
+      "properties": {
+        "cmd": {
+          "title": "Cmd",
+          "description": "(Required) Command to run (anything your system terminal can run).\n\nCan be a string or a list of commands.",
+          "oneOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "array",
+              "items": {
+                "type": "string"
+              }
+            }
+          ]
+        },
+        "wdir": {
+          "title": "Wdir",
+          "description": "Working directory for the cmd, relative to `dvc.yaml`",
+          "type": "string"
+        },
+        "deps": {
+          "title": "Deps",
+          "description": "List of the dependencies for the stage.",
+          "allOf": [
+            {
+              "$ref": "#/definitions/Dependencies"
+            }
+          ]
+        },
+        "params": {
+          "title": "Params",
+          "description": "List of dot-separated parameter dependency keys to track from `params.yaml`.\n\nMay contain other YAML/JSON/TOML/Python parameter file names, with a sub-list of the param names to track in them (leave empty to include all).",
+          "allOf": [
+            {
+              "$ref": "#/definitions/Params"
+            }
+          ]
+        },
+        "outs": {
+          "title": "Outs",
+          "description": "List of the outputs of the stage.",
+          "allOf": [
+            {
+              "$ref": "#/definitions/Outs"
+            }
+          ]
+        },
+        "metrics": {
+          "title": "Metrics",
+          "description": "List of metrics of the stage written to JSON/TOML/YAML.",
+          "allOf": [
+            {
+              "$ref": "#/definitions/Outs"
+            }
+          ]
+        },
+        "plots": {
+          "title": "Plots",
+          "description": "List of plots of the stage for visualization.\n\nPlots may be written to JSON/YAML/CSV/TSV for data or JPEG/GIF/PNG for images.",
+          "allOf": [
+            {
+              "$ref": "#/definitions/Plots"
+            }
+          ]
+        },
+        "frozen": {
+          "title": "Frozen",
+          "description": "Assume stage as unchanged",
+          "default": false,
+          "type": "boolean"
+        },
+        "always_changed": {
+          "title": "Always Changed",
+          "description": "Assume stage as always changed",
+          "default": false,
+          "type": "boolean"
+        },
+        "vars": {
+          "title": "Vars",
+          "description": "List of stage-specific values for substitution.\n\nMay include any dict or a path to a params file.\n\nUse in the stage with the `${}` substitution expression.",
+          "allOf": [
+            {
+              "$ref": "#/definitions/Vars"
+            }
+          ]
+        },
+        "desc": {
+          "title": "Desc",
+          "description": "Description of the stage",
+          "type": "string"
+        },
+        "meta": {
+          "title": "Meta",
+          "description": "Additional information/metadata"
+        },
+        "matrix": {
+          "title": "Matrix",
+          "description": "Generate stages based on combination of variables.\n\nThe variable can be a list of values, or a parametrized string referencing a list.",
+          "type": "object",
+          "additionalProperties": {
+            "oneOf": [
+              {
+                "type": "array",
+                "items": {}
+              },
+              {
+                "type": "string",
+                "pattern": "^\\${.*?}$"
+              }
+            ]
+          }
+        }
+      },
+      "required": [
+        "cmd",
+        "matrix"
+      ],
+      "additionalProperties": false
+    },
     "X": {
       "title": "X",
       "type": "object",