Skip to content

fix(pd): harden production cutover lifecycles - #1472

Closed
sufubao wants to merge 23 commits into
ModelTC:mainfrom
sufubao:deploy/production-cutover-readiness-20260819
Closed

fix(pd): harden production cutover lifecycles#1472
sufubao wants to merge 23 commits into
ModelTC:mainfrom
sufubao:deploy/production-cutover-readiness-20260819

Conversation

@sufubao

@sufubao sufubao commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR hardens PD separation for production cutover across request admission, first-token arbitration, WebSocket ownership, and NIXL peer/transfer lifecycles.

  • Preserve prefill cache metadata while preventing duplicate first-token output, with a bounded fallback when the prefill token is missing.
  • Preserve a terminal decode marker so KV-transfer failures cannot be hidden by first-token deduplication.
  • Cancel worker generation tasks even while they are waiting for shared-memory request slots, and release partially allocated slots on cancellation.
  • Serialize NIXL peer admission/removal and protect reconnects with peer generations and active-transfer tracking.
  • Bind PD registrations to their owning WebSocket, fail requests assigned to replaced or disconnected nodes, and ignore stale disconnects.
  • Supervise worker generation, token-forwarding, and heartbeat tasks so background failures tear down the connection instead of becoming silent hangs.
  • Route NIXL status and completion-notification failures through normal task cleanup; make any remaining uncaught daemon-thread failure visible to process supervision.

Production failure modes addressed

Before these changes, a request could remain stuck after its assigned PD node disconnected, a stale WebSocket cleanup could remove a newer registration, and worker generation exceptions could disappear while the receive loop stayed alive. NIXL status-thread exceptions could also strand pages and transfer handles while the transfer subprocess still appeared healthy.

The resulting lifecycle guarantees are:

  1. A WebSocket owns exactly the registration it created.
  2. Requests are awakened immediately when either assigned node is replaced or disconnected.
  3. A worker connection cannot remain healthy after one of its critical background tasks fails.
  4. A handled NIXL task failure reaches the existing failed-task cleanup path.
  5. An unhandled transfer-thread failure terminates the subprocess and becomes observable by existing liveness supervision.

Inline comments document the non-obvious request-slot cancellation and NIXL failure-recovery invariants.

Audited fixes, one commit each

  • f6a44175 fix(pd): fail requests when assigned nodes disconnect
  • 623d640d fix(pd): propagate worker background failures
  • feedc43e fix(pd): recover NIXL transfer task failures

Each commit includes focused regression coverage and a commit body describing its root cause and runtime behavior.

Validation

PYTHONPATH=. python -m pytest -q \
  unit_tests/server/httpserver/test_pd_pending_abort.py \
  unit_tests/server/httpserver/test_pd_connection_tasks.py \
  unit_tests/test_nixl_kv_transporter_thread_safety.py \
  unit_tests/server/httpserver/test_pd_master_cached_tokens.py \
  test/test_pd_selector/test_pd_master_multi_choice.py \
  unit_tests/server/test_pd_master_mode.py \
  unit_tests/server/test_pd_cache_aware.py

51 passed

The touched files also pass the repository flake8 configuration and git diff --check.

No benchmark was run because this change is correctness and lifecycle focused.

sufubao and others added 23 commits August 13, 2026 17:55
Track registrations by their owning WebSocket so a stale disconnect cannot evict a replacement connection. Wake requests assigned to a replaced or disconnected node and surface the connection failure at every PD wait stage.
Supervise token forwarding, heartbeat, and per-request generation tasks while the worker waits for PD master messages. Tear down the connection when a background task fails so the master observes a disconnect instead of waiting on a silently dead request.
Route status-query and completion-notification errors through the existing failed-task cleanup path so pages and transfer handles are reclaimed. Install a fatal thread exception hook in transfer subprocesses so any unhandled daemon-thread failure is visible to process supervision.
Apply the Black 21.12b0 line-width-120 output required by the upstream pre-commit workflow. This commit is mechanical and contains no runtime changes.
@sufubao

sufubao commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator Author

Superseded by three independent PRs, each based directly on upstream main and containing one fix only:

The replacement PRs intentionally omit the added unit-test changes.

@sufubao sufubao closed this Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant