Repository navigation
feat(router): embedded selection partition parity with the runtime scheduler - #35
Closed
ishandhanani wants to merge 4 commits into
Closed
ishandhanani wants to merge 4 commits into
ishandhanani wants to merge 4 commits into
Conversation
KvWorkerMonitor's prefill-token overload detection reads the scheduler's own load snapshots through SchedulerLoadSender. The embedded selection partition never published them, so --active-prefill-tokens-threshold could not trigger 529 under --router-embedded-selection (tests/router test_mocker_kv_router_overload_529[prefill-tokens] failed only on that path; decode-blocks worked because it uses worker-published metrics). Add SelectionHostHooks.scheduler_load_sink (a SchedulerLoadSink) that the partition's ScopedSequencePublisher forwards publish_scheduler_load(_batch) into; the embedded host bridges it to the router's SchedulerLoadSender. tests/router overload_529 (agg + disagg, prefill-tokens + decode-blocks) now pass on the embedded path. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com> (cherry picked from commit beb6ddc)
Phase 1 exit gates now pass: the frontend A/B benchmark (+6% req/s, lower p50/p99), offline KV replay byte-parity against main on four rows, and tests/router e2e green on both scheduler paths. Flip router_embedded_selection to default true (serde default/skip updated, PyO3 constructor, .pyi, --router-embedded-selection default, docs). The runtime-bound KvScheduler path stays reachable through --no-router-embedded-selection for one release as a rollback and logs a deprecation warning at startup; deleting it is the follow-up. A caller-injected custom WorkerSelector type (not DefaultWorkerSelector) cannot run on the partition, so KvRouter::new keeps the runtime scheduler for it with a warning instead of silently changing selection semantics. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com> (cherry picked from commit 568d06c)
--router-replica-sync previously logged that it was not carried over to the embedded partition. Add SelectionHostHooks.replica_sync, a per-partition factory returning HostReplicaChannels (outbound events, inbound peer events, this replica's id); when the service runs no ZMQ peer mesh the partition's publisher forwards to the host's outbound channel and its replica consumer reads the host's inbound channel. The embedded host builds those channels over the runtime event plane (ACTIVE_SEQUENCES_SUBJECT, same wire formats and subscriber as the runtime scheduler used), so multi-replica frontends keep active-sequence sync on the default scheduler path. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com> (cherry picked from commit d586893)
…d partition Install the process-wide WorkerSelectionPolicyRegistry (the linked custom policy catalog the Python bindings register under --features custom-policy) through dynamo_llm::kv_router::install_worker_selection_policy_registry; embedded partitions and the static frontend resolve KvRouterConfig policy instances against it instead of the built-in default registry. A KvRouter built with Sel = WorkerSelectionPolicy therefore runs on the embedded partition too; only an arbitrary caller-injected WorkerSelector type (unit test fakes) still selects the runtime scheduler. Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com> (cherry picked from commit fc22bac)
ishandhanani
had a problem deploying
to
external_collaborator
September 3, 2026 20:23 — with 
GitHub Actions
Failure
GitHub Actions
Failure
github-actions
Bot
added
feat
documentation
This was referenced Sep 3, 2026
Closed
Owner
Author
This branch had an error being deployed
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Stack 1, group A, step 1 of the unification DEP (#33, "Stack 1 DEP r2" comment). Stacked on #18.
Summary
Brings the embedded selection partition to functional parity with the runtime-bound
KvSchedulerso the next step can delete the latter. Four logical commits, each independently reviewable:fix(router): feed scheduler load from the embedded selection partition— newSelectionHostHooks.scheduler_load_sink(SchedulerLoadSink); the partition'sScopedSequencePublisherforwardspublish_scheduler_load(_batch)into it and the embedded host bridges it to the router'sSchedulerLoadSender, soKvWorkerMonitorprefill-token overload (529) works on the embedded path.feat(router): make embedded selection the default scheduler—router_embedded_selectiondefaults totruein Rust, the Python arg group, and the bindings;--no-router-embedded-selectionlogs a deprecation warning.feat(router): replica sync for the embedded selection partition— newSelectionHostHooks.replica_sync(HostReplicaSyncFactory->HostReplicaChannels).lib/llmsupplies channels carried over the runtime event plane onACTIVE_SEQUENCES_SUBJECT, the same subject and wire format as the runtime scheduler, so mixed fleets interoperate.lib/kv-routerstill never importsdynamo-runtime.feat(router): resolve linked worker-selection policies on the embedded partition— a process-wide policy registry installed by the bindings (install_worker_selection_policy_registry); the partition resolvesrouter_prefill_policy/router_decode_policyagainst it, soWorkerSelectionPolicycallers no longer fall back to the runtime scheduler.These are the same commits as #30's scheduler-parity subset, re-based onto #18 without the direct-dispatch, decode-first, and static-frontend material. #30 will be closed pointing here.
Validation
cargo test -p dynamo-kv-router --lib: 949 passed.cargo test -p dynamo-llm --lib -- kv_router::: 332 passed.cargo clippy -p dynamo-llm -p dynamo-kv-router --tests: clean. Bindings cratecargo check: clean.tests/router/test_router_e2e_with_mockers.pyoverload_529 agg + disagg, prefill-tokens + decode-blocks pass on the embedded path; replica-sync e2e passes.Review notes (graham-code-review pass)
embedded.rs: thereplica_syncfactory hands its channels to the first partition only (one-shotMutex<Option<_>>). The router owns exactly one partition, so this is correct today; the API-cleanup step (SelectionHost.replication) removes the factory shape.kv_router.rs: theTypeIdcheck that routes arbitrarySeltypes to the runtime scheduler is transitional and goes away with theSelgeneric in the API-cleanup step.unwrap/expectin non-test code; spawned tasks inhost_replica_channelsobserve the router's cancellation token; channels are bounded (REPLICA_EVENT_CHANNEL_CAPACITY).