fetch_test_configurations.py

# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""
This script determines what test configurations to run.

Outputs (written to $GITHUB_OUTPUT):
  - sanity_component: JSON object for the sanity component, always present as a
    prerequisite that must pass before other components are run.
  - components: JSON array of component configs for the regular test matrix
    (excludes sanity, which is output separately above).
  - platform: lowercase OS name derived from RUNNER_OS.

Required environment variables:
  - RUNNER_OS (https://docs.github.com/en/actions/how-tos/writing-workflows/choosing-what-your-workflow-does/store-information-in-variables#detecting-the-operating-system)
"""

import ast
import json
import logging
import os
import sys
from pathlib import Path

# Add tests directory to path for extended_tests imports
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests"))
from github_actions_api import *
from extended_tests.benchmark.benchmark_test_matrix import benchmark_matrix
from extended_tests.functional.functional_test_matrix import functional_matrix
from amdgpu_family_matrix import get_all_families_for_trigger_types

logging.basicConfig(level=logging.INFO)

# Note: these paths are relative to the repository root. We could make that
# more explicit, or use absolute paths.
SCRIPT_DIR = Path("./build_tools/github_actions/test_executable_scripts")


def _get_script_path(script_name: str) -> str:
    platform_path = SCRIPT_DIR / script_name
    # Convert to posix (using `/` instead of `\\`) so test workflows can use
    # 'bash' as the shell on Linux and Windows.
    posix_path = platform_path.as_posix()
    return str(posix_path)


test_matrix = {
    # Sanity tests - always run first as a prerequisite for other component tests
    "sanity": {
        "job_name": "sanity",
        "fetch_artifact_args": "--base-only",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_sanity.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
        # Running docker with cap-add and -v /lib/modules, by recommendation of GitHub:
        # https://rocm.docs.amd.com/projects/amdsmi/en/amd-staging/how-to/setup-docker-container.html
        "container_options": "--cap-add SYS_MODULE --cap-add SYS_PTRACE -v /lib/modules:/lib/modules",
    },
    # hip-tests
    "hip-tests": {
        "job_name": "hip-tests",
        "fetch_artifact_args": "--tests",
        "timeout_minutes": 120,
        "test_script": f"python {_get_script_path('test_hiptests.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 4,
            "windows": 4,
        },
    },
    # BLAS tests
    "rocblas": {
        "job_name": "rocblas",
        "fetch_artifact_args": "--blas --tests",
        # GHA step timeout: max category timeout in rocBLAS should be 24 hours / 6 shards = 4 hours per shard
        # 240 min + 20% margin = 288 min
        "timeout_minutes": 288,
        "test_script": f"python {_get_script_path('test_runner.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 6,
            "windows": 6,
        },
    },
    "rocroller": {
        "job_name": "rocroller",
        "fetch_artifact_args": "--blas --tests",
        "timeout_minutes": 60,
        "test_script": f"python {_get_script_path('test_rocroller.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 5,
            "windows": 5,
        },
        "exclude_family": {
            # rocroller does not plan to support Linux and Windows gfx115X architectures
            "linux": [
                "gfx1150",
                "gfx1151",
                "gfx1152",
                "gfx1153",
            ],
            "windows": [
                "gfx1150",
                "gfx1151",
                "gfx1152",
                "gfx1153",
            ],
        },
    },
    "hipblas": {
        "job_name": "hipblas",
        "fetch_artifact_args": "--blas --tests",
        "timeout_minutes": 30,
        "test_script": f"python {_get_script_path('test_hipblas.py')}",
        "platform": ["linux", "windows"],
        # TODO(#2616): Enable full tests once known machine issues are resolved
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "hipblaslt": {
        "job_name": "hipblaslt",
        "fetch_artifact_args": "--blas --tests",
        "timeout_minutes": 180,
        "test_script": f"python {_get_script_path('test_hipblaslt.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 6,
            "windows": 1,
        },
    },
    # SOLVER tests
    "hipsolver": {
        "job_name": "hipsolver",
        "fetch_artifact_args": "--blas --tests",
        "timeout_minutes": 5,
        "test_script": f"python {_get_script_path('test_hipsolver.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "rocsolver": {
        "job_name": "rocsolver",
        "fetch_artifact_args": "--blas --tests",
        # 68350(approx) tests needs 48 mins, so 48 mins / 2 shards = 24 mins per shard
        # 24 mins + 20% margin = 30 mins => ~40 mins (considering gpu delays and lags)
        "timeout_minutes": 40,
        "test_script": f"python {_get_script_path('test_rocsolver.py')}",
        # Issue for adding windows tests: https://github.com/ROCm/TheRock/issues/1770
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 2,
            "windows": 2,
        },
    },
    # PRIM tests
    "rocprim": {
        "job_name": "rocprim",
        "fetch_artifact_args": "--prim --tests",
        "timeout_minutes": 30,
        "test_script": f"python {_get_script_path('test_rocprim.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 2,
            "windows": 2,
        },
    },
    "hipcub": {
        "job_name": "hipcub",
        "fetch_artifact_args": "--prim --tests",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_hipcub.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "rocgdb": {
        "job_name": "rocgdb",
        "fetch_artifact_args": "--debug-tools --tests",
        "timeout_minutes": 45,
        "test_script": f"python {_get_script_path('test_rocgdb.py')}",
        "platform": ["linux"],
        "total_shards": 1,
        "container_image": "ghcr.io/rocm/no_rocm_image_ubuntu24_04_rocgdb@sha256:939b8e35887144d1ca4eca928dc2869991339cab869168790e495fc0a5907bbb",
        "container_options": "--cap-add=SYS_PTRACE",
    },
    "rocr-debug-agent": {
        "job_name": "rocr-debug-agent",
        "fetch_artifact_args": "--debug-tools --tests",
        "timeout_minutes": 10,
        "test_script": f"python {_get_script_path('test_rocr-debug-agent.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "rocthrust": {
        "job_name": "rocthrust",
        "fetch_artifact_args": "--prim --tests",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_rocthrust.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # SPARSE tests
    "hipsparse": {
        "job_name": "hipsparse",
        "fetch_artifact_args": "--blas --tests",
        "timeout_minutes": 30,
        "test_script": f"python {_get_script_path('test_hipsparse.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "rocsparse": {
        "job_name": "rocsparse",
        "fetch_artifact_args": "--blas --tests",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_rocsparse.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "hipsparselt": {
        "job_name": "hipsparselt",
        "fetch_artifact_args": "--blas --tests",
        "timeout_minutes": 30,
        "test_script": f"python {_get_script_path('test_hipsparselt.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # RAND tests
    "rocrand": {
        "job_name": "rocrand",
        "fetch_artifact_args": "--rand --tests",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_rocrand.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "hiprand": {
        "job_name": "hiprand",
        "fetch_artifact_args": "--rand --tests",
        "timeout_minutes": 5,
        "test_script": f"python {_get_script_path('test_hiprand.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # FFT tests
    "rocfft": {
        "job_name": "rocfft",
        "fetch_artifact_args": "--fft --rand --tests",
        "timeout_minutes": 60,
        "test_script": f"python {_get_script_path('test_rocfft.py')}",
        # TODO(geomin12): Add windows test (https://github.com/ROCm/TheRock/issues/1391)
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "hipfft": {
        "job_name": "hipfft",
        "fetch_artifact_args": "--fft --rand --tests",
        "timeout_minutes": 60,
        "test_script": f"python {_get_script_path('test_hipfft.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 2,
            "windows": 2,
        },
    },
    # MIOpen tests
    "miopen": {
        "job_name": "miopen",
        "fetch_artifact_args": "--blas --miopen --rand --tests",
        "timeout_minutes": 60,
        "test_script": f"python {_get_script_path('test_runner.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 4,
            "windows": 4,
        },
    },
    # RCCL tests
    "rccl": {
        "job_name": "rccl",
        "fetch_artifact_args": "--rccl --tests",
        "timeout_minutes": 15,
        "test_script": f"pytest {_get_script_path('test_rccl.py')} -v -s --log-cli-level=info",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
        # Architectures that we have multi GPU setup for testing
        "multi_gpu": {"linux": ["gfx94X-dcgpu"]},
    },
    # rocprofiler-sdk tests
    "rocprofiler-sdk": {
        "job_name": "rocprofiler-sdk",
        "fetch_artifact_args": "--tests",
        "timeout_minutes": 15,
        "additional_requirements_files": [
            "share/rocprofiler-sdk/tests/requirements.txt",
        ],
        "test_script": f"python {_get_script_path('test_rocprofiler_sdk.py')}",
        "platform": ["linux"],
        "container_options": "--cap-add=SYS_PTRACE",
        "total_shards_dict": {
            "linux": 1,
        },
    },
    # hipDNN tests
    "hipdnn": {
        "job_name": "hipdnn",
        "fetch_artifact_args": "--hipdnn --tests",
        "timeout_minutes": 5,
        "test_script": f"python {_get_script_path('test_hipdnn.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # hipDNN install/consumption tests
    "hipdnn_install": {
        "job_name": "hipdnn_install",
        "timeout_minutes": 10,
        "test_script": f"python {_get_script_path('test_hipdnn_install.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # hipDNN samples tests
    "hipdnn-samples": {
        "job_name": "hipdnn-samples",
        "fetch_artifact_args": "--blas --miopen --hipdnn --miopenprovider --hipdnn-samples --tests",
        "timeout_minutes": 5,
        "test_script": f"python {_get_script_path('test_hipdnn_samples.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # MIOpen provider tests
    "miopenprovider": {
        "job_name": "miopenprovider",
        "fetch_artifact_args": "--blas --miopen --hipdnn --miopenprovider --tests",
        "timeout_minutes": 20,
        "test_script": f"python {_get_script_path('test_miopenprovider.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # TODO(iree-org/fusilli/issues/57): Enable fusilli tests once build is
    # enabled by default.
    # "fusilliprovider": {
    #     "job_name": "fusilliprovider",
    #     "fetch_artifact_args": "--hipdnn --fusilliprovider --iree-compiler --tests",
    #     "timeout_minutes": 15,
    #     "test_script": f"python {_get_script_path('test_fusilliprovider.py')}",
    #     "platform": ["linux"],
    #     "total_shards": 1,
    # },
    # hipBLASLt provider tests
    "hipblasltprovider": {
        "job_name": "hipblasltprovider",
        "fetch_artifact_args": "--blas --hipdnn --hipblasltprovider --tests",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_hipblasltprovider.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # rocWMMA tests
    "rocwmma": {
        "job_name": "rocwmma",
        "fetch_artifact_args": "--rocwmma --tests --blas",
        "timeout_minutes": 60,
        "test_script": f"python {_get_script_path('test_rocwmma.py')}",
        "platform": ["linux", "windows"],
        "total_shards_dict": {
            "linux": 4,
            "windows": 2,
        },
    },
    # profiler tests
    "rocprofiler-compute": {
        "job_name": "rocprofiler-compute",
        "fetch_artifact_args": "--rocprofiler-compute --rocprofiler-sdk --tests",
        "timeout_minutes": 60,
        "additional_requirements_files": [
            "libexec/rocprofiler-compute/requirements.txt",
            "libexec/rocprofiler-compute/requirements-test.txt",
        ],
        "test_script": f"python {_get_script_path('test_rocprofiler_compute.py')} -v",
        "platform": ["linux"],
        "total_shards_dict": {"linux": 2},
    },
    "rocprofiler-systems": {
        "job_name": "rocprofiler-systems",
        "fetch_artifact_args": "--rocprofiler-systems --rocprofiler-sdk --tests",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_rocprofiler_systems.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # libhipcxx hipcc tests
    "libhipcxx_hipcc": {
        "job_name": "libhipcxx_hipcc",
        "fetch_artifact_args": "--libhipcxx --tests",
        "timeout_minutes": 30,
        "test_script": f"python {_get_script_path('test_libhipcxx_hipcc.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # libhipcxx hiprtc tests
    "libhipcxx_hiprtc": {
        "job_name": "libhipcxx_hiprtc",
        "fetch_artifact_args": "--libhipcxx --tests",
        "timeout_minutes": 20,
        "test_script": f"python {_get_script_path('test_libhipcxx_hiprtc.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    "rocdecode": {
        "job_name": "rocdecode",
        "fetch_artifact_args": "--rocdecode --tests",
        "timeout_minutes": 10,
        "test_script": f"python {_get_script_path('test_rocdecode.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
        },
        # rocdecode requires FFmpeg dev libraries (libavcodec-dev, libavformat-dev,
        # libavutil-dev) for test builds. These are not bundled in TheRock
        # artifacts and are provided via the specialized media image.
        "container_image": "ghcr.io/rocm/no_rocm_image_ubuntu24_04_media@sha256:d715ae2db664b055c90343e00588ce9ac3eec387513fe359396e5e08e75521ca",
    },
    "rocjpeg": {
        "job_name": "rocjpeg",
        "fetch_artifact_args": "--rocjpeg --tests",
        "timeout_minutes": 10,
        "test_script": f"python {_get_script_path('test_rocjpeg.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
        },
    },
    # aqlprofile tests
    "aqlprofile": {
        "job_name": "aqlprofile",
        "fetch_artifact_args": "--aqlprofile --tests",
        "timeout_minutes": 5,
        "test_script": f"python {_get_script_path('test_aqlprofile.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
    # rocrtst tests
    "rocrtst": {
        "job_name": "rocrtst",
        "fetch_artifact_args": "--rocrtst --tests",
        "timeout_minutes": 15,
        "test_script": f"python {_get_script_path('test_rocrtst.py')}",
        "platform": ["linux"],
        "total_shards_dict": {
            "linux": 1,
            "windows": 1,
        },
    },
}


def run():
    platform = os.getenv("RUNNER_OS").lower()
    projects_to_test = os.getenv("PROJECTS_TO_TEST", "*")
    amdgpu_families = os.getenv("AMDGPU_FAMILIES")
    test_type = os.getenv("TEST_TYPE", "full")
    test_labels = ast.literal_eval(os.getenv("TEST_LABELS") or "[]")
    is_benchmark_workflow = str2bool(os.getenv("IS_BENCHMARK_WORKFLOW", "false"))
    run_functional_tests = str2bool(os.getenv("RUN_FUNCTIONAL_TESTS", "false"))

    logging.info(f"Selecting projects: {projects_to_test}")

    # Determine which test matrix to use
    if is_benchmark_workflow:
        # For benchmark workflow, use ONLY benchmark_matrix
        # Benchmarks don't use test_type/test_labels (all have total_shards=1, no filtering)
        logging.info("Using benchmark_matrix only (benchmark tests)")
        selected_matrix = benchmark_matrix.copy()
    else:
        # For regular workflow, use test_matrix
        logging.info("Using test_matrix only (regular tests)")
        selected_matrix = test_matrix.copy()
        # For nightly/scheduled builds, merge functional tests into the test matrix
        if run_functional_tests and functional_matrix:
            logging.info(
                f"Merging {len(functional_matrix)} functional test(s) into test matrix"
            )
            selected_matrix.update(functional_matrix)

    # This string -> array conversion ensures no partial strings are detected during test selection (ex: "hipblas" in ["hipblaslt", "rocblas"] = false)
    project_array = [item.strip() for item in projects_to_test.split(",")]

    all_components = []
    for key in selected_matrix:
        job_name = selected_matrix[key]["job_name"]

        # If the test is disabled for a particular platform, skip the test
        if (
            "exclude_family" in selected_matrix[key]
            and platform in selected_matrix[key]["exclude_family"]
            and amdgpu_families in selected_matrix[key]["exclude_family"][platform]
        ):
            logging.info(
                f"Excluding job {job_name} for platform {platform} and family {amdgpu_families}"
            )
            continue

        # If test labels are populated, and the test job name is not in the test labels, skip the test
        # Note: Benchmarks never use test_labels (always empty list)
        if key != "sanity" and test_labels and key not in test_labels:
            logging.info(f"Excluding job {job_name} since it's not in the test labels")
            continue

        # If the test is enabled for a particular platform and a particular (or all) projects are selected.
        # Note: Sanity goes through the same all_components loop as other components, but is separated
        # into its own sanity_component GHA output after the loop (see gha_set_output below).
        if platform in selected_matrix[key]["platform"] and (
            key == "sanity" or key in project_array or "*" in project_array
        ):
            logging.info(f"Including job {job_name} with test_type {test_type}")
            job_config_data = selected_matrix[key]
            job_config_data["test_type"] = test_type
            # For CI testing, we construct a shard array based on "total_shards" from "fetch_test_configurations.py"
            # This way, the test jobs will be split up into X shards. (ex: [1, 2, 3, 4] = 4 test shards)
            # For display purposes, we add "i + 1" for the job name (ex: 1 of 4). During the actual test sharding in the test executable, this array will become 0th index
            # Note: Benchmarks always have total_shards=1 (no sharding)
            total_shards = job_config_data.get("total_shards_dict", {}).get(platform, 1)
            job_config_data["shard_arr"] = [i + 1 for i in range(total_shards)]
            job_config_data["total_shards"] = total_shards

            # If the test type is quick tests, we only need one shard for the test job
            # Note: Benchmarks always use test_type="full" but have total_shards=1 anyway
            if test_type == "quick":
                job_config_data["total_shards"] = 1
                job_config_data["shard_arr"] = [1]

            # If the test requires multi GPU testing, we use a multi-GPU test runner for this specific test
            # Inside the "multi_gpu" field, we have a mapping of amdgpu_family -> bool (if multi GPU testing is enabled for that family)
            # If the multi GPU test runner is not enabled, we will skip the test
            if "multi_gpu" in selected_matrix[key]:
                amdgpu_families_matrix = get_all_families_for_trigger_types(
                    ["presubmit", "postsubmit", "nightly"]
                )
                if (
                    platform in selected_matrix[key]["multi_gpu"]
                    and amdgpu_families in selected_matrix[key]["multi_gpu"][platform]
                ):
                    # If the architecture is available for multi GPU testing, we indicate that this specific test requires the multi GPU test runner
                    shortened_amdgpu_families_name = amdgpu_families.split("-")[
                        0
                    ].lower()
                    multi_gpu_runner = amdgpu_families_matrix[
                        shortened_amdgpu_families_name
                    ][platform]["test-runs-on-multi-gpu"]
                    logging.info(
                        f"Including job {job_name} since multi GPU testing is available for family {amdgpu_families} with runner {multi_gpu_runner}"
                    )
                    job_config_data["multi_gpu_runner"] = multi_gpu_runner
                else:
                    # If the architecture is not available for multi GPU testing, we skip the test requiring multi GPU
                    logging.info(
                        f"Excluding job {job_name} since multi GPU testing is not available for family {amdgpu_families}"
                    )
                    continue

            all_components.append(job_config_data)

    # Separate sanity (always a prerequisite) from the regular component matrix.
    sanity_component = next(
        (c for c in all_components if c.get("job_name") == "sanity"), None
    )
    output_matrix = [c for c in all_components if c.get("job_name") != "sanity"]

    gha_set_output(
        {
            "sanity_component": json.dumps(sanity_component),
            "components": json.dumps(output_matrix),
            "platform": platform,
        }
    )


if __name__ == "__main__":
    run()