diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 9ec789e2c2..b3bae46df8 100644 --- a/.github/CODEOWNERS +++ b/.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 diff --git a/.github/labeler.yml b/.github/labeler.yml index 718efa7d62..b6d844db84 100644 --- a/.github/labeler.yml +++ b/.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: diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 699ffa43aa..2be1e45f00 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/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) diff --git a/python/mooncake/__init__.py b/python/mooncake/__init__.py index 8bae873c46..4169a5edc3 100644 --- a/python/mooncake/__init__.py +++ b/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"] diff --git a/python/mooncake/_launcher.py b/python/mooncake/_launcher.py new file mode 100644 index 0000000000..092a97c7ab --- /dev/null +++ b/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/``). 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 diff --git a/mooncake-wheel/mooncake/cli.py b/python/mooncake/cli.py similarity index 53% rename from mooncake-wheel/mooncake/cli.py rename to python/mooncake/cli.py index da9af982ba..cfdef22504 100644 --- a/mooncake-wheel/mooncake/cli.py +++ b/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:]) diff --git a/mooncake-wheel/mooncake/cli_bench.py b/python/mooncake/cli_bench.py similarity index 52% rename from mooncake-wheel/mooncake/cli_bench.py rename to python/mooncake/cli_bench.py index f4bf27d766..421ee1400d 100644 --- a/mooncake-wheel/mooncake/cli_bench.py +++ b/python/mooncake/cli_bench.py @@ -3,10 +3,10 @@ Minimal CLI module for transfer_engine_bench. """ -import os -import stat -import sys import subprocess +import sys + +from mooncake._launcher import locate def main(): @@ -14,14 +14,7 @@ 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:]) diff --git a/mooncake-wheel/mooncake/cli_client.py b/python/mooncake/cli_client.py similarity index 53% rename from mooncake-wheel/mooncake/cli_client.py rename to python/mooncake/cli_client.py index adcf1d3bf7..226c943fad 100644 --- a/mooncake-wheel/mooncake/cli_client.py +++ b/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:]) diff --git a/mooncake-wheel/mooncake/transfer_engine_topology_dump.py b/python/mooncake/transfer_engine_topology_dump.py similarity index 100% rename from mooncake-wheel/mooncake/transfer_engine_topology_dump.py rename to python/mooncake/transfer_engine_topology_dump.py diff --git a/mooncake-wheel/tests/test_cli.py b/python/tests/integration/test_cli.py similarity index 85% rename from mooncake-wheel/tests/test_cli.py rename to python/tests/integration/test_cli.py index 885411de24..3a1fc6946b 100644 --- a/mooncake-wheel/tests/test_cli.py +++ b/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 diff --git a/python/tests/packaging/test_installed_wheel.py b/python/tests/packaging/test_installed_wheel.py index 4945f58d26..54d750fef6 100644 --- a/python/tests/packaging/test_installed_wheel.py +++ b/python/tests/packaging/test_installed_wheel.py @@ -2,6 +2,7 @@ import os from pathlib import Path +import stat import subprocess import sys @@ -55,6 +56,22 @@ def test_wheel_imports_outside_the_repository( from pathlib import Path import sys import mooncake + +cli_modules = ( + "mooncake.cli", + "mooncake.cli_client", + "mooncake.cli_bench", + "mooncake.transfer_engine_topology_dump", +) +assert all(module not in sys.modules for module in cli_modules) + +import mooncake.cli +import mooncake.cli_bench +import mooncake.cli_client +import mooncake.transfer_engine_topology_dump + +assert "mooncake.engine" not in sys.modules + import mooncake.engine import mooncake.mooncake_config import mooncake.reshard @@ -74,6 +91,18 @@ def test_wheel_imports_outside_the_repository( "mooncake_elastic_buffer.py", ): assert (package_path.parent / ep_module).is_file(), ep_module +entry_points = {{ + entry_point.name: entry_point.value + for entry_point in metadata.entry_points(group="console_scripts") +}} +expected_entry_points = {{ + "mooncake_master": "mooncake.cli:main", + "mooncake_client": "mooncake.cli_client:main", + "transfer_engine_bench": "mooncake.cli_bench:main", + "transfer_engine_topology_dump": "mooncake.transfer_engine_topology_dump:main", +}} +assert {{name: entry_points[name] for name in expected_entry_points}} == expected_entry_points + for module in ( "mooncake.mooncake_ssd_register", "mooncake.mooncake_ssd_unregister", @@ -98,6 +127,69 @@ def test_wheel_imports_outside_the_repository( check=True, ) + package_query = "from pathlib import Path; import mooncake; print(Path(mooncake.__file__).parent)" + package_result = subprocess.run( + [str(python), "-I", "-c", package_query], + cwd=tmp_path, + env=clean_environment, + check=True, + capture_output=True, + text=True, + ) + package_directory = Path(package_result.stdout.strip()) + fake_binary = "#!/bin/sh\nprintf '%s\\n' \"$*\"\n" + for binary_name in ("mooncake_master", "mooncake_client", "transfer_engine_bench"): + binary = package_directory / binary_name + binary.write_text(fake_binary) + binary.chmod(binary.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + + module_commands = { + "mooncake_master": "mooncake.cli", + "mooncake_client": "mooncake.cli_client", + "transfer_engine_bench": "mooncake.cli_bench", + } + for command_name, module_name in module_commands.items(): + command_result = subprocess.run( + [str(environment / "bin" / command_name), "console-script-smoke"], + cwd=tmp_path, + env=clean_environment, + check=True, + capture_output=True, + text=True, + ) + assert command_result.stdout.strip() == "console-script-smoke" + + module_result = subprocess.run( + [str(python), "-I", "-m", module_name, "module-smoke"], + cwd=tmp_path, + env=clean_environment, + check=True, + capture_output=True, + text=True, + ) + assert module_result.stdout.strip() == "module-smoke" + + topology_commands = ( + [str(environment / "bin" / "transfer_engine_topology_dump"), "--help"], + [ + str(python), + "-I", + "-m", + "mooncake.transfer_engine_topology_dump", + "--help", + ], + ) + for command in topology_commands: + result = subprocess.run( + command, + cwd=tmp_path, + env=clean_environment, + check=True, + capture_output=True, + text=True, + ) + assert "Dump device topology" in result.stdout + def check_help(): for module in ( "mooncake.mooncake_ssd_register", diff --git a/python/tests/packaging/test_project.py b/python/tests/packaging/test_project.py index 0bb6ae5b62..f291dd612c 100644 --- a/python/tests/packaging/test_project.py +++ b/python/tests/packaging/test_project.py @@ -14,6 +14,18 @@ REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +CLI_MODULES = ( + "cli.py", + "cli_client.py", + "cli_bench.py", + "transfer_engine_topology_dump.py", +) +CLI_ENTRY_POINTS = { + "mooncake_master": "mooncake.cli:main", + "mooncake_client": "mooncake.cli_client:main", + "transfer_engine_bench": "mooncake.cli_bench:main", + "transfer_engine_topology_dump": "mooncake.transfer_engine_topology_dump:main", +} def test_scikit_build_core_is_the_only_build_backend() -> None: @@ -103,6 +115,111 @@ def test_ep_modules_have_one_authoritative_source() -> None: assert not legacy_test.exists() +def test_cli_has_one_authoritative_source_and_test_location() -> None: + package_root = REPOSITORY_ROOT / "python" / "mooncake" + legacy_package_root = REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" + + for module in CLI_MODULES: + assert (package_root / module).is_file() + assert not (legacy_package_root / module).exists() + + assert ( + REPOSITORY_ROOT / "python" / "tests" / "integration" / "test_cli.py" + ).is_file() + assert ( + REPOSITORY_ROOT / "python" / "tests" / "unit" / "test_cli_modules.py" + ).is_file() + assert not (REPOSITORY_ROOT / "mooncake-wheel" / "tests" / "test_cli.py").exists() + + +def test_scikit_build_consumes_unified_python_sources() -> None: + """Lock the Phase-2 cut-over invariants for the scikit-build wheel path. + + These are read-only source checks. The scikit-build backend uses the + repository-root CMakeLists.txt, which must pull in python/CMakeLists.txt + (the \"unified\" Python sources) while mooncake-integration installs the + CLI binaries next to the package so the console scripts resolve them. + + The legacy mooncake-wheel tree still supplies the not-yet-migrated pure- + Python modules at install time (Python/CMakeLists.txt DIRECTORY rule); + Phase 2 should move those into python/mooncake and drop the rule. + """ + root_cmake = (REPOSITORY_ROOT / "CMakeLists.txt").read_text() + # SKBUILD must route through the unified python/CMakeLists.txt. + assert "if(SKBUILD)" in root_cmake + assert "add_subdirectory(python)" in root_cmake + + python_cmake = (REPOSITORY_ROOT / "python" / "CMakeLists.txt").read_text() + assert "../mooncake-wheel/mooncake/" in python_cmake + assert 'PATTERN "__init__.py" EXCLUDE' in python_cmake + # 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 ( + "http_metadata_server.py", + "mooncake_store_service.py", + "mooncake_connector_v1.py", + "vllm_v1_proxy_server.py", + "pg.py", + "structured_object_store.py", + "dataproto_catalog.py", + ): + assert (REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" / module).is_file() + + integration_cmake = ( + REPOSITORY_ROOT / "mooncake-integration" / "CMakeLists.txt" + ).read_text() + # The store CLI binaries ship next to the package under SKBUILD so the + # mooncake.cli / mooncake.cli_client wrappers resolve them in place. + assert "install(TARGETS mooncake_master mooncake_client" in integration_cmake + assert "RUNTIME DESTINATION ${MOONCAKE_PYTHON_INSTALL_DIR}" in integration_cmake + assert "if(WITH_TE AND BUILD_EXAMPLES)" in integration_cmake + + +def test_python_package_init_keeps_namespace_extension() -> None: + """The unified python/mooncake package must behave like the legacy wheel + package: both extend the namespace so independently installed mooncake + subpackages (e.g. mooncake-reshard / mooncake-pg) remain importable.""" + init = (REPOSITORY_ROOT / "python" / "mooncake" / "__init__.py").read_text() + assert "extend_path" in init + assert "__path__ = extend_path(__path__, __name__)" in init + + legacy_init = ( + REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" / "__init__.py" + ).read_text() + assert "extend_path" in legacy_init + + +def test_cli_entry_points_remain_stable_across_build_interfaces() -> None: + for project_file in ( + REPOSITORY_ROOT / "pyproject.toml", + REPOSITORY_ROOT / "mooncake-wheel" / "pyproject.toml", + ): + project = tomllib.loads(project_file.read_text()) + scripts = project["project"]["scripts"] + assert {name: scripts[name] for name in CLI_ENTRY_POINTS} == CLI_ENTRY_POINTS + + +def test_cli_build_inputs_use_the_canonical_sources() -> None: + legacy_build = (REPOSITORY_ROOT / "scripts" / "build_wheel.sh").read_text() + assert 'MIGRATED_PYTHON_SOURCE_DIR="python/mooncake"' in legacy_build + for module in CLI_MODULES: + assert module in legacy_build + # The shared launcher must also be staged so the setuptools wheel has it. + assert "_launcher.py" in legacy_build + + integration_cmake = ( + REPOSITORY_ROOT / "mooncake-integration" / "CMakeLists.txt" + ).read_text() + for module in ( + "cli.py", + "cli_client.py", + "cli_bench.py", + "transfer_engine_topology_dump.py", + ): + assert f"../python/mooncake/{module}" in integration_cmake + assert f"../mooncake-wheel/mooncake/{module}" not in integration_cmake + + def test_lightweight_modules_have_one_authoritative_source() -> None: package_root = REPOSITORY_ROOT / "python" / "mooncake" legacy_package_root = REPOSITORY_ROOT / "mooncake-wheel" / "mooncake" diff --git a/python/tests/unit/test_cli_modules.py b/python/tests/unit/test_cli_modules.py new file mode 100644 index 0000000000..b90f4f9dcd --- /dev/null +++ b/python/tests/unit/test_cli_modules.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path +import stat +import subprocess +import sys +from types import ModuleType, SimpleNamespace + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "mooncake" + + +@pytest.fixture +def mooncake_pkg(monkeypatch: pytest.MonkeyPatch) -> ModuleType: + """Make ``mooncake`` importable as a real package backed by ``python/mooncake``. + + The wrappers do ``from mooncake._launcher import locate``. A stub is injected + into ``sys.modules`` so importing ``mooncake`` does not run the real + ``python/mooncake/__init__.py`` (which imports ``mooncake.buffer_pool``, + absent in the source tree). + """ + pkg = ModuleType("mooncake") + pkg.__path__ = [str(PACKAGE_ROOT)] + pkg.__package__ = "mooncake" + monkeypatch.setitem(sys.modules, "mooncake", pkg) + return pkg + + +def _load_module(filename: str) -> ModuleType: + module_path = PACKAGE_ROOT / filename + spec = importlib.util.spec_from_file_location( + f"_mooncake_cli_test_{module_path.stem}", module_path + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_CLI_CASES = [ + ("cli.py", "mooncake_master", True), + ("cli_client.py", "mooncake_client", True), + ("cli_bench.py", "transfer_engine_bench", False), +] + + +@pytest.mark.parametrize(("filename", "binary_name", "uses_exec"), _CLI_CASES) +def test_cli_wrappers_forward_arguments_via_launcher( + monkeypatch: pytest.MonkeyPatch, + mooncake_pkg: ModuleType, + filename: str, + binary_name: str, + uses_exec: bool, +) -> None: + module = _load_module(filename) + located = f"/pkg/mooncake/{binary_name}" + locate_calls: list[str] = [] + runner_calls: list[tuple[str, list[str]]] = [] + + def fake_locate(name: str) -> str: + locate_calls.append(name) + return located + + monkeypatch.setattr(module, "locate", fake_locate) + monkeypatch.setattr(module.sys, "argv", [filename, "--flag", "value"]) + + expected_arguments = [located, "--flag", "value"] + if uses_exec: + monkeypatch.setattr( + module.os, "execv", lambda path, args: runner_calls.append((path, args)) + ) + assert module.main() is None + else: + monkeypatch.setattr( + module.subprocess, + "call", + lambda args: runner_calls.append((args[0], args)) or 17, + ) + assert module.main() == 17 + + assert locate_calls == [binary_name] + assert runner_calls == [(located, expected_arguments)] + + +def test_launcher_locates_binary_and_sets_exec_bit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + launcher = _load_module("_launcher.py") + + class FakePath: + def __init__(self, path: str) -> None: + self.path = path + + def __truediv__(self, name: str) -> "FakePath": + return FakePath(os.path.join(self.path, name)) + + def is_file(self) -> bool: + return True + + def __str__(self) -> str: + return self.path + + def __fspath__(self) -> str: + return self.path + + chmod_calls: list[tuple[str, int]] = [] + mode = stat.S_IFREG | 0o640 + monkeypatch.setattr(launcher, "files", lambda package: FakePath("/pkg/mooncake")) + monkeypatch.setattr(launcher.os, "access", lambda _path, _mode: False) + monkeypatch.setattr( + launcher.os, "stat", lambda _path: SimpleNamespace(st_mode=mode) + ) + monkeypatch.setattr( + launcher.os, "chmod", lambda path, m: chmod_calls.append((path, m)) + ) + + assert launcher.locate("mooncake_master") == "/pkg/mooncake/mooncake_master" + expected_bits = stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + assert chmod_calls == [("/pkg/mooncake/mooncake_master", mode | expected_bits)] + + +def test_launcher_raises_when_binary_is_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + launcher = _load_module("_launcher.py") + + class FakePath: + def __truediv__(self, name: str) -> "FakePath": + return FakePath() + + def is_file(self) -> bool: + return False + + monkeypatch.setattr(launcher, "files", lambda package: FakePath()) + + with pytest.raises(FileNotFoundError): + launcher.locate("mooncake_master") + + +def test_topology_cli_import_keeps_native_backends_lazy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delitem(sys.modules, "mooncake.engine", raising=False) + monkeypatch.delitem(sys.modules, "tent", raising=False) + + module = _load_module("transfer_engine_topology_dump.py") + + assert module.resolve_backend("te") == "te" + assert "mooncake.engine" not in sys.modules + assert "tent" not in sys.modules + + +def test_topology_cli_runs_from_the_canonical_source_tree() -> None: + result = subprocess.run( + [ + sys.executable, + str(PACKAGE_ROOT / "transfer_engine_topology_dump.py"), + "--help", + ], + check=True, + capture_output=True, + text=True, + ) + + assert "Dump device topology" in result.stdout + assert "--custom-topo-json" in result.stdout diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index 9851dd0297..7b29417cf7 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -185,6 +185,11 @@ echo "Building wheel package..." MIGRATED_PYTHON_SOURCE_DIR="python/mooncake" MIGRATED_PYTHON_STAGING_DIR="$(pwd)/mooncake-wheel/mooncake" MIGRATED_PYTHON_MODULES=( + _launcher.py + cli.py + cli_bench.py + cli_client.py + transfer_engine_topology_dump.py buffer_pool.py mooncake_config.py ep.py diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh index 7da85c0daf..8e61dc6827 100755 --- a/scripts/run_tests.sh +++ b/scripts/run_tests.sh @@ -243,7 +243,7 @@ else fi echo "Running CLI entry point tests..." -python test_cli.py +python ../../python/tests/integration/test_cli.py ENABLE_EMBEDDED_METADATA_SERVER=true if [ "$RUN_TESTS_METADATA_SERVER_MODE" = "managed" ]; then diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh index ed900eeadd..04a9798efa 100755 --- a/scripts/test_installation.sh +++ b/scripts/test_installation.sh @@ -33,6 +33,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.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; } echo "Running import structure test..." # Run the import structure test @@ -68,6 +69,11 @@ 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 transfer_engine_topology_dump entry point..." +transfer_engine_topology_dump --help >/dev/null +python -m mooncake.transfer_engine_topology_dump --help >/dev/null +echo "Success: transfer_engine_topology_dump entry point and module execution work" + cd .. echo "Installation test completed successfully!"