Fix JIT cache crash for rename-race losers on shared filesystems

[Oseltamivir]

Fixes #683.

Problem

On a shared filesystem (the default $HOME/.deep_ep is typically NFS on multi-node clusters), the losing side of the compile/rename race in Compiler::build() crashes at EP_HOST_ASSERT(runtime != nullptr): its rename fails authoritatively at the server, but the winner's directory may not be visible to its NFS client yet, the initial cache miss primed a negative lookup cache entry that can outlive the failed rename. Full mechanism in #683.

Change

Make the losing path self-sufficient instead of requiring cross-client visibility:

• csrc/jit/compiler.hpp: on rename failure, load the runtime from this process's own just-compiled temp directory (equivalent artifacts. The cache key hashes kernel name, source, flags, and compiler signature), cache it under the canonical dir_path key, then delete the temp directory.
• csrc/jit/cache.hpp: add KernelRuntimeCache::put(key_path, load_path) to construct from one path and cache under another.

This is safe because the KernelRuntime constructor consumes the files eagerly (cuobjdump -symbols, then cudaLibraryLoadFromFile/cuModuleLoad) and keeps only GPU handles. Warm starts still read the shared cache exactly as before; other rename failures (e.g. permissions) now degrade gracefully instead of asserting (scope: the rename failure itself — a subsequent failure to load the loser's own artifacts remains fatal, with the same check_validity diagnostics as any other cache entry).

Testing

Two 8-GPU H200 nodes (16 ranks), JIT cache on shared NFS, cold start, ElasticBuffer construction plus a dispatch/combine round trip per rank (EP_JIT_DEBUG=1):

• Unpatched: reproduced: all 8 ranks of the losing node crash at compiler.hpp:158 while JIT-building the dispatch/combine kernels; the other node hangs in the next collective.
• Patched, same scenario: 16/16 pass; the new branch executed 73 times across 80 rank×kernel builds (Loading CUBIN: <cache>/tmp/<uuid>/kernel.cubin), and the shared tmp/ is empty afterwards.
• Warm start: 16/16 pass with all 80 loads served from the canonical cache directories; the new branch is never taken.
• The full extension builds cleanly with the patch (torch 2.10, CUDA 13.0, NCCL 2.30.4, sm90).
• Review commit ef2ae56 re-validated on the same setup: the new EP_JIT_DEBUG fallback log fired 75/75 times, each paired 1:1 with a tmp-dir CUBIN load, 16/16 ranks green.

[ds-review-bot]

🤖 ds-review-bot Code Review

v4

The PR makes the rename-race LOSER self-sufficient instead of reading the winner's directory across NFS clients. The cache miss in build() primes a negative lookup on the loser's client that can outlive the failed atomic rename, so a stat of the winner's dir is still answered "does not exist" until the parent attribute cache expires; the old get(dir_path) then returns nullptr and EP_HOST_ASSERT kills the rank (and hangs the job). Loading the loser's own just-compiled artifacts (which are equivalent by cache key: kernel name, source, flags and compiler signature), caching them under the canonical dir_path, and deleting the tmp dir is the correct, root-cause-aligned fix. Warm starts and the per-process cache are unchanged; graceful on non-rename failures. The implementation in compiler.hpp/cache.hpp is correct and compiles.

v3

This PR fixes a real crash on the losing side of the JIT compile/rename race when the cache lives on a shared filesystem (NFS). Previously, after a failed rename the losing rank called get(dir_path) and asserted runtime != nullptr; on NFS the initial cache-miss can prime a negative lookup cache entry that outlives the failed rename, so the stat reports the winner's directory as missing and the rank dies. The fix makes the losing path self-sufficient: it loads the runtime from its own just-compiled temporary directory (equivalent artifacts, since the cache key hashes name+signature+flags+code), caches it under the canonical dir_path key via the new KernelRuntimeCache::put(), then removes the temp dir with safe_remove_all.

I verified the correctness claims against the code:

• KernelRuntime's constructor (kernel_runtime.hpp) consumes files eagerly: cuobjdump -symbols reads the cubin, then load_kernel() calls cudaLibraryLoadFromFile/cuModuleLoad and retains only GPU handles. So deleting the temp dir immediately after put() is safe.
• dir_path is derived solely from kernel_signature (name+signature+flags+code), so the losing rank's temp artifacts are semantically equivalent to the winner's — caching them under dir_path is sound.
• The winner path still uses get(dir_path)+assert, which is correct because the winner just renamed into dir_path and its own client has fresh positive attributes.
• safe_remove_all is used (consistent with existing usage and the segfault note) rather than std::filesystem::remove_all.

The change is minimal (2 files), well-documented, turns a fatal crash into graceful recovery on any rename failure, and leaves the writer-side fsync/atomic-rename hardening and warm-start behavior unchanged. LGTM.

Files reviewed: 2
Issues found: 🟡 1 warning | 🔵 3 suggestion
Inline comments posted: 4