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

@@ -80,6 +80,90 @@
         "medium_ops_count": 6,
         "notes": ""
       }
+    },
+    {
+      "version": "0.1.95",
+      "date": "2026-10-09",
+      "platform": "MUSA",
+      "pytorch_version": "2.11.0.post2+musa5.2.0",
+      "torch_musa_version": "2.11.0.post2+musa5.2.0",
+      "operations": {
+        "torch.cuda.device_count()": {
+          "mean_ns": 209,
+          "median_ns": 202,
+          "min_ns": 182
+        },
+        "torch.cuda.current_device()": {
+          "mean_ns": 519,
+          "median_ns": 515,
+          "min_ns": 474
+        },
+        "torch.cuda.is_available() [NOT redirected]": {
+          "mean_ns": 609,
+          "median_ns": 606,
+          "min_ns": 575
+        },
+        "torch.cuda.Stream (attr)": {
+          "mean_ns": 177,
+          "median_ns": 172,
+          "min_ns": 161
+        },
+        "torch.cuda.Event (attr)": {
+          "mean_ns": 177,
+          "median_ns": 172,
+          "min_ns": 161
+        },
+        "cudart.cudaHostRegister (attr)": {
+          "mean_ns": 109,
+          "median_ns": 111,
+          "min_ns": 100
+        },
+        "torch.device('cuda')": {
+          "mean_ns": 599,
+          "median_ns": 595,
+          "min_ns": 535
+        },
+        "torch.device('cuda:0')": {
+          "mean_ns": 610,
+          "median_ns": 596,
+          "min_ns": 545
+        },
+        "torch.device('cuda', 0)": {
+          "mean_ns": 634,
+          "median_ns": 626,
+          "min_ns": 565
+        },
+        "cpu_tensor.is_cuda (property)": {
+          "mean_ns": 438,
+          "median_ns": 434,
+          "min_ns": 393
+        },
+        "gpu_tensor.is_cuda (property)": {
+          "mean_ns": 434,
+          "median_ns": 434,
+          "min_ns": 384
+        },
+        "_translate_device('cuda')": {
+          "mean_ns": 181,
+          "median_ns": 182,
+          "min_ns": 161
+        },
+        "_translate_device('cuda:0')": {
+          "mean_ns": 179,
+          "median_ns": 172,
+          "min_ns": 161
+        },
+        "torch.backends.cuda.is_built()": {
+          "mean_ns": 263,
+          "median_ns": 262,
+          "min_ns": 242
+        }
+      },
+      "summary": {
+        "fast_ops_count": 5,
+        "medium_ops_count": 9,
+        "notes": ""
+      }
     }
   ]
 }

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):
+                os.remove(baton)
+                warnings.warn(f"torchada: removed stale JIT build lock {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,111 @@
+"""Tests for the flock guard around torchada's JIT C++ extension builds."""
+
+import errno
+import fcntl
+import sys
+import types
+import warnings
+
+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()
+
+
+def test_baton_is_removed_even_when_warnings_are_errors(ext_dir):
+    """Under ``-W error`` the warning raises, but only after the stale baton is gone."""
+    ext_dir.mkdir()
+    (ext_dir / "lock").write_bytes(b"")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        with pytest.raises(UserWarning, match="stale JIT build lock"):
+            _load(lambda **kwargs: None)
+    assert not (ext_dir / "lock").exists()
+
+
+def test_musa_loader_build_directory_is_used(tmp_path, monkeypatch):
+    """MUSA sources lock and build in torch_musa's build directory."""
+    musa_dir = tmp_path / "musa" / EXT_NAME
+
+    def get_build_directory(name, verbose):
+        musa_dir.mkdir(parents=True, exist_ok=True)
+        return str(musa_dir)
+
+    loader = types.ModuleType("torch_musa.utils.musa_extension")
+    loader.FileBaton = object
+    loader._get_build_directory = get_build_directory
+    utils = types.ModuleType("torch_musa.utils")
+    utils.musa_extension = loader
+    package = types.ModuleType("torch_musa")
+    package.utils = utils
+    for module in (package, utils, loader):
+        monkeypatch.setitem(sys.modules, module.__name__, module)
+
+    musa_dir.mkdir(parents=True)
+    (musa_dir / "lock").write_bytes(b"")
+    with pytest.warns(UserWarning, match="stale JIT build lock"):
+        kwargs = _cpp_ops._locked_load(lambda **kw: kw, name=EXT_NAME, musa=True)
+    assert kwargs["build_directory"] == str(musa_dir)
+    assert not (musa_dir / "lock").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):