From 01608553b2937a23c0828f166f9310e4cf26ab64 Mon Sep 17 00:00:00 2001 From: Aionw Date: Mon, 24 Aug 2026 16:35:30 +0800 Subject: [PATCH] [TransferEngine] Migrate shared segment to unified Python package --- .github/CODEOWNERS | 4 +- .github/labeler.yml | 6 +- mooncake-integration/CMakeLists.txt | 2 +- .../mooncake}/shared_segment.py | 48 +++- python/tests/engine/test_shared_segment.py | 264 ++++++++++++++++++ .../tests/packaging/test_installed_wheel.py | 7 + python/tests/packaging/test_project.py | 25 +- scripts/build_wheel.sh | 2 +- 8 files changed, 339 insertions(+), 19 deletions(-) rename {mooncake-integration => python/mooncake}/shared_segment.py (90%) create mode 100644 python/tests/engine/test_shared_segment.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2c9f901b16..860b3ef034 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -12,9 +12,11 @@ /mooncake-integration/transfer_engine @ShangmingCai @alogfans /mooncake-integration/store @ykwd @stmatengss @zxpdemonio /mooncake-pg @UNIDY2002 @ympcMark @yuechen-sys -/mooncake-reshard @ShangmingCai @stmatengss @Bo-Vincent @zxpdemonio +/mooncake-reshard @ShangmingCai @stmatengss @Bo-Vincent @zxpdemonio /mooncake-store @ykwd @stmatengss @XucSh @YiXR /mooncake-store/*/ha/ @Libotry @YiXR @00fish0 @Icedcoco @Aionw +/python/mooncake/shared_segment.py @ShangmingCai @alogfans +/python/tests/engine/ @ShangmingCai @alogfans /mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q @staryxchen /mooncake-transfer-engine/tent @alogfans @doujiang24 @chestnut-Q @staryxchen @00fish0 @dtcccc /mooncake-transfer-engine/*/transport/hip_transport/ @alogfans @amd-arozanov diff --git a/.github/labeler.yml b/.github/labeler.yml index 92650848b7..970644152b 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -9,7 +9,10 @@ store: Transfer Engine: - changed-files: - - any-glob-to-any-file: 'mooncake-transfer-engine/**/*' + - any-glob-to-any-file: + - 'mooncake-transfer-engine/**/*' + - 'python/mooncake/shared_segment.py' + - 'python/tests/engine/**/*' PyTorch Backend: - changed-files: @@ -57,6 +60,7 @@ Tests: - 'scripts/test_*' - 'mooncake-wheel/tests/**/*' - 'mooncake-reshard/tests/**/*' + - 'python/tests/**/*' - 'scripts/tone_tests/**/*' Ascend/NPU: diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 0ef9dfde04..d9d89a21e9 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -187,7 +187,7 @@ install( if(WITH_TE) install( - FILES "${CMAKE_CURRENT_SOURCE_DIR}/shared_segment.py" + FILES "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/shared_segment.py" DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR} COMPONENT python) endif() diff --git a/mooncake-integration/shared_segment.py b/python/mooncake/shared_segment.py similarity index 90% rename from mooncake-integration/shared_segment.py rename to python/mooncake/shared_segment.py index 5483c0cb09..a550525e78 100644 --- a/mooncake-integration/shared_segment.py +++ b/python/mooncake/shared_segment.py @@ -32,8 +32,14 @@ from dataclasses import dataclass from typing import Any -import torch -import torch.distributed as dist +try: + import torch + import torch.distributed as dist +except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + torch = None + dist = None from .engine import SharedSegment as _CppSharedSegment @@ -51,6 +57,15 @@ class SharedSegmentError(RuntimeError): """Raised when a shared segment cannot be created or addressed.""" +def _require_torch() -> Any: + if torch is None: + raise SharedSegmentError( + "PyTorch is required to create shared-segment tensor views; " + "install mooncake-transfer-engine[hardware]" + ) + return torch + + def create_shared_segment( name: str, blocks: Mapping[str, Mapping[str, Any]], @@ -84,6 +99,7 @@ def create_shared_segment( comm_group = _select_comm_group(comm_group, tp_group) if host_register and not mmap: raise SharedSegmentError("host_register requires mmap=True") + _require_torch() if not shared_segment_supported(mmap=mmap, host_register=host_register): if not mmap: raise SharedSegmentError( @@ -127,9 +143,7 @@ def create_shared_segment( except RuntimeError as exc: create_error = str(exc) - _raise_if_any_rank_failed( - create_error, world_size, rank_id, comm_group, "create" - ) + _raise_if_any_rank_failed(create_error, world_size, rank_id, comm_group, "create") if segment is None: raise SharedSegmentError("Shared segment create returned no segment") @@ -150,7 +164,10 @@ def create_shared_segment( class SharedSegment: - """Handle to one shared host segment. Alive for as long as the mapping is needed.""" + """Handle to one shared host segment. + + The handle must remain alive for as long as its mapping is needed. + """ def __init__( self, @@ -248,7 +265,7 @@ def _parse_block(name: str, spec: Mapping[str, Any]) -> _BlockSpec: if not shape or any(dim <= 0 for dim in shape): raise SharedSegmentError(f"Block {name!r} needs a positive shape") - element_size = torch.empty(0, dtype=dtype).element_size() + element_size = _require_torch().empty(0, dtype=dtype).element_size() nbytes = element_size for dim in shape: nbytes *= dim @@ -272,14 +289,15 @@ def _build_layout( def _current_device_index() -> int: """The accelerator this rank runs on, whatever vendor it is.""" - accelerator = getattr(torch, "accelerator", None) + torch_module = _require_torch() + accelerator = getattr(torch_module, "accelerator", None) if accelerator is not None: try: return int(accelerator.current_device_index()) except (RuntimeError, AttributeError): pass for backend in ("npu",): - module = getattr(torch, backend, None) + module = getattr(torch_module, backend, None) if module is not None and module.is_available(): return int(module.current_device()) return 0 @@ -335,8 +353,9 @@ def _raise_if_any_rank_failed( def _all_gather_blob(blob: bytes, world_size: int, comm_group: Any) -> list[bytes]: - local = torch.frombuffer(bytearray(blob), dtype=torch.uint8) - gathered = [torch.empty_like(local) for _ in range(world_size)] + torch_module = _require_torch() + local = torch_module.frombuffer(bytearray(blob), dtype=torch_module.uint8) + gathered = [torch_module.empty_like(local) for _ in range(world_size)] dist.all_gather(gathered, local, group=_resolve_cpu_group(comm_group)) return [bytes(tensor.numpy()) for tensor in gathered] @@ -357,6 +376,7 @@ def _tensor_from_device_ptr( holder: Any, ) -> torch.Tensor: """Wrap a device-accessible VA (HostRegister or Ascend VMM) as an NPU tensor.""" + torch_module = _require_torch() try: import torch_npu except ImportError as exc: @@ -364,11 +384,11 @@ def _tensor_from_device_ptr( "torch_npu is required to expose device-mapped shared-segment tensors" ) from exc - device = torch.device(f"npu:{device_id}") + device = torch_module.device(f"npu:{device_id}") storage = torch_npu._C._construct_storage_from_data_pointer( int(device_ptr), device, int(nbytes) ) - view = torch.empty(0, dtype=dtype, device=device) + view = torch_module.empty(0, dtype=dtype, device=device) view.set_(storage, 0, shape, _contiguous_strides(shape)) view._mooncake_segment = holder return view @@ -378,6 +398,6 @@ def _tensor_from_host_addr( host_addr: int, nbytes: int, shape: tuple[int, ...], dtype: torch.dtype, holder: Any ) -> torch.Tensor: buffer = (ctypes.c_int8 * nbytes).from_address(host_addr) - view = torch.frombuffer(buffer, dtype=dtype).view(shape) + view = _require_torch().frombuffer(buffer, dtype=dtype).view(shape) view._mooncake_segment = holder return view diff --git a/python/tests/engine/test_shared_segment.py b/python/tests/engine/test_shared_segment.py new file mode 100644 index 0000000000..c0f47ad7a9 --- /dev/null +++ b/python/tests/engine/test_shared_segment.py @@ -0,0 +1,264 @@ +from __future__ import annotations + +import builtins +import gc +import importlib.util +from pathlib import Path +import sys +import types +from typing import Any +import weakref + +import pytest + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +PACKAGE_ROOT = REPOSITORY_ROOT / "python" / "mooncake" +MODULE_PATH = PACKAGE_ROOT / "shared_segment.py" + + +class _FakeDType: + def __init__(self, name: str = "torch.float16", element_size: int = 2) -> None: + self.name = name + self._element_size = element_size + + def __str__(self) -> str: + return self.name + + +class _FakeEmptyTensor: + def __init__(self, dtype: _FakeDType) -> None: + self._dtype = dtype + + def element_size(self) -> int: + return self._dtype._element_size + + +class _NativeSegment: + def __init__(self, harness: types.SimpleNamespace, size: int) -> None: + self._harness = harness + self._size = size + + def __del__(self) -> None: + self._harness.release_count += 1 + + def complete(self, blobs: list[bytes]) -> None: + self._harness.completed_blobs.append(blobs) + if self._harness.complete_error: + raise RuntimeError(self._harness.complete_error) + + def base_addr(self) -> int: + return 0x1000 + + def device_addr(self) -> int: + return 0 + + def size(self) -> int: + return self._size + + +def _fake_torch_modules() -> tuple[types.ModuleType, types.ModuleType]: + distributed = types.ModuleType("torch.distributed") + distributed.all_gather = lambda *args, **kwargs: None + distributed.all_gather_object = lambda *args, **kwargs: None + + torch = types.ModuleType("torch") + torch.__path__ = [] + torch.Tensor = object + torch.dtype = _FakeDType + torch.uint8 = _FakeDType("torch.uint8", 1) + torch.distributed = distributed + torch.empty = lambda size, *, dtype, **kwargs: _FakeEmptyTensor(dtype) + return torch, distributed + + +def _load_module( + monkeypatch: pytest.MonkeyPatch, *, with_torch: bool = True +) -> tuple[types.ModuleType, types.SimpleNamespace]: + harness = types.SimpleNamespace( + complete_error="", + completed_blobs=[], + create_calls=[], + created_ref=None, + release_count=0, + supported_calls=[], + supported_result=True, + ) + + class NativeBinding: + @staticmethod + def supported(mmap: bool, host_register: bool) -> bool: + harness.supported_calls.append((mmap, host_register)) + return harness.supported_result + + @staticmethod + def create(*args: Any) -> tuple[_NativeSegment, bytes]: + harness.create_calls.append(args) + segment = _NativeSegment(harness, args[1]) + harness.created_ref = weakref.ref(segment) + return segment, b"owner-blob" + + package = types.ModuleType("mooncake") + package.__path__ = [str(PACKAGE_ROOT)] + engine = types.ModuleType("mooncake.engine") + engine.SharedSegment = NativeBinding + monkeypatch.setitem(sys.modules, "mooncake", package) + monkeypatch.setitem(sys.modules, "mooncake.engine", engine) + monkeypatch.delitem(sys.modules, "mooncake.shared_segment", raising=False) + + torch_modules = [ + name for name in sys.modules if name == "torch" or name.startswith("torch.") + ] + for name in torch_modules: + monkeypatch.delitem(sys.modules, name) + + if with_torch: + torch, distributed = _fake_torch_modules() + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "torch.distributed", distributed) + + spec = importlib.util.spec_from_file_location( + "mooncake.shared_segment", MODULE_PATH + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + + if with_torch: + spec.loader.exec_module(module) + else: + original_import = builtins.__import__ + + def import_without_torch( + name: str, + globals: dict[str, Any] | None = None, + locals: dict[str, Any] | None = None, + fromlist: tuple[str, ...] = (), + level: int = 0, + ) -> Any: + if name == "torch" or name.startswith("torch."): + raise ModuleNotFoundError("No module named 'torch'", name="torch") + return original_import(name, globals, locals, fromlist, level) + + with monkeypatch.context() as context: + context.setattr(builtins, "__import__", import_without_torch) + spec.loader.exec_module(module) + + return module, harness + + +def _blocks() -> dict[str, dict[str, Any]]: + dtype = _FakeDType() + return { + "z": {"count": 3, "shape": (4,), "dtype": dtype}, + "a": {"count": 2, "shape": (2, 3), "dtype": dtype}, + } + + +def test_create_shared_segment_builds_sorted_layout_and_completes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module, harness = _load_module(monkeypatch) + + segment = module.create_shared_segment( + "cache", + _blocks(), + world_size=1, + rank_id=0, + device_id=7, + ) + + assert harness.create_calls == [ + ( + "cache|a:2:2x3:torch.float16:12;z:3:4:torch.float16:8", + 24_576, + 1, + 0, + 0, + 7, + True, + False, + ) + ] + assert harness.completed_blobs == [[b"owner-blob"]] + assert segment.base_addr() == 0x1000 + assert segment.total_size() == 24_576 + + +@pytest.mark.parametrize( + ("mmap", "host_register", "message"), + [ + (False, False, "no VMM backend"), + (True, True, "cannot HostRegister"), + (True, False, "no mmap shared-segment backend"), + ], +) +def test_create_shared_segment_reports_missing_native_capabilities( + monkeypatch: pytest.MonkeyPatch, + mmap: bool, + host_register: bool, + message: str, +) -> None: + module, harness = _load_module(monkeypatch) + harness.supported_result = False + + with pytest.raises(module.SharedSegmentError, match=message): + module.create_shared_segment( + "cache", + _blocks(), + world_size=1, + rank_id=0, + mmap=mmap, + host_register=host_register, + ) + + assert harness.supported_calls == [(mmap, host_register)] + assert not harness.create_calls + + +def test_host_registration_requires_mmap(monkeypatch: pytest.MonkeyPatch) -> None: + module, harness = _load_module(monkeypatch) + + with pytest.raises(module.SharedSegmentError, match="requires mmap=True"): + module.create_shared_segment( + "cache", + _blocks(), + world_size=1, + rank_id=0, + mmap=False, + host_register=True, + ) + + assert not harness.supported_calls + + +def test_failed_completion_releases_native_segment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module, harness = _load_module(monkeypatch) + harness.complete_error = "mapping failed" + + def create() -> None: + with pytest.raises(module.SharedSegmentError, match="mapping failed"): + module.create_shared_segment( + "cache", _blocks(), world_size=1, rank_id=0, device_id=0 + ) + + create() + gc.collect() + + assert harness.created_ref() is None + assert harness.release_count == 1 + + +def test_capability_check_does_not_require_optional_torch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module, harness = _load_module(monkeypatch, with_torch=False) + + assert module.shared_segment_supported(mmap=True, host_register=False) + with pytest.raises(module.SharedSegmentError, match=r"\[hardware\]"): + module.create_shared_segment("cache", {}, world_size=1, rank_id=0, device_id=0) + + assert harness.supported_calls == [(True, False)] + assert not harness.create_calls diff --git a/python/tests/packaging/test_installed_wheel.py b/python/tests/packaging/test_installed_wheel.py index 74422d9ebb..cecb28cc73 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/python/tests/packaging/test_installed_wheel.py @@ -4,6 +4,7 @@ from pathlib import Path import subprocess import sys +import zipfile import pytest @@ -28,6 +29,8 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: wheel = Path(wheel_value).resolve() assert wheel.is_file(), f"wheel does not exist: {wheel}" + with zipfile.ZipFile(wheel) as archive: + assert "mooncake/shared_segment.py" in archive.namelist() environment = tmp_path / "environment" subprocess.run([sys.executable, "-m", "venv", str(environment)], check=True) @@ -46,6 +49,7 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: import mooncake import mooncake.engine import mooncake.reshard +import mooncake.shared_segment import mooncake.store package_path = Path(mooncake.__file__).resolve() @@ -54,6 +58,9 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: assert metadata.version("mooncake-transfer-engine") == {_project_version()!r} assert mooncake.BufferPool is mooncake.store.BufferPool assert mooncake.engine.TransferEngine is not None +assert mooncake.shared_segment.SharedSegment is not None +assert isinstance(mooncake.shared_segment.shared_segment_supported(), bool) +assert Path(mooncake.shared_segment.__file__).resolve().parent == package_path.parent """ subprocess.run( [str(python), "-I", "-c", smoke_script], diff --git a/python/tests/packaging/test_project.py b/python/tests/packaging/test_project.py index 4ed4b376b8..97bb50741d 100644 --- a/python/tests/packaging/test_project.py +++ b/python/tests/packaging/test_project.py @@ -40,13 +40,16 @@ def test_dependency_boundaries_are_declared() -> None: metadata = project["project"] assert set(metadata["dependencies"]) == {"aiohttp", "msgpack", "requests"} - assert set(metadata["optional-dependencies"]) == { + optional_dependencies = metadata["optional-dependencies"] + assert set(optional_dependencies) == { "administration", "dev", "hardware", "structured", "vllm", } + assert "torch" not in metadata["dependencies"] + assert "torch" in optional_dependencies["hardware"] def test_tracked_source_roots_contain_no_generated_native_artifacts() -> None: @@ -57,6 +60,26 @@ def test_tracked_source_roots_contain_no_generated_native_artifacts() -> None: assert not list((REPOSITORY_ROOT / "mooncake-pg" / "torch").rglob("*.so")) +def test_shared_segment_has_one_authoritative_source() -> None: + canonical_source = REPOSITORY_ROOT / "python" / "mooncake" / "shared_segment.py" + legacy_source = REPOSITORY_ROOT / "mooncake-integration" / "shared_segment.py" + + assert canonical_source.is_file() + assert not legacy_source.exists() + assert ( + REPOSITORY_ROOT / "python" / "tests" / "engine" / "test_shared_segment.py" + ).is_file() + + integration_cmake = ( + REPOSITORY_ROOT / "mooncake-integration" / "CMakeLists.txt" + ).read_text() + legacy_build_script = (REPOSITORY_ROOT / "scripts" / "build_wheel.sh").read_text() + assert "../python/mooncake/shared_segment.py" in integration_cmake + assert "mooncake-integration/shared_segment.py" not in integration_cmake + assert "cp python/mooncake/shared_segment.py" in legacy_build_script + assert "mooncake-integration/shared_segment.py" not in legacy_build_script + + def test_pg_extension_build_stages_outside_the_source_tree( tmp_path: Path, ) -> None: diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index eed257699d..851deb919c 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -49,7 +49,7 @@ if [ -n "$EP_HOST_SO" ]; then fi # Copy the shared segment wrapper, which builds on engine.so -cp mooncake-integration/shared_segment.py mooncake-wheel/mooncake/shared_segment.py +cp python/mooncake/shared_segment.py mooncake-wheel/mooncake/shared_segment.py # Copy libasio.so to mooncake directory (runtime dependency of engine.so) cp ${BUILD_DIR}/mooncake-common/libasio.so mooncake-wheel/mooncake/libasio.so