From 08003cf17ffb26e9f10e1f0b31c62eec80f1cd8a Mon Sep 17 00:00:00 2001 From: Aionw Date: Mon, 24 Aug 2026 16:44:34 +0800 Subject: [PATCH] [Store] Migrate Store REST service to unified Python package --- .github/CODEOWNERS | 4 +- .github/labeler.yml | 6 +- docs/source/zh_archive/mooncake-store.md | 2 +- mooncake-integration/CMakeLists.txt | 2 +- .../mooncake/mooncake_store_service.py | 16 ++--- .../tests/packaging/test_installed_wheel.py | 16 +++++ python/tests/packaging/test_project.py | 26 +++++++++ .../store}/test_mooncake_store_service_api.py | 58 ++++++++++++++----- scripts/build_wheel.sh | 22 +++++-- scripts/test_installation.sh | 6 ++ 10 files changed, 121 insertions(+), 37 deletions(-) rename {mooncake-wheel => python}/mooncake/mooncake_store_service.py (98%) rename {mooncake-wheel/tests => python/tests/store}/test_mooncake_store_service_api.py (96%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2c9f901b16..922dfad50d 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/mooncake_store_service.py @ykwd @stmatengss +/python/tests/store/test_mooncake_store_service_api.py @ykwd @stmatengss /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..384060e79d 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -5,7 +5,10 @@ run-ci: store: - changed-files: - - any-glob-to-any-file: 'mooncake-store/**/*' + - any-glob-to-any-file: + - 'mooncake-store/**/*' + - 'python/mooncake/mooncake_store_service.py' + - 'python/tests/store/test_mooncake_store_service_api.py' Transfer Engine: - 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/docs/source/zh_archive/mooncake-store.md b/docs/source/zh_archive/mooncake-store.md index b5cffe90d4..8bb4854f86 100644 --- a/docs/source/zh_archive/mooncake-store.md +++ b/docs/source/zh_archive/mooncake-store.md @@ -910,7 +910,7 @@ retcode = store.setup( ### 将 Client 作为独立进程启动并通过 HTTP 访问 -使用 `mooncake-wheel/mooncake/mooncake_store_service.py` 可以以独立进程的形式启动 **真实**`Client` 并通过 HTTP 访问。 +使用 `python/mooncake/mooncake_store_service.py` 可以以独立进程的形式启动 **真实**`Client` 并通过 HTTP 访问。 首先，创建并保存一个 JSON 格式的配置文件。例如： diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 0ef9dfde04..06bd30b4d6 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -207,7 +207,7 @@ if(NOT SKBUILD) if(WITH_STORE) install( FILES - "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_store_service.py" + "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/mooncake_store_service.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_config.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli.py" DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR}) diff --git a/mooncake-wheel/mooncake/mooncake_store_service.py b/python/mooncake/mooncake_store_service.py similarity index 98% rename from mooncake-wheel/mooncake/mooncake_store_service.py rename to python/mooncake/mooncake_store_service.py index b5bc66b2de..2314f7bb37 100644 --- a/mooncake-wheel/mooncake/mooncake_store_service.py +++ b/python/mooncake/mooncake_store_service.py @@ -41,9 +41,7 @@ def _shm_name_to_path(name): def _unblock_shutdown_signals(): try: - signal.pthread_sigmask( - signal.SIG_UNBLOCK, {signal.SIGINT, signal.SIGTERM} - ) + signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGINT, signal.SIGTERM}) except AttributeError: pass @@ -126,9 +124,7 @@ def _setup_logging(self): format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) - async def start_store_service( - self, max_wait_time: float = 60, shutdown_event=None - ): + async def start_store_service(self, max_wait_time: float = 60, shutdown_event=None): """ Start the store service with retry mechanism. @@ -191,9 +187,7 @@ async def start_store_service( # chance to publish a shutdown requested while it ran. await asyncio.sleep(0) if shutdown_event.is_set(): - logging.info( - "Store startup cancelled by shutdown request" - ) + logging.info("Store startup cancelled by shutdown request") await self.stop() return False @@ -228,9 +222,7 @@ async def start_store_service( shutdown_event.wait(), timeout=actual_sleep_time, ) - logging.info( - "Store startup cancelled by shutdown request" - ) + logging.info("Store startup cancelled by shutdown request") return False except asyncio.TimeoutError: pass diff --git a/python/tests/packaging/test_installed_wheel.py b/python/tests/packaging/test_installed_wheel.py index 74422d9ebb..771ee07de7 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/python/tests/packaging/test_installed_wheel.py @@ -43,8 +43,16 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: smoke_script = f""" from importlib import metadata from pathlib import Path +import sys +import types + +aiohttp = types.ModuleType("aiohttp") +aiohttp.web = types.SimpleNamespace() +sys.modules["aiohttp"] = aiohttp + import mooncake import mooncake.engine +import mooncake.mooncake_store_service import mooncake.reshard import mooncake.store @@ -54,6 +62,14 @@ 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.mooncake_store_service.MooncakeStoreService is not None +store_entry_points = metadata.entry_points( + group="console_scripts", name="mc_store_rest_server" +) +assert len(store_entry_points) == 1 +store_entry_point = store_entry_points[0] +assert store_entry_point.value == "mooncake.mooncake_store_service:sync_main" +assert store_entry_point.load() is mooncake.mooncake_store_service.sync_main """ 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..5571833520 100644 --- a/python/tests/packaging/test_project.py +++ b/python/tests/packaging/test_project.py @@ -57,6 +57,32 @@ def test_tracked_source_roots_contain_no_generated_native_artifacts() -> None: assert not list((REPOSITORY_ROOT / "mooncake-pg" / "torch").rglob("*.so")) +def test_store_rest_service_has_one_authoritative_source() -> None: + package_root = REPOSITORY_ROOT / "python" / "mooncake" + tests_root = REPOSITORY_ROOT / "python" / "tests" + legacy_package_root = REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" + legacy_tests_root = REPOSITORY_ROOT / "mooncake-wheel" / "tests" + + assert (package_root / "mooncake_store_service.py").is_file() + assert (tests_root / "store" / "test_mooncake_store_service_api.py").is_file() + assert not (legacy_package_root / "mooncake_store_service.py").exists() + assert not (legacy_tests_root / "test_mooncake_store_service_api.py").exists() + + project = tomllib.loads((REPOSITORY_ROOT / "pyproject.toml").read_text()) + assert ( + project["project"]["scripts"]["mc_store_rest_server"] + == "mooncake.mooncake_store_service:sync_main" + ) + + integration_cmake = ( + REPOSITORY_ROOT / "mooncake-integration" / "CMakeLists.txt" + ).read_text() + assert "../python/mooncake/mooncake_store_service.py" in integration_cmake + assert ( + "../mooncake-wheel/mooncake/mooncake_store_service.py" not in integration_cmake + ) + + def test_pg_extension_build_stages_outside_the_source_tree( tmp_path: Path, ) -> None: diff --git a/mooncake-wheel/tests/test_mooncake_store_service_api.py b/python/tests/store/test_mooncake_store_service_api.py similarity index 96% rename from mooncake-wheel/tests/test_mooncake_store_service_api.py rename to python/tests/store/test_mooncake_store_service_api.py index f6b49cad9f..980656f984 100644 --- a/mooncake-wheel/tests/test_mooncake_store_service_api.py +++ b/python/tests/store/test_mooncake_store_service_api.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import asyncio +import importlib.util import json import logging import signal @@ -13,7 +14,13 @@ from unittest import mock from unittest.mock import patch -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +for module_name in tuple(sys.modules): + if module_name == "mooncake" or module_name.startswith("mooncake."): + del sys.modules[module_name] +mooncake_package = types.ModuleType("mooncake") +mooncake_package.__path__ = [str(REPOSITORY_ROOT / "mooncake-wheel" / "mooncake")] +sys.modules["mooncake"] = mooncake_package try: from aiohttp import web as _unused_web # noqa: F401 @@ -44,12 +51,29 @@ class MooncakeDistributedStore: store_module.MooncakeDistributedStore = MooncakeDistributedStore sys.modules["mooncake.store"] = store_module -from mooncake.mooncake_store_service import ( - MooncakeStoreService, - _install_shutdown_signal_handlers, - _shm_name_to_path, - main as store_service_main, +STORE_SERVICE_PATH = ( + REPOSITORY_ROOT / "python" / "mooncake" / "mooncake_store_service.py" ) +STORE_SERVICE_SPEC = importlib.util.spec_from_file_location( + "mooncake.mooncake_store_service", STORE_SERVICE_PATH +) +assert STORE_SERVICE_SPEC is not None +assert STORE_SERVICE_SPEC.loader is not None +store_service_module = importlib.util.module_from_spec(STORE_SERVICE_SPEC) +sys.modules[STORE_SERVICE_SPEC.name] = store_service_module +STORE_SERVICE_SPEC.loader.exec_module(store_service_module) +mooncake_package.mooncake_store_service = store_service_module + +MooncakeStoreService = store_service_module.MooncakeStoreService +_install_shutdown_signal_handlers = ( + store_service_module._install_shutdown_signal_handlers +) +_shm_name_to_path = store_service_module._shm_name_to_path +store_service_main = store_service_module.main + + +def test_store_service_loads_from_canonical_source() -> None: + assert Path(store_service_module.__file__).resolve() == STORE_SERVICE_PATH class FakeStore: @@ -319,7 +343,9 @@ async def test_reconfigure_decode_remount_failure_keeps_previous_segments(self): # subsequent same-path detection keeps working. self.assertEqual(self.service.last_mount_info["path"], "/dev/shm/old") - async def test_reconfigure_decode_same_path_remount_failure_keeps_previous_segments(self): + async def test_reconfigure_decode_same_path_remount_failure_keeps_previous_segments( + self, + ): # A remount to the SAME path that fails to mount must not destroy the # still-healthy previous segments: the node keeps serving from them and # stays in decode mode (make-before-break). Same-path MBB is safe here @@ -397,7 +423,9 @@ def mount_new(path, size, offset, protocol, location): self.assertEqual(self.service.current_mode, "decode") self.assertEqual(self.service.last_mount_info["path"], "/dev/shm/new") - async def test_reconfigure_decode_partial_unmount_failure_keeps_only_failed_ids(self): + async def test_reconfigure_decode_partial_unmount_failure_keeps_only_failed_ids( + self, + ): # New mount succeeds, but retiring the previous segments only PARTIALLY # fails. unmount_segment reports the first error for a batch, so the old # ids must be unmounted individually; mounted_segment_ids must then hold @@ -596,7 +624,7 @@ async def test_handle_put_missing_key(self): async def test_handle_put_missing_value(self): self.fake_store.put = lambda key, value: 0 resp = await self.service.handle_put(FakeRequest({"key": "k"})) - self.assertEqual(resp.status, 500) + self.assertEqual(resp.status, 400) async def test_handle_put_store_failure(self): self.fake_store.put = lambda key, value: -1 @@ -847,7 +875,7 @@ async def test_reconfigure_prefill_unmounts_segments(self): self.service.mounted_segment_ids = [sid] resp = await self.service.handle_reconfigure(FakeRequest({"mode": "prefill"})) self.assertEqual(resp.status, 200) - self.assertEqual(self.fake_store.unmount_calls, [[sid]]) + self.assertEqual(self.fake_store.unmount_calls, [([sid], 0)]) self.assertEqual(self.service.mounted_segment_ids, []) self.assertEqual(self.service.current_mode, "prefill") @@ -888,7 +916,7 @@ async def test_reconfigure_decode_remount_unmounts_previous(self): ) ) self.assertEqual(resp.status, 200) - self.assertEqual(self.fake_store.unmount_calls, [[old_id]]) + self.assertEqual(self.fake_store.unmount_calls, [([old_id], 0)]) self.assertEqual(self.service.current_mode, "decode") # ==================== /api/mount edge cases ==================== @@ -985,7 +1013,7 @@ async def test_unmount_shm_string_segment_id_coercion(self): self.assertEqual(resp.status, 200) self.assertEqual( self.fake_store.unmount_calls, - [["00000000-0000-0000-0000-000000000001"]], + [(["00000000-0000-0000-0000-000000000001"], 0)], ) async def test_unmount_shm_empty_list(self): @@ -1008,6 +1036,8 @@ async def asyncSetUp(self): enable_ssd_offload=False, ssd_offload_path="", tenant_id="", + enable_client_http_server=False, + client_http_port=9300, ) async def test_shutdown_event_stops_startup_retry_sleep(self): @@ -1077,9 +1107,7 @@ async def test_stop_logs_nonzero_close_return(self): await self.service.stop() self.assertIsNone(self.service.store) - self.assertTrue( - any("close returned 7" in message for message in logs.output) - ) + self.assertTrue(any("close returned 7" in message for message in logs.output)) await self.service.stop() store.close.assert_called_once_with() diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index eed257699d..0e127bf291 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -180,15 +180,25 @@ 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 root Python modules and the legacy Reshard package for the +# combined-wheel builder. Each tracked source remains in its authoritative tree. +MIGRATED_PYTHON_SOURCE_DIR="python/mooncake" +MIGRATED_PYTHON_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake" +MIGRATED_PYTHON_MODULES=(mooncake_store_service.py) RESHARD_SOURCE_DIR="mooncake-reshard/python/mooncake/reshard" RESHARD_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake/reshard" -cleanup_reshard_staging() { +cleanup_migrated_python_staging() { + for module in "${MIGRATED_PYTHON_MODULES[@]}"; do + rm -f "${MIGRATED_PYTHON_STAGING_DIR}/${module}" + done rm -rf "${RESHARD_STAGING_DIR}" } -trap cleanup_reshard_staging EXIT -rm -rf "${RESHARD_STAGING_DIR}" +trap cleanup_migrated_python_staging EXIT +cleanup_migrated_python_staging +for module in "${MIGRATED_PYTHON_MODULES[@]}"; do + cp "${MIGRATED_PYTHON_SOURCE_DIR}/${module}" \ + "${MIGRATED_PYTHON_STAGING_DIR}/${module}" +done cp -R "${RESHARD_SOURCE_DIR}" "${RESHARD_STAGING_DIR}" # Build the wheel package @@ -204,7 +214,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_migrated_python_staging } trap cleanup_wheel_metadata_state EXIT diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh index dd8a7bb6b1..df38b1fb90 100755 --- a/scripts/test_installation.sh +++ b/scripts/test_installation.sh @@ -32,6 +32,7 @@ sudo apt-get install -y $SYSTEM_PACKAGES echo "Verifying that import succeeds after installation..." python -c "import mooncake.engine" && echo "Success: Import succeeded after installation" || { echo "ERROR: Import failed after installation!"; exit 1; } +python -c "import mooncake.mooncake_store_service" && echo "Success: Store REST service import succeeded after installation" || { echo "ERROR: Store REST service import failed after installation!"; exit 1; } python -c "import mooncake.reshard.weight" && echo "Success: Reshard import succeeded after installation" || { echo "ERROR: Reshard import failed after installation!"; exit 1; } echo "Running import structure test..." @@ -66,6 +67,11 @@ echo "Verifying transfer_engine_bench entry point..." which transfer_engine_bench || { echo "ERROR: transfer_engine_bench entry point not found!"; exit 1; } echo "Success: transfer_engine_bench entry point found" +echo "Verifying mc_store_rest_server entry point..." +which mc_store_rest_server || { echo "ERROR: mc_store_rest_server entry point not found!"; exit 1; } +mc_store_rest_server --help >/dev/null +echo "Success: mc_store_rest_server entry point found" + cd .. echo "Installation test completed successfully!"