feat: continuous metadata publish retry via heartbeat thread#214
Open
zhengluo-nv wants to merge 1 commit into
Open
feat: continuous metadata publish retry via heartbeat thread#214zhengluo-nv wants to merge 1 commit into
zhengluo-nv wants to merge 1 commit into
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
WalkthroughThis PR refactors metadata publication to defer gRPC calls from startup into an async callback. The worker gRPC server now supports optional/deferred ChangesDeferred metadata publication flow
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelexpress_client/python/modelexpress/worker_server.py (1)
34-49: 🛠️ Refactor suggestion | 🟠 MajorHandle the pre-
mx_source_idwindow inGetTensorManifest().With the new late-bound flow, the gRPC server can accept manifest requests before
set_mx_source_id()runs.fetch_tensor_manifest()always sends a non-empty id, so the current check turns that legitimate startup race intoFAILED_PRECONDITION, and the target-side RDMA load can fail before it ever retries another source.One simple way to make bootstrap requests succeed
def GetTensorManifest(self, request, context): - if request.mx_source_id and request.mx_source_id != self._mx_source_id: + if ( + self._mx_source_id is not None + and request.mx_source_id + and request.mx_source_id != self._mx_source_id + ): context.abort( grpc.StatusCode.FAILED_PRECONDITION, f"mx_source_id mismatch: expected {self._mx_source_id}, " f"got {request.mx_source_id}", )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelexpress_client/python/modelexpress/worker_server.py` around lines 34 - 49, GetTensorManifest currently aborts when request.mx_source_id differs from self._mx_source_id, which breaks the startup race where the server hasn't run set_mx_source_id yet; modify GetTensorManifest so that if self._mx_source_id is None and request.mx_source_id is provided, it accepts the request (and optionally initializes self._mx_source_id = request.mx_source_id) instead of aborting, otherwise keep the existing mismatch check and abort behavior; update logic in GetTensorManifest (referencing GetTensorManifest, set_mx_source_id, and _mx_source_id) so bootstrap manifest requests succeed when the server-side id is not yet set.
🧹 Nitpick comments (1)
modelexpress_client/python/tests/test_vllm_loader.py (1)
453-458: Force centralized mode in the non-P2P tests.Both tests inherit
MX_P2P_METADATAfrom the runner because the env patch only setsMX_CONTIGUOUS_REG. If CI or a local shell already has P2P enabled,start_metadata_publisher()takes the other branch and these assertions become flaky.Proposed test hardening
- with patch.dict("os.environ", {"MX_CONTIGUOUS_REG": "0"}), \ + with patch.dict("os.environ", {"MX_CONTIGUOUS_REG": "0", "MX_P2P_METADATA": "0"}), \ patch("modelexpress.metadata.HeartbeatThread", return_value=mock_hb) as hb_cls: start_metadata_publisher( mx_client, nixl_manager, tensors, global_rank=2, device_id=0, identity=identity, worker_id="inst-uuid", )- with patch.dict("os.environ", {"MX_CONTIGUOUS_REG": "0"}), \ + with patch.dict("os.environ", {"MX_CONTIGUOUS_REG": "0", "MX_P2P_METADATA": "0"}), \ patch("modelexpress.metadata.HeartbeatThread", return_value=mock_hb) as hb_cls: start_metadata_publisher( mx_client, nixl_manager, {}, global_rank=0, device_id=0, identity=identity, worker_id="w-1", )Also applies to: 483-488
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelexpress_client/python/tests/test_vllm_loader.py` around lines 453 - 458, The test is flaky because it only patches MX_CONTIGUOUS_REG and leaves MX_P2P_METADATA inherited from the environment; update the test patch in modelexpress_client/python/tests/test_vllm_loader.py so the with patch.dict(...) includes "MX_P2P_METADATA": "0" (in the same context where MX_CONTIGUOUS_REG is set) to force centralized mode before calling start_metadata_publisher(mx_client, nixl_manager, tensors, ...); apply the same change to the other similar test block that calls start_metadata_publisher around the 483-488 region.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@modelexpress_client/python/modelexpress/heartbeat.py`:
- Around line 169-180: After a successful call to _publish_fn() the code sets
_mx_source_id but leaves _started False until the next READY heartbeat, so
_mark_stale() will return early on shutdown and not remove the published source;
fix by marking the worker as started (set self._started = True) immediately
after assigning self._mx_source_id in the try block (right after the successful
publish log) so that shutdown/_mark_stale() can proceed to mark the source STALE
if the process exits before the next READY heartbeat.
---
Outside diff comments:
In `@modelexpress_client/python/modelexpress/worker_server.py`:
- Around line 34-49: GetTensorManifest currently aborts when
request.mx_source_id differs from self._mx_source_id, which breaks the startup
race where the server hasn't run set_mx_source_id yet; modify GetTensorManifest
so that if self._mx_source_id is None and request.mx_source_id is provided, it
accepts the request (and optionally initializes self._mx_source_id =
request.mx_source_id) instead of aborting, otherwise keep the existing mismatch
check and abort behavior; update logic in GetTensorManifest (referencing
GetTensorManifest, set_mx_source_id, and _mx_source_id) so bootstrap manifest
requests succeed when the server-side id is not yet set.
---
Nitpick comments:
In `@modelexpress_client/python/tests/test_vllm_loader.py`:
- Around line 453-458: The test is flaky because it only patches
MX_CONTIGUOUS_REG and leaves MX_P2P_METADATA inherited from the environment;
update the test patch in modelexpress_client/python/tests/test_vllm_loader.py so
the with patch.dict(...) includes "MX_P2P_METADATA": "0" (in the same context
where MX_CONTIGUOUS_REG is set) to force centralized mode before calling
start_metadata_publisher(mx_client, nixl_manager, tensors, ...); apply the same
change to the other similar test block that calls start_metadata_publisher
around the 483-488 region.
🪄 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: 02994a4f-1d8d-427a-9b79-669465eb2783
📒 Files selected for processing (9)
examples/p2p_transfer_k8s/client/vllm/vllm-single-node-p2p.yamlmodelexpress_client/python/modelexpress/heartbeat.pymodelexpress_client/python/modelexpress/load_strategy/base.pymodelexpress_client/python/modelexpress/load_strategy/default_strategy.pymodelexpress_client/python/modelexpress/load_strategy/rdma_strategy.pymodelexpress_client/python/modelexpress/metadata.pymodelexpress_client/python/modelexpress/nixl_transfer.pymodelexpress_client/python/modelexpress/worker_server.pymodelexpress_client/python/tests/test_vllm_loader.py
💤 Files with no reviewable changes (1)
- modelexpress_client/python/modelexpress/nixl_transfer.py
zhengluo-nv
changed the title
zhengluo-nv
force-pushed
the
zheluo/continuous-metadata-publish
branch
from
31f3ccf to
909ac74
Compare
zhengluo-nv
force-pushed
the
zheluo/continuous-metadata-publish
branch
from
909ac74 to
f9910de
Compare
Contributor
Author
|
Addressed code rabbit comments |
zhengluo-nv
force-pushed
the
zheluo/continuous-metadata-publish
branch
from
40dbf7f to
674a442
Compare
zhengluo-nv
force-pushed
the
zheluo/continuous-metadata-publish
branch
from
674a442 to
0a6fb27
Compare
zhengluo-nv
force-pushed
the
zheluo/continuous-metadata-publish
branch
from
0d52585 to
d35f7e5
Compare
Contributor
Author
|
@coderabbitai review |
Contributor
✅ Actions performedReview triggered.
|
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
zhengluo-nv
force-pushed
the
zheluo/continuous-metadata-publish
branch
from
d35f7e5 to
99f4ad0
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
main.HeartbeatThreadand retry on heartbeat ticks until success orMX_PUBLISH_TIMEOUT_SECSexpires.WorkerGrpcServerto start beforemx_source_idis known, then late-bind it after publish succeeds.Test plan
uv run --extra dev pytest tests/test_heartbeat.py tests/test_vllm_loader.py tests/test_k8s_service_client.pySummary by CodeRabbit
New Features
Refactor