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

name: Test PyTorch Wheels (Full Suite)

on:
  workflow_dispatch:
    inputs:
      amdgpu_family:
        description: GPU family to test
        required: true
        type: string
        default: "gfx94X-dcgpu"
      test_runs_on:
        description: Runner label to use. The selected runner should have a GPU supported by amdgpu_family
        required: true
        type: string
        default: "linux-gfx942-1gpu-ccs-csp-ossci-rocm"
      test_runs_on_multi_gpu:
        description: Runner label for multi-GPU test configs (e.g. distributed). Corresponds to test-runs-on-multi-gpu in amdgpu_family_matrix.py
        type: string
        default: "linux-gfx942-8gpu-ossci-rocm"
      package_index_url:
        description: Python package index URL to test, typically nightly/dev URL with a "whl-multi-arch" subdir
        required: true
        type: string
        default: "https://rocm.nightlies.amd.com/whl-multi-arch/"
      python_version:
        required: true
        type: string
        default: "3.12"
      torch_version:
        description: torch package version to install. (e.g. "2.7.1+rocm7.10.0a20251120")
        required: true
        type: string
      pytorch_git_ref:
        description: PyTorch ref to checkout test sources from. (e.g. "nightly", or "release/2.9")
        type: string
        default: "nightly"
      pytorch_git_repo:
        description: "PyTorch git repo to checkout test sources from. (Options: rocm/pytorch, pytorch/pytorch)"
        type: string
        default: "pytorch/pytorch"
      test_configs:
        description: Space-separated test configs to run (e.g. "default", "default distributed", "default distributed inductor")
        type: string
        default: "default distributed inductor"
      tests_to_include:
        description: Space-separated list of test modules to run (passed as TESTS_TO_INCLUDE to run_test.py). Leave empty to run all tests.
        type: string
        default: ""
      device_query:
        description: >-
          GPU candidate-set selection (--device-query, stage 1). "unique": one
          device per architecture. "all": every device of each architecture.
          "auto": derive from the test config ("all" for distributed, "unique"
          otherwise).
        type: choice
        options:
          - "auto"
          - "unique"
          - "all"
        default: "auto"
      gpu_policy:
        description: >-
          How many candidate GPUs are made visible (--gpu-policy, stage 2).
          "single": one GPU visible at a time. "all": all candidate GPUs visible
          at once. "auto": derive from the test config ("all" for distributed,
          "single" otherwise).
        type: choice
        options:
          - "auto"
          - "single"
          - "all"
        default: "auto"
      repository:
        description: "Repository to checkout. Otherwise, defaults to `github.repository`."
        type: string
      ref:
        description: "Branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow."
        type: string

  workflow_call:
    inputs:
      amdgpu_family:
        required: true
        type: string
      test_runs_on:
        required: false
        type: string
        default: "linux-gfx942-1gpu-ccs-csp-ossci-rocm"
      test_runs_on_multi_gpu:
        type: string
        default: "linux-gfx942-8gpu-ossci-rocm"
      package_index_url:
        required: true
        type: string
      python_version:
        required: true
        type: string
      torch_version:
        required: true
        type: string
      pytorch_git_ref:
        type: string
        default: "nightly"
      pytorch_git_repo:
        type: string
        default: "pytorch/pytorch"
      repository:
        description: "Repository to checkout. Otherwise, defaults to `github.repository`."
        type: string
      ref:
        description: "Branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow."
        type: string
      test_configs:
        type: string
        default: "default distributed inductor"
      tests_to_include:
        description: Space-separated list of test modules to run (passed as TESTS_TO_INCLUDE to run_test.py). Leave empty to run all tests.
        type: string
        default: ""
      device_query:
        description: >-
          --device-query. "unique": one device per architecture; "all": every
          device of each architecture; "auto": derive from the test config.
        type: string
        default: "auto"
      gpu_policy:
        description: >-
          --gpu-policy. "single": one GPU visible at a time; "all": all candidate
          GPUs visible at once; "auto": derive from the test config.
        type: string
        default: "auto"

permissions:
  contents: read

