Skip to content

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

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

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

Conversation

@justinwlin

@justinwlin justinwlin commented Oct 16, 2026 •

Copy link
Copy Markdown
Contributor

Flash → SDK 2.0 fix stack

PRs 1–9 are independent PRs off main, and 10 is stacked on 9. Merge in this order. Hover a link for its status.

# PR Fixes Needs
1 #627 ✅ merged: CI e2e fixed on main —
2 👉 #623 (this PR) GPU crash loop: torch pulled in by peft/accelerate gets bundled nothing
3 #624 Workers land on GPUs you didn't request (whole pools) nothing
4 #625 rp flash dev leaks endpoints when stopped by a signal nothing
5 #626 vol.path says "not mounted" on deployed workers nothing
6 #628 Account API key copied into every endpoint nothing
7 — Release runpod 2.0.0rc2 with 1–6 1–6 merged
8 runpod-workers/sdk-runtime#1 Rebuild runtime images (dev sessions run an old SDK) independent
9 runpod/runpod-plugins-official#89 Plugin docs for the fixed SDK (flash lane → SDK 2.0) 2.0.0rc2 on PyPI
10 runpod/runpod-plugins-official#90 Plugin VolumeCache guidance: global volumes, measured numbers, golden paths 08/09 (stacked on #89) runpod/runpod-plugins-official#89

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 (a dependency behind an extra counts only when that extra is requested), 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.

Testing

GPU app, dependencies=["peft"], 4090 pool main runpod-flash 1.21 This PR
Artifact 645.5 MB, PyPI torch bundled 47.8 MB, torch/ stripped at archive time 80.7 MB, no torch
Worker crash loop: libcudart.so.13: cannot open shared object file image torch 2.9.1+cu128, CUDA available image torch 2.9.1+cu128, CUDA available
Mixed app: CPU transformers + GPU peft (--build-only) First version of this PR (398e5f0) This PR
Manifest excludedPackages ["torch"], plus "cpu resources need torch" warning none
CPU worker at cold start would pip-install CUDA torch onto a 10 GB disk installs nothing
How we tested and compared
  • Same app on three builds. One queue function returns torch.__version__, torch.__file__ and torch.cuda.is_available(). It was deployed with runpod-flash 1.21, with SDK main, and with this branch, all on the same GPU pool. We compared artifact size, bundled packages, and which torch the worker loaded.
    • Flash's worker loaded /usr/local/lib/python3.12/dist-packages/torch. This PR's worker loaded /usr/local/lib/python3.12/site-packages/torch/__init__.py. Both are the image's build.
  • Mixed apps, live. 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.
  • Extras, real wheels. transformers 5.20.2 lists torch>=2.6; extra == "torch". The --build-only comparison above used real pip resolution and showed the first version recording torch for a CPU stage that never imports it.
  • CPU images can't use vendored torch either. docker run python:3.12-slim with a macOS-vendored PyPI torch on PYTHONPATH fails with libcublasLt.so.*[0-9] not found.
  • Unit tests cover pruning (shared nvidia/ dirs, ../../bin script paths, missing RECORD/METADATA, unnamed requirements, user --exclude dependencies, custom images) and mixed-app manifests (CPU need, direct GPU torch, an unrequested extra, a requested extra, and a marker branch that applies without the extra). pytest tests: 1496 passed, coverage 91.94%.
Review (Codex, two rounds)
Finding Resolution
The CPU-needs closure followed every Requires-Dist line, so plain transformers recorded torch Fixed in 1889c3f. A line gated by an extra counts only when that extra was asked for.
Round 2: extra == "torch" or python_version >= "3.12" was treated as gated Fixed in 1889c3f. A line counts as gated only when every or branch names an extra.
A torch wheel without RECORD falls back to removing only torch/ Not changed. pip always writes RECORD for the installs this build makes, and the fallback exists only to fail safe.

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.
…or it

transformers lists torch under its [torch] extra. the cpu-needs closure followed every Requires-Dist line, so a mixed app whose cpu stage used plain transformers recorded torch and its workers installed it at cold start.
@justinwlin
justinwlin force-pushed the fix/build-prune-transitive-torch branch from 9005f22 to 1889c3f Compare October 20, 2026 15:48
@justinwlin
justinwlin marked this pull request as ready for review October 20, 2026 15:48

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.

2 participants