From c901b40b3296f6c705f93da3d2bf813af1e23cab Mon Sep 17 00:00:00 2001 From: Aionw Date: Mon, 24 Aug 2026 16:38:20 +0800 Subject: [PATCH] [Integration] Migrate allocators to unified Python package --- .github/CODEOWNERS | 6 +- .github/labeler.yml | 10 +- .gitignore | 1 + mooncake-integration/CMakeLists.txt | 7 +- python/CMakeLists.txt | 5 +- .../mooncake}/allocator.py | 6 +- .../mooncake}/allocator_ascend_npu.py | 2 + .../mooncake}/fabric_allocator_utils.py | 2 + python/tests/allocators/conftest.py | 75 ++++++++++++ python/tests/allocators/test_allocator.py | 113 ++++++++++++++++++ .../allocators/test_allocator_ascend_npu.py | 105 ++++++++++++++++ .../allocators/test_fabric_allocator_utils.py | 70 +++++++++++ .../tests/packaging/test_installed_wheel.py | 51 ++++++++ python/tests/packaging/test_project.py | 34 ++++++ scripts/build_wheel.sh | 44 ++++--- 15 files changed, 508 insertions(+), 23 deletions(-) rename {mooncake-integration => python/mooncake}/allocator.py (99%) rename {mooncake-integration => python/mooncake}/allocator_ascend_npu.py (98%) rename {mooncake-integration => python/mooncake}/fabric_allocator_utils.py (97%) create mode 100644 python/tests/allocators/conftest.py create mode 100644 python/tests/allocators/test_allocator.py create mode 100644 python/tests/allocators/test_allocator_ascend_npu.py create mode 100644 python/tests/allocators/test_fabric_allocator_utils.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2c9f901b16..d0f679c494 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -12,9 +12,13 @@ /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/allocator.py @ShangmingCai @alogfans +/python/mooncake/allocator_ascend_npu.py @ShangmingCai @alogfans @ascend-direct-dev +/python/mooncake/fabric_allocator_utils.py @ShangmingCai @alogfans +/python/tests/allocators/ @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..0d75744f6a 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -29,7 +29,12 @@ P2P Store: Integration: - changed-files: - - any-glob-to-any-file: 'mooncake-integration/**/*' + - any-glob-to-any-file: + - 'mooncake-integration/**/*' + - 'python/mooncake/allocator.py' + - 'python/mooncake/allocator_ascend_npu.py' + - 'python/mooncake/fabric_allocator_utils.py' + - 'python/tests/allocators/**/*' Common: - changed-files: @@ -57,11 +62,14 @@ Tests: - 'scripts/test_*' - 'mooncake-wheel/tests/**/*' - 'mooncake-reshard/tests/**/*' + - 'python/tests/**/*' - 'scripts/tone_tests/**/*' Ascend/NPU: - changed-files: - any-glob-to-any-file: - '**/ascend*/**' + - 'python/mooncake/allocator_ascend_npu.py' + - 'python/tests/allocators/test_allocator_ascend_npu.py' - 'scripts/ascend/**' - '.github/workflows/*ascend*' diff --git a/.gitignore b/.gitignore index 35a253a732..e12d0583d3 100644 --- a/.gitignore +++ b/.gitignore @@ -195,6 +195,7 @@ libetcd_wrapper.h mooncake-wheel/mooncake/allocator.py mooncake-wheel/mooncake/allocator_ascend_npu.py +mooncake-wheel/mooncake/fabric_allocator_utils.py mooncake-wheel/mooncake/mooncake_master mooncake-wheel/mooncake/transfer_engine_bench diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 0ef9dfde04..fe1afff2fd 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -1,5 +1,6 @@ file(GLOB SOURCES "*.cpp") set(PYTHON_PACKAGE_NAME "mooncake") +set(MOONCAKE_PYTHON_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake") if(SKBUILD) # scikit-build-core sets CMAKE_INSTALL_PREFIX to its private wheel tree. Keep # every Python artifact relative to that prefix. @@ -152,7 +153,7 @@ endif() if(MOONCAKE_HAS_NVLINK_ALLOCATOR) message(STATUS "allocator.py will be installed in the Python package") install( - FILES "${CMAKE_CURRENT_SOURCE_DIR}/allocator.py" + FILES "${MOONCAKE_PYTHON_SOURCE_DIR}/allocator.py" DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR} COMPONENT python) endif() @@ -161,7 +162,7 @@ if(USE_UBSHMEM) message( STATUS "allocator_ascend_npu.py will be installed in the Python package") install( - FILES "${CMAKE_CURRENT_SOURCE_DIR}/allocator_ascend_npu.py" + FILES "${MOONCAKE_PYTHON_SOURCE_DIR}/allocator_ascend_npu.py" DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR} COMPONENT python) endif() @@ -170,7 +171,7 @@ if(MOONCAKE_HAS_NVLINK_ALLOCATOR OR USE_UBSHMEM) message( STATUS "fabric_allocator_utils.py will be installed in the Python package") install( - FILES "${CMAKE_CURRENT_SOURCE_DIR}/fabric_allocator_utils.py" + FILES "${MOONCAKE_PYTHON_SOURCE_DIR}/fabric_allocator_utils.py" DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR} COMPONENT python) endif() diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 6cb7b26ebf..b7cd1270a4 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -22,7 +22,10 @@ install( COMPONENT python FILES_MATCHING PATTERN "*.py" - PATTERN "__init__.py" EXCLUDE) + PATTERN "__init__.py" EXCLUDE + PATTERN "allocator.py" EXCLUDE + PATTERN "allocator_ascend_npu.py" EXCLUDE + PATTERN "fabric_allocator_utils.py" EXCLUDE) install( DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-reshard/python/mooncake/reshard/" diff --git a/mooncake-integration/allocator.py b/python/mooncake/allocator.py similarity index 99% rename from mooncake-integration/allocator.py rename to python/mooncake/allocator.py index f8804a1f31..20a7dc310b 100644 --- a/mooncake-integration/allocator.py +++ b/python/mooncake/allocator.py @@ -1,10 +1,12 @@ +from __future__ import annotations + import ctypes +from enum import IntEnum +from importlib import resources import logging import os import threading -from importlib import resources from typing import Dict, Final -from enum import IntEnum from torch import device as torch_device from torch.cuda.memory import CUDAPluggableAllocator diff --git a/mooncake-integration/allocator_ascend_npu.py b/python/mooncake/allocator_ascend_npu.py similarity index 98% rename from mooncake-integration/allocator_ascend_npu.py rename to python/mooncake/allocator_ascend_npu.py index bc4e041992..490111588c 100644 --- a/mooncake-integration/allocator_ascend_npu.py +++ b/python/mooncake/allocator_ascend_npu.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import ctypes import logging import threading diff --git a/mooncake-integration/fabric_allocator_utils.py b/python/mooncake/fabric_allocator_utils.py similarity index 97% rename from mooncake-integration/fabric_allocator_utils.py rename to python/mooncake/fabric_allocator_utils.py index 03292aa505..e52e95ffe8 100644 --- a/mooncake-integration/fabric_allocator_utils.py +++ b/python/mooncake/fabric_allocator_utils.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import ctypes import logging import os diff --git a/python/tests/allocators/conftest.py b/python/tests/allocators/conftest.py new file mode 100644 index 0000000000..547ef0245e --- /dev/null +++ b/python/tests/allocators/conftest.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +from types import ModuleType +from typing import Callable + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "mooncake" + + +class FakePluggableAllocator: + created: list[tuple[str, str, str]] = [] + + def __init__(self, so_path: str, malloc_symbol: str, free_symbol: str): + self.arguments = (so_path, malloc_symbol, free_symbol) + self.created.append(self.arguments) + + +@pytest.fixture +def import_allocator_module( + monkeypatch: pytest.MonkeyPatch, +) -> Callable[[str], ModuleType]: + package = ModuleType("mooncake") + package.__path__ = [str(PACKAGE_ROOT)] + package_spec = ModuleSpec("mooncake", loader=None, is_package=True) + package_spec.submodule_search_locations = package.__path__ + package.__spec__ = package_spec + + class FakeDevice: + pass + + torch = ModuleType("torch") + torch.__path__ = [] + torch.device = FakeDevice + torch_cuda = ModuleType("torch.cuda") + torch_cuda.__path__ = [] + torch_cuda_memory = ModuleType("torch.cuda.memory") + torch_cuda_memory.CUDAPluggableAllocator = FakePluggableAllocator + + torch_npu = ModuleType("torch_npu") + torch_npu.__path__ = [] + torch_npu_npu = ModuleType("torch_npu.npu") + torch_npu_npu.__path__ = [] + torch_npu_memory = ModuleType("torch_npu.npu.memory") + torch_npu_memory.NPUPluggableAllocator = FakePluggableAllocator + + modules = { + "mooncake": package, + "torch": torch, + "torch.cuda": torch_cuda, + "torch.cuda.memory": torch_cuda_memory, + "torch_npu": torch_npu, + "torch_npu.npu": torch_npu_npu, + "torch_npu.npu.memory": torch_npu_memory, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + FakePluggableAllocator.created.clear() + + def load(module_name: str) -> ModuleType: + for name in ( + "mooncake.allocator", + "mooncake.allocator_ascend_npu", + "mooncake.fabric_allocator_utils", + ): + monkeypatch.delitem(sys.modules, name, raising=False) + return importlib.import_module(module_name) + + return load diff --git a/python/tests/allocators/test_allocator.py b/python/tests/allocators/test_allocator.py new file mode 100644 index 0000000000..d8dd3a4978 --- /dev/null +++ b/python/tests/allocators/test_allocator.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from types import ModuleType +from typing import Callable + +import pytest + +from conftest import FakePluggableAllocator + + +@pytest.mark.parametrize( + ("probe_value", "expected_name"), + [ + (0, "USE_CUDAMALLOC"), + (1, "USE_CUMEMCREATE"), + (-2, "UNSUPPORTED"), + (37, "UNKNOWN"), + ], +) +def test_nvlink_probe_maps_backend_values( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, + probe_value: int, + expected_name: str, +) -> None: + allocator = import_allocator_module("mooncake.allocator") + + def probe_backend(*_args: object) -> int: + return probe_value + + monkeypatch.setattr(allocator, "probe_allocator_backend", probe_backend) + + backend = allocator.NVLinkAllocator._probe_fabric_memory_support("fake.so") + + assert backend.name == expected_name + + +def test_nvlink_detection_probes_once( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator = import_allocator_module("mooncake.allocator") + allocator_type = allocator.NVLinkAllocator + allocator_type._probe_done = False + allocator_type._supports_fabric = allocator.MemoryBackend.UNKNOWN + calls = 0 + + def get_path(_cls: type[object]) -> str: + return "nvlink_allocator.so" + + def probe(_cls: type[object], so_path: str) -> object: + nonlocal calls + calls += 1 + assert so_path == "nvlink_allocator.so" + return allocator.MemoryBackend.USE_CUMEMCREATE + + monkeypatch.setattr(allocator_type, "_get_so_path", classmethod(get_path)) + monkeypatch.setattr( + allocator_type, "_probe_fabric_memory_support", classmethod(probe) + ) + + assert ( + allocator_type.detect_mem_backend() is allocator.MemoryBackend.USE_CUMEMCREATE + ) + assert ( + allocator_type.detect_mem_backend() is allocator.MemoryBackend.USE_CUMEMCREATE + ) + assert calls == 1 + + +def test_nvlink_detection_caches_setup_failure( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator = import_allocator_module("mooncake.allocator") + allocator_type = allocator.NVLinkAllocator + allocator_type._probe_done = False + allocator_type._supports_fabric = allocator.MemoryBackend.UNKNOWN + calls = 0 + + def get_path(_cls: type[object]) -> str: + nonlocal calls + calls += 1 + raise ImportError("allocator unavailable") + + monkeypatch.setattr(allocator_type, "_get_so_path", classmethod(get_path)) + + assert allocator_type.detect_mem_backend() is allocator.MemoryBackend.UNSUPPORTED + assert allocator_type.detect_mem_backend() is allocator.MemoryBackend.UNSUPPORTED + assert calls == 1 + + +def test_nvlink_allocator_instance_is_cached_per_device( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator = import_allocator_module("mooncake.allocator") + allocator_type = allocator.NVLinkAllocator + allocator_type._instances = {} + monkeypatch.setattr( + allocator_type, + "_get_so_path", + classmethod(lambda _cls: "nvlink_allocator.so"), + ) + device = object() + + first = allocator_type.get_allocator(device) + second = allocator_type.get_allocator(device) + + assert first is second + assert FakePluggableAllocator.created == [ + ("nvlink_allocator.so", "mc_allocator_malloc", "mc_allocator_free") + ] diff --git a/python/tests/allocators/test_allocator_ascend_npu.py b/python/tests/allocators/test_allocator_ascend_npu.py new file mode 100644 index 0000000000..50d908c605 --- /dev/null +++ b/python/tests/allocators/test_allocator_ascend_npu.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from types import ModuleType +from typing import Callable + +import pytest + +from conftest import FakePluggableAllocator + + +@pytest.mark.parametrize(("probe_value", "expected"), [(0, False), (1, True)]) +def test_ubshmem_probe_maps_capability( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, + probe_value: int, + expected: bool, +) -> None: + allocator = import_allocator_module("mooncake.allocator_ascend_npu") + + def probe_backend(*_args: object) -> int: + return probe_value + + monkeypatch.setattr(allocator, "probe_allocator_backend", probe_backend) + + assert ( + allocator.UBShmemAllocator._probe_fabric_memory_support("fake.so") is expected + ) + + +def test_ubshmem_detection_probes_once( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator = import_allocator_module("mooncake.allocator_ascend_npu") + allocator_type = allocator.UBShmemAllocator + allocator_type._probe_done = False + allocator_type._supports_fabric = False + calls = 0 + + def get_path(_cls: type[object]) -> str: + return "ubshmem_fabric_allocator.so" + + def probe(_cls: type[object], so_path: str) -> bool: + nonlocal calls + calls += 1 + assert so_path == "ubshmem_fabric_allocator.so" + return True + + monkeypatch.setattr(allocator_type, "_get_so_path", classmethod(get_path)) + monkeypatch.setattr( + allocator_type, "_probe_fabric_memory_support", classmethod(probe) + ) + + assert allocator_type.detect_mem_backend() is True + assert allocator_type.detect_mem_backend() is True + assert calls == 1 + + +def test_ubshmem_detection_caches_setup_failure( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator = import_allocator_module("mooncake.allocator_ascend_npu") + allocator_type = allocator.UBShmemAllocator + allocator_type._probe_done = False + allocator_type._supports_fabric = True + calls = 0 + + def get_path(_cls: type[object]) -> str: + nonlocal calls + calls += 1 + raise ImportError("allocator unavailable") + + monkeypatch.setattr(allocator_type, "_get_so_path", classmethod(get_path)) + + assert allocator_type.detect_mem_backend() is False + assert allocator_type.detect_mem_backend() is False + assert calls == 1 + + +def test_ubshmem_allocator_instance_is_cached_per_device( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + allocator = import_allocator_module("mooncake.allocator_ascend_npu") + allocator_type = allocator.UBShmemAllocator + allocator_type._instances = {} + monkeypatch.setattr( + allocator_type, + "_get_so_path", + classmethod(lambda _cls: "ubshmem_fabric_allocator.so"), + ) + device = object() + + first = allocator_type.get_allocator(device) + second = allocator_type.get_allocator(device) + + assert first is second + assert FakePluggableAllocator.created == [ + ( + "ubshmem_fabric_allocator.so", + "mc_allocator_malloc", + "mc_allocator_free", + ) + ] diff --git a/python/tests/allocators/test_fabric_allocator_utils.py b/python/tests/allocators/test_fabric_allocator_utils.py new file mode 100644 index 0000000000..46d0aa72f3 --- /dev/null +++ b/python/tests/allocators/test_fabric_allocator_utils.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from types import ModuleType +from typing import Callable + +import pytest + + +class FakeProbe: + def __init__(self, result: int): + self.result = result + self.argtypes: list[object] | None = None + self.restype: object | None = None + + def __call__(self, device_id: int) -> int: + assert device_id == 0 + return self.result + + +def test_probe_allocator_backend_calls_the_exported_symbol( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + utilities = import_allocator_module("mooncake.fabric_allocator_utils") + probe = FakeProbe(7) + library = ModuleType("allocator_library") + library.mc_allocator_probe = probe + monkeypatch.setattr(utilities.ctypes, "CDLL", lambda _path: library) + + result = utilities.probe_allocator_backend( + "allocator.so", "mc_allocator_probe", utilities.ctypes.c_int, -2 + ) + + assert result == 7 + assert probe.argtypes == [utilities.ctypes.c_int] + assert probe.restype is utilities.ctypes.c_int + + +def test_probe_allocator_backend_uses_fallback_for_missing_symbol( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + utilities = import_allocator_module("mooncake.fabric_allocator_utils") + monkeypatch.setattr(utilities.ctypes, "CDLL", lambda _path: object()) + + assert ( + utilities.probe_allocator_backend( + "allocator.so", "mc_allocator_probe", utilities.ctypes.c_int, -2 + ) + == -2 + ) + + +def test_probe_allocator_backend_uses_fallback_for_loader_failure( + import_allocator_module: Callable[[str], ModuleType], + monkeypatch: pytest.MonkeyPatch, +) -> None: + utilities = import_allocator_module("mooncake.fabric_allocator_utils") + + def fail_to_load(_path: str) -> object: + raise OSError("not loadable") + + monkeypatch.setattr(utilities.ctypes, "CDLL", fail_to_load) + + assert ( + utilities.probe_allocator_backend( + "allocator.so", "mc_allocator_probe", utilities.ctypes.c_int, False + ) + is False + ) diff --git a/python/tests/packaging/test_installed_wheel.py b/python/tests/packaging/test_installed_wheel.py index 74422d9ebb..280901bc95 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/python/tests/packaging/test_installed_wheel.py @@ -4,11 +4,17 @@ from pathlib import Path import subprocess import sys +from zipfile import ZipFile import pytest REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +ALLOCATOR_MODULES = ( + "allocator.py", + "allocator_ascend_npu.py", + "fabric_allocator_utils.py", +) def _project_version() -> str: @@ -28,6 +34,10 @@ 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(wheel) as archive: + members = set(archive.namelist()) + for module in ALLOCATOR_MODULES: + assert f"mooncake/{module}" in members environment = tmp_path / "environment" subprocess.run([sys.executable, "-m", "venv", str(environment)], check=True) @@ -42,9 +52,13 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: clean_environment["PYTHONNOUSERSITE"] = "1" smoke_script = f""" from importlib import metadata +from importlib.util import find_spec from pathlib import Path +import sys +from types import ModuleType import mooncake import mooncake.engine +import mooncake.fabric_allocator_utils import mooncake.reshard import mooncake.store @@ -54,6 +68,43 @@ 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 +for module in {ALLOCATOR_MODULES!r}: + spec = find_spec(f"mooncake.{{module.removesuffix('.py')}}") + assert spec is not None and spec.origin is not None + assert Path(spec.origin).resolve().parent == package_path.parent +assert "torch" not in sys.modules +assert "torch_npu" not in sys.modules + +class FakePluggableAllocator: + pass + +torch = ModuleType("torch") +torch.__path__ = [] +torch.device = object +torch_cuda = ModuleType("torch.cuda") +torch_cuda.__path__ = [] +torch_cuda_memory = ModuleType("torch.cuda.memory") +torch_cuda_memory.CUDAPluggableAllocator = FakePluggableAllocator +torch_npu = ModuleType("torch_npu") +torch_npu.__path__ = [] +torch_npu_npu = ModuleType("torch_npu.npu") +torch_npu_npu.__path__ = [] +torch_npu_memory = ModuleType("torch_npu.npu.memory") +torch_npu_memory.NPUPluggableAllocator = FakePluggableAllocator +sys.modules.update({{ + "torch": torch, + "torch.cuda": torch_cuda, + "torch.cuda.memory": torch_cuda_memory, + "torch_npu": torch_npu, + "torch_npu.npu": torch_npu_npu, + "torch_npu.npu.memory": torch_npu_memory, +}}) + +import mooncake.allocator +import mooncake.allocator_ascend_npu + +assert mooncake.allocator.NVLinkAllocator is not None +assert mooncake.allocator_ascend_npu.UBShmemAllocator is not None """ 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..c9327adc45 100644 --- a/python/tests/packaging/test_project.py +++ b/python/tests/packaging/test_project.py @@ -14,6 +14,11 @@ REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +ALLOCATOR_MODULES = ( + "allocator.py", + "allocator_ascend_npu.py", + "fabric_allocator_utils.py", +) def test_scikit_build_core_is_the_only_build_backend() -> None: @@ -57,6 +62,35 @@ def test_tracked_source_roots_contain_no_generated_native_artifacts() -> None: assert not list((REPOSITORY_ROOT / "mooncake-pg" / "torch").rglob("*.so")) +def test_allocators_have_one_authoritative_source() -> None: + package_root = REPOSITORY_ROOT / "python" / "mooncake" + legacy_root = REPOSITORY_ROOT / "mooncake-integration" + + for module in ALLOCATOR_MODULES: + assert (package_root / module).is_file() + assert not (legacy_root / module).exists() + + tests_root = REPOSITORY_ROOT / "python" / "tests" / "allocators" + assert (tests_root / "test_allocator.py").is_file() + assert (tests_root / "test_allocator_ascend_npu.py").is_file() + assert (tests_root / "test_fabric_allocator_utils.py").is_file() + + +def test_allocator_build_inputs_use_the_canonical_source() -> None: + integration_cmake = ( + REPOSITORY_ROOT / "mooncake-integration" / "CMakeLists.txt" + ).read_text() + python_cmake = (REPOSITORY_ROOT / "python" / "CMakeLists.txt").read_text() + legacy_builder = (REPOSITORY_ROOT / "scripts" / "build_wheel.sh").read_text() + + assert "../python/mooncake" in integration_cmake + assert 'MIGRATED_ALLOCATOR_SOURCE_DIR="python/mooncake"' in legacy_builder + for module in ALLOCATOR_MODULES: + assert f'"${{MOONCAKE_PYTHON_SOURCE_DIR}}/{module}"' in integration_cmake + assert f'PATTERN "{module}" EXCLUDE' in python_cmake + assert f"mooncake-integration/{module}" not in legacy_builder + + 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..65658bc3fd 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -15,6 +15,21 @@ OUTPUT_DIR=${OUTPUT_DIR:-${2:-"dist"}} # their normal CMake output paths before auditwheel, like the PG core library. BUILD_DIR="${BUILD_DIR:-build}" BUILD_DIR_ABS="$(pwd)/${BUILD_DIR}" +MIGRATED_ALLOCATOR_SOURCE_DIR="python/mooncake" +MIGRATED_ALLOCATOR_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake" +MIGRATED_ALLOCATOR_MODULES=(fabric_allocator_utils.py) +MIGRATED_ALLOCATOR_ALL_MODULES=( + allocator.py + allocator_ascend_npu.py + fabric_allocator_utils.py +) + +cleanup_allocator_staging() { + for module in "${MIGRATED_ALLOCATOR_ALL_MODULES[@]}"; do + rm -f "${MIGRATED_ALLOCATOR_STAGING_DIR}/${module}" + done +} +trap cleanup_allocator_staging EXIT echo "Building wheel for Python ${PYTHON_VERSION} with output directory ${OUTPUT_DIR}" # Ensure LD_LIBRARY_PATH includes /usr/local/lib @@ -24,11 +39,7 @@ echo "Cleaning wheel-build directory" rm -rf mooncake-wheel/mooncake_transfer_engine* rm -rf mooncake-wheel/build/ rm -f mooncake-wheel/mooncake/*.so - -echo "Creating directory structure..." - -# Copy shared allocator helper used by both CUDA and Ascend pluggable allocators. -cp mooncake-integration/fabric_allocator_utils.py mooncake-wheel/mooncake/fabric_allocator_utils.py +cleanup_allocator_staging # Copy engine.so to mooncake directory (will be imported by transfer module) cp ${BUILD_DIR}/mooncake-integration/engine.*.so mooncake-wheel/mooncake/engine.so @@ -114,8 +125,7 @@ if [ -f ${BUILD_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator. cp ${BUILD_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so mooncake-wheel/mooncake/nvlink_allocator.so fi echo "Copying allocator libraries..." - # Copy allocator.py - cp mooncake-integration/allocator.py mooncake-wheel/mooncake/allocator.py + MIGRATED_ALLOCATOR_MODULES+=(allocator.py) else echo "Skipping nvlink_allocator.so (not built - likely ARM64 or non-CUDA build)" fi @@ -125,8 +135,7 @@ if [ -f ${BUILD_DIR}/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_a echo "Copying NPU ubshmem_fabric_allocator.so..." cp ${BUILD_DIR}/mooncake-transfer-engine/ubshmem-allocator/ubshmem_fabric_allocator.so mooncake-wheel/mooncake/ubshmem_fabric_allocator.so echo "Copying NPU allocator libraries..." - # Copy allocator_ascend_npu.py - cp mooncake-integration/allocator_ascend_npu.py mooncake-wheel/mooncake/allocator_ascend_npu.py + MIGRATED_ALLOCATOR_MODULES+=(allocator_ascend_npu.py) else echo "Skipping ubshmem_fabric_allocator.so (not built - likely CUDA or non-NPU build)" fi @@ -180,15 +189,20 @@ if [ "$NPU_BUILD" = "1" ]; then fi echo "Building wheel package..." -# Stage the reshard Python package for the combined Mooncake wheel. The tracked -# source of truth remains in the top-level module. +# Stage migrated allocator modules and the Reshard package for the combined +# Mooncake wheel. Each tracked source remains in its authoritative tree. RESHARD_SOURCE_DIR="mooncake-reshard/python/mooncake/reshard" RESHARD_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake/reshard" -cleanup_reshard_staging() { +cleanup_python_staging() { + cleanup_allocator_staging rm -rf "${RESHARD_STAGING_DIR}" } -trap cleanup_reshard_staging EXIT -rm -rf "${RESHARD_STAGING_DIR}" +trap cleanup_python_staging EXIT +cleanup_python_staging +for module in "${MIGRATED_ALLOCATOR_MODULES[@]}"; do + cp "${MIGRATED_ALLOCATOR_SOURCE_DIR}/${module}" \ + "${MIGRATED_ALLOCATOR_STAGING_DIR}/${module}" +done cp -R "${RESHARD_SOURCE_DIR}" "${RESHARD_STAGING_DIR}" # Build the wheel package @@ -204,7 +218,7 @@ WHEEL_DIR="$(pwd)" cleanup_wheel_metadata_state() { [[ -f "${WHEEL_DIR}/pyproject.toml.backup" ]] && mv "${WHEEL_DIR}/pyproject.toml.backup" "${WHEEL_DIR}/pyproject.toml" rm -f "${WHEEL_DIR}/README.md" - cleanup_reshard_staging + cleanup_python_staging } trap cleanup_wheel_metadata_state EXIT