From 35908bb989f4dd5e38a807752c6a0860d5edee01 Mon Sep 17 00:00:00 2001 From: "yafeng.li" Date: Wed, 30 Sep 2026 15:51:18 +0800 Subject: [PATCH 1/4] fix(cpp_ops): recover from stale JIT build locks torch's JIT loader guards every build with a FileBaton, a plain file created with O_EXCL, and holds it while ninja runs in every new process. A process killed in that window leaves the file behind, and every later `import torchada` spins in FileBaton.wait() forever, e.g. an sglang launch that hangs before logging anything. Route the torchada_cpp_ops and graph-rotation load() calls through _locked_load. It takes an fcntl.flock on /.torchada.flock, which the kernel releases when its holder dies, so whoever holds it can delete a leftover `lock` (with a warning) before building. The build directory comes from the JIT loader in use (torch_musa's for MUSA sources, else torch's) and is passed to load() as build_directory, so the lock and the build share one directory; cache paths are unchanged. Where flock is unsupported (ENOLCK on some NFS) it warns and leaves the baton alone. Loaders without FileBaton (newer torch uses a FileLock that dies with its holder) are passed straight through. Every process building into the same directory must carry this change: an older torchada holds no flock, so its live baton would look stale. Verified on MTT S5000 with torch_musa 2.11.0.post2+musa5.2.0: - kill -9 during a cold build, then 8 concurrent imports: all load, one compile, one warning; unpatched code hangs in FileBaton.wait; - a pre-seeded stale lock recovers without a rebuild; - warm load_cpp_ops 0.110 s median over 8 processes (0.19 s before). The loader API used is present and identical on torch 2.7.1, 2.9.0 and 2.11.0.post1. --- src/torchada/_cpp_ops.py | 52 +++++++++++++++++++++++++-- tests/test_cpp_ops_guard.py | 70 +++++++++++++++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 3 deletions(-) create mode 100644 tests/test_cpp_ops_guard.py diff --git a/src/torchada/_cpp_ops.py b/src/torchada/_cpp_ops.py index 54b6921..26c2817 100644 --- a/src/torchada/_cpp_ops.py +++ b/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"], diff --git a/tests/test_cpp_ops_guard.py b/tests/test_cpp_ops_guard.py new file mode 100644 index 0000000..7c935b7 --- /dev/null +++ b/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() From 930d7f66838447d5ae0060875794ca1484bf71da Mon Sep 17 00:00:00 2001 From: "yafeng.li" Date: Fri, 9 Oct 2026 15:33:11 +0800 Subject: [PATCH 2/4] chore: bump version to 0.1.95 --- README.md | 2 +- README_CN.md | 2 +- benchmarks/benchmark_history.json | 2 +- examples/extension_setup.py | 2 +- pyproject.toml | 2 +- src/torchada/__init__.py | 2 +- src/torchada/csrc/ops.h | 2 +- tests/test_platform.py | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 242dce1..5faa68a 100644 --- a/README.md +++ b/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 diff --git a/README_CN.md b/README_CN.md index 77326ab..9cdf612 100644 --- a/README_CN.md +++ b/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：条件导入 diff --git a/benchmarks/benchmark_history.json b/benchmarks/benchmark_history.json index 6365975..3c5ea64 100644 --- a/benchmarks/benchmark_history.json +++ b/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", diff --git a/examples/extension_setup.py b/examples/extension_setup.py index 68d1c6c..babc9dd 100755 --- a/examples/extension_setup.py +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index d4379e1..b1b5f83 100644 --- a/pyproject.toml +++ b/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"} diff --git a/src/torchada/__init__.py b/src/torchada/__init__.py index 3e12a18..66b522b 100644 --- a/src/torchada/__init__.py +++ b/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 diff --git a/src/torchada/csrc/ops.h b/src/torchada/csrc/ops.h index 397a47f..0b4198b 100644 --- a/src/torchada/csrc/ops.h +++ b/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) { diff --git a/tests/test_platform.py b/tests/test_platform.py index d5457d6..c936f99 100644 --- a/tests/test_platform.py +++ b/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): From 861f7dd4400d491805312b2865c373488374ec74 Mon Sep 17 00:00:00 2001 From: "yafeng.li" Date: Fri, 9 Oct 2026 17:02:49 +0800 Subject: [PATCH 3/4] fix(cpp_ops): remove the stale baton before warning about it Under `-W error` the warning raises, so warning first skipped the removal and left the baton for the next start too. Remove it first, then warn. Also cover the MUSA loader branch with a fake torch_musa module, and the `-W error` ordering. --- src/torchada/_cpp_ops.py | 2 +- tests/test_cpp_ops_guard.py | 41 +++++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/torchada/_cpp_ops.py b/src/torchada/_cpp_ops.py index 26c2817..b8250d5 100644 --- a/src/torchada/_cpp_ops.py +++ b/src/torchada/_cpp_ops.py @@ -105,8 +105,8 @@ def _locked_load(load, name: str, musa: bool, **kwargs): 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) + warnings.warn(f"torchada: removed stale JIT build lock {baton}") return load(name=name, build_directory=build_dir, **kwargs) diff --git a/tests/test_cpp_ops_guard.py b/tests/test_cpp_ops_guard.py index 7c935b7..65e181b 100644 --- a/tests/test_cpp_ops_guard.py +++ b/tests/test_cpp_ops_guard.py @@ -2,6 +2,9 @@ import errno import fcntl +import sys +import types +import warnings import pytest import torch.utils.cpp_extension as cpp_ext @@ -68,3 +71,41 @@ def test_filelock_based_torch_is_left_alone(ext_dir, monkeypatch): 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() From 566d79530853e55fd12ccf0cc52d61ee77a4300f Mon Sep 17 00:00:00 2001 From: "yafeng.li" Date: Fri, 9 Oct 2026 17:06:04 +0800 Subject: [PATCH 4/4] chore(benchmarks): keep the 0.1.94 row and add a measured 0.1.95 row Restore the 0.1.94 entry the version bump had relabelled, and append a 0.1.95 run of benchmark_overhead.py --save on MTT S5000 with torch / torch_musa 2.11.0.post2+musa5.2.0. The 0.1.94 row was taken on torch 2.7.1, so the two rows are not a like-for-like comparison. --- benchmarks/benchmark_history.json | 86 ++++++++++++++++++++++++++++++- 1 file changed, 85 insertions(+), 1 deletion(-) diff --git a/benchmarks/benchmark_history.json b/benchmarks/benchmark_history.json index 3c5ea64..b084604 100644 --- a/benchmarks/benchmark_history.json +++ b/benchmarks/benchmark_history.json @@ -3,7 +3,7 @@ "description": "Historical benchmark results for torchada performance tracking", "results": [ { - "version": "0.1.95", + "version": "0.1.94", "date": "2026-01-29", "platform": "MUSA", "pytorch_version": "2.7.1", @@ -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": "" + } } ] }