Key the C++ test JIT cache per test shard

[bdice]

#2665 keeps the CUDA JIT cache (CUDA_CACHE_PATH) of the C++ tests between CI runs. libcuvs links its JIT-LTO kernels at run time with nvJitLink, and nvJitLink caches the linked kernels there. A CAGRA single-CTA search kernel takes about 5 s to link, and a cached one takes 20–40 ms.

The C++ tests now run in 4 shards (#2720), and each shard links a different set of kernels. All four use the same cache key, though. Caches are immutable, so only one shard can save its archive under that key. The other three restore kernels they don't use and link their own again on every run.

Evidence from CI

Nightly test run https://github.com/rapidsai/cuvs/actions/runs/37894118256 (main at f751cad):

• Within each configuration, all four conda-cpp-tests shards restore the same archive (prefix libcuvs-cudajit-v1, about 9 MB).
• At the end only one shard can save it. In that run it was shard 2, the shortest shard and the one with the fewest JIT kernels.
• Shards 1, 3 and 4 log "Failed to save: Unable to reserve cache with key …, another job may be creating this cache." So the three long shards link their JIT kernels from scratch on every run.

Example jobs, 13.3.0 amd64 rtxpro6000: shard 1, 2, 3, 4.

configuration	shard 1	shard 2	shard 3	shard 4
13.3.0 amd64 rtxpro6000	876 s	216 s	835 s	655 s
12.9.2 amd64 ubuntu22.04 L4	1,151 s	225 s	1,314 s	1,062 s

This PR

• Adds the shard to cache-key-prefix in pr.yaml and test.yaml (libcuvs-cudajit-v1-shard${{ matrix.shard }}of${{ matrix.num_shards }}). Each shard then saves and restores its own archive.
• Logs the number of cache files before and after the tests in ci/test_cpp.sh, so it's visible which shards hit the cache. A shard whose count grows by hundreds of files linked its kernels again.

Measurements

RTX 6000 Ada (48 GB), 36-thread host, otherwise idle. Each shard's exact test list from the run above, with ctest -j8: first with an empty CUDA_CACHE_PATH, then again with the cache it had just filled.

shard	empty cache	own cache	uncompressed cache size
1	899 s	623 s	350 MB
2	230 s	140 s	70 MB
3	964 s	667 s	219 MB
4	748 s	502 s	175 MB

Shard 2 already gets its own cache today. For shards 1, 3 and 4 this is about 820 s less GPU-runner time per configuration, and about 300 s off the longest shard. The real gain may be somewhat smaller, because the restored shard-2 archive can cover a few kernels of the other shards.

For reviewers

• Per-shard caches add up to roughly 800 MB per configuration before compression, so cache storage per configuration grows.
• The first run after this change creates the four per-shard caches; later runs restore them.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[bdice]

/ok to test 87d358c

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/cuvs/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: f796d87d-8098-4640-a07b-709d011240c3

📥 Commits

Reviewing files that changed from the base of the PR and between 04508ab and 38e25f4.

📒 Files selected for processing: 1

• ci/test_cpp.sh

🚧 Files skipped from review as they are similar to previous changes: 1

• ci/test_cpp.sh

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

📝 Summary

Summary by CodeRabbit

• CI Improvements
  • CUDA JIT caches are separated by test shard.
  • Test runs report CUDA cache file counts before and after testing, along with the final cache size.

The C++ test workflows now include the shard index and total shard count in their CUDA JIT cache keys. The C++ test script reports CUDA JIT cache file counts before and after tests, and reports total cache disk usage after tests.

Changes: CUDA test cache

Layer / File(s)	Summary
Shard-specific cache keys .github/workflows/pr.yaml, .github/workflows/test.yaml	Both C++ test workflows include the shard index and total shard count in the CUDA JIT cache key.
Cache usage logging ci/test_cpp.sh	The script reports the cache file count before tests, then reports the file count and total disk usage after tests.

Priority: ➖ Normal

Estimated code review effort: 1 (Trivial) | ~4 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to 38e25

The cache-key and logging changes have no identified issue requiring resolution before merge.

Pre-merge checks: 5

Passed checks: 5 passed

Check name	Status	Explanation
Description check		The description clearly explains the shared CUDA JIT cache problem, the per-shard cache-key change, the added cache logging, and the measured impact.
Title check		The title clearly and concisely identifies the main change: assigning the C++ test JIT cache to each test shard.
Docstring Coverage		No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 1…
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)

• 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.}

[bdice]

/merge