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 12b50af)

* revert: drop ROOT_CONTAINER addition from Makefile

(cherry picked from commit f423145)

* style: use consistent three empty lines in Makefile

(cherry picked from commit 446b45a)

* refactor: add default value for parent-image

(cherry picked from commit 32955ce)

* revert: use original workflow structure

(cherry picked from commit 68c6744)

* refactor: use single build image step

(cherry picked from commit 5f1ac0a)

* fix: run merge tags regardless of repository owner

(cherry picked from commit 3fce366)

* refactor: build cuda12 instead of cuda tag

(cherry picked from commit 217144e)

* 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 <[email protected]>

Author: johanna-reiml-hpi

Diff for: .github/actions/load-image/action.yml

@@ -8,17 +8,20 @@ inputs:
   platform:
     description: Image platform
     required: true
+  variant:
+    description: Variant tag prefix
+    required: true
 
 runs:
   using: composite
   steps:
     - 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

Diff for: .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,14 +75,15 @@ jobs:
         with:
           image: ${{ inputs.parent-image }}
           platform: ${{ inputs.platform }}
+          variant: ${{ inputs.parent-variant }}
 
       - name: Pull ubuntu:22.04 image 📥
         if: inputs.parent-image == ''
         run: docker pull ubuntu:22.04
         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
-          path: /tmp/jupyter/tags/${{ inputs.platform }}-${{ inputs.image }}.txt
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-tags
+          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
-          path: /tmp/jupyter/manifests/${{ inputs.platform }}-${{ inputs.image }}-*.md
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-manifest
+          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
-          path: /tmp/jupyter/hist_lines/${{ inputs.platform }}-${{ inputs.image }}-*.txt
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-history_line
+          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 }}
-          path: /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}.tar.zst
+          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}
+          path: /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst
           retention-days: 3

Diff for: .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

Diff for: .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