[Wheel] Migrate Python CLI modules to unified package

.github/CODEOWNERS

@@ -21,6 +21,12 @@
 /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/cli.py @ykwd @stmatengss @zxpdemonio
+/python/mooncake/cli_client.py @ykwd @stmatengss @zxpdemonio
+/python/mooncake/cli_bench.py @ShangmingCai @alogfans @stmatengss
+/python/mooncake/transfer_engine_topology_dump.py @ShangmingCai @alogfans @stmatengss
+/python/tests/integration/test_cli.py @ShangmingCai @stmatengss @zxpdemonio
+/python/tests/unit/test_cli_modules.py @ShangmingCai @stmatengss @zxpdemonio
 /python/tests/store/ @ykwd @stmatengss
 /python/tests/unit/test_mooncake_config.py @ykwd @stmatengss
 /python/mooncake/mooncake_ssd_*.py @ykwd @stmatengss

.github/labeler.yml

@@ -7,6 +7,8 @@ store:
   - changed-files:
     - any-glob-to-any-file:
       - 'mooncake-store/**/*'
+      - 'python/mooncake/cli.py'
+      - 'python/mooncake/cli_client.py'
       - 'python/mooncake/buffer_pool.py'
       - 'python/mooncake/mooncake_config.py'
       - 'python/tests/store/**/*'
@@ -21,7 +23,10 @@ Mooncake Conductor:
 
 Transfer Engine:
   - changed-files:
-    - any-glob-to-any-file: 'mooncake-transfer-engine/**/*'
+    - any-glob-to-any-file:
+      - 'mooncake-transfer-engine/**/*'
+      - 'python/mooncake/cli_bench.py'
+      - 'python/mooncake/transfer_engine_topology_dump.py'
 
 PyTorch Backend:
   - changed-files:
@@ -38,7 +43,12 @@ Mooncake EP:
 
 Installation:
   - changed-files:
-    - any-glob-to-any-file: 'mooncake-wheel/**/*'
+    - any-glob-to-any-file:
+      - 'mooncake-wheel/**/*'
+      - 'python/mooncake/cli.py'
+      - 'python/mooncake/cli_client.py'
+      - 'python/mooncake/cli_bench.py'
+      - 'python/mooncake/transfer_engine_topology_dump.py'
 
 P2P Store:
   - changed-files:

mooncake-integration/CMakeLists.txt

@@ -220,8 +220,9 @@ if(NOT SKBUILD)
   install(
     FILES
       "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/http_metadata_server.py"
-      "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli_bench.py"
-      "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/transfer_engine_topology_dump.py"
+      "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/_launcher.py"
+      "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/cli_bench.py"
+      "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/transfer_engine_topology_dump.py"
       "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_connector_v1.py"
       "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/vllm_v1_proxy_server.py"
     DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR})
@@ -231,7 +232,8 @@ if(NOT SKBUILD)
       FILES
         "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_store_service.py"
         "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/mooncake_config.py"
-        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli.py"
+        "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/cli.py"
+        "${CMAKE_CURRENT_SOURCE_DIR}/../python/mooncake/cli_client.py"
       DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR})
 
     if(USE_NOF)

python/mooncake/__init__.py

@@ -1,5 +1,9 @@
 """Public Python package for Mooncake."""
 
+from pkgutil import extend_path
+
 from mooncake.buffer_pool import BufferPool, RegisteredBufferPool
 
+__path__ = extend_path(__path__, __name__)
+
 __all__ = ["BufferPool", "RegisteredBufferPool"]

python/mooncake/_launcher.py

@@ -0,0 +1,40 @@
+"""
+Locate and launch native binaries bundled with the Mooncake wheel.
+
+The CLI entry points in ``mooncake.cli`` / ``mooncake.cli_client`` /
+``mooncake.cli_bench`` are Python shims that ``exec`` a compiled C++ binary
+shipped inside the wheel. The binary is resolved through
+:func:`importlib.resources.files` so it stays agnostic to where the package is
+installed, and is made executable on the fly because ``setuptools``
+``package-data`` does not preserve the execute bit.
+"""
+
+from __future__ import annotations
+
+import os
+import stat
+from importlib.resources import files
+
+
+def locate(name: str) -> str:
+    """
+    Return the absolute path of a native binary bundled in this package.
+
+    The named binary must live next to the package (``mooncake/<name>``). It is
+    chmodded to be executable if needed. A clear :class:`FileNotFoundError` is
+    raised when the binary is absent instead of silently falling back to
+    ``PATH`` (which could run a mismatched system build).
+    """
+    path = files("mooncake") / name
+    if not path.is_file():
+        raise FileNotFoundError(
+            f"{name!r} not found in the mooncake package; "
+            "reinstall the package so the native CLI binary is bundled"
+        )
+    file_path = str(path)
+    if not os.access(file_path, os.X_OK):
+        os.chmod(
+            file_path,
+            os.stat(file_path).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH,
+        )
+    return file_path

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

@@ -4,23 +4,17 @@
 """
 
 import os
-import stat
 import sys
 
+from mooncake._launcher import locate
+
 
 def main():
     """
     Main entry point for the mooncake_master command.
     Simply runs the mooncake_master binary with all arguments passed through.
     """
-    # Get the path to the mooncake_master binary
-    package_dir = os.path.dirname(os.path.abspath(__file__))
-    bin_path = os.path.join(package_dir, "mooncake_master")
-
-    # Make sure the binary is executable
-    if not os.access(bin_path, os.X_OK):
-        st = os.stat(bin_path)
-        os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+    bin_path = locate("mooncake_master")
 
     # Preserve the CLI process ID so callers can reliably stop the server.
     os.execv(bin_path, [bin_path] + sys.argv[1:])

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

@@ -3,25 +3,18 @@
 Minimal CLI module for transfer_engine_bench.
 """
 
-import os
-import stat
-import sys
 import subprocess
+import sys
+
+from mooncake._launcher import locate
 
 
 def main():
     """
     Main entry point for the transfer_engine_bench command.
     Simply runs the transfer_engine_bench binary with all arguments passed through.
     """
-    # Get the path to the transfer_engine_bench binary
-    package_dir = os.path.dirname(os.path.abspath(__file__))
-    bin_path = os.path.join(package_dir, "transfer_engine_bench")
-
-    # Make sure the binary is executable
-    if not os.access(bin_path, os.X_OK):
-        st = os.stat(bin_path)
-        os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+    bin_path = locate("transfer_engine_bench")
 
     # Run the binary with all arguments passed through
     return subprocess.call([bin_path] + sys.argv[1:])

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

@@ -4,23 +4,17 @@
 """
 
 import os
-import stat
 import sys
 
+from mooncake._launcher import locate
+
 
 def main():
     """
     Main entry point for the mooncake_client command.
     Simply runs the mooncake_client binary with all arguments passed through.
     """
-    # Get the path to the mooncake_client binary
-    package_dir = os.path.dirname(os.path.abspath(__file__))
-    bin_path = os.path.join(package_dir, "mooncake_client")
-
-    # Make sure the binary is executable
-    if not os.access(bin_path, os.X_OK):
-        st = os.stat(bin_path)
-        os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+    bin_path = locate("mooncake_client")
 
     # Preserve the CLI process ID so callers can reliably stop the client.
     os.execv(bin_path, [bin_path] + sys.argv[1:])

mooncake-wheel/tests/test_cli.py → python/tests/integration/test_cli.py

@@ -13,9 +13,7 @@ def test_entry_point_installed():
     try:
         # Check if mooncake_master is in PATH
         result = subprocess.run(
-            ["which", "mooncake_master"],
-            capture_output=True,
-            text=True
+            ["which", "mooncake_master"], capture_output=True, text=True
         )
 
         if result.returncode != 0:
@@ -24,9 +22,7 @@ def test_entry_point_installed():
 
         print(f"✅ mooncake_master entry point found at: {result.stdout.strip()}")
         result = subprocess.run(
-            ["which", "mooncake_client"],
-            capture_output=True,
-            text=True
+            ["which", "mooncake_client"], capture_output=True, text=True
         )
 
         if result.returncode != 0:
@@ -35,9 +31,7 @@ def test_entry_point_installed():
 
         print(f"✅ mooncake_client entry point found at: {result.stdout.strip()}")
         result = subprocess.run(
-            ["which", "transfer_engine_bench"],
-            capture_output=True,
-            text=True
+            ["which", "transfer_engine_bench"], capture_output=True, text=True
         )
 
         if result.returncode != 0:
@@ -67,10 +61,14 @@ def test_run_master_and_client():
         #    and the default 9003 is still held by a lingering master from the
         #    prior CXL test step when this smoke test runs.
         process = subprocess.Popen(
-            ["mooncake_master", "--port=61351", "--max_threads=2",
-             "--metrics_port=19003"],
+            [
+                "mooncake_master",
+                "--port=61351",
+                "--max_threads=2",
+                "--metrics_port=19003",
+            ],
             stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE
+            stderr=subprocess.PIPE,
         )
 
         # Give it a moment to start
@@ -80,9 +78,13 @@ def test_run_master_and_client():
         if process.poll() is None:
             print("✅ mooncake_master process started successfully")
             client_process = subprocess.Popen(
-                ["mooncake_client", "--master_server_address=127.0.0.1:61351", "--port=61352"],
+                [
+                    "mooncake_client",
+                    "--master_server_address=127.0.0.1:61351",
+                    "--port=61352",
+                ],
                 stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE
+                stderr=subprocess.PIPE,
             )
 
             # Give the client some time to connect
@@ -95,7 +97,7 @@ def test_run_master_and_client():
                 print("✅ mooncake_client process terminated successfully")
             else:
                 stdout, stderr = client_process.communicate()
-                print(f"❌ mooncake_client failed to start")
+                print("❌ mooncake_client failed to start")
                 print(f"stdout: {stdout.decode()}")
                 print(f"stderr: {stderr.decode()}")
             # Terminate the process
@@ -105,7 +107,7 @@ def test_run_master_and_client():
             return True
         else:
             stdout, stderr = process.communicate()
-            print(f"❌ mooncake_master process failed to start")
+            print("❌ mooncake_master process failed to start")
             print(f"stdout: {stdout.decode()}")
             print(f"stderr: {stderr.decode()}")
             return False