Add pytorch-notebook image variants with cuda 11 and 12 (x86_64 versions only) (#2091)

* feat: build cuda variants of pytorch

* feat: build with variant tag

* style: remove unused import

* refactor: rename get_prefix params

(cherry picked from commit 12b50af258c2f331d4100fb63fd41ad1a30acb1d)

* revert: drop ROOT_CONTAINER addition from Makefile

(cherry picked from commit f42314513df2855957a05c6ba0c748d2df26d7b0)

* style: use consistent three empty lines in Makefile

(cherry picked from commit 446b45aab37a37720462b5df305ce96b139cf67a)

* refactor: add default value for parent-image

(cherry picked from commit 32955cec99c7202f0ce50647dfc61ec98f57f741)

* revert: use original workflow structure

(cherry picked from commit 68c6744513636ec93d14f9bd0bbd123907efd13b)

* refactor: use single build image step

(cherry picked from commit 5f1ac0aeedcb5969a6d4b2a5bc939817378ab55d)

* fix: run merge tags regardless of repository owner

(cherry picked from commit 3fce366a98adc5db0d127f28ddf3157d13297a0f)

* refactor: build cuda12 instead of cuda tag

(cherry picked from commit 217144ecd322356376f04efb92792a20b4380177)

* docs: add note about CUDA tags to documentation

* refactor: add default value for variant in build-test-upload

* refactor: swap ordering of cuda11/cuda12 variants

* refactor: remove optional str type in arg parser

* fix: add proper env variables to CUDA Dockerfiles

* fix: remove CUDA build for aarch64

* fix: use latest NVIDIA documentation link

* fix: skip aarch64 tags file for CUDA variants

---------

Co-authored-by: zynaa <7562909-zynaa@users.noreply.gitlab.com>

.github/actions/load-image/action.yml

@@ -8,6 +8,9 @@ inputs:
   platform:
     description: Image platform
     required: true
+  variant:
+    description: Variant tag prefix
+    required: true
 
 runs:
   using: composite
@@ -15,10 +18,10 @@ runs:
     - name: Download built image 📥
       uses: actions/download-artifact@v4
       with:
-        name: ${{ inputs.image }}-${{ inputs.platform }}
+        name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}
         path: /tmp/jupyter/images/
     - name: Load downloaded image to docker 📥
       run: |
-        zstd --uncompress --stdout --rm /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}.tar.zst | docker load
+        zstd --uncompress --stdout --rm /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst | docker load
         docker image ls --all
       shell: bash

.github/workflows/docker-build-test-upload.yml

@@ -11,10 +11,20 @@ on:
         description: Parent image name
         required: true
         type: string
+      parent-variant:
+        description: Parent variant tag prefix
+        required: false
+        type: string
+        default: default
       image:
         description: Image name
         required: true
         type: string
+      variant:
+        description: Variant tag prefix
+        required: false
+        type: string
+        default: default
       platform:
         description: Image platform
         required: true
@@ -29,6 +39,19 @@ jobs:
     runs-on: ${{ inputs.runs-on }}
 
     steps:
+      # Image with CUDA needs extra disk space
+      - name: Free disk space 🧹
+        if: contains(inputs.variant, 'cuda') && inputs.platform == 'x86_64'
+        uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be
+        with:
+          tool-cache: false
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: false
+          docker-images: false
+          swap-storage: false
+
       - name: Checkout Repo ⚡️
         uses: actions/checkout@v4
       - name: Create dev environment 📦
@@ -52,6 +75,7 @@ jobs:
         with:
           image: ${{ inputs.parent-image }}
           platform: ${{ inputs.platform }}
+          variant: ${{ inputs.parent-variant }}
 
       - name: Pull ubuntu:22.04 image 📥
         if: inputs.parent-image == ''
@@ -59,7 +83,7 @@ jobs:
         shell: bash
 
       - name: Build image 🛠
-        run: docker build --rm --force-rm --tag ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} images/${{ inputs.image }}/ --build-arg REGISTRY=${{ env.REGISTRY }} --build-arg OWNER=${{ env.OWNER }}
+        run: docker build --rm --force-rm --tag ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} images/${{ inputs.image }}/${{ inputs.variant != 'default' && inputs.variant || '.' }}/ --build-arg REGISTRY=${{ env.REGISTRY }} --build-arg OWNER=${{ env.OWNER }}
         env:
           DOCKER_BUILDKIT: 1
           # Full logs for CI build
@@ -72,39 +96,39 @@ jobs:
 
       - name: Write tags file 🏷
         run: |
-          python3 -m tagging.write_tags_file --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
+          python3 -m tagging.write_tags_file --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }} --variant ${{ inputs.variant }}
         shell: bash
       - name: Upload tags file 💾
         uses: actions/upload-artifact@v4
         with:
-          name: ${{ inputs.image }}-${{ inputs.platform }}-tags
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-tags
-          path: /tmp/jupyter/tags/${{ inputs.platform }}-${{ inputs.image }}.txt
+          path: /tmp/jupyter/tags/${{ inputs.platform }}-${{ inputs.variant }}-${{ inputs.image }}.txt
           retention-days: 3
 
       - name: Write manifest and build history file 🏷
-        run: python3 -m tagging.write_manifest --short-image-name ${{ inputs.image }} --hist-lines-dir /tmp/jupyter/hist_lines/ --manifests-dir /tmp/jupyter/manifests/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
+        run: python3 -m tagging.write_manifest --short-image-name ${{ inputs.image }} --hist-lines-dir /tmp/jupyter/hist_lines/ --manifests-dir /tmp/jupyter/manifests/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }} --variant ${{ inputs.variant }}
         shell: bash
       - name: Upload manifest file 💾
         uses: actions/upload-artifact@v4
         with:
-          name: ${{ inputs.image }}-${{ inputs.platform }}-manifest
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-manifest
-          path: /tmp/jupyter/manifests/${{ inputs.platform }}-${{ inputs.image }}-*.md
+          path: /tmp/jupyter/manifests/${{ inputs.platform }}-${{ inputs.variant }}-${{ inputs.image }}-*.md
           retention-days: 3
       - name: Upload build history line 💾
         uses: actions/upload-artifact@v4
         with:
-          name: ${{ inputs.image }}-${{ inputs.platform }}-history_line
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-history_line
-          path: /tmp/jupyter/hist_lines/${{ inputs.platform }}-${{ inputs.image }}-*.txt
+          path: /tmp/jupyter/hist_lines/${{ inputs.platform }}-${{ inputs.variant }}-${{ inputs.image }}-*.txt
           retention-days: 3
 
       - name: Save image as a tar for later use 💾
         run: |
           mkdir -p /tmp/jupyter/images/
-          docker save ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} | zstd > /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}.tar.zst
+          docker save ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} | zstd > /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst
         shell: bash
       - name: Upload image as artifact 💾
         uses: actions/upload-artifact@v4
         with:
-          name: ${{ inputs.image }}-${{ inputs.platform }}
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}
-          path: /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}.tar.zst
+          path: /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst
           retention-days: 3

.github/workflows/docker-merge-tags.yml

@@ -7,6 +7,10 @@ env:
 on:
   workflow_call:
     inputs:
+      variant:
+        description: Variant tag prefix
+        required: true
+        type: string
       image:
         description: Image name
         required: true
@@ -30,13 +34,14 @@ jobs:
       - name: Download x86_64 tags file 📥
         uses: actions/download-artifact@v4
         with:
-          name: ${{ inputs.image }}-x86_64-tags
+          name: ${{ inputs.image }}-x86_64-${{ inputs.variant }}-tags
           path: /tmp/jupyter/tags/
       - name: Download aarch64 tags file 📥
         uses: actions/download-artifact@v4
         with:
-          name: ${{ inputs.image }}-aarch64-tags
+          name: ${{ inputs.image }}-aarch64-${{ inputs.variant }}-tags
           path: /tmp/jupyter/tags/
