Persist torch.compile cache across GPU CI jobs

[spaingnv]

Summary

• GPU CI now restores and saves $FLASHDREAMS_CACHE_DIR/torchinductor, keyed by GPU arch and uv.lock.
• compile_module already pinned Inductor under that path locally; this is what was missing for CI.
• Add a ci_gpu two-process check that a warm compile is at least twice as fast as a cold one.

Fixes #500.

Test plan

• pytest flashdreams/tests/test_compile.py -m ci_cpu
• pytest flashdreams/tests/test_compile_cache_gpu.py -m ci_gpu on RTX PRO 6000 (cold ~6s, warm ~1s)
• First GPU CI job on this PR (cold compile + cache upload)
• A later GPU job on the same arch/uv.lock shows an inductor cache restore

[copy-pr-bot]

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.

[greptile-apps]

Greptile Summary

The PR persists PyTorch Inductor artifacts between GPU CI jobs and adds cross-process validation of cache reuse.

• Configures the GPU job to cache $FLASHDREAMS_CACHE_DIR/torchinductor, keyed by GPU architecture and uv.lock.
• Expands cache configuration coverage to verify PyTorch resolves the expected directory.
• Adds a GPU test that compiles in separate processes, verifies cache artifacts persist, and logs cold and warm durations without gating on timing.

Confidence Score: 3/5

The PR should not merge until the GPU cache action is pinned to a reviewed immutable commit SHA.

The timing-gate issue is fixed, but every GPU job still executes actions/cache through the mutable v4 tag, allowing the workflow implementation to change without a repository review.

Files Needing Attention: .github/workflows/ci.yml

Important Files Changed

Filename	Overview
.github/workflows/ci.yml	Adds architecture- and lockfile-keyed persistence for GPU Inductor artifacts; the cache action remains referenced through a mutable tag.
flashdreams/flashdreams/infra/compile.py	Clarifies the existing persistent Inductor cache behavior without changing executable logic.
flashdreams/tests/test_compile.py	Confirms both the environment variable and PyTorch cache resolver use the configured directory.
flashdreams/tests/test_compile_cache_gpu.py	Adds cross-process GPU cache validation and removes the previously unstable timing threshold from the required gate.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[GPU CI job] --> B[Restore architecture-keyed cache]
  B --> C[Set FLASHDREAMS_CACHE_DIR]
  C --> D[First worker compiles TinyMLP]
  D --> E[Write Inductor and FX-graph artifacts]
  E --> F[Second worker compiles TinyMLP]
  F --> G[Verify cached artifacts remain]
  G --> H[Save cache for later GPU jobs]

Loading

_{Reviews (2): Last reviewed commit: "Stop gating GPU compile-cache CI on a 2x..." | Re-trigger Greptile}

[greptile-apps]

Mutable cache action reference

The new GPU CI step references actions/cache@v4, so an upstream tag mutation changes the code executed by the workflow without a repository change. Pin the action to a reviewed full commit SHA to make this dependency immutable.

How this was verified: The workflow directly executes the mutable actions/cache@v4 reference on every GPU CI job.

[fangjunzhou-nv]

#500 is not a CI specific issue. This issue happens to torch.compile even outside the CI. Specifically, running compile_module from flashdreams/flashdreams/infra/compile.py usually takes a long time.

This indicates we're not using torch.compile cache properly and we need to investigate this before we propose a fix.

[spaingnv]

Agreed — the CI cache was the wrong problem. compile.py already pins TORCHINDUCTOR_CACHE_DIR; a tiny MLP hitting that cache does not explain a slow real compile_module / first forward.

Converting to draft. Next I will measure a real OmniDreams compile (wrap vs first forward vs CUDA-graph warmup, cold vs second process) and report cache hit/miss before proposing a fix.

[jmccaffrey-nv]

The change reuses the existing cache-root abstraction, stays narrowly scoped, and the added cache_dir() assertion is useful.

[P1 correctness/reliability] The GPU test can pass without a cache hit. It writes to a fresh tmp_path, while CI persists /tmp/flashdreams-cache/torchinductor; moreover, the assertions only prove files remain, not that the warm worker reused them ([lines 61–103 flashdreams/tests/test_compile_cache_gpu.py#L61-L103)) . Can the test assert PyTorch’s fxgraph_cache_hit counter? Test the desired outcome directly. .

Comment P1 security: GitHub also recommends treating restored cache contents as untrusted and keeping secrets out of cached paths (GitHub cache security guidance).