[Integration] Migrate allocators to unified Python package

.github/CODEOWNERS

@@ -15,6 +15,10 @@
 /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
 /python/mooncake/http_metadata_server.py @ykwd @stmatengss @zxpdemonio
 /python/mooncake/async_store.py @ykwd @stmatengss @zxpdemonio
 /python/tests/store/async_store_integration.py @ykwd @stmatengss @zxpdemonio

.github/labeler.yml

@@ -60,7 +60,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:
@@ -97,5 +102,7 @@ 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*'

.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
 

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.
@@ -174,7 +175,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()
@@ -183,7 +184,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()
@@ -192,7 +193,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()

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/"

mooncake-integration/allocator.py → 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

mooncake-integration/allocator_ascend_npu.py → python/mooncake/allocator_ascend_npu.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import ctypes
 import logging
 import threading

mooncake-integration/fabric_allocator_utils.py → python/mooncake/fabric_allocator_utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import ctypes
 import logging
 import os

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

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")
+    ]