From 1ecefe720df77ca5f473c21fe425bd96cf355a43 Mon Sep 17 00:00:00 2001 From: Aionw Date: Mon, 24 Aug 2026 16:35:25 +0800 Subject: [PATCH] [Wheel] Migrate DataProto catalog to unified Python package --- .github/CODEOWNERS | 4 ++- .github/labeler.yml | 6 +++- python/CMakeLists.txt | 5 ++- .../mooncake/dataproto_catalog.py | 0 .../dataproto}/test_dataproto_catalog.py | 32 +++++++++++++++++++ .../tests/packaging/test_installed_wheel.py | 4 +++ python/tests/packaging/test_project.py | 21 ++++++++++++ scripts/build_wheel.sh | 22 +++++++++---- scripts/test_installation.sh | 4 +++ 9 files changed, 89 insertions(+), 9 deletions(-) rename {mooncake-wheel => python}/mooncake/dataproto_catalog.py (100%) rename {mooncake-wheel/tests => python/tests/integration/dataproto}/test_dataproto_catalog.py (83%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2c9f901b16..ec4058d8cd 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/dataproto_catalog.py @ShangmingCai @stmatengss @zxpdemonio +/python/tests/integration/dataproto/ @ShangmingCai @stmatengss @zxpdemonio /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..5626f60963 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -21,7 +21,10 @@ Mooncake EP: Installation: - changed-files: - - any-glob-to-any-file: 'mooncake-wheel/**/*' + - any-glob-to-any-file: + - 'mooncake-wheel/**/*' + - 'python/mooncake/dataproto_catalog.py' + - 'python/tests/integration/dataproto/**/*' P2P Store: - changed-files: @@ -57,6 +60,7 @@ Tests: - 'scripts/test_*' - 'mooncake-wheel/tests/**/*' - 'mooncake-reshard/tests/**/*' + - 'python/tests/integration/dataproto/**/*' - 'scripts/tone_tests/**/*' Ascend/NPU: diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 6cb7b26ebf..5d44e41b71 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -22,7 +22,10 @@ install( COMPONENT python FILES_MATCHING PATTERN "*.py" - PATTERN "__init__.py" EXCLUDE) + PATTERN "__init__.py" EXCLUDE + # Migrated modules come from tool.scikit-build.wheel.packages. Exclude any + # stale file left by the transitional legacy wheel staging script. + PATTERN "dataproto_catalog.py" EXCLUDE) install( DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-reshard/python/mooncake/reshard/" diff --git a/mooncake-wheel/mooncake/dataproto_catalog.py b/python/mooncake/dataproto_catalog.py similarity index 100% rename from mooncake-wheel/mooncake/dataproto_catalog.py rename to python/mooncake/dataproto_catalog.py diff --git a/mooncake-wheel/tests/test_dataproto_catalog.py b/python/tests/integration/dataproto/test_dataproto_catalog.py similarity index 83% rename from mooncake-wheel/tests/test_dataproto_catalog.py rename to python/tests/integration/dataproto/test_dataproto_catalog.py index ac81c420cf..9b417260e1 100644 --- a/mooncake-wheel/tests/test_dataproto_catalog.py +++ b/python/tests/integration/dataproto/test_dataproto_catalog.py @@ -1,10 +1,42 @@ from __future__ import annotations +import os +from pathlib import Path +import subprocess +import sys + import pytest from mooncake.dataproto_catalog import DataProtoCatalog +def test_catalog_source_import_keeps_structured_object_import_lazy() -> None: + repository_root = Path(__file__).resolve().parents[4] + environment = os.environ.copy() + environment["PYTHONPATH"] = os.pathsep.join( + ( + str(repository_root / "mooncake-wheel"), + str(repository_root / "python"), + ) + ) + environment["PYTHONNOUSERSITE"] = "1" + + subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; import mooncake.dataproto_catalog as catalog; " + "assert catalog.DataProtoCatalog is not None; " + "assert 'mooncake.structured_object_store' not in sys.modules" + ), + ], + cwd=repository_root, + env=environment, + check=True, + ) + + def handle(name: str, fields: list[str], batch_size: int = 2) -> dict: return { "type": "mooncake_dataproto_ref", diff --git a/python/tests/packaging/test_installed_wheel.py b/python/tests/packaging/test_installed_wheel.py index 74422d9ebb..0ab3a7cbe4 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/python/tests/packaging/test_installed_wheel.py @@ -43,7 +43,9 @@ 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 mooncake +import mooncake.dataproto_catalog import mooncake.engine import mooncake.reshard import mooncake.store @@ -53,7 +55,9 @@ def test_wheel_imports_outside_the_repository(tmp_path: Path) -> None: assert not package_path.is_relative_to(repository_path), (package_path, repository_path) assert metadata.version("mooncake-transfer-engine") == {_project_version()!r} assert mooncake.BufferPool is mooncake.store.BufferPool +assert mooncake.dataproto_catalog.DataProtoCatalog is not None assert mooncake.engine.TransferEngine is not None +assert "mooncake.structured_object_store" not in sys.modules """ 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..5fa9f2e93a 100644 --- a/python/tests/packaging/test_project.py +++ b/python/tests/packaging/test_project.py @@ -47,6 +47,10 @@ def test_dependency_boundaries_are_declared() -> None: "structured", "vllm", } + assert set(metadata["optional-dependencies"]["structured"]) == { + "numpy", + "pillow", + } def test_tracked_source_roots_contain_no_generated_native_artifacts() -> None: @@ -57,6 +61,23 @@ def test_tracked_source_roots_contain_no_generated_native_artifacts() -> None: assert not list((REPOSITORY_ROOT / "mooncake-pg" / "torch").rglob("*.so")) +def test_dataproto_catalog_has_one_authoritative_source() -> None: + package_root = REPOSITORY_ROOT / "python" / "mooncake" + tests_root = REPOSITORY_ROOT / "python" / "tests" + legacy_root = REPOSITORY_ROOT / "mooncake-wheel" + + assert (package_root / "dataproto_catalog.py").is_file() + assert ( + tests_root / "integration" / "dataproto" / "test_dataproto_catalog.py" + ).is_file() + assert not (legacy_root / "mooncake" / "dataproto_catalog.py").exists() + assert not (legacy_root / "tests" / "test_dataproto_catalog.py").exists() + assert ( + 'PATTERN "dataproto_catalog.py" EXCLUDE' + in (REPOSITORY_ROOT / "python" / "CMakeLists.txt").read_text() + ) + + def test_pg_extension_build_stages_outside_the_source_tree( tmp_path: Path, ) -> None: diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index eed257699d..49280b3a53 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=(dataproto_catalog.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..91e6868d03 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/integration/dataproto test_env/dataproto_tests cp -r mooncake-reshard/tests test_env/reshard_tests cd test_env pip install torch numpy @@ -51,6 +52,9 @@ echo "Running reshard contract tests..." python -m pip install pytest hypothesis==6.141.0 python -m pytest reshard_tests -q +echo "Running DataProto catalog tests..." +python -m pytest dataproto_tests -q + echo "Verifying mooncake_master entry point..." # Check if the mooncake_master entry point is installed and executable which mooncake_master || { echo "ERROR: mooncake_master entry point not found!"; exit 1; }