run-name: Test PyTorch Full (${{ inputs.amdgpu_family }}, ${{ inputs.torch_version}}, ${{ inputs.test_runs_on }})

jobs:
  prepare_matrix:
    runs-on: ubuntu-24.04
    outputs:
      matrix: ${{ steps.generate.outputs.matrix }}
    steps:
      - name: Checkout
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || '' }}

      - id: generate
        run: |
          python external-builds/pytorch/generate_test_sharding_matrix.py \
            --test-configs '${{ inputs.test_configs }}' \
            --default-runner '${{ inputs.test_runs_on }}' \
            --multi-gpu-runner '${{ inputs.test_runs_on_multi_gpu }}'

  test_wheels:
    needs: prepare_matrix
    name: Test PyTorch Full | ${{ inputs.amdgpu_family }} / ${{ matrix.test_config }} (shard ${{ matrix.shard }}/${{ matrix.num_shards }})
    runs-on: ${{ matrix.runs_on }}
    strategy:
      fail-fast: false
      matrix: ${{ fromJSON(needs.prepare_matrix.outputs.matrix) }}
    container:
      image: ${{ contains(inputs.test_runs_on, 'linux') && 'ghcr.io/rocm/no_rocm_image_ubuntu24_04@sha256:405945a40deaff9db90b9839c0f41d4cba4a383c1a7459b28627047bf6302a26' || null }}
      # SYS_PTRACE is required by PyTorch's test infra which uses gdb to
      # collect stack traces on segfaults and hangs (torch/testing/_internal).
      options: --shm-size=10g
        --cap-add=SYS_PTRACE
        --group-add video
        --device /dev/kfd
        --device /dev/dri
        --group-add 992
        --group-add 110
        --env-file /etc/podinfo/gha-gpu-isolation-settings
        --user 0:0 # Running as root, by recommendation of GitHub: https://docs.github.com/en/actions/reference/workflows-and-actions/dockerfile-support#user
    defaults:
      run:
        shell: bash
    # These env vars are set at the job level so they are available to all steps
    # (e.g. export_test_times.py, run_pytorch_tests_full.py). The Python script
    # uses os.environ.setdefault() so values set here take precedence.
    env:
      VENV_DIR: ${{ github.workspace }}/.venv
      AMDGPU_FAMILY: ${{ inputs.amdgpu_family }}
      SHARD_NUMBER: ${{ matrix.shard }}
      NUM_TEST_SHARDS: ${{ matrix.num_shards }}
      CI: "1"
      PYTORCH_TEST_WITH_ROCM: "1"
      PYTORCH_TESTING_DEVICE_ONLY_FOR: "cuda"
      # Pin OpenBLAS threads to 1 for deterministic reductions; without this,
      # OpenBLAS defaults to the number of CPU cores, which causes
      # non-deterministic results in tests like test_entropy_monte_carlo.
      OPENBLAS_NUM_THREADS: "1"
      PYTORCH_NUM_PYTEST_RERUNS: "2"
      TEST_CONFIG: ${{ matrix.test_config }}
      TESTS_TO_INCLUDE: ${{ inputs.tests_to_include }}

    steps:
      - name: Checkout
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || '' }}

      - name: Set up Python
        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
        with:
          python-version: ${{ inputs.python_version }}

      # Multi-arch packages use pip extras (e.g. torch[device-gfx942]) to
      # select GPU-specific device packages. Expand the family to get the
      # comma-separated device extras string for pip install.
      - name: Expand GPU family to device extras
        id: expand
        run: |
          python build_tools/github_actions/expand_amdgpu_families.py \
            --amdgpu-families "${{ inputs.amdgpu_family }}" \
            --output-mode=device-extras

      # safe.directory must be set before Runner Health Status
      - name: Adjust git config
        run: |
          git config --global --add safe.directory $PWD
          git config fetch.parallel 10
          git config --global core.longpaths true

      # Here we checkout the same version of PyTorch that wheels were built from
      # so we have the right set of test source files. We _probably_ don't need
      # to run HIPIFY, apply patches, or fetch submodules, so we skip those
      # steps to save time.
      - name: Checkout PyTorch Source
        run: |
          python external-builds/pytorch/pytorch_torch_repo.py checkout \
            --gitrepo-origin https://github.com/${{ inputs.pytorch_git_repo }}.git \
            --repo-hashtag ${{ inputs.pytorch_git_ref }} \
            --no-hipify \
            --no-submodules

      - name: Set up virtual environment
        timeout-minutes: 15
        run: |
          python build_tools/setup_venv.py ${VENV_DIR} \
            --packages "torch[${{ steps.expand.outputs.device_extras }}]==${{ inputs.torch_version }} rocm[devel,${{ steps.expand.outputs.device_extras }}]" \
            --index-url=${{ inputs.package_index_url }} \
            --activate-in-future-github-actions-steps

      - name: Initialize ROCm SDK and configure environment
        run: |
          rocm-sdk init

          ROCM_HOME="$(rocm-sdk path --root)"
          ROCM_BIN="$(rocm-sdk path --bin)"
          ROCM_CMAKE_PREFIX="$(rocm-sdk path --cmake)"

          echo "ROCM_HOME=${ROCM_HOME}"
          echo "ROCM_BIN=${ROCM_BIN}"
          echo "ROCM_CMAKE_PREFIX=${ROCM_CMAKE_PREFIX}"

          ROCM_SYSDEPS="${ROCM_HOME}/lib/rocm_sysdeps"
          ROCM_SYSDEPS_INCLUDE="${ROCM_SYSDEPS}/include"
          ROCM_SYSDEPS_PKGCONFIG="${ROCM_SYSDEPS}/lib/pkgconfig"

          # Add ROCm bin to PATH for all subsequent steps
          echo "${ROCM_BIN}" >> "${GITHUB_PATH}"

          # Persist ROCm env vars across workflow steps (mirrors install_rocm.sh)
          echo "ROCM_PATH=${ROCM_HOME}" >> "${GITHUB_ENV}"
          echo "ROCM_HOME=${ROCM_HOME}" >> "${GITHUB_ENV}"
          echo "ROCM_SOURCE_DIR=${ROCM_HOME}" >> "${GITHUB_ENV}"
          echo "ROCM_BIN=${ROCM_BIN}" >> "${GITHUB_ENV}"
          echo "ROCM_CMAKE=${ROCM_CMAKE_PREFIX}" >> "${GITHUB_ENV}"
          echo "CMAKE_PREFIX_PATH=${ROCM_CMAKE_PREFIX}:${CMAKE_PREFIX_PATH:-}" >> "${GITHUB_ENV}"
          echo "HIP_DEVICE_LIB_PATH=${ROCM_HOME}/lib/llvm/amdgcn/bitcode" >> "${GITHUB_ENV}"
          echo "ROCM_DEVICE_LIB_PATH=${ROCM_HOME}/lib/llvm/amdgcn/bitcode" >> "${GITHUB_ENV}"
          echo "ROCM_SYSDEPS_INCLUDE=${ROCM_SYSDEPS_INCLUDE}" >> "${GITHUB_ENV}"
          echo "CPLUS_INCLUDE_PATH=${ROCM_SYSDEPS_INCLUDE}:${CPLUS_INCLUDE_PATH:-}" >> "${GITHUB_ENV}"
          echo "C_INCLUDE_PATH=${ROCM_SYSDEPS_INCLUDE}:${C_INCLUDE_PATH:-}" >> "${GITHUB_ENV}"
          echo "PKG_CONFIG_PATH=${ROCM_SYSDEPS_PKGCONFIG}:${PKG_CONFIG_PATH:-}" >> "${GITHUB_ENV}"
          echo "LD_LIBRARY_PATH=${ROCM_HOME}/lib/host-math/lib:${ROCM_SYSDEPS}/lib:${LD_LIBRARY_PATH:-}" >> "${GITHUB_ENV}"
          echo "LIBRARY_PATH=${ROCM_HOME}/lib/host-math/lib:${ROCM_SYSDEPS}/lib:${LIBRARY_PATH:-}" >> "${GITHUB_ENV}"
          echo "USE_MSLK=0" >> "${GITHUB_ENV}"

          # Use ROCm's clang for JIT C++ compilation (torch.utils.cpp_extension)
          # instead of requiring gcc/build-essential.
          echo "CC=${ROCM_HOME}/lib/llvm/bin/clang" >> "${GITHUB_ENV}"
          echo "CXX=${ROCM_HOME}/lib/llvm/bin/clang++" >> "${GITHUB_ENV}"

      # export_test_times.py must run from the PyTorch repo root with
      # PYTHONPATH=. because it imports internal PyTorch modules (e.g.
      # tools.stats.*) that expect the repo root on sys.path.
      - name: Export test times for sharding
        run: |
          pushd external-builds/pytorch/pytorch
          PYTHONPATH=. python tools/stats/export_test_times.py
          popd

      - name: Install test requirements
        run: |
          python -m pip install -r external-builds/pytorch/requirements-test.txt
          python -m pip install -r external-builds/pytorch/pytorch/.ci/docker/requirements-ci.txt
          pip freeze

      - name: Runner health status
        run: |
          ./build_tools/health_status.py

      - name: Run rocm-sdk sanity tests
        timeout-minutes: 5
        run: |
          rocm-sdk test

      - name: Run PyTorch smoketests
        timeout-minutes: 10
        run: |
          python ./external-builds/pytorch/run_pytorch_smoke_tests.py -- \
            --log-cli-level=INFO \
            -v

      - name: (Linux) Run PyTorch full test suite
        if: ${{ runner.os == 'Linux' }}
        run: |
          python ./external-builds/pytorch/run_pytorch_tests_full.py \
            --device-query "${{ inputs.device_query }}" \
            --gpu-policy "${{ inputs.gpu_policy }}" \
            -- \
            --continue-on-collection-errors

      - name: Upload test reports
        if: always()
        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
        with:
          name: test-reports-${{ matrix.test_config }}-${{ matrix.shard }}-${{ matrix.num_shards }}
          path: external-builds/pytorch/pytorch/test/test-reports/
          if-no-files-found: ignore
          retention-days: 30

      # HACK: The test process does not terminate on its own gracefully,
      # so we write to run_pytorch_tests_exit_code.txt and then kill the process.
      # After killing the process we read the return code to signal it normally.
      # See https://github.com/ROCm/TheRock/issues/999.
      - name: (Windows) Run PyTorch tests
        if: ${{ runner.os == 'Windows'}}
        continue-on-error: true
        run: |
          python ./external-builds/pytorch/run_pytorch_tests.py -- \
            --continue-on-collection-errors \
            -v

      - name: (Windows) Read and propagate exit code
        if: ${{ runner.os == 'Windows'}}
        run: |
          if [ -f run_pytorch_tests_exit_code.txt ]; then
            EXIT_CODE=$(cat run_pytorch_tests_exit_code.txt)
            echo "Exit code from file: ${EXIT_CODE}"
            exit ${EXIT_CODE}
          else
            echo "No run_pytorch_tests_exit_code.txt found"
            exit 1
          fi

  # Aggregate every shard's JUnit XML test reports into a single failure table
  # on the run's summary page, so all failures across configs/shards are visible
  # in one place without digging into individual matrix jobs or artifacts.
  test_summary:
    if: always()
    needs: test_wheels
    name: Test Summary
    runs-on: ubuntu-24.04
    permissions:
      contents: read
    steps:
      - name: Checkout
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || '' }}
          sparse-checkout: external-builds/pytorch/summarize_test_results.py

      - name: Download all test report artifacts
        uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
        with:
          path: artifacts
          pattern: test-reports-*

      - name: Generate combined summary
        env:
          AMDGPU_FAMILY: ${{ inputs.amdgpu_family }}
          TORCH_VERSION: ${{ inputs.torch_version }}
          PYTHON_VERSION: ${{ inputs.python_version }}
          PACKAGE_INDEX_URL: ${{ inputs.package_index_url }}
          PYTORCH_GIT_REF: ${{ inputs.pytorch_git_ref }}
        run: |
          python external-builds/pytorch/summarize_test_results.py \
            --combined-dir artifacts