Skip to content

[Wheel] Migrate Python CLI modules to unified package - #3626

Merged
Aionw merged 5 commits into
mainfrom
codex/migrate-python-cli
Sep 18, 2026
Merged

Aionw merged 5 commits into
mainfrom
codex/migrate-python-cli

Conversation

@Aionw

@Aionw Aionw commented Aug 24, 2026 •

Copy link
Copy Markdown
Collaborator

Description

Implements the general Python CLI subsystem slice of Phase 2 in RFC #3534.

  • move mooncake.cli, mooncake.cli_client, mooncake.cli_bench, and mooncake.transfer_engine_topology_dump into python/mooncake as their sole tracked source;
  • move the owned CLI integration test under python/tests and add focused source-tree wrapper/import coverage;
  • preserve all existing console-script mappings and python -m behavior;
  • keep Transfer Engine/TENT imports lazy until topology execution;
  • point direct CMake install inputs at the canonical sources;
  • stage the canonical CLI files only temporarily for the legacy Phase 2 wheel builder;
  • add source-location, build-input, wheel-manifest, isolated import, console-script, and module-execution packaging coverage;
  • update CLI ownership and labeler paths.

Part of #3534.

Scope and overlap

No open PR was found that migrates this CLI slice. PR #3585 migrates Reshard, and PR #3586 migrates lightweight Store utilities/configuration; this branch is based directly on main and contains neither subsystem. Those PRs touch some shared migration surfaces (scripts/build_wheel.sh, packaging tests, CODEOWNERS, and labeler configuration), but this PR changes those surfaces only for the four CLI modules.

Service entry points, SSD administration, native executables/bindings, Phase 3 native-extension renames, and the Phase 4/5 release cutover remain out of scope.

Module

  • Transfer Engine (mooncake-transfer-engine)
  • Mooncake Store (mooncake-store)
  • Reshard (mooncake-reshard)
  • Mooncake EP (mooncake-ep)
  • Mooncake PG (mooncake-pg)
  • Integration (mooncake-integration)
  • P2P Store (mooncake-p2p-store)
  • Python Wheel (mooncake-wheel)
  • Common (mooncake-common)
  • Mooncake RL (mooncake-rl)
  • CI/CD
  • Docs
  • Other

Type of Change

  • Bug fix
  • New feature
  • Refactor
  • Breaking change
  • Documentation update
  • Performance improvement
  • Other

How Has This Been Tested?

Test commands:

python -m pytest -q python/tests/unit/test_cli_modules.py python/tests/packaging/test_project.py
python -m py_compile python/mooncake/cli.py python/mooncake/cli_client.py python/mooncake/cli_bench.py python/mooncake/transfer_engine_topology_dump.py python/tests/integration/test_cli.py python/tests/unit/test_cli_modules.py python/tests/packaging/test_installed_wheel.py
bash -n scripts/build_wheel.sh scripts/run_tests.sh scripts/test_installation.sh
CMAKE_BUILD_PARALLEL_LEVEL=128 python -m build --wheel
python -m twine check dist/mooncake_transfer_engine-0.3.12.post1-cp310-cp310-linux_x86_64.whl
MOONCAKE_TEST_WHEEL="$PWD/dist/mooncake_transfer_engine-0.3.12.post1-cp310-cp310-linux_x86_64.whl" python -m pytest -q python/tests/unit/test_cli_modules.py python/tests/packaging
git fetch origin main
pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main...HEAD)

Test results:

  • Unit tests pass (focused wrapper, lazy-import, and source execution coverage)
  • Integration tests pass (13 focused/source/packaging tests, including isolated installed-wheel console-script and python -m smoke tests)
  • Manual testing done (root PEP 517 wheel build, wheel manifest inspection, and twine check passed)

No environment-only validation gaps remain. The build emitted one pre-existing GCC warning from RDMA control-frame code outside this Python-only diff.

Checklist

  • I have performed a self-review of my own code
  • I have formatted my code using ./scripts/code_format.sh (no C/C++ source changes; PR-scoped pre-commit passed)
  • I have run pre-commit on the files changed in this PR and all hooks pass
  • I have updated the documentation (not applicable: public imports and command behavior are unchanged)
  • I have added tests to prove my changes are effective
  • For changes >500 LOC: I have filed RFC issue [RFC]: Unify Mooncake's Python package, build, and release architecture #3534 (this scoped diff is below 500 LOC)

AI Assistance Disclosure

  • No AI tools were used
  • AI tools were used (specify below)

OpenAI Codex assisted with repository analysis, the scoped migration, packaging and command-smoke coverage, validation, and PR preparation. The human submitter remains responsible for reviewing and defending every changed line.

