Fix(cpp_ops): recover from stale JIT build locks

[yafengio]

Summary

torch's JIT loader creates a FileBaton (<build_dir>/lock, a plain O_EXCL file) in every process and holds it while ninja runs. If a process is killed during that window, the file is left behind, and every later import torchada blocks forever in FileBaton.wait().

 load_cpp_ops / load_graph_rotation_ops
-  load(name=..., ...)
+  _locked_load(load, name=..., musa=..., ...)
+    loader has no FileBaton (newer torch: FileLock)  → load(...) unchanged
+    build_dir = loader._get_build_directory(name)
+    flock(<build_dir>/.torchada.flock)                 # released by the kernel when the holder dies
+      leftover <build_dir>/lock → remove + warn        # its holder must be dead
+      load(..., build_directory=build_dir)             # lock and build share one dir
     _jit_compile → FileBaton(lock) → ninja

• If flock is unsupported (ENOLCK on some NFS), it warns and leaves the lock alone, so behaviour matches torch's original.
• There are no new switches, no public API changes and no cache path changes. Scope: src/torchada/_cpp_ops.py and tests/test_cpp_ops_guard.py, +119/−3.

Evidence

Measured on MTT S5000 with torch_musa 2.11.0.post2+musa5.2.0.

kill -9 during a cold build, then 8 concurrent `import torchada`
  Before: all 8 block forever in FileBaton.wait (py-spy)
  After:  8/8 loaded, 1 compile (.ninja_log = 3 entries), 1 "removing stale JIT build lock" warning

zero-byte <build_dir>/lock left over a finished build
  Before: hangs (30 s, still in FileBaton.wait)
  After:  loads, no rebuild, lock removed

warm load_cpp_ops, 8 concurrent processes, median
  Before: 0.19 s   After: 0.11 s

pytest tests/test_cpp_ops_guard.py                  4 passed
pytest tests/test_cuda_patching.py -k CppOps        4 passed
pytest tests/test_graph_rotation.py                 7 passed

The same scenarios also pass on torch 2.7.1, 2.9.0 and 2.11.0.post1, where the loader API used is identical. An 8-GPU sglang prefill server starts and serves requests. These runs used earlier versions of this patch that rely on the same flock mechanism.

[yeahdongcn]

Merge verdict: no blocking findings — ready to merge.

Loader sources read at this head: torch 2.5.1–2.12 leave a lock baton, torch 2.14 does not; no earlier round raised findings, so auto-review for this commit closes here.

Findings: 0 blocking · 4 nit

nit	location	problem	fix
nit	src/torchada/_cpp_ops.py:108	warn precedes os.remove, so -W error skips the recovery	remove, then warn
nit	src/torchada/_cpp_ops.py:89	other JIT build dirs stay unguarded	share the guard
nit	tests/test_cpp_ops_guard.py:18	the musa=True branch is never executed	add a fake musa-loader case
nit	benchmarks/benchmark_history.json:6	the 0.1.94 row is relabelled 0.1.95	restore it; add a measured row

What changed · what I could not verify

Diff map (head 930d7f668384, base main): src/torchada/_cpp_ops.py +49/−3 (guard at head lines 75–110), new tests/test_cpp_ops_guard.py +70, and 0.1.94→0.1.95 in pyproject.toml, src/torchada/__init__.py, src/torchada/csrc/ops.h, README.md, README_CN.md, examples/extension_setup.py, benchmarks/benchmark_history.json, tests/test_platform.py. A code search on main reports 0.1.94 in exactly those 8 files, so the bump leaves no stale string. Head is the fix commit plus chore: bump version to 0.1.95 (35908bb).

Verified statically, from the loader sources:

loader	FileBaton	baton path	build-dir helper
torch 2.5.1 / 2.6.0 / 2.11.0 / 2.12.0	present	build_directory/'lock'	_get_build_directory(name, verbose), creates the dir
torch 2.14.0	absent (torch.utils._filelock)	with FileLock(...'lock')	same
torch_musa v2.7.1 / v2.9.1 / main	present	build_directory + "/lock"	_get_build_directory(name, verbose), creates the dir

• The hasattr(loader, "FileBaton") discriminator is real: torch switched to FileLock after the 2.12 line, and only the baton version leaves a lock file behind.
• FileBaton.wait() in 2.5.1–2.12 only spins on os.path.exists and never creates the file, so a leftover lock can only come from a builder killed between try_acquire and release — the PR's premise holds.
• Deleting the baton under the flock is safe for torchada's own processes: the baton is acquired and released strictly inside load() (try/finally in _jit_compile) and the flock spans that whole window, so a free flock means no live torchada builder. Both call sites pass sources as a keyword, so forwarding to load(..., build_directory=...) is valid.
• Passing build_directory is behaviour-preserving: both loaders do build_directory or _get_build_directory(name, verbose), and verbose there only controls logging.
• The MUSA branch resolves the same module it later builds with (torch_musa.utils.musa_extension, then torchada load → musa_ext.load), so flock and baton land in one directory.
• gh pr checks 128, verbatim: no checks reported on the 'fix/jit-stale-build-lock' branch. The repo has no .github/ (API 404), so there is no CI to pass.
• No instruction-like content in the PR body, commit messages, or comments; nothing attempted to steer this review.

Not verified: no MUSA run at this exact head. The body's S5000 numbers (kill -9 then 8 concurrent imports, pytest tests/test_cpp_ops_guard.py 4 passed) are, in its own words, from "earlier versions of this patch", so I checked the mechanism against the pinned loader sources instead. The new .torchada.flock file stays behind in the torch extension build dir; torch hashes sources rather than directory contents, so it is inert.

auto-review round 1/3 · verdict: clean

[yafengio]

@yeahdongcn Thanks for the review! Addressed:

1. Remove-then-warn, so -W error still clears the stale lock (861f7dd, with a test).
2. Out of scope: this PR only guards torchada's own extensions. Wrapping the public load/load_inline would change behaviour for every caller.
3. Added a fake torch_musa loader test for the musa=True path (861f7dd).
4. Restored the 0.1.94 row and added a measured 0.1.95 row from S5000 (566d795). The two rows use different torch versions (2.7.1 vs 2.11.0.post2), so they are not directly comparable.

Re-verified on S5000: test_cpp_ops_guard.py 6 passed, -k CppOps 4 passed.

[yeahdongcn]

Merge verdict: no blocking findings — ready to merge.

Guard, FileBaton discriminator and build dir all match the loader sources at head 566d79530853.

Findings: 0 blocking · 4 nit
Earlier: round 1/3 (930d7f66) raised 0 blocking · 4 nit; three fixed at this head, the scope nit re-listed.

nit	location	problem	fix
nit	src/torchada/_cpp_ops.py:107	exists-then-remove race can disable the operator overrides silently	catch FileNotFoundError
nit	src/torchada/_cpp_ops.py:89	other JIT build dirs stay unguarded	move the guard into torchada.utils.cpp_extension.load
nit	src/torchada/_cpp_ops.py:100	unbounded flock wait: a fork-inherited fd hangs every later import	LOCK_NB retry, then warn
nit	benchmarks/benchmark_history.json:85	row presented as measured, no run recorded	paste the --save output or drop the row

What changed · what I could not verify

Diff map (head 566d79530853, base main). src/torchada/_cpp_ops.py +49/−3 (guard at head lines 75–110); new tests/test_cpp_ops_guard.py (6 tests); 0.1.94→0.1.95 in pyproject.toml, src/torchada/__init__.py, src/torchada/csrc/ops.h, README.md, README_CN.md, examples/extension_setup.py, tests/test_platform.py; one new row in benchmarks/benchmark_history.json. Four commits: 35908bb (guard), 930d7f66 (version bump — the commit the earlier round reviewed), 861f7dd (remove before warn), 566d7953 (benchmark row).

Verified from the loader sources, not from the PR body:

loader	FileBaton in module	baton path	build-dir helper
torch 2.5.1 / 2.12.0 / 2.13.0	yes, from .file_baton import FileBaton	build_directory/'lock'	_get_build_directory(name, verbose), creates the dir
torch 2.14.0 / main	no, from ._filelock import FileLock	with FileLock(build_directory/'lock')	present
torch_musa v2.7.1 / v2.9.1 / main	yes, from torch.utils.file_baton import FileBaton	build_directory + "/lock"	present, creates the dir

• Premise holds: try_acquire is os.open(..., O_CREAT|O_EXCL) and wait() spins on os.path.exists(lock_file_path) with warn_after_seconds=None, so a leftover lock blocks forever — and _jit_compile takes the baton even when version == old_version, so a stale lock hangs a cached load too.
• Discriminator holds: 2.13 still leaves a lock file behind; 2.14+ replaced it with a kernel-released FileLock, so the pass-through at line 96 is correct there.
• Directory agreement holds: torchada's load accepts and forwards build_directory; torch's and torch_musa's _jit_compile both use build_directory or _get_build_directory(name, verbose); torch_musa's load signature carries the keyword. Flock and baton land in the directory the build actually uses.
• The flock spans all of load(), so no live torchada builder for that extension can hold the baton when it is removed; -W error no longer skips the removal (861f7dd).
• No instruction-like content in the body, commits, or comments; nothing attempted to steer this review.

Not verified. No test was run here — this sandbox has no importable torch. The body's pytest tests/test_cpp_ops_guard.py 4 passed predates the two commits that took the file to 6 tests, as the body itself says. gh pr checks 128 says verbatim: "no checks reported on the 'fix/jit-stale-build-lock' branch"; the tree has no .github/, so no CI exists for this head. The 0.1.95 benchmark row is plausible rather than reproduced: its counts are self-consistent with the script's mean_ns < 200 / 200 <= mean_ns < 800 thresholds (5 fast / 9 medium of 14 ops), its keys match save_results_to_history, and the 0.1.94 row is intact.

Unconfirmed risk — would need a concurrent pre-0.1.95 builder to confirm. The flock only proves that no torchada builder holds the baton; the baton itself is torch's, and build dirs are shared per user/python/accelerator. A concurrent process running an older torchada could hold a live baton that this code deletes, leaving two ninja writers in one directory. The periodic-cleanup alternative would be a bounded wait that degrades to today's behaviour instead of removing the file.

auto-review round 1/3 · verdict: clean