feat(kv-router): DisaggCoordinator with prefill-anchored and decode-anchored orders by ishandhanani · Pull Request #40 · ishandhanani/dynamo

ishandhanani · 2026-09-03T22:33:34Z

Stack 2 (disaggregation coordinator), first PR. Stacked on #53, the tip of the upstream-shaped Stack 1 staging series (#45 to #53), rebased onto ai-dynamo/main. Spec: DEP #33, "DisaggCoordinator spec" and "Stack 1 DEP r2" comments.

Summary

Three logical commits:

feat(kv-router): DisaggCoordinator with linked-booking compensation — feat(kv-router): DisaggCoordinator with linked-booking compensation #20 brought onto this stack unchanged: the runtime-free coordinator over two SelectionPools with the LinkedBookingState machine and decode-anchored planning.
feat(kv-router): prefill-anchored order for DisaggCoordinator — the spec's two orders. CoordinationOrder::PrefillAnchored is what the frontend PrefillRouter does today: plan_prefill books prefill KV-aware on the prompt (honoring pin, session-affinity target, allow-set), plan_decode books decode load-only inside the prefill worker's KV-transfer domain whenever the host is ready (before or after prefill_complete), release frees what is left. DecodeAnchored is the existing preview/bypass/commit flow. plan() runs the configured order end to end for hosts that book both at once. DisaggPlan carries optional prefill and decode bookings behind accessors plus its request; one state machine, one release, one KV-transfer constraint helper derived from whichever worker anchors. New state PrefillCommitted; compensation table per the spec, covered by focused tests for both orders (11 coordinator tests).
feat(epp): disaggregated coordination over two pools in the standalone EPP — feat(epp): decode-first disaggregated coordination in the standalone EPP #26 brought onto this stack and adapted: the EPP hosts the coordinator over a prefill and a decode SelectionService; DYN_EPP_DISAGG_ORDER selects decode-anchored (default, previous behavior) or prefill-anchored. Reflector feeding goes through the group B WorkerCatalogSource path.

Not in this PR (spec deliverable 2)

Moving the frontend PrefillRouter onto plan_prefill / prefill_complete / plan_decode / release with no behavior change. Inventory from reading the frontend path for this work:

The frontend fuses prefill selection, booking, and dispatch in RoutingHost::select_and_dispatch_prefill, and its migration retry re-selects inside the dispatch loop; decode goes through the decode RoutingHost's own generate with the KV-transfer constraints merged into the request. Bookings live in each KvRouter's embedded partition and are released by RequestGuard on stream end.
So the coordinator can own the linked bookkeeping and the constraint derivation, but the frontend's SelectionPool implementation has to be the partition-backed plan/dispatch API (preview_kv_route / plan_kv_route_from_preview / dispatch_kv_plan, plus a prefill counterpart) rather than select_and_reserve, or booking and dispatch would separate and migration semantics would change.
Gates per the spec: prefill-router unit tests, tests/router disaggregated e2e identical before/after, replay byte parity on the two disaggregated rows, host disaggregated smoke with engine kill. That is its own PR.

Validation

cargo test -p dynamo-kv-router --features standalone-selection --lib -- coordinator: 11 passed.
cargo test -p dynamo-ext-proc: 142 passed (includes the two-pool coordinator test through the new reflector source).
cargo clippy for dynamo-kv-router (feature standalone-selection) and dynamo-ext-proc (tests): clean.

Review notes (graham-code-review pass)

DisaggPlan keeps its DisaggRequest so plan_decode needs no second copy of the prompt from the host; the decode-anchored path moves the request into the plan once (no clone). The EPP's decode_booking helper carries one expect whose invariant (plan() books decode in both orders) is stated at the call.
CoordinatorError messages no longer say "in the decode pool" because the same constraint helper now runs against either anchor.
EPP default order stays decode-anchored so the fork's existing two-pool behavior is unchanged; prefill-anchored is opt-in until the frontend migration lands and the two can be compared.

