From 2103548548c29e6a5700245f5da64052ee5bd840 Mon Sep 17 00:00:00 2001 From: "WENZHI CHEN (SII)" Date: Tue, 22 Sep 2026 19:15:18 -0400 Subject: [PATCH] feat(router): ship opt-in soft affinity load guard Retain an eligible soft-affinity target until it exceeds a configured active-request limit and another worker has a sufficient load advantage. Register the policy in the shipped catalog and document its YAML configuration. Signed-off-by: WENZHI CHEN (SII) --- .../router/configuration-and-tuning.md | 31 +++ lib/router-plugins/builtin/src/lib.rs | 26 +- .../builtin/src/soft_affinity_load_guard.rs | 261 ++++++++++++++++++ 3 files changed, 317 insertions(+), 1 deletion(-) create mode 100644 lib/router-plugins/builtin/src/soft_affinity_load_guard.rs 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 d93b3a217a68..8c7db1f6f4d6 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 @@ -112,6 +112,7 @@ link no catalog and reject a configured policy type at startup. |---|---| | `default` | Dynamo's built-in selector and cost model. Reserved; always available. | | `dynamo-two-tier-cost-fn` | Ranks on two tiers instead of one additive cost: active-request load first, then device-KV prefix overlap. Prefers the worker holding the largest prefix overlap unless load is badly imbalanced. Thresholds and selection order ported from the experimental SGLang router's `cache_aware_zmq` policy. Thresholds are tunable; the defaults reproduce it exactly. | +| `dynamo-soft-affinity-load-guard` | Retains an eligible soft-affinity target until its active-request count exceeds a threshold and another worker has a sufficient load advantage. Without a target, selects the least-loaded eligible worker. | Write the instance into the same YAML file that `--router-policy-config` already points at: @@ -160,6 +161,36 @@ 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. +#### Guard Soft Affinity Under Load + +Use `dynamo-soft-affinity-load-guard` with soft session affinity to retain a session's worker while +its active-request count is at or below `max_active_requests`. Above that threshold, the policy +moves only when the least-loaded alternative has at least `move_margin` fewer active requests. +The same rule applies when requests carrying a parent session ID share the parent-group binding. + +```yaml +worker_selection: + aggregated: soft-guard + instances: + - name: soft-guard + type: dynamo-soft-affinity-load-guard + parameters: + max_active_requests: 32 + move_margin: 2 +``` + +Run the frontend with `--router-mode kv`, `--router-session-affinity-ttl-secs 60`, +`--router-session-affinity-mode soft`, and `--router-policy-config` pointing to that YAML file. +Both parameters are optional; the values above are the defaults. `move_margin` must be at least +one. When the affinity target is missing from the eligible candidate set, the policy falls back +to the least-loaded worker. If no better alternative exists, it retains the target. Equal-load +workers never trigger a move. Hard affinity and explicit worker targets remain exact and bypass +this advisory choice. + +The active-request count is a routing load signal, not a prediction of queue wait or remaining +work. This policy does not inspect cache residency. Moving to a cold worker may recompute a prefix; +choose thresholds with a workload comparison of request latency and cached-token usage. + #### Override the Selection `DYN_ROUTER_WORKER_SELECTION_POLICY` overrides every stage. `--router-prefill-policy` and diff --git a/lib/router-plugins/builtin/src/lib.rs b/lib/router-plugins/builtin/src/lib.rs index 06e8993e9303..83eecd77c554 100644 --- a/lib/router-plugins/builtin/src/lib.rs +++ b/lib/router-plugins/builtin/src/lib.rs @@ -18,6 +18,7 @@ //! beyond `dynamo-kv-router`, put it behind its own default-on Cargo feature so a build can drop //! it; every policy registered here is compiled into every artifact that links this crate. +mod soft_affinity_load_guard; mod two_tier_cost_fn; use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegistryError}; @@ -30,7 +31,8 @@ use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegis pub fn register( registry: &mut RouterPluginRegistry, ) -> Result<(), WorkerSelectionPolicyRegistryError> { - two_tier_cost_fn::register(registry) + two_tier_cost_fn::register(registry)?; + soft_affinity_load_guard::register(registry) } #[cfg(test)] @@ -90,6 +92,28 @@ worker_selection: } } + #[test] + fn resolves_soft_affinity_load_guard_yaml() { + let (config, resolved) = resolve( + r#" +worker_selection: + aggregated: soft-guard + instances: + - name: soft-guard + type: dynamo-soft-affinity-load-guard + parameters: + max_active_requests: 7 + move_margin: 2 +"#, + ); + let factory = resolved.unwrap().expect("the shipped policy resolves"); + factory( + &config, + WorkerType::Aggregated, + RoutingPartitionRef::new("model", "default"), + ); + } + /// 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] diff --git a/lib/router-plugins/builtin/src/soft_affinity_load_guard.rs b/lib/router-plugins/builtin/src/soft_affinity_load_guard.rs new file mode 100644 index 000000000000..8576421174c5 --- /dev/null +++ b/lib/router-plugins/builtin/src/soft_affinity_load_guard.rs @@ -0,0 +1,261 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Retain an eligible soft-affinity target until its active load warrants moving. +//! +//! Dynamo owns the affinity binding. This picker only selects an eligible worker; a successful +//! dispatch updates a soft binding through the normal coordinator path. + +use std::sync::Arc; + +use dynamo_kv_router::plugins::worker_selection::WorkerSelectionPolicyFactory; +use dynamo_kv_router::plugins::{ + RouterPluginRegistry, WorkerSelectionPolicyParameters, WorkerSelectionPolicyProviderError, + WorkerSelectionPolicyRegistryError, +}; +use dynamo_kv_router::protocols::WorkerAffinityTarget; +use dynamo_kv_router::{ + KvRouterConfig, ScoredWorkerCandidate, WorkerInputView, WorkerInputs, WorkerLoadInput, + WorkerPicker, WorkerSelectionContext, WorkerSelectionPolicy, WorkerSelectionPolicyError, +}; + +pub const POLICY_TYPE: &str = "dynamo-soft-affinity-load-guard"; + +#[derive(Clone, Copy, serde::Deserialize)] +#[serde(deny_unknown_fields, default)] +struct Parameters { + max_active_requests: usize, + move_margin: usize, +} + +impl Default for Parameters { + fn default() -> Self { + Self { + max_active_requests: 32, + move_margin: 2, + } + } +} + +impl Parameters { + fn validate(self) -> Result { + if self.move_margin == 0 { + return Err(WorkerSelectionPolicyProviderError::new( + "move_margin must be at least 1", + )); + } + Ok(self) + } +} + +struct SoftAffinityLoadGuardPicker { + parameters: Parameters, +} + +fn matches_target(candidate: &ScoredWorkerCandidate, target: WorkerAffinityTarget) -> bool { + let worker = candidate.worker(); + worker.worker_id == target.worker_id && target.dp_rank.is_none_or(|rank| worker.dp_rank == rank) +} + +fn least_loaded( + candidates: &[ScoredWorkerCandidate], + loads: &[WorkerLoadInput], + include: impl Fn(&ScoredWorkerCandidate) -> bool, +) -> Option { + candidates + .iter() + .zip(loads) + .enumerate() + .filter(|(_, (candidate, _))| include(candidate)) + .min_by_key(|(_, (candidate, load))| (load.active_requests(), candidate.worker())) + .map(|(row, _)| row) +} + +impl WorkerPicker for SoftAffinityLoadGuardPicker { + fn required_worker_inputs(&self) -> WorkerInputs { + WorkerInputs::LOAD + } + + fn pick( + &mut self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + let candidates = input.candidates(); + let loads = input + .load() + .ok_or_else(|| WorkerSelectionPolicyError::failed("active load input unavailable"))?; + + if let Some(target) = context.affinity_target() + && let Some(target_row) = least_loaded(candidates, loads, |candidate| { + matches_target(candidate, target) + }) + { + let target_load = loads[target_row].active_requests(); + if target_load <= self.parameters.max_active_requests { + return Ok(target_row); + } + + if let Some(alternative_row) = least_loaded(candidates, loads, |candidate| { + !matches_target(candidate, target) + }) { + let alternative_load = loads[alternative_row].active_requests(); + if alternative_load < target_load + && target_load - alternative_load >= self.parameters.move_margin + { + return Ok(alternative_row); + } + } + return Ok(target_row); + } + + least_loaded(candidates, loads, |_| true) + .ok_or_else(|| WorkerSelectionPolicyError::failed("no eligible worker")) + } +} + +fn provider( + parameters: &WorkerSelectionPolicyParameters, +) -> Result { + let parameters: Parameters = parameters.deserialize()?; + let parameters = parameters.validate()?; + + Ok(Arc::new( + move |config: &KvRouterConfig, worker_type, _partition| { + WorkerSelectionPolicy::new( + config.clone(), + worker_type.as_str(), + Vec::new(), + Box::new(SoftAffinityLoadGuardPicker { parameters }), + ) + }, + )) +} + +pub fn register( + registry: &mut RouterPluginRegistry, +) -> Result<(), WorkerSelectionPolicyRegistryError> { + registry.register_worker_selection(POLICY_TYPE, Arc::new(provider)) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use dynamo_kv_router::protocols::{RoutingConstraints, WorkerConfigLike, WorkerWithDpRank}; + use dynamo_kv_router::scheduling::{OverlapSignals, ScheduleMode}; + use dynamo_kv_router::{ + SchedulingRequest, WorkerLoadProjection, WorkerSelectionInput, WorkerSelector, + }; + + use super::*; + + struct TestWorker; + + impl WorkerConfigLike for TestWorker { + 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 { + Some(1024) + } + } + + fn worker(id: u64) -> WorkerWithDpRank { + WorkerWithDpRank::from_worker_id(id) + } + + fn select(parameters: Parameters, target: Option, a_load: usize, b_load: usize) -> u64 { + let mut request = SchedulingRequest { + mode: ScheduleMode::QueryOnly { request_id: None }, + token_seq: None, + isl_tokens: 16, + lora_name: None, + expected_output_tokens: None, + affinity_target: target.map(|id| worker(id).into()), + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), + router_config_override: None, + track_prefill_tokens: true, + priority_jump: 0.0, + strict_priority: 0, + policy_class: None, + session_context: None, + overlap: OverlapSignals::default(), + kv_transfer_candidates: None, + retain_kv_transfer_chain: false, + shared_cache_hits: None, + worker_loads: Default::default(), + resp_tx: None, + }; + for (id, active_requests) in [(7, a_load), (9, b_load)] { + request.worker_loads.insert( + worker(id), + WorkerLoadProjection { + active_requests, + ..Default::default() + }, + ); + } + let workers = HashMap::from([(7, TestWorker), (9, TestWorker)]); + WorkerSelectionPolicy::new( + KvRouterConfig::default(), + "test", + Vec::new(), + Box::new(SoftAffinityLoadGuardPicker { parameters }), + ) + .select_worker(WorkerSelectionInput::configured( + &workers, + &request, + request.eligibility(), + 16, + )) + .unwrap() + .worker + .worker_id + } + + #[test] + fn retains_the_group_until_both_load_conditions_hold() { + let parameters = Parameters { + max_active_requests: 7, + move_margin: 2, + }; + assert_eq!(select(parameters, Some(7), 7, 0), 7); + assert_eq!(select(parameters, Some(7), 8, 7), 7); + assert_eq!(select(parameters, Some(7), 8, 6), 9); + assert_eq!(select(parameters, Some(9), 6, 8), 7); + assert_eq!(select(parameters, Some(9), 7, 8), 9); + } + + #[test] + fn selects_an_eligible_worker_when_no_group_target_exists() { + let parameters = Parameters::default(); + assert_eq!(select(parameters, None, 3, 0), 9); + assert_eq!(select(parameters, Some(99), 3, 0), 9); + assert_eq!(select(parameters, Some(7), 33, 33), 7); + assert_eq!(select(parameters, Some(7), 33, 34), 7); + } + + #[test] + fn rejects_zero_move_margin() { + assert!( + Parameters { + move_margin: 0, + ..Parameters::default() + } + .validate() + .is_err() + ); + } +}