kvcache-ai / kvcache-ai/Mooncake

[RFC]: Unify Mooncake's Python package, build, and release architecture

Open
#3,534 3 comments 1 reaction 0 assignees View on GitHub
RFC
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Changes proposed

## Summary

Establish `python/mooncake` as the single source of truth for all tracked Python code distributed as the `mooncake` import package.

Use `scikit-build-core` as the single PEP 517 build backend. It will combine tracked Python sources with native artifacts supplied through explicit CMake install rules, without copying generated files into the source tree.

Use `cibuildwheel` as the standard release-wheel execution layer for supported manylinux variants. Hardware-specific builders may retain specialized runners and toolchains, but must consume the same PEP 517 backend and produce the same validated artifact contract.

This RFC also defines:

- subsystem-oriented Python module boundaries;
- private native extension modules behind stable public Python facades;
- explicit core, optional, development, and hardware dependency groups;
- isolated installed-wheel validation;
- a checked invariant between Git tags, package metadata, and wheel metadata;
- centralized publication to PyPI and GitHub Releases;
- independently retryable hardware-specific release variants.

Distribution names such as `mooncake-transfer-engine` and its hardware variants remain unchanged. The installed import package remains `mooncake`.

Background and a detailed walkthrough are available in [Mooncake Python Packaging Refactor](https://aionw.github.io/posts/mooncake-python-package-refactoring/).

## Relationship to existing RFCs

This proposal overlaps with two existing RFCs:

- #3425, **Refactor the `mooncake-wheel` package layout**, identifies the same fragmented Python source layout and build-time copying problems. Its target architecture keeps `mooncake-wheel/mooncake` as the source root and retains public native module names. This proposal goes further by introducing `python/mooncake` as a source-only package root, private native extensions, `scikit-build-core`, explicit dependency groups, and installed-wheel validation. If accepted, this RFC should supersede the narrower target architecture in #3425.

- #3130, **Release Branch Workflow for Stable Releases**, overlaps with the release-validation part of this proposal. It focuses on release branches and TestPyPI validation, while this RFC defines the package build contract, artifact identity, wheel validation, and centralized publication. The two proposals are complementary: #3130 may define how a release commit is selected and stabilized, while this RFC defines how that commit becomes validated and published Python distributions.

## Motivation

### Python has no single source root

Tracked Python code currently lives under several component directories, including:

- `mooncake-wheel/mooncake`;
- `mooncake-integration`;
- `mooncake-reshard/python/mooncake`;
- EP and PG component directories.

Whether a Python file enters the wheel can depend on imperative copy logic in `scripts/build_wheel.sh`. Development and tests may consequently import a different file from the one installed from a wheel.

### The source package is also used as a staging directory

Wheel construction copies shared libraries, executables, metadata, and Python files into the tracked package directory. Interrupted or consecutive variant builds can leave stale or variant-specific artifacts in the source tree.

A source directory should not also be a mutable wheel assembly workspace.

### Module boundaries no longer match subsystem boundaries

The package root currently contains Store utilities, structured object storage, EP helpers, services, SSD administration, command-line entry points, and framework integrations.

Large modules such as `structured_object_store.py` combine models, manifests, serialization, buffer management, transport behavior, and optional codecs. The flat layout obscures ownership and allows optional dependencies to leak into core imports.

### Dependencies are implicit

Different modules use NumPy, PyTorch, Pillow, msgspec, ZeroMQ, FastAPI, httpx, Paramiko, vLLM, and hardware runtimes.

Making all of them mandatory would produce an unnecessarily large core installation. Leaving them undeclared produces incomplete environments and late `ModuleNotFoundError` failures.

### Release artifacts need one checked identity

Actions artifacts, PyPI distributions, and GitHub Release assets serve different purposes, but they must refer to the same immutable source and package version.

The build and publication workflows should check this explicitly rather than derive similar version strings independently.

## Design principles

1. All tracked Python package code lives under `python/mooncake`.
2. PEP 517 is the only public Python build interface.
3. Generated files never enter tracked Python source directories.
4. Public API modules are stable Python facades; native filenames are private implementation details.
5. Core imports do not load unrelated integrations or hardware runtimes.
6. Release gates test an installed wheel outside the source tree.
7. Git tags, `pyproject.toml`, and wheel metadata must agree.
8. Builders build and validate; only a centralized publisher may publish.
9. Obsolete source locations, namespace extensions, staging scripts, and internal import paths are removed rather than retained as compatibility layers.

## Proposed repository layout

```text
Mooncake/
pyproject.toml

python/
pyrightconfig.json

mooncake/
__init__.py

engine/
__init__.py
shared_segment.py

store/
__init__.py
async_client.py
buffer_pool.py
config.py
structured/
__init__.py
api.py
models.py
manifest.py
transport.py
buffer.py
codecs/

ep/
__init__.py
buffer.py
elastic_buffer.py

pg/
__init__.py

reshard/
__init__.py
contracts/
weight/

allocators/
integrations/
dataproto/
vllm/
services/
ssd/
cli/

tests/
unit/
integration/
e2e/
hardware/
packaging/
typecheck/

mooncake-integration/ # C++ and binding implementation only
mooncake-ep/ # native component implementation
mooncake-pg/ # native component implementation
mooncake-store/ # native component implementation
mooncake-transfer-engine/ # native component implementation

build/ # ignored, backend-owned output
```

Native component directories continue to own their native implementations. Only tracked Python code belonging to the installed import package moves under `python/mooncake`.

## Public API and native extensions

Native extensions become private implementation modules:

```text
mooncake._engine
mooncake._store
mooncake._ep
mooncake._pg_
```

Stable public APIs are ordinary Python packages:

```python
# python/mooncake/engine/__init__.py
from mooncake._engine import TransferEngine, TransferOpcode

__all__ = ["TransferEngine", "TransferOpcode"]
```

Existing documented integration imports remain valid through these permanent facades:

```python
from mooncake.engine import TransferEngine
from mooncake.store import MooncakeDistributedStore
from mooncake.pg import set_transfer_engine
```

The facade is the canonical API, not a temporary compatibility layer. It owns exports, typing, documentation, capability checks, and native implementation selection.

Old source-only or internal import paths are removed when the new layout lands. If a documented public path must genuinely change, it should be handled as an intentional breaking change with coordinated downstream updates, not with a second implementation, wildcard re-export, namespace extension, or fallback.

## Build contract

The root `pyproject.toml` declares `scikit-build-core` as the only PEP 517 backend:

```toml
[build-system]
requires = ["scikit-build-core", "pybind11"]
build-backend = "scikit_build_core.build"

[tool.scikit-build]
wheel.packages = ["python/mooncake"]
build-dir = "build/python/{wheel_tag}"
```

Native components expose explicit CMake install rules:

```cmake
pybind11_add_module(_engine ...)
install(TARGETS _engine LIBRARY DESTINATION mooncake)
```

Python development and builds use standard commands:

```bash
pip install -e '.[dev]'
python -m build
```

`scikit-build-core` owns the temporary configure, build, install, and wheel trees. Python developers do not need to manually select a CMake install prefix or copy artifacts into the package.

`scripts/build_wheel.sh` and source-tree copy/restore behavior are removed after the backend provides equivalent end-to-end functionality.

## Dependency groups

The exact list requires an import audit, but dependencies should be separated into explicit groups:

| Group | Examples |
|---|---|
| Core | aiohttp, requests, msgpack |
| Structured objects | NumPy, Pillow |
| vLLM integration | msgspec, pyzmq, FastAPI, httpx, vLLM |
| Administration | Paramiko |
| Development | pytest, ruff, pyright, build, twine |
| Hardware | Torch and CUDA/ROCm/CANN/MUSA-specific requirements |

Importing a core subsystem must not fail because an unrelated integration or hardware runtime is absent.

## Wheel validation

Release gates must validate the installed distribution rather than only the source tree:

- `twine check`;
- distribution name, version, ABI, and platform tag;
- wheel file manifest;
- clean-environment `pip install --no-deps`;
- public `mooncake.engine`, `store`, `ep`, and `pg` imports;
- CLI entry points;
- private native extension loading;
- ELF dependencies and RPATH;
- declared optional-dependency boundaries;
- confirmation that the repository source tree is not on `PYTHONPATH`.

## Version and publication contract

`pyproject.toml` is the package-version source of truth. A Git tag selects the immutable release commit. They must match under PEP 440 normalization:

```text
0.3.13 <-> v0.3.13
0.3.13.post1 <-> v0.3.13.post1
0.3.14rc1 <-> v0.3.14rc1
```

Builders upload validated Actions artifacts but hold no PyPI publication authority.

A single publication job:

1. collects all required core artifacts;
2. verifies filenames, versions, tags, and digests;
3. publishes through PyPI Trusted Publishing;
4. attaches the same files to the GitHub Release.

PyPI files are immutable. Retries must compare existing digests instead of using `skip-existing` to hide conflicts.

Default CUDA, CUDA 13, and non-CUDA variants form the core release gate. Specialized variants may retry or backfill from an existing immutable tag without moving or recreating that tag.

## Implementation phases

### Phase 1: Establish the new Python project foundation

- [x] Finished in https://github.com/kvcache-ai/Mooncake/pull/3577

Create the target project structure before moving subsystem implementations:

```text
pyproject.toml
python/
mooncake/
__init__.py
tests/
packaging/
```

This phase defines:

- the root `pyproject.toml`;
- `scikit-build-core` as the PEP 517 backend;
- the `python/mooncake` package root;
- core, optional, development, and hardware dependency groups;
- editable installation through `pip install -e .`;
- wheel construction through `python -m build`;
- the initial CMake install contract for native artifacts;
- minimal packaging and installed-wheel smoke tests.

The new build must work end to end before large-scale module movement begins.

During the transition, the backend may explicitly consume unmigrated modules from their current source locations. It must not create a second tracked copy of those modules under `python/mooncake`, and it must not use the tracked source tree as a staging directory.

Acceptance criteria for this phase:

- a complete wheel can be produced through the new PEP 517 entry point;
- editable installation works through the same backend;
- dependencies are declared in the new `pyproject.toml`;
- generated native artifacts remain outside tracked source directories;
- the installed-wheel smoke test does not import from the repository through `PYTHONPATH`.

### Phase 2: Migrate Python modules by subsystem

Move modules into `python/mooncake` incrementally. Each subsystem migration is an independently reviewable change and includes:

- moving the implementation and tests;
- updating internal imports;
- declaring its core or optional dependencies;
- updating documentation and type-check configuration;
- validating source and installed-wheel imports;
- removing the subsystem's old source location and temporary build mapping in the same change.

A suggested migration order is:

1. lightweight pure-Python utilities and shared configuration;
2. allocators and shared-segment helpers;
3. `mooncake.reshard`;
4. structured object storage and its codecs;
5. framework integrations such as DataProto and vLLM;
6. services, CLI commands, and SSD administration;
7. Store and Transfer Engine public facades;
8. EP and PG Python APIs and native loaders.

Pure-Python and leaf modules move first because they exercise the new package, dependency, test, and editable-install structure without immediately combining the migration with the most hardware-sensitive native extensions.

At every point:

- each module has exactly one authoritative source location;
- the new backend remains capable of building a complete package;
- migrated and unmigrated modules are covered by the same installed-wheel tests;
- no compatibility copy or fallback implementation is introduced.

### Phase 3: Complete the native-module boundary

After the surrounding Python modules have moved:

- rename native extensions to private implementation names such as `_engine`, `_store`, `_ep`, and `_pg_`;
- expose their APIs through permanent Python facades;
- move capability detection and Torch ABI selection into those facades;
- validate extension loading, ELF dependencies, RPATH, and auditwheel output;
- remove build rules that install native extensions under public module names.

This phase changes implementation ownership without requiring users to import native filenames directly.

### Phase 4: Cut release workflows over to the new backend

Once all modules build from the new project structure:

- make release builders invoke the PEP 517 backend;
- adopt `cibuildwheel` for supported manylinux variants;
- require specialized hardware builders to implement the same artifact and validation contract;
- run isolated installed-wheel tests for every released variant;
- enforce Git tag, `pyproject.toml`, and wheel-version identity;
- centralize PyPI and GitHub Release publication;
- preserve independent retry and backfill for specialized variants.

Release workflows should not switch to the new backend until the migrated package has reached feature and import parity.

### Phase 5: Remove the legacy package and build paths

After the release cutover:

- remove `mooncake-wheel` as a source and staging directory;
- remove temporary mappings for legacy module locations;
- remove `scripts/build_wheel.sh` and copy/restore logic;
- remove `pkgutil.extend_path`;
- remove obsolete CMake install paths and variant-specific metadata mutation;
- update contributor documentation so the supported Python workflows are only `pip install -e .`, `python -m build`, and the relevant test commands.

The repository is considered migrated only when `python/mooncake` is the sole tracked source of the installed package and no legacy assembly path remains.

## Non-goals

- Moving native C++/CUDA implementations under `python/`.
- Replacing CMake or auditwheel.
- Making every optional integration a core dependency.
- Allowing multiple hardware wheel variants to coexist in one environment; they continue to provide the same `mooncake` import package.
- Hiding real hardware differences behind one universal builder.
- Redesigning Store, Transfer Engine, EP, PG, or reshard behavior.
- Keeping old source directories, namespace-package extensions, or legacy wheel entry points after the new architecture works end to end.

## Acceptance criteria

The refactor is complete when:

- `python/mooncake` is the only tracked source of the installed package;
- a clean checkout can build through `python -m build`;
- editable installs and wheels use the same backend;
- no build modifies tracked source files;
- all required variants produce validated wheels through the common contract;
- installed-wheel tests run without the source tree on `PYTHONPATH`;
- tag, package metadata, artifact metadata, and GitHub Release agree;
- obsolete staging scripts and source paths are removed.

### Before submitting a new issue...

- [x] Searched existing issues and pull requests.
- [x] Reviewed the overlapping RFCs #3425 and #3130.
- [x] Read the project documentation and contribution guidelines.

Contributor guide

Open the contributing guide

Research direction

Start with the root pyproject.toml, python/mooncake layout, and the packaging tests described under python/tests/packaging; run python -m build to inspect the current PEP 517 path. Read the CMake install rules and release-validation requirements, then compare them with the completed Phase 1 work in pull request #3577. Done means the proposed source, build, wheel-validation, and publication contracts are implemented across the stated phases.

Written by the indexing model from the issue text.

Assessment

Tech stack
cmake, python
Domain
build-system, ci-cd, release
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.