@Aionw
Aionw force-pushed the codex/migrate-python-cli branch from 8aae0e0 to b792be8 Compare September 8, 2026 07:55
@Icedcoco

Icedcoco commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the detailed migration work. I have one question about the non-SKBUILD CMake installation path.

In mooncake-integration/CMakeLists.txt, the direct install rules include cli.py, cli_bench.py, and transfer_engine_topology_dump.py from python/mooncake, but I could not find a corresponding entry for cli_client.py. Since mooncake-wheel/mooncake/cli_client.py is also removed by this PR, could the mooncake_client module be missing from non-SKBUILD installations?

The project metadata still maps the entry point to:

mooncake_client = "mooncake.cli_client:main"

Would it make sense to add python/mooncake/cli_client.py to the direct CMake install list and include a small non-SKBUILD installation check? I may be overlooking another installation path, so I would appreciate your thoughts.

- Add python/mooncake/_launcher.py that resolves bundled native binaries via
  importlib.resources, makes them executable on the fly, and raises a clear
  error instead of falling back to PATH (avoids running a mismatched build).
- Slim cli/cli_client/cli_bench to thin pass-throughs of _launcher.locate,
  removing the triplicated _locate_binary helper.
- Restore pkgutil.extend_path on python/mooncake/__init__.py so the unified
  package keeps the legacy namespace behaviour for independently installed
  mooncake subpackages.
- Install python/mooncake/cli_client.py in the non-SKBUILD CMake install list.
- Rework the unit/packaging tests for the launcher and thin wrappers, and
  extend the project tests to guard the scikit-build wiring and the
  build_wheel.sh staging.
@Icedcoco

Icedcoco commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the follow-up changes. The cli_client.py installation path now looks covered.

I noticed one related packaging question around transfer_engine_bench. The root pyproject.toml sets BUILD_BENCHMARK = false, while the project metadata still exposes:

transfer_engine_bench = "mooncake.cli_bench:main"

The CMake installation of the native transfer_engine_bench binary is also conditional on BUILD_BENCHMARK/BUILD_EXAMPLES. As a result, a default scikit-build wheel may contain the console-script wrapper but not the native binary that _launcher.locate() expects. Running transfer_engine_bench from such a wheel would then fail with FileNotFoundError.

I understand this may predate this PR and may be outside the intended migration scope. Would you prefer to address it here by either enabling and testing the benchmark binary in the default wheel, or by excluding/limiting this entry point for wheel variants that do not build benchmarks? Or should we track it separately as a follow-up issue?

@Aionw

Aionw commented Sep 11, 2026

Copy link
Copy Markdown
Collaborator Author

Thanks for checking this. I rechecked the CMake options: transfer_engine_bench is built under BUILD_EXAMPLES, not BUILD_BENCHMARK. mooncake-common/common.cmake defaults BUILD_EXAMPLES to ON, and the root pyproject.toml does not override it. WITH_TE also defaults to ON. The SKBUILD install rule includes the binary in the Python package under WITH_TE AND BUILD_EXAMPLES.

Therefore, BUILD_BENCHMARK = false does not disable this CLI binary in the default scikit-build configuration, so no build-option or entry-point change is needed here. This conclusion is based on the build/install rules; I have not built a fresh scikit-build wheel for this follow-up.

The earlier “known gap” comment I added to the packaging test was incorrect and likely contributed to the confusion—sorry about that. I have removed it locally. A custom build explicitly disabling BUILD_EXAMPLES is a separate case; handling entry points for such reduced variants can be considered separately from this directory migration.

@Icedcoco

Copy link
Copy Markdown
Collaborator

Thanks for the follow-up changes. I rechecked the non-SKBUILD CMake installation path and noticed one small point that may be worth discussing.

The CLI wrappers now import:

from mooncake._launcher import locate

The wrappers themselves are included in the direct CMake install rules, but I could not find a corresponding install entry for python/mooncake/_launcher.py. Since the legacy directory installation no longer contains this file, would a direct CMake installation still include _launcher.py?

If not, commands such as mooncake_master, mooncake_client, or transfer_engine_bench might fail with ModuleNotFoundError: No module named 'mooncake._launcher' in that installation mode. This does not affect the scikit-build wheel path, which stages the launcher separately.

I may be missing another installation path, so I wanted to raise this for discussion and get your view on whether it should be handled as part of this PR or tracked separately.

Include _launcher.py alongside the non-SKBUILD CLI wrappers and remove the incorrect comment about the default benchmark binary configuration.

@Icedcoco Icedcoco left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM pending CI

@ShangmingCai ShangmingCai left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants