From f0a0016470802b40e5aad4b9504b6c6512f06c90 Mon Sep 17 00:00:00 2001 From: Aionw Date: Sat, 22 Aug 2026 00:36:00 +0800 Subject: [PATCH 1/2] [Store] Migrate lightweight Python utilities --- .github/CODEOWNERS | 4 ++++ .github/labeler.yml | 8 ++++++- mooncake-integration/CMakeLists.txt | 2 +- .../mooncake/buffer_pool.py | 0 .../mooncake/mooncake_config.py | 0 .../tests/packaging/test_installed_wheel.py | 2 ++ python/tests/packaging/test_project.py | 22 +++++++++++++++++++ .../tests/store}/test_buffer_pool_native.py | 4 +--- .../tests/unit}/test_mooncake_config.py | 0 scripts/build_wheel.sh | 22 ++++++++++++++----- scripts/test_installation.sh | 3 ++- 11 files changed, 55 insertions(+), 12 deletions(-) rename {mooncake-wheel => python}/mooncake/buffer_pool.py (100%) rename {mooncake-wheel => python}/mooncake/mooncake_config.py (100%) rename {mooncake-wheel/tests => python/tests/store}/test_buffer_pool_native.py (98%) rename {mooncake-wheel/tests => python/tests/unit}/test_mooncake_config.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2c9f901b16..1a2a95a3ee 100644 --- a/.github/CODEOWNERS +++ b/.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/buffer_pool.py @ykwd @stmatengss +/python/mooncake/mooncake_config.py @ykwd @stmatengss +/python/tests/store/ @ykwd @stmatengss +/python/tests/unit/test_mooncake_config.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..3a1df330df 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -5,7 +5,12 @@ run-ci: store: - changed-files: - - any-glob-to-any-file: 'mooncake-store/**/*' + - any-glob-to-any-file: + - 'mooncake-store/**/*' + - 'python/mooncake/buffer_pool.py' + - 'python/mooncake/mooncake_config.py' + - 'python/tests/store/**/*' + - 'python/tests/unit/test_mooncake_config.py' Transfer Engine: - changed-files: @@ -57,6 +62,7 @@ Tests: - 'scripts/test_*' - 'mooncake-wheel/tests/**/*' - 'mooncake-reshard/tests/**/*' + - 'python/tests/**/*' - 'scripts/tone_tests/**/*' Ascend/NPU: diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 0ef9dfde04..f264b356b9 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -208,7 +208,7 @@ if(NOT SKBUILD) install( FILES "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_store_service.py" - "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_config.py" + "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/mooncake_config.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli.py" DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR}) diff --git a/mooncake-wheel/mooncake/buffer_pool.py b/python/mooncake/buffer_pool.py similarity index 100% rename from mooncake-wheel/mooncake/buffer_pool.py rename to python/mooncake/buffer_pool.py diff --git a/mooncake-wheel/mooncake/mooncake_config.py b/python/mooncake/mooncake_config.py similarity index 100% rename from mooncake-wheel/mooncake/mooncake_config.py rename to python/mooncake/mooncake_config.py diff --git a/python/tests/packaging/test_installed_wheel.py b/python/tests/packaging/test_installed_wheel.py index 74422d9ebb..b0921c832f 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/python/tests/packaging/test_installed_wheel.py @@ -45,6 +45,7 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: from pathlib import Path import mooncake import mooncake.engine +import mooncake.mooncake_config import mooncake.reshard import mooncake.store @@ -54,6 +55,7 @@ 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_config.MooncakeConfig 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..adb6fc3f96 100644 --- a/python/tests/packaging/test_project.py +++ b/python/tests/packaging/test_project.py @@ -57,6 +57,28 @@ def test_tracked_source_roots_contain_no_generated_native_artifacts() -> None: assert not list((REPOSITORY_ROOT / "mooncake-pg" / "torch").rglob("*.so")) +def test_lightweight_modules_have_one_authoritative_source() -> None: + package_root = REPOSITORY_ROOT / "python" / "mooncake" + legacy_package_root = REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" + + for module in ("buffer_pool.py", "mooncake_config.py"): + assert (package_root / module).is_file() + assert not (legacy_package_root / module).exists() + + assert ( + REPOSITORY_ROOT / "python" / "tests" / "unit" / "test_mooncake_config.py" + ).is_file() + assert ( + REPOSITORY_ROOT / "python" / "tests" / "store" / "test_buffer_pool_native.py" + ).is_file() + assert not ( + REPOSITORY_ROOT / "mooncake-wheel" / "tests" / "test_mooncake_config.py" + ).exists() + assert not ( + REPOSITORY_ROOT / "mooncake-wheel" / "tests" / "test_buffer_pool_native.py" + ).exists() + + def test_pg_extension_build_stages_outside_the_source_tree( tmp_path: Path, ) -> None: diff --git a/mooncake-wheel/tests/test_buffer_pool_native.py b/python/tests/store/test_buffer_pool_native.py similarity index 98% rename from mooncake-wheel/tests/test_buffer_pool_native.py rename to python/tests/store/test_buffer_pool_native.py index 14067aedc0..0099f5c33b 100644 --- a/mooncake-wheel/tests/test_buffer_pool_native.py +++ b/python/tests/store/test_buffer_pool_native.py @@ -98,9 +98,7 @@ def test_buffer_pool_rejects_invalid_alignment() -> None: @pytest.mark.parametrize("size", [0, 1, 128 * 1024 + 1]) def test_buffer_pool_supports_arbitrary_sizes(size: int) -> None: store = create_store() - pool = BufferPool( - store, 1024 * 1024, max_size_class=128 * 1024, alignment=4096 - ) + pool = BufferPool(store, 1024 * 1024, max_size_class=128 * 1024, alignment=4096) lease = pool.acquire(size) assert lease.size == size diff --git a/mooncake-wheel/tests/test_mooncake_config.py b/python/tests/unit/test_mooncake_config.py similarity index 100% rename from mooncake-wheel/tests/test_mooncake_config.py rename to python/tests/unit/test_mooncake_config.py diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index eed257699d..ad45733d18 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=(buffer_pool.py mooncake_config.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..22f5f606cd 100755 --- a/scripts/test_installation.sh +++ b/scripts/test_installation.sh @@ -37,6 +37,7 @@ python -c "import mooncake.reshard.weight" && echo "Success: Reshard import succ echo "Running import structure test..." # Run the import structure test cp -r mooncake-wheel/tests test_env/ +cp -r python/tests/unit test_env/unit_tests cp -r mooncake-reshard/tests test_env/reshard_tests cd test_env pip install torch numpy @@ -45,7 +46,7 @@ python tests/test_fast_copy.py python tests/test_import_structure.py echo "Running mooncake config test..." -python tests/test_mooncake_config.py +python unit_tests/test_mooncake_config.py echo "Running reshard contract tests..." python -m pip install pytest hypothesis==6.141.0 From f07bd0774de51118dfbf86b921167f3dcedec635 Mon Sep 17 00:00:00 2001 From: Aionw Date: Sat, 22 Aug 2026 00:36:00 +0800 Subject: [PATCH 2/2] [Store] Migrate lightweight Python utilities --- .github/CODEOWNERS | 4 ++++ .github/labeler.yml | 7 +++++- mooncake-integration/CMakeLists.txt | 2 +- .../mooncake/buffer_pool.py | 0 .../mooncake/mooncake_config.py | 0 .../tests/packaging/test_installed_wheel.py | 2 ++ python/tests/packaging/test_project.py | 22 +++++++++++++++++++ .../tests/store}/test_buffer_pool_native.py | 4 +--- .../tests/unit}/test_mooncake_config.py | 0 scripts/build_wheel.sh | 4 +++- scripts/test_installation.sh | 3 ++- 11 files changed, 41 insertions(+), 7 deletions(-) rename {mooncake-wheel => python}/mooncake/buffer_pool.py (100%) rename {mooncake-wheel => python}/mooncake/mooncake_config.py (100%) rename {mooncake-wheel/tests => python/tests/store}/test_buffer_pool_native.py (98%) rename {mooncake-wheel/tests => python/tests/unit}/test_mooncake_config.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 3d83b130e1..8a9451bd76 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -19,6 +19,10 @@ /python/mooncake/mooncake_elastic_buffer.py @UNIDY2002 @ympcMark @yuechen-sys /python/mooncake/mooncake_ep_buffer.py @UNIDY2002 @ympcMark @yuechen-sys /python/tests/ep/ @UNIDY2002 @ympcMark @yuechen-sys +/python/mooncake/buffer_pool.py @ykwd @stmatengss +/python/mooncake/mooncake_config.py @ykwd @stmatengss +/python/tests/store/ @ykwd @stmatengss +/python/tests/unit/test_mooncake_config.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 36fc9db2ee..b21d3b2cfb 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -5,7 +5,12 @@ run-ci: store: - changed-files: - - any-glob-to-any-file: 'mooncake-store/**/*' + - any-glob-to-any-file: + - 'mooncake-store/**/*' + - 'python/mooncake/buffer_pool.py' + - 'python/mooncake/mooncake_config.py' + - 'python/tests/store/**/*' + - 'python/tests/unit/test_mooncake_config.py' Transfer Engine: - changed-files: diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index ccccd0179f..0e20ba2359 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -208,7 +208,7 @@ if(NOT SKBUILD) install( FILES "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_store_service.py" - "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_config.py" + "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/mooncake_config.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli.py" DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR}) diff --git a/mooncake-wheel/mooncake/buffer_pool.py b/python/mooncake/buffer_pool.py similarity index 100% rename from mooncake-wheel/mooncake/buffer_pool.py rename to python/mooncake/buffer_pool.py diff --git a/mooncake-wheel/mooncake/mooncake_config.py b/python/mooncake/mooncake_config.py similarity index 100% rename from mooncake-wheel/mooncake/mooncake_config.py rename to python/mooncake/mooncake_config.py diff --git a/python/tests/packaging/test_installed_wheel.py b/python/tests/packaging/test_installed_wheel.py index 61fd311f12..652aaab137 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/python/tests/packaging/test_installed_wheel.py @@ -45,6 +45,7 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: from pathlib import Path import mooncake import mooncake.engine +import mooncake.mooncake_config import mooncake.reshard import mooncake.store @@ -60,6 +61,7 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: "mooncake_elastic_buffer.py", ): assert (package_path.parent / ep_module).is_file(), ep_module +assert mooncake.mooncake_config.MooncakeConfig 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 73ee00e893..50a754b8b7 100644 --- a/python/tests/packaging/test_project.py +++ b/python/tests/packaging/test_project.py @@ -89,6 +89,28 @@ def test_ep_modules_have_one_authoritative_source() -> None: assert not legacy_test.exists() +def test_lightweight_modules_have_one_authoritative_source() -> None: + package_root = REPOSITORY_ROOT / "python" / "mooncake" + legacy_package_root = REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" + + for module in ("buffer_pool.py", "mooncake_config.py"): + assert (package_root / module).is_file() + assert not (legacy_package_root / module).exists() + + assert ( + REPOSITORY_ROOT / "python" / "tests" / "unit" / "test_mooncake_config.py" + ).is_file() + assert ( + REPOSITORY_ROOT / "python" / "tests" / "store" / "test_buffer_pool_native.py" + ).is_file() + assert not ( + REPOSITORY_ROOT / "mooncake-wheel" / "tests" / "test_mooncake_config.py" + ).exists() + assert not ( + REPOSITORY_ROOT / "mooncake-wheel" / "tests" / "test_buffer_pool_native.py" + ).exists() + + def test_pg_extension_build_stages_outside_the_source_tree( tmp_path: Path, ) -> None: diff --git a/mooncake-wheel/tests/test_buffer_pool_native.py b/python/tests/store/test_buffer_pool_native.py similarity index 98% rename from mooncake-wheel/tests/test_buffer_pool_native.py rename to python/tests/store/test_buffer_pool_native.py index 14067aedc0..0099f5c33b 100644 --- a/mooncake-wheel/tests/test_buffer_pool_native.py +++ b/python/tests/store/test_buffer_pool_native.py @@ -98,9 +98,7 @@ def test_buffer_pool_rejects_invalid_alignment() -> None: @pytest.mark.parametrize("size", [0, 1, 128 * 1024 + 1]) def test_buffer_pool_supports_arbitrary_sizes(size: int) -> None: store = create_store() - pool = BufferPool( - store, 1024 * 1024, max_size_class=128 * 1024, alignment=4096 - ) + pool = BufferPool(store, 1024 * 1024, max_size_class=128 * 1024, alignment=4096) lease = pool.acquire(size) assert lease.size == size diff --git a/mooncake-wheel/tests/test_mooncake_config.py b/python/tests/unit/test_mooncake_config.py similarity index 100% rename from mooncake-wheel/tests/test_mooncake_config.py rename to python/tests/unit/test_mooncake_config.py diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index 005f91aad7..5ddcafdb2f 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -180,11 +180,13 @@ if [ "$NPU_BUILD" = "1" ]; then fi echo "Building wheel package..." -# Stage migrated EP modules and the legacy Reshard package for the combined +# 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=( + buffer_pool.py + mooncake_config.py ep.py mooncake_ep_buffer.py mooncake_elastic_buffer.py diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh index dd8a7bb6b1..22f5f606cd 100755 --- a/scripts/test_installation.sh +++ b/scripts/test_installation.sh @@ -37,6 +37,7 @@ python -c "import mooncake.reshard.weight" && echo "Success: Reshard import succ echo "Running import structure test..." # Run the import structure test cp -r mooncake-wheel/tests test_env/ +cp -r python/tests/unit test_env/unit_tests cp -r mooncake-reshard/tests test_env/reshard_tests cd test_env pip install torch numpy @@ -45,7 +46,7 @@ python tests/test_fast_copy.py python tests/test_import_structure.py echo "Running mooncake config test..." -python tests/test_mooncake_config.py +python unit_tests/test_mooncake_config.py echo "Running reshard contract tests..." python -m pip install pytest hypothesis==6.141.0