ci: Fixes for cuda13

.github/workflows/_build_in_container.yml

@@ -39,11 +39,11 @@ jobs:
           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/'
+          remove-dotnet: "true"
+          remove-android: "true"
+          remove-haskell: "true"
+          remove-codeql: "true"
+          build-mount-path: "/var/lib/docker/"
 
       - name: Restore /var/lib/docker/
         run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"
@@ -63,9 +63,9 @@ jobs:
         run: |
           echo "Free space:"
           df -h
-  
+
       - name: Pull the container
-        run: docker pull ${{ inputs.container-image }} 
+        run: docker pull ${{ inputs.container-image }}
 
       - name: Set CUDA and PyTorch versions
         run: |
@@ -94,7 +94,7 @@ jobs:
 
       - name: Build wheel
         id: build_wheel
-        env: 
+        env:
           CXX11_ABI: ${{ env.CXX11_ABI }}
           MATRIX_TORCH_VERSION: ${{ env.MATRIX_TORCH_VERSION}}
           WHEEL_CUDA_VERSION: ${{ env.WHEEL_CUDA_VERSION }}
@@ -106,6 +106,7 @@ jobs:
             --workdir /workspace \
             --volume $(pwd):/workspace \
             --volume $GITHUB_ENV:$GITHUB_ENV \
+            -e PIP_CONSTRAINT= \
             -e GITHUB_ENV=$GITHUB_ENV \
             -e CXX11_ABI=$CXX11_ABI \
             -e MATRIX_TORCH_VERSION=$MATRIX_TORCH_VERSION \

.github/workflows/publish.yaml

@@ -40,32 +40,21 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
+        # Using ubuntu-22.04 instead of 24.04 for more compatibility (glibc). Ideally we'd use the
         # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
         os: [ubuntu-22.04]
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-        torch-version: ["2.1.2", "2.2.2", "2.3.1", "2.4.0", "2.5.1", "2.6.0.dev20241001"]
-        cuda-version: ["11.8.0", "12.3.2"]
+        torch-version: ["2.4.0", "2.5.1", "2.6.0", "2.7.1"]
+        cuda-version: ["11.8.0", "12.9.1"]
         # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
         # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
         # Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs)
         # when building without C++11 ABI and using it on nvcr images.
         cxx11_abi: ["FALSE", "TRUE"]
         exclude:
           # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
-          # Pytorch < 2.2 does not support Python 3.12
-          - torch-version: "2.1.2"
-            python-version: "3.12"
-          # Pytorch < 2.5 does not support Python 3.13
-          - torch-version: "2.1.2"
-            python-version: "3.13"
-          - torch-version: "2.2.2"
-            python-version: "3.13"
-          - torch-version: "2.3.1"
-            python-version: "3.13"
           - torch-version: "2.4.0"
             python-version: "3.13"
-
     uses: ./.github/workflows/_build.yml
     with:
       runs-on: ${{ matrix.os }}

setup.py

@@ -293,7 +293,15 @@ def get_wheel_url():
         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
         # 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")
+        if torch_cuda_version.major == 11:
+            torch_cuda_version = parse("11.8")
+        elif 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")
+
         cuda_version = f"{torch_cuda_version.major}"
 
     gpu_compute_version = hip_ver if HIP_BUILD else cuda_version