ci: Build and attach bdist wheels to release page (#2138)

* ci: Build and attach bdist wheels to release page
Signed-off-by: oliver könig <okoenig@nvidia.com>

* free up space
Signed-off-by: oliver könig <okoenig@nvidia.com>

* cleanup
Signed-off-by: oliver könig <okoenig@nvidia.com>

* test
Signed-off-by: oliver könig <okoenig@nvidia.com>

* test
Signed-off-by: oliver könig <okoenig@nvidia.com>

* test
Signed-off-by: oliver könig <okoenig@nvidia.com>

* fix
Signed-off-by: oliver könig <okoenig@nvidia.com>

* test
Signed-off-by: oliver könig <okoenig@nvidia.com>

* fix
Signed-off-by: oliver könig <okoenig@nvidia.com>

* fix
Signed-off-by: oliver könig <okoenig@nvidia.com>

* fix
Signed-off-by: oliver könig <okoenig@nvidia.com>

* fix
Signed-off-by: oliver könig <okoenig@nvidia.com>

* c28619d8999a147d5e09c1199f84ff6af6ad5794
Signed-off-by: oliver könig <okoenig@nvidia.com>

* c28619d8999a147d5e09c1199f84ff6af6ad5794
Signed-off-by: oliver könig <okoenig@nvidia.com>

* Reduce months to check from 7 to 5
Signed-off-by: oliver könig <okoenig@nvidia.com>

* Update .github/scripts/check_for_ngc_images.sh
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update .github/actions/build-pytorch-wheel/build.sh
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: oliver könig <okoenig@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

.github/actions/build-pytorch-wheel/Dockerfile

+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+FROM ubuntu:22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=$PATH:$CUDA_HOME/bin
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;9.0"
+
+ARG PYTHON_VERSION=3.12
+ARG TORCH_VERSION=2.9.1
+ARG CUDA_VERSION=12.9.1
+ARG CUDNN_MAJOR_VERSION=9
+ENV PATH=/opt/venv/bin:$PATH
+ENV PYTHONUNBUFFERED=1
+ARG AARCH=x86_64
+
+# Install Python
+RUN apt-get update && \
+    apt-get install -y software-properties-common wget && \
+    add-apt-repository ppa:deadsnakes/ppa -y && \
+    apt-get install -y python$PYTHON_VERSION-dev python$PYTHON_VERSION-venv python3-pip && \
+    python$PYTHON_VERSION -m venv /opt/venv
+
+
+# Install cuda-toolkit
+RUN CUDA_MAJOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'}) && \
+    CUDA_MINOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $2'}) && \
+    rm /etc/apt/sources.list.d/cuda*.list || true && \
+    rm /etc/apt/sources.list.d/nvidia-cuda.list || true && \
+    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${AARCH}/cuda-keyring_1.1-1_all.deb && \
+    dpkg -i cuda-keyring_1.1-1_all.deb && \
+    rm cuda-keyring_1.1-1_all.deb && \
+    apt-get update && \
+    apt-get install -y cuda-toolkit-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} cudnn-cuda-$CUDA_MAJOR_VERSION libcudnn$CUDNN_MAJOR_VERSION-cuda-$CUDA_MAJOR_VERSION libnccl2 libnccl-dev cmake
+
+# Install PyTorch
+RUN export MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'}) && \
+    export MATRIX_TORCH_VERSION=$(echo $TORCH_VERSION | awk -F \. {'print $1 "." $2'}) && \
+    export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \ 
+    minv = {'2.5': 118, '2.6': 118, '2.7': 118, '2.8': 126, '2.9': 126}[env['MATRIX_TORCH_VERSION']]; \
+    maxv = {'2.5': 124, '2.6': 126, '2.7': 128, '2.8': 129, '2.9': 130}[env['MATRIX_TORCH_VERSION']]; \
+    print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)" \
+    ) && \
+    pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
\ No newline at end of file

.github/actions/build-pytorch-wheel/action.yml