The frontend router now always schedules on an embedded selection-service partition. Remove `KvScheduler` (kv_router/scheduler.rs keeps only its runtime-free re-exports), the runtime sequence publisher and its constructors, the direct-ZMQ active-sequence ingress, the two-variant `RouterScheduler` dispatch, and the LoRA load estimator's uncalled scheduler polling. `KvRouter<Sel>` keeps its type parameter for API compatibility (`PhantomData<fn() -> Sel>`). A caller-injected `WorkerSelector<ModelRuntimeConfig>` that is neither the default selector nor a registry `WorkerSelectionPolicy` now runs inside the partition: `WorkerSelectionPolicy::delegated` (new `DelegatedWorkerSelector` trait in dynamo-kv-router) hands the selector the scheduling request and a config-erased view of the eligible workers (`ErasedWorkerConfig`), which the embedded host projects back onto `ModelRuntimeConfig`. The unit tests that inject fake selectors run unchanged on that path. The per-worker `router_worker_registered` gauge is now maintained by the embedded catalog feeder. `router_embedded_selection` is deprecated and ignored; `false` logs a warning. Four sequence-level tests that built the runtime sequences directly are removed; their behavior is covered by the replica-lifecycle and session-affinity e2e tests on the embedded path. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com> (cherry picked from commit 71cea67)

The embedded selection partition is the only scheduler, so the flag has been a no-op since the runtime scheduler was deleted. It never shipped in a release, so no wire compatibility shim is needed: drop it from KvRouterConfig (Rust and serde), the DYN_ROUTER_EMBEDDED_SELECTION environment override, the Python argument group, the bindings signature, and the frontend configuration reference. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

…router KvRouter, RoutingHost, RoutePlan, RequestGuard, PrefillRouter, ManagedKvRouter, ModelWatcher, and PreprocessedRouting were generic over a caller-injected WorkerSelector type that only ever named the policy the embedded partition would run. Replace the parameter with SelectionPolicySource: Registry (the linked policy the router config names, else Dynamo's default) or Factory (an explicit WorkerSelectionPolicyFactory). The router resolves it once at construction and instantiates one policy to learn the worker inputs it needs, so shared-cache and indexer gating are unchanged. ModelManager constructors take the policy source in place of a selector (kv_chooser_for_with_policy[_and_client], managed_kv_router_for_with_policy); HttpFrontend, ModelWatcher, and PrefillRouter carry it instead of a per-model selector factory. The bindings' custom-policy feature no longer needs a separate router type. The delegating-selector adapter in the embedded backend is removed with its only caller. Unit tests that injected fake selectors now inject picker or filter policies through the public worker-selection API; the overload test becomes a filtered-workers test because a policy cannot report overload directly. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

DelegatedWorkerSelector, ErasedWorkerConfig, and the Delegated policy state existed only so the frontend could run an arbitrary caller-injected WorkerSelector type on a partition. That path is gone with the selector type parameter; custom policies use WorkerSelectionPolicy's filter, scorer, and picker API like every other host. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

…ionHost SelectionHostHooks was a flat bag of eight optional fields. Regroup it by purpose as SelectionHost { load, cache, eligibility, telemetry, replication }, each group defaulting to the standalone service's behavior so a host sets only the groups it owns. The builder method becomes SelectionServiceBuilder::host. No behavior change; the frontend's embedded partition sets the same values. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

Phase 4 of the runtime-optional frontend roadmap: the EPP pod reflector is the Kubernetes catalog feeder, and it registered every pod as a single data-parallel rank (the hardcoded data_parallel_size: 1 in selector.rs). - EppStandaloneConfig gains DYN_EPP_DATA_PARALLEL_SIZE (default 1) and DYN_EPP_KV_EVENT_PORT_STRIDE (default 1). Rank r publishes KV events on kv_event_port + r * stride; the reflector emits one KV-event endpoint per rank (RawWorker.kv_events_endpoints) and the registration carries data_parallel_size, which the selector passes through with data_parallel_start_rank 0. Tests: strided per-rank endpoints for a 3-rank pod; existing reflector, adapter, and selector tests updated to the per-rank shape (144 pass). Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

… membership A host owns discovery and exposes it as a WorkerCatalogSource: a stream of complete desired membership snapshots. CatalogReconciler applies each snapshot to a SelectionCore: it upserts new or changed records, retries records the core reported not schedulable, deletes workers that left (including ones that never converged), rejects duplicate ids before mutating anything, and clears the catalog when the source closes so selection fails closed. A CatalogObserver lets a host keep per-worker metrics in step. This is the reconcile loop the EPP standalone selector and the frontend's embedded partition each carried; both switch to it in the next commits. SelectionService exposes its core for hosts that drive the catalog directly. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

…ogReconciler The runtime-config watch becomes a WorkerCatalogSource (RuntimeDiscoverySource) and the partition's catalog is driven by the shared CatalogReconciler instead of the router-local CatalogFeeder. The per-worker router_worker_registered gauge is kept through a CatalogObserver. Behavior is unchanged except that a worker the core reports not schedulable is now retried on every membership snapshot rather than only when its runtime config changes. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

The standalone EPP's topology adapter yields the reflector's Ready pods as WorkerRequest snapshots (PodReflectorSource) and lets CatalogReconciler own the desired-versus-actual diff. Selector::reconcile, WorkerRegistration, and the EPP-local reconcile state are removed; their tests moved to dynamo-kv-router with the reconciler. A closed reflector channel still yields one empty snapshot so selection fails closed, as before. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

…ex comes from Replace three knobs that together described the same thing (the builder's external_kv_events() and remote_indexer(), and SelectionHost's OverlapRefreshSource) with one KvIndexSource on HostCache: - Owned: the core owns the index and subscribes to worker KV events over ZMQ (the standalone default; workers need KV-event endpoints). - External { refresh }: the host owns and feeds the index; workers need no KV-event metadata; refresh optionally lets the core re-score queued requests against the host's index (the frontend's embedded partition). - Remote(url): a standalone indexer serves the primary index. The listener decision, schedulability metadata, and dequeue-time refresh all derive from this one value. SelectionServiceBuilder::kv_index sets it; SelectionServiceConfig and the Python bindings map remote_indexer_url onto Remote. No behavior change. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

