[Weight] Add runtime weight manifest contracts - #3187
Conversation
zxpdemonio
left a comment
There was a problem hiding this comment.
Review — [Weight] Add runtime weight manifest contracts
I reviewed the contract module, docs, wheel staging, and tests against RFC #3111. Locally: 139 passed with PYTHONPATH=mooncake-model-weight/python.
Verdict
Solid first-phase contract: placement / binding split is clear, validation is fail-closed in the right places, and the PR stays within the stated non-goals. CI is currently red (spell check), and there are a few contract/docs nits worth fixing before merge. No deep semantic blockers beyond that if the deferred planner checks are intentional.
Blocker / merge gate
- Spell Check fails on
nd_*identifiers
Typos treatsndasandintest_model_weight_manifest.py(nd_descriptor,nd_placement,nd_runtime). CI Gate is red because of this.
Fix by renaming (n_dim_*/multi_dim_*) or adding an allowlist entry in.typos.toml.
Contract review (what looks good)
- Clear split:
PlacementManifest(serializable, address-free) vsRuntimeBindingManifest(ephemeral, generation/lease-fenced). - Canonicalization of
partition_dim/shard_dims, stableplacement_id+ digest, and bind-time digest checks. - Geometry checks: bounds, non-shard axis integrity, nbytes =
prod(local_shape) * itemsize. - Same-rank N-D box overlap rejection (sweep + AABB) and runtime address-range overlap with explicit alias opt-in.
- Contiguity / canonical stride /
address_semantics="view"rules are documented and tested. - u64 address-range overflow checks; strict JSON schema for placement round-trip.
- Wheel staging + EXIT cleanup;
pkgutil.extend_pathenablesmooncake.model_weightas a split package. - Scope matches RFC phase-1 (no planner / TE / Store).
Issues / questions
-
README local test recipe is wrong (
mooncake-model-weight/README.md)
It suggests:PYTHONPATH=mooncake-wheel python -m pytest ...
That fails with
No module named 'mooncake.model_weight'unless the package has been staged into the wheel. The working recipes are:PYTHONPATH=mooncake-model-weight/python python -m pytest -q mooncake-model-weight/tests # or, with extend_path: PYTHONPATH=mooncake-wheel:mooncake-model-weight/python ... -
Cross-rank logical overlap is intentionally unchecked
Sametensor_idwith overlapping boxes on differentParallelRankvalues is accepted (verified). That is fine if coverage/conflict belongs to the planner, but please call it out explicitly indocs/source/design/model-weight-manifest.mdso adapters do not assume the contract enforces a global partition. -
Incomplete shard coverage is also accepted
A TP shard that only covers part ofglobal_shapevalidates today. Same ask: document as planner-owned, so this is not mistaken for a full placement invariant. -
Strict
type(x) is intwill reject NumPy integers
Framework inventories often carrynumpy.int64extents. TorchSizeis fine (Pythonint), but raw NumPy metadata will raise. Consideroperator.index()/numbers.Integral(still rejectingbool), or document that adapters must coerce. -
addressrejects0(minimum=1)
Reasonable as a null sentinel for GPU views; worth one sentence in the design doc so CPU/testing harnesses do not treat it as accidental. -
Empty placement is allowed
PlacementManifest(model_id=..., revision=..., tensors=(), fragments=())succeeds and gets a digest. Harmless, but if empty manifests should be invalid at the control-plane boundary, reject them here or document “empty is a valid sentinel”. -
Generated
placement_fragment_idis content-hashed, but explicit IDs become part ofplacement_id
Adapters that sometimes omit and sometimes supply IDs for the same geometry will see identity churn. Document “IDs must be stable across restarts” for framework exporters.
Non-blocking nits
RuntimeManifestcan exist withoutlease_id; only projection to binding requires it — good, already documented.build_wheel.shinstalls a second EXIT trap that replaces the first; safe today because the later trap also cleans staging, but a single composed cleanup would be clearer.- No public schema version on serialized placement beyond the internal
"schema": "weight-placement"used forplacement_id. Fine for v0; consider surfacing it in JSON before Store persistence lands. - PR checklist marks docs updated (good); spell-check + README fix should land with this PR.
Test coverage
Strong for construction, projection, bind, aliases, overflow, and type strictness. Missing (optional follow-ups):
- README/
PYTHONPATHsmoke as documented - explicit “cross-rank overlap allowed” / “incomplete coverage allowed” tests if those are permanent contract guarantees
- installed-wheel
import mooncake.model_weightassertion intest_installation.sh(tests help, but a one-liner import check is cheaper signal)
Summary
Approve directionally after fixing the typos CI failure and correcting the README import path. The remaining items are mostly “make deferred invariants explicit” so the next planner PR does not rediscover them.
b849312 to
c71dc62
Compare
c71dc62 to
6a19b22
Compare
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
The corresponding test file is also approximately 1,700 lines. Since future planner, Store, and framework-integration work will likely extend this module, it may be worth splitting it along these existing responsibility boundaries to keep the implementation easier to review and maintain. |
Thanks, this was a very helpful suggestion. I have split both the implementation and tests along the proposed responsibility boundaries, while keeping the public import surface unchanged. |
Description
Related RFC: #3111
This PR introduces the framework-neutral runtime weight manifest contract for
Mooncake model-weight reuse and heterogeneous resharding.
It is the first implementation phase of the RFC. This PR defines how frameworks
describe the logical placement and live physical location of model weights,
without coupling Mooncake Core to SGLang, vLLM, PyTorch, model-specific naming
rules, or a particular transfer backend.
What This PR Adds
1. Top-level Model Weight module
Add
mooncake-model-weightas a top-level Mooncake module.The public Python API is exposed through:
The tracked source of truth remains in:
During release-wheel construction, this package is temporarily staged into the
combined Mooncake wheel and removed after the build.
2. Logical placement contract
Add
PlacementManifestto describe reusable and serializable model-weightplacement:
global_offsetandlocal_shape;partition_dimandshard_dimssemantics;Logical placement does not contain GPU addresses, endpoints, leases, or runtime
allocation owners.
3. Runtime binding contract
Add
RuntimeBindingManifestto describe one live runtime generation:A runtime restart can reuse the same logical placement while publishing a new
runtime binding.
4. Validated runtime snapshot
Add
RuntimeManifestand framework-neutral inventory adapters for exportingruntime tensor facts.
The contract validates:
Mooncake does not infer layer, expert, partition, or layout semantics from
parameter names. These facts must be supplied by the framework adapter.
5. Projection and binding APIs
Add APIs to:
Why Split Placement and Binding?
Logical placement is stable and reusable, while GPU addresses and allocations
belong to one process lifetime.
Separating them allows future reshard planners and transfer executors to reuse
the same logical placement while still requiring a fresh, generation-fenced
runtime binding before accessing memory.
Wheel Integration
The main Mooncake wheel now stages
mooncake.model_weightfrom the top-levelmodule during
scripts/build_wheel.sh.The installation test runs the manifest contract tests against the installed
wheel. The staging directory is cleaned after both successful and failed
builds.
Non-Goals
This PR intentionally does not:
N-D reshard planning, TE execution, Store integration, and framework adapters
will be submitted as follow-up changes built on this contract.
Module
mooncake-transfer-engine)mooncake-store)mooncake-ep)mooncake-pg)mooncake-integration)mooncake-p2p-store)mooncake-wheel)mooncake-common)mooncake-rl)mooncake-model-weight)This PR contains the minimal release-wheel staging required to publish the new
module, but feature ownership remains under
mooncake-model-weight.Type of Change
How Has This Been Tested?
Test commands
Test results
The tests cover:
Real framework integration, planner, Transfer Engine, Store, and SGLang serving
E2E coverage will be added in follow-up PRs.
Commit Structure
The change is organized into reviewable commits:
mooncake.model_weightAPI.Checklist
./scripts/code_format.sh(not applicable; no C/C++ code changed)pre-commit run --all-filesAI Assistance Disclosure
AI assistance was used to inspect the existing Mooncake packaging and manifest
boundaries, implement the contract and tests, organize the commit history, and
run code review and validation. The submitter has reviewed the changes and is
responsible for understanding and maintaining the code.