ci: warm the GPU uv cache on a CPU runner before GPU jobs by zywind · Pull Request #798 · NVIDIA-NeMo/Safe-Synthesizer

zywind · 2026-10-27T14:49:54Z

Summary

Follow-up to #795. GPU jobs restore one uv cache entry per CUDA extra, keyed on a hash of pyproject.toml and .python-version. When that entry is missing (the first nightly after any pyproject.toml change, about weekly, or after eviction), every GPU job downloads ~280 packages itself. On slow nights that has run past the 60-minute smoke job timeout while holding an A100: nightly 40437182386, the first after #795 merged, lost 4 smoke jobs this way (one cu129 job spent 44m 44s downloading; three cu130 jobs never finished).

This PR adds a warm-uv-cache job (one per CUDA extra) on linux-amd64-cpu4 with a 120-minute timeout. The GPU smoke and e2e jobs now needs: it and keep their 60-minute limit.

How it works:

Runs setup-uv with the same key-affecting inputs as setup-gpu-test-env (enable-cache, cache-suffix, cache-dependency-glob, no python-version), with restore and save off, to get the exact cache key.
Looks the key up through the Actions caches API (this ref or the default branch). On a hit it exits in seconds, without downloading the multi-GB entry.
On a miss: runs setup-uv again with saving on, builds the wheel, and runs install_nss.sh for Python 3.11-3.14 so every version's torch wheels land in one entry. The setup-uv post step saves it.

Details:

Saving is split into a second setup-uv step so a hit never uploads an empty cache; on a branch that would shadow main's entry.
Uses the uv version pinned in .mise.toml (0.13.4), which is what the GPU install actually runs (mise shims precede the setup-uv binary on PATH). The uv version is not part of the cache key.
CPU and GPU runners report the same image (VM-Images version: 78c50ae) and cache path (/home/runner/.cache/uv), so keys and cache versions match.
If the warming job fails, GPU jobs still run (always()) and fill the cache themselves as before.
PYTHON_VERSIONS duplicates the smoke matrix; a comment marks it to keep in sync.

Human review

Select exactly one. The default is no human review; uncheck it when choosing another level.

⚪ No human review; automated review only
🟡 Partially reviewed or spot-checked by a human
🔵 Complete diff reviewed by a human
🟢 Complete diff reviewed and verified by a human

Verification performed: manual GPU smoke run 40518808190 on this branch. Both warming jobs computed keys identical to main's -cu129/-cu130 entries, found them, and exited in 17 s; all 8 GPU smoke jobs then hit the cache and passed. actionlint reports only the self-hosted runner labels it doesn't know.

Not yet exercised: the cache-miss path (install on a CPU runner, save, GPU jobs restoring that entry). Testing it on a branch creates ~12 GB of branch-scoped entries that would evict main's, so it is left to the first nightly after the next pyproject.toml change. If it misbehaves, GPU jobs fall back to filling the cache as they do today.

Pre-Review Checklist

Ensure that the following pass:

mise run format leaves only intended formatting changes.
mise run check ::: test passes locally.
mise run test:e2e passes locally
mise run test:ci-container passes locally (recommended)
GPU CI status check passes -- comment /sync on this PR to trigger a run (auto-triggers on ready-for-review)

Pre-Merge Checklist

New or updated tests for any fix or new behavior
Updated documentation for new features and behaviors, including docstrings for API docs.

Other Notes

CI-config-only change; no tests apply.

Summary by CodeRabbit

Chores
- Improved preparation for GPU test runs by reusing cached test dependencies or preparing them when needed across supported CUDA and Python versions.
- GPU smoke and end-to-end checks can still run if dependency preparation fails. These workflow updates do not change application behavior.

When the per-CUDA uv cache entry is missing (first run after a pyproject.toml change, or after eviction), every GPU job downloads ~280 packages. On slow nights that has pushed setup past the 60-minute GPU job timeout while holding an A100. Add a warm-uv-cache job (one per CUDA extra) on a CPU runner with a 120-minute timeout. It computes the same setup-uv cache key as the GPU jobs, checks for an existing entry through the API, and on a miss installs the wheel for every smoke-test Python version and saves the cache. GPU jobs wait for it but still run if it fails. Signed-off-by: Yunfeng Zhang <yunzhang@nvidia.com>

coderabbitai · 2026-10-27T14:50:12Z

No actionable comments were generated in the recent review.

Recent review info

Run configuration

Source mode: Auto
Repository: NVIDIA-NeMo/Safe-Synthesizer
Branch: yunfeng/gpu-ci-cache-warm
File: .coderabbit.yaml
Commit: d833f4f
Layers: Repository YAML
Review profile: CHILL
Plan: Enterprise
Run ID: 68d2999a-c47f-453e-a1c5-ee29b3b59b0f

Commits

Reviewing files that changed from the base of the PR and between 5424643 and d833f4f.

Files selected for processing: 1

.github/workflows/gpu-tests.yml

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Recent review details

Context from checks skipped due to timeout: 7

GitHub Check: CodeRabbit
GitHub Check: End-user Wheel Install
GitHub Check: Unit Tests (3.14)
GitHub Check: Unit Tests (3.12)
GitHub Check: Unit Tests (3.11)
GitHub Check: Unit Tests (3.13)
GitHub Check: Greptile Review

Additional context used

Path-based instructions: 1

