Update base (venv)

.github/workflows/docker-build.yml

@@ -8,9 +8,9 @@ on:
 env:
   UBUNTU_VERSION: 22.04
   BUILDX_NO_DEFAULT_ATTESTATIONS: 1
-  LATEST_CUDA: "2.2.1-py3.10-cuda-11.8.0-runtime-22.04"
-  LATEST_ROCM: "2.2.1-py3.10-rocm-5.7-runtime-22.04"
-  LATEST_CPU: "2.2.1-py3.10-cpu-22.04"
+  LATEST_CUDA: "2.3.1-py3.10-v2-cuda-11.8.0-runtime-22.04"
+  LATEST_ROCM: "2.3.1-py3.10-v2-rocm-5.7-runtime-22.04"
+  LATEST_CPU: "2.3.1-py3.10-v2-cpu-22.04"
 
 jobs:
   cpu-base:
@@ -23,12 +23,8 @@ jobs:
           - "3.11"
           - "3.12"
         pytorch:
-          - "2.1.2"
-          - "2.2.0"
           - "2.2.1"
-        exclude:
-          - python: "3.12"
-            pytorch: "2.1.2"
+          - "2.3.1"
     steps:
       -
         name: Free Space
@@ -65,7 +61,7 @@ jobs:
         name: Set tags
         run: |
           img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          ver_tag="${{ matrix.pytorch }}-py${{ matrix.python }}-cpu-${{ env.UBUNTU_VERSION }}"
+          ver_tag="${{ matrix.pytorch }}-py${{ matrix.python }}-v2-cpu-${{ env.UBUNTU_VERSION }}"
           
           if [[ $ver_tag == ${{ env.LATEST_CPU }} ]]; then
               TAGS="${img_path}:latest-cpu, ${img_path}:$ver_tag"
@@ -79,7 +75,7 @@ jobs:
         with:
           context: build
           build-args: |
-            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.python }}-cpu-${{ env.UBUNTU_VERSION }}
+            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.python }}-v2-cpu-${{ env.UBUNTU_VERSION }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
           push: true
           # Avoids unknown/unknown architecture and extra metadata
@@ -96,19 +92,15 @@ jobs:
           - "3.11"
           - "3.12"
         pytorch:
-          - "2.1.2"
-          - "2.2.0"
           - "2.2.1"
+          - "2.3.1"
         cuda:
           - "11.8.0"
-          - "12.1.0"
+          - "12.1.1"
         level:
           - "base"
           - "runtime"
           - "devel"
-        exclude:
-          - python: "3.12"
-            pytorch: "2.1.2"
 
     steps:
       -
@@ -146,7 +138,7 @@ jobs:
         name: Set tags
         run: |
           img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          ver_tag="${{ matrix.pytorch }}-py${{ matrix.python }}-cuda-${{ matrix.cuda }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}"
+          ver_tag="${{ matrix.pytorch }}-py${{ matrix.python }}-v2-cuda-${{ matrix.cuda }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}"
           
           if [[ $ver_tag == ${{ env.LATEST_CUDA }} ]]; then
               TAGS="${img_path}:latest, ${img_path}:latest-cuda, ${img_path}:$ver_tag"
@@ -160,7 +152,7 @@ jobs:
         with:
           context: build
           build-args: |
-            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.python }}-cuda-${{ matrix.cuda }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}
+            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.python }}-v2-cuda-${{ matrix.cuda }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
           push: true
           provenance: false
@@ -176,23 +168,18 @@ jobs:
           - "3.11"
           - "3.12"
         pytorch:
-          - "2.1.2"
-          - "2.2.0"
           - "2.2.1"
+          - "2.3.1"
         rocm:
-          - "5.6"
           - "5.7"
+          - "6.0.2"
         level:
           - "runtime"
         exclude:
-          - pytorch: "2.2.1"
-            rocm: "5.6"
-          - pytorch: "2.2.0"
-            rocm: "5.6"
-          - pytorch: "2.1.2"
+          - pytorch: "2.3.1"
             rocm: "5.7"
-          - python: "3.12"
-            rocm: "5.6"
+          - pytorch: "2.2.1"
+            rocm: "6.0.2"
 
     steps:
       -
@@ -230,7 +217,7 @@ jobs:
         name: Set tags
         run: |
           img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          ver_tag="${{ matrix.pytorch }}-py${{ matrix.python }}-rocm-${{ matrix.rocm }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}"
+          ver_tag="${{ matrix.pytorch }}-py${{ matrix.python }}-v2-rocm-${{ matrix.rocm }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}"
           
           if [[ $ver_tag == ${{ env.LATEST_ROCM }} ]]; then
               TAGS="${img_path}:latest-rocm, ${img_path}:$ver_tag"
@@ -244,7 +231,7 @@ jobs:
         with:
           context: build
           build-args: |
-            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.python }}-rocm-${{ matrix.rocm }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}
+            IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.python }}-v2-rocm-${{ matrix.rocm }}-${{ matrix.level }}-${{ env.UBUNTU_VERSION }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
           push: true
           provenance: false

.gitignore

@@ -1,5 +1,4 @@
 workspace
-build/COPY_ROOT_EXTRA/
 config/authorized_keys
 config/rclone
 .env

build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/amd.sh

@@ -0,0 +1,16 @@
+#!/bin/false
+
+build_amd_main() {
+    build_amd_install_torch
+    build_common_run_tests
+}
+
+build_amd_install_torch() {
+    "$PYTHON_VENV_PIP" install --no-cache-dir \
+        torch==${PYTORCH_VERSION} \
+        torchvision \
+        torchaudio \
+        --extra-index-url=https://download.pytorch.org/whl/rocm${ROCM_VERSION}
+}
+
+build_amd_main "$@"

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/clean.sh → build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/clean.sh

@@ -2,7 +2,6 @@
 
 # Tidy up and keep image small
 apt-get clean -y
-micromamba clean -ay
 
 fix-permissions.sh -o container
 

build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/common.sh

@@ -0,0 +1,17 @@
+#!/bin/false
+
+# For logic common to cpu, cuda & rocm
+
+build_common_main() {
+    :
+}
+
+build_common_run_tests() {
+    installed_pytorch_version=$("$PYTHON_VENV_PYTHON" -c "import torch; print(torch.__version__)")
+    if [[ "$installed_pytorch_version" != "$PYTORCH_VERSION"* ]]; then
+        echo "Expected PyTorch ${PYTORCH_VERSION} but found ${installed_pytorch_version}\n"
+        exit 1
+    fi
+}
+
+build_common_main "$@"

build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/cpu.sh

@@ -0,0 +1,16 @@
+#!/bin/false
+
+build_cpu_main() {
+    build_cpu_install_torch
+    build_common_run_tests
+}
+
+build_cpu_install_deps() {
+    "$PYTHON_VENV_PIP" install --no-cache-dir \
+        torch==${PYTORCH_VERSION} \
+        torchvision \
+        torchaudio \
+        --extra-index-url=https://download.pytorch.org/whl/cpu
+}
+
+build_cpu_main "$@"

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh → build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/init.sh

@@ -4,8 +4,13 @@
 set -eo pipefail
 umask 002
 
+source /opt/ai-dock/etc/environment.sh
 source /opt/ai-dock/bin/build/layer0/common.sh
 
+PYTHON_VENV="$VENV_DIR/$PYTHON_DEFAULT_VENV"
+PYTHON_VENV_PYTHON="$PYTHON_VENV/bin/python"
+PYTHON_VENV_PIP="$PYTHON_VENV/bin/pip"
+
 if [[ "$XPU_TARGET" == "NVIDIA_GPU" ]]; then
     source /opt/ai-dock/bin/build/layer0/nvidia.sh
 elif [[ "$XPU_TARGET" == "AMD_GPU" ]]; then
@@ -17,6 +22,4 @@ else
     exit 1
 fi
 
-$MAMBA_DEFAULT_RUN python /opt/ai-dock/tests/assert-torch-version.py
-
 source /opt/ai-dock/bin/build/layer0/clean.sh

build/COPY_ROOT_0/opt/ai-dock/bin/build/layer0/nvidia.sh

@@ -0,0 +1,27 @@
+#!/bin/false
+
+build_nvidia_main() {
+    build_nvidia_install_torch
+    build_common_run_tests
+    build_nvidia_run_tests
+}
+
+build_nvidia_install_torch() {
+    short_cuda_version="cu$(cut -d '.' -f 1,2 <<< "${CUDA_VERSION}" | tr -d '.')"
+    "$PYTHON_VENV_PIP" install --no-cache-dir \
+        nvidia-ml-py3 \
+        torch==${PYTORCH_VERSION} \
+        torchvision \
+        torchaudio \
+        --extra-index-url=https://download.pytorch.org/whl/$short_cuda_version
+}
+
+build_nvidia_run_tests() {
+    installed_pytorch_cuda_version=$("$PYTHON_VENV_PYTHON" -c "import torch; print(torch.version.cuda)")
+    if [[ "$CUDA_VERSION" != "$installed_pytorch_cuda"* ]]; then
+        echo "Expected PyTorch CUDA ${CUDA_VERSION} but found ${installed_pytorch_cuda}\n"
+        exit 1
+    fi
+}
+
+build_nvidia_main "$@"

build/COPY_ROOT_99/opt/ai-dock/bin/build/layer1/init.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Override this file to add extras to your build
+
+## START
+#umask 002
+
+
+
+## END
+#fix-permissions.sh -o container

build/Dockerfile

@@ -1,6 +1,6 @@
 # For build automation - Allows building from any ai-dock base image
 # Use a *cuda*base* image as default because pytorch brings the libs
-ARG IMAGE_BASE="ghcr.io/ai-dock/python:3.10-cuda-11.8.0-base-22.04"
+ARG IMAGE_BASE="ghcr.io/ai-dock/python:3.10-v2-cuda-11.8.0-base-22.04"
 FROM ${IMAGE_BASE}
 
 LABEL org.opencontainers.image.source https://github.com/ai-dock/pytorch
@@ -13,7 +13,7 @@ ENV IMAGE_SLUG="pytorch"
 # Note: the default mamba environment is set by ai-dock/python
 
 # Copy early so we can use scripts in the build - Changes to these files will invalidate the cache and cause a rebuild.
-COPY --chown=0:1111 ./COPY_ROOT/ /
+COPY --chown=0:1111 ./COPY_ROOT_0/ /
 
 # Use build scripts to ensure we can build all targets from one Dockerfile in a single layer.
 # Don't put anything heavy in here - We can use multi-stage building above if necessary.
@@ -22,9 +22,5 @@ ENV PYTORCH_VERSION=${PYTORCH_VERSION}
 
 RUN set -eo pipefail && /opt/ai-dock/bin/build/layer0/init.sh | tee /var/log/build.log
 
-# Copy overrides and new files into a final layer for fast rebuilds. Uncomment below
-#COPY --chown=0:1111 ./COPY_ROOT_EXTRA/ /
-#RUN set -eo pipefail && /opt/ai-dock/bin/build/layer1/init.sh | tee -a /var/log/build.log
-
 # Keep init.sh as-is and place additional logic in /opt/ai-dock/bin/preflight.sh
 CMD ["init.sh"]

docker-compose.yaml

@@ -6,12 +6,12 @@ services:
     build:
       context: ./build
       args:
-        IMAGE_BASE: ${IMAGE_BASE:-ghcr.io/ai-dock/python:3.10-cuda-11.8.0-runtime-22.04}
-        PYTORCH_VERSION: ${PYTORCH_VERSION:-2.1.1}
+        IMAGE_BASE: ${IMAGE_BASE:-ghcr.io/ai-dock/python:3.10-v2-cuda-11.8.0-runtime-22.04}
+        PYTORCH_VERSION: ${PYTORCH_VERSION:-2.3.0}
       tags:
-        - "ghcr.io/ai-dock/pytorch:${IMAGE_TAG:-2.1.1-py3.10-cuda-11.8.0-runtime-22.04}"
+        - "ghcr.io/ai-dock/pytorch:${IMAGE_TAG:-2.3.0-py3.10-v2-cuda-11.8.0-runtime-22.04}"
 
-    image: ghcr.io/ai-dock/pytorch:${IMAGE_TAG:-2.1.1-py3.10-cuda-11.8.0-runtime-22.04}
+    image: ghcr.io/ai-dock/pytorch:${IMAGE_TAG:-2.3.0-py3.10-v2-cuda-11.8.0-runtime-22.04}
 
     ## For Nvidia GPU's - You probably want to uncomment this
     #deploy: