Add file-level model weight cache control plane - #2829
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an experimental file-level model weight cache control plane for Mooncake Store, adding a command-line interface, a model cache client, keyspace helpers, and comprehensive unit tests. The implementation is solid, but several improvements should be made to enhance robustness, correctness, and performance. Specifically, the store setup result check should safely handle None returns, and the size parser should support TB and PB units. Additionally, file materialization should write to a temporary file and atomically rename it to prevent corruption, index updates should be made atomic to avoid race conditions during concurrent imports, redundant store existence checks should be removed to optimize round-trips, and chunk cleanup during import failures should be wrapped in try-except blocks to ensure all orphaned chunks are removed.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| except TypeError: | ||
| result = store.setup(config) | ||
|
|
||
| if result != 0: |
There was a problem hiding this comment.
The setup method of MooncakeDistributedStore might return None on success in some environments or wrappers. Checking result != 0 could lead to a false-positive RuntimeError if None is returned. It is safer and more consistent with the rest of the codebase to check result not in (0, None).
| if result != 0: | |
| if result not in (0, None): |
| units = { | ||
| "KB": 1024, | ||
| "MB": 1024**2, | ||
| "GB": 1024**3, | ||
| } |
There was a problem hiding this comment.
Adding support for Terabytes (TB) and Petabytes (PB) to the size parser makes it more robust and future-proof, especially when dealing with extremely large models or storage segments.
units = {
"KB": 1024,
"MB": 1024**2,
"GB": 1024**3,
"TB": 1024**4,
"PB": 1024**5,
}| WEIGHT_SUFFIXES = { | ||
| ".bin", | ||
| ".gguf", | ||
| ".pt", | ||
| ".pth", | ||
| ".safetensors", | ||
| } |
There was a problem hiding this comment.
Consider adding .ckpt to the list of recognized weight suffixes. It is a very common format used by many models (e.g., Stable Diffusion checkpoints).
| WEIGHT_SUFFIXES = { | |
| ".bin", | |
| ".gguf", | |
| ".pt", | |
| ".pth", | |
| ".safetensors", | |
| } | |
| WEIGHT_SUFFIXES = { | |
| ".bin", | |
| ".ckpt", | |
| ".gguf", | |
| ".pt", | |
| ".pth", | |
| ".safetensors", | |
| } |
| with output.open("wb") as output_file: | ||
| for chunk_key in self._record_chunks(record): | ||
| chunk = self._get(chunk_key) | ||
| output_file.write(chunk) | ||
| digest.update(chunk) | ||
| size += len(chunk) | ||
| if size != record.size: | ||
| raise ValueError( | ||
| f"file size mismatch for {record.path}: " | ||
| f"expected {record.size}, got {size}" | ||
| ) | ||
| if digest.hexdigest() != record.sha256: | ||
| raise ValueError(f"sha256 mismatch for {record.path}") |
There was a problem hiding this comment.
Writing directly to the target output_path can leave a partially written or corrupted file if the operation is interrupted or fails halfway through. A more robust pattern is to write to a temporary file in the same directory and atomically rename it to the target path once the size and checksum verification succeeds.
| with output.open("wb") as output_file: | |
| for chunk_key in self._record_chunks(record): | |
| chunk = self._get(chunk_key) | |
| output_file.write(chunk) | |
| digest.update(chunk) | |
| size += len(chunk) | |
| if size != record.size: | |
| raise ValueError( | |
| f"file size mismatch for {record.path}: " | |
| f"expected {record.size}, got {size}" | |
| ) | |
| if digest.hexdigest() != record.sha256: | |
| raise ValueError(f"sha256 mismatch for {record.path}") | |
| temp_output = output.with_suffix(f".{time.time_ns()}.tmp") | |
| try: | |
| with temp_output.open("wb") as output_file: | |
| for chunk_key in self._record_chunks(record): | |
| chunk = self._get(chunk_key) | |
| output_file.write(chunk) | |
| digest.update(chunk) | |
| size += len(chunk) | |
| if size != record.size: | |
| raise ValueError( | |
| f"file size mismatch for {record.path}: " | |
| f"expected {record.size}, got {size}" | |
| ) | |
| if digest.hexdigest() != record.sha256: | |
| raise ValueError(f"sha256 mismatch for {record.path}") | |
| temp_output.rename(output) | |
| except BaseException: | |
| if temp_output.exists(): | |
| temp_output.unlink() | |
| raise |
| def _add_to_indexes(self, manifest: ModelFileManifest) -> None: | ||
| all_models = self._read_index(model_index_key()) | ||
| all_models.add(manifest.checkpoint_id) | ||
| self._write_index(model_index_key(), all_models) | ||
|
|
||
| per_model = self._read_index(model_id_index_key(manifest.model_id)) | ||
| per_model.add(manifest.checkpoint_id) | ||
| self._write_index(model_id_index_key(manifest.model_id), per_model) |
There was a problem hiding this comment.
The read-modify-write cycle for updating index files (weight/index/models and weight/index/model/{sha256}) is not atomic. If multiple clients import checkpoints concurrently, they can overwrite each other's updates, leading to lost index entries. Consider using a proper metadata database/catalog backend or prefix-based key listing if supported by the store to avoid a single global index object.
| def _read_index(self, key: str) -> set[str]: | ||
| if not self._exists(key): | ||
| return set() | ||
| value = self.store.get(key) | ||
| if value is None: | ||
| return set() | ||
| if isinstance(value, str): | ||
| payload = value.encode("utf-8") | ||
| else: | ||
| payload = bytes(value) | ||
| if not payload: | ||
| return set() | ||
| return set(json.loads(payload.decode("utf-8"))) |
There was a problem hiding this comment.
Calling _exists before get causes an unnecessary extra round-trip to the store. Since self.store.get(key) already returns None if the key does not exist, we can directly call get and handle the None case, which is much more efficient.
def _read_index(self, key: str) -> set[str]:
value = self.store.get(key)
if value is None:
return set()
if isinstance(value, str):
payload = value.encode("utf-8")
else:
payload = bytes(value)
if not payload:
return set()
return set(json.loads(payload.decode("utf-8")))| except BaseException: | ||
| for chunk_key in chunks: | ||
| self._remove(chunk_key, force=True) | ||
| raise |
There was a problem hiding this comment.
If self._remove raises an exception during cleanup (e.g., due to a temporary network issue or store error), the loop will be interrupted, leaving the remaining chunks orphaned in the store. Wrapping the individual chunk deletion in a try...except block ensures that we attempt to clean up as many chunks as possible.
except BaseException:
for chunk_key in chunks:
try:
self._remove(chunk_key, force=True)
except Exception:
pass
raise
github-actions
Bot
added
documentation
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
The CI failed for the non-consistent docs, could you take a look? Thanks |
|
|
@qoder Review it |
|
I've finished a pass over PR #2829. Here’s a concise summary of my review: Strengths
Issues / Risks
Recommended tests / follow-ups
🤖 Generated by Qoder • View workflow run |
184294f to
10fee62
Compare
10fee62 to
1721c66
Compare
Aionw
left a comment
Aionw
left a comment
There was a problem hiding this comment.
Re-review of the latest head. These four findings are new and intentionally exclude issues already covered by existing review threads.
| try: | ||
| self._remove(chunk_key, force=True) | ||
| except BaseException: | ||
| pass |
There was a problem hiding this comment.
[P2] Preserve keys whose cleanup fails. Continuing through all removals is useful, but swallowing a removal failure loses the only cleanup record for that partial file: the current file is never appended to records, so the FAILED manifest and a later delete_model() cannot find the surviving chunk. I reproduced an orphan that remains even after deleting the failed checkpoint. Please retain every partial chunk key in a cleanup journal/failed record until removal succeeds.
There was a problem hiding this comment.
Fixed. Partial chunk keys are now preserved in a FAILED record until cleanup succeeds, so delete_model() can reliably reclaim any survivors after a failed import.
Resolve all 5 review comments from Aionw on PR kvcache-ai#2829: - F1 (model.py): verify chunks BEFORE publishing the manifest, so a READY status is never visible for unverified data. - F2 (model.py): persist per-file chunk_size in the manifest; readers refuse to guess a multi-chunk layout when the field is absent (0=legacy). - F3 (model.py): on a failed import, record already-written chunk keys in a journal so a later delete can reclaim them; nothing is silently leaked. - F4 (model.py): remove the manifest LAST during delete (de-index first, then chunks, then manifest), making delete crash-safe and retryable. - F5 (model_keyspace.py): add isinstance + length bound (<=255) to id validation in addition to the existing charset regex. Also add regression tests covering crash-recovery, write-once stores, and chunk_size round-trip (14 new cases in test_weight_model_cache.py).
6 participants
Description
Add a file-level weight caching control plane on top of MooncakeStore, so model checkpoints can be imported into the store and served to inference engines (e.g. SGLang) as plain files — without a tensor-level, GPU-coupled save path.
The idea is to treat MooncakeStore as a generic model cache that is decoupled from the inference engine: weights are stored file-level (HuggingFace safetensors chunked with a per-model manifest), and engines load files exactly as they would from disk, only the source changes (disk → MooncakeStore). This control plane provides the management side — import / list / inspect / verify / delete — independently of any serving process.
The matching SGLang connector (data plane) is in sgl-project/sglang#30728. Design discussion / RFC: #2282 (Unified KVCache and Model Weight Management); this PR implements Phase 1 (store-backed weight loading) plus the import/manage control plane.
Module
mooncake-wheel)Type of Change
How Has This Been Tested?
Test commands:
cd mooncake-wheel python -m pytest tests/test_weight_model_cache.py tests/test_weight_model_cli.py -vEnd-to-end (cross-machine, one storage node + one 8×H20 node):
Then loaded Qwen3-235B-A22B from the store via the SGLang connector in ~30s
with correct inference output.
Test results:
Checklist
./scripts/code_format.shpre-commit run --all-filesand all hooks passAI Assistance Disclosure
AI assistance (Claude Code) was used for parts of the implementation, tests, and documentation. All changes have been reviewed and are understood by the human submitter.