Skip to content

fix(docker): install torchvision in GPU load-balancer image - #113

Open
justinwlin wants to merge 1 commit into
mainfrom
fix/lb-image-torchvision
Open

justinwlin wants to merge 1 commit into
mainfrom
fix/lb-image-torchvision

Conversation

@justinwlin

Copy link
Copy Markdown

Installs the pinned torchvision in Dockerfile-lb, matching what #101 did for Dockerfile.

Why

flash build never bundles torchvision, because a bundled copy would shadow one matched to the image's torch. The queue-based GPU image has installed torchvision==0.24.1 since #101. The load-balanced GPU image never got it, so LB endpoints that import torchvision fail with ModuleNotFoundError.

Change

  • ARG TORCHVISION_VERSION=0.24.1, the release paired with the base torch 2.9.1.
  • Install from the cu128 index, the same install block as Dockerfile.
  • Add import torchvision to the build-time verify step.
Verification
  • The install block is identical to Dockerfile (diff of the two blocks is empty).
  • The published images confirm the gap: runpod/flash:py3.12-latest has the torchvision step in its build history, and runpod/flash-lb:py3.12-latest doesn't.
  • torch==2.9.1+cu128 + torchvision==0.24.1 resolves cleanly against the cu128 index. PyPI's latest torchvision 0.29.1 needs torch>=2.14, which is why the pin matters.
  • I didn't run a local image build (amd64 GPU image). CI builds Dockerfile-lb for Python 3.10 through 3.13, and the non-3.12 path reinstalls torch before this step.

Dockerfile got a pinned torchvision in #101 (SLS-377), but Dockerfile-lb did
not. flash build auto-excludes torchvision from tarballs, so load-balanced GPU
endpoints that import torchvision fail with ModuleNotFoundError.

Mirror the Dockerfile install: torchvision 0.24.1 from the cu128 index, paired
with the base image torch 2.9.1, plus the build-time import check.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant