diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 0fecc2ff09..b2944ac72f 100644 --- a/.github/CODEOWNERS +++ b/.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 diff --git a/.github/labeler.yml b/.github/labeler.yml index 7815e2dd20..016e64fc5d 100644 --- a/.github/labeler.yml +++ b/.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' 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 e7d3c3699e..acc6de27fb 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/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" 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 08d3d66b92..139252d76b 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/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 ( diff --git a/python/tests/packaging/test_project.py b/python/tests/packaging/test_project.py index ce2a018645..03198eff57 100644 --- a/python/tests/packaging/test_project.py +++ b/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", 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..956f152bc5 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 / "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() diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index 600b3de956..50380e8b9f 100755 --- a/scripts/build_wheel.sh +++ b/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 diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh index 8de054332a..8a08f7f47a 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; } 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