build_wheels_linux_v2.yml

name: Build Linux Wheels V2 (uv)

# V2 of build_wheels_linux.yml that provisions Python and the build toolchain
# with `uv` instead of conda.
#
# What changed vs. v1 (build_wheels_linux.yml):
#   - No conda / miniforge. `uv` installs the requested Python (including
#     free-threaded builds like 3.13t/3.14t and pre-releases) via
#     python-build-standalone and creates an isolated venv.
#   - The build env is a uv `--seed` venv (ships pip) that is activated for every
#     build/test step, so the existing `pip`/`python`/`setup.py` invocations work
#     unchanged. `CONDA_RUN` is kept (as an empty string) only so the step bodies
#     stay diff-minimal against v1.
#   - `pytorch_pkg_helpers` (which produces BUILD_ENV_FILE) is run through uv.
#
# Notes / caveats:
#   - conda previously also supplied some *system* libraries (openssl, libpng,
#     libwebp, pkg-config). Those are expected to come from the manylinux builder
#     image (or the caller's pre-script). Only pip-installable build tools
#     (cmake/ninja/wheel/build) are installed here.
#   - BUILD_ENV_FILE (from pytorch_pkg_helpers) exports
#     `PATH="/opt/python/<abi>/bin:$PATH"` for the in-image CPython. Since v2 uses
#     a uv venv instead, every step re-prepends the venv bin *after* sourcing
#     BUILD_ENV_FILE so the uv interpreter wins.

on:
  workflow_call:
    inputs:
      repository:
        description: 'Repository to checkout, defaults to ""'
        default: ""
        type: string
      ref:
        description: 'Reference to checkout, defaults to "nightly"'
        default: "nightly"
        type: string
      test-infra-repository:
        description: "Test infra repository to use"
        default: "pytorch/test-infra"
        type: string
      test-infra-ref:
        description: "Test infra reference to use"
        default: ""
        type: string
      build-matrix:
        description: "Build matrix to utilize"
        default: ""
        type: string
      pre-script:
        description: "Pre script to run prior to build"
        default: ""
        type: string
      post-script:
        description: "Post script to run prior to build"
        default: ""
        type: string
      smoke-test-script:
        description: "Script for Smoke Test for a specific domain"
        default: ""
        type: string
      env-var-script:
        description: "Script that sets Domain-Specific Environment Variables"
        default: ""
        type: string
      package-name:
        description: "Name of the actual python package that is imported in the smoke test"
        default: ""
        type: string
      build-target:
        description: "The target to build and publish (for repos that build multiple packages)"
        default: ""
        type: string
      trigger-event:
        description: "Trigger Event in caller that determines whether or not to upload"
        default: ""
        type: string
      cache-path:
        description: "The path(s) on the runner to cache or restore. The path is relative to repository."
        default: ""
        type: string
      cache-key:
        description: "The key created when saving a cache and the key used to search for a cache."
        default: ""
        type: string
      architecture:
        description: Architecture to build for x86_64 for default Linux, or aarch64 for Linux aarch64 builds
        required: false
        type: string
        default: x86_64
      runner:
        description: "Override runner for builds"
        default: ""
        type: string
      submodules:
        description: Works as stated in actions/checkout, but the default value is recursive
        required: false
        type: string
        default: recursive
      setup-miniconda:
        description: "Kept for caller compatibility; ignored in the uv-based v2."
        required: false
        type: boolean
        default: false
      build-platform:
        description: Platform to build wheels, choose from 'python-build-package' or 'setup-py'
        required: false
        type: string
        default: 'setup-py'
      upload-to-pypi:
        description: The comma-separated list of CUDA arch to be uploaded to pypi
        default: ""
        type: string
      build-command:
        description: The build command to use if build-platform is python-build-package
        required: false
        default: "python -m build --wheel"
        type: string
      wheel-upload-path:
        description: Custom wheel upload path
        required: false
        type: string
      wheel-nightly-policy:
        description: Custom wheel upload policy for nightly
        type: string
        default: ''
      also-upload-to-gpu-arch-versions:
        description: "Comma-separated list of extra gpu_arch_versions to upload the wheel to (e.g. 'xpu' to also upload a cpu wheel to the xpu index)"
        required: false
        type: string
        default: ''
      run-smoke-test:
        description: Set to true if smoke test is needed
        required: false
        type: boolean
        default: true
      python-version:
        description: "Override the uv-provisioned Python version (blank = use matrix.python_version)"
        required: false
        type: string
        default: ""
      pip-install-torch-extra-args:
        description: Extra arguments to pass to the command that install base torch dependency
        required: false
        default: ""
        type: string
      timeout:
        description: 'Timeout for the job (in minutes)'
        default: 120
        type: number
    secrets:
      PYPI_API_TOKEN:
        description: An optional token to upload to pypi
        required: false
      R2_ACCOUNT_ID:
        description: Cloudflare R2 account ID
        required: false
      R2_ACCESS_KEY_ID:
        description: Cloudflare R2 access key ID
        required: false
      R2_SECRET_ACCESS_KEY:
        description: Cloudflare R2 secret access key
        required: false

