[Store] Migrate Store REST service to unified Python package

.github/CODEOWNERS

@@ -15,6 +15,8 @@
 /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
 /python/mooncake/_fast_copy.c @ykwd @stmatengss @zxpdemonio
 /python/mooncake/structured_object_store.py @ykwd @stmatengss @zxpdemonio
 /python/mooncake/http_metadata_server.py @ykwd @stmatengss @zxpdemonio

.github/labeler.yml

@@ -7,6 +7,7 @@ store:
   - changed-files:
     - any-glob-to-any-file:
       - 'mooncake-store/**/*'
+      - 'python/mooncake/mooncake_store_service.py'
       - 'python/mooncake/_fast_copy.c'
       - 'python/mooncake/structured_object_store.py'
       - 'python/mooncake/http_metadata_server.py'

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 格式的配置文件。例如：
 

mooncake-integration/CMakeLists.txt

@@ -246,7 +246,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}/../python/mooncake/mooncake_config.py"
         "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/cli.py"
         "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/cli_client.py"

mooncake-wheel/mooncake/mooncake_store_service.py → 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

python/tests/packaging/test_installed_wheel.py

@@ -59,6 +59,7 @@ def test_wheel_imports_outside_the_repository(
 from importlib import import_module, metadata, util
 from pathlib import Path
 import sys
+
 import mooncake
 
 cli_modules = (
@@ -78,6 +79,7 @@ def test_wheel_imports_outside_the_repository(
 
 import mooncake.async_store
 import mooncake.engine
+import mooncake.mooncake_store_service
 import mooncake.http_metadata_server
 import mooncake.mooncake_config
 import mooncake.reshard
@@ -95,6 +97,14 @@ def test_wheel_imports_outside_the_repository(
 )
 assert Path(mooncake.async_store.__file__).resolve().parent == package_path.parent
 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
 assert mooncake.http_metadata_server.KVBootstrapServer is not None
 assert mooncake.mooncake_config.MooncakeConfig is not None
 for ep_module in (

python/tests/packaging/test_project.py

@@ -85,6 +85,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_structured_object_storage_has_one_authoritative_source() -> None:
     package_root = REPOSITORY_ROOT / "python" / "mooncake"
     test_root = REPOSITORY_ROOT / "python" / "tests" / "store"
@@ -269,7 +295,6 @@ def test_scikit_build_consumes_unified_python_sources() -> None:
     # Modules still supplied from the legacy tree at build time. Phase 2 should
     # migrate these into python/mooncake and drop the DIRECTORY rule above.
     for module in (
-        "mooncake_store_service.py",
         "mooncake_connector_v1.py",
         "vllm_v1_proxy_server.py",
         "pg.py",

mooncake-wheel/tests/test_mooncake_store_service_api.py → 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 / "python" / "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()
 

scripts/build_wheel.sh

@@ -190,6 +190,7 @@ echo "Building wheel package..."
 MIGRATED_PYTHON_SOURCE_DIR="python/mooncake"
 MIGRATED_PYTHON_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake"
 MIGRATED_PYTHON_MODULES=(
+    mooncake_store_service.py
     structured_object_store.py
     _fast_copy.c
     http_metadata_server.py

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; }
 python -c "import mooncake.cli, mooncake.cli_bench, mooncake.cli_client, mooncake.transfer_engine_topology_dump" && echo "Success: CLI imports succeeded after installation" || { echo "ERROR: CLI imports failed after installation!"; exit 1; }
 
@@ -75,6 +76,10 @@ 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"
 echo "Verifying mooncake_http_metadata_server entry point..."
 which mooncake_http_metadata_server || { echo "ERROR: mooncake_http_metadata_server entry point not found!"; exit 1; }
 mooncake_http_metadata_server --help >/dev/null