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..b084604 100644 --- a/benchmarks/benchmark_history.json +++ b/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": "" + } } ] } 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/_cpp_ops.py b/src/torchada/_cpp_ops.py index 54b6921..b8250d5 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): + 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"], 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_cpp_ops_guard.py b/tests/test_cpp_ops_guard.py new file mode 100644 index 0000000..65e181b --- /dev/null +++ b/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() 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):