Skip to content

[Integration] Migrate allocators to unified Python package - #3628

Open
Aionw wants to merge 3 commits into
mainfrom
codex/migrate-allocator-python
Open

Aionw wants to merge 3 commits into
mainfrom
codex/migrate-allocator-python

Conversation

@Aionw

@Aionw Aionw commented Aug 24, 2026 •

Copy link
Copy Markdown
Collaborator

Description

Implements the allocator subsystem portion of Phase 2 in #3534.

  • move mooncake.allocator, mooncake.allocator_ascend_npu, and mooncake.fabric_allocator_utils into python/mooncake as their sole tracked source;
  • preserve the existing flat public module imports and allocator behavior;
  • update direct CMake installation, root scikit-build packaging exclusions, and the transitional scripts/build_wheel.sh staging path;
  • add hardware-free CUDA/NVLink and Ascend/UBSHMEM probing coverage using synthetic runtime modules;
  • add source-location, build-input, wheel-manifest, and isolated installed-wheel checks;
  • update ownership and labeling for the migrated implementation and tests.

Overlap check: Since this PR was opened, #3586, #3624, #3626, #3630, #3632, and #3635 have landed on main; this branch includes them through its main syncs, most recently the 2026-09-28 merge of main (d9152a4). #3585 (Reshard) is still open and remains out of scope. #2404 (now merged) and #3396 (still open) change unsupported allocator protocol handling in the native Transfer Engine binding; they do not move or modify these Python allocator implementations. #2305 (open) adds a logging hook to the legacy mooncake-integration/allocator.py; if it lands, that change must move to python/mooncake/allocator.py.

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?

Original submission (before merging main on 2026-09-28; superseded by "After merging main" below):

Test commands:

python -m pytest -q python/tests/allocators python/tests/packaging/test_project.py
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/allocators python/tests/packaging
bash -n scripts/build_wheel.sh
git fetch origin main
pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main...HEAD)
git diff --check origin/main...HEAD

Test results:

  • Unit tests pass (21 focused allocator and source-packaging tests; 22 tests with the installed-wheel check enabled)
  • Integration tests pass (root PEP 517 wheel build, twine check, and isolated --no-deps wheel install/import outside the repository)
  • Manual testing done (wheel manifest contains all three canonical allocator modules; core imports did not load torch or torch_npu; installed hardware modules imported with synthetic runtimes)

After merging main (d9152a4, 2026-09-28):

PYTHONPATH=python:mooncake-wheel python -m pytest -q python/tests/allocators python/tests/packaging
pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main...HEAD)
  • 32 passed, 2 skipped; the skips are the installed-wheel smoke tests, which need MOONCAKE_TEST_WHEEL / MOONCAKE_TEST_LEGACY_WHEEL.
  • The merge's conflicts in .github/CODEOWNERS, python/tests/packaging/test_installed_wheel.py, python/tests/packaging/test_project.py were resolved by keeping both sides.
  • CI's changed-lines Python check (ruff check plus scripts/ci/check_changed_python_format.py, ruff 0.6.9) passes.
  • PR-scoped pre-commit passes every hook except cmake-format, which only rewrites the tent_metrics block in mooncake-integration/CMakeLists.txt inherited from main; main at d9152a4 fails the same way, and this PR's own CMake lines pass.
  • The root PEP 517 wheel build, twine check, and installed-wheel smoke runs from the original submission were not repeated after the merge. CI does not run them either (it builds only the legacy scripts/build_wheel.sh wheel and runs scripts/test_installation.sh), so they still need a rerun.

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 formatting passed)
  • I have run pre-commit on the files changed in this PR and all hooks pass (after merging main, except the pre-existing cmake-format issue noted under "After merging main")
  • I have updated the documentation (not applicable: public imports and behavior are unchanged, and no allocator source paths are documented)
  • 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

AI Assistance Disclosure

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

OpenAI Codex assisted with repository and overlap analysis, the module migration, packaging/build updates, test authoring, validation, and PR preparation. The human submitter remains responsible for reviewing every changed line and understanding and defending the change end to end.

Claude Code (Anthropic) merged main into this branch on 2026-09-28, resolved the conflicts, re-ran the focused tests listed under "After merging main", and updated this description.

@github-actions github-actions Bot added the Tests label Sep 17, 2026
Conflict resolution (all same-spot appends; both sides kept):
- .github/CODEOWNERS: kept the allocator owner entries and main's
  http_metadata_server/async_store/buffer_pool entries, allocator lines
  first after the ha/ entry, matching how siblings were inserted.
- python/tests/packaging/test_installed_wheel.py: kept the
  fabric_allocator_utils import alongside main's http_metadata_server and
  mooncake_config imports (alphabetical), and kept main's simple
  KVBootstrapServer/MooncakeConfig asserts followed by the allocator
  wheel-location and fake-torch import checks.
- python/tests/packaging/test_project.py: kept ALLOCATOR_MODULES next to
  main's CLI_MODULES/CLI_ENTRY_POINTS, and kept the allocator source/build
  input tests ahead of main's http_metadata_server and async_store tests.

scripts/build_wheel.sh, mooncake-integration/CMakeLists.txt,
python/CMakeLists.txt and .github/labeler.yml merged cleanly.

This branch has not been deployed

No deployments
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.

1 participant