From b47c27535e58d2868a97b26f149efdb0bbd880d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= <dean.quinanola@runpod.io>
Date: Thu, 9 Jul 2026 09:31:29 -0700
Subject: [PATCH] fix(docker): install torchvision in GPU worker image

The runpod/pytorch base image ships torch/torchaudio/triton but not
torchvision, and `flash build` auto-excludes torchvision from the
deployment tarball (SIZE_PROHIBITIVE_PACKAGES) on the assumption the
base image provides it. Nothing backfilled it, so any deployed Flash
app importing torchvision failed at runtime with ModuleNotFoundError.

Install torchvision in the worker image, mirroring the existing numpy
backfill: pin to the release paired with the base torch (2.9.1 ->
0.24.1) and use the CUDA wheel index to stay CUDA-aligned. Add a build
-time import assertion so a future torch bump that breaks the pairing
fails the build instead of the customer.

Fixes SLS-377.
---
 Dockerfile | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 01a49ca..ef0d4ad 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -8,6 +8,11 @@ FROM runpod/pytorch:1.0.3-cu1281-torch291-ubuntu2204
 # Target Python version for the worker runtime.
 ARG PYTHON_VERSION=3.12
 ARG TORCH_VERSION=2.9.1+cu128
+# torchvision is NOT shipped by the runpod/pytorch base image (unlike torch and
+# torchaudio) and is auto-excluded from Flash build tarballs, so it must be
+# installed here. Pin to the release paired with TORCH_VERSION (torch 2.9.1 ->
+# torchvision 0.24.1); update both together on a torch bump. See SLS-377.
+ARG TORCHVISION_VERSION=0.24.1
 ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
 
 # Expose the target version to the running worker for startup validation.
@@ -75,9 +80,18 @@ RUN uv export --format requirements-txt --no-dev --no-hashes > requirements.txt
 # must be provided here in the base image.
 RUN python -m pip install --no-cache-dir numpy
 
-# Verify torch, numpy, and the expected Python version are available.
+# Install torchvision for the active Python version.
+# The runpod/pytorch base ships torch/torchaudio but NOT torchvision, and Flash
+# build auto-excludes torchvision from tarballs, so it must be provided here.
+# Use the CUDA wheel index to stay aligned with the base image torch build.
+RUN python -m pip install --no-cache-dir \
+      --index-url ${TORCH_INDEX_URL} \
+      "torchvision==${TORCHVISION_VERSION}"
+
+# Verify torch, torchvision, 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 torchvision; print(f'torchvision {torchvision.__version__}')" \
  && python -c "import numpy; print(f'numpy {numpy.__version__}')"
 
 CMD ["python", "handler.py"]