TheRock/external-builds/pytorch/run_pytorch_tests.py

#!/usr/bin/env python3
# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""PyTorch ROCm Pytest Runner with additional test exclusion capabilities.

This script runs PyTorch unit tests using pytest with additional test exclusion
capabilities tailored for AMD ROCm GPUs.

Test Exclusion Criteria
------------------------
Tests may be skipped based on:
- AMDGPU family compatibility (e.g., gfx942, gfx1151)
- PyTorch version-specific issues
- Platform (Linux, Windows)
- Known failures not yet upstreamed to PyTorch

Environment Variables
---------------------
THEROCK_ROOT_DIR : str, optional
                   Root directory of TheRock project.
                   If not set, auto-detected from script location.
AMDGPU_FAMILY :     str, optional
                    Target AMDGPU family for testing (e.g., "gfx942", "gfx94X").
                    Names should match those in "TheRock/cmake/therock_amdgpu_targets.cmake".
                    Supports wildcards (e.g., "gfx94X" matches any gfx94* architecture).
                    If not set, auto-detects from available hardware using PyTorch.
PYTORCH_VERSION :   str, optional
                    PyTorch version for version-specific test filtering (e.g., "2.10").
                    Format: "major.minor" as string.
                    If not set, auto-detects from installed PyTorch package.
HIP_VISIBLE_DEVICES : str, optional (read/write)
                      If already set, the script respects this constraint and only selects
                      from the GPUs visible within this limitation (e.g., in containers).
                      The script will further filter and update this variable based on
                      the AMDGPU_FAMILY selection or auto-detection.

Usage Examples
--------------
Basic usage (auto-detect everything):
    $ python run_pytorch_tests.py

Debug mode (run only skipped tests):
    $ python run_pytorch_tests.py --debug

Custom test selection with pytest -k:
    $ python run_pytorch_tests.py -k "test_nn and not test_dropout"

Pass additional pytest arguments after "--":
    $ python run_pytorch_tests.py -- -m "slow"
    $ python run_pytorch_tests.py -- --tb=short -x

GPU selection options:
    $ python run_pytorch_tests.py --gpu-policy all --device-query all
    $ python run_pytorch_tests.py --gpu-policy single --device-query all

Exit Codes
----------
0 : All tests passed
1 : Test failures or collection errors
? : Other exit codes from pytest
Other : Pytest-specific error codes

Side-effects
------------
- This script modifies PYTHONPATH and sys.path to include PyTorch test directory
- Creates a temporary MIOpen cache directory for each run
- Sets HIP_VISIBLE_DEVICES environment variable to select specific GPU(s) for testing
- Runs tests sequentially (--numprocesses=0) by default
"""

import argparse
import os
import platform
import sys
import tempfile

from skip_tests.create_skip_tests import *
from pathlib import Path

import pytest

from pytorch_utils import (
    check_pytorch_source_version,
    configure_gpu_visibility,
    detect_pytorch_version,
    reconcile_agent_visibility_env,
)

THIS_SCRIPT_DIR = Path(__file__).resolve().parent


def setup_env(pytorch_dir: str) -> None:
    """Set up environment variables required for PyTorch testing with ROCm.

    Args:
        pytorch_dir: Path to the PyTorch directory containing test files.

    Side effects:
        - Sets multiple environment variables for PyTorch testing
        - Creates a temporary directory for MIOpen cache
        - Modifies sys.path to include the test directory
    """
    reconcile_agent_visibility_env()

    os.environ["PYTORCH_PRINT_REPRO_ON_FAILURE"] = "0"
    os.environ["PYTORCH_TEST_WITH_ROCM"] = "1"
    os.environ["MIOPEN_CUSTOM_CACHE_DIR"] = tempfile.mkdtemp()
    os.environ["PYTORCH_TESTING_DEVICE_ONLY_FOR"] = "cuda"

    old_pythonpath = os.getenv("PYTHONPATH", "")
    test_dir = f"{pytorch_dir}/test"

    if old_pythonpath:
        os.environ["PYTHONPATH"] = f"{test_dir}:{old_pythonpath}"
    else:
        os.environ["PYTHONPATH"] = test_dir

    # Force update the PYTHONPATH to be part of the sys path
    # Otherwise our current python process that will run pytest will NOT
    # find it and pytest will crash!
    if test_dir not in sys.path:
        sys.path.insert(0, test_dir)


def cmd_arguments(argv: list[str]) -> tuple[argparse.Namespace, list[str]]:
    """Parse command line arguments.

    Args:
        argv: Command line arguments (without program name).

    Returns:
        Tuple of (parsed args, passthrough pytest args passed after "--").
    """
    # Extract passthrough pytest args after "--"
    try:
        rest_pos = argv.index("--")
    except ValueError:
        passthrough_pytest_args = []
    else:
        passthrough_pytest_args = argv[rest_pos + 1 :]
        argv = argv[:rest_pos]

    parser = argparse.ArgumentParser(
        description="""
Runs PyTorch pytest for AMD GPUs. Skips additional tests compared to upstream.
Additional tests to be skipped can be tuned by PyTorch version and amdgpu family.
All arguments after "--" are passed directly to pytest.
"""
    )

    parser.add_argument(
        "--amdgpu-family",
        type=str,
        default=os.getenv("AMDGPU_FAMILY", ""),
        required=False,
        help="""Amdgpu family (e.g. "gfx942").
Select (potentially) additional tests to be skipped based on the amdgpu family""",
    )

    pytorch_version = os.getenv("PYTORCH_VERSION")
    parser.add_argument(
        "--pytorch-version",
        type=str,
        default=pytorch_version if pytorch_version is not None else "",
        required=False,
        help="""Pytorch version (e.g. "2.7" or "all").
Select (potentially) additional tests to be skipped based on the Pytorch version.
'All' is also possible. Then all skip tests for all pytorch versions are included.
If no PyTorch version is given, it is auto-determined by the PyTorch used to run pytest.""",
    )

    default_pytorch_dir = THIS_SCRIPT_DIR / "pytorch"
    parser.add_argument(
        "--pytorch-dir",
        type=Path,
        default=default_pytorch_dir,
        help="""Path for the pytorch repository, where tests will be sourced from
By default the pytorch directory is determined based on this script's location
""",
    )

    parser.add_argument(
        "--debug",
        default=False,
        required=False,
        action=argparse.BooleanOptionalAction,
        help="""Inverts the selection. Only runs skipped tests.""",
    )

    parser.add_argument(
        "-k",
        default="",
        required=False,
        help="""Overwrites the pytest -k option that decides which tests should be run or skipped""",
    )

    parser.add_argument(
        "--cache",
        default=True,
        required=False,
        action=argparse.BooleanOptionalAction,
        help="""Enable pytest caching (default). Use --no-cache when only having read-only access to pytorch directory""",
    )

    # GPU selection happens in two stages:
    #   1. --device-query  decides which GPUs enter the candidate set.
    #   2. --gpu-policy    decides how many candidates are made visible to tests.
    parser.add_argument(
        "--device-query",
        type=str,
        choices=["unique", "all"],
        default="unique",
        help="""Stage 1: which GPUs enter the candidate set (see --gpu-policy for stage 2).
- "unique": one device per architecture (default). E.g. {gfx942:[0], gfx1100:[2]}.
- "all": every device of each architecture. E.g. {gfx942:[0,1], gfx1100:[2]}.""",
    )

    parser.add_argument(
        "--gpu-policy",
        type=str,
        choices=["single", "all"],
        default="single",
        help="""Stage 2: how many candidate GPUs to make visible (see --device-query for stage 1).
- "single": one GPU visible at a time (default). Suitable for most unit tests.
- "all": all candidate GPUs visible at once. Useful for multi-GPU tests.""",
    )

    parser.add_argument(
        "--allow-version-mismatch",
        default=False,
        required=False,
        action=argparse.BooleanOptionalAction,
        help="""Allows version mismatches between pytorch test sources and installed packages. Defaults to False, so mismatched versions block running tests""",
    )

    args = parser.parse_args(argv)

    if not args.pytorch_dir.exists():
        parser.error(
            f"Directory at '{args.pytorch_dir}' does not exist, checkout pytorch and then set the path via --pytorch-dir or check it out in TheRock/external-build/pytorch/<your pytorch directory>"
        )

    return args, passthrough_pytest_args


def main() -> int:
    """Main entry point for the PyTorch test runner.

    Returns:
        Exit code from pytest (0 for success, non-zero for failures).
    """
    try:
        args, passthrough_pytest_args = cmd_arguments(sys.argv[1:])

        pytorch_dir = args.pytorch_dir
        check_pytorch_source_version(
            pytorch_dir=pytorch_dir, allow_mismatch=args.allow_version_mismatch
        )

        setup_env(pytorch_dir)

        # CRITICAL: Determine AMDGPU family and set HIP_VISIBLE_DEVICES
        # BEFORE importing torch/running pytest. Once torch.cuda is initialized,
        # changing HIP_VISIBLE_DEVICES has no effect.
        selected_archs = configure_gpu_visibility(
            args.amdgpu_family, args.device_query, args.gpu_policy
        )

        # Determine PyTorch version
        pytorch_version = args.pytorch_version
        if not pytorch_version:
            pytorch_version = detect_pytorch_version()
        print(f"Using PyTorch version: {pytorch_version}")

        # Get tests to skip
        tests_to_skip = get_tests(
            amdgpu_family=selected_archs,
            pytorch_version=pytorch_version,
            platform=platform.system(),
            create_skip_list=not args.debug,
        )

        # Allow manual override of test selection
        if args.k:
            tests_to_skip = args.k

        pytest_args = [
            f"{pytorch_dir}/test/test_nn.py",
            f"{pytorch_dir}/test/test_torch.py",
            f"{pytorch_dir}/test/test_cuda.py",
            f"{pytorch_dir}/test/test_unary_ufuncs.py",
            f"{pytorch_dir}/test/test_binary_ufuncs.py",
            f"{pytorch_dir}/test/test_autograd.py",
            f"-k={tests_to_skip}",
            # "-n 0",  # TODO does this need rework? why should we not run this multithreaded? this does not seem to exist?
            # -n numprocesses, --numprocesses=numprocesses
            #         Shortcut for '--dist=load --tx=NUM*popen'.
            #         With 'logical', attempt to detect logical CPU count (requires psutil, falls back to 'auto').
            #         With 'auto', attempt to detect physical CPU count. If physical CPU count cannot be determined, falls back to 1.
            #         Forced to 0 (disabled) when used with --pdb.
        ]

        if not args.cache:
            pytest_args += [
                "-p",
                "no:cacheprovider",  # Disable caching: useful when running in a container
            ]
        # Append any passthrough pytest args passed after "--"
        pytest_args.extend(passthrough_pytest_args)

        retcode = pytest.main(pytest_args)
        print(f"Pytest finished with return code: {retcode}")
        return retcode
    except (ValueError, IndexError) as e:
        print(f"[ERROR] Exception in PyTorch unit-tests runner: {e}")
        sys.exit(1)


if __name__ == "__main__":
    sys.exit(main())