build_portable_linux_pytorch_wheels.yml

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

name: Build Portable Linux PyTorch Wheels

on:
  workflow_call:
    inputs:
      amdgpu_family:
        required: true
        type: string
      python_version:
        required: true
        type: string
      release_type:
        description: The type of release to build ("dev", "nightly", or "prerelease"). All developer-triggered jobs should use "dev"!
        required: true
        type: string
      s3_subdir:
        description: S3 subdirectory, not including the GPU-family
        required: true
        type: string
      s3_staging_subdir:
        description: S3 staging subdirectory, not including the GPU-family
        required: true
        type: string
      cloudfront_url:
        description: CloudFront URL pointing to Python index
        required: true
        type: string
      cloudfront_staging_url:
        description: CloudFront base URL pointing to staging Python index
        required: true
        type: string
      rocm_version:
        description: ROCm version to pip install (e.g. "7.10.0a20251124")
        type: string
      pytorch_git_ref:
        description: PyTorch ref to checkout. (typically "nightly", or "release/X.Y")
        required: true
        type: string
      cache_type:
        description: "Compiler cache type (sccache, ccache, or none)"
        required: false
        type: string
        default: "none"
      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_dispatch:
    inputs:
      amdgpu_family:
        type: choice
        options:
          - gfx101X-dgpu
          - gfx103X-all
          - gfx110X-all
          - gfx1150
          - gfx1151
          - gfx1152
          - gfx1153
          - gfx120X-all
          - gfx900
          - gfx906
          - gfx908
          - gfx90a
          - gfx94X-dcgpu
          - gfx950-dcgpu
        default: gfx94X-dcgpu
      python_version:
        required: true
        type: string
        default: "3.12"
      release_type:
        description: The type of release to build ("dev", "nightly", or "prerelease"). All developer-triggered jobs should use "dev"!
        type: string
        default: "dev"
      s3_subdir:
        description: S3 subdirectory, not including the GPU-family
        type: string
        default: "v2"
      s3_staging_subdir:
        description: S3 staging subdirectory, not including the GPU-family
        type: string
        default: "v2-staging"
      cloudfront_url:
        description: CloudFront base URL pointing to Python index
        type: string
        default: "https://rocm.devreleases.amd.com/v2"
      cloudfront_staging_url:
        description: CloudFront base URL pointing to staging Python index
        type: string
        default: "https://rocm.devreleases.amd.com/v2-staging"
      rocm_version:
        description: ROCm version to pip install (e.g. "7.10.0a20251124")
        type: string
      pytorch_git_ref:
        description: PyTorch ref to checkout. (typically "nightly", or "release/X.Y")
        required: true
        type: string
        default: "release/2.8"
      cache_type:
        description: "Compiler cache type"
        type: choice
        options:
          - none
          - sccache
          - ccache
        default: none

run-name: Build portable Linux PyTorch Wheels (${{ inputs.amdgpu_family }}, ${{ inputs.python_version }}, ${{ inputs.release_type }})