…ion service lib/llm and lib/kv-router each carried an Indexer enum with the same Single/Concurrent/Remote shape (primary, lower tiers, predict-on-route side overlay, primary_records_routing_decisions) and diverging method sets. Fold the frontend's richer implementation into dynamo_kv_router::services::indexer: - lookup.rs: the borrowed/owned HashInput lookup pipeline (primary, lower tiers seeded from confirmed primary matches, side overlay merged by per-worker max), find_matches_by_tier[_ref][_with_options], find_primary_*. - recording.rs: routing-decision recording targets (side overlay, approximate primary, remote primary) and the per-request ApproximateRequestLease. - backend.rs: the merged enum gains None, residency projection setters, try_apply_event (the backpressure-aware admit with a completion barrier on Cleared that recovery relies on), reset_worker_dp_rank_and_wait, and the approximate-LRU capacity and stats API. apply_event_routed and remove_worker[_dp_rank] stay as the standalone listener's path. Remote primaries go behind RemotePrimary so the HTTP standalone client here and the request-plane RemoteIndexer in lib/llm plug into the same enum. dump_events returns KvRouterError like the rest of the surface. - The frontend keeps only construction (indexer::build from KvRouterConfig with component metrics and the request-plane remote), the served-indexer service, and recovery. Its 16 indexer tests move with the code. No behavior change: both hosts run the exact code they ran before for their own paths; TieredMatchProvider keeps the capability-gated variant. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

KvIndexSource::External { refresh } gave the core only a refresh handle onto an index the host kept elsewhere. Replace it with Provided(Arc<Indexer>): the host builds and feeds the index, and the core holds it as every partition's index, reading it for selection and dequeue-time refresh (when the index supports refresh) and never writing to it. The frontend's embedded partition now shares the KvRouter's Indexer object instead of a TieredMatchProvider view of it. RefreshProvider loses its External variant. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

KvIndexSource::Owned carries the ingress that feeds the core-built index. The trait decides what metadata a worker needs to be schedulable, attaches its ranks when the worker becomes schedulable, and detaches them when it leaves. ZmqDirectIngress is the existing standalone behavior moved out of SelectionCore: endpoint validation, one ZMQ listener per rank on the registry, replay endpoint, deregistration on delete. The core keeps the lifecycle decisions (when to attach, how to mark Incomplete) and no longer knows about ZMQ endpoints. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

…tend KvEventIngress now builds the partition's index (open) as well as feeding it, so KvIndexSource::Owned covers both hosts and Provided goes away: - ZmqDirectIngress opens the registry's index and, on detach, deregisters the worker's listeners or, when there are none (events off or a remote primary), drops the worker from the index: the two behaviors SelectionCore used to branch on itself. - RuntimeIngress (lib/llm/src/kv_router/indexer/ingress.rs) builds the frontend's index from the router config, starts the event-plane subscription with per-rank recovery and the served-indexer endpoint, and hands the index to the partition. Its per-worker hooks are no-ops because membership comes from the KV-source watch. KvRouter::new shrinks to constructing it; task guard and shutdown go through it. The core no longer special-cases a host-fed index anywhere: every Owned index is opened by and detached through its ingress. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

