Repository navigation
fix(build): keep torch's dependency tree out of tarballs - #390
Open
justinwlin wants to merge 5 commits into
Open
justinwlin wants to merge 5 commits into
justinwlin wants to merge 5 commits into
Conversation
2 tasks done
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Pruning can remove explicitly requested wheels and dependencies of user-excluded packages, while some ignored pins still bypass warnings.
3 open findings
What changed in this PR
Reduces flash build archive sizes by pruning dependencies needed only by excluded torch-family packages.
Changes:
- Adds dependency-graph pruning and file ownership filtering.
- Displays exclusion notices and ignored-version warnings.
- Adds tests for pruning safeguards and notices.
| File | Description |
|---|---|
| tests/unit/cli/commands/test_build.py | Tests pruning, file preservation, and notices. |
| src/runpod_flash/cli/commands/build.py | Implements dependency pruning, archive filtering, and notices. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
deanq
reviewed
Oct 16, 2026
…exclusions torch-family packages stay excluded: the GPU worker image ships torch 2.9.1 and a torchvision pinned to it, and the tarball sits ahead of site-packages, so a bundled PyPI torchvision (0.29.1 needs torch>=2.14) would shadow the paired one. But pip still installs their whole dependency tree, and only exact package names were stripped: - On Linux x86_64, any requirement that depends on torch (timm, accelerate, sentence-transformers) pulls ~2 GB of nvidia/cuda wheels. A one-dependency timm app built to 1957.5 MB and failed the 1500 MB limit. - functorch/, torchgen/ and torchvision.libs/ leaked into the archive. Drop wheels that only the torch family depends on (from Requires-Dist), and every file their RECORD lists unless a kept wheel also installed it. Files no RECORD lists stay. Pruning is skipped when any wheel lacks METADATA or RECORD, or a requirement line has no readable name (bare URL, path, -r). A user --exclude still drops only the named package. The same timm app now builds to 32.5 MB. Builds also print which torch-family requirements were not bundled, and warn when a version pin on them is ignored, instead of a debug-only log.
- Detect torch-family requirements with _requirement_name, so "~=" pins and "name @ url" forms are excluded and warned about instead of reaching pip. - Treat bare archive filenames (*.whl, *.tar.gz) as unnamed, so pruning fails closed instead of dropping the wheel they install. - Protect dependencies of user --exclude packages even when they overlap torch's tree, so --exclude removes only the named package.
… or @ - Name wheel paths and URLs by their filename, so a pinned download.pytorch.org torch/torchvision wheel is excluded and warned about instead of reaching pip. - extract_package_name now splits on ~, @, whitespace and (, so --exclude sympy matches "sympy~=1.14" and "sympy @ url". - Only treat a requirement as an archive filename when the whole line is one, so a distribution named "my.whl" still counts as a root.
…name normalization _requirement_name now returns "" for an archive filename followed by an environment marker, so the _wheel_distribution_name fallback runs and a torch wheel with a marker is filtered and warned about. Replace _normalize_package_name with the PEP 503 _canonical_name so tarball filtering uses a single normalization rule, and hoist the canonical torch-family set into SIZE_PROHIBITIVE_CANONICAL.
justinwlin
force-pushed
the
fix/torch-dependency-tree
branch
from
October 21, 2026 22:10
d3cc1c7 to
7a951f1
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stops
flash buildfrom bundling the wheels that onlytorchandtorchvisionneed, and tells users what was left out. Supersedes #376. Refs #353.Why not un-exclude torchvision (#376)
The GPU worker image installs
torchvision==0.24.1, paired with itstorch 2.9.1. The tarball is extracted ahead of site-packages, so a bundled PyPItorchvision(0.29.1, needstorch>=2.14) would override it and fail to load.torchvisionstays excluded. The load-balanced image is missing it, and that's fixed in runpod-workers/flash#113.What was actually breaking
torch(timm,accelerate,sentence-transformers) installs about 2 GB ofnvidia-*/cuda-*wheels. Only the exact namestorch/torchvision/... were stripped. An app whose only dependency wastimmbuilt to 1957.5 MB and failed the 1500 MB limit.functorch/,torchgen/andtorchvision.libs/leaked into every archive that pulled intorch.Change
Requires-Distgraph), plus every file theirRECORDlists unless a kept wheel installed it too. Files noRECORDlists stay.METADATA/RECORD, aRequires-Distline is unreadable, or a requirement line has no name (bare URL, path, archive filename,-r). Your own requirements and--excludedependencies are never pruned.not bundling torchvision: provided by the Runpod GPU worker image, and warn when a pin liketorchvision==0.29.1,~=0.24.1,@ urlor adownload.pytorch.orgwheel URL is ignored.--excludenow also matches~=and@ urlrequirements.timmappVerification
flash buildof a one-dependencytimmapp in alinux/amd64python:3.12-slimcontainer, before and after: 1957.5 MB failing, then 32.5 MB with nothing from torch/nvidia/cuda/triton left in the archive.sympy,mpmath,networkx,functorch,torchgendropped).-n auto -m "not serial") + 53 serial, coverage 86.82%.ruff checkandruff format --checkclean. No new mypy errors inbuild.py.nvidia/dirs, missingRECORD/METADATA, unreadableRequires-Dist, direct-URL, bare-URL and archive-filename requirements,~=/@pins,RECORDpath normalization, unrecorded files,--excludedependencies (including onestorchalso needs), and the notices. Each review fix's test was confirmed to fail on the earlier code.torch/is dropped whole, and nested-r other.txtfiles go to pip unparsed (pruning is skipped for them).