jobs:
  build_pytorch_wheels:
    name: Build | ${{ inputs.amdgpu_family }} | py ${{ inputs.python_version }} | torch ${{ inputs.pytorch_git_ref }}
    runs-on: ${{ github.repository_owner == 'ROCm' && 'azure-linux-scale-rocm' || 'ubuntu-24.04' }}
    permissions:
      id-token: write
      contents: read
    container:
      image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:702a5133851e6d1daf1207d2c9fbb01c2667914a5b6dc5a01faeb3ce66ea6421
    env:
      OUTPUT_DIR: ${{ github.workspace }}/output
      PACKAGE_DIST_DIR: ${{ github.workspace }}/output/packages/dist
      RELEASE_TYPE: ${{ inputs.release_type }}
      S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
      optional_build_prod_arguments: ""
      SCCACHE_BUCKET: "therock-${{ inputs.release_type }}-pytorch-sccache"
      SCCACHE_REGION: us-east-2
      SCCACHE_S3_KEY_PREFIX: "linux/${{ inputs.amdgpu_family }}/"
      SCCACHE_S3_USE_SSL: "true"
      SCCACHE_S3_SERVER_SIDE_ENCRYPTION: "true"
      SCCACHE_IDLE_TIMEOUT: "600"
      SCCACHE_LOG: warn
    outputs:
      cp_version: ${{ env.cp_version }}
      # The following are python package versions produced by the build. The
      # exact versions will depend on workflow inputs and the underlying code.
      # For example:
      #   Inputs
      #     rocm_version       : 7.10.0a20251120
      #     pytorch_git_ref    : release/2.9
      #   Outputs
      #     torch_version      : 2.9.1+rocm7.10.0a20251120
      #     torchaudio_version : 2.9.0+rocm7.10.0a20251120
      #     torchvision_version: 0.24.0+rocm7.10.0a20251120
      #     triton_version     : 3.5.1+rocm7.10.0a20251120
      # Future jobs can use these version outputs to identify newly built
      # packages, for example via `pip install torch==${TORCH_VERSION}`.
      torch_version: ${{ steps.build-pytorch-wheels.outputs.torch_version }}
      torchaudio_version: ${{ steps.build-pytorch-wheels.outputs.torchaudio_version }}
      apex_version: ${{ steps.build-pytorch-wheels.outputs.apex_version }}
      torchvision_version: ${{ steps.build-pytorch-wheels.outputs.torchvision_version }}
      triton_version: ${{ steps.build-pytorch-wheels.outputs.triton_version }}
    steps:
      - name: Checkout
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || '' }}

      - name: Configure Git Identity
        run: |
          git config --global user.name "therockbot"
          git config --global user.email "therockbot@amd.com"

      - name: Select Python version
        run: |
          python build_tools/github_actions/python_to_cp_version.py \
            --python-version ${{ inputs.python_version }}

      - name: Add selected Python version to PATH
        run: |
          python_dir="/opt/python/${{ env.cp_version }}"
          if ! [ -x "${python_dir}/bin/python" ]; then
            echo "ERROR: Could not find python: ${python_dir}"
            exit 1
          fi
          echo "${python_dir}/bin" >> "$GITHUB_PATH"

      - name: Install python deps for CI
        run: |
          pip install -r external-builds/pytorch/requirements-ci.txt

      # Checkout nightly sources from https://github.com/pytorch/pytorch
      - name: Checkout PyTorch Source Repos from nightly branch
        if: ${{ inputs.pytorch_git_ref == 'nightly' }}
        run: |
          ./external-builds/pytorch/pytorch_torch_repo.py checkout --repo-hashtag nightly
          ./external-builds/pytorch/pytorch_audio_repo.py checkout --repo-hashtag nightly
          ./external-builds/pytorch/pytorch_apex_repo.py checkout --repo-hashtag master
          ./external-builds/pytorch/pytorch_vision_repo.py checkout --repo-hashtag nightly
          ./external-builds/pytorch/pytorch_triton_repo.py checkout

      # Checkout stable sources from https://github.com/ROCm/pytorch
      - name: Checkout PyTorch Source Repos from stable branch
        if: ${{ inputs.pytorch_git_ref != 'nightly' }}
        run: |
          ./external-builds/pytorch/pytorch_torch_repo.py checkout --gitrepo-origin https://github.com/ROCm/pytorch.git --repo-hashtag ${{ inputs.pytorch_git_ref }}
          ./external-builds/pytorch/pytorch_audio_repo.py checkout --require-related-commit
          ./external-builds/pytorch/pytorch_apex_repo.py checkout --require-related-commit
          ./external-builds/pytorch/pytorch_vision_repo.py checkout --require-related-commit
          ./external-builds/pytorch/pytorch_triton_repo.py checkout

      # Note: determine_version.py sets optional_build_prod_arguments in
      # GITHUB_ENV, which includes --rocm-sdk-version and --version-suffix.
      - name: Determine optional arguments passed to `build_prod_wheels.py`
        if: ${{ inputs.rocm_version }}
        run: |
          python build_tools/github_actions/determine_version.py \
            --rocm-version ${{ inputs.rocm_version }}

      - name: Configure AWS Credentials for sccache
        if: ${{ inputs.cache_type == 'sccache' && github.repository_owner == 'ROCm' }}
        uses: aws-actions/configure-aws-credentials@ec61189d14ec14c8efccab744f656cffd0e33f37 # v6.1.0
        with:
          aws-region: us-east-2
          role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}

      - name: Verify sccache
        if: ${{ inputs.cache_type == 'sccache' }}
        run: |
          echo "sccache version: $(sccache --version)"
          echo "SCCACHE_BUCKET=${SCCACHE_BUCKET}"
          echo "SCCACHE_S3_KEY_PREFIX=${SCCACHE_S3_KEY_PREFIX}"

      - name: Build PyTorch Wheels
        id: build-pytorch-wheels
        run: |
          cache_flag=""
          if [ "${{ inputs.cache_type }}" = "sccache" ]; then
            cache_flag="--use-sccache"
          elif [ "${{ inputs.cache_type }}" = "ccache" ]; then
            cache_flag="--use-ccache"
          fi

          echo "Building PyTorch wheels for ${{ inputs.amdgpu_family }} (cache: ${{ inputs.cache_type }})"
          ./external-builds/pytorch/build_prod_wheels.py \
            build \
            --install-rocm \
            --index-url "${{ inputs.cloudfront_url }}/${{ inputs.amdgpu_family }}/" \
            --clean \
            --output-dir ${{ env.PACKAGE_DIST_DIR }} \
            ${cache_flag} \
            ${{ env.optional_build_prod_arguments }}
          python ./build_tools/github_actions/write_torch_versions.py --dist-dir ${{ env.PACKAGE_DIST_DIR }}

      - name: Report cache stats
        if: ${{ !cancelled() && inputs.cache_type != 'none' }}
        run: |
          if [ "${{ inputs.cache_type }}" = "sccache" ]; then
            echo "sccache stats:"
            echo "--------------"
            sccache --show-stats || true
          elif [ "${{ inputs.cache_type }}" = "ccache" ]; then
            echo "ccache stats:"
            echo "-------------"
            ccache -s -v || true
          fi

      - name: Sanity Check Wheel
        run: |
          python external-builds/pytorch/sanity_check_wheel.py ${{ env.PACKAGE_DIST_DIR }}/

      - name: Generate PyTorch manifest
        run: |
          python build_tools/github_actions/generate_pytorch_manifest.py \
            --manifest-dir "${{ env.PACKAGE_DIST_DIR }}/manifests" \
            --python-version "${{ inputs.python_version }}" \
            --pytorch-git-ref "${{ inputs.pytorch_git_ref }}" \
            --pytorch-dir "external-builds/pytorch/pytorch" \
            --pytorch-audio-dir "external-builds/pytorch/pytorch_audio" \
            --pytorch-vision-dir "external-builds/pytorch/pytorch_vision" \
            --triton-dir "external-builds/pytorch/triton" \
            --apex-dir "external-builds/pytorch/apex"

      - name: Configure AWS Credentials
        if: always()
        uses: aws-actions/configure-aws-credentials@ec61189d14ec14c8efccab744f656cffd0e33f37 # v6.1.0
        with:
          aws-region: us-east-2
          role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}

      - name: Upload wheels to S3 staging
        if: ${{ github.repository_owner == 'ROCm' }}
        run: |
          aws s3 cp ${{ env.PACKAGE_DIST_DIR }}/ s3://${{ env.S3_BUCKET_PY }}/${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}/ \
            --recursive --exclude "*" --include "*.whl"

      - name: Upload PyTorch manifest to S3 artifacts
        if: ${{ github.repository_owner == 'ROCm' }}
        run: |
          python build_tools/github_actions/upload_pytorch_manifest.py \
            --dist-dir "${{ env.PACKAGE_DIST_DIR }}" \
            --run-id "${{ github.run_id }}" \
            --amdgpu-family "${{ inputs.amdgpu_family }}" \
            --python-version "${{ inputs.python_version }}" \
            --pytorch-git-ref "${{ inputs.pytorch_git_ref }}"

      - name: (Re-)Generate Python package release index for staging
        if: ${{ github.repository_owner == 'ROCm' }}
        env:
          # Environment variables to be set for `manage.py`
          CUSTOM_PREFIX: "${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}"
        run: |
          python ./build_tools/third_party/s3_management/manage.py ${{ env.CUSTOM_PREFIX }}

  generate_target_to_run:
    name: Generate target_to_run
    runs-on: ubuntu-24.04
    permissions:
      contents: read
    outputs:
      test_runs_on: ${{ steps.configure.outputs.test-runs-on }}
      bypass_tests_for_releases: ${{ steps.configure.outputs.bypass_tests_for_releases }}
    steps:
      - name: Checking out repository
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || '' }}

      - name: Generating target to run
        id: configure
        env:
          TARGET: ${{ inputs.amdgpu_family }}
          PLATFORM: "linux"
        run: python ./build_tools/github_actions/configure_target_run.py

  test_pytorch_wheels:
    name: Test | ${{ inputs.amdgpu_family }} | ${{ needs.generate_target_to_run.outputs.test_runs_on }}
    if: ${{ needs.generate_target_to_run.outputs.test_runs_on != '' }}
    needs: [build_pytorch_wheels, generate_target_to_run]
    permissions:
      contents: read
    uses: ./.github/workflows/test_pytorch_wheels.yml
    with:
      amdgpu_family: ${{ inputs.amdgpu_family }}
      test_runs_on: ${{ needs.generate_target_to_run.outputs.test_runs_on }}
      package_index_url: ${{ inputs.cloudfront_staging_url }}
      python_version: ${{ inputs.python_version }}
      torch_version: ${{ needs.build_pytorch_wheels.outputs.torch_version }}
      pytorch_git_ref: ${{ inputs.pytorch_git_ref }}
      repository: ${{ inputs.repository || github.repository }}
      ref: ${{ inputs.ref || '' }}

  upload_pytorch_wheels:
    name: Release PyTorch Wheels to S3
    needs: [build_pytorch_wheels, generate_target_to_run, test_pytorch_wheels]
    if: ${{ !cancelled() }}
    runs-on: ubuntu-24.04
    permissions:
      id-token: write
      contents: read
    env:
      S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
      CP_VERSION: "${{ needs.build_pytorch_wheels.outputs.cp_version }}"
      APEX_VERSION: "${{ needs.build_pytorch_wheels.outputs.apex_version }}"
      TORCH_VERSION: "${{ needs.build_pytorch_wheels.outputs.torch_version }}"
      TORCHAUDIO_VERSION: "${{ needs.build_pytorch_wheels.outputs.torchaudio_version }}"
      TORCHVISION_VERSION: "${{ needs.build_pytorch_wheels.outputs.torchvision_version }}"
      TRITON_VERSION: "${{ needs.build_pytorch_wheels.outputs.triton_version }}"

    steps:
      - name: Checkout
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref || '' }}

      - name: Configure AWS Credentials
        if: always()
        uses: aws-actions/configure-aws-credentials@ec61189d14ec14c8efccab744f656cffd0e33f37 # v6.1.0
        with:
          aws-region: us-east-2
          role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}

      - name: Determine upload flag
        env:
          BUILD_RESULT: ${{ needs.build_pytorch_wheels.result }}
          TEST_RESULT: ${{ needs.test_pytorch_wheels.result }}
          TEST_RUNS_ON: ${{ needs.generate_target_to_run.outputs.test_runs_on }}
          BYPASS_TESTS_FOR_RELEASES: ${{ needs.generate_target_to_run.outputs.bypass_tests_for_releases }}
        run: python ./build_tools/github_actions/promote_wheels_based_on_policy.py

      - name: Copy PyTorch wheels from staging to release S3
        if: ${{ env.upload == 'true' }}
        run: |
          echo "Copying exact tested wheels to release S3 bucket..."
          aws s3 cp \
            s3://${S3_BUCKET_PY}/${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}/ \
            s3://${S3_BUCKET_PY}/${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}/ \
            --recursive \
            --exclude "*" \
            --include "torch-${TORCH_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
            --include "torchaudio-${TORCHAUDIO_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
            --include "torchvision-${TORCHVISION_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
            --include "triton-${TRITON_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
            --include "apex-${APEX_VERSION}-${CP_VERSION}-linux_x86_64.whl"


      - name: Install python deps for CI
        run: |
          pip install -r external-builds/pytorch/requirements-ci.txt

      - name: (Re-)Generate Python package release index
        if: ${{ env.upload == 'true' }}
        env:
          # Environment variables to be set for `manage.py`
          CUSTOM_PREFIX: "${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}"
        run: |
          python ./build_tools/third_party/s3_management/manage.py ${{ env.CUSTOM_PREFIX }}