Repository navigation
fix(build): keep torch's dependency tree out of tarballs - #390
justinwlin wants to merge 4 commits into
Conversation
…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.
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.
- 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.
| if should_exclude_package(req, excluded_packages): | ||
| # _requirement_name also reads "name @ url" forms, and wheel | ||
| # paths or URLs are named by their filename. | ||
| family_name = _requirement_name(req) or _wheel_distribution_name(req) |
There was a problem hiding this comment.
Follow-up to the resolved bare-wheel thread: with a marker, _requirement_name("torch-2.9.1-cp312-cp312-manylinux_2_28_x86_64.whl; platform_system=='Linux'") returns torch_2_9_1_cp312_..._whl (the regex accepts the whole filename because ; follows it), so the or _wheel_distribution_name(req) fallback never runs. The torch wheel isn't recognized, isn't filtered, and gets no ignored-pin warning, so pip installs the ~2.5 GB torch + CUDA tree.
Checking _wheel_distribution_name(req) first (_wheel_distribution_name(req) or _requirement_name(req)), or having _requirement_name return "" when the name ends in .whl regardless of what follows, covers it.
| return False | ||
|
|
||
|
|
||
| def _canonical_name(name: str) -> str: |
There was a problem hiding this comment.
_canonical_name duplicates _normalize_package_name (line 261), and they disagree on .: _is_excluded_top_dir checks package dirs with _normalize_package_name and dist-info dirs with _canonical_name, so a name containing . is normalized two different ways within the same function. Since _canonical_name is the PEP 503 one, replacing _normalize_package_name with it keeps a single rule.
Related: {_canonical_name(p) for p in SIZE_PROHIBITIVE_PACKAGES} is rebuilt in run_build and twice in create_tarball. A module-level frozenset next to SIZE_PROHIBITIVE_PACKAGES keeps the requirement filter and the tarball pruning on the same definition.
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).