From b6ca9d596b1848bb251c24ca0f1435bf398b1d68 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Mon, 21 Sep 2026 02:13:42 -0700 Subject: [PATCH 1/7] feat(router): configure two-tier soft affinity Signed-off-by: Guan Luo --- .../pages/cli/kv-aware-routing/overview.mdx | 6 + .../router/configuration-and-tuning.md | 8 + .../router/custom-worker-selection.mdx | 22 +- .../soft-pin-repin/README.md | 5 +- .../soft-pin-repin/src/lib.rs | 2 + lib/kv-router/src/plugins/worker_selection.rs | 5 +- lib/kv-router/src/scheduling/queue.rs | 254 ++++++++++++++++-- lib/kv-router/src/scheduling/selector/mod.rs | 4 +- .../src/scheduling/selector/policy.rs | 30 ++- lib/llm/src/kv_router/routing_host/tests.rs | 153 ++++++++++- lib/router-plugins/builtin/src/lib.rs | 66 ++++- .../builtin/src/two_tier_cost_fn.rs | 5 + 12 files changed, 526 insertions(+), 34 deletions(-) diff --git a/docs/fern/pages/cli/kv-aware-routing/overview.mdx b/docs/fern/pages/cli/kv-aware-routing/overview.mdx index 6f483149b63a..b61819d9feab 100644 --- a/docs/fern/pages/cli/kv-aware-routing/overview.mdx +++ b/docs/fern/pages/cli/kv-aware-routing/overview.mdx @@ -96,6 +96,8 @@ worker_selection: instances: - name: dynamo-two-tier-cost-fn type: dynamo-two-tier-cost-fn + parameters: + respect_soft_affinity: true ``` Then point the frontend at it: @@ -109,6 +111,10 @@ python3 -m dynamo.frontend \ To compare against the built-in selector, restart with `DYN_ROUTER_WORKER_SELECTION_POLICY=default` — no config change needed. +`respect_soft_affinity: true` keeps an eligible soft session-affinity target as the exclusive +candidate. The default is `false`, which leaves the target advisory. Ineligible soft targets fall +back to normal two-tier selection; hard affinity and explicit worker targets remain exact. + For the available policy types and per-stage prefill/decode selection, see [Worker-Selection Policies](../../developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md#worker-selection-policies). diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md index 08026aa76981..d316fea51cf3 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md @@ -146,6 +146,7 @@ experimental router exactly. Add any subset to tune it: cache_threshold: 0.5 balance_abs_threshold: 32 balance_rel_threshold: 1.1 + respect_soft_affinity: true ``` | Parameter | Default | Meaning | @@ -153,6 +154,7 @@ experimental router exactly. Add any subset to tune it: | `cache_threshold` | `0.5` | Fraction of the request's blocks that must be device-resident on the best worker before the cache tier applies. Compared strictly. Must be in `[0.0, 1.0]`. | | `balance_abs_threshold` | `32` | Minimum active-request spread before the load tier applies. | | `balance_rel_threshold` | `1.1` | Minimum ratio of largest to smallest active-request count before the load tier applies. Must be at least `1.0`. | +| `respect_soft_affinity` | `false` | When `true`, an eligible soft session-affinity target is the only candidate passed to the two-tier picker. When `false`, the target remains advisory and the picker can select another eligible worker. | Both load gates must hold before the load tier displaces the cache tier. Parameters are validated at startup, so an out-of-range value or an unknown key fails the process immediately, naming the key, @@ -160,6 +162,12 @@ rather than being silently ignored. It selects the least-loaded worker once the largest count is more than 1.1 times the smallest; otherwise it prefers the worker holding the largest device-KV overlap when that overlap covers more than 50% of the request's blocks. +`respect_soft_affinity` affects soft affinity only. A target is eligible when it remains discovered, +available, below the overload threshold, allowed by the request's worker allowlist, and compatible +with its routing constraints. If the target is ineligible, the policy falls back to normal two-tier +selection. Hard session affinity and explicit request targets remain exact constraints for either +value. + #### Override the Selection `DYN_ROUTER_WORKER_SELECTION_POLICY` overrides every stage. `--router-prefill-policy` and diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index 5827ada9569d..e622f2fdc76a 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -412,7 +412,7 @@ Request context and worker identity are always available. Dynamo calculates opti | `block_size()` | Tokens in one KV block | | `tracks_prefill_tokens()` | Whether the request contributes to prefill-load tracking | | `session_context()` | Optional session metadata described below | -| `affinity_target()` | Optional advisory worker and data-parallel rank from soft session affinity | +| `affinity_target()` | Optional worker and data-parallel rank from soft session affinity | | `expected_output_tokens()` | Optional expected output length | | `priority_jump()` | Scheduler priority boost. Queue policies treat negative values as zero | | `strict_priority()` | Strict integer priority. The queue orders larger values first | @@ -433,7 +433,7 @@ The [custom policy examples](https://github.com/ai-dynamo/dynamo/blob/main/examp ### Soft Session Affinity -Set `--router-session-affinity-mode soft` with a session-affinity TTL to let a custom policy influence an existing session binding. The current binding enters the normal selection pipeline as `affinity_target()`. The policy still receives the full host-eligible candidate set, subject to its own filters, and may select another worker. Dynamo rebinds the session after dispatch returns a response stream. +Set `--router-session-affinity-mode soft` with a session-affinity TTL to let a custom policy influence an existing session binding. The current binding enters the normal selection pipeline as `affinity_target()`. By default, the policy receives the full host-eligible candidate set, subject to its own filters, and may select another worker. Dynamo rebinds the session after dispatch returns a response stream. ```mermaid flowchart LR @@ -445,6 +445,24 @@ flowchart LR Match both fields when a policy wants to recognize the target. A target with `dp_rank: None` matches every rank on its worker; a populated rank matches only that worker-rank pair. The target can be absent from the candidate table when the worker is unavailable or a custom filter rejects it. +To retain an eligible target without duplicating affinity logic in the picker, opt into scheduler +candidate narrowing when you build the policy: + +```rust +WorkerSelectionPolicy::new( + config.clone(), + worker_type.as_str(), + scorers, + Box::new(LowestCostPicker), +) +.with_exclusive_affinity_target(true) +``` + +The default is `false`. When enabled, Dynamo checks discovery, availability, overload, request +allowlists, routing constraints, and the target's data-parallel rank before narrowing. An ineligible +target leaves the normal candidate set intact. Policy filters still run after narrowing, so rejecting +the sole narrowed candidate returns the policy's usual all-filtered error. + Hard affinity remains the default. A hard binding and every explicit request target take the exact-target path instead of advisory custom selection. A selection, setup, or dispatch failure before a response stream leaves the previous soft binding intact. An error or cancellation after the stream is returned does not roll back a completed rebind. The [custom policy examples](https://github.com/ai-dynamo/dynamo/blob/main/examples/router/custom-policy-example/README.md) include an overload-aware soft pinning policy that retains the advisory target until its active-request count exceeds a configured threshold, then selects the least-loaded alternative. Its two-Mocker walkthrough holds the first request open and verifies an `A -> B -> B` worker sequence: overload moves the second request to B, and the third request retains B after load drains, proving that the plugin-selected dispatch updated the binding. diff --git a/examples/router/custom-policy-example/soft-pin-repin/README.md b/examples/router/custom-policy-example/soft-pin-repin/README.md index cb153757bac7..fd94b3416c74 100644 --- a/examples/router/custom-policy-example/soft-pin-repin/README.md +++ b/examples/router/custom-policy-example/soft-pin-repin/README.md @@ -7,7 +7,10 @@ SPDX-License-Identifier: Apache-2.0 The `soft-pin-repin` policy keeps a session on its current worker until that worker exceeds an active-request threshold. It then selects the least-loaded alternative and lets Dynamo store the selected worker as the new soft binding. -The policy implements only `WorkerPicker`. It requests `WorkerInputs::LOAD`; Dynamo continues to own worker discovery, eligibility, reservations, accounting, dispatch, and session binding state. +The policy implements only `WorkerPicker`. It requests `WorkerInputs::LOAD` and explicitly keeps +`with_exclusive_affinity_target(false)` so the picker can compare the current soft target with every +eligible alternative. Dynamo continues to own worker discovery, eligibility, reservations, +accounting, dispatch, and session binding state. ## Policy Behavior diff --git a/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs b/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs index 87ab114593bd..b47eda31af96 100644 --- a/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs +++ b/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs @@ -108,6 +108,8 @@ fn provider( max_active_requests, }), ) + // This policy needs every eligible candidate so it can decide when to repin. + .with_exclusive_affinity_target(false) }, )) } diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index b16c5db7cb1e..59e138463d1a 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -191,8 +191,9 @@ impl WorkerSelectionContext<'_> { /// Return the session-affinity target resolved by the request host. /// - /// The default selector treats an eligible target as exclusive. Custom policies receive it as - /// advisory context; it may be absent from their candidate set when unavailable or filtered. + /// Custom policies receive the target as advisory context by default and may opt into exclusive + /// handling through [`WorkerSelectionPolicy::with_exclusive_affinity_target`]. An ineligible + /// target does not narrow the candidate set and may be absent from it. pub fn affinity_target(&self) -> Option { self.request.affinity_target } diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index d120f186198c..ffeeb91682de 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -30,6 +30,8 @@ use super::types::{ NonMaxOverlapSelection, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, SchedulingContext, SchedulingRequest, SchedulingResponse, WorkerAvailabilityProvider, }; +#[cfg(test)] +use crate::protocols::WorkerAffinityTarget; use crate::protocols::{ LocalBlockHash, PrefillLoadHint, WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank, @@ -2119,9 +2121,14 @@ mod tests { #[derive(Clone)] struct MinDecodeSelector { rendezvous: Option>, + exclusive_affinity_target: bool, } impl WorkerSelector for MinDecodeSelector { + fn uses_exclusive_affinity_target(&self) -> bool { + self.exclusive_affinity_target + } + fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::CACHE | WorkerInputs::LOAD } @@ -2199,8 +2206,50 @@ mod tests { Arc>, Arc>, ) { - let dp_range: HashMap = - (0..num_workers as u64).map(|id| (id, (0, 1))).collect(); + let configs = (0..num_workers as u64) + .map(|id| { + ( + id, + SimpleWorkerConfig { + max_num_batched_tokens: Some(isl as u64), + ..Default::default() + }, + ) + }) + .collect(); + make_queue_with_custom_configs_and_selector( + configs, + block_size, + threshold_frac, + selector, + None, + None, + ) + } + + #[allow(clippy::type_complexity)] + fn make_queue_with_custom_configs_and_selector< + Sel: WorkerSelector + Send + 'static, + >( + configs: HashMap, + block_size: u32, + threshold_frac: Option, + selector: Sel, + overloaded_worker_provider: Option, + available_worker_provider: Option, + ) -> ( + Arc>, + Arc>, + ) { + let dp_range: HashMap = configs + .iter() + .map(|(&id, config)| { + ( + id, + (config.data_parallel_start_rank, config.data_parallel_size), + ) + }) + .collect(); let slots = Arc::new(ActiveSequencesMultiWorker::new( NoopSequencePublisher, block_size as usize, @@ -2209,17 +2258,6 @@ mod tests { 0, "test", )); - - let mut configs: HashMap = HashMap::new(); - for id in 0..num_workers as u64 { - configs.insert( - id, - SimpleWorkerConfig { - max_num_batched_tokens: Some(isl as u64), - ..Default::default() - }, - ); - } let (_cfg_tx, cfg_rx) = watch::channel(configs); let queue = Arc::new(SchedulerQueue::new( @@ -2230,8 +2268,8 @@ mod tests { selector, None, None, - None, - None, + overloaded_worker_provider, + available_worker_provider, )); (queue, slots) @@ -2623,6 +2661,176 @@ mod tests { (req, rx) } + #[tokio::test] + async fn affinity_target_narrows_candidates_only_when_selector_opts_in() { + let target = WorkerWithDpRank::new(1, 0); + for (exclusive_affinity_target, expected) in + [(false, WorkerWithDpRank::new(0, 0)), (true, target)] + { + let (queue, _slots) = make_queue_with_custom_selector( + 2, + 16, + 64, + None, + MinDecodeSelector { + rendezvous: None, + exclusive_affinity_target, + }, + ); + let (mut request, _response_rx) = make_request("soft-affinity", 64); + request.affinity_target = Some(target.into()); + + let selected = queue.select_without_admission(request).await.unwrap(); + + assert_eq!(selected.response.best_worker, expected); + } + } + + #[tokio::test] + async fn exclusive_affinity_falls_back_when_target_is_ineligible() { + let selector = || MinDecodeSelector { + rendezvous: None, + exclusive_affinity_target: true, + }; + let target = WorkerWithDpRank::new(1, 0); + + let overloaded: OverloadedWorkerProvider = Arc::new(|| Some(HashSet::from([1]))); + let configs = (0..2) + .map(|id| { + ( + id, + SimpleWorkerConfig { + max_num_batched_tokens: Some(64), + ..Default::default() + }, + ) + }) + .collect(); + let (queue, _slots) = make_queue_with_custom_configs_and_selector( + configs, + 16, + None, + selector(), + Some(overloaded), + None, + ); + let (mut request, _response_rx) = make_request("overloaded-affinity", 64); + request.affinity_target = Some(target.into()); + assert_eq!( + queue + .select_without_admission(request) + .await + .unwrap() + .response + .best_worker, + WorkerWithDpRank::new(0, 0) + ); + + let available: WorkerAvailabilityProvider = + Arc::new(|_| Some(Arc::new(HashSet::from([0])))); + let configs = (0..2) + .map(|id| { + ( + id, + SimpleWorkerConfig { + max_num_batched_tokens: Some(64), + ..Default::default() + }, + ) + }) + .collect(); + let (queue, _slots) = make_queue_with_custom_configs_and_selector( + configs, + 16, + None, + selector(), + None, + Some(available), + ); + let (mut request, _response_rx) = make_request("unavailable-affinity", 64); + request.affinity_target = Some(target.into()); + assert_eq!( + queue + .select_without_admission(request) + .await + .unwrap() + .response + .best_worker, + WorkerWithDpRank::new(0, 0) + ); + + let (queue, _slots) = make_queue_with_custom_selector(2, 16, 64, None, selector()); + let (mut request, _response_rx) = make_request("disallowed-affinity", 64); + request.affinity_target = Some(target.into()); + request.allowed_worker_ids = Some(HashSet::from([0])); + assert_eq!( + queue + .select_without_admission(request) + .await + .unwrap() + .response + .best_worker, + WorkerWithDpRank::new(0, 0) + ); + } + + #[tokio::test] + async fn exclusive_affinity_respects_exact_and_worker_only_dp_targets() { + let configs = HashMap::from([ + ( + 7, + SimpleWorkerConfig { + data_parallel_start_rank: 2, + data_parallel_size: 2, + max_num_batched_tokens: Some(64), + ..Default::default() + }, + ), + ( + 8, + SimpleWorkerConfig { + max_num_batched_tokens: Some(64), + ..Default::default() + }, + ), + ]); + let (queue, _slots) = make_queue_with_custom_configs_and_selector( + configs, + 16, + None, + MinDecodeSelector { + rendezvous: None, + exclusive_affinity_target: true, + }, + None, + None, + ); + + let (mut exact, _response_rx) = make_request("exact-rank-affinity", 64); + exact.affinity_target = Some(WorkerWithDpRank::new(7, 3).into()); + assert_eq!( + queue + .select_without_admission(exact) + .await + .unwrap() + .response + .best_worker, + WorkerWithDpRank::new(7, 3) + ); + + let (mut worker_only, _response_rx) = make_request("worker-only-affinity", 64); + worker_only.affinity_target = Some(WorkerAffinityTarget::new(7, None)); + assert_eq!( + queue + .select_without_admission(worker_only) + .await + .unwrap() + .response + .best_worker, + WorkerWithDpRank::new(7, 2) + ); + } + #[tokio::test] async fn classifier_overrides_are_validated_before_enqueue() { let profile = policy_profile( @@ -3110,7 +3318,10 @@ policy_classes: 16, 64, None, - MinDecodeSelector { rendezvous: None }, + MinDecodeSelector { + rendezvous: None, + exclusive_affinity_target: false, + }, ); let worker0 = WorkerWithDpRank::new(0, 0); let worker1 = WorkerWithDpRank::new(1, 0); @@ -3158,7 +3369,10 @@ policy_classes: 16, 64, None, - MinDecodeSelector { rendezvous: None }, + MinDecodeSelector { + rendezvous: None, + exclusive_affinity_target: false, + }, ); let observer_gate = Arc::new((StdMutex::new(false), Condvar::new())); let callback_gate = Arc::clone(&observer_gate); @@ -3218,7 +3432,10 @@ policy_classes: 16, 64, None, - MinDecodeSelector { rendezvous: None }, + MinDecodeSelector { + rendezvous: None, + exclusive_affinity_target: false, + }, ); let worker0 = WorkerWithDpRank::new(0, 0); let worker1 = WorkerWithDpRank::new(1, 0); @@ -3465,6 +3682,7 @@ policy_classes: async fn test_concurrent_immediate_admissions_see_prior_booking() { let selector = MinDecodeSelector { rendezvous: Some(Arc::new(SelectorRendezvous::default())), + exclusive_affinity_target: false, }; let (queue, slots) = make_queue_with_custom_selector(2, 16, 512, None, selector); let barrier = Arc::new(Barrier::new(3)); diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index 61b4f5c1c33b..c82b1e97b0f8 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -36,8 +36,8 @@ pub trait WorkerSelector { /// Whether an eligible affinity target exclusively constrains worker selection. /// - /// The default selector uses exclusive affinity. Custom policies receive affinity as - /// advisory context and may choose another eligible worker. + /// The default selector uses exclusive affinity. Custom policies receive affinity as advisory + /// context unless they opt into exclusive candidate narrowing. fn uses_exclusive_affinity_target(&self) -> bool { false } diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index 01213e781f67..a4d50659e380 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -49,6 +49,7 @@ pub(super) struct CustomWorkerSelectionState { pub struct WorkerSelectionPolicy { kv_router_config: KvRouterConfig, worker_label: &'static str, + exclusive_affinity_target: bool, state: WorkerSelectionPolicyState, } @@ -87,6 +88,7 @@ impl WorkerSelectionPolicy { Self { kv_router_config, worker_label, + exclusive_affinity_target: false, state: WorkerSelectionPolicyState::Custom(RefCell::new(CustomWorkerSelectionState { filters, scorers, @@ -110,9 +112,20 @@ impl WorkerSelectionPolicy { Self { kv_router_config, worker_label, + exclusive_affinity_target: true, state: WorkerSelectionPolicyState::Default(picker), } } + + /// Configure whether an eligible affinity target exclusively constrains candidate selection. + /// + /// Custom policies are advisory by default. When enabled, the scheduler narrows the eligible + /// candidate set to the affinity target before filters, scorers, and the picker run. If the + /// target is ineligible, normal selection continues over the remaining eligible candidates. + pub fn with_exclusive_affinity_target(mut self, enabled: bool) -> Self { + self.exclusive_affinity_target = enabled; + self + } } #[inline(always)] @@ -269,7 +282,7 @@ pub(super) fn collect_custom_candidates( impl WorkerSelector for WorkerSelectionPolicy { fn uses_exclusive_affinity_target(&self) -> bool { - matches!(&self.state, WorkerSelectionPolicyState::Default(_)) + self.exclusive_affinity_target } fn required_worker_inputs(&self) -> WorkerInputs { @@ -827,12 +840,21 @@ mod tests { vec![Box::new(LoadScorer)], Box::new(FirstPicker), ); - let inputs = + let advisory_inputs = >::required_worker_inputs( &policy, ); + assert!(!uses_exclusive_affinity(&policy)); - assert!(inputs.contains(WorkerInputs::CACHE)); - assert!(inputs.contains(WorkerInputs::LOAD)); + let policy = policy.with_exclusive_affinity_target(true); + let exclusive_inputs = + >::required_worker_inputs( + &policy, + ); + + assert!(uses_exclusive_affinity(&policy)); + assert_eq!(exclusive_inputs, advisory_inputs); + assert!(exclusive_inputs.contains(WorkerInputs::CACHE)); + assert!(exclusive_inputs.contains(WorkerInputs::LOAD)); } } diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 25eaa3eb288f..e7932d9751f1 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -13,6 +13,10 @@ use std::{ use crate::kv_router::SelectionPolicySource; use dynamo_kv_router::{ config::KvRouterConfig, + plugins::worker_selection::{ + WorkerInputView, WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicy, + WorkerSelectionPolicyError, + }, protocols::RoutingConstraints, scheduling::{ ClassifierError, ClassifyEvent, ClassifyFuture, ClassifyRequest, RequestClassifier, @@ -467,6 +471,36 @@ fn affinity_request( request } +struct HighestWorkerPicker; + +impl WorkerPicker for HighestWorkerPicker { + fn pick( + &mut self, + _context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + input + .candidates() + .iter() + .enumerate() + .max_by_key(|(_, candidate)| candidate.worker()) + .map(|(row, _)| row) + .ok_or_else(|| WorkerSelectionPolicyError::failed("no eligible worker")) + } +} + +fn highest_worker_policy(exclusive_affinity_target: bool) -> SelectionPolicySource { + SelectionPolicySource::Factory(Arc::new(move |config, worker_type, _| { + WorkerSelectionPolicy::new( + config.clone(), + worker_type.as_str(), + Vec::new(), + Box::new(HighestWorkerPicker), + ) + .with_exclusive_affinity_target(exclusive_affinity_target) + })) +} + #[tokio::test] #[serial_test::serial] async fn builtin_affinity_uses_common_host_for_every_policy() { @@ -1189,6 +1223,39 @@ async fn router_with_worker_configs_and_classifier( session_affinity_ttl: Option, workers: HashMap, classifier: Option, +) -> (RoutingHost, Runtime) { + router_with_worker_configs_classifier_policy_and_mode( + session_affinity_ttl, + workers, + classifier, + SelectionPolicySource::Registry, + crate::session_affinity::SessionAffinityMode::Hard, + ) + .await +} + +async fn router_with_worker_configs_and_policy( + session_affinity_ttl: Option, + workers: HashMap, + policy: SelectionPolicySource, + affinity_mode: crate::session_affinity::SessionAffinityMode, +) -> (RoutingHost, Runtime) { + Box::pin(router_with_worker_configs_classifier_policy_and_mode( + session_affinity_ttl, + workers, + None::, + policy, + affinity_mode, + )) + .await +} + +async fn router_with_worker_configs_classifier_policy_and_mode( + session_affinity_ttl: Option, + workers: HashMap, + classifier: Option, + policy: SelectionPolicySource, + affinity_mode: crate::session_affinity::SessionAffinityMode, ) -> (RoutingHost, Runtime) { let runtime = Runtime::from_current().unwrap(); // Each runtime has its own in-memory discovery store, so fixture namespaces can repeat. @@ -1216,7 +1283,7 @@ async fn router_with_worker_configs_and_classifier( workers, None, 16, - SelectionPolicySource::Registry, + policy, Some(config), None, "decode", @@ -1233,7 +1300,11 @@ async fn router_with_worker_configs_and_classifier( let inner = PushRouter::from_client(client, RouterMode::KV) .await .unwrap(); - let mut router = RoutingHost::new(inner, Arc::new(chooser), session_affinity_ttl).unwrap(); + let affinity = session_affinity_ttl + .map(|ttl| chooser.affinity_coordinator(ttl, affinity_mode)) + .transpose() + .unwrap(); + let mut router = RoutingHost::new_with_coordinator(inner, Arc::new(chooser), affinity); // Parallel tests must not share the process-global request counters. router.request_metrics = router.request_metrics.with_isolated_counters_for_test(); router @@ -1244,6 +1315,84 @@ async fn router_with_worker_configs_and_classifier( (router, runtime) } +#[tokio::test] +async fn policy_affinity_control_rebinds_soft_sessions_and_preserves_hard_pins() { + const TARGET_WORKER: u64 = 7; + const PREFERRED_WORKER: u64 = 8; + + for exclusive_affinity_target in [false, true] { + for affinity_mode in [ + crate::session_affinity::SessionAffinityMode::Soft, + crate::session_affinity::SessionAffinityMode::Hard, + ] { + let workers = HashMap::from([ + (TARGET_WORKER, ModelRuntimeConfig::default()), + (PREFERRED_WORKER, ModelRuntimeConfig::default()), + ]); + let (router, runtime) = router_with_worker_configs_and_policy( + Some(Duration::from_secs(10)), + workers, + highest_worker_policy(exclusive_affinity_target), + affinity_mode, + ) + .await; + let session_id = SessionAffinityId::new(format!( + "policy-{affinity_mode:?}-{exclusive_affinity_target}" + )); + bind_affinity_target(&router, &session_id, AffinityTarget::worker(TARGET_WORKER)).await; + let mut request = Context::new(request()); + request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone()); + + let (selection, operation) = router + .select_with_affinity( + &request, + RequestPhase::Aggregated, + false, + &CleanupBudget::default(), + ) + .await + .unwrap(); + let expected = if affinity_mode == crate::session_affinity::SessionAffinityMode::Soft + && !exclusive_affinity_target + { + PREFERRED_WORKER + } else { + TARGET_WORKER + }; + assert_eq!(selection.worker.worker_id, expected); + + // Model a successful dispatch response through the same affinity binding path used + // by `generate`: advisory soft affinity may rebind, while hard affinity remains pinned. + let output = Annotated::from_data(LLMEngineOutput { + finish_reason: Some(FinishReason::Stop), + ..Default::default() + }); + let stream = ResponseStream::new( + Box::pin(stream::once(async move { output })), + request.context().clone(), + ); + let mut stream = router + .bind_affinity(operation, route_target(selection.worker), stream) + .unwrap(); + while stream.next().await.is_some() {} + router.kv_router().free(request.id()).await.unwrap(); + + assert_eq!( + router + .affinity + .as_ref() + .unwrap() + .query_target(&session_id, None) + .unwrap(), + Some(AffinityTarget::worker(expected)) + ); + + drop(router); + runtime.shutdown(); + } + } +} + /// A dispatch plane that yields before completing, unlike [`CompletedBuiltinDispatch`]. /// A recorded worker id therefore means the dispatch survived, not merely started. #[derive(Default)] diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index 06e8993e9303..dff50302282e 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -36,7 +36,8 @@ pub fn register( #[cfg(test)] mod tests { use dynamo_kv_router::plugins::worker_selection::WorkerSelectionPolicyFactory; - use dynamo_kv_router::{KvRouterConfig, RoutingPartitionRef, WorkerType}; + use dynamo_kv_router::protocols::WorkerConfigLike; + use dynamo_kv_router::{KvRouterConfig, RoutingPartitionRef, WorkerSelector, WorkerType}; use super::*; @@ -61,6 +62,41 @@ mod tests { (config, resolved) } + struct PolicyTestWorker; + + impl WorkerConfigLike for PolicyTestWorker { + fn data_parallel_start_rank(&self) -> u32 { + 0 + } + + fn data_parallel_size(&self) -> u32 { + 1 + } + + fn max_num_batched_tokens(&self) -> Option { + None + } + + fn total_kv_blocks(&self) -> Option { + None + } + } + + fn resolves_exclusive_affinity(yaml: &str) -> bool { + let (config, resolved) = resolve(yaml); + let factory = resolved + .unwrap() + .expect("a configured instance resolves to a factory"); + let policy = factory( + &config, + WorkerType::Aggregated, + RoutingPartitionRef::new("model", "default"), + ); + >::uses_exclusive_affinity_target( + &policy, + ) + } + /// Catches a policy type name that drifts from its documentation, and proves the documented /// instance shape constructs for every stage it selects. #[test] @@ -90,6 +126,30 @@ worker_selection: } } + #[test] + fn resolves_soft_affinity_parameter() { + let yaml = |parameter: &str| { + format!( + r#" +worker_selection: + aggregated: two-tier + instances: + - name: two-tier + type: dynamo-two-tier-cost-fn +{parameter} +"# + ) + }; + + assert!(!resolves_exclusive_affinity(&yaml(""))); + assert!(!resolves_exclusive_affinity(&yaml( + " parameters:\n respect_soft_affinity: false" + ))); + assert!(resolves_exclusive_affinity(&yaml( + " parameters:\n respect_soft_affinity: true" + ))); + } + /// An unknown parameter key is a mistake, most often a misremembered threshold name. It must /// fail startup rather than silently leaving the default in place. #[test] @@ -102,7 +162,7 @@ worker_selection: - name: dynamo-two-tier-cost-fn type: dynamo-two-tier-cost-fn parameters: - cache_affinity_threshold: 0.8 + respect_soft_affinitty: true "#, ); @@ -115,7 +175,7 @@ worker_selection: "unexpected error: {error}" ); assert!( - error.to_string().contains("cache_affinity_threshold"), + error.to_string().contains("respect_soft_affinitty"), "the error should name the offending key: {error}" ); } diff --git a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs index f4a2a3f3bf37..e26f5612cc98 100644 --- a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs +++ b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs @@ -48,6 +48,7 @@ pub const POLICY_TYPE: &str = "dynamo-two-tier-cost-fn"; const DEFAULT_CACHE_THRESHOLD: f64 = 0.5; const DEFAULT_BALANCE_ABS_THRESHOLD: usize = 32; const DEFAULT_BALANCE_REL_THRESHOLD: f64 = 1.1; +const DEFAULT_RESPECT_SOFT_AFFINITY: bool = false; /// Tunables for [`POLICY_TYPE`], named after their `sgl-router` counterparts. /// @@ -63,6 +64,8 @@ struct Parameters { balance_abs_threshold: usize, /// Minimum ratio of largest to smallest active-request count before the load tier applies. balance_rel_threshold: f64, + /// Whether an eligible soft-affinity target exclusively constrains candidate selection. + respect_soft_affinity: bool, } impl Default for Parameters { @@ -71,6 +74,7 @@ impl Default for Parameters { cache_threshold: DEFAULT_CACHE_THRESHOLD, balance_abs_threshold: DEFAULT_BALANCE_ABS_THRESHOLD, balance_rel_threshold: DEFAULT_BALANCE_REL_THRESHOLD, + respect_soft_affinity: DEFAULT_RESPECT_SOFT_AFFINITY, } } } @@ -173,6 +177,7 @@ fn provider( Vec::new(), Box::new(TwoTierCostFnPicker { parameters }), ) + .with_exclusive_affinity_target(parameters.respect_soft_affinity) }, )) } From 20a4f841c1c07759ea888187afb22ca60ce58c34 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Mon, 21 Sep 2026 10:21:03 -0700 Subject: [PATCH 2/7] test(router): streamline soft affinity coverage Signed-off-by: Guan Luo --- lib/kv-router/src/scheduling/queue.rs | 152 +++++--------------- lib/llm/src/kv_router/routing_host/tests.rs | 119 +++++++-------- 2 files changed, 86 insertions(+), 185 deletions(-) diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index ffeeb91682de..7dee5adfeca8 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -2217,14 +2217,7 @@ mod tests { ) }) .collect(); - make_queue_with_custom_configs_and_selector( - configs, - block_size, - threshold_frac, - selector, - None, - None, - ) + make_queue_with_custom_configs_and_selector(configs, block_size, threshold_frac, selector) } #[allow(clippy::type_complexity)] @@ -2235,8 +2228,6 @@ mod tests { block_size: u32, threshold_frac: Option, selector: Sel, - overloaded_worker_provider: Option, - available_worker_provider: Option, ) -> ( Arc>, Arc>, @@ -2268,8 +2259,8 @@ mod tests { selector, None, None, - overloaded_worker_provider, - available_worker_provider, + None, + None, )); (queue, slots) @@ -2687,95 +2678,28 @@ mod tests { } #[tokio::test] - async fn exclusive_affinity_falls_back_when_target_is_ineligible() { - let selector = || MinDecodeSelector { - rendezvous: None, - exclusive_affinity_target: true, - }; - let target = WorkerWithDpRank::new(1, 0); - - let overloaded: OverloadedWorkerProvider = Arc::new(|| Some(HashSet::from([1]))); - let configs = (0..2) - .map(|id| { - ( - id, - SimpleWorkerConfig { - max_num_batched_tokens: Some(64), - ..Default::default() - }, - ) - }) - .collect(); - let (queue, _slots) = make_queue_with_custom_configs_and_selector( - configs, - 16, - None, - selector(), - Some(overloaded), - None, - ); - let (mut request, _response_rx) = make_request("overloaded-affinity", 64); - request.affinity_target = Some(target.into()); - assert_eq!( - queue - .select_without_admission(request) - .await - .unwrap() - .response - .best_worker, - WorkerWithDpRank::new(0, 0) - ); - - let available: WorkerAvailabilityProvider = - Arc::new(|_| Some(Arc::new(HashSet::from([0])))); - let configs = (0..2) - .map(|id| { - ( - id, - SimpleWorkerConfig { - max_num_batched_tokens: Some(64), - ..Default::default() - }, - ) - }) - .collect(); - let (queue, _slots) = make_queue_with_custom_configs_and_selector( - configs, + async fn exclusive_affinity_falls_back_when_target_is_disallowed() { + let (queue, _slots) = make_queue_with_custom_selector( + 2, 16, + 64, None, - selector(), - None, - Some(available), - ); - let (mut request, _response_rx) = make_request("unavailable-affinity", 64); - request.affinity_target = Some(target.into()); - assert_eq!( - queue - .select_without_admission(request) - .await - .unwrap() - .response - .best_worker, - WorkerWithDpRank::new(0, 0) + MinDecodeSelector { + rendezvous: None, + exclusive_affinity_target: true, + }, ); - - let (queue, _slots) = make_queue_with_custom_selector(2, 16, 64, None, selector()); let (mut request, _response_rx) = make_request("disallowed-affinity", 64); - request.affinity_target = Some(target.into()); + request.affinity_target = Some(WorkerWithDpRank::new(1, 0).into()); request.allowed_worker_ids = Some(HashSet::from([0])); - assert_eq!( - queue - .select_without_admission(request) - .await - .unwrap() - .response - .best_worker, - WorkerWithDpRank::new(0, 0) - ); + + let selected = queue.select_without_admission(request).await.unwrap(); + + assert_eq!(selected.response.best_worker, WorkerWithDpRank::new(0, 0)); } #[tokio::test] - async fn exclusive_affinity_respects_exact_and_worker_only_dp_targets() { + async fn exclusive_affinity_respects_dp_rank_scope() { let configs = HashMap::from([ ( 7, @@ -2802,33 +2726,25 @@ mod tests { rendezvous: None, exclusive_affinity_target: true, }, - None, - None, - ); - - let (mut exact, _response_rx) = make_request("exact-rank-affinity", 64); - exact.affinity_target = Some(WorkerWithDpRank::new(7, 3).into()); - assert_eq!( - queue - .select_without_admission(exact) - .await - .unwrap() - .response - .best_worker, - WorkerWithDpRank::new(7, 3) ); - let (mut worker_only, _response_rx) = make_request("worker-only-affinity", 64); - worker_only.affinity_target = Some(WorkerAffinityTarget::new(7, None)); - assert_eq!( - queue - .select_without_admission(worker_only) - .await - .unwrap() - .response - .best_worker, - WorkerWithDpRank::new(7, 2) - ); + for (name, target, expected) in [ + ( + "exact-rank-affinity", + WorkerAffinityTarget::new(7, Some(3)), + WorkerWithDpRank::new(7, 3), + ), + ( + "worker-only-affinity", + WorkerAffinityTarget::new(7, None), + WorkerWithDpRank::new(7, 2), + ), + ] { + let (mut request, _response_rx) = make_request(name, 64); + request.affinity_target = Some(target); + let selected = queue.select_without_admission(request).await.unwrap(); + assert_eq!(selected.response.best_worker, expected); + } } #[tokio::test] diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index e7932d9751f1..966de6643c6f 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -1316,80 +1316,65 @@ async fn router_with_worker_configs_classifier_policy_and_mode( } #[tokio::test] -async fn policy_affinity_control_rebinds_soft_sessions_and_preserves_hard_pins() { +async fn policy_affinity_control_selects_and_rebinds_soft_sessions() { const TARGET_WORKER: u64 = 7; const PREFERRED_WORKER: u64 = 8; - for exclusive_affinity_target in [false, true] { - for affinity_mode in [ + for (exclusive_affinity_target, expected) in [(false, PREFERRED_WORKER), (true, TARGET_WORKER)] + { + let workers = HashMap::from([ + (TARGET_WORKER, ModelRuntimeConfig::default()), + (PREFERRED_WORKER, ModelRuntimeConfig::default()), + ]); + let (router, runtime) = router_with_worker_configs_and_policy( + Some(Duration::from_secs(10)), + workers, + highest_worker_policy(exclusive_affinity_target), crate::session_affinity::SessionAffinityMode::Soft, - crate::session_affinity::SessionAffinityMode::Hard, - ] { - let workers = HashMap::from([ - (TARGET_WORKER, ModelRuntimeConfig::default()), - (PREFERRED_WORKER, ModelRuntimeConfig::default()), - ]); - let (router, runtime) = router_with_worker_configs_and_policy( - Some(Duration::from_secs(10)), - workers, - highest_worker_policy(exclusive_affinity_target), - affinity_mode, + ) + .await; + let session_id = SessionAffinityId::new(format!("policy-{exclusive_affinity_target}")); + bind_affinity_target(&router, &session_id, AffinityTarget::worker(TARGET_WORKER)).await; + let mut request = Context::new(request()); + request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone()); + + let (selection, operation) = router + .select_with_affinity( + &request, + RequestPhase::Aggregated, + false, + &CleanupBudget::default(), ) - .await; - let session_id = SessionAffinityId::new(format!( - "policy-{affinity_mode:?}-{exclusive_affinity_target}" - )); - bind_affinity_target(&router, &session_id, AffinityTarget::worker(TARGET_WORKER)).await; - let mut request = Context::new(request()); - request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone()); + .await + .unwrap(); + assert_eq!(selection.worker.worker_id, expected); - let (selection, operation) = router - .select_with_affinity( - &request, - RequestPhase::Aggregated, - false, - &CleanupBudget::default(), - ) - .await - .unwrap(); - let expected = if affinity_mode == crate::session_affinity::SessionAffinityMode::Soft - && !exclusive_affinity_target - { - PREFERRED_WORKER - } else { - TARGET_WORKER - }; - assert_eq!(selection.worker.worker_id, expected); + let output = Annotated::from_data(LLMEngineOutput { + finish_reason: Some(FinishReason::Stop), + ..Default::default() + }); + let stream = ResponseStream::new( + Box::pin(stream::once(async move { output })), + request.context().clone(), + ); + let mut stream = router + .bind_affinity(operation, route_target(selection.worker), stream) + .unwrap(); + while stream.next().await.is_some() {} + router.kv_router().free(request.id()).await.unwrap(); - // Model a successful dispatch response through the same affinity binding path used - // by `generate`: advisory soft affinity may rebind, while hard affinity remains pinned. - let output = Annotated::from_data(LLMEngineOutput { - finish_reason: Some(FinishReason::Stop), - ..Default::default() - }); - let stream = ResponseStream::new( - Box::pin(stream::once(async move { output })), - request.context().clone(), - ); - let mut stream = router - .bind_affinity(operation, route_target(selection.worker), stream) - .unwrap(); - while stream.next().await.is_some() {} - router.kv_router().free(request.id()).await.unwrap(); - - assert_eq!( - router - .affinity - .as_ref() - .unwrap() - .query_target(&session_id, None) - .unwrap(), - Some(AffinityTarget::worker(expected)) - ); - - drop(router); - runtime.shutdown(); - } + assert_eq!( + router + .affinity + .as_ref() + .unwrap() + .query_target(&session_id, None) + .unwrap(), + Some(AffinityTarget::worker(expected)) + ); + + drop(router); + runtime.shutdown(); } } From c11cc08191e9f2ffedc59798961fa766d76a6bd3 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Mon, 21 Sep 2026 17:03:02 -0700 Subject: [PATCH 3/7] fix(router): keep soft affinity policy-owned Signed-off-by: Guan Luo --- .../pages/cli/kv-aware-routing/overview.mdx | 7 +- .../router/configuration-and-tuning.md | 12 +- .../router/custom-worker-selection.mdx | 22 +-- .../soft-pin-repin/README.md | 5 +- .../soft-pin-repin/src/lib.rs | 2 - lib/kv-router/src/plugins/worker_selection.rs | 5 +- lib/kv-router/src/scheduling/AGENTS.md | 6 + lib/kv-router/src/scheduling/queue.rs | 166 ++---------------- lib/kv-router/src/scheduling/selector/mod.rs | 4 +- .../src/scheduling/selector/policy.rs | 30 +--- lib/llm/src/kv_router/routing_host/tests.rs | 138 +-------------- lib/router-plugins/builtin/src/lib.rs | 56 ++---- .../builtin/src/two_tier_cost_fn.rs | 96 ++++++++-- 13 files changed, 138 insertions(+), 411 deletions(-) diff --git a/docs/fern/pages/cli/kv-aware-routing/overview.mdx b/docs/fern/pages/cli/kv-aware-routing/overview.mdx index b61819d9feab..0b7dbdf45301 100644 --- a/docs/fern/pages/cli/kv-aware-routing/overview.mdx +++ b/docs/fern/pages/cli/kv-aware-routing/overview.mdx @@ -111,9 +111,10 @@ python3 -m dynamo.frontend \ To compare against the built-in selector, restart with `DYN_ROUTER_WORKER_SELECTION_POLICY=default` — no config change needed. -`respect_soft_affinity: true` keeps an eligible soft session-affinity target as the exclusive -candidate. The default is `false`, which leaves the target advisory. Ineligible soft targets fall -back to normal two-tier selection; hard affinity and explicit worker targets remain exact. +`respect_soft_affinity: true` makes the two-tier picker retain a soft session-affinity target when +that target is among its eligible candidates. The default is `false`, which lets cache and load +choose among the full eligible set. An absent target falls back to normal two-tier selection; hard +affinity and explicit worker targets remain exact. For the available policy types and per-stage prefill/decode selection, see [Worker-Selection Policies](../../developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md#worker-selection-policies). diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md index d316fea51cf3..10adf527152f 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md @@ -154,7 +154,7 @@ experimental router exactly. Add any subset to tune it: | `cache_threshold` | `0.5` | Fraction of the request's blocks that must be device-resident on the best worker before the cache tier applies. Compared strictly. Must be in `[0.0, 1.0]`. | | `balance_abs_threshold` | `32` | Minimum active-request spread before the load tier applies. | | `balance_rel_threshold` | `1.1` | Minimum ratio of largest to smallest active-request count before the load tier applies. Must be at least `1.0`. | -| `respect_soft_affinity` | `false` | When `true`, an eligible soft session-affinity target is the only candidate passed to the two-tier picker. When `false`, the target remains advisory and the picker can select another eligible worker. | +| `respect_soft_affinity` | `false` | When `true`, the two-tier picker retains a soft session-affinity target found in its eligible candidates. When `false`, cache and load choose among the full eligible set. | Both load gates must hold before the load tier displaces the cache tier. Parameters are validated at startup, so an out-of-range value or an unknown key fails the process immediately, naming the key, @@ -162,11 +162,11 @@ rather than being silently ignored. It selects the least-loaded worker once the largest count is more than 1.1 times the smallest; otherwise it prefers the worker holding the largest device-KV overlap when that overlap covers more than 50% of the request's blocks. -`respect_soft_affinity` affects soft affinity only. A target is eligible when it remains discovered, -available, below the overload threshold, allowed by the request's worker allowlist, and compatible -with its routing constraints. If the target is ineligible, the policy falls back to normal two-tier -selection. Hard session affinity and explicit request targets remain exact constraints for either -value. +`respect_soft_affinity` affects soft affinity only. Custom policies always receive the full eligible +candidate set plus the affinity target as context; this parameter tells the two-tier picker whether +to retain a matching target before considering other workers. If the target is absent from that +set, the policy falls back to normal two-tier selection. Hard session affinity and explicit request +targets remain exact constraints for either value. #### Override the Selection diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx index e622f2fdc76a..5827ada9569d 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/custom-worker-selection.mdx @@ -412,7 +412,7 @@ Request context and worker identity are always available. Dynamo calculates opti | `block_size()` | Tokens in one KV block | | `tracks_prefill_tokens()` | Whether the request contributes to prefill-load tracking | | `session_context()` | Optional session metadata described below | -| `affinity_target()` | Optional worker and data-parallel rank from soft session affinity | +| `affinity_target()` | Optional advisory worker and data-parallel rank from soft session affinity | | `expected_output_tokens()` | Optional expected output length | | `priority_jump()` | Scheduler priority boost. Queue policies treat negative values as zero | | `strict_priority()` | Strict integer priority. The queue orders larger values first | @@ -433,7 +433,7 @@ The [custom policy examples](https://github.com/ai-dynamo/dynamo/blob/main/examp ### Soft Session Affinity -Set `--router-session-affinity-mode soft` with a session-affinity TTL to let a custom policy influence an existing session binding. The current binding enters the normal selection pipeline as `affinity_target()`. By default, the policy receives the full host-eligible candidate set, subject to its own filters, and may select another worker. Dynamo rebinds the session after dispatch returns a response stream. +Set `--router-session-affinity-mode soft` with a session-affinity TTL to let a custom policy influence an existing session binding. The current binding enters the normal selection pipeline as `affinity_target()`. The policy still receives the full host-eligible candidate set, subject to its own filters, and may select another worker. Dynamo rebinds the session after dispatch returns a response stream. ```mermaid flowchart LR @@ -445,24 +445,6 @@ flowchart LR Match both fields when a policy wants to recognize the target. A target with `dp_rank: None` matches every rank on its worker; a populated rank matches only that worker-rank pair. The target can be absent from the candidate table when the worker is unavailable or a custom filter rejects it. -To retain an eligible target without duplicating affinity logic in the picker, opt into scheduler -candidate narrowing when you build the policy: - -```rust -WorkerSelectionPolicy::new( - config.clone(), - worker_type.as_str(), - scorers, - Box::new(LowestCostPicker), -) -.with_exclusive_affinity_target(true) -``` - -The default is `false`. When enabled, Dynamo checks discovery, availability, overload, request -allowlists, routing constraints, and the target's data-parallel rank before narrowing. An ineligible -target leaves the normal candidate set intact. Policy filters still run after narrowing, so rejecting -the sole narrowed candidate returns the policy's usual all-filtered error. - Hard affinity remains the default. A hard binding and every explicit request target take the exact-target path instead of advisory custom selection. A selection, setup, or dispatch failure before a response stream leaves the previous soft binding intact. An error or cancellation after the stream is returned does not roll back a completed rebind. The [custom policy examples](https://github.com/ai-dynamo/dynamo/blob/main/examples/router/custom-policy-example/README.md) include an overload-aware soft pinning policy that retains the advisory target until its active-request count exceeds a configured threshold, then selects the least-loaded alternative. Its two-Mocker walkthrough holds the first request open and verifies an `A -> B -> B` worker sequence: overload moves the second request to B, and the third request retains B after load drains, proving that the plugin-selected dispatch updated the binding. diff --git a/examples/router/custom-policy-example/soft-pin-repin/README.md b/examples/router/custom-policy-example/soft-pin-repin/README.md index fd94b3416c74..cb153757bac7 100644 --- a/examples/router/custom-policy-example/soft-pin-repin/README.md +++ b/examples/router/custom-policy-example/soft-pin-repin/README.md @@ -7,10 +7,7 @@ SPDX-License-Identifier: Apache-2.0 The `soft-pin-repin` policy keeps a session on its current worker until that worker exceeds an active-request threshold. It then selects the least-loaded alternative and lets Dynamo store the selected worker as the new soft binding. -The policy implements only `WorkerPicker`. It requests `WorkerInputs::LOAD` and explicitly keeps -`with_exclusive_affinity_target(false)` so the picker can compare the current soft target with every -eligible alternative. Dynamo continues to own worker discovery, eligibility, reservations, -accounting, dispatch, and session binding state. +The policy implements only `WorkerPicker`. It requests `WorkerInputs::LOAD`; Dynamo continues to own worker discovery, eligibility, reservations, accounting, dispatch, and session binding state. ## Policy Behavior diff --git a/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs b/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs index b47eda31af96..87ab114593bd 100644 --- a/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs +++ b/examples/router/custom-policy-example/soft-pin-repin/src/lib.rs @@ -108,8 +108,6 @@ fn provider( max_active_requests, }), ) - // This policy needs every eligible candidate so it can decide when to repin. - .with_exclusive_affinity_target(false) }, )) } diff --git a/lib/kv-router/src/plugins/worker_selection.rs b/lib/kv-router/src/plugins/worker_selection.rs index 59e138463d1a..b16c5db7cb1e 100644 --- a/lib/kv-router/src/plugins/worker_selection.rs +++ b/lib/kv-router/src/plugins/worker_selection.rs @@ -191,9 +191,8 @@ impl WorkerSelectionContext<'_> { /// Return the session-affinity target resolved by the request host. /// - /// Custom policies receive the target as advisory context by default and may opt into exclusive - /// handling through [`WorkerSelectionPolicy::with_exclusive_affinity_target`]. An ineligible - /// target does not narrow the candidate set and may be absent from it. + /// The default selector treats an eligible target as exclusive. Custom policies receive it as + /// advisory context; it may be absent from their candidate set when unavailable or filtered. pub fn affinity_target(&self) -> Option { self.request.affinity_target } diff --git a/lib/kv-router/src/scheduling/AGENTS.md b/lib/kv-router/src/scheduling/AGENTS.md index 8f80b3312e41..d76bff9a0383 100644 --- a/lib/kv-router/src/scheduling/AGENTS.md +++ b/lib/kv-router/src/scheduling/AGENTS.md @@ -122,6 +122,12 @@ PolicyClassQueue("agents") ## Public Worker-Selection API +- `WorkerSelector::uses_exclusive_affinity_target()` is a one-off compatibility hook that + preserves the default selector's historical soft-affinity behavior. Custom + `WorkerSelectionPolicy` implementations must not override or make it configurable. They receive + the full eligible candidate set and use `WorkerSelectionContext::affinity_target()` to retain, + ignore, or replace a soft target. + The `selector` module contains the public Rust contract for custom worker filters, scorers, and pickers. Treat each public item as a versioned external API. - Do not add a public field, accessor, input group, type, or re-export unless the task explicitly requires a new external policy capability. diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 7dee5adfeca8..d120f186198c 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -30,8 +30,6 @@ use super::types::{ NonMaxOverlapSelection, NonMaxOverlapSelectionObserver, OverloadedWorkerProvider, SchedulingContext, SchedulingRequest, SchedulingResponse, WorkerAvailabilityProvider, }; -#[cfg(test)] -use crate::protocols::WorkerAffinityTarget; use crate::protocols::{ LocalBlockHash, PrefillLoadHint, WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank, @@ -2121,14 +2119,9 @@ mod tests { #[derive(Clone)] struct MinDecodeSelector { rendezvous: Option>, - exclusive_affinity_target: bool, } impl WorkerSelector for MinDecodeSelector { - fn uses_exclusive_affinity_target(&self) -> bool { - self.exclusive_affinity_target - } - fn required_worker_inputs(&self) -> WorkerInputs { WorkerInputs::CACHE | WorkerInputs::LOAD } @@ -2206,41 +2199,8 @@ mod tests { Arc>, Arc>, ) { - let configs = (0..num_workers as u64) - .map(|id| { - ( - id, - SimpleWorkerConfig { - max_num_batched_tokens: Some(isl as u64), - ..Default::default() - }, - ) - }) - .collect(); - make_queue_with_custom_configs_and_selector(configs, block_size, threshold_frac, selector) - } - - #[allow(clippy::type_complexity)] - fn make_queue_with_custom_configs_and_selector< - Sel: WorkerSelector + Send + 'static, - >( - configs: HashMap, - block_size: u32, - threshold_frac: Option, - selector: Sel, - ) -> ( - Arc>, - Arc>, - ) { - let dp_range: HashMap = configs - .iter() - .map(|(&id, config)| { - ( - id, - (config.data_parallel_start_rank, config.data_parallel_size), - ) - }) - .collect(); + let dp_range: HashMap = + (0..num_workers as u64).map(|id| (id, (0, 1))).collect(); let slots = Arc::new(ActiveSequencesMultiWorker::new( NoopSequencePublisher, block_size as usize, @@ -2249,6 +2209,17 @@ mod tests { 0, "test", )); + + let mut configs: HashMap = HashMap::new(); + for id in 0..num_workers as u64 { + configs.insert( + id, + SimpleWorkerConfig { + max_num_batched_tokens: Some(isl as u64), + ..Default::default() + }, + ); + } let (_cfg_tx, cfg_rx) = watch::channel(configs); let queue = Arc::new(SchedulerQueue::new( @@ -2652,101 +2623,6 @@ mod tests { (req, rx) } - #[tokio::test] - async fn affinity_target_narrows_candidates_only_when_selector_opts_in() { - let target = WorkerWithDpRank::new(1, 0); - for (exclusive_affinity_target, expected) in - [(false, WorkerWithDpRank::new(0, 0)), (true, target)] - { - let (queue, _slots) = make_queue_with_custom_selector( - 2, - 16, - 64, - None, - MinDecodeSelector { - rendezvous: None, - exclusive_affinity_target, - }, - ); - let (mut request, _response_rx) = make_request("soft-affinity", 64); - request.affinity_target = Some(target.into()); - - let selected = queue.select_without_admission(request).await.unwrap(); - - assert_eq!(selected.response.best_worker, expected); - } - } - - #[tokio::test] - async fn exclusive_affinity_falls_back_when_target_is_disallowed() { - let (queue, _slots) = make_queue_with_custom_selector( - 2, - 16, - 64, - None, - MinDecodeSelector { - rendezvous: None, - exclusive_affinity_target: true, - }, - ); - let (mut request, _response_rx) = make_request("disallowed-affinity", 64); - request.affinity_target = Some(WorkerWithDpRank::new(1, 0).into()); - request.allowed_worker_ids = Some(HashSet::from([0])); - - let selected = queue.select_without_admission(request).await.unwrap(); - - assert_eq!(selected.response.best_worker, WorkerWithDpRank::new(0, 0)); - } - - #[tokio::test] - async fn exclusive_affinity_respects_dp_rank_scope() { - let configs = HashMap::from([ - ( - 7, - SimpleWorkerConfig { - data_parallel_start_rank: 2, - data_parallel_size: 2, - max_num_batched_tokens: Some(64), - ..Default::default() - }, - ), - ( - 8, - SimpleWorkerConfig { - max_num_batched_tokens: Some(64), - ..Default::default() - }, - ), - ]); - let (queue, _slots) = make_queue_with_custom_configs_and_selector( - configs, - 16, - None, - MinDecodeSelector { - rendezvous: None, - exclusive_affinity_target: true, - }, - ); - - for (name, target, expected) in [ - ( - "exact-rank-affinity", - WorkerAffinityTarget::new(7, Some(3)), - WorkerWithDpRank::new(7, 3), - ), - ( - "worker-only-affinity", - WorkerAffinityTarget::new(7, None), - WorkerWithDpRank::new(7, 2), - ), - ] { - let (mut request, _response_rx) = make_request(name, 64); - request.affinity_target = Some(target); - let selected = queue.select_without_admission(request).await.unwrap(); - assert_eq!(selected.response.best_worker, expected); - } - } - #[tokio::test] async fn classifier_overrides_are_validated_before_enqueue() { let profile = policy_profile( @@ -3234,10 +3110,7 @@ policy_classes: 16, 64, None, - MinDecodeSelector { - rendezvous: None, - exclusive_affinity_target: false, - }, + MinDecodeSelector { rendezvous: None }, ); let worker0 = WorkerWithDpRank::new(0, 0); let worker1 = WorkerWithDpRank::new(1, 0); @@ -3285,10 +3158,7 @@ policy_classes: 16, 64, None, - MinDecodeSelector { - rendezvous: None, - exclusive_affinity_target: false, - }, + MinDecodeSelector { rendezvous: None }, ); let observer_gate = Arc::new((StdMutex::new(false), Condvar::new())); let callback_gate = Arc::clone(&observer_gate); @@ -3348,10 +3218,7 @@ policy_classes: 16, 64, None, - MinDecodeSelector { - rendezvous: None, - exclusive_affinity_target: false, - }, + MinDecodeSelector { rendezvous: None }, ); let worker0 = WorkerWithDpRank::new(0, 0); let worker1 = WorkerWithDpRank::new(1, 0); @@ -3598,7 +3465,6 @@ policy_classes: async fn test_concurrent_immediate_admissions_see_prior_booking() { let selector = MinDecodeSelector { rendezvous: Some(Arc::new(SelectorRendezvous::default())), - exclusive_affinity_target: false, }; let (queue, slots) = make_queue_with_custom_selector(2, 16, 512, None, selector); let barrier = Arc::new(Barrier::new(3)); diff --git a/lib/kv-router/src/scheduling/selector/mod.rs b/lib/kv-router/src/scheduling/selector/mod.rs index c82b1e97b0f8..61b4f5c1c33b 100644 --- a/lib/kv-router/src/scheduling/selector/mod.rs +++ b/lib/kv-router/src/scheduling/selector/mod.rs @@ -36,8 +36,8 @@ pub trait WorkerSelector { /// Whether an eligible affinity target exclusively constrains worker selection. /// - /// The default selector uses exclusive affinity. Custom policies receive affinity as advisory - /// context unless they opt into exclusive candidate narrowing. + /// The default selector uses exclusive affinity. Custom policies receive affinity as + /// advisory context and may choose another eligible worker. fn uses_exclusive_affinity_target(&self) -> bool { false } diff --git a/lib/kv-router/src/scheduling/selector/policy.rs b/lib/kv-router/src/scheduling/selector/policy.rs index a4d50659e380..01213e781f67 100644 --- a/lib/kv-router/src/scheduling/selector/policy.rs +++ b/lib/kv-router/src/scheduling/selector/policy.rs @@ -49,7 +49,6 @@ pub(super) struct CustomWorkerSelectionState { pub struct WorkerSelectionPolicy { kv_router_config: KvRouterConfig, worker_label: &'static str, - exclusive_affinity_target: bool, state: WorkerSelectionPolicyState, } @@ -88,7 +87,6 @@ impl WorkerSelectionPolicy { Self { kv_router_config, worker_label, - exclusive_affinity_target: false, state: WorkerSelectionPolicyState::Custom(RefCell::new(CustomWorkerSelectionState { filters, scorers, @@ -112,20 +110,9 @@ impl WorkerSelectionPolicy { Self { kv_router_config, worker_label, - exclusive_affinity_target: true, state: WorkerSelectionPolicyState::Default(picker), } } - - /// Configure whether an eligible affinity target exclusively constrains candidate selection. - /// - /// Custom policies are advisory by default. When enabled, the scheduler narrows the eligible - /// candidate set to the affinity target before filters, scorers, and the picker run. If the - /// target is ineligible, normal selection continues over the remaining eligible candidates. - pub fn with_exclusive_affinity_target(mut self, enabled: bool) -> Self { - self.exclusive_affinity_target = enabled; - self - } } #[inline(always)] @@ -282,7 +269,7 @@ pub(super) fn collect_custom_candidates( impl WorkerSelector for WorkerSelectionPolicy { fn uses_exclusive_affinity_target(&self) -> bool { - self.exclusive_affinity_target + matches!(&self.state, WorkerSelectionPolicyState::Default(_)) } fn required_worker_inputs(&self) -> WorkerInputs { @@ -840,21 +827,12 @@ mod tests { vec![Box::new(LoadScorer)], Box::new(FirstPicker), ); - let advisory_inputs = + let inputs = >::required_worker_inputs( &policy, ); - assert!(!uses_exclusive_affinity(&policy)); - let policy = policy.with_exclusive_affinity_target(true); - let exclusive_inputs = - >::required_worker_inputs( - &policy, - ); - - assert!(uses_exclusive_affinity(&policy)); - assert_eq!(exclusive_inputs, advisory_inputs); - assert!(exclusive_inputs.contains(WorkerInputs::CACHE)); - assert!(exclusive_inputs.contains(WorkerInputs::LOAD)); + assert!(inputs.contains(WorkerInputs::CACHE)); + assert!(inputs.contains(WorkerInputs::LOAD)); } } diff --git a/lib/llm/src/kv_router/routing_host/tests.rs b/lib/llm/src/kv_router/routing_host/tests.rs index 966de6643c6f..25eaa3eb288f 100644 --- a/lib/llm/src/kv_router/routing_host/tests.rs +++ b/lib/llm/src/kv_router/routing_host/tests.rs @@ -13,10 +13,6 @@ use std::{ use crate::kv_router::SelectionPolicySource; use dynamo_kv_router::{ config::KvRouterConfig, - plugins::worker_selection::{ - WorkerInputView, WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicy, - WorkerSelectionPolicyError, - }, protocols::RoutingConstraints, scheduling::{ ClassifierError, ClassifyEvent, ClassifyFuture, ClassifyRequest, RequestClassifier, @@ -471,36 +467,6 @@ fn affinity_request( request } -struct HighestWorkerPicker; - -impl WorkerPicker for HighestWorkerPicker { - fn pick( - &mut self, - _context: &WorkerSelectionContext<'_>, - input: WorkerInputView<'_>, - ) -> Result { - input - .candidates() - .iter() - .enumerate() - .max_by_key(|(_, candidate)| candidate.worker()) - .map(|(row, _)| row) - .ok_or_else(|| WorkerSelectionPolicyError::failed("no eligible worker")) - } -} - -fn highest_worker_policy(exclusive_affinity_target: bool) -> SelectionPolicySource { - SelectionPolicySource::Factory(Arc::new(move |config, worker_type, _| { - WorkerSelectionPolicy::new( - config.clone(), - worker_type.as_str(), - Vec::new(), - Box::new(HighestWorkerPicker), - ) - .with_exclusive_affinity_target(exclusive_affinity_target) - })) -} - #[tokio::test] #[serial_test::serial] async fn builtin_affinity_uses_common_host_for_every_policy() { @@ -1223,39 +1189,6 @@ async fn router_with_worker_configs_and_classifier( session_affinity_ttl: Option, workers: HashMap, classifier: Option, -) -> (RoutingHost, Runtime) { - router_with_worker_configs_classifier_policy_and_mode( - session_affinity_ttl, - workers, - classifier, - SelectionPolicySource::Registry, - crate::session_affinity::SessionAffinityMode::Hard, - ) - .await -} - -async fn router_with_worker_configs_and_policy( - session_affinity_ttl: Option, - workers: HashMap, - policy: SelectionPolicySource, - affinity_mode: crate::session_affinity::SessionAffinityMode, -) -> (RoutingHost, Runtime) { - Box::pin(router_with_worker_configs_classifier_policy_and_mode( - session_affinity_ttl, - workers, - None::, - policy, - affinity_mode, - )) - .await -} - -async fn router_with_worker_configs_classifier_policy_and_mode( - session_affinity_ttl: Option, - workers: HashMap, - classifier: Option, - policy: SelectionPolicySource, - affinity_mode: crate::session_affinity::SessionAffinityMode, ) -> (RoutingHost, Runtime) { let runtime = Runtime::from_current().unwrap(); // Each runtime has its own in-memory discovery store, so fixture namespaces can repeat. @@ -1283,7 +1216,7 @@ async fn router_with_worker_configs_classifier_policy_and_mode( workers, None, 16, - policy, + SelectionPolicySource::Registry, Some(config), None, "decode", @@ -1300,11 +1233,7 @@ async fn router_with_worker_configs_classifier_policy_and_mode( let inner = PushRouter::from_client(client, RouterMode::KV) .await .unwrap(); - let affinity = session_affinity_ttl - .map(|ttl| chooser.affinity_coordinator(ttl, affinity_mode)) - .transpose() - .unwrap(); - let mut router = RoutingHost::new_with_coordinator(inner, Arc::new(chooser), affinity); + let mut router = RoutingHost::new(inner, Arc::new(chooser), session_affinity_ttl).unwrap(); // Parallel tests must not share the process-global request counters. router.request_metrics = router.request_metrics.with_isolated_counters_for_test(); router @@ -1315,69 +1244,6 @@ async fn router_with_worker_configs_classifier_policy_and_mode( (router, runtime) } -#[tokio::test] -async fn policy_affinity_control_selects_and_rebinds_soft_sessions() { - const TARGET_WORKER: u64 = 7; - const PREFERRED_WORKER: u64 = 8; - - for (exclusive_affinity_target, expected) in [(false, PREFERRED_WORKER), (true, TARGET_WORKER)] - { - let workers = HashMap::from([ - (TARGET_WORKER, ModelRuntimeConfig::default()), - (PREFERRED_WORKER, ModelRuntimeConfig::default()), - ]); - let (router, runtime) = router_with_worker_configs_and_policy( - Some(Duration::from_secs(10)), - workers, - highest_worker_policy(exclusive_affinity_target), - crate::session_affinity::SessionAffinityMode::Soft, - ) - .await; - let session_id = SessionAffinityId::new(format!("policy-{exclusive_affinity_target}")); - bind_affinity_target(&router, &session_id, AffinityTarget::worker(TARGET_WORKER)).await; - let mut request = Context::new(request()); - request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone()); - - let (selection, operation) = router - .select_with_affinity( - &request, - RequestPhase::Aggregated, - false, - &CleanupBudget::default(), - ) - .await - .unwrap(); - assert_eq!(selection.worker.worker_id, expected); - - let output = Annotated::from_data(LLMEngineOutput { - finish_reason: Some(FinishReason::Stop), - ..Default::default() - }); - let stream = ResponseStream::new( - Box::pin(stream::once(async move { output })), - request.context().clone(), - ); - let mut stream = router - .bind_affinity(operation, route_target(selection.worker), stream) - .unwrap(); - while stream.next().await.is_some() {} - router.kv_router().free(request.id()).await.unwrap(); - - assert_eq!( - router - .affinity - .as_ref() - .unwrap() - .query_target(&session_id, None) - .unwrap(), - Some(AffinityTarget::worker(expected)) - ); - - drop(router); - runtime.shutdown(); - } -} - /// A dispatch plane that yields before completing, unlike [`CompletedBuiltinDispatch`]. /// A recorded worker id therefore means the dispatch survived, not merely started. #[derive(Default)] diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index dff50302282e..0750bc389530 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -36,8 +36,7 @@ pub fn register( #[cfg(test)] mod tests { use dynamo_kv_router::plugins::worker_selection::WorkerSelectionPolicyFactory; - use dynamo_kv_router::protocols::WorkerConfigLike; - use dynamo_kv_router::{KvRouterConfig, RoutingPartitionRef, WorkerSelector, WorkerType}; + use dynamo_kv_router::{KvRouterConfig, RoutingPartitionRef, WorkerType}; use super::*; @@ -62,41 +61,6 @@ mod tests { (config, resolved) } - struct PolicyTestWorker; - - impl WorkerConfigLike for PolicyTestWorker { - fn data_parallel_start_rank(&self) -> u32 { - 0 - } - - fn data_parallel_size(&self) -> u32 { - 1 - } - - fn max_num_batched_tokens(&self) -> Option { - None - } - - fn total_kv_blocks(&self) -> Option { - None - } - } - - fn resolves_exclusive_affinity(yaml: &str) -> bool { - let (config, resolved) = resolve(yaml); - let factory = resolved - .unwrap() - .expect("a configured instance resolves to a factory"); - let policy = factory( - &config, - WorkerType::Aggregated, - RoutingPartitionRef::new("model", "default"), - ); - >::uses_exclusive_affinity_target( - &policy, - ) - } - /// Catches a policy type name that drifts from its documentation, and proves the documented /// instance shape constructs for every stage it selects. #[test] @@ -141,13 +105,17 @@ worker_selection: ) }; - assert!(!resolves_exclusive_affinity(&yaml(""))); - assert!(!resolves_exclusive_affinity(&yaml( - " parameters:\n respect_soft_affinity: false" - ))); - assert!(resolves_exclusive_affinity(&yaml( - " parameters:\n respect_soft_affinity: true" - ))); + for parameter in [ + "", + " parameters:\n respect_soft_affinity: false", + " parameters:\n respect_soft_affinity: true", + ] { + let (_config, resolved) = resolve(&yaml(parameter)); + assert!( + resolved.unwrap().is_some(), + "parameter variant should resolve: {parameter}" + ); + } } /// An unknown parameter key is a mistake, most often a misremembered threshold name. It must diff --git a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs index e26f5612cc98..bd1f634e21f4 100644 --- a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs +++ b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs @@ -64,7 +64,7 @@ struct Parameters { balance_abs_threshold: usize, /// Minimum ratio of largest to smallest active-request count before the load tier applies. balance_rel_threshold: f64, - /// Whether an eligible soft-affinity target exclusively constrains candidate selection. + /// Whether to retain an eligible soft-affinity target before considering other workers. respect_soft_affinity: bool, } @@ -99,27 +99,28 @@ fn least_loaded(load: &[WorkerLoadInput], rows: impl Iterator) -> rows.min_by_key(|&row| load[row].active_requests()) } -fn select_row( +fn select_row_where( parameters: &Parameters, cache: &[WorkerCacheInput], load: &[WorkerLoadInput], request_blocks: u64, + include: impl Fn(usize) -> bool + Copy, ) -> Option { if cache.is_empty() || cache.len() != load.len() { return None; } - let min_load = load.iter().map(|item| item.active_requests()).min()?; - let max_load = load.iter().map(|item| item.active_requests()).max()?; + let rows = || (0..load.len()).filter(|&row| include(row)); + let min_load = rows().map(|row| load[row].active_requests()).min()?; + let max_load = rows().map(|row| load[row].active_requests()).max()?; if max_load.saturating_sub(min_load) > parameters.balance_abs_threshold && (max_load as f64) > parameters.balance_rel_threshold * (min_load as f64) { - return least_loaded(load, 0..load.len()); + return least_loaded(load, rows()); } - let max_overlap = cache - .iter() - .map(|item| item.device_overlap_blocks()) + let max_overlap = rows() + .map(|row| cache[row].device_overlap_blocks()) .max_by(f64::total_cmp)?; let cache_ratio = if request_blocks == 0 { 0.0 @@ -129,13 +130,20 @@ fn select_row( if cache_ratio > parameters.cache_threshold { return least_loaded( load, - cache.iter().enumerate().filter_map(|(row, item)| { - (item.device_overlap_blocks() == max_overlap).then_some(row) - }), + rows().filter(|&row| cache[row].device_overlap_blocks() == max_overlap), ); } - least_loaded(load, 0..load.len()) + least_loaded(load, rows()) +} + +fn select_row( + parameters: &Parameters, + cache: &[WorkerCacheInput], + load: &[WorkerLoadInput], + request_blocks: u64, +) -> Option { + select_row_where(parameters, cache, load, request_blocks, |_| true) } struct TwoTierCostFnPicker { @@ -158,6 +166,24 @@ impl WorkerPicker for TwoTierCostFnPicker { let load = input .load() .ok_or_else(|| WorkerSelectionPolicyError::failed("load input unavailable"))?; + + if self.parameters.respect_soft_affinity + && let Some(target) = context.affinity_target() + && let Some(row) = select_row_where( + &self.parameters, + cache, + load, + context.request_blocks(), + |row| { + let worker = input.candidates()[row].worker(); + worker.worker_id == target.worker_id + && target.dp_rank.is_none_or(|rank| worker.dp_rank == rank) + }, + ) + { + return Ok(row); + } + select_row(&self.parameters, cache, load, context.request_blocks()) .ok_or_else(|| WorkerSelectionPolicyError::failed("no eligible worker")) } @@ -177,7 +203,6 @@ fn provider( Vec::new(), Box::new(TwoTierCostFnPicker { parameters }), ) - .with_exclusive_affinity_target(parameters.respect_soft_affinity) }, )) } @@ -192,7 +217,9 @@ pub fn register( mod tests { use std::collections::HashMap; - use dynamo_kv_router::protocols::{RoutingConstraints, WorkerConfigLike, WorkerWithDpRank}; + use dynamo_kv_router::protocols::{ + RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, WorkerWithDpRank, + }; use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode}; use dynamo_kv_router::{ SchedulingRequest, WorkerLoadProjection, WorkerSelectionInput, WorkerSelector, @@ -233,13 +260,21 @@ mod tests { } fn select_with(parameters: Parameters, workers: [(u64, usize, usize); 2]) -> WorkerWithDpRank { + select_with_affinity(parameters, workers, None) + } + + fn select_with_affinity( + parameters: Parameters, + workers: [(u64, usize, usize); 2], + affinity_target: Option, + ) -> WorkerWithDpRank { let mut request = SchedulingRequest { mode: ScheduleMode::QueryOnly { request_id: None }, token_seq: None, isl_tokens: TEN_BLOCKS, lora_name: None, expected_output_tokens: None, - affinity_target: None, + affinity_target, pinned_worker: None, allowed_worker_ids: None, routing_constraints: RoutingConstraints::default(), @@ -312,6 +347,37 @@ mod tests { assert_eq!(select_with(tuned, workers), worker(B)); } + #[test] + fn soft_affinity_is_policy_owned() { + let workers = [(A, 0, 0), (B, 6, 4)]; + let target = Some(WorkerAffinityTarget::new(A, None)); + + assert_eq!( + select_with_affinity(Parameters::default(), workers, target), + worker(B), + "the default ignores advisory soft affinity" + ); + + let respecting = Parameters { + respect_soft_affinity: true, + ..Parameters::default() + }; + assert_eq!( + select_with_affinity(respecting, workers, target), + worker(A), + "an eligible target is retained by the picker" + ); + assert_eq!( + select_with_affinity( + respecting, + workers, + Some(WorkerAffinityTarget::new(999, None)), + ), + worker(B), + "an absent target falls back to normal two-tier selection" + ); + } + #[test] fn rejects_out_of_range_parameters() { let cache = |v| { From 57be4852dbd67866368ef622d5366cb59a7179ef Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Mon, 21 Sep 2026 17:46:34 -0700 Subject: [PATCH 4/7] test(router): pass affinity through select helper Signed-off-by: Guan Luo --- .../builtin/src/two_tier_cost_fn.rs | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs index bd1f634e21f4..7832f5f76dee 100644 --- a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs +++ b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs @@ -256,14 +256,10 @@ mod tests { /// Select among workers given as `(worker_id, device_overlap_blocks, active_requests)`. fn select(workers: [(u64, usize, usize); 2]) -> WorkerWithDpRank { - select_with(Parameters::default(), workers) + select_with(Parameters::default(), workers, None) } - fn select_with(parameters: Parameters, workers: [(u64, usize, usize); 2]) -> WorkerWithDpRank { - select_with_affinity(parameters, workers, None) - } - - fn select_with_affinity( + fn select_with( parameters: Parameters, workers: [(u64, usize, usize); 2], affinity_target: Option, @@ -344,7 +340,7 @@ mod tests { cache_threshold: 0.2, ..Parameters::default() }; - assert_eq!(select_with(tuned, workers), worker(B)); + assert_eq!(select_with(tuned, workers, None), worker(B)); } #[test] @@ -353,7 +349,7 @@ mod tests { let target = Some(WorkerAffinityTarget::new(A, None)); assert_eq!( - select_with_affinity(Parameters::default(), workers, target), + select_with(Parameters::default(), workers, target), worker(B), "the default ignores advisory soft affinity" ); @@ -363,12 +359,12 @@ mod tests { ..Parameters::default() }; assert_eq!( - select_with_affinity(respecting, workers, target), + select_with(respecting, workers, target), worker(A), "an eligible target is retained by the picker" ); assert_eq!( - select_with_affinity( + select_with( respecting, workers, Some(WorkerAffinityTarget::new(999, None)), From a704a75f0d091718ae8c2bfafcb87ea7a21efde5 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Mon, 21 Sep 2026 18:08:40 -0700 Subject: [PATCH 5/7] test(router): cover rank-specific soft affinity --- .../builtin/src/two_tier_cost_fn.rs | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs index 7832f5f76dee..e4d67a373025 100644 --- a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs +++ b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs @@ -372,6 +372,26 @@ mod tests { worker(B), "an absent target falls back to normal two-tier selection" ); + + let rank_sensitive_workers = [(A, 0, 0), (B, 0, 4)]; + assert_eq!( + select_with( + respecting, + rank_sensitive_workers, + Some(WorkerAffinityTarget::new(B, Some(0))), + ), + worker(B), + "a matching rank-specific target is retained" + ); + assert_eq!( + select_with( + respecting, + rank_sensitive_workers, + Some(WorkerAffinityTarget::new(B, Some(1))), + ), + worker(A), + "a missing target rank falls back instead of retaining another rank" + ); } #[test] From e5eb8650aebe8fc394db18716b13e7deb577672f Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Wed, 23 Sep 2026 14:36:50 -0700 Subject: [PATCH 6/7] fix(router): adapt two-tier policy to cache input view Signed-off-by: Guan Luo --- lib/router-plugins/builtin/src/two_tier_cost_fn.rs | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs index e2ed888e292b..cab66b8fb775 100644 --- a/lib/router-plugins/builtin/src/two_tier_cost_fn.rs +++ b/lib/router-plugins/builtin/src/two_tier_cost_fn.rs @@ -110,6 +110,12 @@ fn select_row_where( return None; } + let device_overlap = |row| { + cache + .get(row) + .expect("cache and load lengths were checked") + .device_overlap_blocks() + }; let rows = || (0..load.len()).filter(|&row| include(row)); let min_load = rows().map(|row| load[row].active_requests()).min()?; let max_load = rows().map(|row| load[row].active_requests()).max()?; @@ -119,9 +125,7 @@ fn select_row_where( return least_loaded(load, rows()); } - let max_overlap = rows() - .map(|row| cache[row].device_overlap_blocks()) - .max_by(f64::total_cmp)?; + let max_overlap = rows().map(device_overlap).max_by(f64::total_cmp)?; let cache_ratio = if request_blocks == 0 { 0.0 } else { @@ -130,7 +134,7 @@ fn select_row_where( if cache_ratio > parameters.cache_threshold { return least_loaded( load, - rows().filter(|&row| cache[row].device_overlap_blocks() == max_overlap), + rows().filter(|&row| device_overlap(row) == max_overlap), ); } @@ -139,7 +143,7 @@ fn select_row_where( fn select_row( parameters: &Parameters, - cache: &[WorkerCacheInput], + cache: WorkerCacheInputs<'_>, load: &[WorkerLoadInput], request_blocks: u64, ) -> Option { From 93e76bffd06f521f299b0a69b587346f1f35f9eb Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Tue, 29 Sep 2026 00:07:18 -0700 Subject: [PATCH 7/7] test(router): remove duplicate policy resolution case Signed-off-by: Guan Luo --- lib/router-plugins/builtin/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index a42a50d5a4f3..a1b19f9698ae 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -106,7 +106,6 @@ worker_selection: }; for parameter in [ - "", " parameters:\n respect_soft_affinity: false", " parameters:\n respect_soft_affinity: true", ] {