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

name: Test PyTorch Wheels

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"
      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.7")
        type: string
        default: "release/2.7"
      device_query:
        description: >-
          GPU candidate-set selection (--device-query, stage 1). "unique": one
          device per architecture. "all": every device of each architecture.
        type: choice
        options:
          - "unique"
          - "all"
        default: "unique"
      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.
        type: choice
        options:
          - "single"
          - "all"
        default: "single"

  workflow_call:
    inputs:
      amdgpu_family:
        required: true
        type: string
      test_runs_on:
        required: true
        type: string
      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: "release/2.7"
      device_query:
        description: >-
          --device-query. "unique": one device per architecture; "all": every
          device of each architecture.
        type: string
        default: "unique"
      gpu_policy:
        description: >-
          --gpu-policy. "single": one GPU visible at a time; "all": all candidate
          GPUs visible at once.
        type: string
        default: "single"
      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

permissions:
  contents: read

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

jobs:
  test_wheels:
    name: Test PyTorch | ${{ inputs.amdgpu_family }}
    runs-on: ${{ inputs.test_runs_on }}
    container:
      image: ${{ contains(inputs.test_runs_on, 'linux') && 'ghcr.io/rocm/no_rocm_image_ubuntu24_04@sha256:405945a40deaff9db90b9839c0f41d4cba4a383c1a7459b28627047bf6302a26' || null }}
      options: --ipc host
        --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
    env:
      VENV_DIR: ${{ github.workspace }}/.venv
      AMDGPU_FAMILY: ${{ inputs.amdgpu_family }}

    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

      # TODO: also upload and reference test report together with this logging?
      - name: Summarize workflow inputs
        run: |
          python build_tools/github_actions/summarize_test_pytorch_workflow.py \
            --torch-version=${{ inputs.torch_version }} \
            --pytorch-git-ref=${{ inputs.pytorch_git_ref }} \
            --index-url=${{ inputs.package_index_url }} \
            --device-extras="${{ steps.expand.outputs.device_extras }}"

      - name: Set git options
        run: |
          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.
      # TODO(#5110): Accept manifest_url and use checkout_from_manifest.py
      # --no-hipify --no-submodules so tests use the exact source commit and
      # related repository pins that produced the package under test.
      - name: Checkout PyTorch Source Repos from nightly branch
        if: ${{ (inputs.pytorch_git_ref == 'nightly') }}
        run: |
          python external-builds/pytorch/pytorch_torch_repo.py checkout \
            --gitrepo-origin https://github.com/pytorch/pytorch.git \
            --repo-hashtag nightly \
            --no-hipify \
            --no-submodules

      - name: Checkout PyTorch Source Repos from stable branch
        if: ${{ (inputs.pytorch_git_ref != 'nightly') }}
        run: |
          python external-builds/pytorch/pytorch_torch_repo.py checkout \
            --gitrepo-origin https://github.com/ROCm/pytorch.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 }}" \
            --index-url=${{ inputs.package_index_url }} \
            --activate-in-future-github-actions-steps

      - 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: 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: Run PyTorch tests
        timeout-minutes: 120
        run: |
          python ./external-builds/pytorch/run_pytorch_tests.py \
            --device-query "${{ inputs.device_query }}" \
            --gpu-policy "${{ inputs.gpu_policy }}" \
            -- \
            --continue-on-collection-errors \
            --import-mode=importlib \
            -v