kvcache-ai / kvcache-ai/Mooncake

[RFC]: Refactor the `mooncake-wheel` package layout

Open
#3,425 1 comment 2 reactions 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

Make `mooncake-wheel` the source of truth for every Python module that ships
in the Mooncake wheel, and group that Python by subsystem instead of leaving
it as a flat directory that also receives build-time copies of native
artifacts.

## Motivation

Today the installed Python package is one wheel, but the source layout does
not match that fact.

As a general packaging principle, Python source that ships as part of a
single wheel should live, as much as practical, under one complete and
self-contained source tree. A developer inspecting that tree should be able
to understand what Python code belongs to the package without first reading
the wheel build script or reconstructing which files are copied from other
top-level directories.

This is particularly important for an open-source project. The repository
layout is part of the project's interface to contributors: it affects code
discovery, local development, testing, static analysis, IDE behavior, and
the cost of understanding where a change should be made. When the wheel is
assembled from Python sources scattered across several unrelated directory
trees, the repository no longer represents the package that users actually
install. Build-time copying also creates two notions of package structure:
the source layout seen by contributors and the assembled layout seen by
users.

Mooncake currently has exactly this mismatch:

* `mooncake-reshard/` is a top-level tree whose only distribution path is the
combined wheel. `scripts/build_wheel.sh` copies
`mooncake-reshard/python/mooncake/reshard` into
`mooncake-wheel/mooncake/reshard` and deletes the copy on exit. Developers
need `PYTHONPATH=mooncake-wheel:mooncake-reshard/python` to test it.
* `mooncake-wheel/mooncake/` is a flat mix of unrelated Python: Store
helpers, Transfer Engine CLI, vLLM connectors, SSD/SPDK scripts, RL
DataProto transfer (`structured_object_store.py`, ~6.5k lines), and EP
loaders. Native `.so` files and binaries are copied into the same
directory at build time.
* Some Python files are tracked in `mooncake-integration/` and copied into
the wheel (`async_store.py`, `allocator.py`,
`allocator_ascend_npu.py`, `fabric_allocator_utils.py`). The wheel is the
install surface, but not the source of truth.
* The wheel already uses `pkgutil.extend_path` so that an out-of-tree
`mooncake.reshard` can join the `mooncake` namespace. That split exists
only because reshard source does not live in the wheel.

The result is extra staging, multiple places to look for Python, and no
obvious home for the next Python-only subsystem. It also makes the package
boundary less explicit: whether a Python module is part of the Mooncake
wheel is determined partly by repository location and partly by imperative
copy logic in the build script.

The preferred model is that `mooncake-wheel/mooncake/` describes the Python
package directly. Build-time assembly should be reserved for artifacts that
are genuinely produced by the build, such as native libraries and binaries,
rather than for relocating tracked Python source code.

## Goals

- One source tree for Python that is installed as `import mooncake`.
- Keep every existing public import working.
- Give `mooncake.reshard` a real subpackage inside the wheel, with its tests
and typechecks next to it.
- Stop copying Python into the wheel at pack time.
- Leave native artifacts (`.so`, `mooncake_master`, and friends) at the
`mooncake/` package root, because the wheel sets ELF `RPATH=$ORIGIN`.
- Make later grouping of RL / CLI / SSD / integration Python possible
without another packaging redesign.

## Non-goals

- A separate PyPI package for reshard, RL helpers, or CLI.
- Changing reshard manifest contracts, planner behavior, Store, or Transfer
Engine.
- Moving C++/CUDA sources into the wheel.

## Current state

Tracked Python in `mooncake-wheel/mooncake/` includes, among others:

| Area | Files |
|------|--------|
| Native loaders | `ep.py`, `pg.py`, `buffer_pool.py` |
| RL / structured objects | `structured_object_store.py`, `dataproto_catalog.py`, `_fast_copy.c` |
| CLI / services | `cli.py`, `cli_client.py`, `cli_bench.py`, `http_metadata_server.py`, `mooncake_store_service.py`, `transfer_engine_topology_dump.py` |
| SSD | `mooncake_ssd_register.py`, `mooncake_ssd_unregister.py`, `spdk_tgt_create.py` |
| Serving glue | `mooncake_connector_v1.py`, `vllm_v1_proxy_server.py` |
| Config / EP Python | `mooncake_config.py`, `mooncake_ep_buffer.py`, `mooncake_elastic_buffer.py` |

Copied at wheel build time, not tracked under `mooncake-wheel/`:

- `mooncake-reshard/python/mooncake/reshard/` → `mooncake/reshard/`
- `mooncake-integration/store/async_store.py`
- `mooncake-integration/allocator.py`, `allocator_ascend_npu.py`,
`fabric_allocator_utils.py`
- CMake outputs: `engine.so`, `store.so`, `lib*.so`, `mooncake_master`,
`mooncake_client`, `transfer_engine_bench`, optional CUDA/NPU/EP `.so`

Public imports that must not break:

```python
from mooncake.engine import TransferEngine
from mooncake.store import MooncakeDistributedStore
from mooncake.reshard.weight import WeightPlacementManifest
from mooncake.structured_object_store import MooncakeBundleTransfer
from mooncake.mooncake_connector_v1 import ... # vLLM kv_connector_module_path
from mooncake import BufferPool
```

Console scripts in `pyproject.toml` (`mooncake_master`, `mooncake_client`,
`mc_store_rest_server`, and others) are also public.

## Constraints

1. **Frozen binding names.** `mooncake.store`, `mooncake.engine`,
`mooncake.ep`, and `mooncake.pg` are extension modules. They stay at the
package root.
2. **ELF RPATH.** `build_wheel.sh` sets `RPATH=$ORIGIN` on ELF files under
`mooncake/`. Shared libraries and binaries stay in that directory. This
RFC does not invent a `_native/` subdirectory that would require RPATH
changes.
3. **Strip / patchelf / CUDA variants do not rewrite `.py` files.** Variant
builds only change wheel metadata and which `.so` files are included.
Pure Python is identical in every variant. Layout is not blocked by
those steps.
4. **Compatibility shims.** A shim is a tiny module at the old import path
that re-exports the new location (`from mooncake.rl.structured_object_store import *`).
Shims exist so callers and docs do not have to move in the same PR as
the files.

## Proposed layout

```
mooncake-wheel/
pyproject.toml
setup.py
pyrightconfig.json
mooncake/
__init__.py
ep.py
pg.py
buffer_pool.py
# build artifacts, gitignored: engine.so, store.so, lib*.so, binaries

reshard/
contracts/
weight/

rl/
structured_object_store.py
dataproto_catalog.py
_fast_copy.c

cli/
services/
ssd/
integration/vllm/
store_py/

# shims at the original module names until callers migrate

tests/
engine/
store/
reshard/
rl/
cli/
integration/
typecheck/negative/
```

`mooncake-wheel` becomes both the Python source root and the install
assembly directory. `scripts/build_wheel.sh` only copies CMake artifacts
into `mooncake/` and applies variant metadata.

`mooncake-reshard/` as a top-level tree goes away after phase 1.
`mooncake-rl/examples/` can stay as examples; RL library code that ships in
the wheel lives under `mooncake/rl/`.

Namespace-package `extend_path` can be removed once reshard source lives
inside the wheel. Keep it only if an external package still needs to extend
`mooncake`.

## Compatibility

Moving Python source into subsystem subpackages must not require downstream
users to update their imports in the same release.

For modules whose public import path changes, the old path remains as a thin
compatibility shim that re-exports the implementation from its new home

Shims contain no business logic. After a documented deprecation window they
can be removed in a follow-up.

### Before submitting a new issue...

- [ ] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Contributor guide

Open the contributing guide

Research direction

Start with scripts/build_wheel.sh, mooncake-wheel/pyproject.toml, setup.py, and pyrightconfig.json, then inspect the listed Python sources under mooncake-reshard, mooncake-integration, and mooncake-wheel. Trace the wheel build and existing tests, typechecks, and public imports before proposing a phased layout. Done means Python source has one wheel-owned tree, compatibility imports remain working, and build-time copying is limited to native artifacts.

Written by the indexing model from the issue text.

Assessment

Tech stack
cmake, python, shell
Domain
backend, build-system, developer-experience
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.