Repository: meta-llama/llama-stack
File: scripts/generate_ci_matrix.py (Python)

Code before:
    matrix_config = json.load(f)

DEFAULT_MATRIX = matrix_config["default"]
SCHEDULE_MATRICES: dict[str, list[dict[str, str]]] = matrix_config.get("schedules", {})


def generate_matrix(schedule="", test_setup="", matrix_key="default"):
    """
    Generate test matrix based on schedule, manual input, or matrix key.

    Args:
        schedule: GitHub cron schedule string (e.g., "1 0 * * 0" for weekly)
        test_setup: Manual test setup input (e.g., "ollama-vision")
        matrix_key: Matrix configuration key from ci_matrix.json (e.g., "default", "stainless")

    Returns:
        Matrix configuration as JSON string
    """
    # Weekly scheduled test matrices (highest priority)
    if schedule and schedule in SCHEDULE_MATRICES:
        matrix = SCHEDULE_MATRICES[schedule]
    # Manual input for specific setup
    elif test_setup == "ollama-vision":
        matrix = [{"suite": "vision", "setup": "ollama-vision"}]
    # Use specified matrix key from ci_matrix.json
    elif matrix_key and matrix_key in matrix_config:
        matrix = matrix_config[matrix_key]
    # Default: use JSON-defined default matrix
    else:
        matrix = DEFAULT_MATRIX

    # GitHub Actions expects {"include": [...]} format
    return json.dumps({"include": matrix})


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate CI test matrix")
    parser.add_argument("--schedule", default="", help="GitHub schedule cron string")
    parser.add_argument("--test-setup", default="", help="Manual test setup input")
    parser.add_argument("--matrix-key", default="default", help="Matrix configuration key from ci_matrix.json")

    args = parser.parse_args()

    print(generate_matrix(args.schedule, args.test_setup, args.matrix_key))


Review comment (suggestion):
you should error if the key is not in the config, otherwise you are going to just skip over into the `else` branch. silent error eating is not a good idea.

Code after:
    matrix_config = json.load(f)

DEFAULT_MATRIX = matrix_config["default"]
SCHEDULE_MATRICES: dict[str, list[dict[str, str]]] = matrix_config.get("schedules", {})


def generate_matrix(schedule="", test_setup="", matrix_key="default"):
    """
    Generate test matrix based on schedule, manual input, or matrix key.

    Args:
        schedule: GitHub cron schedule string (e.g., "1 0 * * 0" for weekly)
        test_setup: Manual test setup input (e.g., "ollama-vision")
        matrix_key: Matrix configuration key from ci_matrix.json (e.g., "default", "stainless")

    Returns:
        Matrix configuration as JSON string
    """
    # Weekly scheduled test matrices (highest priority)
    if schedule and schedule in SCHEDULE_MATRICES:
        matrix = SCHEDULE_MATRICES[schedule]
    # Manual input for specific setup
    elif test_setup == "ollama-vision":
        matrix = [{"suite": "vision", "setup": "ollama-vision"}]
    # Use specified matrix key from ci_matrix.json
    elif matrix_key:
        if matrix_key not in matrix_config:
            raise ValueError(f"Invalid matrix_key '{matrix_key}'. Available keys: {list(matrix_config.keys())}")
        matrix = matrix_config[matrix_key]
    # Default: use JSON-defined default matrix
    else:
        matrix = DEFAULT_MATRIX

    # GitHub Actions expects {"include": [...]} format
    return json.dumps({"include": matrix})


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate CI test matrix")
    parser.add_argument("--schedule", default="", help="GitHub schedule cron string")
    parser.add_argument("--test-setup", default="", help="Manual test setup input")
    parser.add_argument("--matrix-key", default="default", help="Matrix configuration key from ci_matrix.json")

    args = parser.parse_args()

    print(generate_matrix(args.schedule, args.test_setup, args.matrix_key))