[CI] Add ability to configure release tests with a matrix (#51721)

bveeramani · web-flow · commit 75941e7165bd · 2025-03-27T22:58:14.000Z
To ensure Ray Data works performantly in a wide range of scenarios, we want to run release tests with many variants. For example, for our aggregation release tests, we want to run release tests with and without autoscaling, with different cardinalities, and with different implementations. To decrease the maintenance cost and improve the readability of the release test config, this PR introduces a matrix syntax based on [Buildkite's](https://buildkite.com/docs/pipelines/configure/workflows/build-matrix). --------- Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
diff --git a/release/ray_release/config.py b/release/ray_release/config.py
@@ -1,22 +1,22 @@
 import copy
+import itertools
 import json
 import os
 import re
-from typing import Dict, List, Optional, Tuple, Any
+from typing import Any, Dict, List, Optional, Tuple
 
 import jsonschema
 import yaml
-from ray_release.test import (
-    Test,
-    TestDefinition,
-)
 from ray_release.anyscale_util import find_cloud_by_name
 from ray_release.bazel import bazel_runfile
 from ray_release.exception import ReleaseTestCLIError, ReleaseTestConfigError
 from ray_release.logger import logger
+from ray_release.test import (
+    Test,
+    TestDefinition,
+)
 from ray_release.util import DeferredEnvVar, deep_update
 
-
 DEFAULT_WHEEL_WAIT_TIMEOUT = 7200  # Two hours
 DEFAULT_COMMAND_TIMEOUT = 1800
 DEFAULT_BUILD_TIMEOUT = 3600
@@ -84,6 +84,11 @@ def parse_test_definition(test_definitions: List[TestDefinition]) -> List[Test]:
     default_definition = {}
     tests = []
     for test_definition in test_definitions:
+        if "matrix" in test_definition and "variations" in test_definition:
+            raise ReleaseTestConfigError(
+                "You can't specify both 'matrix' and 'variations' in a test definition"
+            )
+
         if test_definition["name"] == "DEFAULTS":
             default_definition = copy.deepcopy(test_definition)
             continue
@@ -93,29 +98,102 @@ def parse_test_definition(test_definitions: List[TestDefinition]) -> List[Test]:
             copy.deepcopy(default_definition), test_definition
         )
 
-        if "variations" not in test_definition:
+        if "variations" in test_definition:
+            tests.extend(_parse_test_definition_with_variations(test_definition))
+        elif "matrix" in test_definition:
+            tests.extend(_parse_test_definition_with_matrix(test_definition))
+        else:
             tests.append(Test(test_definition))
-            continue
 
-        variations = test_definition.pop("variations")
+    return tests
+
+
+def _parse_test_definition_with_variations(
+    test_definition: TestDefinition,
+) -> List[Test]:
+    tests = []
+
+    variations = test_definition.pop("variations")
+    _test_definition_invariant(
+        test_definition,
+        variations,
+        "variations field cannot be empty in a test definition",
+    )
+    for variation in variations:
         _test_definition_invariant(
             test_definition,
-            variations,
-            "variations field cannot be empty in a test definition",
+            "__suffix__" in variation,
+            "missing __suffix__ field in a variation",
         )
-        for variation in variations:
-            _test_definition_invariant(
-                test_definition,
-                "__suffix__" in variation,
-                "missing __suffix__ field in a variation",
+        test = copy.deepcopy(test_definition)
+        test["name"] = f'{test["name"]}.{variation.pop("__suffix__")}'
+        test = deep_update(test, variation)
+        tests.append(Test(test))
+
+    return tests
+
+
+def _parse_test_definition_with_matrix(test_definition: TestDefinition) -> List[Test]:
+    tests = []
+
+    matrix = test_definition.pop("matrix")
+    variables = tuple(matrix["setup"].keys())
+    combinations = itertools.product(*matrix["setup"].values())
+    for combination in combinations:
+        test = test_definition
+        for variable, value in zip(variables, combination):
+            test = _substitute_variable(test, variable, str(value))
+        tests.append(Test(test))
+
+    adjustments = matrix.pop("adjustments", [])
+    for adjustment in adjustments:
+        if not set(adjustment["with"]) == set(variables):
+            raise ReleaseTestConfigError(
+                "You need to specify all matrix variables in the adjustment."
             )
-            test = copy.deepcopy(test_definition)
-            test["name"] = f'{test["name"]}.{variation.pop("__suffix__")}'
-            test = deep_update(test, variation)
-            tests.append(Test(test))
+
+        test = test_definition
+        for variable, value in adjustment["with"].items():
+            test = _substitute_variable(test, variable, str(value))
+        tests.append(Test(test))
+
     return tests
 
 
+def _substitute_variable(data: Dict, variable: str, replacement: str) -> Dict:
+    """Substitute a variable in the provided dictionary with a replacement value.
+
+    This function traverses dict and list values, and substitutes the variable in
+    string values. The syntax for variables is `{{variable}}`.
+
+    Args:
+        data: The dictionary to substitute the variable in.
+        variable: The variable to substitute.
+        replacement: The replacement value.
+
+    Returns:
+        A new dictionary with the variable substituted.
+
+    Examples:
+        >>> test_definition = {"name": "test-{{arg}}"}
+        >>> _substitute_variable(test_definition, "arg", "1")
+        {'name': 'test-1'}
+    """
+    # Create a copy to avoid mutating the original.
+    data = copy.deepcopy(data)
+
+    pattern = r"\{\{\s*" + re.escape(variable) + r"\s*\}\}"
+    for key, value in data.items():
+        if isinstance(value, dict):
+            data[key] = _substitute_variable(value, variable, replacement)
+        elif isinstance(value, list):
+            data[key] = [re.sub(pattern, replacement, string) for string in value]
+        elif isinstance(value, str):
+            data[key] = re.sub(pattern, replacement, value)
+
+    return data
+
+
 def load_schema_file(path: Optional[str] = None) -> Dict:
     path = path or RELEASE_TEST_SCHEMA_FILE
     with open(path, "rt") as fp:
diff --git a/release/ray_release/tests/test_config.py b/release/ray_release/tests/test_config.py
@@ -1,16 +1,17 @@
 import sys
-import yaml
-import pytest
 
-from ray_release.test import Test
+import pytest
+import yaml
 from ray_release.config import (
-    read_and_validate_release_test_collection,
-    validate_cluster_compute,
+    _substitute_variable,
     load_schema_file,
     parse_test_definition,
+    read_and_validate_release_test_collection,
+    validate_cluster_compute,
     validate_test,
 )
 from ray_release.exception import ReleaseTestConfigError
+from ray_release.test import Test
 
 _TEST_COLLECTION_FILES = [
     "release/release_tests.yaml",
@@ -129,6 +130,110 @@ def test_parse_test_definition_with_defaults():
     assert test_with_override["working_dir"] == "overridden_working_dir"
 
 
+def test_parse_test_definition_with_matrix_and_variations_raises():
+    # Matrix and variations are mutually exclusive.
+    test_definitions = yaml.safe_load(
+        """
+        - name: test
+          frequency: nightly
+          team: team
+          working_dir: sample_dir
+          cluster:
+            byod:
+              type: gpu
+            cluster_compute: "{{os}}.yaml"
+          matrix:
+            setup:
+              os: [windows, linux]
+          run:
+            timeout: 100
+            script: python script.py
+          variations:
+            - __suffix__: amd64
+            - __suffix__: arm64
+    """
+    )
+    with pytest.raises(ReleaseTestConfigError):
+        parse_test_definition(test_definitions)
+
+
+def test_parse_test_definition_with_matrix_and_adjustments():
+    test_definitions = yaml.safe_load(
+        """
+        - name: "test-{{compute}}-{{arg}}"
+          matrix:
+            setup:
+              compute: [fixed, autoscaling]
+              arg: [0, 1]
+            adjustments:
+                - with:
+                    # Only run arg 2 with fixed compute
+                    compute: fixed
+                    arg: 2
+          frequency: nightly
+          team: team
+          working_dir: sample_dir
+          cluster:
+            byod:
+              type: gpu
+              runtime_env:
+                - SCALING_MODE={{compute}}
+            cluster_compute: "{{compute}}.yaml"
+          run:
+            timeout: 100
+            script: python script.py --arg "{{arg}}"
+    """
+    )
+    tests = parse_test_definition(test_definitions)
+    schema = load_schema_file()
+
+    assert len(tests) == 5  # 4 from matrix, 1 from adjustments
+    assert not any(validate_test(test, schema) for test in tests)
+    for i, (compute, arg) in enumerate(
+        [
+            ("fixed", 0),
+            ("fixed", 1),
+            ("autoscaling", 0),
+            ("autoscaling", 1),
+            ("fixed", 2),
+        ]
+    ):
+        assert tests[i]["name"] == f"test-{compute}-{arg}"
+        assert tests[i]["cluster"]["cluster_compute"] == f"{compute}.yaml"
+        assert tests[i]["cluster"]["byod"]["runtime_env"] == [f"SCALING_MODE={compute}"]
+
+
+class TestSubstituteVariable:
+    def test_does_not_mutate_original(self):
+        test_definition = {"name": "test-{{arg}}"}
+
+        substituted = _substitute_variable(test_definition, "arg", "1")
+
+        assert substituted is not test_definition
+        assert test_definition == {"name": "test-{{arg}}"}
+
+    def test_substitute_variable_in_string(self):
+        test_definition = {"name": "test-{{arg}}"}
+
+        substituted = _substitute_variable(test_definition, "arg", "1")
+
+        assert substituted == {"name": "test-1"}
+
+    def test_substitute_variable_in_list(self):
+        test_definition = {"items": ["item-{{arg}}"]}
+
+        substituted = _substitute_variable(test_definition, "arg", "1")
+
+        assert substituted == {"items": ["item-1"]}
+
+    def test_substitute_variable_in_dict(self):
+        test_definition = {"outer": {"inner": "item-{{arg}}"}}
+
+        substituted = _substitute_variable(test_definition, "arg", "1")
+
+        assert substituted == {"outer": {"inner": "item-1"}}
+
+
 def test_schema_validation():
     test = VALID_TEST.copy()
 
