Upgrade to torch 2.9.1 and Python 3.14.2

The upstream torch arm64 wheels have been compiled for CUDA as well
since 2.6.0. This makes the upgrade a lot easier.

Different dependencies in the images caused difficulties in building
the derived images. The dependencies are now the same.

This now changes the support from sm_72 (Xavier) and sm_87 (Orin) to
sm_80 (Ampere) and sm_90 (Hopper) for the arm64 image.

Author: kohtala

.github/workflows/main.yml

@@ -8,11 +8,6 @@ on:
   push:
     branches:
     - main
-    paths-ignore:
-    - 'torch-wheel.Dockerfile'
-    - '.github/workflows/torch-wheel.yml'
-    - 'torchvision-wheel.Dockerfile'
-    - '.github/workflows/torchvision-wheel.yml'
   # We want to monthly update the base image for security
   # TODO Can we avoid building the same image?
   schedule:
@@ -28,42 +23,35 @@ jobs:
       matrix:
         include:
         # Check torch + torchvision compatibility from https://github.com/pytorch/vision?tab=readme-ov-file#installation
-        - torch: '2.3.1'
-          torch-wheel-tag: '2.3.1'
-          torchvision: '0.18.1'
-          python: '3.11.9-slim'
+        - torch: '2.9.1'
+          python: '3.14.2'
+          # We need to use project index for CUDA on aarch64 as pypi wheels are for CPU only
+          index: cu126
           platforms: linux/amd64,linux/arm64
-          constraints: constraints-2.3.1.txt
+          constraints: constraints-2.9.1.txt
 
-        - torch: '2.3.1'
-          torch-wheel-tag: '2.3.1'
-          torchvision: '0.18.1'
-          python: '3.11.10'
+        - torch: '2.9.1'
+          python: '3.14.2-slim'
+          index: cu126
           platforms: linux/amd64,linux/arm64
-          constraints: constraints-2.3.1.txt
-        - torch: '2.3.1'
-          torch-wheel-tag: '2.3.1'
-          torchvision: '0.18.1'
-          python: '3.11.10-slim'
-          platforms: linux/amd64,linux/arm64
-          constraints: constraints-2.3.1.txt
-        - torch: '2.3.1+cpu'
-          torch-wheel-tag: '2.3.1'
-          python: '3.11.10'
-          extra-index-url: 'https://download.pytorch.org/whl/cpu'
+          constraints: constraints-2.9.1.txt
+
+        - torch: '2.9.1'
+          python: '3.14.2'
+          index: cpu
           platforms: linux/amd64
-          constraints: constraints-2.3.1.txt
+          constraints: constraints-2.9.1.txt
     steps:
     - id: created
       run: echo "created=$(date --utc +%Y-%m-%dT%H:%M:%SZ)" >> $GITHUB_OUTPUT
 
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v6
 
     - name: Set up Docker Buildx
-      uses: docker/setup-buildx-action@b5ca514318bd6ebac0fb2aedd5d36ec1b5c232a2 # v3.10.0
+      uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
 
     - name: Log in to Docker Hub
-      uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772 # v3.4.0
+      uses: docker/login-action@5e57cd118135c172c3672efd75eb46360885c0ef # v3.6.0
       if: github.ref == 'refs/heads/main'
       with:
         username: ${{ secrets.DOCKERHUB_USERNAME }}
@@ -74,9 +62,9 @@ jobs:
       # Tags have limited set of valid character, '+' not included
       # https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pulling-manifests
       run: |
-        echo "tag=$(echo -n "${{ matrix.torch }}-${{ matrix.python }}" | tr -c 'a-zA-Z0-9._-' '[-*]')" >> $GITHUB_OUTPUT
+        echo "tag=$(echo -n "${{ matrix.torch }}${{ matrix.index && format('+{0}', matrix.index) || '' }}-${{ matrix.python }}" | tr -c 'a-zA-Z0-9._-' '[-*]')" >> $GITHUB_OUTPUT
 
-    - uses: docker/build-push-action@1dc73863535b631f98b2378be8619f83b136f4a0 # v6.17.0
+    - uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # v6.18.0
       with:
         context: .
         push: ${{ github.ref == 'refs/heads/main' }}
@@ -88,7 +76,7 @@ jobs:
           TORCH_REQUIREMENT=${{ matrix.url || format('torch=={0}', matrix.torch) }}
           TORCH_WHEEL_SOURCE=${{ matrix.torch-wheel-tag && format('{0}/torch-wheels:{1}', vars.DOCKERHUB_USERNAME, matrix.torch-wheel-tag) || 'scratch'}}
           TORCHVISION_WHEEL_SOURCE=${{ matrix.torchvision && format('{0}/torchvision-wheels:{1}', vars.DOCKERHUB_USERNAME, matrix.torchvision) || 'scratch'}}
-          EXTRA_INDEX_URL=${{ matrix.extra-index-url }}
+          EXTRA_INDEX_URL=${{ matrix.extra-index-url || format('https://download.pytorch.org/whl/{0}/', matrix.index) }}
           ${{ matrix.constraints && format('CONSTRAINTS={0}', matrix.constraints) || '' }}
         # TODO add the latest tag, maybe somehow via docker/metadata-action
         tags: |
@@ -105,7 +93,7 @@ jobs:
     # https://octokit.github.io/rest.js/v19#git-update-ref
     - name: Tag the commit or update tag
       if: github.ref == 'refs/heads/main'
-      uses: actions/github-script@v7
+      uses: actions/github-script@v8
       with:
         script: |
           try {
@@ -133,9 +121,9 @@ jobs:
     runs-on: ubuntu-latest
     if: github.ref == 'refs/heads/main'
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v6
     - name: Docker Hub Description
-      uses: peter-evans/dockerhub-description@432a30c9e07499fd01da9f8a49f0faf9e0ca5b77 # v4.0.2
+      uses: peter-evans/dockerhub-description@1b9a80c056b620d92cedb9d9b5a223409c68ddfa # v5.0.0
       with:
         username: ${{ secrets.DOCKERHUB_USERNAME }}
         password: ${{ secrets.DOCKERHUB_TOKEN }}

.github/workflows/torch-wheel.yml

@@ -1,14 +1,14 @@
 # syntax=docker/dockerfile:1
 
-ARG PYTHON=3.11.10
-ARG TORCH=2.3.1
+ARG PYTHON=3.14.2
+ARG TORCH=2.9.1
 ARG TORCH_REQUIREMENT="torch==${TORCH}"
-ARG EXTRA_INDEX_URL
+ARG EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu126/"
 ARG TORCH_WHEEL_SOURCE="scratch"
 ARG TORCHVISION_WHEEL_SOURCE="scratch"
 ARG CREATED
 ARG SOURCE_COMMIT
-ARG CONSTRAINTS=constraints-2.3.1.txt
+ARG CONSTRAINTS=constraints-2.9.1.txt
 
 # Using variable in RUN --mount=from gives error 'from' doesn't support variable expansion, define alias stage instead
 FROM ${TORCH_WHEEL_SOURCE} AS torch-wheel-image
@@ -18,8 +18,8 @@ FROM ${TORCHVISION_WHEEL_SOURCE} AS torchvision-wheel-image
 FROM python:${PYTHON}
 
 ARG TARGETPLATFORM
-RUN --mount=type=cache,target=/var/cache/apt,id=bookworm-/var/cache/apt-${TARGETPLATFORM} \
-    --mount=type=cache,target=/var/lib/apt,sharing=locked,id=bookworm-/var/lib/apt-${TARGETPLATFORM} \
+RUN --mount=type=cache,target=/var/cache/apt,id=trixie-/var/cache/apt-${TARGETPLATFORM} \
+    --mount=type=cache,target=/var/lib/apt,sharing=locked,id=trixie-/var/lib/apt-${TARGETPLATFORM} \
     <<NUR
     set -ex
 # To keep cache of downloaded .debs, replace docker configuration
@@ -28,15 +28,6 @@ RUN --mount=type=cache,target=/var/cache/apt,id=bookworm-/var/cache/apt-${TARGET
     apt-get update
     DEBIAN_FRONTEND=noninteractive \
     apt-get upgrade -y --no-install-recommends
-
-    case ${TARGETPLATFORM} in
-        # Extra dependencies needed to run pytorch on Jetson
-        # Retreived from: https://github.com/dusty-nv/jetson-containers/blob/master/packages/pytorch/Dockerfile
-        "linux/arm64")
-            DEBIAN_FRONTEND=noninteractive \
-            apt-get install -y libopenblas-dev libopenmpi-dev openmpi-common openmpi-bin gfortran libomp-dev
-        ;;
-    esac
 NUR
 
 COPY README.md LICENSE /

.github/workflows/torchvision-wheel.yml

@@ -0,0 +1,10 @@
+filelock==3.20.2
+fsspec==2025.12.0
+Jinja2==3.1.6
+MarkupSafe==3.0.3
+mpmath==1.3.0
+networkx==3.6.1
+setuptools==80.9.0; python_version >= "3.12"
+sympy==1.14.0
+triton==3.5.1
+typing_extensions==4.15.0