… host Move the session table out of lib/llm/src/session_affinity into dynamo_kv_router::services::selection::affinity: SessionAffinity (session id -> bound worker/rank, versioned, TTL reaper, entry and id limits), the AffinityInitialization / AffinityLease lifecycle (commit, rebind, invalidate, release on drop), hard and soft validation, and replica apply with the same outcome table. Replication becomes an AffinityReplicaSink the host installs; the table stays transport-free. The frontend AffinityCoordinator is now a thin wrapper: it converts the pipeline RouteTarget at the boundary, abandons an initialization wait when the request context stops, holds a lease for the life of a response stream, and keeps its runtime event-plane replica sync (same subject and wire format) behind the sink. SessionAffinityMode moves with the table; lib/llm re-exports it. All 37 session-affinity tests and the routing-host affinity tests run unchanged against the wrapper. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

SelectionServiceBuilder::session_affinity(ttl) gives the core its own SessionAffinity table. A booked selection whose request carries a session id (session_id or session_context) and no explicit pin or affinity target acquires the session first: a bound session steers selection to its worker (exclusive for the default selector) and a new one is bound to the worker the scheduler chose. The lease is held with the booking and released by free_reservation; a bound worker that was not selected (it left) is invalidated and the session rebound. Advisory selects read the binding. Bindings replicate over the existing ZMQ peer mesh on their own topic (dynamo.session-affinity.v1) so peers that do not subscribe never see them; the active-sequence topic and wire format are untouched. The publisher multiplexes both channels onto one socket; PeerManager subscribes to the second topic only when a handler is given. SelectionServiceConfig carries the TTL for the standalone binary. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

DYN_EPP_SESSION_AFFINITY_TTL_SECS turns on the selector's session table. The EPP passes x-dynamo-session-id, the same header the frontend reads, as the request's session id, so repeat requests of a session land on the worker that served it and peer EPPs learn the binding over the replica mesh. Hard mode only: the EPP routes by headers and never observes where a request ran. The standalone selection binary and the Python SelectionService gain the same TTL setting (--session-affinity-ttl-secs, session_affinity_ttl_secs). Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

Phase 3 step 1 of the runtime-optional frontend roadmap: the contract-5 disaggregated coordination as a runtime-free component over two selection pools, shared by every host (frontend, EPP, standalone service). DisaggCoordinator books a request decode-first: advisory decode preview (potential_decode_blocks, decode_busy, cached tokens), optional advisory prefill busy read, conditional-disaggregation policy plus the decode busy gate, decode commit pinned to the previewed worker, then prefill commit constrained to the decode worker's KV-transfer domain (its kv_transfer_domain and topology_domains become a required or preferred dynamo.topology taint per kv_transfer_enforcement). Cross-pool atomicity is explicitly waived: the two ledgers are independent and every partial state is compensated and recorded in the plan's LinkedBookingState transitions. Prefill failure after decode is booked follows PrefillFailurePolicy (free decode and fail, or bypass on the decode worker); prefill_complete frees prefill early; release frees the rest. SelectionPool is the seam: SelectionService implements it; the EPP can implement it over its own client. Not yet wired into PrefillRouter (the production path still books prefill first); that switch is gated on the replay-parity disaggregated lifecycle campaign. Tests: linked path with exact transition sequence and per-pool load checks, policy bypass, decode-gate veto, prefill failure under both compensation policies, decode failure booking nothing, KV-transfer-domain constraint steering prefill, and incomplete KV-transfer metadata failing closed. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com> (cherry picked from commit dc94857)

The coordinator gains a CoordinationOrder. PrefillAnchored is what the frontend PrefillRouter does today: book prefill KV-aware on the prompt (plan_prefill), then book decode load-only inside the prefill worker's KV-transfer domain (plan_decode) whenever the host is ready, before or after prefill_complete. DecodeAnchored is the existing preview/bypass/commit flow. plan() runs the configured order end to end for hosts that book both pools at once (the EPP). DisaggPlan now holds optional prefill and decode bookings behind accessors and carries its request, so the two phases share one state machine, one release, and one KV-transfer constraint helper (derived from whichever worker anchors). DisaggRequest carries the pin, affinity target, and allow-set the frontend passes today. Compensation table and the new PrefillCommitted state are covered by focused tests for both orders. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

…e EPP The standalone EPP hosts a DisaggCoordinator over a prefill and a decode SelectionService when DYN_EPP_PREFILL_INFERENCE_POOL_NAME names a second InferencePool. Selection, booking, cancellation, and error semantics come from the coordinator; headers name both workers (x-dynamo-prefill-instance-id and x-dynamo-prefill-dp-rank) and on_prefill_complete/on_request_complete map to prefill_complete/release. DYN_EPP_DISAGG_ORDER selects the CoordinationOrder: decode-anchored (default, the previous behavior) or prefill-anchored, the order the frontend uses. DYN_EPP_DISAGG_PREFILL_FAILURE keeps its free-decode / bypass-on-decode compensation for the decode-anchored order. (Adapted from the p3-epp-disagg branch onto the two-order coordinator API and the WorkerCatalogSource-based reflector.) Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

ishandhanani · 2026-09-04T04:20:25Z

Gates at `d364d29` (rebased onto #53 / `idhanani/up-9-epp-affinity`)

cargo test -p dynamo-kv-router --features standalone-selection,standalone-slot-tracker --lib: 1118 passed, 2 ignored.
cargo test -p dynamo-ext-proc: 142 passed.
cargo clippy -p dynamo-kv-router -p dynamo-ext-proc --all-targets --features dynamo-kv-router/standalone-selection -- -D warnings: clean.

Rebase from the previous tip was conflict-free (three commits).

ishandhanani had a problem deploying to external_collaborator September 3, 2026 22:33 — with GitHub Actions Failure

github-actions Bot added feat documentation Improvements or additions to documentation labels Sep 3, 2026

ishandhanani changed the base branch from idhanani/s1c-kv-index-source to idhanani/s1c-kv-index-unify September 4, 2026 00:04

ishandhanani force-pushed the idhanani/s2-coordinator branch from b20380d to 8ea88e6 Compare September 4, 2026 00:04

ishandhanani had a problem deploying to external_collaborator September 4, 2026 00:04 — with GitHub Actions Failure

ishandhanani changed the base branch from idhanani/s1c-kv-index-unify to idhanani/s1d-affinity September 4, 2026 00:32

ishandhanani force-pushed the idhanani/s2-coordinator branch from 8ea88e6 to b3cd827 Compare September 4, 2026 00:32

ishandhanani had a problem deploying to external_collaborator September 4, 2026 00:32 — with GitHub Actions Failure

ishandhanani had a problem deploying to external_collaborator September 4, 2026 00:33 — with GitHub Actions Failure

ishandhanani changed the base branch from idhanani/s1d-affinity to idhanani/s1c-runtime-ingress September 4, 2026 01:25

ishandhanani force-pushed the idhanani/s2-coordinator branch from b3cd827 to 7fe3398 Compare September 4, 2026 01:25

ishandhanani had a problem deploying to external_collaborator September 4, 2026 01:25 — with GitHub Actions Failure

ishandhanani changed the base branch from idhanani/s1c-runtime-ingress to idhanani/s1d-epp-affinity September 4, 2026 01:52

ishandhanani had a problem deploying to external_collaborator September 4, 2026 01:52 — with GitHub Actions Failure

ishandhanani force-pushed the idhanani/s2-coordinator branch from 7fe3398 to 1ca17dd Compare September 4, 2026 01:53

ishandhanani added 17 commits September 4, 2026 03:59

ishandhanani added 3 commits September 4, 2026 04:15

ishandhanani force-pushed the idhanani/s2-coordinator branch from 3e0d869 to d364d29 Compare September 4, 2026 04:15

ishandhanani had a problem deploying to external_collaborator September 4, 2026 04:15 — with GitHub Actions Failure

ishandhanani added this pull request to stack #54 September 9, 2026 06:20

This was referenced Sep 10, 2026

feat(router): unify frontend and standalone selection core #55

Closed

feat(router): unify frontend and standalone selection core ai-dynamo/dynamo#14570

Merged

ishandhanani closed this Sep 19, 2026

Navigation Menu

Repository navigation

feat(kv-router): DisaggCoordinator with prefill-anchored and decode-anchored orders - #40

feat(kv-router): DisaggCoordinator with prefill-anchored and decode-anchored orders#40
ishandhanani wants to merge 27 commits into
idhanani/up-9-epp-affinityfrom
idhanani/s2-coordinator

ishandhanani commented Sep 3, 2026 •

edited

Loading

Uh oh!

ishandhanani commented Sep 4, 2026

Uh oh!

This branch had an error being deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

ishandhanani commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Not in this PR (spec deliverable 2)

Validation

Review notes (graham-code-review pass)

Uh oh!

ishandhanani commented Sep 4, 2026

Gates at d364d29 (rebased onto #53 / idhanani/up-9-epp-affinity)

Uh oh!

This branch had an error being deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ishandhanani commented Sep 3, 2026 •

edited

Loading

Gates at `d364d29` (rebased onto #53 / `idhanani/up-9-epp-affinity`)