feat(llm): add pool-scoped KV DC Relay runtime

[Kaonael]

Summary

This PR formalizes the Dynamo-local pool model and per-pool runtime ownership for the KV DC Relay.

The central invariant is:

Each materializable serving endpoint maps to one atomic Dynamo PoolId. Each active pool owns one Relay actor, one CKF state, and one producer generation.

Independent endpoints or deployments remain separate pools even when they serve the same canonical model. A canonical model may therefore appear in multiple pool descriptors.
Comparing those pools and selecting between them remains consumer or global-router policy.

Details

Pool discovery and model projection

• Projects Dynamo ModelDeploymentCards into endpoint-centric pool memberships.
• Requires each materializable endpoint to resolve to:
  • one IndexerDomainId;
  • one canonical base-model identity;
  • one unambiguous KV-state endpoint.
• Derives PoolId from the resolved indexer domain and DC identity.
• Normalizes canonical model and alias registrations.
• Represents LoRA adapters explicitly while keeping them attached to the pool of their backing base model.
• Treats multiple domains, multiple base models, conflicting request-facing names, and duplicate live PoolId ownership as fail-closed materialization conflicts.
• Keeps invalid aliases and orphan adapter cards as soft discovery errors when the remaining endpoint projection is valid.
• Reconciles discovery periodically, rebinds all scoped streams after a stream terminates, and opens configured namespace queries concurrently.
• Suppresses duplicate discovery upserts before rebuilding the complete membership projection.

Pool runtime and producer contract

• Adds PoolRegistry as the owner of active pool actors, layout generations, reservations, and catalog publication.
• Exposes the producer-side contract through DcPoolCatalog and DcPoolDescriptor.
• Publishes the serving endpoint, ProducerIdentity, and normalized base/LoRA registrations for each active pool.
• Keeps derived cross-pool model and alias lookup indexes out of the Relay contract.
• Separates the backend-specific DRT instance identity from a random per-KvDcRelay::start Relay incarnation.
• Uses the Relay incarnation and monotonic layout generation to prevent producer identity reuse across Relay and pool lifetimes.
• Updates catalog entries incrementally and materializes their ordered representation lazily.
• Allocates and initializes CKF state through spawn_blocking, outside the registry mutex.
• Uses bounded concurrent CKF allocation, two-phase reservations, and RAII rollback so cancellation, shutdown, or allocation failure cannot leave stale reservations or publish
  incomplete pools.

Lifecycle, recovery, and fault handling

• Reuses Dynamo’s worker-query recovery model, with (worker_id, dp_rank) identifying a source and SourceEpoch fencing source replacement.
• Refreshes registration-only changes in place when the pool binding remains unchanged.
• Rebuilds the producer when the indexer domain, KV-state endpoint, or worker/rank topology changes.
• Withdraws a pool from DcPoolCatalog before waiting for recovery and actor teardown.
• Routes terminal actor and recovery failures through the same fenced-retirement path.
• Makes actor fault publication cancellation-aware and drains faults concurrently during teardown.
• Coalesces pending source faults by (worker_id, dp_rank), retains the newest epoch and strongest action, and keeps the pending state bounded.
• Gives producer-fencing faults priority over weaker recovery work; overflow also fails safe by fencing the producer.

Out of scope

This PR intentionally does not add:

• WAN transport or a global consumer;
• consumer-side model/alias lookup policy;

Where should the reviewer start?

1. components/src/dynamo/kv_dc_relay/README.md
   • The intended endpoint-scoped pool model.
2. lib/llm/src/kv_dc_relay/identity.rs
   • Relay identity, producer catalog, canonical registrations, and LoRA targets.
3. lib/llm/src/kv_dc_relay/discovery.rs
   • Projection from model cards into materializable endpoint memberships.
4. lib/llm/src/kv_dc_relay/pool_registry.rs
   • Actor ownership, allocation reservations, catalog publication, and retirement.
5. lib/llm/src/kv_dc_relay/host.rs
   • Discovery reconciliation, pool transitions, recovery, fencing, and fault coalescing.
6. lib/llm/src/kv_dc_relay/actor.rs and lib/llm/src/discovery/kv_source_membership.rs
   • Producer identity construction, cancellation-aware fault delivery, and binding-relevant source membership.

Validation

• cargo fmt --all -- --check
• cargo test -p dynamo-llm kv_dc_relay:: --lib
  • 65 tests passed.
• cargo test -p dynamo-llm --features ckf-diagnostics kv_dc_relay:: --lib
  • 66 tests passed.
• cargo test -p dynamo-llm kv_source_membership::tests --lib
  • 6 tests passed.
• cargo clippy -p dynamo-llm --all-targets --features ckf-diagnostics -- -D warnings
• git diff --check origin/main...HEAD

Related Issues

• Relates to DEP: Multi-DC KV-Aware Request Routing #11225

---

Summary by CodeRabbit

• New Features

  • Added pool-based KV cache relay management with contributor-aware ownership and lifecycle controls.
  • Added model identity, alias, routing, and pool-binding support.
  • Added configurable discovery across namespaces and endpoint prefixes.
  • Added publication subscriptions with incremental updates and snapshot recovery.
  • Added options to track publication activity without materializing payloads.

• Bug Fixes

  • Improved handling of conflicting memberships, stale sources, publication gaps, and departed endpoints.
  • Added validation for model identities, aliases, and format metadata.

• Documentation

  • Added an overview of KV cache relay pools and their runtime responsibilities.

[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]

🌿 Fern Docs Preview: https://nvidia-preview-019fae29-9a2f-7522-9a6a-e9a8e87388e7.docs.buildwithfern.com/dynamo/dev

[coderabbitai]

Walkthrough

The relay runtime now models CKF ownership by contributor and physical pool, discovers domain-slot memberships, manages shared pool actors through a registry, distributes recoverable publication updates, and updates host lifecycle, diagnostics, and configuration handling.

Changes

KV DC Relay runtime

Layer / File(s)	Summary
Contributor-scoped CKF state lib/kv-router/src/indexer/cuckoo.rs, lib/kv-router/src/indexer/cuckoo/dc.rs	CKF membership, removals, format validation, publication materialization, telemetry, and exports now use contributor-scoped source identities.
Relay identity and domain discovery lib/llm/src/kv_dc_relay.rs, lib/llm/src/kv_dc_relay/identity.rs, lib/llm/src/kv_dc_relay/discovery.rs, lib/llm/src/kv_dc_relay/resolution.rs, lib/llm/src/kv_dc_relay/README.md	Adds validated model and pool binding types, domain-based membership discovery, configurable query scopes, alias handling, conditional compilation updates, and runtime documentation.
Contributor-aware actor flow lib/llm/src/kv_dc_relay/actor.rs	Actor commands, recovery, source epochs, faults, resets, replacements, removals, and statistics now carry contributor-scoped identities.
Pool attachment and catalog lifecycle lib/llm/src/kv_dc_relay/pool_registry.rs	Adds pool creation, attachment and detachment, registration validation, catalog publication, fencing, fault routing, publication subscriptions, shutdown, and lifecycle tests.
Snapshot and delta publication hub lib/llm/src/kv_dc_relay/publication_hub.rs	Adds snapshot tracking, delta validation, subscriber fanout, lag recovery, retirement handling, and publication tests.
Pool-oriented host orchestration lib/llm/src/kv_dc_relay/host.rs	Reworks host slots and lifecycle management around domain slots and physical pools, including pool recovery, diagnostics, retired-slot reaping, and startup validation.

Estimated code review effort: 5 (Critical) | ~120 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 28.64% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: a pool-scoped KV DC Relay runtime.
Description check	✅ Passed	The PR description covers the required areas and includes a related issue link; section names differ slightly but the content is complete.

---

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

[coderabbitai]

Actionable comments posted: 5

🧹 Nitpick comments (8)

lib/llm/src/kv_dc_relay/host.rs (1)

745-759: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Failed registration refresh warns on every loop iteration.

On error active.registrations stays stale, so the mismatch persists and this branch re-runs on each SlotInput — including the periodic Health tick at Line 911 — emitting an identical warning indefinitely for a durable conflict (e.g. an alias claimed by another pool). Consider rate-limiting (first occurrence at warn, subsequent at debug, like the start_failures pattern at Line 856).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_dc_relay/host.rs` around lines 745 - 759, The registration
refresh error path in the loop around pools.replace_registrations must avoid
warning on every SlotInput/Health iteration. Add persistent failure tracking for
this refresh, following the existing start_failures pattern near the relevant
host logic: log the first failure at warn, subsequent identical failures at
debug, and reset the tracking after a successful replacement.

lib/llm/src/kv_dc_relay/pool_registry.rs (2)

592-601: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

publication_hubs is written but never read.

refresh_publication_hubs is called from attach, detach, replace_registrations, and fence, and the vector is cleared in shutdown, but nothing ever reads self.publication_hubs. Every mutation also sorts and clones hubs while holding the registry state lock. If this is staged for the next layer, mark it explicitly (like PoolCatalog at Line 101) or drop it until a consumer exists.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_dc_relay/pool_registry.rs` around lines 592 - 601, The
publication_hubs field updated by refresh_publication_hubs is currently unused;
either mark it explicitly as intentionally staged, matching the existing
PoolCatalog convention, or remove publication_hubs, refresh_publication_hubs,
and their mutation/cleanup calls until a consumer is implemented. Prefer
removing the unused sorting and hub cloning if no near-term consumer is
required.

---

808-840: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Cross-pool validation is O(pools × contributors × registrations) per registration.

For each incoming registration the loop rescans every contributor's registrations in every pool. With a large fleet this makes attach/replace_registrations quadratic while holding the global state mutex. A HashMap<CanonicalModelId, ModelTarget> maintained alongside canonical_refcounts would make the target-conflict check O(1).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_dc_relay/pool_registry.rs` around lines 808 - 840, Update
validate_registrations to avoid traversing every pool, contributor, and existing
registration for each incoming registration. Maintain and reuse a
HashMap<CanonicalModelId, ModelTarget> alongside canonical_refcounts in
PoolRegistryState, and perform existing-registration target conflict checks
through that index in O(1) per registration while preserving duplicate and
conflict errors for both the current attachment and prior registrations.

lib/llm/src/kv_dc_relay/publication_hub.rs (1)

163-196: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

subscribe/snapshot are async without awaiting.

Both take a synchronous parking_lot lock and never await. Making them sync keeps the API honest and avoids callers assuming a yield point.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_dc_relay/publication_hub.rs` around lines 163 - 196, Make
PublicationHub::subscribe and the test-only snapshot method synchronous by
removing async from their signatures. Update all call sites and tests to stop
awaiting these methods while preserving their existing locking, readiness
checks, and return values.

lib/llm/src/kv_dc_relay/discovery.rs (3)

841-866: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

list_queries and open_query_streams serialize one round trip per namespace.

Both run sequentially, so startup and every 30s reconcile cost N sequential discovery calls. futures::future::try_join_all would bound this to a single round trip; worth doing if the namespace list is expected to be more than a handful.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_dc_relay/discovery.rs` around lines 841 - 866, Update
list_queries and open_query_streams to initiate discovery operations for all
queries concurrently using futures::future::try_join_all rather than awaiting
each request inside the loop. Preserve error propagation and aggregate the
returned instances or DiscoveryStream values into the existing result types.

---

1181-1208: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

No test covers the ambiguous-name Binding conflict path.

Lines 618-657 are the most intricate new logic — one request-facing name resolving to two different BindingIdentity targets must produce a MembershipConflict::Binding and drop the claim — and it also has to stay in step with validate_registrations in pool_registry.rs. A case where an alias on one card collides with a canonical model name on another would lock that contract in.

Want me to draft that test?

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_dc_relay/discovery.rs` around lines 1181 - 1208, Add a test
covering the ambiguous-name Binding conflict path in the discovery logic around
the existing alias tests: create two cards where an alias on one card matches
the canonical model name on the other, apply both registrations, and assert the
shared request-facing name produces MembershipConflict::Binding with the claim
dropped. Keep the expectations aligned with validate_registrations in
pool_registry.rs.

---

483-503: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Adapter cards with no backing base card disappear without any diagnostic.

When worker_bases is empty, worker_domain_ids is empty too, so the conflict loop at Line 491 iterates zero times and the adapter is dropped silently — no conflict entry, no log. That is the likely steady state while a base card is still propagating or has just been removed, so a debug log (or a conflict on every domain of the endpoint) would make the transient visible in membership_conflicts.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_dc_relay/discovery.rs` around lines 483 - 503, The conflict
handling around worker_domain_ids must preserve diagnostics when worker_bases is
empty instead of silently dropping the adapter. Update the branch using
worker_domain_ids, ambiguous_workers, and conflicting_worker_facts so the
no-backing-base case records a MembershipConflict::Card for an appropriate
endpoint domain (or otherwise emits the requested debug diagnostic), while
retaining existing conflict behavior for populated domain IDs.

lib/kv-router/src/indexer/cuckoo/dc.rs (1)

828-836: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Leftover "worker" vocabulary after the source-keyed rewrite. The membership key changed from WorkerWithDpRank to DcCkfSourceId, but two surfaces still speak in worker terms, which is exactly the aliasing this PR eliminates.

• lib/kv-router/src/indexer/cuckoo/dc.rs#L828-L836: rename the worker: DcCkfSourceId parameter on store (and the matching one on remove/remove_member) to source, matching remove_source and source_counts.
• lib/kv-router/src/indexer/cuckoo/dc.rs#L793-L801: document member_counts as the legacy unscoped view, or fold duplicates, since two contributors reusing a worker/rank now collapse to identical WorkerWithDpRank entries.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/kv-router/src/indexer/cuckoo/dc.rs` around lines 828 - 836, The
source-keyed rewrite still exposes worker terminology and an ambiguous
member_counts view. In lib/kv-router/src/indexer/cuckoo/dc.rs lines 828-836,
rename the store parameter and corresponding remove/remove_member parameters
from worker to source, updating their uses to match remove_source and
source_counts; in lines 793-801, document member_counts as the legacy unscoped
view or fold duplicate entries caused by contributors sharing a worker/rank.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/llm/src/kv_dc_relay/discovery.rs`:
- Around line 702-720: Remove the per-view pruning of self.domain_generations
after rebuilding domains so counters remain monotonic for slots that temporarily
disappear and later reappear. Keep generation updates in the existing changed
branch and preserve the current cleanup of warned_invalid_models and
warned_invalid_aliases.
- Around line 49-81: Normalize each namespace and endpoint prefix by trimming
surrounding whitespace before validation, deduplication, containment checks, and
any downstream query construction. Update the validation logic in the relevant
discovery configuration method so `" prod"` and `"prod"` are treated
identically, while preserving the existing empty, duplicate, and
namespace-boundary errors.
- Around line 736-737: Update the stream handling around open_query_streams so
termination of any individual query stream causes the outer discovery/watch loop
to break and rebind all streams. Wrap or otherwise monitor each substream before
it is consumed by SelectAll, while preserving normal message handling and
cancellation behavior.

In `@lib/llm/src/kv_dc_relay/identity.rs`:
- Around line 277-321: Update the pool helper and
pool_identity_sources_report_derived_components test to construct a fully
explicit pool with an Explicit cache-semantics source as well as an Explicit
routing source, then assert is_derived() is false for that pool. Preserve the
existing derived-source assertion using DefaultDerived cache semantics and
routing.

In `@lib/llm/src/kv_dc_relay/resolution.rs`:
- Around line 50-55: Gate the PoolBinding and EndpointLocator types, along with
their test-only constructors and accessor methods, behind #[cfg(test)] in the
relay resolution module. Keep production relay code unchanged and ensure these
helpers remain available to tests without dead-code warnings.

---

Nitpick comments:
In `@lib/kv-router/src/indexer/cuckoo/dc.rs`:
- Around line 828-836: The source-keyed rewrite still exposes worker terminology
and an ambiguous member_counts view. In lib/kv-router/src/indexer/cuckoo/dc.rs
lines 828-836, rename the store parameter and corresponding remove/remove_member
parameters from worker to source, updating their uses to match remove_source and
source_counts; in lines 793-801, document member_counts as the legacy unscoped
view or fold duplicate entries caused by contributors sharing a worker/rank.

In `@lib/llm/src/kv_dc_relay/discovery.rs`:
- Around line 841-866: Update list_queries and open_query_streams to initiate
discovery operations for all queries concurrently using
futures::future::try_join_all rather than awaiting each request inside the loop.
Preserve error propagation and aggregate the returned instances or
DiscoveryStream values into the existing result types.
- Around line 1181-1208: Add a test covering the ambiguous-name Binding conflict
path in the discovery logic around the existing alias tests: create two cards
where an alias on one card matches the canonical model name on the other, apply
both registrations, and assert the shared request-facing name produces
MembershipConflict::Binding with the claim dropped. Keep the expectations
aligned with validate_registrations in pool_registry.rs.
- Around line 483-503: The conflict handling around worker_domain_ids must
preserve diagnostics when worker_bases is empty instead of silently dropping the
adapter. Update the branch using worker_domain_ids, ambiguous_workers, and
conflicting_worker_facts so the no-backing-base case records a
MembershipConflict::Card for an appropriate endpoint domain (or otherwise emits
the requested debug diagnostic), while retaining existing conflict behavior for
populated domain IDs.

In `@lib/llm/src/kv_dc_relay/host.rs`:
- Around line 745-759: The registration refresh error path in the loop around
pools.replace_registrations must avoid warning on every SlotInput/Health
iteration. Add persistent failure tracking for this refresh, following the
existing start_failures pattern near the relevant host logic: log the first
failure at warn, subsequent identical failures at debug, and reset the tracking
after a successful replacement.

In `@lib/llm/src/kv_dc_relay/pool_registry.rs`:
- Around line 592-601: The publication_hubs field updated by
refresh_publication_hubs is currently unused; either mark it explicitly as
intentionally staged, matching the existing PoolCatalog convention, or remove
publication_hubs, refresh_publication_hubs, and their mutation/cleanup calls
until a consumer is implemented. Prefer removing the unused sorting and hub
cloning if no near-term consumer is required.
- Around line 808-840: Update validate_registrations to avoid traversing every
pool, contributor, and existing registration for each incoming registration.
Maintain and reuse a HashMap<CanonicalModelId, ModelTarget> alongside
canonical_refcounts in PoolRegistryState, and perform existing-registration
target conflict checks through that index in O(1) per registration while
preserving duplicate and conflict errors for both the current attachment and
prior registrations.

In `@lib/llm/src/kv_dc_relay/publication_hub.rs`:
- Around line 163-196: Make PublicationHub::subscribe and the test-only snapshot
method synchronous by removing async from their signatures. Update all call
sites and tests to stop awaiting these methods while preserving their existing
locking, readiness checks, and return values.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5e9ec6d4-fcbf-44ff-a4ca-859b199d25c6

📥 Commits

Reviewing files that changed from the base of the PR and between 769ee07 and 81f9a6c.

📒 Files selected for processing (11)

• lib/kv-router/src/indexer/cuckoo.rs
• lib/kv-router/src/indexer/cuckoo/dc.rs
• lib/llm/src/kv_dc_relay.rs
• lib/llm/src/kv_dc_relay/README.md
• lib/llm/src/kv_dc_relay/actor.rs
• lib/llm/src/kv_dc_relay/discovery.rs
• lib/llm/src/kv_dc_relay/host.rs
• lib/llm/src/kv_dc_relay/identity.rs
• lib/llm/src/kv_dc_relay/pool_registry.rs
• lib/llm/src/kv_dc_relay/publication_hub.rs
• lib/llm/src/kv_dc_relay/resolution.rs

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.

[datadog-official]

🎯 Code Coverage (details)
• Patch Coverage: 100.00%
• Overall Coverage: 42.40% (-14.00%)

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: e6ec122 | Docs | Datadog PR Page | Give us feedback!}

[PeaBrane]

/ok to test 6e06e64

[PeaBrane]

Suggestion on keeping the producer/consumer catalog boundary minimal.

[PeaBrane]

Neutral architectural follow-up—not a blocking verdict. The common pattern in the remaining findings is that PoolRegistry state currently serves both as local actor ownership and as externally visible “active catalog” truth. A useful invariant would be: every advertised descriptor names one unique, currently active producer generation. Concretely, give each Relay start a unique incarnation, publish or refresh catalog state only after validating the binding transition, and withdraw a generation before draining it. Keep consumer routing indexes derived from the pool-scoped catalog as discussed separately. This adds explicit activation/withdrawal steps, but centralizes lifecycle consistency now, before consumers depend on the contract.

[PeaBrane]

/ok to test 6ab1c26

[hhzhang16]

/ok to test e6ec122