ci: pull CPU-only torch wheels on GPU-less runners#163
Open
bkowshik wants to merge 1 commit into
Open
Conversation
CI installs torch/torchvision/torchaudio via uv on GPU-less 8-core-ubuntu runners with no CPU pin, so uv resolves the default CUDA builds and drags in ~5 GB of CUDA libraries that never execute. That bloat is paid twice: once on every install, and again on every monthly .venv cache upload and restore (the venv cache is already 5-6 GB). Set UV_TORCH_BACKEND=cpu on the shared "Install / sync packages" step in the setup-python-env composite action, which every env-using job calls, so all four jobs (test-neuralset, test-downstream, typecheck, build-docs) get the CPU-only index in one place. Scoped to the CI action rather than the package pyprojects: neuralbench pins torch to keep CUDA capability sm_60 support for GPU users, and that intent stays intact since the env var only affects this workflow.
meta-cla
Bot
added
the
CLA Signed
kingjr
reviewed
Jun 18, 2026
| # instead of the default CUDA builds (~5 GB of unused CUDA libs that also | ||
| # inflate every monthly .venv cache upload/restore). Scoped to CI here, | ||
| # not the pyprojects, so GPU users (neuralbench sm_60) are unaffected. | ||
| UV_TORCH_BACKEND: cpu |
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
2 participants
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.
What
Add
UV_TORCH_BACKEND: cputo theInstall / sync packagesstep in.github/actions/setup-python-env/action.ymlso CI installs CPU-onlytorch/torchvision/torchaudio.Why
CI installs torch via
uvon GPU-less8-core-ubunturunners with no CPU pin. torch comes from the sub-project pyprojects (torch>=2.5.1in neuralset/neuraltrain, theallextra adding torchvision/torchaudio;torch==2.6in neuralbench). With no backend hint,uvresolves the CUDA builds and pulls ~5 GB of CUDA libraries that never run. The cost is paid twice: on every install, and again on every monthly.venvcache upload/restore (already 5-6 GB).Fix
The composite action's install step is the single chokepoint — every env-using job (
test-neuralset,test-downstream,typecheck,build-docs) calls it. One env var fixes all four.Scoped to CI, not the pyprojects:
neuralbenchpins torch forCUDA capability sm_60, so GPU users must stay on CUDA builds. The env var only affects this workflow.Notes
UV_TORCH_BACKENDrequiresuv >= 0.6.0.setup-uv@v4is unpinned and defaults to latest, so this works today; pinninguv < 0.6.0would silently no-op it.YAML validated; no change to local/GPU installs.