[Wheel] Migrate DataProto catalog to unified Python package

.github/CODEOWNERS

@@ -37,6 +37,8 @@
 /python/mooncake/mooncake_ssd_*.py @ykwd @stmatengss
 /python/mooncake/spdk_tgt_create.py @ykwd @stmatengss
 /python/tests/ssd/ @ykwd @stmatengss
+/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

.github/labeler.yml

@@ -53,6 +53,8 @@ Installation:
       - 'python/mooncake/cli_client.py'
       - 'python/mooncake/cli_bench.py'
       - 'python/mooncake/transfer_engine_topology_dump.py'
+      - 'python/mooncake/dataproto_catalog.py'
+      - 'python/tests/integration/dataproto/**/*'
 
 P2P Store:
   - changed-files:

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/"

mooncake-wheel/mooncake/dataproto_catalog.py → python/mooncake/dataproto_catalog.py

@@ -629,9 +629,7 @@ def _publish(self, publication: Mapping[str, Any]) -> dict[str, Any]:
             *publication["args"],
             **publication["kwargs"],
         )
-        if publication["method"] == "publish_append" and result.get(
-            "append_rejected"
-        ):
+        if publication["method"] == "publish_append" and result.get("append_rejected"):
             self.transfer.cleanup_dataproto_append(
                 publication["kwargs"]["previous_handle"],
                 publication["kwargs"]["handle"],

mooncake-wheel/tests/test_dataproto_catalog.py → python/tests/integration/dataproto/test_dataproto_catalog.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
 import copy
+import os
+from pathlib import Path
+import subprocess
+import sys
 from types import SimpleNamespace
 
 import pytest
@@ -9,6 +13,33 @@
 from mooncake.structured_object_store import import_dataproto_ref
 
 
+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 / "python"),
+            str(repository_root / "mooncake-wheel"),
+        )
+    )
+    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",
@@ -178,14 +209,17 @@ def test_catalog_append_replaces_managed_handle_ownership() -> None:
         "member": "batch.score",
         "section": "batch",
     }
-    assert "fragment identity" in catalog.publish_append(
-        "wrong-partition",
-        "manifest/base",
-        "other",
-        ["a"],
-        previous_handle=base,
-        handle=appended,
-    )["append_rejected"]
+    assert (
+        "fragment identity"
+        in catalog.publish_append(
+            "wrong-partition",
+            "manifest/base",
+            "other",
+            ["a"],
+            previous_handle=base,
+            handle=appended,
+        )["append_rejected"]
+    )
     catalog.publish_append(
         "append-op",
         "manifest/base",
@@ -202,14 +236,17 @@ def test_catalog_append_replaces_managed_handle_ownership() -> None:
         previous_handle=base,
         handle=appended,
     )["fields"] == ["input", "value", "score"]
-    assert "stale" in catalog.publish_append(
-        "stale-op",
-        "manifest/base",
-        "train",
-        ["a"],
-        previous_handle=base,
-        handle=appended,
-    )["append_rejected"]
+    assert (
+        "stale"
+        in catalog.publish_append(
+            "stale-op",
+            "manifest/base",
+            "train",
+            ["a"],
+            previous_handle=base,
+            handle=appended,
+        )["append_rejected"]
+    )
 
     plan = catalog.resolve("train", ["a"])
     assert plan["fields"] == ["input", "value", "score"]

python/tests/packaging/test_installed_wheel.py

@@ -77,6 +77,7 @@ def test_wheel_imports_outside_the_repository(
 assert "mooncake.engine" not in sys.modules
 
 import mooncake.async_store
+import mooncake.dataproto_catalog
 import mooncake.engine
 import mooncake.http_metadata_server
 import mooncake.mooncake_config
@@ -94,9 +95,11 @@ def test_wheel_imports_outside_the_repository(
     mooncake.store.MooncakeDistributedStore,
 )
 assert Path(mooncake.async_store.__file__).resolve().parent == package_path.parent
+assert mooncake.dataproto_catalog.DataProtoCatalog is not None
 assert mooncake.engine.TransferEngine is not None
 assert mooncake.http_metadata_server.KVBootstrapServer is not None
 assert mooncake.mooncake_config.MooncakeConfig is not None
+assert "mooncake.structured_object_store" not in sys.modules
 for ep_module in (
     "ep.py",
     "mooncake_ep_buffer.py",

python/tests/packaging/test_project.py

@@ -61,6 +61,10 @@ def test_dependency_boundaries_are_declared() -> None:
         "vllm",
     }
     assert metadata["optional-dependencies"]["administration"] == ["paramiko"]
+    assert set(metadata["optional-dependencies"]["structured"]) == {
+        "numpy",
+        "pillow",
+    }
 
 
 def test_legacy_wheel_administration_extra_matches_root_project() -> None:
@@ -247,7 +251,6 @@ def test_scikit_build_consumes_unified_python_sources() -> None:
         "vllm_v1_proxy_server.py",
         "pg.py",
         "structured_object_store.py",
-        "dataproto_catalog.py",
     ):
         assert (REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" / module).is_file()
 
@@ -357,6 +360,23 @@ def test_ssd_administration_modules_have_one_authoritative_source() -> None:
     ).exists()
 
 
+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:

scripts/build_wheel.sh

@@ -200,6 +200,7 @@ MIGRATED_PYTHON_MODULES=(
     mooncake_ssd_register.py
     mooncake_ssd_unregister.py
     spdk_tgt_create.py
+    dataproto_catalog.py
 )
 RESHARD_SOURCE_DIR="mooncake-reshard/python/mooncake/reshard"
 RESHARD_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake/reshard"

scripts/test_installation.sh

@@ -42,6 +42,7 @@ cp -r python/tests/services test_env/service_tests
 cp -r python/tests/unit test_env/unit_tests
 cp -r mooncake-reshard/tests test_env/reshard_tests
 cp -r python/tests/ssd test_env/ssd_tests
+cp -r python/tests/integration/dataproto test_env/dataproto_tests
 cd test_env
 pip install torch numpy
 python -c "import mooncake._fast_copy"
@@ -58,6 +59,9 @@ echo "Running reshard contract tests..."
 python -m pip install pytest hypothesis==6.141.0 paramiko
 python -m pytest reshard_tests ssd_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; }