Add Mooncake hidden states extraction backend#836
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughChangesMooncake hidden-states flow
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
The quality checks have failed. Please run |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/e2e/smoke/test_mooncake_e2e_smoke.py (1)
1-135: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winApply Ruff formatting.
The Quality workflow currently fails because
ruff format --checkwould reformat this file.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/e2e/smoke/test_mooncake_e2e_smoke.py` around lines 1 - 135, Apply Ruff formatting to the entire test_mooncake_hidden_states_roundtrip module, using the repository’s configured Ruff formatter so it passes ruff format --check. Preserve the existing test behavior and logic while updating only formatting.Source: Pipeline failures
🧹 Nitpick comments (2)
tests/unit/data_generation/test_mooncake_store.py (1)
15-18: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for
MooncakeStoreConfig.from_dictfield filtering.The connector builds
MooncakeStoreConfigfrom an untrusted extra-config dict viafrom_dict, which silently drops unknown keys. That contract has no direct test here.✅ Example test to add
def test_from_dict_filters_unknown_keys(): cfg = MooncakeStoreConfig.from_dict({"protocol": "rdma", "bogus": 1}) assert cfg.protocol == "rdma" assert not hasattr(cfg, "bogus")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/data_generation/test_mooncake_store.py` around lines 15 - 18, Add a unit test in the Mooncake store tests covering MooncakeStoreConfig.from_dict with a valid protocol field and an unknown key, asserting the valid value is retained and the unknown key is not present on the resulting configuration.hs_connectors/src/hs_connectors/transfer.py (1)
233-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate CLI-arg wiring and hostname/config-resolution logic within
MooncakeBackend.
add_train_args/add_launch_argsdefine identical argument sets, andfrom_train_args/build_kv_transfer_configindependently duplicate the hostname-resolution +MooncakeStoreConfigconstruction. Extracting shared helpers would reduce drift risk when new Mooncake options are added.♻️ Sketch: shared helpers
+def _add_mooncake_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--mooncake-master", type=str, default="127.0.0.1:50051", ...) + parser.add_argument("--mooncake-metadata-server", type=str, default="P2PHANDSHAKE", ...) + parser.add_argument("--mooncake-protocol", choices=["tcp", "rdma"], default="tcp", ...) + +def _resolve_mooncake_config(args: argparse.Namespace) -> "MooncakeStoreConfig": + local_hostname = os.environ.get("MOONCAKE_LOCAL_HOSTNAME") or socket.gethostbyname(socket.gethostname()) + return MooncakeStoreConfig( + local_hostname=local_hostname, + metadata_server=args.mooncake_metadata_server, + master_server_address=args.mooncake_master, + protocol=args.mooncake_protocol, + )Then
add_train_args/add_launch_argscall_add_mooncake_args(parser), andfrom_train_args/build_kv_transfer_configcall_resolve_mooncake_config(args).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@hs_connectors/src/hs_connectors/transfer.py` around lines 233 - 346, Extract the duplicated Mooncake CLI option definitions from add_train_args and add_launch_args into a shared _add_mooncake_args(parser) helper, and have both methods delegate to it. Extract hostname resolution and MooncakeStoreConfig construction from from_train_args and build_kv_transfer_config into a shared _resolve_mooncake_config(args) helper, then reuse that config in both call sites while preserving their existing behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@hs_connectors/src/hs_connectors/mooncake_hidden_states_connector.py`:
- Around line 199-213: Update the completion handling around _req_futures and
done_sending so failed futures are not reported as successfully finished. When
future.exception() is non-null, propagate the failure through the connector’s
established error mechanism (or write the distinguishable store failure marker
used by get_sample), and only add the request ID to done_sending after a
successful write; preserve cleanup of the future and accumulated request state.
- Around line 88-91: Update the assertion in MooncakeHiddenStatesConnector to
validate that speculative_config specifies the extract_hidden_states method, not
merely that the configuration exists. Keep the error message aligned with this
explicit requirement.
In `@src/speculators/train/trainer.py`:
- Around line 618-627: Move the eager val_loader iterator creation from the
post-model-setup portion of run_training() to before setup_model() initializes
or moves the model. Preserve the existing guards on val_loader, num_workers, and
_iterator so workers are spawned only when needed.
In `@tests/e2e/smoke/test_mooncake_e2e_smoke.py`:
- Around line 26-31: Skip the Mooncake smoke tests when the underlying Mooncake
dependency is unavailable: in tests/e2e/smoke/test_mooncake_e2e_smoke.py, add
pytest.importorskip("mooncake.store", ...) before importing
hs_connectors.mooncake_store; in
tests/e2e/smoke/test_mooncake_online_training.py, add the same guard before
starting mooncake_master. Keep the existing connector import and startup flow
unchanged after the dependency checks.
---
Outside diff comments:
In `@tests/e2e/smoke/test_mooncake_e2e_smoke.py`:
- Around line 1-135: Apply Ruff formatting to the entire
test_mooncake_hidden_states_roundtrip module, using the repository’s configured
Ruff formatter so it passes ruff format --check. Preserve the existing test
behavior and logic while updating only formatting.
---
Nitpick comments:
In `@hs_connectors/src/hs_connectors/transfer.py`:
- Around line 233-346: Extract the duplicated Mooncake CLI option definitions
from add_train_args and add_launch_args into a shared _add_mooncake_args(parser)
helper, and have both methods delegate to it. Extract hostname resolution and
MooncakeStoreConfig construction from from_train_args and
build_kv_transfer_config into a shared _resolve_mooncake_config(args) helper,
then reuse that config in both call sites while preserving their existing
behavior.
In `@tests/unit/data_generation/test_mooncake_store.py`:
- Around line 15-18: Add a unit test in the Mooncake store tests covering
MooncakeStoreConfig.from_dict with a valid protocol field and an unknown key,
asserting the valid value is retained and the unknown key is not present on the
resulting configuration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 47b26def-a3cf-4544-97e2-cf5a47f42319
📒 Files selected for processing (12)
hs_connectors/pyproject.tomlhs_connectors/src/hs_connectors/__init__.pyhs_connectors/src/hs_connectors/mooncake_hidden_states_connector.pyhs_connectors/src/hs_connectors/mooncake_store.pyhs_connectors/src/hs_connectors/transfer.pypyproject.tomlsrc/speculators/data_generation/vllm_client.pysrc/speculators/train/trainer.pytests/e2e/smoke/test_mooncake_e2e_smoke.pytests/e2e/smoke/test_mooncake_online_training.pytests/e2e/utils.pytests/unit/data_generation/test_mooncake_store.py
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
Design looks good -- the Mooncake connector follows the ExampleHiddenStatesConnector patterns closely (group/block-size resolution via spec, scheduler-side pending-saves, worker-side async extraction) and the backend registration through HiddenStatesBackend.register is clean. One missing safety check below.
🤖 Generated with Claude Code using the /pr-review skill
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
WindChimeRan
left a comment
WindChimeRan
left a comment
There was a problem hiding this comment.
I love the direction of single node without using disk as an intermediary and I think we should use mooncake as the default.
However, I couldn't get it runnable (our A100 machine):
My claude code's diagnosis:
0.3.12's prebuilt mooncake_master segfaults on startup, and the last working master (0.3.11) is manylinux_2_35-only — so mooncake can't run on glibc-2.34 hosts (RHEL 9); please ship a fixed master on a lower-glibc wheel or pin a known-good version.
Also Please run some benchmarks to show the speedup from the saving of the original disk lock data transfer.
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
|
@fynnsu #841 just landed — the Mooncake backend needs its args mirrored into the schema 👋 Heads up: #841 (config-file-first CLI) is now in This PR doesn't touch
The fix — mirror the 3 args into the schemaYour Option A — mirror into # src/speculators/train/config/schema.py — in class DataArgs
mooncake_master: str = Field(
default="127.0.0.1:50051",
description="Mooncake master server address. "
"Used with --hidden-states-backend mooncake.",
)
mooncake_metadata_server: str = Field(
default="P2PHANDSHAKE",
description="Mooncake metadata server (or P2PHANDSHAKE). "
"Used with --hidden-states-backend mooncake.",
)
mooncake_protocol: Literal["tcp", "rdma"] = Field(
default="tcp",
description="Mooncake transport protocol. "
"Used with --hidden-states-backend mooncake.",
)Cheapest change, but Option B — a dedicated # src/speculators/train/config/schema.py
class BackendArgs(_Group):
"""Per-backend transport args, mirrored from each HiddenStatesBackend's
add_train_args (hs_connectors stays argparse-based). Only the fields for the
selected --hidden-states-backend apply."""
hidden_states_path: str | None = Field(default=None, ...) # 'file' (moved from DataArgs)
mooncake_master: str = Field(default="127.0.0.1:50051", ...) # 'mooncake'
mooncake_metadata_server: str = Field(default="P2PHANDSHAKE", ...)
mooncake_protocol: Literal["tcp", "rdma"] = Field(default="tcp", ...)
# register in _GROUPS (order defines --help + run.yaml layout), right after "data"
"data": DataArgs,
"backend": BackendArgs,
# and add the field on TrainConfig
backend: BackendArgs = Field(default_factory=BackendArgs)More churn on the just-landed schema, and it shifts My take: since the reconciliation test itself notes NIXL/RDMA are coming, I'd lean B (it's cheapest to extract now, with only one field to move) — but A is totally fine if you want to keep this PR minimal; just file a follow-up to extract
Sanity check after adding thempython -m pytest tests/unit/train/config/test_backend_reconciliation.py
python scripts/train.py --help | grep mooncake # all three flags show upPing me if you'd like the change as a ready-to-apply patch. Or if you'd like me to push the commit |
4 participants
Purpose
Adds a mooncake backend to the hs_connectors sub-package. This will enable training on multi-node setups, or single node without using disk as an intermediary.
Requires the
mooncake-transfer-enginepackage.Tests
Comprehensive testing is still in progress.
I have successfully completed a basic regression test training using this backend.
Checklist
I have filled in: