Skip to content

[Bugfix] Replace assertions with defensive guards in KV transfer finished callbacks - #6

Open
Jackie2049 wants to merge 1 commit into
mainfrom
fix/kv-transfer-race-condition-v2
Open

[Bugfix] Replace assertions with defensive guards in KV transfer finished callbacks#6
Jackie2049 wants to merge 1 commit into
mainfrom
fix/kv-transfer-race-condition-v2

Conversation

@Jackie2049

@Jackie2049 Jackie2049 commented Jun 4, 2026

Copy link
Copy Markdown
Owner

Summary

Closes #1

In disaggregated PD deployments, a race condition exists where a request may be aborted (removed from self.requests) while KV transfer is still in flight. The assert statements in _update_from_kv_xfer_finished() crash the entire serving process instead of gracefully handling the race.

Reference: This fix addresses a production crash scenario in PD-separated deployments, similar to issues discussed in the NIXL KV connector roadmap (vllm-project#33702).

Changes

Replace 3 assertions with defensive guard checks in vllm/v1/core/sched/scheduler.py:

  1. finished_recving — request not found: assert req_id in self.requestsif req_id not in self.requests: logger.warning(...); continue
  2. finished_recving — unexpected status: assert RequestStatus.is_finished(req.status)if not RequestStatus.is_finished(...): logger.warning(...); continue
  3. finished_sending — request not found: assert req_id in self.requestsif req_id not in self.requests: logger.warning(...); continue

Why this is safe

  • If the request was already removed, its blocks were already freed during the abort path (_free_request)
  • Skipping the duplicate free is the correct behavior
  • The continue ensures no other cleanup logic is skipped for remaining requests in the loop
  • Warning logs preserve observability for debugging

Testing plan

  • Code review: minimal, focused change (3 assertions → if/continue)
  • Unit test: simulate abort + finished_sending race condition
    • 6 scenarios validated (normal operation, recv race, send race, unexpected status, mixed, empty lists)
    • All scenarios PASS: old code crashes on every race, new code gracefully handles with warnings
    • Normal operation identical: guard version produces same recv_ids and freed_blocks as assert version
    • Validation script: tools/test_kv_transfer_guard.py in rollout-infra repo (standalone, no vLLM dependency)
  • Verify no memory leak (blocks are not double-freed)
    • Aborted requests have blocks freed during abort path; guard correctly skips duplicate free
  • Verify normal KV transfer flow is unaffected
    • When all req_ids exist, guard version produces identical results to assert version (0 warnings)

@Jackie2049
Jackie2049 force-pushed the fix/kv-transfer-race-condition-v2 branch from 1c78df3 to 91a3f26 Compare June 4, 2026 05:24
…shed callbacks

In disaggregated PD deployments, a race condition exists where a request
may be aborted (removed from self.requests) while KV transfer is still
in flight. The assertions in _update_from_kv_xfer_finished crash the
entire serving process instead of gracefully handling the race.

Replace assertions with if/continue guard checks and warning logs.
If a request was already removed, its blocks were already freed during
the abort path, so skipping is the correct behavior.

Closes: #1

Co-authored-by: Boundless <ruihang_wu@163.com>
@Jackie2049
Jackie2049 force-pushed the fix/kv-transfer-race-condition-v2 branch from 91a3f26 to d81d4c7 Compare June 4, 2026 05:43
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.

[Bug] Race condition in scheduler finished_sending KV transfer causes production crash

1 participant