permissions:
  id-token: write
  contents: read

jobs:
  build:
    strategy:
      fail-fast: false
      matrix: ${{ fromJSON(inputs.build-matrix) }}
    env:
      PYTHON_VERSION: ${{ inputs.python-version != '' && inputs.python-version || matrix.python_version }}
      PACKAGE_TYPE: wheel
      REPOSITORY: ${{ inputs.repository }}
      REF: ${{ inputs.ref }}
      CU_VERSION: ${{ matrix.desired_cuda }}
      UPLOAD_TO_BASE_BUCKET: ${{ matrix.upload_to_base_bucket }}
      ARCH: ${{ inputs.architecture }}
      BUILD_TARGET: ${{ inputs.build-target }}
      # Where the uv-managed build venv lives; activated in every build/test step.
      BUILD_VENV: ${{ runner.temp }}/build_venv
    name: build-${{ matrix.build_name }}
    runs-on: ${{ (inputs.runner != '' && inputs.runner) || matrix.validation_runner }}
    environment: ${{(inputs.trigger-event == 'schedule' || (inputs.trigger-event == 'push' && (startsWith(github.event.ref, 'refs/heads/nightly') || startsWith(github.event.ref, 'refs/tags/v')))) && 'pytorchbot-env' || ''}}
    container:
      image: ${{ matrix.container_image }}
      options: ${{ matrix.gpu_arch_type == 'cuda' && '--gpus all -e NVIDIA_DRIVER_CAPABILITIES=compute,utility,video' || ' ' }}
    timeout-minutes: ${{ inputs.timeout }}
    steps:
      - name: Clean workspace
        shell: bash -l {0}
        run: |
          set -euxo pipefail
          echo "::group::Cleanup debug output"
          rm -rf "${GITHUB_WORKSPACE}"
          mkdir -p "${GITHUB_WORKSPACE}"

          if [[ "${{ inputs.architecture }}" = "aarch64" ]]; then
            rm -rf "${RUNNER_TEMP}/*"
          fi
          echo "::endgroup::"

      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          # Support the use case where we need to checkout someone's fork
          repository: ${{ inputs.test-infra-repository }}
          ref: ${{ inputs.test-infra-ref }}
          path: test-infra

      - uses: ./test-infra/.github/actions/set-channel

      - name: Set PYTORCH_VERSION
        if: ${{ env.CHANNEL == 'test' }}
        run: |
          # When building RC, set the version to be the current candidate version,
          # otherwise, leave it alone so nightly will pick up the latest
          echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}"

      # Checkout the target repository (v1 did this inside setup-binary-builds).
      - name: Checkout target repository
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: ${{ inputs.repository }}
          ref: ${{ inputs.ref }}
          submodules: ${{ inputs.submodules }}
          path: ${{ inputs.repository }}

      - name: Set artifact name
        shell: bash -l {0}
        run: |
          set -euxo pipefail
          # "/" is not allowed in artifact names; replace all forward slashes.
          # Convert e.g. 3.10.19 -> 3.10 while preserving suffixes like 3.13t.
          SHORT_PYTHON_VERSION=$(echo "${PYTHON_VERSION}" | sed -E 's/^([0-9]+\.[0-9]+[a-z]?).*/\1/')
          echo "ARTIFACT_NAME=${REPOSITORY//\//_}_${REF//\//_}_${SHORT_PYTHON_VERSION}_${CU_VERSION}_${ARCH}" >> "${GITHUB_ENV}"

      - name: Install uv
        shell: bash -l {0}
        run: |
          set -euxo pipefail
          # Container-safe install (no node/JS action needed inside the builder image).
          if ! command -v uv >/dev/null 2>&1; then
            curl -LsSf https://astral.sh/uv/install.sh | sh
          fi
          # uv installs to ~/.local/bin
          echo "${HOME}/.local/bin" >> "${GITHUB_PATH}"

      - name: Provision Python + build venv with uv
        shell: bash -l {0}
        run: |
          set -euxo pipefail
          export PATH="${HOME}/.local/bin:${PATH}"

          # uv understands the free-threaded suffix (e.g. 3.13t / 3.14t) and can
          # fetch pre-release interpreters, so no special-casing is needed.
          uv python install "${PYTHON_VERSION}"

          # --seed installs pip/setuptools so the existing `pip`/`python -m pip`
          # step bodies work unchanged.
          uv venv --seed --python "${PYTHON_VERSION}" "${BUILD_VENV}"

          # Build toolchain that conda used to provide (pip-installable subset).
          # System libs (openssl/libpng/libwebp/pkg-config) come from the builder
          # image or the caller's pre-script.
          VIRTUAL_ENV="${BUILD_VENV}" uv pip install \
            "cmake==3.31.2" "ninja==1.12.1" "wheel==0.37" "build==1.2.2"

          # Activate the venv for all subsequent steps. CONDA_RUN is kept as an
          # empty string only to keep the step bodies diff-minimal vs v1.
          echo "VIRTUAL_ENV=${BUILD_VENV}" >> "${GITHUB_ENV}"
          echo "${BUILD_VENV}/bin" >> "${GITHUB_PATH}"
          echo "CONDA_RUN=" >> "${GITHUB_ENV}"
          {
            echo "Provisioned interpreter:"
            "${BUILD_VENV}/bin/python" --version
          }

      - name: Generate BUILD_ENV_FILE from pytorch_pkg_helpers
        working-directory: ${{ inputs.repository }}
        shell: bash -l {0}
        run: |
          set -euxo pipefail
          export PATH="${HOME}/.local/bin:${PATH}"
          BUILD_ENV_FILE="${RUNNER_TEMP}/build_env_${GITHUB_RUN_ID}"
          # Run the helper isolated via uv (it only needs a stable 3.10 runtime;
          # it does not need the build interpreter).
          uv run --python 3.10 --with "${GITHUB_WORKSPACE}/test-infra/tools/pkg-helpers" \
            python -m pytorch_pkg_helpers > "${BUILD_ENV_FILE}"
          cat "${BUILD_ENV_FILE}"
          echo "BUILD_ENV_FILE=${BUILD_ENV_FILE}" >> "${GITHUB_ENV}"

      - name: Combine Env Var and Build Env Files
        if: ${{ inputs.env-var-script != '' }}
        working-directory: ${{ inputs.repository }}
        run: |
          set -euxo pipefail
          cat "${{ inputs.env-var-script }}" >> "${BUILD_ENV_FILE}"

      - name: Add XPU Env Vars in Build Env File
        if: ${{ matrix.gpu_arch_type == 'xpu' }}
        run: |
          {
            echo "set +u"
            echo "source /opt/intel/oneapi/compiler/latest/env/vars.sh"
            echo "source /opt/intel/oneapi/pti/latest/env/vars.sh"
          } >> "${BUILD_ENV_FILE}"

      - name: Install torch dependency
        run: |
          set -euxo pipefail
          # shellcheck disable=SC1090
          source "${BUILD_ENV_FILE}"
          # Re-prepend the uv venv so it wins over the /opt/python PATH that
          # BUILD_ENV_FILE exports.
          export PATH="${BUILD_VENV}/bin:${PATH}"
          # shellcheck disable=SC2086
          ${CONDA_RUN} ${PIP_INSTALL_TORCH} ${{ inputs.pip-install-torch-extra-args }}

      - name: Run Pre-Script with Caching
        if: ${{ inputs.pre-script != '' }}
        uses: ./test-infra/.github/actions/run-script-with-cache
        with:
          cache-path: ${{ inputs.cache-path }}
          cache-key: ${{ inputs.cache-key }}
          repository: ${{ inputs.repository  }}
          script: ${{ inputs.pre-script }}

      - name: Build the wheel (python-build-package)
        if: ${{ inputs.build-platform == 'python-build-package' }}
        working-directory: ${{ inputs.repository }}
        shell: bash -l {0}
        run: |
          set -euxo pipefail
          source "${BUILD_ENV_FILE}"
          export PATH="${BUILD_VENV}/bin:${PATH}"
          export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
          echo "Successfully installed Python build package"
          ${CONDA_RUN} ${{ inputs.build-command }}

      - name: Build the wheel (setup-py)
        if: ${{ inputs.build-platform == 'setup-py' }}
        working-directory: ${{ inputs.repository }}
        shell: bash -l {0}
        run: |
          set -euxo pipefail
          source "${BUILD_ENV_FILE}"
          export PATH="${BUILD_VENV}/bin:${PATH}"
          export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
          ${CONDA_RUN} python setup.py clean
          echo "Successfully ran python setup.py clean"
          ${CONDA_RUN} python setup.py bdist_wheel

      - name: Repair Manylinux_2_28 Wheel
        shell: bash -l {0}
        env:
          PACKAGE_NAME: ${{ inputs.package-name }}
          SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }}
        run: |
          set -euxo pipefail
          source "${BUILD_ENV_FILE}"
          export PATH="${BUILD_VENV}/bin:${PATH}"
          for pkg in ${{ inputs.repository }}/dist/*-linux_*.whl; do
            # if the glob didn't match anything
            if [[ ! -e $pkg ]]; then
                continue
            fi
            abs_pkg=$(realpath $pkg)
            ./test-infra/.github/scripts/repair_manylinux_2_28.sh $abs_pkg
          done

      - name: Run Post-Script
        if: ${{ inputs.post-script != '' }}
        uses: ./test-infra/.github/actions/run-script-with-cache
        with:
          repository: ${{ inputs.repository  }}
          script: ${{ inputs.post-script }}

      - name: Smoke Test
        if: ${{ inputs.run-smoke-test }}
        shell: bash -l {0}
        env:
          PACKAGE_NAME: ${{ inputs.package-name }}
          SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }}
        run: |
          set -euxo pipefail
          source "${BUILD_ENV_FILE}"
          export PATH="${BUILD_VENV}/bin:${PATH}"
          WHEEL_NAME=$(ls "${{ inputs.repository }}/dist/")
          echo "$WHEEL_NAME"

          ${CONDA_RUN} pip install "${{ inputs.repository }}/dist/$WHEEL_NAME"
          # Checking that we have a pinned version of torch in our dependency tree
          (
            pushd "${RUNNER_TEMP}"
            unzip -o "${GITHUB_WORKSPACE}/${{ inputs.repository }}/dist/$WHEEL_NAME"
            # Ensure that pytorch version is pinned, should output file where it was found
            grep "Requires-Dist: torch (==.*)" -r .
          )

          if [[ (! -f "${{ inputs.repository }}"/${SMOKE_TEST_SCRIPT}) ]]; then
            echo "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT} not found"
            ${CONDA_RUN} python -c "import ${PACKAGE_NAME}; print('package version is ', ${PACKAGE_NAME}.__version__)"
          else
            echo "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT} found"
            ${CONDA_RUN} python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}"
          fi
      # NB: Only upload to GitHub after passing smoke tests

      - name: Upload wheel to GitHub
        continue-on-error: true
        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
        with:
          name: ${{ env.ARTIFACT_NAME }}
          path: ${{ inputs.repository }}/dist/

  upload:
    needs: build
    uses: ./.github/workflows/_binary_upload.yml
    if: always()
    with:
      repository: ${{ inputs.repository }}
      ref: ${{ inputs.ref }}
      test-infra-repository: ${{ inputs.test-infra-repository }}
      test-infra-ref: ${{ inputs.test-infra-ref }}
      build-matrix: ${{ inputs.build-matrix }}
      architecture: ${{ inputs.architecture }}
      trigger-event: ${{ inputs.trigger-event }}
      upload-to-pypi: ${{ inputs.upload-to-pypi }}
      wheel-upload-path: ${{ inputs.wheel-upload-path }}
      wheel-nightly-policy: ${{ inputs.wheel-nightly-policy }}
      also-upload-to-gpu-arch-versions: ${{ inputs.also-upload-to-gpu-arch-versions }}
    secrets:
      PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
      R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
      R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
      R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
  cancel-in-progress: true