Skip to content

fix(ci): cache heavy CUDA wheels in install-test (#1796)#1798

Open
KevinSailema wants to merge 2 commits intoNVIDIA-NeMo:mainfrom
KevinSailema:fix/issue-1796-optimize-cuda-install-ci
Open

fix(ci): cache heavy CUDA wheels in install-test (#1796)#1798
KevinSailema wants to merge 2 commits intoNVIDIA-NeMo:mainfrom
KevinSailema:fix/issue-1796-optimize-cuda-install-ci

Conversation

@KevinSailema
Copy link
Copy Markdown

Reduce Python+CUDA installation time in CI by prebuilding and caching heavyweight CUDA dependency wheels, then reusing those wheels in install jobs.

Changelog

  • Add a dedicated CUDA wheelhouse job to build wheels for heavy dependencies used by CUDA extras.
  • Add GitHub Actions cache for the wheelhouse, keyed by OS, Python version, dependency lock inputs, and workflow file hash.
  • Add fail-fast validation so the workflow errors early if expected heavyweight wheels are not produced.
  • Upload built wheels as an artifact for downstream jobs in the same workflow run.
  • Update the CUDA pip install job to depend on wheelhouse generation and download wheel artifacts before install.
  • Update pip install behavior to prefer local prebuilt wheels first, while preserving fallback behavior when cache/artifact is cold.
  • Update the install summary job dependencies to include the new wheelhouse job so gating stays consistent.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?

Notes:

  • This is a CI workflow optimization change; no application/runtime logic was modified.
  • Validation is based on workflow execution success and install-time reduction in CI runs.

Additional Information

Signed-off-by: Kevin Sailema <108644636+KevinSailema@users.noreply.github.com>
@KevinSailema KevinSailema requested a review from a team as a code owner April 13, 2026 04:57
@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot bot commented Apr 13, 2026

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@akoumpa
Copy link
Copy Markdown
Contributor

akoumpa commented Apr 13, 2026

Thanks a lot for making these @KevinSailema ! I'm not an area expert, so I wanted to ask you, can this also be used when the docker container is made that run on the gha ci workers? because initially, it tests installation on a few targets [20 min], then builds the docker container [20 min], so I was wondering if your fix could be reused?

Hi, @thomasdhc can you review as the domain expert? Thank you.

@thomasdhc
Copy link
Copy Markdown
Contributor

/ok to test cc8cb2d

@thomasdhc
Copy link
Copy Markdown
Contributor

@KevinSailema Please review failures

@akoumpa akoumpa linked an issue Apr 14, 2026 that may be closed by this pull request
@chtruong814 chtruong814 added waiting-for-customer Waiting for response from the original author and removed waiting-for-customer Waiting for response from the original author labels Apr 14, 2026
Signed-off-by: Kevin Sailema <108644636+KevinSailema@users.noreply.github.com>
@chtruong814 chtruong814 added the needs-follow-up Issue needs follow-up label Apr 18, 2026
@chtruong814 chtruong814 added waiting-on-customer Waiting on the original author to respond and removed needs-follow-up Issue needs follow-up labels Apr 18, 2026
@akoumpa
Copy link
Copy Markdown
Contributor

akoumpa commented Apr 19, 2026

/ok to test 5368547

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request waiting-on-customer Waiting on the original author to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Optimize causal-conv1d/transformer engine installation in CI

5 participants