+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+name: Build PyTorch Wheel
+description: Builds a PyTorch wheel for TransformerEngine
+
+inputs:
+  release-version:
+    description: 'The release version to use for the build'
+    required: true
+  python-version:
+    description: 'The Python version to use for the build'
+    required: true
+  cuda-version:
+    description: 'The CUDA version to use for the build'
+    required: true
+  cudnn-version:
+    description: 'The cuDNN version to use for the build'
+    required: true
+  torch-version:
+    description: 'The PyTorch version to use for the build'
+    required: true
+  cxx11_abi:
+    description: 'Enable torch flag C++11 ABI (TRUE/FALSE)'
+    required: true
+  base-image:
+    description: 'The base image to use for the build'
+    required: false
+  aarch:
+    description: 'The architecture to use for the build'
+    required: true
+outputs:
+  wheel_name:
+    description: 'The name of the built wheel'
+    value: ${{ steps.build_wheel.outputs.wheel_name }}
+
+runs:
+  using: 'composite'
+  steps:
+    - name: Move /var/lib/docker/
+      shell: bash -euxo pipefail {0}
+      run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker"
+
+    - name: Maximize build space
+      uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794
+      with:
+        root-reserve-mb: 5120
+        temp-reserve-mb: 32
+        swap-size-mb: 10240
+        remove-dotnet: 'true'
+        remove-android: 'true'
+        remove-haskell: 'true'
+        remove-codeql: 'true'
+        build-mount-path: '/var/lib/docker/'
+
+    - name: Restore /var/lib/docker/
+      shell: bash -euxo pipefail {0}
+      run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"
+
+    - name: Checkout
+      uses: actions/checkout@v4
+      with:
+        ref: ${{ inputs.release-version }}
+        submodules: recursive
+
+    - name: Checkout build tools
+      uses: actions/checkout@v4
+      with:
+        path: build-tools
+        submodules: recursive
+
+    - name: Build image
+      shell: bash -euxo pipefail {0}
+      env:
+        BASE_IMAGE: ${{ inputs.base-image }}
+      run: |
+        if [[ "${BASE_IMAGE}" == "" ]]; then
+          docker build \
+            -t transformer-engine-build \
+            -f build-tools/.github/actions/build-pytorch-wheel/Dockerfile \
+            --build-arg PYTHON_VERSION=${{ inputs.python-version }} \
+            --build-arg TORCH_VERSION=${{ inputs.torch-version }} \
+            --build-arg CUDA_VERSION=${{ inputs.cuda-version }} \
+            --build-arg CUDNN_MAJOR_VERSION=${{ inputs.cudnn-version }} \
+            --build-arg AARCH=${{ inputs.aarch }} \
+            .
+        else
+          docker pull ${BASE_IMAGE}
+          docker tag ${BASE_IMAGE} transformer-engine-build
+        fi
+    - name: Build wheel
+      shell: bash -euxo pipefail {0}
+      id: build_wheel
+      env:
+        CXX11_ABI: ${{ inputs.cxx11_abi }}
+      run: |
+        echo ::group::Build wheel
+
+        EXIT_CODE=$(docker run \
+            --rm \
+            --shm-size=64g \
+            --workdir /workspace/transformer_engine/pytorch \
+            --volume $(pwd):/workspace \
+            --volume $GITHUB_OUTPUT:$GITHUB_OUTPUT \
+            -e PIP_CONSTRAINT= \
+            -e CXX11_ABI=$CXX11_ABI \
+            -e GITHUB_OUTPUT=$GITHUB_OUTPUT \
+            transformer-engine-build bash /workspace/build-tools/.github/actions/build-pytorch-wheel/build.sh | tail -n 1)
+
+        # Do not fail the job if timeout killed the build
+        exit $EXIT_CODE
+        echo ::endgroup::
+
+    - name: Log Built Wheels
+      shell: bash -euxo pipefail {0}
+      run: |
+        ls transformer_engine/pytorch/dist

.github/actions/build-pytorch-wheel/build.sh

