feat: non blocking radix snapshot upload

[PeaBrane]

Overview:

titled

Summary by CodeRabbit

• Refactor
  • Snapshot uploads are now performed asynchronously in the background instead of blocking the caller.
  • Upload errors are now logged internally rather than propagating as failures.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Walkthrough

The purge_then_snapshot function in the NATS object store interaction now executes snapshot uploads asynchronously in a background task via tokio::spawn instead of blocking the caller. Error handling shifts from propagating errors to logging warnings within the spawned task.

Changes

Cohort / File(s)	Summary
Async snapshot upload refactor lib/llm/src/kv_router/subscriber.rs	Moved snapshot upload from blocking call to asynchronous background task; computed event_count before spawn and captured bucket_name for logging; changed error handling from error propagation to warning logs within spawned task; caller no longer awaits result.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

• Verify the safety of spawning background task without awaiting completion and whether losing immediate error propagation is acceptable in this context
• Ensure bucket_name and event_count captures are correct within the spawned closure
• Confirm that logging warnings on failure provides sufficient observability for operational monitoring

Poem

🐰 Uploads now fly swift on winds of spawn,
No wait, no block—the task moves on!
Errors logged where shadows creep,
The caller rests in async sleep. 🌙

Pre-merge checks

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The pull request description is severely incomplete, containing only 'titled' under the Overview section with no Details, reviewer guidance, or related issue references.	Complete the description by adding Details explaining the non-blocking implementation, specify which files to review (kv_router/subscriber.rs), and reference any related GitHub issues using action keywords.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: making radix snapshot upload non-blocking, which aligns with the code modification that moves snapshot upload to an async background task.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

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)

lib/llm/src/kv_router/subscriber.rs (1)

406-412: Success log is misleading—snapshot upload is still in progress.

Line 410 logs "Successfully performed purge and snapshot" immediately after purge_then_snapshot returns Ok(()), but the snapshot upload is now asynchronous and may not have completed (or could still fail). This gives operators false confidence that the snapshot succeeded.

Update the log to reflect what actually completed:

-                    Ok(_) => tracing::info!("Successfully performed purge and snapshot"),
+                    Ok(_) => tracing::info!("Successfully performed purge and initiated snapshot upload"),
                     Err(e) => tracing::debug!("Could not perform purge and snapshot: {e:?}"),

The actual upload success is now logged inside the spawned task (line 208), so operators will see both messages if the upload succeeds.

🧹 Nitpick comments (1)

lib/llm/src/kv_router/subscriber.rs (1)

208-211: Consider logging upload-specific timing.

The elapsed time includes the entire operation from line 134 (purge, dump request, spawn overhead) rather than just the upload duration. If uploads queue up or there's spawn delay, this timing could be misleading.

Optionally measure upload-specific duration:

+        let upload_start = std::time::Instant::now();
         if let Err(e) = nats_client.object_store_upload_data(&events, &url).await {
             tracing::warn!("Failed to upload snapshot: {e:?}");
             return;
         }

         tracing::info!(
-            "Successfully uploaded snapshot with {event_count} events to bucket {bucket_name} in {}ms",
-            start_time.elapsed().as_millis()
+            "Successfully uploaded snapshot with {event_count} events to bucket {bucket_name} in {}ms (total operation: {}ms)",
+            upload_start.elapsed().as_millis(),
+            start_time.elapsed().as_millis()
         );

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c9e445a and 62ed8a4.

📒 Files selected for processing (1)

• lib/llm/src/kv_router/subscriber.rs (1 hunks)

🧰 Additional context used

🧠 Learnings (10)

📓 Common learnings

Learnt from: oandreeva-nv
Repo: ai-dynamo/dynamo PR: 2989
File: lib/llm/src/block_manager/distributed/transfer.rs:6-6
Timestamp: 2025-09-18T21:47:44.143Z
Learning: For PR ai-dynamo/dynamo#2989, the ConnectorTransferBatcher architectural issues will be addressed in a follow-up PR by removing the duplicate batching logic and integrating distributed transfers with the existing TransferBatcher + LocalTransferManager pipeline, rather than adding bounded concurrency primitives like Semaphore.

Learnt from: kthui
Repo: ai-dynamo/dynamo PR: 1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.

Learnt from: ryanolson
Repo: ai-dynamo/dynamo PR: 1093
File: lib/llm/src/block_manager/block/registry.rs:98-122
Timestamp: 2025-05-29T06:20:12.901Z
Learning: In lib/llm/src/block_manager/block/registry.rs, the background task spawned for handling unregister notifications uses detached concurrency by design. The JoinHandle is intentionally not stored as this represents a reasonable architectural tradeoff for a long-running cleanup task.

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3077
File: lib/llm/src/kv_router/subscriber.rs:334-336
Timestamp: 2025-09-17T01:00:50.937Z
Learning: PeaBrane suggested using tokio::select! arm ordering with the existing biased directive in the indexer to create a natural barrier for dump requests, ensuring KV events are drained before snapshotting. This approach leverages existing architecture (biased select) to solve race conditions with minimal code changes, which aligns with their preference for contained solutions.

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3077
File: lib/llm/src/kv_router/subscriber.rs:334-336
Timestamp: 2025-09-17T01:00:50.937Z
Learning: PeaBrane suggested using tokio::select! arm ordering with the existing biased directive in the indexer to create a natural barrier for dump requests, ensuring KV events are drained before snapshotting. This approach leverages existing architecture (biased select) to solve race conditions with minimal code changes, which aligns with their preference for contained solutions.

Learnt from: oandreeva-nv
Repo: ai-dynamo/dynamo PR: 1195
File: lib/llm/tests/block_manager.rs:150-152
Timestamp: 2025-06-02T19:37:27.666Z
Learning: In Rust/Tokio applications, when background tasks use channels for communication, dropping the sender automatically signals task termination when the receiver gets `None`. The `start_batching_publisher` function in `lib/llm/tests/block_manager.rs` demonstrates this pattern: when the `KVBMDynamoRuntimeComponent` is dropped, its `batch_tx` sender is dropped, causing `rx.recv()` to return `None`, which triggers cleanup and task termination.

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3184
File: docs/architecture/kv_cache_routing.md:70-73
Timestamp: 2025-09-23T20:08:37.105Z
Learning: PeaBrane prefers to keep documentation diagrams simplified to avoid visual overload, even when this means sacrificing some technical precision for the sake of clarity and comprehension. They prioritize pedagogical effectiveness over exhaustive technical detail in architectural diagrams.

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 2756
File: lib/llm/src/kv_router/subscriber.rs:36-44
Timestamp: 2025-08-29T10:03:48.330Z
Learning: PeaBrane prefers to keep PRs contained in scope and is willing to defer technical improvements to future PRs when the current implementation works for the immediate use case. They acknowledge technical debt but prioritize deliverability over completeness in individual PRs.

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3095
File: lib/llm/src/kv_router/subscriber.rs:200-223
Timestamp: 2025-09-17T20:55:41.416Z
Learning: In the dynamo codebase, PeaBrane prefers to maintain consistency with existing etcd key parsing patterns (like splitting on '/' and parsing the last segment) rather than introducing more robust parsing approaches, even when the current approach might be brittle, to keep the codebase aligned and avoid divergent patterns.

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3095
File: lib/llm/src/kv_router/indexer.rs:0-0
Timestamp: 2025-09-17T20:55:06.333Z
Learning: When PeaBrane encounters a complex implementation issue that would significantly expand PR scope (like the remove_worker_sender method in lib/llm/src/kv_router/indexer.rs that required thread-safe map updates and proper shard targeting), they prefer to remove the problematic implementation entirely rather than rush a partial fix, deferring the proper solution to a future PR.

📚 Learning: 2025-05-29T06:20:12.901Z

Learnt from: ryanolson
Repo: ai-dynamo/dynamo PR: 1093
File: lib/llm/src/block_manager/block/registry.rs:98-122
Timestamp: 2025-05-29T06:20:12.901Z
Learning: In lib/llm/src/block_manager/block/registry.rs, the background task spawned for handling unregister notifications uses detached concurrency by design. The JoinHandle is intentionally not stored as this represents a reasonable architectural tradeoff for a long-running cleanup task.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-05-29T00:02:35.018Z

Learnt from: alec-flowers
Repo: ai-dynamo/dynamo PR: 1181
File: lib/llm/src/kv_router/publisher.rs:379-425
Timestamp: 2025-05-29T00:02:35.018Z
Learning: In lib/llm/src/kv_router/publisher.rs, the functions `create_stored_blocks` and `create_stored_block_from_parts` are correctly implemented and not problematic duplications of existing functionality elsewhere in the codebase.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-09-17T01:00:50.937Z

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3077
File: lib/llm/src/kv_router/subscriber.rs:334-336
Timestamp: 2025-09-17T01:00:50.937Z
Learning: PeaBrane suggested using tokio::select! arm ordering with the existing biased directive in the indexer to create a natural barrier for dump requests, ensuring KV events are drained before snapshotting. This approach leverages existing architecture (biased select) to solve race conditions with minimal code changes, which aligns with their preference for contained solutions.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-09-17T01:00:50.937Z

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3077
File: lib/llm/src/kv_router/subscriber.rs:334-336
Timestamp: 2025-09-17T01:00:50.937Z
Learning: PeaBrane identified that reordering tokio::select! arms in the indexer (moving dump_rx.recv() to be after event_rx.recv()) creates a natural barrier that ensures RouterEvents are always processed before dump requests, solving the ack-before-commit race condition. This leverages the existing biased directive and requires minimal code changes, aligning with their preference for contained solutions.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-06-02T19:37:27.666Z

Learnt from: oandreeva-nv
Repo: ai-dynamo/dynamo PR: 1195
File: lib/llm/tests/block_manager.rs:150-152
Timestamp: 2025-06-02T19:37:27.666Z
Learning: In Rust/Tokio applications, when background tasks use channels for communication, dropping the sender automatically signals task termination when the receiver gets `None`. The `start_batching_publisher` function in `lib/llm/tests/block_manager.rs` demonstrates this pattern: when the `KVBMDynamoRuntimeComponent` is dropped, its `batch_tx` sender is dropped, causing `rx.recv()` to return `None`, which triggers cleanup and task termination.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-06-13T22:07:24.843Z

Learnt from: kthui
Repo: ai-dynamo/dynamo PR: 1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-10-14T00:58:05.744Z

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3597
File: lib/llm/src/kv_router/indexer.rs:437-441
Timestamp: 2025-10-14T00:58:05.744Z
Learning: In lib/llm/src/kv_router/indexer.rs, when a KvCacheEventData::Cleared event is received, the system intentionally clears all dp_ranks for the given worker_id by calling clear_all_blocks(worker.worker_id). This is the desired behavior and should not be scoped to individual dp_ranks.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-09-17T20:55:06.333Z

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3095
File: lib/llm/src/kv_router/indexer.rs:0-0
Timestamp: 2025-09-17T20:55:06.333Z
Learning: When PeaBrane encounters a complex implementation issue that would significantly expand PR scope (like the remove_worker_sender method in lib/llm/src/kv_router/indexer.rs that required thread-safe map updates and proper shard targeting), they prefer to remove the problematic implementation entirely rather than rush a partial fix, deferring the proper solution to a future PR.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

📚 Learning: 2025-09-03T19:31:32.621Z

Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 2840
File: lib/llm/src/kv_router/sequence.rs:86-88
Timestamp: 2025-09-03T19:31:32.621Z
Learning: PeaBrane chose to defer fixing the corner case where a single late-arriving request might never expire in the ActiveSequences expiry mechanism (lib/llm/src/kv_router/sequence.rs). They prefer to avoid adding a background loop for periodic cleanup at this time, accepting the technical debt to keep the current PR scope contained.

Applied to files:

• lib/llm/src/kv_router/subscriber.rs

🧬 Code graph analysis (1)

lib/llm/src/kv_router/subscriber.rs (1)

lib/runtime/src/transports/nats.rs (1)

• addr (72-75)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)

• GitHub Check: trtllm (arm64)
• GitHub Check: operator (amd64)
• GitHub Check: Build and Test - dynamo
• GitHub Check: tests (launch/dynamo-run)
• GitHub Check: clippy (.)
• GitHub Check: tests (lib/bindings/python)
• GitHub Check: clippy (lib/bindings/python)
• GitHub Check: tests (.)
• GitHub Check: tests (lib/runtime/examples)
• GitHub Check: clippy (launch/dynamo-run)