Review GitHub configuration for branch protection expectations, CODEOWNERS alignment, least privilege permissions, pinned actions where practical, and consistency with CONTRIBUTING.md.

CodeRabbit configuration file

Files:

.github/workflows/gpu-tests.yml

actionlint (1.7.12)

.github/workflows/gpu-tests.yml

[error] 93-93: label "linux-amd64-cpu4" is unknown. available labels are "windows-latest", "windows-latest-8-cores", "windows-2025", "windows-2025-vs2026", "windows-2022", "windows-11-arm", "ubuntu-slim", "ubuntu-latest", "ubuntu-latest-4-cores", "ubuntu-latest-8-cores", "ubuntu-latest-16-cores", "ubuntu-24.04", "ubuntu-24.04-arm", "ubuntu-22.04", "ubuntu-22.04-arm", "macos-latest", "macos-latest-xlarge", "macos-latest-large", "macos-26-intel", "macos-26-xlarge", "macos-26-large", "macos-26", "macos-15-intel", "macos-15-xlarge", "macos-15-large", "macos-15", "macos-14-xlarge", "macos-14-large", "macos-14", "self-hosted", "x64", "arm", "arm64", "linux", "macos", "windows". if it is a custom label for self-hosted runner, set list of labels in actionlint.yaml config file

(runner-label)

Additional comments: 1

.github/workflows/gpu-tests.yml: 1

70-197: LGTM!

Also applies to: 200-202, 287-287

Walkthrough

The GPU test workflow adds a CPU job that checks and warms CUDA-specific dependency caches. The GPU smoke and E2E jobs now depend on that job.

Changes

GPU Test Cache Warming

Layer / File(s)	Summary
Warm CUDA dependency caches `.github/workflows/gpu-tests.yml`	A CPU job checks for matching `cu129` and `cu130` caches. On a cache miss, it builds the wheel and installs GPU test dependencies for Python 3.11–3.14.
Wire cache warming to GPU jobs `.github/workflows/gpu-tests.yml`	The smoke-test and E2E jobs now depend on cache warming. The smoke-test job retains its existing `always()` condition.

Priority: Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Other

Possibly related PRs

NVIDIA-NeMo/Safe-Synthesizer#795: Changes the setup-uv cache key to separate entries for cu129 and cu130, which this workflow warms.

Suggested labels: perf, chore

Suggested reviewers: binaryaaron

Merge Risk: Minimal · up to d833f

GPU jobs wait for cache warming but can proceed if it fails. No merge-blocking issue was identified.

Pre-merge checks: 4

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: warming the GPU uv cache on a CPU runner before GPU jobs.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Finishing Touches

🧪 Generate unit tests (beta)

Commit to this branch
Create a new PR

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

greptile-apps · 2026-10-27T14:54:08Z

Confidence Score: 5/5

[Low impact] The PR appears safe to merge; no blocking issue was found.

Summary

Adds a CPU job that warms the GPU download cache separately for each CUDA extra.

GPU test runs warm each CUDA cache on a CPU runner first.

Acknowledged scope:

zywind explicitly deferred testing the cache-miss path because branch-scoped entries could evict the default branch’s cache.
zywind explicitly accepted the duplicated Python-version list and added a reminder to keep it aligned with the smoke matrix.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  Changes["Detect changes or manual dispatch"] --> Warm["CPU warm job per CUDA extra"]
  Warm --> Lookup{"Cache entry exists?"}
  Lookup -->|Yes| Done["Skip downloads"]
  Lookup -->|No| Install["Build wheel and install for Python 3.11–3.14"]
  Install --> Save["Save download cache"]
  Done --> GPU["Run selected GPU suites"]
  Save --> GPU
  Warm -->|Failure| Fallback["GPU jobs download what they need"]
  Fallback --> GPU
  GPU --> Status["Report GPU results"]

_{Reviews (1) · Last reviewed commit: "ci: warm the GPU uv cache on a CPU runne..." · Reviewed by Greptile}

codecov · 2026-10-27T14:57:01Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

zywind requested a review from a team as a code owner October 27, 2026 14:49

github-actions Bot added the area:ci label Oct 27, 2026

coderabbitai Bot added chore Maintenance not tied to a user-visible change perf Performance improvement labels Oct 27, 2026

mckornfield approved these changes Oct 27, 2026

View reviewed changes

zywind added this pull request to the merge queue Oct 27, 2026

Merged via the queue into main with commit e9e20e6 Oct 27, 2026
41 checks passed

zywind deleted the yunfeng/gpu-ci-cache-warm branch October 27, 2026 17:01

Navigation Menu

Repository navigation

ci: warm the GPU uv cache on a CPU runner before GPU jobs - #798

ci: warm the GPU uv cache on a CPU runner before GPU jobs#798
zywind merged 1 commit into
mainfrom
yunfeng/gpu-ci-cache-warm

zywind commented Oct 27, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Oct 27, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Oct 27, 2026 •

edited

Loading

Diagram

Uh oh!

codecov Bot commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

zywind commented Oct 27, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Human review

Pre-Review Checklist

Pre-Merge Checklist

Other Notes

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

greptile-apps Bot commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

codecov Bot commented Oct 27, 2026

Codecov Report

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zywind commented Oct 27, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Oct 27, 2026 •

edited

Loading

greptile-apps Bot commented Oct 27, 2026 •

edited

Loading