[surajssd]

Overview:

KvRouter::new hands the scheduler a strong Arc to the RequestLeaseManager as its replica observer, and the manager holds the scheduler's SchedulerBookingCleanup handle that it receives from set_scheduler. Both edges are strong, so neither side can drop once construction has wired them. A router that fails to construct after that point, or one that is dropped, leaks its scheduler actor and lease state. This PR replaces the strong observer with a weak one so the cycle is broken while callbacks still reach the manager for as long as it is alive.

This is the first of three independent PRs split out of #14965. That PR's Mooncake Store test for failed router construction depends on RequestLeaseManager::replica_observer from this change. The leak itself exists on main today and is not specific to Mooncake Store.

Details:

• Add RequestLeaseManager::replica_observer, which returns an Arc<dyn ReplicaRequestLeaseObserver> backed by a Weak<RequestLeaseManagerInner>.
• Replace the impl ReplicaRequestLeaseObserver for RequestLeaseManager with a private WeakRequestLeaseObserver. Each of admitted, progressed, and completed upgrades the weak reference and forwards to the existing register_remote, touch_booking, and complete_remote methods, or becomes a no-op once the manager is gone.
• KvRouter::new passes request_leases.replica_observer() to the scheduler instead of Some(Arc::new(request_leases.clone())).
• Add the regression test replica_observer_forwards_live_callbacks_without_retaining_manager. It admits a booking on a tracked router, drives the three callbacks through the observer and checks the lease record at each step, drops the router, and asserts the manager is released. It then calls the callbacks again on the dangling observer to confirm they are no-ops.
• Make the tracked_router test helper in kv_router.rs pub(super) so the request_lease test module can reuse it.

Where should the reviewer start?

• lib/llm/src/kv_router/request_lease.rs: replica_observer and WeakRequestLeaseObserver, then the new test at the top of the tests module.
• lib/llm/src/kv_router.rs: the one-line change in KvRouter::new where the observer is handed to the scheduler.

Validation

cargo test -p dynamo-llm --no-default-features --lib -- kv_router::request_lease

Both tests in the module pass. The regression test was also run against the previous strong observer in a throwaway worktree, where the manager.upgrade().is_none() assertion fails, so the test does detect the cycle. cargo fmt --all --check is clean.

Related Issues

🔗 This PR is linked to an issue:

• Relates to feat(kv-router): add Mooncake Store shared-cache routing #14965

Summary by CodeRabbit

• Bug Fixes
  • Improved reliability when replica activity updates arrive after request tracking has ended. These late updates are now safely ignored, while updates for active requests continue to be recorded and completed as expected.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

👋 Hi surajssd! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: ai-dynamo/dynamo/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c11153fa-9239-476c-8e41-05ccfa4ea400

📥 Commits

Reviewing files that changed from the base of the PR and between a76e12e and c6c92af.

📒 Files selected for processing (2)

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

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The router now passes a weak-reference-backed replica observer to EmbeddedSelection. The observer forwards replica callbacks only while the request-lease manager exists. An async test checks lease updates and callbacks after manager destruction.

Changes

Replica Lease Observer

Layer / File(s)	Summary
Observer wiring and lifecycle lib/llm/src/kv_router/request_lease.rs, lib/llm/src/kv_router.rs	The request-lease manager creates a weak observer for replica callbacks, and the router passes it to EmbeddedSelection. The test checks lease registration, clock updates, completion, and callbacks after manager destruction. The tracked_router test helper is now visible to sibling modules.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to c6c92

The observer change appears ready to merge after normal checks.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 45.45% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description includes all required sections: Overview, Details, reviewer guidance, and Related Issues. It identifies issue #14965 and provides implementation details and validation results.
Title check	✅ Passed	The title clearly and concisely describes the main change: breaking the request-lease observer ownership cycle in the KV router.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

---

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

[jthomson04]

Reviewed the split from #14965. The weak observer breaks the ownership cycle and preserves live callbacks. The change is self-contained. No code findings from source review.

[jthomson04]

/ok to test c6c92af

[jthomson04]

/ok to test 8737d74