feat: multi-Python worker images with startup version check (AE-2827) (#89)

* feat: multi-Python worker images with startup version check (AE-2827)

Add Python 3.10 and 3.11 support to GPU worker images via side-by-side
torch install in the existing runpod/pytorch base. 3.12 keeps the fast
path (torch pre-installed) to avoid the ~7 GB reinstall cost on hot
deployments; 3.10/3.11 images pay that cost once per cold start per DC.

Sibling to flash#322 which landed the SDK-level plumbing. Tags follow
the same ``py${VERSION}-${TAG}`` scheme already in use for CPU images.

- Dockerfile / Dockerfile-lb (GPU): accept PYTHON_VERSION build arg;
  install torch from download.pytorch.org/whl/cu128 and repoint
  /usr/local/bin/python for non-3.12 targets; validate interpreter
  matches the arg during build.
- Dockerfile-cpu / Dockerfile-lb-cpu (CPU): surface PYTHON_VERSION at
  runtime via FLASH_PYTHON_VERSION env so the worker's startup check
  can read it.
- src/version.py: new ``assert_python_version_matches_image`` — raises
  PythonVersionMismatchError at handler boot when ``sys.version_info``
  disagrees with the image's stamped FLASH_PYTHON_VERSION. Caught
  before user code runs; skipped when the env var is unset (local dev).
- src/handler.py / src/lb_handler.py: call the assertion immediately
  after logging setup, before ``maybe_unpack()`` and handler import.
- tests/unit/test_version.py: 4 new cases covering env-unset skip,
  match, mismatch raise, and message contents.
- tests/unit/test_lb_handler.py: extend the mocked ``version`` module
  with ``assert_python_version_matches_image`` so fresh-import tests
  don't break.
- .github/workflows/ci.yml: expand CI to build GPU and LB images
  across {3.10, 3.11, 3.12}; align prod CPU and LB-CPU default to
  3.12 (matches flash's DEFAULT_PYTHON_VERSION).

* fix(dockerfile): bootstrap pip via get-pip.py for non-3.12 GPU builds

Ubuntu 22.04's system python3.10 has ensurepip disabled by Debian
policy, which broke the side-by-side torch install for 3.10 GPU images
(CI: docker-test-gpu (3.10), docker-test-lb (3.10)). python3.11 is a
separate interpreter without the disable, so only 3.10 was affected.

Use urllib+get-pip.py instead of ensurepip — works for any interpreter
regardless of distro patching, and urllib is stdlib so no curl dep.

Also corrects the outdated deadsnakes comment on both Dockerfiles: the
runpod/pytorch base image layers alt-Python 3.11/3.12 on top of the
system 3.10, not via deadsnakes.

Author: deanq

.github/workflows/ci.yml

@@ -73,6 +73,9 @@ jobs:
   docker-test:
     runs-on: ubuntu-latest
     if: github.event_name != 'pull_request' || github.head_ref != 'release-please--branches--main'
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -89,19 +92,22 @@ jobs:
           push: false
           tags: flash-cpu:test
           build-args: |
-            PYTHON_VERSION=3.11
-          cache-from: type=gha
-          cache-to: type=gha,mode=max
+            PYTHON_VERSION=${{ matrix.python-version }}
+          cache-from: type=gha,scope=cpu-test-py${{ matrix.python-version }}
+          cache-to: type=gha,mode=max,scope=cpu-test-py${{ matrix.python-version }}
           load: true
 
       - name: Test CPU handler execution in Docker environment
         run: |
-          echo "Testing CPU handler in Docker environment..."
+          echo "Testing CPU handler (Python ${{ matrix.python-version }})..."
           docker run --rm flash-cpu:test ./test-handler.sh
 
   docker-test-lb-cpu:
     runs-on: ubuntu-latest
     if: github.event_name != 'pull_request' || github.head_ref != 'release-please--branches--main'
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -118,24 +124,97 @@ jobs:
           push: false
           tags: flash-lb-cpu:test
           build-args: |
-            PYTHON_VERSION=3.11
-          cache-from: type=gha
-          cache-to: type=gha,mode=max
+            PYTHON_VERSION=${{ matrix.python-version }}
+          cache-from: type=gha,scope=lb-cpu-test-py${{ matrix.python-version }}
+          cache-to: type=gha,mode=max,scope=lb-cpu-test-py${{ matrix.python-version }}
           load: true
 
       - name: Test LB handler execution in Docker environment
         run: |
-          echo "Testing LB handler in Docker environment..."
+          echo "Testing LB handler (Python ${{ matrix.python-version }})..."
           docker run --rm flash-lb-cpu:test ./test-lb-handler.sh
 
+  docker-test-gpu:
+    runs-on: ubuntu-latest
+    if: github.event_name != 'pull_request' || github.head_ref != 'release-please--branches--main'
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - name: Clear space
+        run: |
+          rm -rf /usr/share/dotnet /opt/ghc /usr/local/share/boost "$AGENT_TOOLSDIRECTORY"
+          docker system prune -af
+          df -h
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build GPU Docker image
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: ./Dockerfile
+          platforms: linux/amd64
+          push: false
+          tags: flash-gpu:test
+          build-args: |
+            PYTHON_VERSION=${{ matrix.python-version }}
+          cache-from: type=gha,scope=gpu-test-py${{ matrix.python-version }}
+          cache-to: type=gha,mode=max,scope=gpu-test-py${{ matrix.python-version }}
+
+  docker-test-lb:
+    runs-on: ubuntu-latest
+    if: github.event_name != 'pull_request' || github.head_ref != 'release-please--branches--main'
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - name: Clear space
+        run: |
+          rm -rf /usr/share/dotnet /opt/ghc /usr/local/share/boost "$AGENT_TOOLSDIRECTORY"
+          docker system prune -af
+          df -h
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build GPU Load Balancer Docker image
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: ./Dockerfile-lb
+          platforms: linux/amd64
+          push: false
+          tags: flash-lb:test
+          build-args: |
+            PYTHON_VERSION=${{ matrix.python-version }}
+          cache-from: type=gha,scope=lb-test-py${{ matrix.python-version }}
+          cache-to: type=gha,mode=max,scope=lb-test-py${{ matrix.python-version }}
+
   docker-validation:
     runs-on: ubuntu-latest
-    needs: [test, lint, docker-test, docker-test-lb-cpu]
+    needs: [test, lint, docker-test, docker-test-lb-cpu, docker-test-gpu, docker-test-lb]
     if: always()
     steps:
       - name: Check all jobs succeeded
         run: |
-          results=("${{ needs.test.result }}" "${{ needs.lint.result }}" "${{ needs.docker-test.result }}" "${{ needs.docker-test-lb-cpu.result }}")
+          results=(
+            "${{ needs.test.result }}"
+            "${{ needs.lint.result }}"
+            "${{ needs.docker-test.result }}"
+            "${{ needs.docker-test-lb-cpu.result }}"
+            "${{ needs.docker-test-gpu.result }}"
+            "${{ needs.docker-test-lb.result }}"
+          )
           for result in "${results[@]}"; do
             if [[ "$result" != "success" && "$result" != "skipped" ]]; then
               echo "One or more quality checks failed (got: $result)"
@@ -168,8 +247,13 @@ jobs:
     needs: [release]
     if: needs.release.outputs.release_created
     strategy:
+      fail-fast: false
       matrix:
         include:
+          - python-version: "3.10"
+            is-default: false
+          - python-version: "3.11"
+            is-default: false
           - python-version: "3.12"
             is-default: true
     steps:
@@ -226,22 +310,25 @@ jobs:
           platforms: linux/amd64
           push: true
           tags: ${{ steps.tags.outputs.tags }}
-          cache-from: type=gha,scope=gpu
-          cache-to: type=gha,mode=max,scope=gpu
+          build-args: |
+            PYTHON_VERSION=${{ matrix.python-version }}
+          cache-from: type=gha,scope=gpu-py${{ matrix.python-version }}
+          cache-to: type=gha,mode=max,scope=gpu-py${{ matrix.python-version }}
 
   docker-prod-cpu:
     runs-on: ubuntu-latest
     needs: [release]
     if: needs.release.outputs.release_created
     strategy:
+      fail-fast: false
       matrix:
         include:
           - python-version: "3.10"
             is-default: false
           - python-version: "3.11"
-            is-default: true
-          - python-version: "3.12"
             is-default: false
+          - python-version: "3.12"
+            is-default: true
     steps:
       - name: Clear Space
         run: |
@@ -306,8 +393,13 @@ jobs:
     needs: [release]
     if: needs.release.outputs.release_created
     strategy:
+      fail-fast: false
       matrix:
         include:
+          - python-version: "3.10"
+            is-default: false
+          - python-version: "3.11"
+            is-default: false
           - python-version: "3.12"
             is-default: true
     steps:
@@ -364,22 +456,25 @@ jobs:
           platforms: linux/amd64
           push: true
           tags: ${{ steps.tags.outputs.tags }}
-          cache-from: type=gha,scope=lb
-          cache-to: type=gha,mode=max,scope=lb
+          build-args: |
+            PYTHON_VERSION=${{ matrix.python-version }}
+          cache-from: type=gha,scope=lb-py${{ matrix.python-version }}
+          cache-to: type=gha,mode=max,scope=lb-py${{ matrix.python-version }}
 
   docker-prod-lb-cpu:
     runs-on: ubuntu-latest
     needs: [release]
     if: needs.release.outputs.release_created
     strategy:
+      fail-fast: false
       matrix:
         include:
           - python-version: "3.10"
             is-default: false
           - python-version: "3.11"
-            is-default: true
-          - python-version: "3.12"
             is-default: false
+          - python-version: "3.12"
+            is-default: true
     steps:
       - name: Clear Space
         run: |

Dockerfile

@@ -1,15 +1,37 @@
-# Base image provides Python 3.12 (from runpod/pytorch:1.0.3-cu1281-torch291-ubuntu2204)
+# Base image (runpod/pytorch:ubuntu2204) ships Ubuntu 22.04 system python3.10
+# plus alt-Python interpreters for 3.11/3.12 with torch pre-installed on 3.12.
+# For non-3.12 targets we reinstall torch from the CUDA 12.8 wheel index
+# (~7 GB overhead) and repoint /usr/local/bin/python so the worker CMD picks
+# up the correct interpreter.
 FROM runpod/pytorch:1.0.3-cu1281-torch291-ubuntu2204
 
-# Use the base image's Python as-is to preserve pre-installed packages (torch, cuda libs).
-# The pytorch base image provides its own Python with torch already installed.
-# Symlinking to /usr/bin/python3.X would switch to a bare system Python without torch.
-# Validate that the base image provides the expected Python version.
-ARG EXPECTED_PYTHON_VERSION=3.12
-RUN python --version && \
-    actual=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") && \
-    if [ "$actual" != "$EXPECTED_PYTHON_VERSION" ]; then \
-      echo "ERROR: Expected Python $EXPECTED_PYTHON_VERSION but base image provides $actual" && exit 1; \
+# Target Python version for the worker runtime.
+ARG PYTHON_VERSION=3.12
+ARG TORCH_VERSION=2.9.1+cu128
+ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
+
+# Expose the target version to the running worker for startup validation.
+ENV FLASH_PYTHON_VERSION=${PYTHON_VERSION}
+
+# Validate the base image provides the requested interpreter and activate it.
+# For non-3.12 targets, install torch for the selected Python and repoint
+# /usr/local/bin/python and python3 so downstream `python` invocations use it.
+# For 3.12 we keep the base image's python/torch untouched to avoid the
+# ~7 GB reinstall cost.
+#
+# pip bootstrap: Ubuntu 22.04's system python3.10 has ensurepip disabled by
+# Debian policy, so we install pip via get-pip.py (works for any interpreter
+# regardless of distro patching). urllib is stdlib, avoiding a curl dependency.
+RUN python${PYTHON_VERSION} --version \
+ && if [ "${PYTHON_VERSION}" != "3.12" ]; then \
+      python${PYTHON_VERSION} -c "import urllib.request; urllib.request.urlretrieve('https://bootstrap.pypa.io/get-pip.py', '/tmp/get-pip.py')" \
+      && python${PYTHON_VERSION} /tmp/get-pip.py --no-cache-dir \
+      && rm -f /tmp/get-pip.py \
+      && python${PYTHON_VERSION} -m pip install --no-cache-dir \
+           --index-url ${TORCH_INDEX_URL} \
+           "torch==${TORCH_VERSION}" \
+      && ln -sf "$(which python${PYTHON_VERSION})" /usr/local/bin/python \
+      && ln -sf "$(which python${PYTHON_VERSION})" /usr/local/bin/python3; \
     fi
 
 WORKDIR /app
@@ -41,20 +63,21 @@ RUN DEBIAN_FRONTEND=noninteractive apt-get update && apt-get install -y --no-ins
  && rm -rf /var/lib/apt/lists/*
 
 # Copy app code and install dependencies
-# Use --python to target the base image's Python (preserves torch in its site-packages)
+# Use --python to target the active interpreter (preserves torch in its site-packages)
 COPY README.md pyproject.toml uv.lock ./
 COPY src/ ./
 RUN uv export --format requirements-txt --no-dev --no-hashes > requirements.txt \
  && uv pip install --python $(which python) --break-system-packages -r requirements.txt
 
-# Install numpy for the base image's Python version.
+# Install numpy for the active Python version.
 # The runpod/pytorch image ships torch but not numpy. Flash build excludes numpy
 # from tarballs (BASE_IMAGE_PACKAGES) to save tarball space (~30 MB), so numpy
 # must be provided here in the base image.
 RUN python -m pip install --no-cache-dir numpy
 
-# Verify torch and numpy are available from the base image
-RUN python -c "import torch; print(f'torch {torch.__version__} CUDA {torch.cuda.is_available()}')" \
+# Verify torch, numpy, and the expected Python version are available.
+RUN python -c "import sys; actual = f'{sys.version_info.major}.{sys.version_info.minor}'; expected = '${PYTHON_VERSION}'; assert actual == expected, f'Expected Python {expected}, got {actual}'; print(f'Python {actual} OK')" \
+ && python -c "import torch; print(f'torch {torch.__version__} CUDA {torch.cuda.is_available()}')" \
  && python -c "import numpy; print(f'numpy {numpy.__version__}')"
 
 CMD ["python", "handler.py"]

Dockerfile-cpu

@@ -1,6 +1,11 @@
 ARG PYTHON_VERSION=3.12
 FROM python:${PYTHON_VERSION}-slim
 
+# Re-declare after FROM so the value is visible in this build stage, and
+# expose it at runtime for the worker's startup version check.
+ARG PYTHON_VERSION
+ENV FLASH_PYTHON_VERSION=${PYTHON_VERSION}
+
 WORKDIR /app
 
 # Prevent interactive prompts during package installation

Dockerfile-lb

@@ -1,13 +1,37 @@
-# Base image provides Python 3.12 (from runpod/pytorch:1.0.3-cu1281-torch291-ubuntu2204)
+# Base image (runpod/pytorch:ubuntu2204) ships Ubuntu 22.04 system python3.10
+# plus alt-Python interpreters for 3.11/3.12 with torch pre-installed on 3.12.
+# For non-3.12 targets we reinstall torch from the CUDA 12.8 wheel index
+# (~7 GB overhead) and repoint /usr/local/bin/python so the worker CMD picks
+# up the correct interpreter.
 FROM runpod/pytorch:1.0.3-cu1281-torch291-ubuntu2204
 
-# Use the base image's Python as-is to preserve pre-installed packages (torch, cuda libs).
-# Validate that the base image provides the expected Python version.
-ARG EXPECTED_PYTHON_VERSION=3.12
-RUN python --version && \
-    actual=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") && \
-    if [ "$actual" != "$EXPECTED_PYTHON_VERSION" ]; then \
-      echo "ERROR: Expected Python $EXPECTED_PYTHON_VERSION but base image provides $actual" && exit 1; \
+# Target Python version for the worker runtime.
+ARG PYTHON_VERSION=3.12
+ARG TORCH_VERSION=2.9.1+cu128
+ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
+
+# Expose the target version to the running worker for startup validation.
+ENV FLASH_PYTHON_VERSION=${PYTHON_VERSION}
+
+# Validate the base image provides the requested interpreter and activate it.
+# For non-3.12 targets, install torch for the selected Python and repoint
+# /usr/local/bin/python and python3 so downstream `python` invocations use it.
+# For 3.12 we keep the base image's python/torch untouched to avoid the
+# ~7 GB reinstall cost.
+#
+# pip bootstrap: Ubuntu 22.04's system python3.10 has ensurepip disabled by
+# Debian policy, so we install pip via get-pip.py (works for any interpreter
+# regardless of distro patching). urllib is stdlib, avoiding a curl dependency.
+RUN python${PYTHON_VERSION} --version \
+ && if [ "${PYTHON_VERSION}" != "3.12" ]; then \
+      python${PYTHON_VERSION} -c "import urllib.request; urllib.request.urlretrieve('https://bootstrap.pypa.io/get-pip.py', '/tmp/get-pip.py')" \
+      && python${PYTHON_VERSION} /tmp/get-pip.py --no-cache-dir \
+      && rm -f /tmp/get-pip.py \
+      && python${PYTHON_VERSION} -m pip install --no-cache-dir \
+           --index-url ${TORCH_INDEX_URL} \
+           "torch==${TORCH_VERSION}" \
+      && ln -sf "$(which python${PYTHON_VERSION})" /usr/local/bin/python \
+      && ln -sf "$(which python${PYTHON_VERSION})" /usr/local/bin/python3; \
     fi
 
 WORKDIR /app
@@ -39,20 +63,21 @@ RUN DEBIAN_FRONTEND=noninteractive apt-get update && apt-get install -y --no-ins
  && rm -rf /var/lib/apt/lists/*
 
 # Copy app code and install dependencies
-# Use --python to target the base image's Python (preserves torch in its site-packages)
+# Use --python to target the active interpreter (preserves torch in its site-packages)
 COPY README.md pyproject.toml uv.lock ./
 COPY src/ ./
 RUN uv export --format requirements-txt --no-dev --no-hashes > requirements.txt \
  && uv pip install --python $(which python) --break-system-packages -r requirements.txt
 
-# Install numpy for the base image's Python version.
+# Install numpy for the active Python version.
 # The runpod/pytorch image ships torch but not numpy. Flash build excludes numpy
 # from tarballs (BASE_IMAGE_PACKAGES) to save tarball space (~30 MB), so numpy
 # must be provided here in the base image.
 RUN python -m pip install --no-cache-dir numpy
 
-# Verify torch and numpy are available from the base image
-RUN python -c "import torch; print(f'torch {torch.__version__} CUDA {torch.cuda.is_available()}')" \
+# Verify torch, numpy, and the expected Python version are available.
+RUN python -c "import sys; actual = f'{sys.version_info.major}.{sys.version_info.minor}'; expected = '${PYTHON_VERSION}'; assert actual == expected, f'Expected Python {expected}, got {actual}'; print(f'Python {actual} OK')" \
+ && python -c "import torch; print(f'torch {torch.__version__} CUDA {torch.cuda.is_available()}')" \
  && python -c "import numpy; print(f'numpy {numpy.__version__}')"
 
 EXPOSE 80

Dockerfile-lb-cpu

@@ -1,6 +1,11 @@
 ARG PYTHON_VERSION=3.12
 FROM python:${PYTHON_VERSION}-slim
 
+# Re-declare after FROM so the value is visible in this build stage, and
+# expose it at runtime for the worker's startup version check.
+ARG PYTHON_VERSION
+ENV FLASH_PYTHON_VERSION=${PYTHON_VERSION}
+
 WORKDIR /app
 
 # Prevent interactive prompts during package installation