+#!/bin/bash
+
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+set -eoxu pipefail
+
+export NVTE_PYTORCH_FORCE_BUILD=TRUE
+export NVTE_NO_LOCAL_VERSION=1
+export NVTE_PYTORCH_FORCE_CXX11_ABI=$CXX11_ABI
+export PIP_CONSTRAINT=
+
+pip install wheel packaging nvidia-mathdx ninja pybind11
+
+# 5h timeout since GH allows max 6h and we want some buffer
+EXIT_CODE=0
+timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$?
+
+if [ $EXIT_CODE -eq 0 ]; then
+    wheel_name=$(python -c "import setup; print(setup.get_wheel_url()[1])" | tail -n 1)
+    ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
+    echo "wheel_name=${wheel_name}" | tee -a "$GITHUB_OUTPUT"
+fi
+
+echo $EXIT_CODE

.github/scripts/check_for_ngc_images.sh

+#!/bin/bash
+
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+# Configuration
+BASE_IMAGE="nvcr.io/nvidia/pytorch"
+TAG_SUFFIX="-py3"
+MONTHS_TO_CHECK=5 # Check current month and previous 4 months (total 5)
+
+# Initialize an array to store existing tags
+EXISTING_TAGS=()
+
+echo "Checking for existence of the last ${MONTHS_TO_CHECK} NGC PyTorch images: ${BASE_IMAGE}:YY.MM${TAG_SUFFIX}"
+echo "---------------------------------------------------------------------"
+
+# Loop through the last N months
+for i in $(seq 0 $((MONTHS_TO_CHECK - 1))); do
+    # Calculate Year and Month for the tag
+    CURRENT_YEAR=$(date +%Y)
+    CURRENT_MONTH=$(date +%m)
+
+    # Calculate target month and year
+    TARGET_DATE=$(date -d "$CURRENT_YEAR-$CURRENT_MONTH-01 -$i months" +%y.%m)
+
+    # Construct the full image tag and the tag-only string
+    IMAGE_TAG="${TARGET_DATE}${TAG_SUFFIX}"
+    FULL_IMAGE="${BASE_IMAGE}:${IMAGE_TAG}"
+
+    echo "Checking: ${FULL_IMAGE}"
+
+    # Use 'docker manifest inspect' to check for image existence without pulling.
+    if docker manifest inspect "${FULL_IMAGE}" > /dev/null 2>&1; then
+        echo "✅ EXISTS: Found."
+        # Add the tag-only string to the array
+        EXISTING_TAGS+=("nvcr.io/nvidia/pytorch:${IMAGE_TAG}")
+    else
+        echo "❌ MISSING: Not found."
+    fi
+done
+
+echo "---------------------------------------------------------------------"
+
+## JSON Output Generation
+# This uses the collected array to build a JSON string.
+
+# 1. Convert the shell array to a newline-separated string.
+TAGS_NL_SEP=$(printf "%s\n" "${EXISTING_TAGS[@]}")
+
+# 2. Use jq to read the newline-separated list and format it into a JSON array.
+# . | split("\n") | .[:-1] reads the input, splits it by newline, and removes the trailing empty element.
+if command -v jq &> /dev/null; then
+    JSON_STRING=$(echo -e "${TAGS_NL_SEP}" | jq -R -s 'split("\n") | .[:-1]')
+
+    echo "Generated JSON String of Existing Tags:"
+    echo "${JSON_STRING}"
+
+    # Optional: Save the JSON string to a variable for further use
+    # echo "JSON_STRING is now available in the shell if you source this script."
+else
+    echo "WARNING: 'jq' is not installed. Cannot format output as JSON."
+    echo "Found Tags: ${EXISTING_TAGS[*]}"
+fi
+
+echo "---"
+echo "Check complete."
+
+echo "${JSON_STRING}" > ngc_images.json

.github/workflows/attach-wheels-to-release.yml

+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+# This workflow will:
+# - Create a new Github release
+# - Build wheels for supported architectures
+# - Deploy the wheels to the Github release
+# - Release the static code to PyPi
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+name: Attach wheels to release
+
+on:
+  release:
+    types: [published]
+  workflow_dispatch:
+    inputs:
+      runs-on:
+        description: 'The runner to use for the build'
+        required: true
+        type: string
+        default: ubuntu-22.04
+      release-version:
+        description: 'Release version'
+        required: true
+        default: '0.1.0'
+      python-version:
+        description: 'Python version'
+        required: true
+        default: '3.12'
+      torch-version:
+        description: 'Torch version'
+        required: true
+        default: '2.8.0'
+      cuda-version:
+        description: 'CUDA version'
+        required: true
+        default: '12.9.1'
+      cudnn-version:
+        description: 'CUDNN version'
+        required: true
+        default: '9'
+      cxx11_abi:
+        description: 'C++11 ABI'
+        required: true
+        type: choice
+        default: 'TRUE'
+        options:
+          - 'TRUE'
+          - 'FALSE'
+      ngc-image:
+        description: 'NGC PyTorch image (will take precedence over the source build)'
+        required: false
+        type: string
+        default: ''
+jobs:
+  pre-flight:
+    runs-on: ubuntu-latest
+    outputs:
+      build-wheel-matrix: ${{ steps.matrix.outputs.matrix }}
+      release-assets-url: ${{ steps.release-assets-url.outputs.upload_url }}
+      ngc-images: ${{ steps.check_for_ngc_images.outputs.IMAGES }}
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Build release matrix
+        id: matrix
+        env:
+          EVENT: ${{ github.event_name }}
+        run: |
+          if [[ "$EVENT" == "release" ]]; then
+            MATRIX=$(echo '{
+              "os": ["ubuntu-22.04", "ubuntu-22.04-arm"],
+              "release-version": ["${{ github.event.release.tag_name }}"],
+              "python-version": ["3.12"], 
+              "torch-version": ["2.8.0"], 
+              "cuda-version": ["12.9.1"], 
+              "cudnn-version": ["9"], 
+              "cxx11_abi": ["TRUE"]
+            }' | jq -rc)
+          else
+            MATRIX=$(echo '{
+              "os": ["${{ inputs.runs-on }}"],
+              "release-version": ["${{ inputs.release-version }}"],
+              "python-version": ["${{ inputs.python-version }}"], 
+              "torch-version": ["${{ inputs.torch-version }}"], 
+              "cuda-version": ["${{ inputs.cuda-version }}"], 
+              "cudnn-version": ["${{ inputs.cudnn-version }}"], 
+              "cxx11_abi": ["${{ inputs.cxx11_abi }}"]
+            }' | jq -rc)
+          fi
+
+          echo "matrix=$MATRIX" | tee -a "$GITHUB_OUTPUT"
+
+      - name: Get Release with tag
+        id: get_current_release
+        uses: joutvhu/get-release@v1
+        if: ${{ github.event_name == 'workflow_dispatch' }}
+        with:
+          tag_name: ${{ inputs.release-version }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Get release assets url
+        env:
+          EVENT: ${{ github.event_name }}
+        if: ${{ (success() || !failure()) && !cancelled()}}
+        id: release-assets-url
+        run: |
+          if [[ "$EVENT" == "release" ]]; then
+            echo "upload_url=${{ github.event.release.upload_url }}" | tee -a "$GITHUB_OUTPUT"
+          else
+            echo "upload_url=${{ steps.get_current_release.outputs.upload_url }}" | tee -a "$GITHUB_OUTPUT"
+          fi
+
+      - name: Check for NGC PyTorch images
+        id: check_for_ngc_images
+        if: ${{ (success() || !failure()) && !cancelled()}}
+        env:
+          EVENT: ${{ github.event_name }}
+        run: |
+          if [[ "$EVENT" == "release" ]]; then
+            bash ./.github/scripts/check_for_ngc_images.sh
+            echo "IMAGES=$(cat ngc_images.json | jq -cr)" | tee -a $GITHUB_OUTPUT
+          else
+            echo 'IMAGES=["${{ inputs.ngc-image }}"]' | tee -a "$GITHUB_OUTPUT"
+          fi
+
+  build_wheels:
+    name: Build Wheel
+    runs-on: ${{ matrix.os }}
+    needs: pre-flight
+    if: ${{ github.event_name == 'release' || inputs.ngc-image == '' }}
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJson(needs.pre-flight.outputs.build-wheel-matrix) }}
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+
+      - name: 'Build PyTorch Wheel'
+        uses: ./.github/actions/build-pytorch-wheel
+        id: build-pytorch-wheel
+        with:
+          release-version: ${{ matrix.release-version }}
+          python-version: ${{ matrix.python-version }}
+          cuda-version: ${{ matrix.cuda-version }}
+          cudnn-version: ${{ matrix.cudnn-version }}
+          torch-version: ${{ matrix.torch-version }}
+          cxx11_abi: ${{ matrix.cxx11_abi }}
+          aarch: ${{ matrix.os == 'ubuntu-22.04' && 'x86_64' || 'sbsa' }}
+        env:
+          NVTE_FRAMEWORK: pytorch
+          MAX_JOBS: 1
+
+      - name: Upload Release Asset
+        id: upload_release_asset
+        uses: actions/upload-release-asset@v1
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          upload_url: ${{ needs.pre-flight.outputs.release-assets-url }}
+          asset_path: ./transformer_engine/pytorch/dist/${{ steps.build-pytorch-wheel.outputs.wheel_name }}
+          asset_name: ${{ steps.build-pytorch-wheel.outputs.wheel_name }}
+          asset_content_type: application/*
+
+  build_wheels_for_ngc:
+    name: Build Wheels for NGC PyTorch images
+    runs-on: ${{ matrix.os }}
+    needs: pre-flight
+    if: ${{ github.event_name == 'release' || inputs.ngc-image != '' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-22.04]
+        container-image: ${{ fromJson(needs.pre-flight.outputs.ngc-images) }}
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+
+      - name: 'Build PyTorch Wheel'
+        uses: ./.github/actions/build-pytorch-wheel
+        id: build-pytorch-wheel
+        with:
+          base-image: ${{ matrix.container-image }}
+
+      - name: Upload Release Asset
+        id: upload_release_asset
+        uses: actions/upload-release-asset@v1
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          upload_url: ${{ needs.pre-flight.outputs.release-assets-url }}
+          asset_path: ./transformer_engine/pytorch/dist/${{ steps.build-pytorch-wheel.outputs.wheel_name }}
+          asset_name: ${{ steps.build-pytorch-wheel.outputs.wheel_name }}
+          asset_content_type: application/*

.github/workflows/build.yml

@@ -35,26 +35,52 @@ jobs:
   pytorch:
     name: 'PyTorch'
     runs-on: ubuntu-latest
-    container:
-      image: nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04
-      options: --user root
     steps:
-      - name: 'Dependencies'
-        run: |
-          apt-get update
-          apt-get install -y git python3.9 pip cudnn9-cuda-12
-          pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript
+      - name: Move /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker"
+
+      - name: Maximize build space
+        uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794
+        with:
+          root-reserve-mb: 5120
+          temp-reserve-mb: 32
+          swap-size-mb: 10240
+          remove-dotnet: 'true'
+          remove-android: 'true'
+          remove-haskell: 'true'
+          remove-codeql: 'true'
+          build-mount-path: '/var/lib/docker/'
+
+      - name: Restore /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"
+
       - name: 'Checkout'
         uses: actions/checkout@v3
         with:
           submodules: recursive
+
+      - name: Start named container
+        run: |
+          docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 sleep infinity
+
+      - name: 'Dependencies'
+        run: |
+          docker exec builder bash -c '\
+            apt-get update && \
+            apt-get install -y git python3.9 pip cudnn9-cuda-12 && \
+            pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript && \
+            apt-get clean \
+          '
+
       - name: 'Build'
-        run: pip install --no-build-isolation . -v --no-deps
+        run: docker exec builder bash -c 'pip install --no-build-isolation . -v --no-deps'
         env:
           NVTE_FRAMEWORK: pytorch
           MAX_JOBS: 1
       - name: 'Sanity check'
-        run: python3 tests/pytorch/test_sanity_import.py
+        run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py'
   jax:
     name: 'JAX'
     runs-on: ubuntu-latest
@@ -78,22 +104,47 @@ jobs:
   all:
     name: 'All'
     runs-on: ubuntu-latest
-    container:
-      image: ghcr.io/nvidia/jax:jax
-      options: --user root
     steps:
-      - name: 'Dependencies'
-        run: |
-          pip install pybind11[global] einops onnxscript
-          pip install torch --index-url https://download.pytorch.org/whl/cu130
+      - name: Move /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker"
+
+      - name: Maximize build space
+        uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794
+        with:
+          root-reserve-mb: 5120
+          temp-reserve-mb: 32
+          swap-size-mb: 10240
+          remove-dotnet: 'true'
+          remove-android: 'true'
+          remove-haskell: 'true'
+          remove-codeql: 'true'
+          build-mount-path: '/var/lib/docker/'
+
+      - name: Restore /var/lib/docker/
+        shell: bash -euxo pipefail {0}
+        run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"
+
       - name: 'Checkout'
         uses: actions/checkout@v3
         with:
           submodules: recursive
+
+      - name: Start named container
+        run: |
+          docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity
+
+      - name: 'Dependencies'
+        run: |
+          docker exec builder bash -c '\
+            pip install pybind11[global] einops onnxscript && \
+            pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130
+          '
+
       - name: 'Build'
-        run: pip install --no-build-isolation . -v --no-deps
+        run: docker exec builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps'
         env:
           NVTE_FRAMEWORK: all
           MAX_JOBS: 1
       - name: 'Sanity check'
-        run: python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py
+        run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py'

.gitignore

+.venv
 *.o
 *.swp
 *.ii

transformer_engine/pytorch/setup.py

@@ -75,21 +75,29 @@ def get_platform():
 
 def get_wheel_url():
     """Construct the wheel URL for the current platform."""
-    torch_version_raw = parse(torch.__version__)
     python_version = f"cp{sys.version_info.major}{sys.version_info.minor}"
     platform_name = get_platform()
     nvte_version = te_version()
-    torch_version = f"{torch_version_raw.major}.{torch_version_raw.minor}"
     cxx11_abi = str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()
 
     # Determine the version numbers that will be used to determine the correct wheel
     # We're using the CUDA version used to build torch, not the one currently installed
     # _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME)
     torch_cuda_version = parse(torch.version.cuda)
-    # For CUDA 11, we only compile for CUDA 11.8, and for CUDA 12 we only compile for CUDA 12.3
+    # For CUDA 12 we only compile for CUDA 12.3
     # to save CI time. Minor versions should be compatible.
-    torch_cuda_version = parse("11.8") if torch_cuda_version.major == 11 else parse("12.3")
-    # cuda_version = f"{cuda_version_raw.major}{cuda_version_raw.minor}"
+    if torch_cuda_version.major == 12:
+        torch_cuda_version = parse("12.3")
+    elif torch_cuda_version.major == 13:
+        torch_cuda_version = parse("13.0")
+    else:
+        raise ValueError(f"CUDA version {torch_cuda_version} not supported")
+
+    if os.environ.get("NVIDIA_PRODUCT_NAME", "") == "PyTorch":
+        torch_version = str(os.environ.get("NVIDIA_PYTORCH_VERSION"))
+    else:
+        torch_version = f"{torch.__version__}"
+
     cuda_version = f"{torch_cuda_version.major}"
 
     # Determine wheel URL based on CUDA version, torch version, python version and OS
@@ -109,8 +117,10 @@ class CachedWheelsCommand(_bdist_wheel):
     """
 
     def run(self):
+        """Acts a proxy before _bdist_wheel.run() and downloads a prebuilt wheel if available."""
         if FORCE_BUILD:
             super().run()
+            return
 
         wheel_url, wheel_filename = get_wheel_url()
         print("Guessing wheel URL: ", wheel_url)
@@ -129,10 +139,12 @@ class CachedWheelsCommand(_bdist_wheel):
             wheel_path = os.path.join(self.dist_dir, archive_basename + ".whl")
             print("Raw wheel path", wheel_path)
             os.rename(wheel_filename, wheel_path)
+            return
         except (urllib.error.HTTPError, urllib.error.URLError):
             print("Precompiled wheel not found. Building from source...")
             # If the wheel could not be downloaded, build from source
             super().run()
+            return
 
 
 if __name__ == "__main__":