feat: Python version matrix for Docker images + error-masking fix

Dockerfile

@@ -1,5 +1,10 @@
-FROM pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
-# Python 3.12 included in this PyTorch image
+ARG PYTHON_VERSION=3.11
+ARG PYTORCH_BASE=pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
+FROM ${PYTORCH_BASE}
+
+# Validate base image Python matches requested version
+ARG PYTHON_VERSION
+RUN python -c "import sys; expected='$PYTHON_VERSION'; actual='.'.join(map(str,sys.version_info[:2])); assert actual==expected, f'Python mismatch: expected {expected}, got {actual}'"
 
 WORKDIR /app
 

Dockerfile-cpu

@@ -1,4 +1,5 @@
-FROM python:3.12-slim
+ARG PYTHON_VERSION=3.11
+FROM python:${PYTHON_VERSION}-slim
 
 WORKDIR /app
 

Dockerfile-lb

@@ -1,5 +1,10 @@
-FROM pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
-# Python 3.12 included in this PyTorch image
+ARG PYTHON_VERSION=3.11
+ARG PYTORCH_BASE=pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
+FROM ${PYTORCH_BASE}
+
+# Validate base image Python matches requested version
+ARG PYTHON_VERSION
+RUN python -c "import sys; expected='$PYTHON_VERSION'; actual='.'.join(map(str,sys.version_info[:2])); assert actual==expected, f'Python mismatch: expected {expected}, got {actual}'"
 
 WORKDIR /app
 

Dockerfile-lb-cpu

@@ -1,4 +1,5 @@
-FROM python:3.12-slim
+ARG PYTHON_VERSION=3.11
+FROM python:${PYTHON_VERSION}-slim
 
 WORKDIR /app
 

Makefile

@@ -19,6 +19,15 @@ endif
 WIP_TAG ?= wip
 MULTI_PLATFORM := linux/amd64,linux/arm64
 
+# Python version matrix for multi-version builds
+GPU_PYTHON_VERSIONS := 3.11 3.12
+CPU_PYTHON_VERSIONS := 3.10 3.11 3.12
+DEFAULT_PYTHON_VERSION := 3.11
+
+# PyTorch base image mapping per Python version (GPU only)
+PYTORCH_BASE_3.11 := pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
+PYTORCH_BASE_3.12 := pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime
+
 .PHONY: setup help
 
 # Check if 'uv' is installed
@@ -130,6 +139,184 @@ build-wip-lb-cpu: setup # Build and push LB CPU image (multi-platform)
 	-t $(IMAGE)-lb-cpu:$(WIP_TAG) \
 	. --push
 
+# Versioned Build Targets (multi-Python-version matrix)
+# GPU images: Python 3.11, 3.12 (with PyTorch base)
+# CPU images: Python 3.10, 3.11, 3.12 (python:X.Y-slim)
+# Tag format: py${VERSION}-${TAG} (e.g., runpod/flash:py3.11-local)
+
+build-gpu-versioned: setup _build-gpu-versioned # Build GPU images for all GPU Python versions
+_build-gpu-versioned:
+	@pytorch_base() { \
+		case "$$1" in \
+			3.11) echo "$(PYTORCH_BASE_3.11)";; \
+			3.12) echo "$(PYTORCH_BASE_3.12)";; \
+			*) echo "ERROR: No PyTorch base image mapped for Python $$1" >&2; exit 1;; \
+		esac; \
+	}; \
+	for pyver in $(GPU_PYTHON_VERSIONS); do \
+		base=$$(pytorch_base $$pyver); \
+		echo "Building GPU image for Python $$pyver (base: $$base)..."; \
+		docker buildx build \
+			--platform $(PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			--build-arg PYTORCH_BASE=$$base \
+			-t $(IMAGE):py$$pyver-$(TAG) \
+			. --load; \
+	done
+
+build-cpu-versioned: setup _build-cpu-versioned # Build CPU images for all CPU Python versions
+_build-cpu-versioned:
+	@for pyver in $(CPU_PYTHON_VERSIONS); do \
+		echo "Building CPU image for Python $$pyver..."; \
+		docker buildx build \
+			--platform $(PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			-f Dockerfile-cpu \
+			-t $(IMAGE)-cpu:py$$pyver-$(TAG) \
+			. --load; \
+	done
+
+build-lb-versioned: setup _build-lb-versioned # Build GPU-LB images for all GPU Python versions
+_build-lb-versioned:
+	@pytorch_base() { \
+		case "$$1" in \
+			3.11) echo "$(PYTORCH_BASE_3.11)";; \
+			3.12) echo "$(PYTORCH_BASE_3.12)";; \
+			*) echo "ERROR: No PyTorch base image mapped for Python $$1" >&2; exit 1;; \
+		esac; \
+	}; \
+	for pyver in $(GPU_PYTHON_VERSIONS); do \
+		base=$$(pytorch_base $$pyver); \
+		echo "Building GPU-LB image for Python $$pyver (base: $$base)..."; \
+		docker buildx build \
+			--platform $(PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			--build-arg PYTORCH_BASE=$$base \
+			-f Dockerfile-lb \
+			-t $(IMAGE)-lb:py$$pyver-$(TAG) \
+			. --load; \
+	done
+
+build-lb-cpu-versioned: setup _build-lb-cpu-versioned # Build CPU-LB images for all CPU Python versions
+_build-lb-cpu-versioned:
+	@for pyver in $(CPU_PYTHON_VERSIONS); do \
+		echo "Building CPU-LB image for Python $$pyver..."; \
+		docker buildx build \
+			--platform $(PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			-f Dockerfile-lb-cpu \
+			-t $(IMAGE)-lb-cpu:py$$pyver-$(TAG) \
+			. --load; \
+	done
+
+build-all-versioned: setup _build-gpu-versioned _build-cpu-versioned _build-lb-versioned _build-lb-cpu-versioned # Build all 10 versioned images (GPU+CPU, QB+LB)
+	@echo "All 10 versioned images built."
+
+# Versioned WIP Push Targets (multi-platform, requires Docker Hub push)
+# Also tags DEFAULT_PYTHON_VERSION images as latest (unversioned tag)
+
+build-wip-versioned: setup # Build and push all versioned images (multi-platform)
+	@echo "Building and pushing all versioned images with tag prefix py*-$(WIP_TAG)..."
+	@pytorch_base() { \
+		case "$$1" in \
+			3.11) echo "$(PYTORCH_BASE_3.11)";; \
+			3.12) echo "$(PYTORCH_BASE_3.12)";; \
+			*) echo "ERROR: No PyTorch base image mapped for Python $$1" >&2; exit 1;; \
+		esac; \
+	}; \
+	for pyver in $(GPU_PYTHON_VERSIONS); do \
+		base=$$(pytorch_base $$pyver); \
+		echo "Pushing GPU QB image for Python $$pyver..."; \
+		tag_args="-t $(IMAGE):py$$pyver-$(WIP_TAG)"; \
+		if [ "$$pyver" = "$(DEFAULT_PYTHON_VERSION)" ]; then \
+			tag_args="$$tag_args -t $(IMAGE):$(WIP_TAG)"; \
+		fi; \
+		docker buildx build \
+			--platform $(MULTI_PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			--build-arg PYTORCH_BASE=$$base \
+			$$tag_args \
+			. --push; \
+	done
+	@for pyver in $(CPU_PYTHON_VERSIONS); do \
+		echo "Pushing CPU QB image for Python $$pyver..."; \
+		tag_args="-t $(IMAGE)-cpu:py$$pyver-$(WIP_TAG)"; \
+		if [ "$$pyver" = "$(DEFAULT_PYTHON_VERSION)" ]; then \
+			tag_args="$$tag_args -t $(IMAGE)-cpu:$(WIP_TAG)"; \
+		fi; \
+		docker buildx build \
+			--platform $(MULTI_PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			-f Dockerfile-cpu \
+			$$tag_args \
+			. --push; \
+	done
+	@pytorch_base() { \
+		case "$$1" in \
+			3.11) echo "$(PYTORCH_BASE_3.11)";; \
+			3.12) echo "$(PYTORCH_BASE_3.12)";; \
+			*) echo "ERROR: No PyTorch base image mapped for Python $$1" >&2; exit 1;; \
+		esac; \
+	}; \
+	for pyver in $(GPU_PYTHON_VERSIONS); do \
+		base=$$(pytorch_base $$pyver); \
+		echo "Pushing GPU LB image for Python $$pyver..."; \
+		tag_args="-t $(IMAGE)-lb:py$$pyver-$(WIP_TAG)"; \
+		if [ "$$pyver" = "$(DEFAULT_PYTHON_VERSION)" ]; then \
+			tag_args="$$tag_args -t $(IMAGE)-lb:$(WIP_TAG)"; \
+		fi; \
+		docker buildx build \
+			--platform $(MULTI_PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			--build-arg PYTORCH_BASE=$$base \
+			-f Dockerfile-lb \
+			$$tag_args \
+			. --push; \
+	done
+	@for pyver in $(CPU_PYTHON_VERSIONS); do \
+		echo "Pushing CPU LB image for Python $$pyver..."; \
+		tag_args="-t $(IMAGE)-lb-cpu:py$$pyver-$(WIP_TAG)"; \
+		if [ "$$pyver" = "$(DEFAULT_PYTHON_VERSION)" ]; then \
+			tag_args="$$tag_args -t $(IMAGE)-lb-cpu:$(WIP_TAG)"; \
+		fi; \
+		docker buildx build \
+			--platform $(MULTI_PLATFORM) \
+			--build-arg PYTHON_VERSION=$$pyver \
+			-f Dockerfile-lb-cpu \
+			$$tag_args \
+			. --push; \
+	done
+	@echo "All versioned images pushed. Default ($(DEFAULT_PYTHON_VERSION)) also tagged as :$(WIP_TAG)."
+
+# Versioned Smoke Tests
+
+smoketest-versioned: build-all-versioned # Verify Python version in each versioned image
+	@echo "Running Python version checks across all versioned images..."
+	@fail=0; \
+	check_version() { \
+		label="$$1"; image="$$2"; pyver="$$3"; \
+		echo -n "$$label py$$pyver: "; \
+		out=$$(docker run --rm $$image python --version 2>&1) || { echo "docker run failed"; fail=1; return; }; \
+		case "$$out" in \
+			Python\ $$pyver*) echo "$$out" ;; \
+			*) echo "Version mismatch (expected $$pyver): $$out"; fail=1 ;; \
+		esac; \
+	}; \
+	for pyver in $(GPU_PYTHON_VERSIONS); do \
+		check_version "GPU QB" "$(IMAGE):py$$pyver-$(TAG)" "$$pyver"; \
+	done; \
+	for pyver in $(CPU_PYTHON_VERSIONS); do \
+		check_version "CPU QB" "$(IMAGE)-cpu:py$$pyver-$(TAG)" "$$pyver"; \
+	done; \
+	for pyver in $(GPU_PYTHON_VERSIONS); do \
+		check_version "GPU LB" "$(IMAGE)-lb:py$$pyver-$(TAG)" "$$pyver"; \
+	done; \
+	for pyver in $(CPU_PYTHON_VERSIONS); do \
+		check_version "CPU LB" "$(IMAGE)-lb-cpu:py$$pyver-$(TAG)" "$$pyver"; \
+	done; \
+	if [ $$fail -ne 0 ]; then echo "FAIL: Some images failed version check"; exit 1; fi; \
+	echo "All 10 images passed Python version check."
+
 # Test commands
 test: # Run all tests in parallel
 	uv run pytest tests/ -v -n auto --dist loadscope