Optimize multi-arch Docker images

.github/workflows/docker-build.yml

@@ -29,6 +29,7 @@ jobs:
     strategy:
       matrix:
         arch: [amd64, arm64]
+        variant: [cpu, cuda]
       fail-fast: false
     steps:
       - name: Checkout code
@@ -54,10 +55,10 @@ jobs:
         with:
           images: ghcr.io/${{ github.repository }}
           tags: |
-            type=raw,value=latest,enable=${{ github.ref == 'refs/heads/main' }}
-            type=ref,event=pr
-            type=ref,event=tag
-            type=sha,format=short
+            type=raw,value=${{ matrix.variant }}-${{ matrix.arch }},enable=${{ github.ref == 'refs/heads/main' }}
+            type=raw,value=${{ matrix.variant }}-${{ matrix.arch }}-pr-${{ github.event.pull_request.number }},enable=${{ github.event_name == 'pull_request' }}
+            type=raw,value=${{ matrix.variant }}-${{ matrix.arch }}-${{ github.ref_name }},enable=${{ startsWith(github.ref, 'refs/tags/') }}
+            type=raw,value=${{ matrix.variant }}-${{ matrix.arch }}-sha-${{ github.sha }}
 
       - name: Build and push
         uses: docker/build-push-action@v6
@@ -66,4 +67,89 @@ jobs:
           platforms: linux/${{ matrix.arch }}
           push: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
           tags: ${{ steps.meta.outputs.tags }}
-          labels: ${{ steps.meta.outputs.labels }}
+          provenance: false
+          build-args: |
+            PYTORCH_VARIANT=${{ matrix.variant }}
+
+  create-manifests:
+    needs: build-and-push
+    runs-on: ubuntu-latest
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository
+    permissions:
+      packages: write
+    steps:
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+
+      - name: Setup Docker Buildx
+        uses: docker/setup-buildx-action@v3
+        with:
+          driver-opts: |
+            image=moby/buildkit:latest
+            network=host
+
+      - name: Log in to GHCR
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Create and push CPU manifest
+        run: |
+          # Determine the correct tag suffixes based on the event type
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            # PR build - use the commit SHA for more predictable references
+            AMD64_TAG="cpu-amd64-sha-${{ github.sha }}"
+            ARM64_TAG="cpu-arm64-sha-${{ github.sha }}"
+            TARGET_TAG="cpu-sha-${{ github.sha }}"
+          elif [[ "${{ startsWith(github.ref, 'refs/tags/') }}" == "true" ]]; then
+            # Tag build - use tag version in the name
+            AMD64_TAG="cpu-amd64-${{ github.ref_name }}"
+            ARM64_TAG="cpu-arm64-${{ github.ref_name }}"
+            TARGET_TAG="cpu-${{ github.ref_name }}"
+          else
+            # Main branch build - use simple arch tags
+            AMD64_TAG="cpu-amd64"
+            ARM64_TAG="cpu-arm64"
+            TARGET_TAG="cpu"
+          fi
+
+          # Create the manifest with the correct tag names
+          echo "Creating CPU manifest using $AMD64_TAG and $ARM64_TAG"
+          docker buildx imagetools create --tag ghcr.io/${{ github.repository }}:${TARGET_TAG} \
+            ghcr.io/${{ github.repository }}:${AMD64_TAG} \
+            ghcr.io/${{ github.repository }}:${ARM64_TAG}
+
+          # If on main branch, also tag as latest
+          if [[ "${{ github.ref }}" == "refs/heads/main" ]]; then
+            docker buildx imagetools create --tag ghcr.io/${{ github.repository }}:latest \
+              ghcr.io/${{ github.repository }}:${AMD64_TAG} \
+              ghcr.io/${{ github.repository }}:${ARM64_TAG}
+          fi
+
+      - name: Create and push CUDA manifest
+        run: |
+          # Determine the correct tag suffixes based on the event type
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            # PR build - use the commit SHA for more predictable references
+            AMD64_TAG="cuda-amd64-sha-${{ github.sha }}"
+            ARM64_TAG="cuda-arm64-sha-${{ github.sha }}"
+            TARGET_TAG="cuda-sha-${{ github.sha }}"
+          elif [[ "${{ startsWith(github.ref, 'refs/tags/') }}" == "true" ]]; then
+            # Tag build - use tag version in the name
+            AMD64_TAG="cuda-amd64-${{ github.ref_name }}"
+            ARM64_TAG="cuda-arm64-${{ github.ref_name }}"
+            TARGET_TAG="cuda-${{ github.ref_name }}"
+          else
+            # Main branch build - use simple arch tags
+            AMD64_TAG="cuda-amd64"
+            ARM64_TAG="cuda-arm64"
+            TARGET_TAG="cuda"
+          fi
+
+          # Create the manifest with the correct tag names
+          echo "Creating CUDA manifest using $AMD64_TAG and $ARM64_TAG"
+          docker buildx imagetools create --tag ghcr.io/${{ github.repository }}:${TARGET_TAG} \
+            ghcr.io/${{ github.repository }}:${AMD64_TAG} \
+            ghcr.io/${{ github.repository }}:${ARM64_TAG}

Dockerfile

@@ -1,19 +1,8 @@
-# Use the jupyter/minimal-notebook as the base image
-FROM quay.io/jupyter/minimal-notebook:latest
+# Stage 1: Build environment
+FROM quay.io/jupyter/minimal-notebook:latest AS builder
 
-# Metadata labels
-LABEL org.opencontainers.image.title="Python Tutorial"
-LABEL org.opencontainers.image.description="A containerized Python tutorial environment with Jupyter Lab."
-LABEL org.opencontainers.image.authors="Empa Scientific IT <scientificit@empa.ch>"
-LABEL org.opencontainers.image.url="https://github.com/empa-scientific-it/python-tutorial"
-LABEL org.opencontainers.image.source="https://github.com/empa-scientific-it/python-tutorial"
-LABEL org.opencontainers.image.version="1.0.0"
-LABEL org.opencontainers.image.licenses="MIT"
-
-# Set environment variables for the tutorial and repository
-ENV BASENAME="python-tutorial"
-ENV REPO=${HOME}/${BASENAME}
-ENV IPYTHONDIR="${HOME}/.ipython"
+# Define build argument for PyTorch variant (cpu or cuda)
+ARG PYTORCH_VARIANT=cpu
 
 # Switch to root user to install additional dependencies
 USER root
@@ -33,16 +22,61 @@ USER ${NB_UID}
 # Set up the Conda environment
 COPY docker/environment.yml /tmp/environment.yml
 RUN mamba env update -n base -f /tmp/environment.yml && \
+    # Force remove any existing PyTorch installations first
+    pip uninstall -y torch torchvision && \
+    # Install PyTorch packages without cache - conditionally based on variant
+    if [ "$PYTORCH_VARIANT" = "cpu" ]; then \
+        echo "Installing CPU-only PyTorch" && \
+        pip install --no-cache-dir --force-reinstall torch torchvision --index-url https://download.pytorch.org/whl/cpu; \
+    else \
+        echo "Installing CUDA-enabled PyTorch" && \
+        pip install --no-cache-dir --force-reinstall torch torchvision; \
+    fi && \
+    # Clean up all package caches to reduce image size
     mamba clean --all -f -y && \
+    # Remove pip cache
+    rm -rf ~/.cache/pip && \
     fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 
-# Prepare IPython configuration (move earlier in the build)
-RUN mkdir -p ${HOME}/.ipython/profile_default
+# Stage 2: Runtime environment - creates a lighter final image
+FROM quay.io/jupyter/minimal-notebook:latest
+
+# Inherit build argument for image labeling
+ARG PYTORCH_VARIANT=cpu
+
+# Metadata labels
+LABEL org.opencontainers.image.title="Python Tutorial"
+LABEL org.opencontainers.image.description="A containerized Python tutorial environment with Jupyter Lab."
+LABEL org.opencontainers.image.authors="Empa Scientific IT <scientificit@empa.ch>"
+LABEL org.opencontainers.image.url="https://github.com/empa-scientific-it/python-tutorial"
+LABEL org.opencontainers.image.source="https://github.com/empa-scientific-it/python-tutorial"
+LABEL org.opencontainers.image.version="1.0.0"
+LABEL org.opencontainers.image.licenses="MIT"
+LABEL org.opencontainers.image.variant="pytorch-${PYTORCH_VARIANT}"
+
+# Switch to root user to install minimal dependencies
+USER root
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    libgl1 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Switch back to the default notebook user
+USER ${NB_UID}
+
+# Copy the conda environment from the builder stage
+COPY --from=builder ${CONDA_DIR} ${CONDA_DIR}
+
+# Copy home directory with configurations
+COPY --from=builder --chown=${NB_UID}:${NB_GID} /home/${NB_USER} /home/${NB_USER}
+
+# Prepare IPython configuration
 COPY --chown=${NB_UID}:${NB_GID} binder/ipython_config.py ${HOME}/.ipython/profile_default/
 
