Fix(cpp_ops): recover from stale JIT build locks

README.md

@@ -408,7 +408,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain.
 
 ```
 # pyproject.toml or requirements.txt
-torchada>=0.1.94
+torchada>=0.1.95
 ```
 
 ### Step 2: Conditional Import

README_CN.md

@@ -387,7 +387,7 @@ if torchada.is_gpu_device(device):  # 在 CUDA 和 MUSA 上都能工作
 
 ```
 # pyproject.toml 或 requirements.txt
-torchada>=0.1.94
+torchada>=0.1.95
 ```
 
 ### 步骤 2：条件导入

benchmarks/benchmark_history.json

@@ -3,7 +3,7 @@
   "description": "Historical benchmark results for torchada performance tracking",
   "results": [
     {
-      "version": "0.1.94",
+      "version": "0.1.95",
       "date": "2026-01-29",
       "platform": "MUSA",
       "pytorch_version": "2.7.1",

examples/extension_setup.py

@@ -69,7 +69,7 @@ def get_extensions():
     if extensions:
         setup(
             name="my_cuda_extension",
-            version="0.1.94",
+            version="0.1.95",
             ext_modules=extensions,
             cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},
             python_requires=">=3.8",

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchada"
-version = "0.1.94"
+version = "0.1.95"
 description = "Adapter package for torch_musa to act exactly like PyTorch CUDA"
 readme = "README.md"
 license = {text = "MIT"}

src/torchada/__init__.py

@@ -24,7 +24,7 @@
     from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME
 """
 
-__version__ = "0.1.94"
+__version__ = "0.1.95"
 
 from . import cuda, utils
 

src/torchada/_cpp_ops.py

@@ -24,6 +24,8 @@
 _graph_rotation_module: Optional[object] = None
 _musa_arch_cached: Optional[str] = None
 
+_FLOCK_FILE = ".torchada.flock"
+
 
 def _detect_musa_arch() -> str:
     """
@@ -70,6 +72,44 @@ def _detect_musa_arch() -> str:
     return arch
 
 
+def _jit_loader_module(musa: bool):
+    """torch_musa's JIT loader for MUSA sources when importable, else torch's."""
+    if musa:
+        try:
+            import torch_musa.utils.musa_extension as loader
+
+            return loader
+        except ImportError:
+            pass
+    import torch.utils.cpp_extension as loader
+
+    return loader
+
+
+def _locked_load(load, name: str, musa: bool, **kwargs):
+    """Call a JIT ``load`` under a flock that, unlike torch's ``lock`` baton, dies with its holder."""
+    import fcntl
+    import warnings
+
+    loader = _jit_loader_module(musa)
+    if not hasattr(loader, "FileBaton"):  # newer torch guards JIT builds with a FileLock itself
+        return load(name=name, **kwargs)
+    build_dir = loader._get_build_directory(name, False)
+    with open(os.path.join(build_dir, _FLOCK_FILE), "a") as lock_file:
+        try:
+            fcntl.flock(lock_file, fcntl.LOCK_EX)
+        except OSError as e:
+            warnings.warn(
+                f"torchada: flock unavailable on {lock_file.name}, building without it: {e}"
+            )
+        else:
+            baton = os.path.join(build_dir, "lock")
+            if os.path.exists(baton):
+                warnings.warn(f"torchada: removing stale JIT build lock {baton}")
+                os.remove(baton)
+        return load(name=name, build_directory=build_dir, **kwargs)
+
+
 def load_cpp_ops(force_reload: bool = False) -> Optional[object]:
     """
     Load the C++ operator overrides extension.
@@ -144,8 +184,10 @@ def load_cpp_ops(force_reload: bool = False) -> Optional[object]:
                 mtgpu_target = _detect_musa_arch()
             extra_cuda_cflags.append(f"--offload-arch={mtgpu_target}")
 
-            _cpp_ops_module = load(
+            _cpp_ops_module = _locked_load(
+                load,
                 name="torchada_cpp_ops",
+                musa=True,
                 sources=all_sources,
                 extra_include_paths=[csrc_dir],
                 extra_cuda_cflags=extra_cuda_cflags,
@@ -155,8 +197,10 @@ def load_cpp_ops(force_reload: bool = False) -> Optional[object]:
             # Pure C++ extension - use torch's loader directly
             from torch.utils.cpp_extension import load
 
-            _cpp_ops_module = load(
+            _cpp_ops_module = _locked_load(
+                load,
                 name="torchada_cpp_ops",
+                musa=False,
                 sources=all_sources,
                 extra_include_paths=[csrc_dir],
                 verbose=verbose,
@@ -211,8 +255,10 @@ def load_graph_rotation_ops(force_reload: bool = False) -> Optional[object]:
         libdir = osp.join(tm, "lib")
         verbose = os.environ.get("TORCHADA_CPP_OPS_VERBOSE") == "1"
 
-        _graph_rotation_module = load(
+        _graph_rotation_module = _locked_load(
+            load,
             name="torchada_graph_rotation_ops",
+            musa=False,
             sources=[src],
             extra_include_paths=include_dirs,
             extra_cflags=["-O2"],

src/torchada/csrc/ops.h

@@ -40,7 +40,7 @@
 namespace torchada {
 
 // Version information
-constexpr const char* VERSION = "0.1.94";
+constexpr const char* VERSION = "0.1.95";
 
 // Check if operator override is enabled via environment variable
 inline bool is_override_enabled(const char* op_name) {

tests/test_cpp_ops_guard.py

@@ -0,0 +1,70 @@
+"""Tests for the flock guard around torchada's JIT C++ extension builds."""
+
+import errno
+import fcntl
+
+import pytest
+import torch.utils.cpp_extension as cpp_ext
+
+from torchada import _cpp_ops
+
+EXT_NAME = "fake_ext"
+
+
+@pytest.fixture
+def ext_dir(tmp_path, monkeypatch):
+    """Point torch's JIT build root at tmp_path and make torch look FileBaton-based."""
+    monkeypatch.setenv("TORCH_EXTENSIONS_DIR", str(tmp_path))
+    monkeypatch.setattr(cpp_ext, "FileBaton", object, raising=False)
+    return tmp_path / EXT_NAME
+
+
+def _load(fake_load):
+    return _cpp_ops._locked_load(fake_load, name=EXT_NAME, musa=False)
+
+
+def test_stale_baton_is_removed(ext_dir):
+    """A torch ``lock`` left by a killed builder is deleted before the build."""
+    ext_dir.mkdir()
+    (ext_dir / "lock").write_bytes(b"")
+
+    def fake_load(**kwargs):
+        assert not (ext_dir / "lock").exists()
+        return kwargs
+
+    with pytest.warns(UserWarning, match="stale JIT build lock"):
+        kwargs = _load(fake_load)
+    assert kwargs == {"name": EXT_NAME, "build_directory": str(ext_dir)}
+
+
+def test_flock_is_held_during_load(ext_dir):
+    """Another opener of the flock file is excluded while the build runs."""
+
+    def fake_load(**kwargs):
+        with open(ext_dir / _cpp_ops._FLOCK_FILE) as other:
+            with pytest.raises(BlockingIOError):
+                fcntl.flock(other, fcntl.LOCK_EX | fcntl.LOCK_NB)
+
+    _load(fake_load)
+
+
+def test_flock_unavailable_keeps_the_baton(ext_dir, monkeypatch):
+    """Without flock (e.g. ENOLCK on NFS) the baton cannot be proven stale, so it stays."""
+    ext_dir.mkdir()
+    (ext_dir / "lock").write_bytes(b"")
+
+    def no_flock(*args):
+        raise OSError(errno.ENOLCK, "No locks available")
+
+    monkeypatch.setattr(fcntl, "flock", no_flock)
+    with pytest.warns(UserWarning, match="flock unavailable"):
+        _load(lambda **kwargs: None)
+    assert (ext_dir / "lock").exists()
+
+
+def test_filelock_based_torch_is_left_alone(ext_dir, monkeypatch):
+    """Newer torch locks JIT builds itself, so the call passes straight through."""
+    monkeypatch.delattr(cpp_ext, "FileBaton")
+
+    assert _load(lambda **kwargs: kwargs) == {"name": EXT_NAME}
+    assert not (ext_dir / _cpp_ops._FLOCK_FILE).exists()

tests/test_platform.py

@@ -59,7 +59,7 @@ def test_get_version(self):
 
         version = torchada.get_version()
         assert version == torchada.__version__
-        assert version == "0.1.94"
+        assert version == "0.1.95"
         assert isinstance(version, str)
 
     def test_project_version_matches_runtime_version(self):