feat: get block ids with load errors #5668
Open
Conversation
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
Add infrastructure to track block IDs that fail during async KV cache loading, enabling vLLM to identify and handle failed blocks gracefully. Changes: - protocol.rs: Add block_ids field to WorkerTransferRequest - scheduler.rs: Update tests with new block_ids field - slot.rs: Populate block_ids for onboard/offload operations - worker.rs: Add failed_block_ids tracking and get_block_ids_with_load_errors method - connector_worker.py: Expose get_block_ids_with_load_errors to Python - dynamo_connector.py: Implement KVConnectorBase_V1 interface method Signed-off-by: Patrick Riel <priel@nvidia.com>
…acking - Add failure_tx/failure_rx channel between Scheduler and WorkerSchedulerClient - Add drain_failures() method to collect pending failure notifications - Track request_id -> uuid -> block_ids mapping in KvConnectorWorker - Populate failed_block_ids from scheduler failure notifications - Clean up block_ids tracking when requests complete Signed-off-by: Patrick Riel <priel@nvidia.com>
Contributor
WalkthroughThis pull request introduces a block load error tracking system that propagates failed block IDs through the KV connector stack. Methods are added to Python and Rust layers to expose and retrieve failed block IDs, with supporting infrastructure to capture, track, and drain failure notifications across workers and schedulers. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 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
🤖 Fix all issues with AI agents
In `@lib/bindings/kvbm/src/block_manager/vllm/connector/worker.rs`:
- Around line 293-305: When enqueuing onboarding_operations we currently stash
request->blocks in request_to_blocks but later remove that mapping in
get_finished before we call get_block_ids_with_load_errors, which can drop the
failure->block correlation; fix by ensuring failures are drained/translated into
block IDs before cleanup: call get_block_ids_with_load_errors (or a shared
helper that drains scheduler failures and maps them to block IDs using
request_to_blocks) prior to removing the request's entry from request_to_blocks
in get_finished; factor the draining/mapping logic into a reusable helper (e.g.,
drain_and_map_load_failures) and use it from both the onboarding loop code paths
and the get_finished cleanup to preserve failure→block mappings.
🧹 Nitpick comments (1)
lib/llm/src/block_manager/connector/protocol.rs (1)
152-154: Confirm whether#[serde(default)]should mask missingblock_ids.
Defaulting to empty accepts older senders but silently drops error tracking. If deployments are atomic and you want strictness, consider making the field required so deserialization fails fast when it’s absent.Based on learnings, keep the spec explicit if backward compatibility isn't required.💡 Possible tweak to keep the protocol explicit
- #[serde(default)] - pub block_ids: Vec<usize>, + pub block_ids: Vec<usize>,
Comment on lines
293
to
+305
| // immediately enqueue the onboarding operations | ||
| for operation in onboarding_operations { | ||
| let request_id = operation.request_id.clone(); | ||
| let uuid = operation.uuid; | ||
| let block_ids = operation.block_ids.clone(); | ||
|
|
||
| // Store block_ids mapping for error tracking (only for load operations) | ||
| if !block_ids.is_empty() { | ||
| self.request_to_blocks | ||
| .entry(request_id.clone()) | ||
| .or_default() | ||
| .insert(uuid, block_ids); | ||
| } |
Contributor
There was a problem hiding this comment.
Avoid losing failure → block mappings on request cleanup.
request_to_blocks is removed in get_finished, but failures are only mapped when get_block_ids_with_load_errors() drains the scheduler. If get_finished runs first, failure notifications for that request can no longer be correlated to block IDs.
🛠️ Suggested fix: drain failures before cleanup and reuse a helper
+ fn absorb_failure_notifications(&mut self) {
+ let failures = self.connector.drain_failures();
+ for (request_id, failed_uuids) in failures {
+ if let Some(uuid_to_blocks) = self.request_to_blocks.get(&request_id) {
+ for uuid in failed_uuids {
+ if let Some(block_ids) = uuid_to_blocks.get(&uuid) {
+ for &block_id in block_ids {
+ self.failed_block_ids.insert(block_id as u32);
+ }
+ tracing::warn!(
+ request_id = %request_id,
+ operation_id = %uuid,
+ block_ids = ?block_ids,
+ "load operation failed; marking blocks as failed"
+ );
+ }
+ }
+ }
+ }
+ }
@@
- for request_id in &is_finished_onboarding {
+ // Drain failures while block-id mappings still exist.
+ self.absorb_failure_notifications();
+ for request_id in &is_finished_onboarding {
self.maybe_finished_onboarding.remove(request_id);
// Clean up block_ids tracking for this request
self.request_to_blocks.remove(request_id);
@@
- fn get_block_ids_with_load_errors(&mut self) -> HashSet<u32> {
- // Drain failures from the scheduler and convert (request_id, uuid) -> block_ids
- let failures = self.connector.drain_failures();
- for (request_id, failed_uuids) in failures {
- if let Some(uuid_to_blocks) = self.request_to_blocks.get(&request_id) {
- for uuid in failed_uuids {
- if let Some(block_ids) = uuid_to_blocks.get(&uuid) {
- for &block_id in block_ids {
- self.failed_block_ids.insert(block_id as u32);
- }
- tracing::warn!(
- request_id = %request_id,
- operation_id = %uuid,
- block_ids = ?block_ids,
- "load operation failed; marking blocks as failed"
- );
- }
- }
- }
- }
-
- // Return and clear the failed block IDs
- std::mem::take(&mut self.failed_block_ids)
- }
+ fn get_block_ids_with_load_errors(&mut self) -> HashSet<u32> {
+ self.absorb_failure_notifications();
+ // Return and clear the failed block IDs
+ std::mem::take(&mut self.failed_block_ids)
+ }Also applies to: 473-509
🤖 Prompt for AI Agents
In `@lib/bindings/kvbm/src/block_manager/vllm/connector/worker.rs` around lines
293 - 305, When enqueuing onboarding_operations we currently stash
request->blocks in request_to_blocks but later remove that mapping in
get_finished before we call get_block_ids_with_load_errors, which can drop the
failure->block correlation; fix by ensuring failures are drained/translated into
block IDs before cleanup: call get_block_ids_with_load_errors (or a shared
helper that drains scheduler failures and maps them to block IDs using
request_to_blocks) prior to removing the request's entry from request_to_blocks
in get_finished; factor the draining/mapping logic into a reusable helper (e.g.,
drain_and_map_load_failures) and use it from both the onboarding loop code paths
and the get_finished cleanup to preserve failure→block mappings.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Overview:
Add infrastructure to track and report block IDs that fail during asynchronous KV cache load operations in the KVBM connector. This enables vLLM to identify failed blocks and take corrective action (e.g., mark for recomputation).
Details:
block_idsfield toWorkerTransferRequestprotocol for tracking which blocks belong to each transfer operationfailed_block_idsfield andget_block_ids_with_load_errors()method toKvConnectorWorkerconnector_worker.py,dynamo_connector.py)failure_tx/failure_rxchannel betweenSchedulerandWorkerSchedulerClientto propagate transfer failuresdrain_failures()method to collect pending failure notificationsrequest_id→uuid→block_idsmapping in worker to convert failure notifications to block IDshandle_immediate_result()to send failure notifications whenresult.status.is_err()Where should the reviewer start?
lib/llm/src/block_manager/connector/scheduler.rs- Core failure notification channel implementationlib/bindings/kvbm/src/block_manager/vllm/connector/worker.rs- Block ID tracking and error aggregation logiclib/llm/src/block_manager/connector/protocol.rs- Protocol changes (block_idsfield)Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.