+        if: github.repository_owner == 'jupyter' && !contains(inputs.variant, 'cuda')
 
       # Docker might be stuck when pulling images
       # https://github.com/docker/for-mac/issues/2083
@@ -57,5 +62,5 @@ jobs:
 
       - name: Merge tags for the images 🔀
         if: env.PUSH_TO_REGISTRY == 'true'
-        run: python3 -m tagging.merge_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/
+        run: python3 -m tagging.merge_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --variant ${{ inputs.variant }}
         shell: bash

.github/workflows/docker-tag-push.yml

@@ -16,6 +16,10 @@ on:
         description: Image platform
         required: true
         type: string
+      variant:
+        description: Variant tag prefix
+        required: true
+        type: string
     secrets:
       REGISTRY_USERNAME:
         required: true
@@ -27,6 +31,19 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
+      # Image with CUDA needs extra disk space
+      - name: Free disk space 🧹
+        if: contains(inputs.variant, 'cuda') && inputs.platform == 'x86_64'
+        uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be
+        with:
+          tool-cache: false
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: false
+          docker-images: false
+          swap-storage: false
+
       - name: Checkout Repo ⚡️
         uses: actions/checkout@v4
       - name: Create dev environment 📦
@@ -36,6 +53,7 @@ jobs:
         with:
           image: ${{ inputs.image }}
           platform: ${{ inputs.platform }}
+          variant: ${{ inputs.variant }}
 
       - name: Login to Registry 🔐
         if: env.PUSH_TO_REGISTRY == 'true'
@@ -48,10 +66,10 @@ jobs:
       - name: Download tags file 📥
         uses: actions/download-artifact@v4
         with:
-          name: ${{ inputs.image }}-${{ inputs.platform }}-tags
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-tags
           path: /tmp/jupyter/tags/
       - name: Apply tags to the loaded image 🏷
-        run: python3 -m tagging.apply_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --platform ${{ inputs.platform }} --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
+        run: python3 -m tagging.apply_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --platform ${{ inputs.platform }} --variant ${{ inputs.variant }} --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
       # This step is needed to prevent pushing non-multi-arch "latest" tag
       - name: Remove the "latest" tag from the image 🗑️
         run: docker image rmi ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }}:latest

.github/workflows/docker.yml

@@ -216,6 +216,28 @@ jobs:
     needs: [x86_64-scipy]
     if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
 
+  x86_64-pytorch-cuda11:
+    uses: ./.github/workflows/docker-build-test-upload.yml
+    with:
+      parent-image: scipy-notebook
+      image: pytorch-notebook
+      variant: cuda11
+      platform: x86_64
+      runs-on: ubuntu-latest
+    needs: [x86_64-scipy]
+    if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
+
+  x86_64-pytorch-cuda12:
+    uses: ./.github/workflows/docker-build-test-upload.yml
+    with:
+      parent-image: scipy-notebook
+      image: pytorch-notebook
+      variant: cuda12
+      platform: x86_64
+      runs-on: ubuntu-latest
+    needs: [x86_64-scipy]
+    if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
+
   aarch64-datascience:
     uses: ./.github/workflows/docker-build-test-upload.yml
     with:
@@ -280,25 +302,26 @@ jobs:
     uses: ./.github/workflows/docker-tag-push.yml
     with:
       platform: aarch64
-      image: ${{ matrix.image }}
+      image: ${{ matrix.image-variant.image }}
+      variant: ${{ matrix.image-variant.variant }}
     secrets:
       REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
       REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
     strategy:
       matrix:
-        image:
+        image-variant:
           [
-            docker-stacks-foundation,
+            { image: docker-stacks-foundation, variant: default },
-            base-notebook,
+            { image: base-notebook, variant: default },
-            minimal-notebook,
+            { image: minimal-notebook, variant: default },
-            scipy-notebook,
+            { image: scipy-notebook, variant: default },
-            r-notebook,
+            { image: r-notebook, variant: default },
-            julia-notebook,
+            { image: julia-notebook, variant: default },
-            tensorflow-notebook,
+            { image: tensorflow-notebook, variant: default },
-            pytorch-notebook,
+            { image: pytorch-notebook, variant: default },
-            datascience-notebook,
+            { image: datascience-notebook, variant: default },
-            pyspark-notebook,
+            { image: pyspark-notebook, variant: default },
-            all-spark-notebook,
+            { image: all-spark-notebook, variant: default },
           ]
     needs:
       [
@@ -320,13 +343,18 @@ jobs:
     uses: ./.github/workflows/docker-tag-push.yml
     with:
       platform: aarch64
-      image: ${{ matrix.image }}
+      image: ${{ matrix.image-variant.image }}
+      variant: ${{ matrix.image-variant.variant }}
     secrets:
       REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
       REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
     strategy:
       matrix:
-        image: [docker-stacks-foundation, base-notebook]
+        image-variant:
+          [
+            { image: docker-stacks-foundation, variant: default },
+            { image: base-notebook, variant: default },
+          ]
     needs: [aarch64-foundation, aarch64-base]
     if: github.repository_owner == 'jupyter' && contains(github.event.pull_request.title, '[FAST_BUILD]')
 
@@ -334,25 +362,28 @@ jobs:
     uses: ./.github/workflows/docker-tag-push.yml
     with:
       platform: x86_64
-      image: ${{ matrix.image }}
+      image: ${{ matrix.image-variant.image }}
+      variant: ${{ matrix.image-variant.variant }}
     secrets:
       REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
       REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
     strategy:
       matrix:
-        image:
+        image-variant:
           [
-            docker-stacks-foundation,
+            { image: docker-stacks-foundation, variant: default },
-            base-notebook,
+            { image: base-notebook, variant: default },
-            minimal-notebook,
+            { image: minimal-notebook, variant: default },
-            scipy-notebook,
+            { image: scipy-notebook, variant: default },
-            r-notebook,
+            { image: r-notebook, variant: default },
-            julia-notebook,
+            { image: julia-notebook, variant: default },
-            tensorflow-notebook,
+            { image: tensorflow-notebook, variant: default },
-            pytorch-notebook,
+            { image: pytorch-notebook, variant: default },
-            datascience-notebook,
+            { image: pytorch-notebook, variant: cuda11 },
-            pyspark-notebook,
+            { image: pytorch-notebook, variant: cuda12 },
-            all-spark-notebook,
+            { image: datascience-notebook, variant: default },
+            { image: pyspark-notebook, variant: default },
+            { image: all-spark-notebook, variant: default },
           ]
     needs:
       [
@@ -374,54 +405,75 @@ jobs:
     uses: ./.github/workflows/docker-tag-push.yml
     with:
       platform: x86_64
-      image: ${{ matrix.image }}
+      image: ${{ matrix.image-variant.image }}
+      variant: ${{ matrix.image-variant.variant }}
     secrets:
       REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
       REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
     strategy:
       matrix:
-        image: [docker-stacks-foundation, base-notebook]
+        image-variant:
+          [
+            { image: docker-stacks-foundation, variant: default },
+            { image: base-notebook, variant: default },
+          ]
     needs: [x86_64-foundation, x86_64-base]
     if: contains(github.event.pull_request.title, '[FAST_BUILD]')
 
   merge-tags:
     uses: ./.github/workflows/docker-merge-tags.yml
     with:
-      image: ${{ matrix.image }}
+      image: ${{ matrix.image-variant.image }}
+      variant: ${{ matrix.image-variant.variant }}
     secrets:
       REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
       REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
     strategy:
       matrix:
-        image:
+        image-variant:
           [
-            docker-stacks-foundation,
+            { image: docker-stacks-foundation, variant: default },
-            base-notebook,
+            { image: base-notebook, variant: default },
-            minimal-notebook,
+            { image: minimal-notebook, variant: default },
-            scipy-notebook,
+            { image: scipy-notebook, variant: default },
-            r-notebook,
+            { image: r-notebook, variant: default },
-            julia-notebook,
+            { image: julia-notebook, variant: default },
-            tensorflow-notebook,
+            { image: tensorflow-notebook, variant: default },
-            pytorch-notebook,
+            { image: pytorch-notebook, variant: default },
-            datascience-notebook,
+            { image: pytorch-notebook, variant: cuda11 },
-            pyspark-notebook,
+            { image: pytorch-notebook, variant: cuda12 },
-            all-spark-notebook,
+            { image: datascience-notebook, variant: default },
+            { image: pyspark-notebook, variant: default },
+            { image: all-spark-notebook, variant: default },
           ]
     needs: [aarch64-images-tag-push, x86_64-images-tag-push]
-    if: github.repository_owner == 'jupyter' && !contains(github.event.pull_request.title, '[FAST_BUILD]')
+    if: |
+      always() &&
+      needs.x86_64-images-tag-push.result == 'success' &&
+      (needs.aarch64-images-tag-push.result == 'success' || needs.aarch64-images-tag-push.result == 'skipped') &&
+      !contains(github.event.pull_request.title, '[FAST_BUILD]')
 
   merge-tags-fast:
     uses: ./.github/workflows/docker-merge-tags.yml
     with:
-      image: ${{ matrix.image }}
+      image: ${{ matrix.image-variant.image }}
+      variant: ${{ matrix.image-variant.variant }}
     secrets:
       REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
       REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
     strategy:
       matrix:
-        image: [docker-stacks-foundation, base-notebook]
+        image-variant:
+          [
+            { image: docker-stacks-foundation, variant: default },
+            { image: base-notebook, variant: default },
+          ]
     needs: [aarch64-images-tag-push-fast, x86_64-images-tag-push-fast]
-    if: github.repository_owner == 'jupyter' && contains(github.event.pull_request.title, '[FAST_BUILD]')
+    if: |
+      always() &&
+      needs.x86_64-images-tag-push-fast.result == 'success' &&
+      (needs.aarch64-images-tag-push-fast.result == 'success' || needs.aarch64-images-tag-push-fast.result == 'skipped') &&
+      contains(github.event.pull_request.title, '[FAST_BUILD]')
 
   wiki-update:
     uses: ./.github/workflows/docker-wiki-update.yml

Makefile

@@ -68,10 +68,11 @@ linkcheck-docs: ## check broken links
 
 
 
+hook/%: VARIANT?=default
 hook/%: ## run post-build hooks for an image
-	python3 -m tagging.write_tags_file --short-image-name "$(notdir $@)" --tags-dir /tmp/jupyter/tags/ --registry "$(REGISTRY)" --owner "$(OWNER)" && \
+	python3 -m tagging.write_tags_file --short-image-name "$(notdir $@)" --tags-dir /tmp/jupyter/tags/ --registry "$(REGISTRY)" --owner "$(OWNER)" --variant "$(VARIANT)" && \
-	python3 -m tagging.write_manifest --short-image-name "$(notdir $@)" --hist-lines-dir /tmp/jupyter/hist_lines/ --manifests-dir /tmp/jupyter/manifests/ --registry "$(REGISTRY)" --owner "$(OWNER)" && \
+	python3 -m tagging.write_manifest --short-image-name "$(notdir $@)" --hist-lines-dir /tmp/jupyter/hist_lines/ --manifests-dir /tmp/jupyter/manifests/ --registry "$(REGISTRY)" --owner "$(OWNER)" --variant "$(VARIANT)" && \
-	python3 -m tagging.apply_tags --short-image-name "$(notdir $@)" --tags-dir /tmp/jupyter/tags/ --platform "$(shell uname -m)" --registry "$(REGISTRY)" --owner "$(OWNER)"
+	python3 -m tagging.apply_tags --short-image-name "$(notdir $@)" --tags-dir /tmp/jupyter/tags/ --platform "$(shell uname -m)" --variant "$(VARIANT)" --registry "$(REGISTRY)" --owner "$(OWNER)"
 hook-all: $(foreach I, $(ALL_IMAGES), hook/$(I)) ## run post-build hooks for all images
 
 

docs/using/selecting.md

@@ -197,6 +197,8 @@ It contains:
 - Everything in `jupyter/scipy-notebook` and its ancestor images
 - [pytorch](https://pytorch.org/) machine learning library
 
+> **GPU Acceleration:** Append a CUDA version prefix (`cuda11-` for CUDA 11 or `cuda12-` for CUDA 12) to the image tag to allow PyTorch operations to use compatible NVIDIA GPUs for accelerated computation.
+
 ### jupyter/datascience-notebook
 
 [Source on GitHub](https://github.com/jupyter/docker-stacks/tree/main/images/datascience-notebook) |
@@ -317,7 +319,7 @@ See the [contributing guide](../contributing/stacks.md) for information about ho
 [almond]: https://almond.sh
 [almond_b]: https://mybinder.org/v2/gh/almond-sh/examples/master?urlpath=lab%2Ftree%2Fnotebooks%2Findex.ipynb
 
-### GPU-accelerated notebooks
+### Other GPU-accelerated notebooks
 
 | Flavor             | Description                                                                                                                                                                                                                                                                                                                                                  |
 | ------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |

images/pytorch-notebook/Dockerfile

@@ -11,7 +11,7 @@ LABEL maintainer="Jupyter Project <jupyter@googlegroups.com>"
 # Fix: https://github.com/koalaman/shellcheck/wiki/SC3014
 SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 
-# Install PyTorch with pip
+# Install PyTorch with pip (https://pytorch.org/get-started/locally/)
 # hadolint ignore=DL3013
 RUN pip install --no-cache-dir --index-url 'https://download.pytorch.org/whl/cpu' \
     'torch' \

images/pytorch-notebook/cuda11/Dockerfile

@@ -0,0 +1,25 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+ARG REGISTRY=quay.io
+ARG OWNER=jupyter
+ARG BASE_CONTAINER=$REGISTRY/$OWNER/scipy-notebook
+FROM $BASE_CONTAINER
+
+LABEL maintainer="Jupyter Project <jupyter@googlegroups.com>"
+
+# Fix: https://github.com/hadolint/hadolint/wiki/DL4006
+# Fix: https://github.com/koalaman/shellcheck/wiki/SC3014
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+# Install PyTorch with pip (https://pytorch.org/get-started/locally/)
+# hadolint ignore=DL3013
+RUN pip install --no-cache-dir --extra-index-url=https://pypi.nvidia.com --index-url 'https://download.pytorch.org/whl/cu118' \
+    'torch' \
+    'torchvision' \
+    'torchaudio'  && \
+    fix-permissions "${CONDA_DIR}" && \
+    fix-permissions "/home/${NB_USER}"
+
+# https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/docker-specialized.html#dockerfiles
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility

images/pytorch-notebook/cuda12/Dockerfile

@@ -0,0 +1,25 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+ARG REGISTRY=quay.io
+ARG OWNER=jupyter
+ARG BASE_CONTAINER=$REGISTRY/$OWNER/scipy-notebook
+FROM $BASE_CONTAINER
+
+LABEL maintainer="Jupyter Project <jupyter@googlegroups.com>"
+
+# Fix: https://github.com/hadolint/hadolint/wiki/DL4006
+# Fix: https://github.com/koalaman/shellcheck/wiki/SC3014
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+# Install PyTorch with pip (https://pytorch.org/get-started/locally/)
+# hadolint ignore=DL3013
+RUN pip install --no-cache-dir --extra-index-url=https://pypi.nvidia.com --index-url 'https://download.pytorch.org/whl/cu121' \
+    'torch' \
+    'torchvision' \
+    'torchaudio'  && \
+    fix-permissions "${CONDA_DIR}" && \
+    fix-permissions "/home/${NB_USER}"
+
+# https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/docker-specialized.html#dockerfiles
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility

tagging/apply_tags.py

@@ -8,6 +8,7 @@ from pathlib import Path
 import plumbum
 
 from tagging.get_platform import unify_aarch64
+from tagging.get_prefix import get_file_prefix_for_platform
 
 docker = plumbum.local["docker"]
 
@@ -20,14 +21,16 @@ def apply_tags(
     owner: str,
     tags_dir: Path,
     platform: str,
+    variant: str,
 ) -> None:
     """
     Tags <registry>/<owner>/<short_image_name>:latest with the tags reported by all taggers for this image
     """
     LOGGER.info(f"Tagging image: {short_image_name}")
 
+    file_prefix = get_file_prefix_for_platform(platform, variant)
     image = f"{registry}/{owner}/{short_image_name}:latest"
-    filename = f"{platform}-{short_image_name}.txt"
+    filename = f"{file_prefix}-{short_image_name}.txt"
     tags = (tags_dir / filename).read_text().splitlines()
 
     for tag in tags:
@@ -69,9 +72,19 @@ if __name__ == "__main__":
         required=True,
         help="Owner of the image",
     )
+    arg_parser.add_argument(
+        "--variant",
+        required=True,
+        help="Variant tag prefix",
+    )
     args = arg_parser.parse_args()
     args.platform = unify_aarch64(args.platform)
 
     apply_tags(
-        args.short_image_name, args.registry, args.owner, args.tags_dir, args.platform
+        args.short_image_name,
+        args.registry,
+        args.owner,
+        args.tags_dir,
+        args.platform,
+        args.variant,
     )

tagging/get_prefix.py

@@ -0,0 +1,25 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+from tagging.get_platform import get_platform
+
+DEFAULT_VARIANT = "default"
+
+
+def get_file_prefix_for_platform(platform: str, variant: str) -> str:
+    return f"{platform}-{variant}"
+
+
+def get_tag_prefix_for_platform(platform: str, variant: str) -> str:
+    if variant == DEFAULT_VARIANT:
+        return platform
+    return f"{platform}-{variant}"
+
+
+def get_file_prefix(variant: str) -> str:
+    platform = get_platform()
+    return get_file_prefix_for_platform(platform, variant)
+
+
+def get_tag_prefix(variant: str) -> str:
+    platform = get_platform()
+    return get_tag_prefix_for_platform(platform, variant)

tagging/merge_tags.py

@@ -8,6 +8,7 @@ from pathlib import Path
 import plumbum
 
 from tagging.get_platform import ALL_PLATFORMS
+from tagging.get_prefix import get_file_prefix_for_platform
 
 docker = plumbum.local["docker"]
 
@@ -16,6 +17,7 @@ LOGGER = logging.getLogger(__name__)
 
 def merge_tags(
     short_image_name: str,
+    variant: str,
     tags_dir: Path,
 ) -> None:
     """
@@ -26,9 +28,12 @@ def merge_tags(
     all_tags: set[str] = set()
 
     for platform in ALL_PLATFORMS:
-        filename = f"{platform}-{short_image_name}.txt"
+        file_prefix = get_file_prefix_for_platform(platform, variant)
-        tags = (tags_dir / filename).read_text().splitlines()
+        filename = f"{file_prefix}-{short_image_name}.txt"
-        all_tags.update(tag.replace(platform + "-", "") for tag in tags)
+        file_path = tags_dir / filename
+        if file_path.exists():
+            tags = file_path.read_text().splitlines()
+            all_tags.update(tag.replace(platform + "-", "") for tag in tags)
 
     LOGGER.info(f"Got tags: {all_tags}")
 
@@ -61,6 +66,11 @@ if __name__ == "__main__":
         required=True,
         help="Short image name",
     )
+    arg_parser.add_argument(
+        "--variant",
+        required=True,
+        help="Variant tag prefix",
+    )
     arg_parser.add_argument(
         "--tags-dir",
         required=True,
@@ -69,4 +79,4 @@ if __name__ == "__main__":
     )
     args = arg_parser.parse_args()
 
-    merge_tags(args.short_image_name, args.tags_dir)
+    merge_tags(args.short_image_name, args.variant, args.tags_dir)

tagging/write_manifest.py

@@ -9,7 +9,7 @@ from pathlib import Path
 from docker.models.containers import Container
 
 from tagging.docker_runner import DockerRunner
-from tagging.get_platform import get_platform
+from tagging.get_prefix import get_file_prefix, get_tag_prefix
 from tagging.get_taggers_and_manifests import get_taggers_and_manifests
 from tagging.git_helper import GitHelper
 from tagging.manifests import ManifestHeader, ManifestInterface
@@ -73,6 +73,7 @@ def write_manifest(
     short_image_name: str,
     registry: str,
     owner: str,
+    variant: str,
     hist_lines_dir: Path,
     manifests_dir: Path,
 ) -> None:
@@ -81,12 +82,12 @@ def write_manifest(
 
     image = f"{registry}/{owner}/{short_image_name}:latest"
 
-    file_prefix = get_platform()
+    file_prefix = get_file_prefix(variant)
     commit_hash_tag = GitHelper.commit_hash_tag()
     filename = f"{file_prefix}-{short_image_name}-{commit_hash_tag}"
 
     with DockerRunner(image) as container:
-        tags_prefix = get_platform()
+        tags_prefix = get_tag_prefix(variant)
         all_tags = [
             tags_prefix + "-" + tagger.tag_value(container) for tagger in taggers
         ]
@@ -137,6 +138,11 @@ if __name__ == "__main__":
         required=True,
         help="Owner of the image",
     )
+    arg_parser.add_argument(
+        "--variant",
+        required=True,
+        help="Variant tag prefix",
+    )
     args = arg_parser.parse_args()
 
     LOGGER.info(f"Current build timestamp: {BUILD_TIMESTAMP}")
@@ -145,6 +151,7 @@ if __name__ == "__main__":
         args.short_image_name,
         args.registry,
         args.owner,
+        args.variant,
         args.hist_lines_dir,
         args.manifests_dir,
     )

tagging/write_tags_file.py

@@ -6,7 +6,7 @@ import logging
 from pathlib import Path
 
 from tagging.docker_runner import DockerRunner
-from tagging.get_platform import get_platform
+from tagging.get_prefix import get_file_prefix, get_tag_prefix
 from tagging.get_taggers_and_manifests import get_taggers_and_manifests
 
 LOGGER = logging.getLogger(__name__)
@@ -16,6 +16,7 @@ def write_tags_file(
     short_image_name: str,
     registry: str,
     owner: str,
+    variant: str,
     tags_dir: Path,
 ) -> None:
     """
@@ -25,9 +26,10 @@ def write_tags_file(
     taggers, _ = get_taggers_and_manifests(short_image_name)
 
     image = f"{registry}/{owner}/{short_image_name}:latest"
-    tags_prefix = get_platform()
+    file_prefix = get_file_prefix(variant)
-    filename = f"{tags_prefix}-{short_image_name}.txt"
+    filename = f"{file_prefix}-{short_image_name}.txt"
 
+    tags_prefix = get_tag_prefix(variant)
     tags = [f"{registry}/{owner}/{short_image_name}:{tags_prefix}-latest"]
     with DockerRunner(image) as container:
         for tagger in taggers:
@@ -70,6 +72,17 @@ if __name__ == "__main__":
         required=True,
         help="Owner of the image",
     )
+    arg_parser.add_argument(
+        "--variant",
+        required=True,
+        help="Variant tag prefix",
+    )
     args = arg_parser.parse_args()
 
-    write_tags_file(args.short_image_name, args.registry, args.owner, args.tags_dir)
+    write_tags_file(
+        args.short_image_name,
+        args.registry,
+        args.owner,
+        args.variant,
+        args.tags_dir,
+    )