-# Set the working directory to the repository
-WORKDIR ${REPO}
+# Set the working directory to user's home (repository will be cloned here by Renku)
+WORKDIR /home/${NB_USER}
 
 # Use the default ENTRYPOINT from the base image to start Jupyter Lab
 ENTRYPOINT ["tini", "-g", "--", "start.sh"]

README.md

@@ -48,7 +48,7 @@ You should now create a new environment with `conda`:
 conda env create -f binder/environment.yml
 ```
 
-> **Warning**
+> [!WARNING]
 >
 > If you are on Windows and using Command Prompt or the PowerShell, please make sure to adjust the paths in the commands above accordingly.
 
@@ -74,7 +74,7 @@ jupyter lab
 
 ### 2. With Docker
 
-> **Note**
+> [!NOTE]
 >
 > The following instructions are for Windows. With minor changes, the steps work on macOS or Linux as well.
 
@@ -84,28 +84,44 @@ jupyter lab
 
 3. Open PowerShell: Once Docker Desktop is installed, open PowerShell on your Windows machine. You can do this by pressing the "Windows" key and typing "PowerShell" in the search bar.
 
-4. Pull the Docker image: In PowerShell, run the following command to pull the "empascientificit/python-tutorial" Docker image:
+4. Pull the Docker image: In PowerShell, run the following command to pull the Docker image:
 
 ```console
 docker pull ghcr.io/empa-scientific-it/python-tutorial:latest
 ```
 
+> [!NOTE]
+>
+> The `latest` tag points to the CPU-only variant of the image, which is optimized for size and compatibility. If you have a CUDA-compatible GPU and want to use GPU acceleration for PyTorch operations, you can use the CUDA-enabled variant by replacing `latest` with `cuda`:
+>
+> ```console
+> docker pull ghcr.io/empa-scientific-it/python-tutorial:cuda
+> ```
+
+> [!IMPORTANT]
+>
+> Using the CUDA variant requires a NVIDIA GPU with compatible drivers properly installed and configured for Docker. See [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html) for setup instructions.
+
 5. Run the Docker container: Once the image is downloaded, run the following command to start a Docker container from the image:
 
 ```console
 docker run -p 8888:8888 --name python_tutorial -v /path/to/python-tutorial:/home/jovyan/python-tutorial ghcr.io/empa-scientific-it/python-tutorial:latest jupyter lab --ip 0.0.0.0 --no-browser
 ```
 
+> [!NOTE]
+>
+> If you pulled the CUDA variant, replace `:latest` with `:cuda` in the command above.
+
 Replace `/path/to/python-tutorial` with the path to the folder you created in step 2, for example `C:/Users/yourusername/Desktop/python-tutorial`.
 
-> **Note**
+> [!NOTE]
 >
-> The above command will **mirror** the content of your local folder (e.g., `C:/Users/yourusername/Desktop/python-tutorial`) to the `work/` folder **inside the container**. In this way, every file or folder you copy or create into `work/` will be saved on your machine, and will remain there **even if you stop Docker**.
+> The above command will **mirror** the content of your local folder (e.g., `C:/Users/yourusername/Desktop/python-tutorial`) to the `~/python-tutorial` folder **inside the container**. In this way, every file or folder you copy or create into `~/python-tutorial` will be saved on your machine, and will remain there **even if you stop Docker**.
 
 6. Access the Jupyter Notebook: Open a web browser and navigate to `http://localhost:8888/lab`. You should see the Jupyter Notebook interface. Enter the token provided in the PowerShell console to access the notebook. Alternatively, you can directly click on the link that appears in the PowerShell after the container has started.
 
 You can now use the Jupyter in the Docker container to run the python-tutorial. When you're done, you can stop the container by pressing `Ctrl+C` in the PowerShell console.
 
-> **Note**
+> [!NOTE]
 >
 > If you want to restart the container, you can simply run the command `docker container start python_tutorial`.

docker/environment.yml

@@ -26,8 +26,5 @@ dependencies:
       - python-dotenv
       - pillow
       - opencv-python
-      - torch
-      - torchaudio
-      - torchvision
       - albumentations
       - grad-cam