Skip to content

fix(apps): keep transitive torch out of GPU artifacts - #623

Draft
justinwlin wants to merge 2 commits into
mainfrom
fix/build-prune-transitive-torch
Draft

justinwlin wants to merge 2 commits into
mainfrom
fix/build-prune-transitive-torch

Conversation

@justinwlin

Copy link
Copy Markdown
Contributor

GPU apps whose dependencies pull in torch (peft, accelerate, sentence-transformers) crash-loop after rp flash deploy. runpod-flash didn't have this problem, so Flash apps migrating to 2.0 break.

Cause

split_exclusions only skips torch-family packages listed as top-level requirements. pip still installs them as dependencies, and the vendored PyPI torch shadows the GPU image's build. A torch built off Linux has no NVIDIA libraries, so it fails to load (libcudart.so.13 / libcublasLt.so not found). runpod-flash removes these at archive time.

Change

  • After vendoring, remove every installed excluded distribution, plus the wheels only the torch family needs (NVIDIA/CUDA libraries, sympy, networkx, and scripts like torchrun). This uses each wheel's RECORD and Requires-Dist, the same approach as fix(build): keep torch's dependency tree out of tarballs聽flash#390. Without complete metadata, or when a requirement has no name (URL, path), only the excluded packages themselves are removed.
  • Auto-exclusion now applies when any resource is GPU on a builtin image and none uses a custom image. Before, every resource had to be GPU, so a mixed CPU+GPU app bundled torch and its GPU stage crash-looped. A vendored PyPI torch can't import on the builtin CPU images either, so CPU resources lose nothing.
  • Workers pip-install manifest excludedPackages that their image lacks. In mixed apps, torch is listed only when the CPU resources' dependencies need it, and the build warns that those workers will install it at cold start. Otherwise a CPU stage that never uses torch tries to install CUDA torch onto its 10 GB disk and fails to boot.
GPU app with dependencies=["peft"] Before After
Artifact 645.6 MB, torch 2.14.1 bundled 80.7 MB, no torch
Worker crash loop image torch 2.9.1+cu128, CUDA available
Verification
  • Live rp flash deploy of a GPU queue app with dependencies=["peft"]. The worker returned {"cuda": true, "file": "/usr/local/lib/python3.12/site-packages/torch/__init__.py", "torch": "2.9.1+cu128", "peft": "0.21.2"}.
  • Live mixed app: a CPU queue (humanize) plus a GPU queue (torch, peft). Both returned results. Before the manifest change, the CPU worker failed bootstrap verify with No space left on device while installing torch.
  • docker run python:3.12-slim with a macOS-vendored PyPI torch on PYTHONPATH fails with libcublasLt.so.*[0-9] not found, so vendoring torch never worked on the CPU images.
  • Unit tests for pruning (shared nvidia/ dirs, ../../bin script paths, missing RECORD/METADATA, unnamed requirements, user --exclude dependencies, custom images, mixed apps, manifest contents). pytest tests: 1494 passed, coverage 91.94%.

Packages that depend on torch (peft, accelerate, sentence-transformers)
installed it into the artifact even on builtin GPU images, because only
top-level requirements were excluded. The vendored PyPI torch shadowed
the image's build and failed to load its CUDA libraries, so workers
crash-looped. runpod-flash stripped these at archive time.

After vendoring, remove every installed excluded distribution, plus the
wheels only the torch family depends on (nvidia/cuda, sympy, networkx),
using each wheel's RECORD and Requires-Dist. Fails closed to name-only
removal when metadata is missing or a requirement has no name. Script
paths pip --target records as ../../bin are mapped onto the target.

Auto-exclusion now applies when any resource is GPU on a builtin image
and none uses a custom image, matching runpod-flash for mixed CPU+GPU
apps. A vendored PyPI torch cannot import on the builtin CPU images
either, so CPU resources lose nothing.
Workers pip-install manifest excludedPackages their image lacks at cold
start. In a mixed CPU+GPU app the GPU images already have the torch
family, so only CPU workers act on the list, and a CPU stage that never
uses torch tried to install CUDA torch onto its default 10 GB disk and
failed to boot.

In mixed apps, keep a torch-family name in the manifest only when the
CPU resources' dependency closure (their decorator dependencies plus
requirements.txt) reaches it, and warn that those workers install it at
cold start. Unreadable metadata keeps every name, as before.

This branch has not been deployed

No deployments
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