diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 0d43bbc0ba91..ff5e67a21bc2 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -450,6 +450,7 @@ impl RouterHandles { /// /// When `allowed_worker_ids` is Some, only workers in that set are considered. /// Returns worker_id on success. + #[expect(clippy::too_many_arguments)] async fn query_prefill_worker( &self, tokens: &[u32], @@ -458,6 +459,7 @@ impl RouterHandles { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, ) -> Result<(u64, Option), QueryRouterResult> { if let Some(ref ids) = allowed_worker_ids { self.prefill_router.register_workers(ids); @@ -471,6 +473,7 @@ impl RouterHandles { lora_name, priority_jump, allowed_worker_ids, + routing_constraints, ) .await .map_err(|e| { @@ -500,6 +503,7 @@ impl RouterHandles { is_disaggregated: bool, priority_jump: f64, allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, ) -> Result<(WorkerWithDpRank, u32), QueryRouterResult> { if let Some(ref ids) = allowed_worker_ids { self.decode_router.register_workers(ids); @@ -529,6 +533,7 @@ impl RouterHandles { priority_jump, None, allowed_worker_ids, + routing_constraints, ) .await .map_err(|e| { @@ -1124,7 +1129,7 @@ pub unsafe extern "C" fn free_routing_result(result: *mut CRoutingResult) { /// Parse a JSON request string, apply the chat template, tokenize, and lift /// the router-relevant `priority_jump` out of `nvext.agent_hints.priority`. /// -/// Returns `(token_ids, priority_jump)` on success, or a `QueryRouterResult` +/// Returns `(token_ids, priority_jump, routing_constraints)` on success, or a `QueryRouterResult` /// error code. `priority_jump` is `0.0` when no hint is present. This mirrors /// the standalone Dynamo preprocessor lift in `lib/llm/src/preprocessor.rs` /// so the GAIE/EPP path produces the same queue ordering as a non-EPP @@ -1132,7 +1137,7 @@ pub unsafe extern "C" fn free_routing_result(result: *mut CRoutingResult) { unsafe fn preprocess_request( handles: &RouterHandles, request_json: *const c_char, -) -> Result<(Vec, f64), QueryRouterResult> { +) -> Result<(Vec, f64, RoutingConstraints), QueryRouterResult> { let preprocessor = match &handles.preprocessor { Some(p) => p, None => { @@ -1156,6 +1161,11 @@ unsafe fn preprocess_request( }; let priority_jump = extract_priority_jump(&request); + let routing_constraints = request + .nvext + .as_ref() + .and_then(|nvext| nvext.routing_constraints.clone()) + .unwrap_or_default(); let formatted_prompt = match preprocessor.apply_template(&request) { Ok(Some(prompt)) => prompt, @@ -1182,7 +1192,7 @@ unsafe fn preprocess_request( "[EPP-TOKENIZE] Tokenized prompt in C bindings (this is the ONLY tokenization)" ); - Ok((token_ids, priority_jump)) + Ok((token_ids, priority_jump, routing_constraints)) } /// Parse pods JSON into an optional set of allowed worker IDs. @@ -1268,10 +1278,11 @@ pub unsafe extern "C" fn route_prefill_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump) = match unsafe { preprocess_request(handles, request_json) } { - Ok(t) => t, - Err(code) => return code, - }; + let (tokens, priority_jump, routing_constraints) = + match unsafe { preprocess_request(handles, request_json) } { + Ok(t) => t, + Err(code) => return code, + }; let allowed_worker_ids = unsafe { parse_pods_filter(pods_json) }; @@ -1284,6 +1295,7 @@ pub unsafe extern "C" fn route_prefill_request( None, priority_jump, allowed_worker_ids, + routing_constraints, ) .await?; @@ -1346,16 +1358,23 @@ pub unsafe extern "C" fn route_decode_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump) = match unsafe { preprocess_request(handles, request_json) } { - Ok(t) => t, - Err(code) => return code, - }; + let (tokens, priority_jump, routing_constraints) = + match unsafe { preprocess_request(handles, request_json) } { + Ok(t) => t, + Err(code) => return code, + }; let allowed_worker_ids = unsafe { parse_pods_filter(pods_json) }; let result = handles.runtime.secondary().block_on(async { let (decode_worker, _overlap_blocks) = handles - .query_decode_worker(&tokens, is_disaggregated, priority_jump, allowed_worker_ids) + .query_decode_worker( + &tokens, + is_disaggregated, + priority_jump, + allowed_worker_ids, + routing_constraints, + ) .await?; tracing::info!( diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 6ddf9f230595..4a26a16153ed 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -39,7 +39,7 @@ use dynamo_llm::{self as llm_rs}; use crate::llm::entrypoint::RouterConfig as PyRouterConfig; -use crate::llm::local_model::ModelRuntimeConfig; +use crate::llm::local_model::{ModelRuntimeConfig, RoutingConstraints}; use crate::llm::preprocessor::{MediaDecoder, MediaFetcher}; #[pyclass(eq, eq_int)] @@ -183,6 +183,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; // Internal: only in _internal, not public API m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 5c7243276559..b05f0b677a6a 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -9,6 +9,7 @@ use std::sync::atomic::AtomicU32; use std::sync::mpsc; use tokio_stream::StreamExt; +use super::local_model::RoutingConstraints; use super::*; use crate::Endpoint; #[cfg(feature = "kv-indexer")] @@ -957,7 +958,7 @@ impl KvRouter { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (token_ids, model, stop_conditions=None, sampling_options=None, output_options=None, router_config_override=None, worker_id=None, dp_rank=None, extra_args=None, block_mm_infos=None, multi_modal_data=None, mm_routing_info=None))] + #[pyo3(signature = (token_ids, model, stop_conditions=None, sampling_options=None, output_options=None, router_config_override=None, worker_id=None, dp_rank=None, extra_args=None, block_mm_infos=None, multi_modal_data=None, mm_routing_info=None, routing_constraints=None))] fn generate<'p>( &self, py: Python<'p>, @@ -973,6 +974,7 @@ impl KvRouter { block_mm_infos: Option, multi_modal_data: Option, mm_routing_info: Option, + routing_constraints: Option, ) -> PyResult> { // Depythonize the options with defaults let stop_conditions: StopConditions = if let Some(obj) = stop_conditions { @@ -1048,10 +1050,11 @@ impl KvRouter { .tracker(Some(tracker.clone())); // Set routing hints if worker_id or dp_rank is provided - if worker_id.is_some() || dp_rank.is_some() { + if worker_id.is_some() || dp_rank.is_some() || routing_constraints.is_some() { let routing = llm_rs::protocols::common::preprocessor::RoutingHints { backend_instance_id: worker_id, dp_rank, + routing_constraints: routing_constraints.map(Into::into), ..Default::default() }; request_builder.routing(Some(routing)); @@ -1087,7 +1090,7 @@ impl KvRouter { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (token_ids, router_config_override=None, request_id=None, update_indexer=false, block_mm_infos=None, lora_name=None))] + #[pyo3(signature = (token_ids, router_config_override=None, request_id=None, update_indexer=false, block_mm_infos=None, lora_name=None, routing_constraints=None))] fn best_worker<'p>( &self, py: Python<'p>, @@ -1097,6 +1100,7 @@ impl KvRouter { update_indexer: bool, block_mm_infos: Option, lora_name: Option, + routing_constraints: Option, ) -> PyResult> { let router_config_override = if let Some(obj) = router_config_override { let override_config: RouterConfigOverride = @@ -1125,6 +1129,7 @@ impl KvRouter { 0.0, None, None, // allowed_worker_ids: pass via RoutingHints in PreprocessedRequest path + routing_constraints.map(Into::into).unwrap_or_default(), ) .await .map_err(to_pyerr)?; diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index 4817b9152032..e03cd9eefb6e 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -1,10 +1,55 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::{HashMap, HashSet}; + use super::*; +use dynamo_kv_router::protocols::RoutingConstraints as RsRoutingConstraints; use llm_rs::local_model::runtime_config::DisaggregatedEndpoint as RsDisaggregatedEndpoint; use llm_rs::local_model::runtime_config::ModelRuntimeConfig as RsModelRuntimeConfig; +#[pyclass] +#[derive(Clone, Debug, Default)] +pub struct RoutingConstraints { + #[pyo3(get, set)] + pub required_taints: HashSet, + #[pyo3(get, set)] + pub preferred_taints: HashMap, +} + +#[pymethods] +impl RoutingConstraints { + #[new] + #[pyo3(signature = (required_taints=None, preferred_taints=None))] + fn new( + required_taints: Option>, + preferred_taints: Option>, + ) -> Self { + Self { + required_taints: required_taints.unwrap_or_default(), + preferred_taints: preferred_taints.unwrap_or_default(), + } + } +} + +impl From for RsRoutingConstraints { + fn from(value: RoutingConstraints) -> Self { + Self { + required_taints: value.required_taints, + preferred_taints: value.preferred_taints, + } + } +} + +impl From for RoutingConstraints { + fn from(value: RsRoutingConstraints) -> Self { + Self { + required_taints: value.required_taints, + preferred_taints: value.preferred_taints, + } + } +} + #[pyclass] #[derive(Clone, Debug, Default)] pub struct ModelRuntimeConfig { @@ -73,6 +118,11 @@ impl ModelRuntimeConfig { self.inner.enable_eagle = enable_eagle; } + #[setter] + fn set_taints(&mut self, taints: HashSet) { + self.inner.taints = taints; + } + fn set_engine_specific(&mut self, key: &str, value: String) -> PyResult<()> { let value: serde_json::Value = serde_json::from_str(&value).map_err(to_pyerr)?; self.inner @@ -182,4 +232,9 @@ impl ModelRuntimeConfig { fn enable_eagle(&self) -> bool { self.inner.enable_eagle } + + #[getter] + fn taints(&self) -> HashSet { + self.inner.taints.clone() + } } diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 37aa2446b37d..59193b76cf7a 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -12,6 +12,7 @@ from typing import ( List, Literal, Optional, + Set, Tuple, ) @@ -514,6 +515,7 @@ class ModelRuntimeConfig: data_parallel_size: int enable_local_indexer: bool enable_eagle: bool + taints: Set[str] runtime_data: dict[str, Any] tensor_model_config: Any | None bootstrap_host: str | None @@ -545,6 +547,26 @@ class ModelRuntimeConfig: """Get the tensor model configuration.""" ... +class RoutingConstraints: + """ + Request-side routing constraints. + + ``required_taints`` is a hard eligibility filter. + ``preferred_taints`` maps taint -> signed weight. + Positive weights prefer matching workers, negative weights avoid them, + and ``0.0`` is neutral. Matching weights are summed and squashed with + ``tanh``, so opposite preferences cancel before Dynamo converts the + bounded bias into a strictly positive score multiplier. + """ + required_taints: Set[str] + preferred_taints: Dict[str, float] + + def __init__( + self, + required_taints: Optional[Set[str]] = None, + preferred_taints: Optional[Dict[str, float]] = None, + ) -> None: ... + class OverlapScores: """ A collection of prefix matching scores of workers for a given token ids. @@ -1998,6 +2020,7 @@ class KvRouter: block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, multi_modal_data: Optional[JsonLike] = None, mm_routing_info: Optional[JsonLike] = None, + routing_constraints: Optional[RoutingConstraints] = None, ) -> AsyncIterator[JsonLike]: """ Generate text using the KV-aware router. @@ -2026,6 +2049,7 @@ class KvRouter: mm_routing_info: Optional structured routing-only multimodal payload (e.g., {"routing_token_ids": [...], "block_mm_infos": [...]}) used by router selection without changing execution token_ids. + routing_constraints: Optional request routing constraints used to constrain or prefer tainted workers. Returns: An async iterator yielding generation responses @@ -2059,6 +2083,7 @@ class KvRouter: update_indexer: bool = False, block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, + routing_constraints: Optional[RoutingConstraints] = None, ) -> Tuple[int, int, int]: """ Find the best matching worker for the given tokens. diff --git a/lib/bindings/python/src/dynamo/llm/__init__.py b/lib/bindings/python/src/dynamo/llm/__init__.py index 8913f334b54a..0338e1c68288 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -34,6 +34,7 @@ from dynamo._core import ReasoningConfig as ReasoningConfig from dynamo._core import RouterConfig as RouterConfig from dynamo._core import RouterMode as RouterMode +from dynamo._core import RoutingConstraints as RoutingConstraints from dynamo._core import SglangArgs as SglangArgs from dynamo._core import WorkerMetricsPublisher as WorkerMetricsPublisher from dynamo._core import compute_block_hash_for_seq as compute_block_hash_for_seq diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index fe9fffc0c797..400ffc9a4fcf 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -1,8 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::ops::Range; +use std::sync::LazyLock; use std::time::Duration; use dynamo_tokens::{SequenceHash, Token, compute_hash_v2, compute_next_sequence_hash}; @@ -164,8 +166,72 @@ pub trait WorkerConfigLike { fn data_parallel_size(&self) -> u32; fn max_num_batched_tokens(&self) -> Option; fn total_kv_blocks(&self) -> Option; + fn taints(&self) -> &HashSet { + &EMPTY_WORKER_TAINTS + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +pub struct RoutingConstraints { + #[serde(default, skip_serializing_if = "HashSet::is_empty")] + pub required_taints: HashSet, + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub preferred_taints: HashMap, +} + +impl RoutingConstraints { + pub fn is_empty(&self) -> bool { + self.required_taints.is_empty() && self.preferred_taints.is_empty() + } + + pub fn has_hard_constraints(&self) -> bool { + !self.required_taints.is_empty() + } + + pub fn is_compatible_with_worker_taints(&self, worker_taints: &HashSet) -> bool { + if self.required_taints.is_empty() { + return true; + } + + self.required_taints + .iter() + .all(|taint| worker_taints.contains(taint)) + } + + pub fn preferred_taint_matches(&self, worker_taints: &HashSet) -> usize { + if self.preferred_taints.is_empty() { + return 0; + } + + self.preferred_taints + .keys() + .filter(|taint| worker_taints.contains(*taint)) + .count() + } + + pub fn preferred_taint_multiplier(&self, worker_taints: &HashSet) -> Option { + if self.preferred_taints.is_empty() { + return None; + } + + // Use exp(-tanh(sum)) so equal-magnitude positive and negative preferences + // have reciprocal effect around the neutral multiplier 1.0, while keeping the + // multiplier strictly positive and bounded to [exp(-1), exp(1)] ~= [0.368, 2.718] + // for numerically stable composition with the existing linear work score. + let bias = self + .preferred_taints + .iter() + .filter(|(taint, _)| worker_taints.contains(*taint)) + .map(|(_, weight)| f64::from(*weight)) + .sum::() + .tanh(); + + Some((-bias).exp()) + } } +static EMPTY_WORKER_TAINTS: LazyLock> = LazyLock::new(HashSet::new); + /// Transport abstraction for publishing batched router-visible KV cache events. pub trait RouterEventSink: Send + Sync { fn publish_event(&self, event: &RouterEvent) @@ -307,6 +373,8 @@ pub enum RouterRequest { tokens: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] block_mm_infos: Option>>, + #[serde(default, skip_serializing_if = "RoutingConstraints::is_empty")] + routing_constraints: RoutingConstraints, }, MarkPrefill, MarkFree { @@ -322,6 +390,7 @@ impl Default for RouterRequest { RouterRequest::New { tokens: vec![], block_mm_infos: None, + routing_constraints: RoutingConstraints::default(), } } } diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 73a438776d3a..e72e4d30e90b 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -17,6 +17,7 @@ use super::selector::{DefaultWorkerSelector, WorkerSelector}; use super::types::{ KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, TierOverlapBlocks, }; +use crate::protocols::RoutingConstraints; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::{ ActiveSequencesMultiWorker, PrefillTokenDeltas, SequenceError, SequencePublisher, @@ -185,6 +186,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, shared_cache_hits: Option, ) -> Result { let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); @@ -201,6 +203,7 @@ where decode_blocks: FxHashMap::default(), prefill_tokens: FxHashMap::default(), track_prefill_tokens, + routing_constraints, router_config_override: router_config_override.cloned(), update_states, lora_name, @@ -457,6 +460,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -501,6 +505,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -545,6 +550,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -590,6 +596,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -613,6 +620,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -657,6 +665,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -680,6 +689,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -738,6 +748,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -761,6 +772,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -818,6 +830,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -841,6 +854,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -896,6 +910,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -996,6 +1011,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -1094,6 +1110,7 @@ mod tests { None, None, None, + crate::protocols::RoutingConstraints::default(), None, ) .await diff --git a/lib/kv-router/src/scheduling/policy.rs b/lib/kv-router/src/scheduling/policy.rs index f9a914a410ed..942da281f1b6 100644 --- a/lib/kv-router/src/scheduling/policy.rs +++ b/lib/kv-router/src/scheduling/policy.rs @@ -4,7 +4,8 @@ use std::time::Duration; use super::config::RouterQueuePolicy; -use super::types::SchedulingRequest; +use super::types::SchedulingContext; +use crate::protocols::WorkerConfigLike; use ordered_float::OrderedFloat; /// Pluggable scheduling policy that determines queue ordering. /// Monomorphized for zero-cost inlining on the hot comparison path. @@ -15,10 +16,19 @@ pub trait SchedulingPolicy: Send + Sync + 'static { type Key: Ord + Eq + Clone + Send + 'static; /// Compute priority key at enqueue time. - fn enqueue_key(&self, arrival_offset: Duration, request: &SchedulingRequest) -> Self::Key; + fn enqueue_key( + &self, + arrival_offset: Duration, + ctx: SchedulingContext<'_, C>, + ) -> Self::Key; /// Recompute priority key during update(). Default: return old key unchanged. - fn rekey(&self, _now: Duration, old_key: &Self::Key, _req: &SchedulingRequest) -> Self::Key { + fn rekey( + &self, + _now: Duration, + old_key: &Self::Key, + _ctx: SchedulingContext<'_, C>, + ) -> Self::Key { old_key.clone() } @@ -37,8 +47,12 @@ pub struct FcfsPolicy; impl SchedulingPolicy for FcfsPolicy { type Key = OrderedFloat; - fn enqueue_key(&self, arrival_offset: Duration, request: &SchedulingRequest) -> Self::Key { - OrderedFloat(request.priority_jump.max(0.0) - arrival_offset.as_secs_f64()) + fn enqueue_key( + &self, + arrival_offset: Duration, + ctx: SchedulingContext<'_, C>, + ) -> Self::Key { + OrderedFloat(ctx.request().priority_jump.max(0.0) - arrival_offset.as_secs_f64()) } } @@ -52,8 +66,12 @@ pub struct LcfsPolicy; impl SchedulingPolicy for LcfsPolicy { type Key = OrderedFloat; - fn enqueue_key(&self, arrival_offset: Duration, request: &SchedulingRequest) -> Self::Key { - OrderedFloat(request.priority_jump.max(0.0) + arrival_offset.as_secs_f64()) + fn enqueue_key( + &self, + arrival_offset: Duration, + ctx: SchedulingContext<'_, C>, + ) -> Self::Key { + OrderedFloat(ctx.request().priority_jump.max(0.0) + arrival_offset.as_secs_f64()) } } @@ -71,9 +89,13 @@ pub struct WsptPolicy; impl SchedulingPolicy for WsptPolicy { type Key = OrderedFloat; - fn enqueue_key(&self, _arrival_offset: Duration, request: &SchedulingRequest) -> Self::Key { - let weight = 1.0 + request.priority_jump.max(0.0); - let new_tokens = request.best_effective_prefill_tokens().max(1); + fn enqueue_key( + &self, + _arrival_offset: Duration, + ctx: SchedulingContext<'_, C>, + ) -> Self::Key { + let weight = 1.0 + ctx.request().priority_jump.max(0.0); + let new_tokens = ctx.best_effective_prefill_tokens().max(1); OrderedFloat(weight / new_tokens as f64) } } @@ -100,21 +122,58 @@ impl RouterSchedulingPolicy { impl SchedulingPolicy for RouterSchedulingPolicy { type Key = OrderedFloat; - fn enqueue_key(&self, arrival_offset: Duration, request: &SchedulingRequest) -> Self::Key { + fn enqueue_key( + &self, + arrival_offset: Duration, + ctx: SchedulingContext<'_, C>, + ) -> Self::Key { match self { - Self::Fcfs(p) => p.enqueue_key(arrival_offset, request), - Self::Lcfs(p) => p.enqueue_key(arrival_offset, request), - Self::Wspt(p) => p.enqueue_key(arrival_offset, request), + Self::Fcfs(p) => p.enqueue_key(arrival_offset, ctx), + Self::Lcfs(p) => p.enqueue_key(arrival_offset, ctx), + Self::Wspt(p) => p.enqueue_key(arrival_offset, ctx), } } } #[cfg(test)] mod tests { + use std::collections::HashMap; + use rustc_hash::FxHashMap; use super::*; + use crate::SchedulingRequest; use crate::protocols::{OverlapScores, WorkerWithDpRank}; + use crate::test_utils::SimpleWorkerConfig; + + fn workers_for_request(request: &SchedulingRequest) -> HashMap { + let mut workers = HashMap::new(); + + for worker in request.effective_cached_tokens.keys() { + workers.entry(worker.worker_id).or_default(); + } + + if let Some(worker) = request.pinned_worker { + workers.entry(worker.worker_id).or_default(); + } + + if let Some(allowed_worker_ids) = request.allowed_worker_ids.as_ref() { + for &worker_id in allowed_worker_ids { + workers.entry(worker_id).or_default(); + } + } + + workers + } + + fn enqueue_key( + policy: &P, + arrival_offset: Duration, + request: &SchedulingRequest, + ) -> P::Key { + let workers = workers_for_request(request); + policy.enqueue_key(arrival_offset, SchedulingContext::new(request, &workers)) + } fn request_with( isl_tokens: usize, @@ -148,6 +207,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: None, } @@ -170,8 +230,8 @@ mod tests { fn fcfs_earlier_arrival_scheduled_first() { let policy = FcfsPolicy; let req = request_with(512, 0.0, OverlapScores::default()); - let early = policy.enqueue_key(Duration::from_secs(1), &req); - let late = policy.enqueue_key(Duration::from_secs(10), &req); + let early = enqueue_key(&policy, Duration::from_secs(1), &req); + let late = enqueue_key(&policy, Duration::from_secs(10), &req); assert!(early > late, "earlier arrival should have higher key"); } @@ -182,8 +242,8 @@ mod tests { let normal = request_with(512, 0.0, OverlapScores::default()); let boosted = request_with(512, 100.0, OverlapScores::default()); let t = Duration::from_secs(10); - let key_normal = policy.enqueue_key(t, &normal); - let key_boosted = policy.enqueue_key(t, &boosted); + let key_normal = enqueue_key(&policy, t, &normal); + let key_boosted = enqueue_key(&policy, t, &boosted); assert!( key_boosted > key_normal, "priority_jump should produce a higher key" @@ -198,8 +258,8 @@ mod tests { // B should be scheduled first despite arriving later. let a = request_with(512, 0.0, OverlapScores::default()); let b = request_with(512, 50.0, OverlapScores::default()); - let key_a = policy.enqueue_key(Duration::from_secs(0), &a); - let key_b = policy.enqueue_key(Duration::from_secs(5), &b); + let key_a = enqueue_key(&policy, Duration::from_secs(0), &a); + let key_b = enqueue_key(&policy, Duration::from_secs(5), &b); assert!(key_b > key_a); } @@ -207,8 +267,8 @@ mod tests { fn lcfs_later_arrival_scheduled_first() { let policy = LcfsPolicy; let req = request_with(512, 0.0, OverlapScores::default()); - let early = policy.enqueue_key(Duration::from_secs(1), &req); - let late = policy.enqueue_key(Duration::from_secs(10), &req); + let early = enqueue_key(&policy, Duration::from_secs(1), &req); + let late = enqueue_key(&policy, Duration::from_secs(10), &req); assert!(late > early, "later arrival should have higher key"); } @@ -218,8 +278,8 @@ mod tests { let normal = request_with(512, 0.0, OverlapScores::default()); let boosted = request_with(512, 100.0, OverlapScores::default()); let t = Duration::from_secs(10); - let key_normal = policy.enqueue_key(t, &normal); - let key_boosted = policy.enqueue_key(t, &boosted); + let key_normal = enqueue_key(&policy, t, &normal); + let key_boosted = enqueue_key(&policy, t, &boosted); assert!( key_boosted > key_normal, "priority_jump should produce a higher key" @@ -233,10 +293,10 @@ mod tests { let late = Duration::from_secs(10); let fcfs = RouterSchedulingPolicy::new(RouterQueuePolicy::Fcfs); - assert!(fcfs.enqueue_key(early, &req) > fcfs.enqueue_key(late, &req)); + assert!(enqueue_key(&fcfs, early, &req) > enqueue_key(&fcfs, late, &req)); let lcfs = RouterSchedulingPolicy::new(RouterQueuePolicy::Lcfs); - assert!(lcfs.enqueue_key(late, &req) > lcfs.enqueue_key(early, &req)); + assert!(enqueue_key(&lcfs, late, &req) > enqueue_key(&lcfs, early, &req)); } // ---- WSPT policy tests ---- @@ -248,7 +308,7 @@ mod tests { let long = request_with(1000, 0.0, OverlapScores::default()); let t = Duration::ZERO; assert!( - policy.enqueue_key(t, &short) > policy.enqueue_key(t, &long), + enqueue_key(&policy, t, &short) > enqueue_key(&policy, t, &long), "shorter request should have higher key" ); } @@ -260,8 +320,8 @@ mod tests { let no_cache = request_with(1024, 0.0, OverlapScores::default()); let cached = request_with(1024, 0.0, overlaps_from(&[(0, 60)])); let t = Duration::ZERO; - let key_no_cache = policy.enqueue_key(t, &no_cache); - let key_cached = policy.enqueue_key(t, &cached); + let key_no_cache = enqueue_key(&policy, t, &no_cache); + let key_cached = enqueue_key(&policy, t, &cached); assert!( key_cached > key_no_cache, "request with overlap should have higher key (fewer new tokens)" @@ -274,7 +334,7 @@ mod tests { let mut req = request_with(1024, 0.0, overlaps_from(&[(0, 60)])); req.track_prefill_tokens = false; - let key = policy.enqueue_key(Duration::ZERO, &req); + let key = enqueue_key(&policy, Duration::ZERO, &req); let expected = OrderedFloat(1.0 / 64.0); assert_eq!(key, expected); } @@ -286,7 +346,7 @@ mod tests { let boosted = request_with(512, 5.0, OverlapScores::default()); let t = Duration::ZERO; assert!( - policy.enqueue_key(t, &boosted) > policy.enqueue_key(t, &normal), + enqueue_key(&policy, t, &boosted) > enqueue_key(&policy, t, &normal), "priority_jump should increase key" ); } @@ -301,7 +361,7 @@ mod tests { 0.0, overlaps_from(&[(0, 10), (1, 20), (2, 50), (3, 60)]), ); - let key = policy.enqueue_key(Duration::ZERO, &req); + let key = enqueue_key(&policy, Duration::ZERO, &req); let expected = OrderedFloat(1.0 / 64.0); assert_eq!(key, expected); } @@ -312,7 +372,7 @@ mod tests { let mut req = request_with(1024, 0.0, overlaps_from(&[(0, 60), (1, 1)])); req.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); - let key = policy.enqueue_key(Duration::ZERO, &req); + let key = enqueue_key(&policy, Duration::ZERO, &req); let expected = OrderedFloat(1.0 / 1008.0); assert_eq!(key, expected); } @@ -323,7 +383,7 @@ mod tests { let mut req = request_with(1024, 0.0, overlaps_from(&[(0, 60)])); req.pinned_worker = Some(WorkerWithDpRank::new(1, 0)); - let key = policy.enqueue_key(Duration::ZERO, &req); + let key = enqueue_key(&policy, Duration::ZERO, &req); let expected = OrderedFloat(1.0 / 1024.0); assert_eq!(key, expected); } @@ -332,7 +392,7 @@ mod tests { fn wspt_no_overlap_falls_back_to_isl() { let policy = WsptPolicy; let req = request_with(512, 0.0, OverlapScores::default()); - let key = policy.enqueue_key(Duration::ZERO, &req); + let key = enqueue_key(&policy, Duration::ZERO, &req); let expected = OrderedFloat(1.0 / 512.0); assert_eq!(key, expected); } @@ -342,8 +402,26 @@ mod tests { let policy = WsptPolicy; // 512 tokens, 64 blocks cached = 1024 cached tokens > ISL → saturating_sub → 0 → max(1) let req = request_with(512, 0.0, overlaps_from(&[(0, 64)])); - let key = policy.enqueue_key(Duration::ZERO, &req); + let key = enqueue_key(&policy, Duration::ZERO, &req); let expected = OrderedFloat(1.0 / 1.0); assert_eq!(key, expected); } + + #[test] + fn wspt_required_taints_ignore_incompatible_overlap() { + let policy = WsptPolicy; + let mut req = request_with(1024, 0.0, overlaps_from(&[(0, 60), (1, 1)])); + req.routing_constraints.required_taints = + std::collections::HashSet::from(["mdc-b".to_string()]); + + let mut workers = workers_for_request(&req); + workers.get_mut(&0).unwrap().taints = + std::collections::HashSet::from(["mdc-a".to_string()]); + workers.get_mut(&1).unwrap().taints = + std::collections::HashSet::from(["mdc-b".to_string()]); + + let key = policy.enqueue_key(Duration::ZERO, SchedulingContext::new(&req, &workers)); + let expected = OrderedFloat(1.0 / 1008.0); + assert_eq!(key, expected); + } } diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 91c803221830..7ef6e5134c72 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use std::cmp::Ordering; -use std::collections::{BinaryHeap, HashMap, HashSet}; +use std::collections::{BinaryHeap, HashMap}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; @@ -13,7 +13,10 @@ use tokio::time::Instant; use super::policy::{FcfsPolicy, SchedulingPolicy}; use super::prefill_load::PrefillLoadEstimator; use super::selector::{DefaultWorkerSelector, WorkerSelector}; -use super::types::{SchedulingRequest, SchedulingResponse, pinned_worker_config}; +use super::types::{ + RoutingEligibility, SchedulingContext, SchedulingRequest, SchedulingResponse, + pinned_worker_config, +}; use crate::protocols::{PrefillLoadHint, WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::{ActiveSequencesMultiWorker, SequencePublisher, SequenceRequest}; @@ -143,7 +146,9 @@ impl< /// When `allowed_worker_ids` is set on the request without an exact pin /// (external routing), the capacity check is skipped. pub async fn enqueue(&self, mut request: SchedulingRequest) { - if let Err(error) = request.validate_worker_constraints() { + let eligibility = request.eligibility(); + + if let Err(error) = eligibility.validate_pinned_worker() { request.respond(Err(error)); return; } @@ -156,20 +161,19 @@ impl< return; }; - if request.bypass_capacity_check() { + if eligibility.bypasses_capacity_check() { self.admit_one(request, decay_now).await; return; } - if self.all_workers_busy( - threshold, - request.allowed_worker_ids.as_ref(), - request.pinned_worker, - decay_now, - ) { + if self.all_workers_busy(threshold, request.eligibility(), decay_now) { tracing::debug!("all workers busy, queueing request"); let arrival_offset = self.start_time.elapsed(); - let key = self.policy.enqueue_key(arrival_offset, &request); + let key = { + let workers = self.workers_with_configs.borrow(); + self.policy + .enqueue_key(arrival_offset, SchedulingContext::new(&request, &workers)) + }; let isl_tokens = request.isl_tokens; self.pending.lock().await.push(QueueEntry { key, request }); self.pending_count.fetch_add(1, AtomicOrdering::Relaxed); @@ -191,11 +195,16 @@ impl< if S::DYNAMIC { let now = self.start_time.elapsed(); let mut heap = self.pending.lock().await; + let workers = self.workers_with_configs.borrow(); let rekeyed: Vec<_> = std::mem::take(&mut *heap) .into_vec() .into_iter() .map(|e| QueueEntry { - key: self.policy.rekey(now, &e.key, &e.request), + key: self.policy.rekey( + now, + &e.key, + SchedulingContext::new(&e.request, &workers), + ), request: e.request, }) .collect(); @@ -213,12 +222,7 @@ impl< // drain overhead bounded to the heap front. A blocked pinned or // otherwise constrained request can temporarily stall later // schedulable entries until we adopt a cheaper non-HOL strategy. - if self.all_workers_busy( - threshold, - front.request.allowed_worker_ids.as_ref(), - front.request.pinned_worker, - decay_now, - ) { + if self.all_workers_busy(threshold, front.request.eligibility(), decay_now) { break; } let entry = heap.pop().expect("heap front vanished before pop"); @@ -350,17 +354,19 @@ impl< fn all_workers_busy( &self, threshold: f64, - allowed: Option<&HashSet>, - pinned_worker: Option, + eligibility: RoutingEligibility<'_>, decay_now: Instant, ) -> bool { let active_tokens = self.slots.active_tokens(decay_now); let configs = self.workers_with_configs.borrow(); - if let Some(worker) = pinned_worker { + if let Some(worker) = eligibility.pinned_worker() { let Ok(config) = pinned_worker_config::(&*configs, worker) else { return false; }; + if !eligibility.allows_worker(worker.worker_id, config) { + return false; + } let max_batched = config .max_num_batched_tokens() @@ -371,9 +377,7 @@ impl< let mut checked_any = false; for (&worker_id, config) in configs.iter() { - if let Some(ids) = allowed - && !ids.contains(&worker_id) - { + if !eligibility.allows_worker(worker_id, config) { continue; } let dp_size = config.data_parallel_size(); @@ -397,7 +401,7 @@ impl< #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Condvar, Mutex as StdMutex}; use std::time::Duration; @@ -635,6 +639,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -1029,6 +1034,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: Some(allowed), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -1064,6 +1070,46 @@ mod tests { )); } + #[tokio::test(flavor = "multi_thread")] + async fn test_disallowed_worker_ids_fail_without_queueing() { + let (queue, _slots) = make_queue(1, 16, 256, Some(0.0)); + let (mut req, rx) = make_request("disallowed", 256); + req.allowed_worker_ids = Some(HashSet::from([999])); + + queue.enqueue(req).await; + + let resp = rx.await.expect("oneshot dropped"); + assert!(matches!(resp, Err(KvSchedulerError::NoEndpoints))); + assert_eq!(queue.pending_count(), 0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_incompatible_required_taints_fail_without_queueing() { + let (queue, _slots, cfg_tx) = make_queue_with_sender(1, 16, 256, Some(0.0), None); + let mut configs = HashMap::new(); + configs.insert( + 0_u64, + SimpleWorkerConfig { + max_num_batched_tokens: Some(256), + taints: HashSet::from(["mdc-a".to_string()]), + ..Default::default() + }, + ); + cfg_tx.send(configs).unwrap(); + + let (mut req, rx) = make_request("tainted", 256); + req.routing_constraints = crate::protocols::RoutingConstraints { + required_taints: HashSet::from(["mdc-b".to_string()]), + preferred_taints: HashMap::new(), + }; + + queue.enqueue(req).await; + + let resp = rx.await.expect("oneshot dropped"); + assert!(matches!(resp, Err(KvSchedulerError::NoEndpoints))); + assert_eq!(queue.pending_count(), 0); + } + #[tokio::test(flavor = "multi_thread")] async fn test_pinned_request_head_of_line_blocks_other_worker_capacity() { let (queue, slots) = make_queue(2, 16, 256, Some(0.0)); diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 78ad979363b7..e35c89b38835 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -203,14 +203,17 @@ impl WorkerSelector for DefaultWorkerSelector { block_size: u32, ) -> Result { assert!(request.isl_tokens > 0); - request.validate_worker_constraints()?; + let eligibility = request.eligibility(); + eligibility.validate_pinned_worker()?; - let pinned_worker = request.pinned_worker; + let pinned_worker = eligibility.pinned_worker(); if pinned_worker.is_none() - && !workers - .keys() - .any(|worker_id| request.is_worker_allowed(*worker_id)) + && !eligibility.has_eligible_worker( + workers + .iter() + .map(|(&worker_id, config)| (worker_id, config)), + ) { return Err(KvSchedulerError::NoEndpoints); } @@ -237,6 +240,12 @@ impl WorkerSelector for DefaultWorkerSelector { if let Some(worker) = pinned_worker { pinned_worker_config(workers, worker)?; + if workers + .get(&worker.worker_id) + .is_some_and(|config| !eligibility.allows_worker(worker.worker_id, config)) + { + return Err(KvSchedulerError::NoEndpoints); + } let logit = self.worker_logit(request, worker, block_size, weights, "Pinned formula"); let effective_overlap_blocks = request.effective_overlap_blocks_for(worker); @@ -264,14 +273,23 @@ impl WorkerSelector for DefaultWorkerSelector { .as_ref() .and_then(|cfg| cfg.router_temperature) .unwrap_or(self.kv_router_config.router_temperature); - let get_score = |worker: WorkerWithDpRank| -> f64 { - self.worker_logit(request, worker, block_size, weights, "Formula") + let base_score = self.worker_logit(request, worker, block_size, weights, "Formula"); + let Some(config) = workers.get(&worker.worker_id) else { + return base_score; + }; + match request + .routing_constraints + .preferred_taint_multiplier(config.taints()) + { + Some(multiplier) => base_score * multiplier, + None => base_score, + } }; let worker_iter = workers .iter() - .filter(move |(worker_id, _)| request.is_worker_allowed(**worker_id)) + .filter(move |(worker_id, config)| eligibility.allows_worker(**worker_id, *config)) .flat_map(|(worker_id, config)| { let data_parallel_size = config.data_parallel_size(); let data_parallel_start_rank = config.data_parallel_start_rank(); @@ -379,8 +397,37 @@ impl WorkerSelector for DefaultWorkerSelector { #[cfg(test)] mod tests { + use std::collections::HashSet; + use super::*; - use crate::protocols::SharedCacheHits; + use crate::protocols::{SharedCacheHits, WorkerConfigLike}; + + #[derive(Clone, Default)] + struct TaintedWorkerConfig { + taints: HashSet, + } + + impl WorkerConfigLike for TaintedWorkerConfig { + 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 taints(&self) -> &HashSet { + &self.taints + } + } #[test] fn test_softmax_sample_single_key() { @@ -522,6 +569,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: None, }; @@ -544,6 +592,268 @@ mod tests { ); } + #[test] + fn test_required_taints_return_no_endpoints_when_no_worker_matches() { + let selector = DefaultWorkerSelector::new(Some(KvRouterConfig::default()), "test"); + let workers = HashMap::from([( + 10, + TaintedWorkerConfig { + taints: HashSet::from(["mdc-a".to_string()]), + }, + )]); + let request = SchedulingRequest { + maybe_request_id: Some("test".into()), + token_seq: None, + isl_tokens: 16, + tier_overlap_blocks: Default::default(), + effective_overlap_blocks: HashMap::default(), + effective_cached_tokens: HashMap::default(), + decode_blocks: FxHashMap::default(), + prefill_tokens: FxHashMap::default(), + track_prefill_tokens: true, + router_config_override: None, + update_states: false, + lora_name: None, + priority_jump: 0.0, + expected_output_tokens: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints { + required_taints: HashSet::from(["mdc-b".to_string()]), + preferred_taints: HashMap::new(), + }, + shared_cache_hits: None, + resp_tx: None, + }; + + let result = selector.select_worker(&workers, &request, 16); + assert!(matches!(result, Err(KvSchedulerError::NoEndpoints))); + } + + #[test] + fn test_required_taints_filter_out_incompatible_workers() { + let selector = DefaultWorkerSelector::new(Some(KvRouterConfig::default()), "test"); + let workers = HashMap::from([ + ( + 10, + TaintedWorkerConfig { + taints: HashSet::from(["mdc-a".to_string()]), + }, + ), + ( + 20, + TaintedWorkerConfig { + taints: HashSet::from(["mdc-b".to_string()]), + }, + ), + ]); + let request = SchedulingRequest { + maybe_request_id: Some("test".into()), + token_seq: None, + isl_tokens: 16, + tier_overlap_blocks: Default::default(), + effective_overlap_blocks: HashMap::default(), + effective_cached_tokens: HashMap::default(), + decode_blocks: FxHashMap::default(), + prefill_tokens: FxHashMap::default(), + track_prefill_tokens: true, + router_config_override: None, + update_states: false, + lora_name: None, + priority_jump: 0.0, + expected_output_tokens: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints { + required_taints: HashSet::from(["mdc-b".to_string()]), + preferred_taints: HashMap::new(), + }, + shared_cache_hits: None, + resp_tx: None, + }; + + let result = selector.select_worker(&workers, &request, 16).unwrap(); + assert_eq!(result.worker.worker_id, 20); + } + + #[test] + fn test_required_taints_switch_matching_worker_sets_by_label() { + let selector = DefaultWorkerSelector::new(Some(KvRouterConfig::default()), "test"); + let name_a = "mdc-a".to_string(); + let name_b = "mdc-b".to_string(); + let name_c = "mdc-c".to_string(); + let taint_a = TaintedWorkerConfig { + taints: HashSet::from([name_a.clone()]), + }; + let taint_b = TaintedWorkerConfig { + taints: HashSet::from([name_b.clone()]), + }; + let taint_c = TaintedWorkerConfig { + taints: HashSet::from([name_c.clone()]), + }; + let workers = HashMap::from([ + (10, taint_a.clone()), + (11, taint_a), + (20, taint_b.clone()), + (21, taint_b), + (30, taint_c.clone()), + (31, taint_c), + ]); + + for (required_taint, expected_worker_id, noisy_worker_id) in [ + (name_a, 10_u64, 11_u64), + (name_b, 20_u64, 21_u64), + (name_c, 30_u64, 31_u64), + ] { + let mut decode_blocks = FxHashMap::default(); + decode_blocks.insert(WorkerWithDpRank::from_worker_id(expected_worker_id), 0); + decode_blocks.insert(WorkerWithDpRank::from_worker_id(noisy_worker_id), 400_000); + + let request = SchedulingRequest { + maybe_request_id: Some("test".into()), + token_seq: None, + isl_tokens: 16, + tier_overlap_blocks: Default::default(), + effective_overlap_blocks: HashMap::default(), + effective_cached_tokens: HashMap::default(), + decode_blocks, + prefill_tokens: FxHashMap::default(), + track_prefill_tokens: true, + router_config_override: None, + update_states: false, + lora_name: None, + priority_jump: 0.0, + expected_output_tokens: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints { + required_taints: HashSet::from([required_taint.clone()]), + preferred_taints: HashMap::new(), + }, + shared_cache_hits: None, + resp_tx: None, + }; + + let result = selector.select_worker(&workers, &request, 16).unwrap(); + assert_eq!( + result.worker.worker_id, expected_worker_id, + "required taint {required_taint} should route only within its compatible worker set" + ); + } + } + + #[test] + fn test_preferred_taints_prefer_matching_worker() { + let selector = DefaultWorkerSelector::new( + Some(KvRouterConfig { + router_temperature: 0.0, + ..Default::default() + }), + "test", + ); + let workers = HashMap::from([ + ( + 10, + TaintedWorkerConfig { + taints: HashSet::from(["mdc-a".to_string()]), + }, + ), + ( + 20, + TaintedWorkerConfig { + taints: HashSet::from(["mdc-b".to_string()]), + }, + ), + ]); + let mut decode_blocks = FxHashMap::default(); + decode_blocks.insert(WorkerWithDpRank::from_worker_id(10), 100); + decode_blocks.insert(WorkerWithDpRank::from_worker_id(20), 90); + + let request = SchedulingRequest { + maybe_request_id: Some("test".into()), + token_seq: None, + isl_tokens: 16, + tier_overlap_blocks: Default::default(), + effective_overlap_blocks: HashMap::default(), + effective_cached_tokens: HashMap::default(), + decode_blocks, + prefill_tokens: FxHashMap::default(), + track_prefill_tokens: true, + router_config_override: None, + update_states: false, + lora_name: None, + priority_jump: 0.0, + expected_output_tokens: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints { + required_taints: HashSet::new(), + preferred_taints: HashMap::from([("mdc-a".to_string(), 0.85)]), + }, + shared_cache_hits: None, + resp_tx: None, + }; + + let result = selector.select_worker(&workers, &request, 16).unwrap(); + assert_eq!(result.worker.worker_id, 10); + } + + #[test] + fn test_negative_preferred_taints_avoid_matching_worker() { + let selector = DefaultWorkerSelector::new( + Some(KvRouterConfig { + router_temperature: 0.0, + ..Default::default() + }), + "test", + ); + let workers = HashMap::from([ + ( + 10, + TaintedWorkerConfig { + taints: HashSet::from(["mdc-a".to_string()]), + }, + ), + ( + 20, + TaintedWorkerConfig { + taints: HashSet::from(["mdc-b".to_string()]), + }, + ), + ]); + let mut decode_blocks = FxHashMap::default(); + decode_blocks.insert(WorkerWithDpRank::from_worker_id(10), 90); + decode_blocks.insert(WorkerWithDpRank::from_worker_id(20), 100); + + let request = SchedulingRequest { + maybe_request_id: Some("test".into()), + token_seq: None, + isl_tokens: 16, + tier_overlap_blocks: Default::default(), + effective_overlap_blocks: HashMap::default(), + effective_cached_tokens: HashMap::default(), + decode_blocks, + prefill_tokens: FxHashMap::default(), + track_prefill_tokens: true, + router_config_override: None, + update_states: false, + lora_name: None, + priority_jump: 0.0, + expected_output_tokens: None, + pinned_worker: None, + allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints { + required_taints: HashSet::new(), + preferred_taints: HashMap::from([("mdc-a".to_string(), -0.25)]), + }, + shared_cache_hits: None, + resp_tx: None, + }; + + let result = selector.select_worker(&workers, &request, 16).unwrap(); + assert_eq!(result.worker.worker_id, 20); + } + /// Test the scoring formula with shared cache hits. /// /// Request [A, B, C, D], shared_cache_multiplier=0.5, block_size=1 @@ -604,6 +914,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: Some(shared_hits), resp_tx: Some(tx), }; @@ -668,6 +979,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -727,6 +1039,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -776,6 +1089,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index bd45ca3bffbb..b1a08c87e953 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -8,7 +8,9 @@ use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use super::config::RouterConfigOverride; -use crate::protocols::{DpRank, SharedCacheHits, WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use crate::protocols::{ + DpRank, RoutingConstraints, SharedCacheHits, WorkerConfigLike, WorkerId, WorkerWithDpRank, +}; use crate::sequences::PrefillTokenDeltas; #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -62,6 +64,7 @@ pub struct SchedulingRequest { // Routing constraints and request-level config. pub pinned_worker: Option, pub allowed_worker_ids: Option>, + pub routing_constraints: RoutingConstraints, pub router_config_override: Option, pub track_prefill_tokens: bool, pub priority_jump: f64, @@ -81,7 +84,141 @@ pub struct SchedulingRequest { pub resp_tx: Option>>, } +#[derive(Clone, Copy)] +pub(crate) struct RoutingEligibility<'a> { + allowed_worker_ids: Option<&'a HashSet>, + pinned_worker: Option, + routing_constraints: &'a RoutingConstraints, +} + +impl<'a> RoutingEligibility<'a> { + #[inline] + pub(crate) fn new( + allowed_worker_ids: Option<&'a HashSet>, + pinned_worker: Option, + routing_constraints: &'a RoutingConstraints, + ) -> Self { + Self { + allowed_worker_ids, + pinned_worker, + routing_constraints, + } + } + + #[inline] + pub(crate) fn pinned_worker(&self) -> Option { + self.pinned_worker + } + + #[inline] + pub(crate) fn allows_worker_id(&self, worker_id: WorkerId) -> bool { + self.allowed_worker_ids + .is_none_or(|worker_ids| worker_ids.contains(&worker_id)) + } + + #[inline] + pub(crate) fn allows_worker( + &self, + worker_id: WorkerId, + config: &C, + ) -> bool { + self.allows_worker_id(worker_id) + && self + .routing_constraints + .is_compatible_with_worker_taints(config.taints()) + } + + #[inline] + pub(crate) fn has_eligible_worker<'w, C, I>(&self, workers: I) -> bool + where + C: WorkerConfigLike + 'w, + I: IntoIterator, + { + for (worker_id, config) in workers { + if !self.allows_worker_id(worker_id) { + continue; + } + + if self.allows_worker(worker_id, config) { + return true; + } + } + + false + } + + #[inline] + pub(crate) fn validate_pinned_worker(&self) -> Result<(), KvSchedulerError> { + let Some(pinned_worker) = self.pinned_worker else { + return Ok(()); + }; + + if self.allows_worker_id(pinned_worker.worker_id) { + return Ok(()); + } + + Err(KvSchedulerError::PinnedWorkerNotAllowed { + worker_id: pinned_worker.worker_id, + }) + } + + #[inline] + pub(crate) fn bypasses_capacity_check(&self) -> bool { + self.pinned_worker.is_none() && self.allowed_worker_ids.is_some() + } +} + +#[derive(Clone, Copy)] +pub struct SchedulingContext<'a, C> { + request: &'a SchedulingRequest, + eligibility: RoutingEligibility<'a>, + workers: &'a HashMap, +} + +impl<'a, C: WorkerConfigLike> SchedulingContext<'a, C> { + pub fn new(request: &'a SchedulingRequest, workers: &'a HashMap) -> Self { + Self { + request, + eligibility: request.eligibility(), + workers, + } + } + + pub fn request(&self) -> &'a SchedulingRequest { + self.request + } + + pub fn best_effective_prefill_tokens(&self) -> usize { + let cached_tokens = match self.eligibility.pinned_worker() { + Some(worker) => self.request.effective_cached_tokens_for(worker), + None => self + .request + .effective_cached_tokens + .iter() + .filter(|(worker, _)| { + self.workers.get(&worker.worker_id).is_some_and(|config| { + self.eligibility.allows_worker(worker.worker_id, config) + }) + }) + .map(|(_, cached_tokens)| *cached_tokens) + .max() + .unwrap_or(0), + }; + + self.request.isl_tokens.saturating_sub(cached_tokens) + } +} + impl SchedulingRequest { + #[inline] + pub(crate) fn eligibility(&self) -> RoutingEligibility<'_> { + RoutingEligibility::new( + self.allowed_worker_ids.as_ref(), + self.pinned_worker, + &self.routing_constraints, + ) + } + pub(crate) fn prefill_token_deltas(&self) -> PrefillTokenDeltas { if !self.track_prefill_tokens { return PrefillTokenDeltas::none(); @@ -109,21 +246,6 @@ impl SchedulingRequest { PrefillTokenDeltas::new(self.isl_tokens, by_worker) } - pub(crate) fn best_effective_prefill_tokens(&self) -> usize { - let cached_tokens = match self.pinned_worker { - Some(worker) => self.effective_cached_tokens_for(worker), - None => self - .effective_cached_tokens - .iter() - .filter(|(worker, _)| self.is_worker_allowed(worker.worker_id)) - .map(|(_, cached_tokens)| *cached_tokens) - .max() - .unwrap_or(0), - }; - - self.isl_tokens.saturating_sub(cached_tokens) - } - pub(crate) fn effective_cached_tokens_for(&self, worker: WorkerWithDpRank) -> usize { self.effective_cached_tokens .get(&worker) @@ -138,12 +260,6 @@ impl SchedulingRequest { .unwrap_or(0.0) } - pub(crate) fn is_worker_allowed(&self, worker_id: WorkerId) -> bool { - self.allowed_worker_ids - .as_ref() - .is_none_or(|ids| ids.contains(&worker_id)) - } - #[cfg(test)] pub(crate) fn prefill_tokens_for(&self, worker: WorkerWithDpRank) -> usize { let default_prefill_tokens = if self.track_prefill_tokens { @@ -175,26 +291,6 @@ impl SchedulingRequest { self.isl_tokens.div_ceil(block_size as usize) as u64 } - pub fn validate_worker_constraints(&self) -> Result<(), KvSchedulerError> { - let Some(pinned_worker) = self.pinned_worker else { - return Ok(()); - }; - let Some(allowed_worker_ids) = self.allowed_worker_ids.as_ref() else { - return Ok(()); - }; - if allowed_worker_ids.contains(&pinned_worker.worker_id) { - return Ok(()); - } - - Err(KvSchedulerError::PinnedWorkerNotAllowed { - worker_id: pinned_worker.worker_id, - }) - } - - pub fn bypass_capacity_check(&self) -> bool { - self.pinned_worker.is_none() && self.allowed_worker_ids.is_some() - } - pub fn respond(&mut self, result: Result) { let Some(tx) = self.resp_tx.take() else { tracing::error!("respond called multiple times on same request"); diff --git a/lib/kv-router/src/test_utils.rs b/lib/kv-router/src/test_utils.rs index 67d50184a267..dc100eb20964 100644 --- a/lib/kv-router/src/test_utils.rs +++ b/lib/kv-router/src/test_utils.rs @@ -3,6 +3,7 @@ //! Shared test utilities for radix tree tests. +use std::collections::HashSet; use std::future; use crate::indexer::KvIndexerInterface; @@ -370,6 +371,7 @@ pub struct SimpleWorkerConfig { pub data_parallel_size: u32, pub max_num_batched_tokens: Option, pub total_kv_blocks: Option, + pub taints: HashSet, } impl Default for SimpleWorkerConfig { @@ -379,6 +381,7 @@ impl Default for SimpleWorkerConfig { data_parallel_size: 1, max_num_batched_tokens: None, total_kv_blocks: None, + taints: HashSet::new(), } } } @@ -399,4 +402,8 @@ impl WorkerConfigLike for SimpleWorkerConfig { fn total_kv_blocks(&self) -> Option { self.total_kv_blocks } + + fn taints(&self) -> &HashSet { + &self.taints + } } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 625a1f29f64e..96c1fa4c5f45 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -15,8 +15,8 @@ use dynamo_kv_router::{ protocols::KV_EVENT_SUBJECT, protocols::{ BlockExtraInfo, BlockHashOptions, DpRank, LocalBlockHash, PrefillLoadHint, RouterEvent, - RouterRequest, RouterResponse, TokensWithHashes, WorkerConfigLike, WorkerId, - WorkerWithDpRank, compute_block_hash_for_seq, + RouterRequest, RouterResponse, RoutingConstraints, TokensWithHashes, WorkerConfigLike, + WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, }, scheduling::TierOverlapBlocks, }; @@ -529,6 +529,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, ) -> anyhow::Result { let start = Instant::now(); @@ -633,6 +634,7 @@ where expected_output_tokens, pinned_worker, allowed_worker_ids, + routing_constraints, shared_cache_hits, ) .instrument(tracing::info_span!("kv_router.schedule")) @@ -696,6 +698,7 @@ where priority_jump: f64, expected_output_tokens: Option, allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, ) -> anyhow::Result<(WorkerWithDpRank, u32)> { let result = self .find_best_match_details( @@ -709,6 +712,7 @@ where expected_output_tokens, None, allowed_worker_ids, + routing_constraints, ) .await?; Ok((result.worker, result.cache_hit.rounded_overlap_blocks())) @@ -1072,6 +1076,7 @@ where RouterRequest::New { tokens, block_mm_infos, + routing_constraints, } => { let (best_worker, overlap_blocks) = self .find_best_match( @@ -1084,6 +1089,7 @@ where 0.0, None, None, + routing_constraints, ) .await?; @@ -1311,6 +1317,7 @@ mod tests { 0.0, None, None, + RoutingConstraints::default(), ) .await .unwrap(); @@ -1344,6 +1351,7 @@ mod tests { 0.0, None, None, + RoutingConstraints::default(), ) .await .unwrap(); diff --git a/lib/llm/src/kv_router/prefill_router/execution.rs b/lib/llm/src/kv_router/prefill_router/execution.rs index c9621f6a144e..eee21ab83869 100644 --- a/lib/llm/src/kv_router/prefill_router/execution.rs +++ b/lib/llm/src/kv_router/prefill_router/execution.rs @@ -8,7 +8,7 @@ use futures::StreamExt; use tokio::sync::OwnedSemaphorePermit; use tracing::Instrument; -use dynamo_kv_router::protocols::{BlockExtraInfo, WorkerId}; +use dynamo_kv_router::protocols::{BlockExtraInfo, RoutingConstraints, WorkerId}; use dynamo_runtime::{pipeline::SingleIn, protocols::maybe_error::MaybeError}; use super::{InnerPrefillRouter, PrefillError, PrefillResolveDecision, PrefillRouter}; @@ -58,6 +58,11 @@ impl PrefillRouter { .routing .as_ref() .and_then(|r| r.allowed_worker_ids.clone()); + let routing_constraints = req + .routing + .as_ref() + .and_then(|r| r.routing_constraints.clone()) + .unwrap_or_default(); let (routing_token_ids, block_mm_infos) = req.block_mm_routing_info(); match self .query_prefill_worker( @@ -67,6 +72,7 @@ impl PrefillRouter { lora_name, priority_jump, allowed_worker_ids, + routing_constraints, ) .await { @@ -264,6 +270,7 @@ impl PrefillRouter { /// /// This is the shared worker selection logic used by both `resolve_prefill_worker` /// and `query_route`. + #[expect(clippy::too_many_arguments)] pub async fn query_prefill_worker( &self, token_ids: &[u32], @@ -272,6 +279,7 @@ impl PrefillRouter { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, ) -> Result<(u64, Option)> { let prefill_router = self .prefill_router @@ -292,6 +300,7 @@ impl PrefillRouter { priority_jump, None, allowed_worker_ids, + routing_constraints, ) .await?; Ok((worker.worker_id, Some(worker.dp_rank))) diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 69bfe6a87059..34f8ab2b7cea 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use anyhow::Result; -use dynamo_kv_router::protocols::{TokensWithHashes, WorkerWithDpRank}; +use dynamo_kv_router::protocols::{TokensWithHashes, WorkerConfigLike, WorkerWithDpRank}; use dynamo_runtime::{ dynamo_nvtx_range, metrics::frontend_perf::{STAGE_DISPATCH, STAGE_ROUTE, StageGuard}, @@ -302,6 +302,9 @@ impl KvPushRouter { let priority_jump = routing.and_then(|r| r.priority_jump).unwrap_or(0.0); let expected_output_tokens = routing.and_then(|r| r.expected_output_tokens); let allowed_worker_ids = routing.and_then(|r| r.allowed_worker_ids.clone()); + let routing_constraints = routing + .and_then(|r| r.routing_constraints.clone()) + .unwrap_or_default(); let (routing_token_ids, block_mm_infos) = request.block_mm_routing_info(); let Some((pinned_worker_id, requested_dp_rank)) = pinned_worker_hint(phase, routing) else { let _nvtx_kv = dynamo_nvtx_range!("route.kv_match"); @@ -318,6 +321,7 @@ impl KvPushRouter { expected_output_tokens, None, allowed_worker_ids, + routing_constraints.clone(), ) .await?; let best_worker = selection.worker; @@ -374,6 +378,7 @@ impl KvPushRouter { expected_output_tokens, Some(pinned_worker), allowed_worker_ids, + routing_constraints.clone(), ) .await?; let best_worker = selection.worker; @@ -404,6 +409,30 @@ impl KvPushRouter { "Routing to specified worker" ); + if routing_constraints.has_hard_constraints() { + let configs = self.chooser.workers_with_configs.borrow(); + match configs.get(&pinned_worker_id) { + Some(config) + if !routing_constraints.is_compatible_with_worker_taints(config.taints()) => + { + return Err(anyhow::anyhow!( + "Pinned worker {} does not satisfy required taints {:?}; worker taints: {:?}", + pinned_worker_id, + routing_constraints.required_taints, + config.taints() + )); + } + None => { + return Err(anyhow::anyhow!( + "Pinned worker {} could not be validated against required taints {:?} because worker config was unavailable", + pinned_worker_id, + routing_constraints.required_taints + )); + } + _ => {} + } + } + // Build a WorkerWithDpRank; use 0 as a fallback dp_rank when it // couldn't be resolved -- this is only used for the cache-hit // estimate query and won't affect scheduler state. diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 04c97aa2b59c..f775d9df06fd 100644 --- a/lib/llm/src/kv_router/scheduler.rs +++ b/lib/llm/src/kv_router/scheduler.rs @@ -20,7 +20,7 @@ use anyhow::Result; use dynamo_kv_router::{ PrefillLoadEstimator, config::{KvRouterConfig, RouterConfigOverride}, - protocols::{WorkerId, WorkerWithDpRank}, + protocols::{RoutingConstraints, WorkerId, WorkerWithDpRank}, }; use dynamo_runtime::component::Component; use dynamo_runtime::traits::DistributedRuntimeProvider; @@ -141,6 +141,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, shared_cache_hits: Option, ) -> Result { let response = self @@ -159,6 +160,7 @@ where expected_output_tokens, pinned_worker, allowed_worker_ids, + routing_constraints, shared_cache_hits, ) .await; diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 56cbd12a7b31..170e9fb8302d 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -1,12 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use serde::{Deserialize, Serialize, de::DeserializeOwned}; use crate::protocols::tensor; - #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] pub struct DisaggregatedEndpoint { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -64,6 +63,9 @@ pub struct ModelRuntimeConfig { #[serde(default = "default_eagle")] pub enable_eagle: bool, + + #[serde(default, skip_serializing_if = "HashSet::is_empty")] + pub taints: HashSet, } const fn default_data_parallel_start_rank() -> u32 { @@ -102,6 +104,7 @@ impl Default for ModelRuntimeConfig { tensor_model_config: None, disaggregated_endpoint: None, enable_eagle: false, + taints: HashSet::new(), } } } @@ -122,6 +125,10 @@ impl dynamo_kv_router::WorkerConfigLike for ModelRuntimeConfig { fn total_kv_blocks(&self) -> Option { self.total_kv_blocks } + + fn taints(&self) -> &HashSet { + &self.taints + } } impl ModelRuntimeConfig { diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index ea850db67c3b..bb6de03520a1 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -545,6 +545,7 @@ impl OpenAIPreprocessor { lora_name, allowed_worker_ids: None, session_control: nvext.session_control.clone(), + routing_constraints: nvext.routing_constraints.clone(), }; builder.routing(Some(routing)); } else if lora_name.is_some() { diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs index 749e6d792e60..89af400f1e58 100644 --- a/lib/llm/src/protocols/common/preprocessor.rs +++ b/lib/llm/src/protocols/common/preprocessor.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use derive_builder::Builder; use dynamo_kv_router::{ config::RouterConfigOverride, - protocols::{BlockExtraInfo, WorkerId}, + protocols::{BlockExtraInfo, RoutingConstraints, WorkerId}, }; use serde::{Deserialize, Serialize}; @@ -68,6 +68,10 @@ pub struct RoutingHints { #[serde(default, skip_serializing_if = "Option::is_none")] pub allowed_worker_ids: Option>, + /// Request routing constraints used for worker compatibility and soft preference. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub routing_constraints: Option, + /// Session control for subagent KV isolation and sticky routing. /// Contains session_id (for affinity) and optional action (open/close). #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 7c561dd4a59d..4aaffebeb2b5 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -3,6 +3,7 @@ use axum::http::HeaderMap; use derive_builder::Builder; +use dynamo_kv_router::protocols::RoutingConstraints; use dynamo_protocols::types::StopReason; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; @@ -284,6 +285,27 @@ impl NvExtResponseFieldSelection { } } +/// OpenAPI-facing schema for request routing constraints. +/// +/// Runtime serialization still uses `dynamo_kv_router::protocols::RoutingConstraints`; +/// this mirror exists so `NvExt` can expose the concrete field shape without +/// making the kv-router crate depend on utoipa. +#[derive(ToSchema, Serialize, Deserialize, Debug, Clone, PartialEq, Default)] +pub struct RoutingConstraintsSchema { + /// Worker taints that must be matched for the request to be eligible. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub required_taints: Vec, + + /// Soft preference weights keyed by worker taint. + /// Positive weights prefer matching workers; negative weights avoid them. + /// A weight of 0.0 is neutral and has no effect. + /// Matching weights are summed and squashed with `tanh`, so opposite + /// preferences cancel before Dynamo converts the bounded bias into a + /// strictly positive score multiplier. + #[serde(default, skip_serializing_if = "std::collections::HashMap::is_empty")] + pub preferred_taints: std::collections::HashMap, +} + /// NVIDIA LLM extensions to the OpenAI API #[derive(ToSchema, Serialize, Deserialize, Builder, Validate, Debug, Clone)] #[validate(schema(function = "validate_nv_ext"))] @@ -384,6 +406,12 @@ pub struct NvExt { #[builder(default, setter(strip_option))] #[serde(default, skip_serializing_if = "Option::is_none")] pub session_control: Option, + + /// Request routing constraints used to constrain or prefer tainted workers. + #[builder(default, setter(strip_option))] + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schema(value_type = RoutingConstraintsSchema)] + pub routing_constraints: Option, } /// Hints from the agent/caller about request characteristics. @@ -514,6 +542,7 @@ mod tests { assert_eq!(nv_ext.agent_context, None); assert_eq!(nv_ext.request_timestamp_ms, None); assert_eq!(nv_ext.session_control, None); + assert_eq!(nv_ext.routing_constraints, None); } // Test valid builder configurations diff --git a/lib/mocker/src/replay/offline/components/router.rs b/lib/mocker/src/replay/offline/components/router.rs index c4a950e16525..1c2ae9a87fa0 100644 --- a/lib/mocker/src/replay/offline/components/router.rs +++ b/lib/mocker/src/replay/offline/components/router.rs @@ -10,10 +10,11 @@ use anyhow::{Context, Result, anyhow}; use dynamo_kv_router::LocalBlockHash; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::protocols::{ - BlockHashOptions, OverlapScores, PrefillLoadHint, RouterEvent, WorkerConfigLike, WorkerId, - WorkerWithDpRank, compute_block_hash_for_seq, + BlockHashOptions, OverlapScores, PrefillLoadHint, RouterEvent, RoutingConstraints, + WorkerConfigLike, WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, }; use dynamo_kv_router::queue::DEFAULT_MAX_BATCHED_TOKENS; +use dynamo_kv_router::scheduling::SchedulingContext; use dynamo_kv_router::{ ActiveSequencesMultiWorker, DefaultWorkerSelector, PrefillTokenDeltas, RadixTree, RouterSchedulingPolicy, SchedulingPolicy, SchedulingRequest, SequenceRequest, WorkerSelector, @@ -173,6 +174,7 @@ impl PendingRequest { expected_output_tokens: self.expected_output_tokens, pinned_worker: None, allowed_worker_ids: None, + routing_constraints: RoutingConstraints::default(), shared_cache_hits: None, resp_tx: None, } @@ -441,13 +443,14 @@ impl OfflineReplayRouter { fn enqueue_key(&self, now_ms: f64, request: &PendingRequest) -> ReplayQueueKey { let arrival_offset = Duration::from_secs_f64((now_ms.max(0.0)) / 1000.0); + let scheduling_request = request.scheduling_request( + self.block_size as usize, + FxHashMap::default(), + FxHashMap::default(), + ); self.policy.enqueue_key( arrival_offset, - &request.scheduling_request( - self.block_size as usize, - FxHashMap::default(), - FxHashMap::default(), - ), + SchedulingContext::new(&scheduling_request, &self.workers_with_configs), ) } diff --git a/lib/mocker/src/replay/online/router.rs b/lib/mocker/src/replay/online/router.rs index ff6a440870ca..37a5e326570e 100644 --- a/lib/mocker/src/replay/online/router.rs +++ b/lib/mocker/src/replay/online/router.rs @@ -12,7 +12,7 @@ use dynamo_kv_router::indexer::{ KvIndexer, KvIndexerInterface, KvIndexerMetrics, ThreadPoolIndexer, }; use dynamo_kv_router::protocols::{ - BlockHashOptions, OverlapScores, RouterEvent, StorageTier, WorkerId, + BlockHashOptions, OverlapScores, RouterEvent, RoutingConstraints, StorageTier, WorkerId, }; use dynamo_kv_router::scheduling::TierOverlapBlocks; use tokio::sync::mpsc; @@ -252,6 +252,7 @@ impl KvReplayRouter { ), None, None, + RoutingConstraints::default(), None, ) .await?;