[https://nvbugs/6566735][fix] recover disaggregated worker heartbeat registration - #17461
Open
reasonsolo wants to merge 3 commits into
Open
[https://nvbugs/6566735][fix] recover disaggregated worker heartbeat registration#17461reasonsolo wants to merge 3 commits into
reasonsolo wants to merge 3 commits into
Conversation
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #64947 [ run ] triggered by Bot. Commit: |
Contributor
WalkthroughThe worker registration API now supports overwriting existing storage entries. Registration retries preserve this option. Heartbeat recovery enables overwriting, and regression tests verify failure state, overwrite behavior, and TTL. ChangesWorker registration recovery
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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
🧹 Nitpick comments (3)
tensorrt_llm/serve/disagg_auto_scaling.py (2)
318-321: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType and document the new public option.
Add
overwrite_if_exists: bool = Falseand document that it replaces an existing storage entry. Include the option in a Google-style docstring. Typevalidatorandretry_intervalwhen updating this public signature.Proposed signature change
- overwrite_if_exists=False) -> bool: + overwrite_if_exists: bool = False) -> bool:As per coding guidelines, Python functions must be annotated and externally usable interfaces must use Google-style docstrings.
🤖 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 `@tensorrt_llm/serve/disagg_auto_scaling.py` around lines 318 - 321, Update the public async register_worker signature to annotate validator and retry_interval, and declare overwrite_if_exists as a bool defaulting to False. Add or extend its Google-style docstring with an Args entry explaining that overwrite_if_exists replaces an existing storage entry, while documenting the other newly typed parameters as needed.Source: Coding guidelines
318-321: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAnnotate and document
register_worker. Add precise types for all parameters and documentoverwrite_if_existssemantics. No incompatible overrides or wrappers exist.🤖 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 `@tensorrt_llm/serve/disagg_auto_scaling.py` around lines 318 - 321, Update register_worker with precise type annotations for validator, retry_interval, and overwrite_if_exists, and retain the bool return annotation. Add a concise docstring documenting each parameter, explicitly stating when overwrite_if_exists replaces an existing worker registration.tests/unittest/disaggregated/test_disagg_cluster_manager_worker.py (1)
30-56: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for registration retry.
- Test coverage summary: one test added; no tests modified or removed.
- CI coverage includes
l0_cpu.ymland the hardware-specific test lists. The QA list contains a different disaggregated unit test, not this test.- Add a test where
storage.set()returnsFalseonce, then succeeds, and assert the retry inregister_worker()attensorrt_llm/serve/disagg_auto_scaling.py:338-347.- Coverage verdict: insufficient. Run
pytest tests/unittest/.🤖 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/unittest/disaggregated/test_disagg_cluster_manager_worker.py` around lines 30 - 56, Add a test for the registration retry path in register_worker: configure storage.set to return False on the first call and True on the second, invoke registration, and assert it is awaited twice with the expected registration arguments. Keep the existing heartbeat recovery test unchanged and verify the full unit-test suite with pytest tests/unittest/.Sources: Coding guidelines, Path instructions
🤖 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 `@tensorrt_llm/serve/disagg_auto_scaling.py`:
- Around line 338-347: Update the failure path in register_worker after the
retry branch: when registration fails and retry_interval is non-positive, return
False immediately before updating _last_heartbeat or starting heartbeat
handling. Preserve the existing recursive retry behavior for positive retry
intervals.
---
Nitpick comments:
In `@tensorrt_llm/serve/disagg_auto_scaling.py`:
- Around line 318-321: Update the public async register_worker signature to
annotate validator and retry_interval, and declare overwrite_if_exists as a bool
defaulting to False. Add or extend its Google-style docstring with an Args entry
explaining that overwrite_if_exists replaces an existing storage entry, while
documenting the other newly typed parameters as needed.
- Around line 318-321: Update register_worker with precise type annotations for
validator, retry_interval, and overwrite_if_exists, and retain the bool return
annotation. Add a concise docstring documenting each parameter, explicitly
stating when overwrite_if_exists replaces an existing worker registration.
In `@tests/unittest/disaggregated/test_disagg_cluster_manager_worker.py`:
- Around line 30-56: Add a test for the registration retry path in
register_worker: configure storage.set to return False on the first call and
True on the second, invoke registration, and assert it is awaited twice with the
expected registration arguments. Keep the existing heartbeat recovery test
unchanged and verify the full unit-test suite with pytest tests/unittest/.
🪄 Autofix
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: Enterprise
Run ID: f5d2b234-1497-4708-9440-c82267a190fe
📒 Files selected for processing (2)
tensorrt_llm/serve/disagg_auto_scaling.pytests/unittest/disaggregated/test_disagg_cluster_manager_worker.py
Shixiaowei02
approved these changes
Aug 10, 2026
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/disaggregated/test_disagg_cluster_manager_worker.py (1)
30-48: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd retry-path coverage for
register_worker.
test_registration_failure_without_retry_returns_falsecovers onlyretry_interval=0. Add a test forstorage.setreturningFalseand thenTrue. Mock the retry sleep. Assert that both calls receiveoverwrite_if_exists=True.Test coverage summary:
- Added:
test_registration_failure_without_retry_returns_false,test_heartbeat_recovery_overwrites_stale_registration.- Modified or removed: none.
- CI coverage:
tests/integration/test_lists/test-db/l0_cpu.ymlincludesunittest/disaggregated.- QA coverage: no matching entry exists in
tests/integration/test_lists/qa/llm_function_core.txt.- Coverage verdict: insufficient until the positive retry path is covered.
🤖 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/unittest/disaggregated/test_disagg_cluster_manager_worker.py` around lines 30 - 48, Add a retry-path test alongside test_registration_failure_without_retry_returns_false where storage.set returns False first and True second, mock the retry sleep to avoid delays, and invoke register_worker with a positive retry interval. Assert registration succeeds and both storage.set calls include overwrite_if_exists=True.Source: Path instructions
🤖 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.
Nitpick comments:
In `@tests/unittest/disaggregated/test_disagg_cluster_manager_worker.py`:
- Around line 30-48: Add a retry-path test alongside
test_registration_failure_without_retry_returns_false where storage.set returns
False first and True second, mock the retry sleep to avoid delays, and invoke
register_worker with a positive retry interval. Assert registration succeeds and
both storage.set calls include overwrite_if_exists=True.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 49c67b18-164d-4c6a-9a7f-23ee17f9a1da
📒 Files selected for processing (2)
tensorrt_llm/serve/disagg_auto_scaling.pytests/unittest/disaggregated/test_disagg_cluster_manager_worker.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/serve/disagg_auto_scaling.py
Collaborator
|
PR_Github #64947 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #64997 [ run ] triggered by Bot. Commit: |
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
3 participants
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.
Dev Engineer Review
overwrite_if_existstoregister_worker.False.QA Engineer Review
test_registration_failure_without_retry_returns_false.test_heartbeat_recovery_overwrites_stale_registration.tests/integration/test_lists/.Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.