From c796c2860eb77ce849ad4bf71e5f8e883928d23a Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 10:00:07 -0700 Subject: [PATCH 01/36] add taints Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/lib.rs | 3 +- lib/bindings/python/rust/llm/kv.rs | 11 ++-- lib/bindings/python/rust/llm/local_model.rs | 53 +++++++++++++++++++ lib/bindings/python/src/dynamo/_core.pyi | 14 +++++ .../python/src/dynamo/llm/__init__.py | 1 + lib/kv-router/src/protocols.rs | 31 +++++++++++ lib/kv-router/src/scheduling/local.rs | 3 ++ lib/kv-router/src/scheduling/selector.rs | 19 +++++++ lib/kv-router/src/scheduling/types.rs | 3 +- lib/llm/src/kv_router.rs | 9 +++- .../src/kv_router/prefill_router/execution.rs | 10 +++- lib/llm/src/kv_router/push_router.rs | 3 ++ lib/llm/src/kv_router/scheduler.rs | 4 +- lib/llm/src/local_model/runtime_config.rs | 9 ++++ lib/llm/src/preprocessor.rs | 1 + lib/llm/src/protocols/common/preprocessor.rs | 6 ++- lib/llm/src/protocols/openai/nvext.rs | 7 +++ .../src/replay/offline/components/router.rs | 5 +- lib/mocker/src/replay/online/router.rs | 3 +- 19 files changed, 183 insertions(+), 12 deletions(-) diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 0c7e219d51f4..93db083eb86a 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, Taints}; 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 229487356d87..3f6fe5b913c4 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -10,6 +10,7 @@ use std::sync::mpsc; use tokio_stream::StreamExt; use super::*; +use super::local_model::Taints; use crate::Endpoint; #[cfg(feature = "kv-indexer")] use clap::Parser; @@ -936,7 +937,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, taints=None))] fn generate<'p>( &self, py: Python<'p>, @@ -952,6 +953,7 @@ impl KvRouter { block_mm_infos: Option, multi_modal_data: Option, mm_routing_info: Option, + taints: Option, ) -> PyResult> { // Depythonize the options with defaults let stop_conditions: StopConditions = if let Some(obj) = stop_conditions { @@ -1027,10 +1029,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() || taints.is_some() { let routing = llm_rs::protocols::common::preprocessor::RoutingHints { backend_instance_id: worker_id, dp_rank, + taints: taints.map(|t| t.inner), ..Default::default() }; request_builder.routing(Some(routing)); @@ -1066,7 +1069,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, taints=None))] fn best_worker<'p>( &self, py: Python<'p>, @@ -1076,6 +1079,7 @@ impl KvRouter { update_indexer: bool, block_mm_infos: Option, lora_name: Option, + taints: Option, ) -> PyResult> { let router_config_override = if let Some(obj) = router_config_override { let override_config: RouterConfigOverride = @@ -1104,6 +1108,7 @@ impl KvRouter { 0.0, None, None, // allowed_worker_ids: pass via RoutingHints in PreprocessedRequest path + taints.map(|t| t.inner).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..6af9b3071117 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -4,6 +4,47 @@ use super::*; use llm_rs::local_model::runtime_config::DisaggregatedEndpoint as RsDisaggregatedEndpoint; use llm_rs::local_model::runtime_config::ModelRuntimeConfig as RsModelRuntimeConfig; +use llm_rs::protocols::Taints as RsTaints; + +#[pyclass] +#[derive(Clone, Debug, Default)] +pub struct Taints { + pub(crate) inner: RsTaints, +} + +#[pymethods] +impl Taints { + #[new] + #[pyo3(signature = (required=None, preferred=None))] + fn new(required: Option>, preferred: Option>) -> Self { + Self { + inner: RsTaints { + required: required.unwrap_or_default(), + preferred: preferred.unwrap_or_default(), + }, + } + } + + #[getter] + fn required(&self) -> Vec { + self.inner.required.clone() + } + + #[setter] + fn set_required(&mut self, required: Vec) { + self.inner.required = required; + } + + #[getter] + fn preferred(&self) -> Vec { + self.inner.preferred.clone() + } + + #[setter] + fn set_preferred(&mut self, preferred: Vec) { + self.inner.preferred = preferred; + } +} #[pyclass] #[derive(Clone, Debug, Default)] @@ -73,6 +114,11 @@ impl ModelRuntimeConfig { self.inner.enable_eagle = enable_eagle; } + #[setter] + fn set_taints(&mut self, taints: &Taints) { + self.inner.taints = taints.inner.clone(); + } + 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 +228,11 @@ impl ModelRuntimeConfig { fn enable_eagle(&self) -> bool { self.inner.enable_eagle } + + #[getter] + fn taints(&self) -> Taints { + Taints { + inner: self.inner.taints.clone(), + } + } } diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index a3c9cc095e81..daac9ced0a2d 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -514,6 +514,7 @@ class ModelRuntimeConfig: data_parallel_size: int enable_local_indexer: bool enable_eagle: bool + taints: Taints runtime_data: dict[str, Any] tensor_model_config: Any | None bootstrap_host: str | None @@ -545,6 +546,16 @@ class ModelRuntimeConfig: """Get the tensor model configuration.""" ... +class Taints: + required: List[str] + preferred: List[str] + + def __init__( + self, + required: Optional[List[str]] = None, + preferred: Optional[List[str]] = None, + ) -> None: ... + class OverlapScores: """ A collection of prefix matching scores of workers for a given token ids. @@ -1993,6 +2004,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, + taints: Optional[Taints] = None, ) -> AsyncIterator[JsonLike]: """ Generate text using the KV-aware router. @@ -2021,6 +2033,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. + taints: Optional request taints used to constrain or prefer tainted workers. Returns: An async iterator yielding generation responses @@ -2054,6 +2067,7 @@ class KvRouter: update_indexer: bool = False, block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, + taints: Optional[Taints] = 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 757e2a6ff012..2786020e15fa 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -33,6 +33,7 @@ from dynamo._core import RouterConfig as RouterConfig from dynamo._core import RouterMode as RouterMode from dynamo._core import SglangArgs as SglangArgs +from dynamo._core import Taints as Taints from dynamo._core import WorkerMetricsPublisher as WorkerMetricsPublisher from dynamo._core import compute_block_hash_for_seq as compute_block_hash_for_seq from dynamo._core import fetch_model as fetch_model diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index ac9e0eb44387..560df4481268 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -4,6 +4,7 @@ use std::future::Future; use std::ops::Range; use std::time::Duration; +use std::sync::LazyLock; use dynamo_tokens::{SequenceHash, Token, compute_hash_v2}; use rustc_hash::FxHashMap; @@ -180,8 +181,38 @@ 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) -> &Taints { + &EMPTY_TAINTS + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +pub struct Taints { + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub required: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub preferred: Vec, } +impl Taints { + pub fn is_empty(&self) -> bool { + self.required.is_empty() && self.preferred.is_empty() + } + + pub fn is_compatible_with(&self, request: &Taints) -> bool { + request + .required + .iter() + .all(|taint| self.required.iter().any(|worker_taint| worker_taint == taint)) + && self + .required + .iter() + .all(|taint| request.required.iter().any(|request_taint| request_taint == taint)) + } +} + +static EMPTY_TAINTS: LazyLock = LazyLock::new(Taints::default); + /// Transport abstraction for publishing batched router-visible KV cache events. pub trait RouterEventSink: Send + Sync { fn publish_event(&self, event: &RouterEvent) diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 73a438776d3a..886a37b1fe9e 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -18,6 +18,7 @@ use super::types::{ KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, TierOverlapBlocks, }; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use crate::protocols::Taints; use crate::sequences::{ ActiveSequencesMultiWorker, PrefillTokenDeltas, SequenceError, SequencePublisher, SequenceRequest, @@ -185,6 +186,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, + taints: Taints, 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, + taints, router_config_override: router_config_override.cloned(), update_states, lora_name, diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index b34911f52a1c..d40106b3b396 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -215,6 +215,16 @@ impl WorkerSelector for DefaultWorkerSelector { return Err(KvSchedulerError::NoEndpoints); } + if pinned_worker.is_none() + && !request.taints.is_empty() + && workers + .iter() + .filter(|(worker_id, _)| request.is_worker_allowed(**worker_id)) + .all(|(_, config)| !config.taints().is_compatible_with(&request.taints)) + { + return Err(KvSchedulerError::NoEndpoints); + } + let request_blocks = request.request_blocks(block_size); let weights = LogitWeights { @@ -237,6 +247,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| !config.taints().is_compatible_with(&request.taints)) + { + 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); @@ -272,6 +288,9 @@ impl WorkerSelector for DefaultWorkerSelector { let worker_iter = workers .iter() .filter(move |(worker_id, _)| request.is_worker_allowed(**worker_id)) + .filter(move |(_, config)| { + request.taints.is_empty() || config.taints().is_compatible_with(&request.taints) + }) .flat_map(|(worker_id, config)| { let data_parallel_size = config.data_parallel_size(); let data_parallel_start_rank = config.data_parallel_start_rank(); diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index bd45ca3bffbb..465a6e9d1455 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -8,7 +8,7 @@ use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use super::config::RouterConfigOverride; -use crate::protocols::{DpRank, SharedCacheHits, WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use crate::protocols::{DpRank, SharedCacheHits, Taints, WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::PrefillTokenDeltas; #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -62,6 +62,7 @@ pub struct SchedulingRequest { // Routing constraints and request-level config. pub pinned_worker: Option, pub allowed_worker_ids: Option>, + pub taints: Taints, pub router_config_override: Option, pub track_prefill_tokens: bool, pub priority_jump: f64, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index b8dfd44fddf8..5e65845fb7f9 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -15,7 +15,7 @@ use dynamo_kv_router::{ protocols::KV_EVENT_SUBJECT, protocols::{ BlockExtraInfo, BlockHashOptions, DpRank, LocalBlockHash, PrefillLoadHint, RouterEvent, - RouterRequest, RouterResponse, TokensWithHashes, WorkerId, WorkerWithDpRank, + RouterRequest, RouterResponse, Taints, TokensWithHashes, WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, }, scheduling::TierOverlapBlocks, @@ -473,6 +473,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, + taints: Taints, ) -> anyhow::Result { let start = Instant::now(); @@ -577,6 +578,7 @@ where expected_output_tokens, pinned_worker, allowed_worker_ids, + taints, shared_cache_hits, ) .instrument(tracing::info_span!("kv_router.schedule")) @@ -640,6 +642,7 @@ where priority_jump: f64, expected_output_tokens: Option, allowed_worker_ids: Option>, + taints: Taints, ) -> anyhow::Result<(WorkerWithDpRank, u32)> { let result = self .find_best_match_details( @@ -653,6 +656,7 @@ where expected_output_tokens, None, allowed_worker_ids, + taints, ) .await?; Ok((result.worker, result.cache_hit.rounded_overlap_blocks())) @@ -897,6 +901,7 @@ where 0.0, None, None, + Taints::default(), ) .await?; @@ -1124,6 +1129,7 @@ mod tests { 0.0, None, None, + Taints::default(), ) .await .unwrap(); @@ -1157,6 +1163,7 @@ mod tests { 0.0, None, None, + Taints::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..f0e2225783d1 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, Taints, 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 taints = req + .routing + .as_ref() + .and_then(|r| r.taints.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, + taints, ) .await { @@ -272,6 +278,7 @@ impl PrefillRouter { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, + taints: Taints, ) -> Result<(u64, Option)> { let prefill_router = self .prefill_router @@ -292,6 +299,7 @@ impl PrefillRouter { priority_jump, None, allowed_worker_ids, + taints, ) .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 0aa59dc3342f..b525ab5c8e9a 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -302,6 +302,7 @@ 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 taints = routing.and_then(|r| r.taints.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 +319,7 @@ impl KvPushRouter { expected_output_tokens, None, allowed_worker_ids, + taints, ) .await?; let best_worker = selection.worker; @@ -374,6 +376,7 @@ impl KvPushRouter { expected_output_tokens, Some(pinned_worker), allowed_worker_ids, + taints, ) .await?; let best_worker = selection.worker; diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 04c97aa2b59c..95944c1908a0 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::{Taints, 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>, + taints: Taints, shared_cache_hits: Option, ) -> Result { let response = self @@ -159,6 +160,7 @@ where expected_output_tokens, pinned_worker, allowed_worker_ids, + taints, 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..38d8d8ea3941 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -6,6 +6,7 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize, de::DeserializeOwned}; use crate::protocols::tensor; +use dynamo_kv_router::protocols::Taints; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] pub struct DisaggregatedEndpoint { @@ -64,6 +65,9 @@ pub struct ModelRuntimeConfig { #[serde(default = "default_eagle")] pub enable_eagle: bool, + + #[serde(default, skip_serializing_if = "Taints::is_empty")] + pub taints: Taints, } const fn default_data_parallel_start_rank() -> u32 { @@ -102,6 +106,7 @@ impl Default for ModelRuntimeConfig { tensor_model_config: None, disaggregated_endpoint: None, enable_eagle: false, + taints: Taints::default(), } } } @@ -122,6 +127,10 @@ impl dynamo_kv_router::WorkerConfigLike for ModelRuntimeConfig { fn total_kv_blocks(&self) -> Option { self.total_kv_blocks } + + fn taints(&self) -> &Taints { + &self.taints + } } impl ModelRuntimeConfig { diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index b4d0d3773e41..5ab8911159f7 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -544,6 +544,7 @@ impl OpenAIPreprocessor { lora_name, allowed_worker_ids: None, session_control: nvext.session_control.clone(), + taints: nvext.taints.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..02da20f29e60 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, Taints, 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 taints used for worker compatibility and soft preference. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub taints: 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..f91046c270a8 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::Taints; use dynamo_protocols::types::StopReason; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; @@ -384,6 +385,11 @@ pub struct NvExt { #[builder(default, setter(strip_option))] #[serde(default, skip_serializing_if = "Option::is_none")] pub session_control: Option, + + /// Request taints used to constrain or prefer tainted workers. + #[builder(default, setter(strip_option))] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub taints: Option, } /// Hints from the agent/caller about request characteristics. @@ -514,6 +520,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.taints, 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..2ecaef61e8d6 100644 --- a/lib/mocker/src/replay/offline/components/router.rs +++ b/lib/mocker/src/replay/offline/components/router.rs @@ -10,8 +10,8 @@ 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, Taints, WorkerConfigLike, + WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, }; use dynamo_kv_router::queue::DEFAULT_MAX_BATCHED_TOKENS; use dynamo_kv_router::{ @@ -173,6 +173,7 @@ impl PendingRequest { expected_output_tokens: self.expected_output_tokens, pinned_worker: None, allowed_worker_ids: None, + taints: Taints::default(), shared_cache_hits: None, resp_tx: None, } diff --git a/lib/mocker/src/replay/online/router.rs b/lib/mocker/src/replay/online/router.rs index ff6a440870ca..785fab3b1368 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, StorageTier, Taints, WorkerId, }; use dynamo_kv_router::scheduling::TierOverlapBlocks; use tokio::sync::mpsc; @@ -252,6 +252,7 @@ impl KvReplayRouter { ), None, None, + Taints::default(), None, ) .await?; From bcb124c927da08725f98496088d238b5ef260291 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 10:14:44 -0700 Subject: [PATCH 02/36] add taints Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/llm/local_model.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index 6af9b3071117..bd6778b85e32 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -2,9 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use super::*; +use dynamo_kv_router::protocols::Taints as RsTaints; use llm_rs::local_model::runtime_config::DisaggregatedEndpoint as RsDisaggregatedEndpoint; use llm_rs::local_model::runtime_config::ModelRuntimeConfig as RsModelRuntimeConfig; -use llm_rs::protocols::Taints as RsTaints; #[pyclass] #[derive(Clone, Debug, Default)] From 9bc8f64842d2a8efc2e06605ed59202f41e7599c Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 10:24:31 -0700 Subject: [PATCH 03/36] add schema Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/llm/src/protocols/openai/nvext.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index f91046c270a8..b716018281e5 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -285,6 +285,22 @@ impl NvExtResponseFieldSelection { } } +/// OpenAPI-facing schema for request taints. +/// +/// Runtime serialization still uses `dynamo_kv_router::protocols::Taints`; +/// 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, Eq, Default)] +pub struct TaintsSchema { + /// Taints that must be matched for the request to be eligible for a worker. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub required: Vec, + + /// Reserved for future soft-preference routing. Currently not used by routing. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub preferred: Vec, +} + /// NVIDIA LLM extensions to the OpenAI API #[derive(ToSchema, Serialize, Deserialize, Builder, Validate, Debug, Clone)] #[validate(schema(function = "validate_nv_ext"))] @@ -389,6 +405,7 @@ pub struct NvExt { /// Request taints used to constrain or prefer tainted workers. #[builder(default, setter(strip_option))] #[serde(default, skip_serializing_if = "Option::is_none")] + #[schema(value_type = TaintsSchema)] pub taints: Option, } From c164501c604f6331407ee1e3ae90b8168b2f6a01 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 10:27:39 -0700 Subject: [PATCH 04/36] make taints simplified Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- .codex | 0 lib/bindings/python/rust/llm/local_model.rs | 10 ++++------ lib/bindings/python/src/dynamo/_core.pyi | 2 +- lib/kv-router/src/protocols.rs | 16 ++++++---------- lib/kv-router/src/scheduling/selector.rs | 17 ++++++++++++++--- lib/llm/src/local_model/runtime_config.rs | 10 ++++------ 6 files changed, 29 insertions(+), 26 deletions(-) create mode 100644 .codex diff --git a/.codex b/.codex new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index bd6778b85e32..e577136c7b14 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -115,8 +115,8 @@ impl ModelRuntimeConfig { } #[setter] - fn set_taints(&mut self, taints: &Taints) { - self.inner.taints = taints.inner.clone(); + fn set_taints(&mut self, taints: Vec) { + self.inner.taints = taints; } fn set_engine_specific(&mut self, key: &str, value: String) -> PyResult<()> { @@ -230,9 +230,7 @@ impl ModelRuntimeConfig { } #[getter] - fn taints(&self) -> Taints { - Taints { - inner: self.inner.taints.clone(), - } + fn taints(&self) -> Vec { + self.inner.taints.clone() } } diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 9c4434fd43d1..5f1e8e772957 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -514,7 +514,7 @@ class ModelRuntimeConfig: data_parallel_size: int enable_local_indexer: bool enable_eagle: bool - taints: Taints + taints: List[str] runtime_data: dict[str, Any] tensor_model_config: Any | None bootstrap_host: str | None diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 560df4481268..8c98c44da6ea 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -181,8 +181,8 @@ 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) -> &Taints { - &EMPTY_TAINTS + fn taints(&self) -> &[String] { + &EMPTY_WORKER_TAINTS } } @@ -199,19 +199,15 @@ impl Taints { self.required.is_empty() && self.preferred.is_empty() } - pub fn is_compatible_with(&self, request: &Taints) -> bool { - request + pub fn is_compatible_with_worker_taints(&self, worker_taints: &[String]) -> bool { + self .required .iter() - .all(|taint| self.required.iter().any(|worker_taint| worker_taint == taint)) - && self - .required - .iter() - .all(|taint| request.required.iter().any(|request_taint| request_taint == taint)) + .all(|taint| worker_taints.iter().any(|worker_taint| worker_taint == taint)) } } -static EMPTY_TAINTS: LazyLock = LazyLock::new(Taints::default); +static EMPTY_WORKER_TAINTS: LazyLock> = LazyLock::new(Vec::new); /// Transport abstraction for publishing batched router-visible KV cache events. pub trait RouterEventSink: Send + Sync { diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index fe090bf994d9..4f4dcc69712e 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -220,7 +220,11 @@ impl WorkerSelector for DefaultWorkerSelector { && workers .iter() .filter(|(worker_id, _)| request.is_worker_allowed(**worker_id)) - .all(|(_, config)| !config.taints().is_compatible_with(&request.taints)) + .all(|(_, config)| { + !request + .taints + .is_compatible_with_worker_taints(config.taints()) + }) { return Err(KvSchedulerError::NoEndpoints); } @@ -249,7 +253,11 @@ impl WorkerSelector for DefaultWorkerSelector { pinned_worker_config(workers, worker)?; if workers .get(&worker.worker_id) - .is_some_and(|config| !config.taints().is_compatible_with(&request.taints)) + .is_some_and(|config| { + !request + .taints + .is_compatible_with_worker_taints(config.taints()) + }) { return Err(KvSchedulerError::NoEndpoints); } @@ -289,7 +297,10 @@ impl WorkerSelector for DefaultWorkerSelector { .iter() .filter(move |(worker_id, _)| request.is_worker_allowed(**worker_id)) .filter(move |(_, config)| { - request.taints.is_empty() || config.taints().is_compatible_with(&request.taints) + request.taints.is_empty() + || request + .taints + .is_compatible_with_worker_taints(config.taints()) }) .flat_map(|(worker_id, config)| { let data_parallel_size = config.data_parallel_size(); diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 38d8d8ea3941..79d9debd7227 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -6,8 +6,6 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize, de::DeserializeOwned}; use crate::protocols::tensor; -use dynamo_kv_router::protocols::Taints; - #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] pub struct DisaggregatedEndpoint { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -66,8 +64,8 @@ pub struct ModelRuntimeConfig { #[serde(default = "default_eagle")] pub enable_eagle: bool, - #[serde(default, skip_serializing_if = "Taints::is_empty")] - pub taints: Taints, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub taints: Vec, } const fn default_data_parallel_start_rank() -> u32 { @@ -106,7 +104,7 @@ impl Default for ModelRuntimeConfig { tensor_model_config: None, disaggregated_endpoint: None, enable_eagle: false, - taints: Taints::default(), + taints: Vec::new(), } } } @@ -128,7 +126,7 @@ impl dynamo_kv_router::WorkerConfigLike for ModelRuntimeConfig { self.total_kv_blocks } - fn taints(&self) -> &Taints { + fn taints(&self) -> &[String] { &self.taints } } From 70c5c71676f0f341363361c50fe022e7f3e5ede6 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 11:02:22 -0700 Subject: [PATCH 05/36] do not, commit the codexx file ... Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- .codex | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 .codex diff --git a/.codex b/.codex deleted file mode 100644 index e69de29bb2d1..000000000000 From 304a5b66b6ef0914f2e33162d96eefc2e6c8025f Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 11:13:12 -0700 Subject: [PATCH 06/36] fix the tests Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 11 ++++++----- lib/kv-router/src/scheduling/local.rs | 16 +++++++++++++++- lib/kv-router/src/scheduling/policy.rs | 1 + lib/kv-router/src/scheduling/queue.rs | 2 ++ lib/kv-router/src/scheduling/selector.rs | 18 ++++++++++-------- lib/kv-router/src/scheduling/types.rs | 4 +++- 6 files changed, 37 insertions(+), 15 deletions(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 8c98c44da6ea..0cac6201a000 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -3,8 +3,8 @@ use std::future::Future; use std::ops::Range; -use std::time::Duration; use std::sync::LazyLock; +use std::time::Duration; use dynamo_tokens::{SequenceHash, Token, compute_hash_v2}; use rustc_hash::FxHashMap; @@ -200,10 +200,11 @@ impl Taints { } pub fn is_compatible_with_worker_taints(&self, worker_taints: &[String]) -> bool { - self - .required - .iter() - .all(|taint| worker_taints.iter().any(|worker_taint| worker_taint == taint)) + self.required.iter().all(|taint| { + worker_taints + .iter() + .any(|worker_taint| worker_taint == taint) + }) } } diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 886a37b1fe9e..671c0db04463 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -17,8 +17,8 @@ use super::selector::{DefaultWorkerSelector, WorkerSelector}; use super::types::{ KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, TierOverlapBlocks, }; -use crate::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::protocols::Taints; +use crate::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::{ ActiveSequencesMultiWorker, PrefillTokenDeltas, SequenceError, SequencePublisher, SequenceRequest, @@ -460,6 +460,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -504,6 +505,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -548,6 +550,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -593,6 +596,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -616,6 +620,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -660,6 +665,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -683,6 +689,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -741,6 +748,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -764,6 +772,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -821,6 +830,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -844,6 +854,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -899,6 +910,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -999,6 +1011,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await @@ -1097,6 +1110,7 @@ mod tests { None, None, None, + crate::protocols::Taints::default(), None, ) .await diff --git a/lib/kv-router/src/scheduling/policy.rs b/lib/kv-router/src/scheduling/policy.rs index f9a914a410ed..74e54c4507ce 100644 --- a/lib/kv-router/src/scheduling/policy.rs +++ b/lib/kv-router/src/scheduling/policy.rs @@ -148,6 +148,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + taints: crate::protocols::Taints::default(), shared_cache_hits: None, resp_tx: None, } diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 91c803221830..300e27f6143e 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -635,6 +635,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + taints: crate::protocols::Taints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -1029,6 +1030,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: Some(allowed), + taints: crate::protocols::Taints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 4f4dcc69712e..d39f352aca22 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -251,14 +251,11 @@ 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| { - !request - .taints - .is_compatible_with_worker_taints(config.taints()) - }) - { + if workers.get(&worker.worker_id).is_some_and(|config| { + !request + .taints + .is_compatible_with_worker_taints(config.taints()) + }) { return Err(KvSchedulerError::NoEndpoints); } @@ -552,6 +549,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + taints: crate::protocols::Taints::default(), shared_cache_hits: None, resp_tx: None, }; @@ -634,6 +632,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + taints: crate::protocols::Taints::default(), shared_cache_hits: Some(shared_hits), resp_tx: Some(tx), }; @@ -698,6 +697,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + taints: crate::protocols::Taints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -757,6 +757,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + taints: crate::protocols::Taints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -806,6 +807,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, + taints: crate::protocols::Taints::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 465a6e9d1455..befbd5828632 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, Taints, WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use crate::protocols::{ + DpRank, SharedCacheHits, Taints, WorkerConfigLike, WorkerId, WorkerWithDpRank, +}; use crate::sequences::PrefillTokenDeltas; #[derive(Debug, Clone, Default, Serialize, Deserialize)] From 848070a1d260f3ea5a2aa816b7af394525e9a8fb Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 11:14:05 -0700 Subject: [PATCH 07/36] fix the fmt Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/llm/kv.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index adaf011b9223..0c1d69a51fe4 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -9,8 +9,8 @@ use std::sync::atomic::AtomicU32; use std::sync::mpsc; use tokio_stream::StreamExt; -use super::*; use super::local_model::Taints; +use super::*; use crate::Endpoint; #[cfg(feature = "kv-indexer")] use clap::Parser; From 6e934bfa04559b27f7572231dd9b62b0e399d43d Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 11:31:00 -0700 Subject: [PATCH 08/36] fix the fmt Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/llm/src/kv_router/prefill_router/execution.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/llm/src/kv_router/prefill_router/execution.rs b/lib/llm/src/kv_router/prefill_router/execution.rs index f0e2225783d1..822dc039dd9a 100644 --- a/lib/llm/src/kv_router/prefill_router/execution.rs +++ b/lib/llm/src/kv_router/prefill_router/execution.rs @@ -270,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], From bd2a68185d0cdb9ce1f1cf9aa6955dfce5ffcece Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 11:41:09 -0700 Subject: [PATCH 09/36] fix the fmt Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/c/src/lib.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 0d43bbc0ba91..3fa2e42beee0 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -471,6 +471,7 @@ impl RouterHandles { lora_name, priority_jump, allowed_worker_ids, + Taints::default(), ) .await .map_err(|e| { @@ -529,6 +530,7 @@ impl RouterHandles { priority_jump, None, allowed_worker_ids, + Taints::default(), ) .await .map_err(|e| { From 7d25ea9dbda33cfced7f096a9c27ebcae2e5d839 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 12:52:22 -0700 Subject: [PATCH 10/36] adding a single good regression test and the complaint about the c bindings Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/c/src/lib.rs | 33 ++++- lib/kv-router/src/scheduling/selector.rs | 179 ++++++++++++++++++++++- 2 files changed, 203 insertions(+), 9 deletions(-) diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 3fa2e42beee0..d8055d32d24d 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>, + taints: Taints, ) -> Result<(u64, Option), QueryRouterResult> { if let Some(ref ids) = allowed_worker_ids { self.prefill_router.register_workers(ids); @@ -471,7 +473,7 @@ impl RouterHandles { lora_name, priority_jump, allowed_worker_ids, - Taints::default(), + taints, ) .await .map_err(|e| { @@ -501,6 +503,7 @@ impl RouterHandles { is_disaggregated: bool, priority_jump: f64, allowed_worker_ids: Option>, + taints: Taints, ) -> Result<(WorkerWithDpRank, u32), QueryRouterResult> { if let Some(ref ids) = allowed_worker_ids { self.decode_router.register_workers(ids); @@ -530,7 +533,7 @@ impl RouterHandles { priority_jump, None, allowed_worker_ids, - Taints::default(), + taints, ) .await .map_err(|e| { @@ -1126,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, taints)` 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 @@ -1134,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, Taints), QueryRouterResult> { let preprocessor = match &handles.preprocessor { Some(p) => p, None => { @@ -1158,6 +1161,11 @@ unsafe fn preprocess_request( }; let priority_jump = extract_priority_jump(&request); + let taints = request + .nvext + .as_ref() + .and_then(|nvext| nvext.taints.clone()) + .unwrap_or_default(); let formatted_prompt = match preprocessor.apply_template(&request) { Ok(Some(prompt)) => prompt, @@ -1184,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, taints)) } /// Parse pods JSON into an optional set of allowed worker IDs. @@ -1270,7 +1278,8 @@ pub unsafe extern "C" fn route_prefill_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump) = match unsafe { preprocess_request(handles, request_json) } { + let (tokens, priority_jump, taints) = match unsafe { preprocess_request(handles, request_json) } + { Ok(t) => t, Err(code) => return code, }; @@ -1286,6 +1295,7 @@ pub unsafe extern "C" fn route_prefill_request( None, priority_jump, allowed_worker_ids, + taints, ) .await?; @@ -1348,7 +1358,8 @@ pub unsafe extern "C" fn route_decode_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump) = match unsafe { preprocess_request(handles, request_json) } { + let (tokens, priority_jump, taints) = match unsafe { preprocess_request(handles, request_json) } + { Ok(t) => t, Err(code) => return code, }; @@ -1357,7 +1368,13 @@ pub unsafe extern "C" fn route_decode_request( 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, + taints, + ) .await?; tracing::info!( diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index d39f352aca22..4355b8d77a90 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -407,7 +407,34 @@ impl WorkerSelector for DefaultWorkerSelector { #[cfg(test)] mod tests { use super::*; - use crate::protocols::SharedCacheHits; + use crate::protocols::{SharedCacheHits, WorkerConfigLike}; + + #[derive(Clone, Default)] + struct TaintedWorkerConfig { + taints: Vec, + } + + 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) -> &[String] { + &self.taints + } + } #[test] fn test_softmax_sample_single_key() { @@ -572,6 +599,156 @@ 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: vec!["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, + taints: crate::protocols::Taints { + required: vec!["mdc-b".to_string()], + preferred: Vec::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: vec!["mdc-a".to_string()], + }, + ), + ( + 20, + TaintedWorkerConfig { + taints: vec!["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, + taints: crate::protocols::Taints { + required: vec!["mdc-b".to_string()], + preferred: Vec::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: vec![name_a.clone()], + }; + let taint_b = TaintedWorkerConfig { + taints: vec![name_b.clone()], + }; + let taint_c = TaintedWorkerConfig { + taints: vec![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, + taints: crate::protocols::Taints { + required: vec![required_taint.clone()], + preferred: Vec::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 the scoring formula with shared cache hits. /// /// Request [A, B, C, D], shared_cache_multiplier=0.5, block_size=1 From 7f731f8404e10e9ebb418c4541943f934f1d83a4 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 13:47:55 -0700 Subject: [PATCH 11/36] added taints to queue, and push-router Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/scheduling/queue.rs | 11 +++++++++- lib/llm/src/kv_router/push_router.rs | 30 ++++++++++++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 300e27f6143e..333927a1f934 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -14,7 +14,7 @@ use super::policy::{FcfsPolicy, SchedulingPolicy}; use super::prefill_load::PrefillLoadEstimator; use super::selector::{DefaultWorkerSelector, WorkerSelector}; use super::types::{SchedulingRequest, SchedulingResponse, pinned_worker_config}; -use crate::protocols::{PrefillLoadHint, WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use crate::protocols::{PrefillLoadHint, Taints, WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::{ActiveSequencesMultiWorker, SequencePublisher, SequenceRequest}; /// Large default for max_num_batched_tokens when not configured (effectively disables queueing for that worker) @@ -165,6 +165,7 @@ impl< threshold, request.allowed_worker_ids.as_ref(), request.pinned_worker, + &request.taints, decay_now, ) { tracing::debug!("all workers busy, queueing request"); @@ -217,6 +218,7 @@ impl< threshold, front.request.allowed_worker_ids.as_ref(), front.request.pinned_worker, + &front.request.taints, decay_now, ) { break; @@ -352,6 +354,7 @@ impl< threshold: f64, allowed: Option<&HashSet>, pinned_worker: Option, + taints: &Taints, decay_now: Instant, ) -> bool { let active_tokens = self.slots.active_tokens(decay_now); @@ -361,6 +364,9 @@ impl< let Ok(config) = pinned_worker_config::(&*configs, worker) else { return false; }; + if !taints.is_empty() && !taints.is_compatible_with_worker_taints(config.taints()) { + return false; + } let max_batched = config .max_num_batched_tokens() @@ -376,6 +382,9 @@ impl< { continue; } + if !taints.is_empty() && !taints.is_compatible_with_worker_taints(config.taints()) { + continue; + } let dp_size = config.data_parallel_size(); let dp_start_rank = config.data_parallel_start_rank(); let max_batched = config diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index d865ad9af935..5a943c24fedc 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}, @@ -407,6 +407,34 @@ impl KvPushRouter { "Routing to specified worker" ); + if !taints.is_empty() { + let configs = self.chooser.workers_with_configs.borrow(); + match configs.get(&pinned_worker_id) { + Some(config) if !taints.is_compatible_with_worker_taints(config.taints()) => { + tracing::warn!( + request_id = %context_id, + worker_id = pinned_worker_id, + dp_rank = ?resolved_dp_rank, + requested_taints = ?taints.required, + worker_taints = ?config.taints(), + ?phase, + "Pinned worker fallback bypassed incompatible required taints" + ); + } + None => { + tracing::warn!( + request_id = %context_id, + worker_id = pinned_worker_id, + dp_rank = ?resolved_dp_rank, + requested_taints = ?taints.required, + ?phase, + "Pinned worker fallback could not validate required taints because worker config was unavailable" + ); + } + _ => {} + } + } + // 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. From da3bce67d0699182f56aa174a23f78c68e236e33 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 13:58:22 -0700 Subject: [PATCH 12/36] added tests to make sure the queue is passing Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/scheduling/queue.rs | 40 +++++++++++++++++++++++++++ lib/kv-router/src/test_utils.rs | 6 ++++ 2 files changed, 46 insertions(+) diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 333927a1f934..cf562b0a8cb6 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -1075,6 +1075,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: vec!["mdc-a".to_string()], + ..Default::default() + }, + ); + cfg_tx.send(configs).unwrap(); + + let (mut req, rx) = make_request("tainted", 256); + req.taints = crate::protocols::Taints { + required: vec!["mdc-b".to_string()], + preferred: Vec::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/test_utils.rs b/lib/kv-router/src/test_utils.rs index 67d50184a267..48fd8f00d8fe 100644 --- a/lib/kv-router/src/test_utils.rs +++ b/lib/kv-router/src/test_utils.rs @@ -370,6 +370,7 @@ pub struct SimpleWorkerConfig { pub data_parallel_size: u32, pub max_num_batched_tokens: Option, pub total_kv_blocks: Option, + pub taints: Vec, } impl Default for SimpleWorkerConfig { @@ -379,6 +380,7 @@ impl Default for SimpleWorkerConfig { data_parallel_size: 1, max_num_batched_tokens: None, total_kv_blocks: None, + taints: Vec::new(), } } } @@ -399,4 +401,8 @@ impl WorkerConfigLike for SimpleWorkerConfig { fn total_kv_blocks(&self) -> Option { self.total_kv_blocks } + + fn taints(&self) -> &[String] { + &self.taints + } } From dd566cc3584764736ea36bb08cb0c7a991bd22bf Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 14:20:29 -0700 Subject: [PATCH 13/36] rename from Taints to RoutingContraints Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/c/src/lib.rs | 42 +++++++++---------- lib/bindings/python/rust/lib.rs | 4 +- lib/bindings/python/rust/llm/kv.rs | 16 +++---- lib/bindings/python/rust/llm/local_model.rs | 10 ++--- lib/bindings/python/src/dynamo/_core.pyi | 8 ++-- .../python/src/dynamo/llm/__init__.py | 2 +- lib/kv-router/src/protocols.rs | 4 +- lib/kv-router/src/scheduling/local.rs | 34 +++++++-------- lib/kv-router/src/scheduling/policy.rs | 2 +- lib/kv-router/src/scheduling/queue.rs | 24 +++++++---- lib/kv-router/src/scheduling/selector.rs | 26 ++++++------ lib/kv-router/src/scheduling/types.rs | 4 +- lib/llm/src/kv_router.rs | 18 ++++---- .../src/kv_router/prefill_router/execution.rs | 12 +++--- lib/llm/src/kv_router/push_router.rs | 18 ++++---- lib/llm/src/kv_router/scheduler.rs | 6 +-- lib/llm/src/preprocessor.rs | 2 +- lib/llm/src/protocols/common/preprocessor.rs | 6 +-- lib/llm/src/protocols/openai/nvext.rs | 18 ++++---- .../src/replay/offline/components/router.rs | 6 +-- lib/mocker/src/replay/online/router.rs | 4 +- 21 files changed, 138 insertions(+), 128 deletions(-) diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index d8055d32d24d..ff5e67a21bc2 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -459,7 +459,7 @@ impl RouterHandles { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, - taints: Taints, + routing_constraints: RoutingConstraints, ) -> Result<(u64, Option), QueryRouterResult> { if let Some(ref ids) = allowed_worker_ids { self.prefill_router.register_workers(ids); @@ -473,7 +473,7 @@ impl RouterHandles { lora_name, priority_jump, allowed_worker_ids, - taints, + routing_constraints, ) .await .map_err(|e| { @@ -503,7 +503,7 @@ impl RouterHandles { is_disaggregated: bool, priority_jump: f64, allowed_worker_ids: Option>, - taints: Taints, + routing_constraints: RoutingConstraints, ) -> Result<(WorkerWithDpRank, u32), QueryRouterResult> { if let Some(ref ids) = allowed_worker_ids { self.decode_router.register_workers(ids); @@ -533,7 +533,7 @@ impl RouterHandles { priority_jump, None, allowed_worker_ids, - taints, + routing_constraints, ) .await .map_err(|e| { @@ -1129,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, taints)` 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 @@ -1137,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, Taints), QueryRouterResult> { +) -> Result<(Vec, f64, RoutingConstraints), QueryRouterResult> { let preprocessor = match &handles.preprocessor { Some(p) => p, None => { @@ -1161,10 +1161,10 @@ unsafe fn preprocess_request( }; let priority_jump = extract_priority_jump(&request); - let taints = request + let routing_constraints = request .nvext .as_ref() - .and_then(|nvext| nvext.taints.clone()) + .and_then(|nvext| nvext.routing_constraints.clone()) .unwrap_or_default(); let formatted_prompt = match preprocessor.apply_template(&request) { @@ -1192,7 +1192,7 @@ unsafe fn preprocess_request( "[EPP-TOKENIZE] Tokenized prompt in C bindings (this is the ONLY tokenization)" ); - Ok((token_ids, priority_jump, taints)) + Ok((token_ids, priority_jump, routing_constraints)) } /// Parse pods JSON into an optional set of allowed worker IDs. @@ -1278,11 +1278,11 @@ pub unsafe extern "C" fn route_prefill_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump, taints) = 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) }; @@ -1295,7 +1295,7 @@ pub unsafe extern "C" fn route_prefill_request( None, priority_jump, allowed_worker_ids, - taints, + routing_constraints, ) .await?; @@ -1358,11 +1358,11 @@ pub unsafe extern "C" fn route_decode_request( let handles = unsafe { &*handle }; - let (tokens, priority_jump, taints) = 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) }; @@ -1373,7 +1373,7 @@ pub unsafe extern "C" fn route_decode_request( is_disaggregated, priority_jump, allowed_worker_ids, - taints, + routing_constraints, ) .await?; diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 93db083eb86a..59ea320296d9 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, Taints}; +use crate::llm::local_model::{ModelRuntimeConfig, RoutingConstraints}; use crate::llm::preprocessor::{MediaDecoder, MediaFetcher}; #[pyclass(eq, eq_int)] @@ -183,7 +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::()?; m.add_class::()?; diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 0c1d69a51fe4..11ee1d5b10cb 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -9,7 +9,7 @@ use std::sync::atomic::AtomicU32; use std::sync::mpsc; use tokio_stream::StreamExt; -use super::local_model::Taints; +use super::local_model::RoutingConstraints; use super::*; use crate::Endpoint; #[cfg(feature = "kv-indexer")] @@ -937,7 +937,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, taints=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>, @@ -953,7 +953,7 @@ impl KvRouter { block_mm_infos: Option, multi_modal_data: Option, mm_routing_info: Option, - taints: Option, + routing_constraints: Option, ) -> PyResult> { // Depythonize the options with defaults let stop_conditions: StopConditions = if let Some(obj) = stop_conditions { @@ -1029,11 +1029,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() || taints.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, - taints: taints.map(|t| t.inner), + routing_constraints: routing_constraints.map(|t| t.inner), ..Default::default() }; request_builder.routing(Some(routing)); @@ -1069,7 +1069,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, taints=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>, @@ -1079,7 +1079,7 @@ impl KvRouter { update_indexer: bool, block_mm_infos: Option, lora_name: Option, - taints: Option, + routing_constraints: Option, ) -> PyResult> { let router_config_override = if let Some(obj) = router_config_override { let override_config: RouterConfigOverride = @@ -1108,7 +1108,7 @@ impl KvRouter { 0.0, None, None, // allowed_worker_ids: pass via RoutingHints in PreprocessedRequest path - taints.map(|t| t.inner).unwrap_or_default(), + routing_constraints.map(|t| t.inner).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 e577136c7b14..5f6918acdfe3 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -2,23 +2,23 @@ // SPDX-License-Identifier: Apache-2.0 use super::*; -use dynamo_kv_router::protocols::Taints as RsTaints; +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 Taints { - pub(crate) inner: RsTaints, +pub struct RoutingConstraints { + pub(crate) inner: RsRoutingConstraints, } #[pymethods] -impl Taints { +impl RoutingConstraints { #[new] #[pyo3(signature = (required=None, preferred=None))] fn new(required: Option>, preferred: Option>) -> Self { Self { - inner: RsTaints { + inner: RsRoutingConstraints { required: required.unwrap_or_default(), preferred: preferred.unwrap_or_default(), }, diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 5f1e8e772957..d72dc3d0a373 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -546,7 +546,7 @@ class ModelRuntimeConfig: """Get the tensor model configuration.""" ... -class Taints: +class RoutingConstraints: required: List[str] preferred: List[str] @@ -2009,7 +2009,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, - taints: Optional[Taints] = None, + routing_constraints: Optional[RoutingConstraints] = None, ) -> AsyncIterator[JsonLike]: """ Generate text using the KV-aware router. @@ -2038,7 +2038,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. - taints: Optional request taints used to constrain or prefer tainted workers. + routing_constraints: Optional request routing constraints used to constrain or prefer tainted workers. Returns: An async iterator yielding generation responses @@ -2072,7 +2072,7 @@ class KvRouter: update_indexer: bool = False, block_mm_infos: Optional[List[Optional[Dict[str, Any]]]] = None, lora_name: Optional[str] = None, - taints: Optional[Taints] = 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 2786020e15fa..f9b0e02b956f 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -33,7 +33,7 @@ from dynamo._core import RouterConfig as RouterConfig from dynamo._core import RouterMode as RouterMode from dynamo._core import SglangArgs as SglangArgs -from dynamo._core import Taints as Taints +from dynamo._core import RoutingConstraints as RoutingConstraints from dynamo._core import WorkerMetricsPublisher as WorkerMetricsPublisher from dynamo._core import compute_block_hash_for_seq as compute_block_hash_for_seq from dynamo._core import fetch_model as fetch_model diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 23d7447d2c54..7199a07545be 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -171,14 +171,14 @@ pub trait WorkerConfigLike { } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] -pub struct Taints { +pub struct RoutingConstraints { #[serde(default, skip_serializing_if = "Vec::is_empty")] pub required: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub preferred: Vec, } -impl Taints { +impl RoutingConstraints { pub fn is_empty(&self) -> bool { self.required.is_empty() && self.preferred.is_empty() } diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 671c0db04463..e72e4d30e90b 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -17,7 +17,7 @@ use super::selector::{DefaultWorkerSelector, WorkerSelector}; use super::types::{ KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, TierOverlapBlocks, }; -use crate::protocols::Taints; +use crate::protocols::RoutingConstraints; use crate::protocols::{WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::{ ActiveSequencesMultiWorker, PrefillTokenDeltas, SequenceError, SequencePublisher, @@ -186,7 +186,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, - taints: Taints, + routing_constraints: RoutingConstraints, shared_cache_hits: Option, ) -> Result { let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); @@ -203,7 +203,7 @@ where decode_blocks: FxHashMap::default(), prefill_tokens: FxHashMap::default(), track_prefill_tokens, - taints, + routing_constraints, router_config_override: router_config_override.cloned(), update_states, lora_name, @@ -460,7 +460,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -505,7 +505,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -550,7 +550,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -596,7 +596,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -620,7 +620,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -665,7 +665,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -689,7 +689,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -748,7 +748,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -772,7 +772,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -830,7 +830,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -854,7 +854,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -910,7 +910,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -1011,7 +1011,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + crate::protocols::RoutingConstraints::default(), None, ) .await @@ -1110,7 +1110,7 @@ mod tests { None, None, None, - crate::protocols::Taints::default(), + 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 74e54c4507ce..f1ac5e4cf518 100644 --- a/lib/kv-router/src/scheduling/policy.rs +++ b/lib/kv-router/src/scheduling/policy.rs @@ -148,7 +148,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints::default(), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: None, } diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index cf562b0a8cb6..7158fbbd8c71 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -14,7 +14,9 @@ use super::policy::{FcfsPolicy, SchedulingPolicy}; use super::prefill_load::PrefillLoadEstimator; use super::selector::{DefaultWorkerSelector, WorkerSelector}; use super::types::{SchedulingRequest, SchedulingResponse, pinned_worker_config}; -use crate::protocols::{PrefillLoadHint, Taints, WorkerConfigLike, WorkerId, WorkerWithDpRank}; +use crate::protocols::{ + PrefillLoadHint, RoutingConstraints, WorkerConfigLike, WorkerId, WorkerWithDpRank, +}; use crate::sequences::{ActiveSequencesMultiWorker, SequencePublisher, SequenceRequest}; /// Large default for max_num_batched_tokens when not configured (effectively disables queueing for that worker) @@ -165,7 +167,7 @@ impl< threshold, request.allowed_worker_ids.as_ref(), request.pinned_worker, - &request.taints, + &request.routing_constraints, decay_now, ) { tracing::debug!("all workers busy, queueing request"); @@ -218,7 +220,7 @@ impl< threshold, front.request.allowed_worker_ids.as_ref(), front.request.pinned_worker, - &front.request.taints, + &front.request.routing_constraints, decay_now, ) { break; @@ -354,7 +356,7 @@ impl< threshold: f64, allowed: Option<&HashSet>, pinned_worker: Option, - taints: &Taints, + routing_constraints: &RoutingConstraints, decay_now: Instant, ) -> bool { let active_tokens = self.slots.active_tokens(decay_now); @@ -364,7 +366,9 @@ impl< let Ok(config) = pinned_worker_config::(&*configs, worker) else { return false; }; - if !taints.is_empty() && !taints.is_compatible_with_worker_taints(config.taints()) { + if !routing_constraints.is_empty() + && !routing_constraints.is_compatible_with_worker_taints(config.taints()) + { return false; } @@ -382,7 +386,9 @@ impl< { continue; } - if !taints.is_empty() && !taints.is_compatible_with_worker_taints(config.taints()) { + if !routing_constraints.is_empty() + && !routing_constraints.is_compatible_with_worker_taints(config.taints()) + { continue; } let dp_size = config.data_parallel_size(); @@ -644,7 +650,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints::default(), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -1039,7 +1045,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: Some(allowed), - taints: crate::protocols::Taints::default(), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -1103,7 +1109,7 @@ mod tests { cfg_tx.send(configs).unwrap(); let (mut req, rx) = make_request("tainted", 256); - req.taints = crate::protocols::Taints { + req.routing_constraints = crate::protocols::RoutingConstraints { required: vec!["mdc-b".to_string()], preferred: Vec::new(), }; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 4355b8d77a90..24db8aefab39 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -216,13 +216,13 @@ impl WorkerSelector for DefaultWorkerSelector { } if pinned_worker.is_none() - && !request.taints.is_empty() + && !request.routing_constraints.is_empty() && workers .iter() .filter(|(worker_id, _)| request.is_worker_allowed(**worker_id)) .all(|(_, config)| { !request - .taints + .routing_constraints .is_compatible_with_worker_taints(config.taints()) }) { @@ -253,7 +253,7 @@ impl WorkerSelector for DefaultWorkerSelector { pinned_worker_config(workers, worker)?; if workers.get(&worker.worker_id).is_some_and(|config| { !request - .taints + .routing_constraints .is_compatible_with_worker_taints(config.taints()) }) { return Err(KvSchedulerError::NoEndpoints); @@ -294,9 +294,9 @@ impl WorkerSelector for DefaultWorkerSelector { .iter() .filter(move |(worker_id, _)| request.is_worker_allowed(**worker_id)) .filter(move |(_, config)| { - request.taints.is_empty() + request.routing_constraints.is_empty() || request - .taints + .routing_constraints .is_compatible_with_worker_taints(config.taints()) }) .flat_map(|(worker_id, config)| { @@ -576,7 +576,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints::default(), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: None, }; @@ -625,7 +625,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints { + routing_constraints: crate::protocols::RoutingConstraints { required: vec!["mdc-b".to_string()], preferred: Vec::new(), }, @@ -671,7 +671,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints { + routing_constraints: crate::protocols::RoutingConstraints { required: vec!["mdc-b".to_string()], preferred: Vec::new(), }, @@ -733,7 +733,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints { + routing_constraints: crate::protocols::RoutingConstraints { required: vec![required_taint.clone()], preferred: Vec::new(), }, @@ -809,7 +809,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints::default(), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: Some(shared_hits), resp_tx: Some(tx), }; @@ -874,7 +874,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints::default(), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -934,7 +934,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints::default(), + routing_constraints: crate::protocols::RoutingConstraints::default(), shared_cache_hits: None, resp_tx: Some(tx), }; @@ -984,7 +984,7 @@ mod tests { expected_output_tokens: None, pinned_worker: None, allowed_worker_ids: None, - taints: crate::protocols::Taints::default(), + 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 befbd5828632..4d5d9857c8c3 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize}; use super::config::RouterConfigOverride; use crate::protocols::{ - DpRank, SharedCacheHits, Taints, WorkerConfigLike, WorkerId, WorkerWithDpRank, + DpRank, RoutingConstraints, SharedCacheHits, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; use crate::sequences::PrefillTokenDeltas; @@ -64,7 +64,7 @@ pub struct SchedulingRequest { // Routing constraints and request-level config. pub pinned_worker: Option, pub allowed_worker_ids: Option>, - pub taints: Taints, + pub routing_constraints: RoutingConstraints, pub router_config_override: Option, pub track_prefill_tokens: bool, pub priority_jump: f64, diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 83025528db5c..cea2610da031 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, Taints, TokensWithHashes, WorkerConfigLike, WorkerId, - WorkerWithDpRank, compute_block_hash_for_seq, + RouterRequest, RouterResponse, RoutingConstraints, TokensWithHashes, WorkerConfigLike, + WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, }, scheduling::TierOverlapBlocks, }; @@ -529,7 +529,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, - taints: Taints, + routing_constraints: RoutingConstraints, ) -> anyhow::Result { let start = Instant::now(); @@ -634,7 +634,7 @@ where expected_output_tokens, pinned_worker, allowed_worker_ids, - taints, + routing_constraints, shared_cache_hits, ) .instrument(tracing::info_span!("kv_router.schedule")) @@ -698,7 +698,7 @@ where priority_jump: f64, expected_output_tokens: Option, allowed_worker_ids: Option>, - taints: Taints, + routing_constraints: RoutingConstraints, ) -> anyhow::Result<(WorkerWithDpRank, u32)> { let result = self .find_best_match_details( @@ -712,7 +712,7 @@ where expected_output_tokens, None, allowed_worker_ids, - taints, + routing_constraints, ) .await?; Ok((result.worker, result.cache_hit.rounded_overlap_blocks())) @@ -1088,7 +1088,7 @@ where 0.0, None, None, - Taints::default(), + RoutingConstraints::default(), ) .await?; @@ -1316,7 +1316,7 @@ mod tests { 0.0, None, None, - Taints::default(), + RoutingConstraints::default(), ) .await .unwrap(); @@ -1350,7 +1350,7 @@ mod tests { 0.0, None, None, - Taints::default(), + 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 822dc039dd9a..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, Taints, 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,10 +58,10 @@ impl PrefillRouter { .routing .as_ref() .and_then(|r| r.allowed_worker_ids.clone()); - let taints = req + let routing_constraints = req .routing .as_ref() - .and_then(|r| r.taints.clone()) + .and_then(|r| r.routing_constraints.clone()) .unwrap_or_default(); let (routing_token_ids, block_mm_infos) = req.block_mm_routing_info(); match self @@ -72,7 +72,7 @@ impl PrefillRouter { lora_name, priority_jump, allowed_worker_ids, - taints, + routing_constraints, ) .await { @@ -279,7 +279,7 @@ impl PrefillRouter { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, - taints: Taints, + routing_constraints: RoutingConstraints, ) -> Result<(u64, Option)> { let prefill_router = self .prefill_router @@ -300,7 +300,7 @@ impl PrefillRouter { priority_jump, None, allowed_worker_ids, - taints, + 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 5a943c24fedc..4b4525b31054 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -302,7 +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 taints = routing.and_then(|r| r.taints.clone()).unwrap_or_default(); + 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"); @@ -319,7 +321,7 @@ impl KvPushRouter { expected_output_tokens, None, allowed_worker_ids, - taints, + routing_constraints.clone(), ) .await?; let best_worker = selection.worker; @@ -376,7 +378,7 @@ impl KvPushRouter { expected_output_tokens, Some(pinned_worker), allowed_worker_ids, - taints, + routing_constraints.clone(), ) .await?; let best_worker = selection.worker; @@ -407,15 +409,17 @@ impl KvPushRouter { "Routing to specified worker" ); - if !taints.is_empty() { + if !routing_constraints.is_empty() { let configs = self.chooser.workers_with_configs.borrow(); match configs.get(&pinned_worker_id) { - Some(config) if !taints.is_compatible_with_worker_taints(config.taints()) => { + Some(config) + if !routing_constraints.is_compatible_with_worker_taints(config.taints()) => + { tracing::warn!( request_id = %context_id, worker_id = pinned_worker_id, dp_rank = ?resolved_dp_rank, - requested_taints = ?taints.required, + requested_taints = ?routing_constraints.required, worker_taints = ?config.taints(), ?phase, "Pinned worker fallback bypassed incompatible required taints" @@ -426,7 +430,7 @@ impl KvPushRouter { request_id = %context_id, worker_id = pinned_worker_id, dp_rank = ?resolved_dp_rank, - requested_taints = ?taints.required, + requested_taints = ?routing_constraints.required, ?phase, "Pinned worker fallback could not validate required taints because worker config was unavailable" ); diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 95944c1908a0..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::{Taints, WorkerId, WorkerWithDpRank}, + protocols::{RoutingConstraints, WorkerId, WorkerWithDpRank}, }; use dynamo_runtime::component::Component; use dynamo_runtime::traits::DistributedRuntimeProvider; @@ -141,7 +141,7 @@ where expected_output_tokens: Option, pinned_worker: Option, allowed_worker_ids: Option>, - taints: Taints, + routing_constraints: RoutingConstraints, shared_cache_hits: Option, ) -> Result { let response = self @@ -160,7 +160,7 @@ where expected_output_tokens, pinned_worker, allowed_worker_ids, - taints, + routing_constraints, shared_cache_hits, ) .await; diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 9939f31cfb8b..bb6de03520a1 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -545,7 +545,7 @@ impl OpenAIPreprocessor { lora_name, allowed_worker_ids: None, session_control: nvext.session_control.clone(), - taints: nvext.taints.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 02da20f29e60..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, Taints, WorkerId}, + protocols::{BlockExtraInfo, RoutingConstraints, WorkerId}, }; use serde::{Deserialize, Serialize}; @@ -68,9 +68,9 @@ pub struct RoutingHints { #[serde(default, skip_serializing_if = "Option::is_none")] pub allowed_worker_ids: Option>, - /// Request taints used for worker compatibility and soft preference. + /// Request routing constraints used for worker compatibility and soft preference. #[serde(default, skip_serializing_if = "Option::is_none")] - pub taints: Option, + pub routing_constraints: Option, /// Session control for subagent KV isolation and sticky routing. /// Contains session_id (for affinity) and optional action (open/close). diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index b716018281e5..5c2209f35568 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -3,7 +3,7 @@ use axum::http::HeaderMap; use derive_builder::Builder; -use dynamo_kv_router::protocols::Taints; +use dynamo_kv_router::protocols::RoutingConstraints; use dynamo_protocols::types::StopReason; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; @@ -285,14 +285,14 @@ impl NvExtResponseFieldSelection { } } -/// OpenAPI-facing schema for request taints. +/// OpenAPI-facing schema for request routing constraints. /// -/// Runtime serialization still uses `dynamo_kv_router::protocols::Taints`; +/// 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, Eq, Default)] -pub struct TaintsSchema { - /// Taints that must be matched for the request to be eligible for a worker. +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: Vec, @@ -402,11 +402,11 @@ pub struct NvExt { #[serde(default, skip_serializing_if = "Option::is_none")] pub session_control: Option, - /// Request taints used to constrain or prefer tainted workers. + /// 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 = TaintsSchema)] - pub taints: Option, + #[schema(value_type = RoutingConstraintsSchema)] + pub routing_constraints: Option, } /// Hints from the agent/caller about request characteristics. @@ -537,7 +537,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.taints, 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 2ecaef61e8d6..e2c0ed56a3ce 100644 --- a/lib/mocker/src/replay/offline/components/router.rs +++ b/lib/mocker/src/replay/offline/components/router.rs @@ -10,8 +10,8 @@ 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, Taints, 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::{ @@ -173,7 +173,7 @@ impl PendingRequest { expected_output_tokens: self.expected_output_tokens, pinned_worker: None, allowed_worker_ids: None, - taints: Taints::default(), + routing_constraints: RoutingConstraints::default(), shared_cache_hits: None, resp_tx: None, } diff --git a/lib/mocker/src/replay/online/router.rs b/lib/mocker/src/replay/online/router.rs index 785fab3b1368..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, Taints, WorkerId, + BlockHashOptions, OverlapScores, RouterEvent, RoutingConstraints, StorageTier, WorkerId, }; use dynamo_kv_router::scheduling::TierOverlapBlocks; use tokio::sync::mpsc; @@ -252,7 +252,7 @@ impl KvReplayRouter { ), None, None, - Taints::default(), + RoutingConstraints::default(), None, ) .await?; From 7c3911a3e0918d4fe2588f31ca8da2924c985782 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 14:20:39 -0700 Subject: [PATCH 14/36] rename from Taints to RoutingContraints Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/src/dynamo/llm/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/bindings/python/src/dynamo/llm/__init__.py b/lib/bindings/python/src/dynamo/llm/__init__.py index f9b0e02b956f..1004f003b38f 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -32,8 +32,8 @@ 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 SglangArgs as SglangArgs 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 from dynamo._core import fetch_model as fetch_model From 04e78c8e310de3b03e41b309ed61e7422f7cb8e7 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 14:26:51 -0700 Subject: [PATCH 15/36] rename required to required_taints Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/llm/local_model.rs | 24 ++++++++++----------- lib/bindings/python/src/dynamo/_core.pyi | 8 +++---- lib/kv-router/src/protocols.rs | 8 +++---- lib/kv-router/src/scheduling/queue.rs | 4 ++-- lib/kv-router/src/scheduling/selector.rs | 12 +++++------ lib/llm/src/kv_router/push_router.rs | 4 ++-- lib/llm/src/protocols/openai/nvext.rs | 4 ++-- 7 files changed, 32 insertions(+), 32 deletions(-) diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index 5f6918acdfe3..f6d0db4e9643 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -15,34 +15,34 @@ pub struct RoutingConstraints { #[pymethods] impl RoutingConstraints { #[new] - #[pyo3(signature = (required=None, preferred=None))] - fn new(required: Option>, preferred: Option>) -> Self { + #[pyo3(signature = (required_taints=None, preferred_taints=None))] + fn new(required_taints: Option>, preferred_taints: Option>) -> Self { Self { inner: RsRoutingConstraints { - required: required.unwrap_or_default(), - preferred: preferred.unwrap_or_default(), + required_taints: required_taints.unwrap_or_default(), + preferred_taints: preferred_taints.unwrap_or_default(), }, } } #[getter] - fn required(&self) -> Vec { - self.inner.required.clone() + fn required_taints(&self) -> Vec { + self.inner.required_taints.clone() } #[setter] - fn set_required(&mut self, required: Vec) { - self.inner.required = required; + fn set_required_taints(&mut self, required_taints: Vec) { + self.inner.required_taints = required_taints; } #[getter] - fn preferred(&self) -> Vec { - self.inner.preferred.clone() + fn preferred_taints(&self) -> Vec { + self.inner.preferred_taints.clone() } #[setter] - fn set_preferred(&mut self, preferred: Vec) { - self.inner.preferred = preferred; + fn set_preferred_taints(&mut self, preferred_taints: Vec) { + self.inner.preferred_taints = preferred_taints; } } diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index d72dc3d0a373..e739e6f35ab9 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -547,13 +547,13 @@ class ModelRuntimeConfig: ... class RoutingConstraints: - required: List[str] - preferred: List[str] + required_taints: List[str] + preferred_taints: List[str] def __init__( self, - required: Optional[List[str]] = None, - preferred: Optional[List[str]] = None, + required_taints: Optional[List[str]] = None, + preferred_taints: Optional[List[str]] = None, ) -> None: ... class OverlapScores: diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 7199a07545be..ba25d46e670d 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -173,18 +173,18 @@ pub trait WorkerConfigLike { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] pub struct RoutingConstraints { #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub required: Vec, + pub required_taints: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub preferred: Vec, + pub preferred_taints: Vec, } impl RoutingConstraints { pub fn is_empty(&self) -> bool { - self.required.is_empty() && self.preferred.is_empty() + self.required_taints.is_empty() && self.preferred_taints.is_empty() } pub fn is_compatible_with_worker_taints(&self, worker_taints: &[String]) -> bool { - self.required.iter().all(|taint| { + self.required_taints.iter().all(|taint| { worker_taints .iter() .any(|worker_taint| worker_taint == taint) diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 7158fbbd8c71..bfc85ed652d2 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -1110,8 +1110,8 @@ mod tests { let (mut req, rx) = make_request("tainted", 256); req.routing_constraints = crate::protocols::RoutingConstraints { - required: vec!["mdc-b".to_string()], - preferred: Vec::new(), + required_taints: vec!["mdc-b".to_string()], + preferred_taints: Vec::new(), }; queue.enqueue(req).await; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 24db8aefab39..7c4627cb732e 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -626,8 +626,8 @@ mod tests { pinned_worker: None, allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { - required: vec!["mdc-b".to_string()], - preferred: Vec::new(), + required_taints: vec!["mdc-b".to_string()], + preferred_taints: Vec::new(), }, shared_cache_hits: None, resp_tx: None, @@ -672,8 +672,8 @@ mod tests { pinned_worker: None, allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { - required: vec!["mdc-b".to_string()], - preferred: Vec::new(), + required_taints: vec!["mdc-b".to_string()], + preferred_taints: Vec::new(), }, shared_cache_hits: None, resp_tx: None, @@ -734,8 +734,8 @@ mod tests { pinned_worker: None, allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { - required: vec![required_taint.clone()], - preferred: Vec::new(), + required_taints: vec![required_taint.clone()], + preferred_taints: Vec::new(), }, shared_cache_hits: None, resp_tx: None, diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 4b4525b31054..f3d1c17b62bc 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -419,7 +419,7 @@ impl KvPushRouter { request_id = %context_id, worker_id = pinned_worker_id, dp_rank = ?resolved_dp_rank, - requested_taints = ?routing_constraints.required, + requested_taints = ?routing_constraints.required_taints, worker_taints = ?config.taints(), ?phase, "Pinned worker fallback bypassed incompatible required taints" @@ -430,7 +430,7 @@ impl KvPushRouter { request_id = %context_id, worker_id = pinned_worker_id, dp_rank = ?resolved_dp_rank, - requested_taints = ?routing_constraints.required, + requested_taints = ?routing_constraints.required_taints, ?phase, "Pinned worker fallback could not validate required taints because worker config was unavailable" ); diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 5c2209f35568..631c73ee8a00 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -294,11 +294,11 @@ impl NvExtResponseFieldSelection { 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: Vec, + pub required_taints: Vec, /// Reserved for future soft-preference routing. Currently not used by routing. #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub preferred: Vec, + pub preferred_taints: Vec, } /// NVIDIA LLM extensions to the OpenAI API From 28e853feffa6e833e4aad388de64dfd58f35e6e6 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 14:31:08 -0700 Subject: [PATCH 16/36] add routing contraints to RouterRequest which is the whole reason we want it Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 3 +++ lib/llm/src/kv_router.rs | 3 ++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index ba25d46e670d..9587df24a8ad 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -335,6 +335,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 { @@ -350,6 +352,7 @@ impl Default for RouterRequest { RouterRequest::New { tokens: vec![], block_mm_infos: None, + routing_constraints: RoutingConstraints::default(), } } } diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index cea2610da031..96c1fa4c5f45 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1076,6 +1076,7 @@ where RouterRequest::New { tokens, block_mm_infos, + routing_constraints, } => { let (best_worker, overlap_blocks) = self .find_best_match( @@ -1088,7 +1089,7 @@ where 0.0, None, None, - RoutingConstraints::default(), + routing_constraints, ) .await?; From 179203bbe2c4fb0b9e202a88725834fc8891245e Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 14:50:02 -0700 Subject: [PATCH 17/36] taints are now hashsets Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/llm/local_model.rs | 24 ++++++++++----- .../python/src/dynamo/llm/__init__.py | 1 - lib/kv-router/src/protocols.rs | 23 +++++++------- lib/kv-router/src/scheduling/queue.rs | 6 ++-- lib/kv-router/src/scheduling/selector.rs | 30 ++++++++++--------- lib/kv-router/src/test_utils.rs | 7 +++-- lib/llm/src/local_model/runtime_config.rs | 10 +++---- 7 files changed, 55 insertions(+), 46 deletions(-) diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index f6d0db4e9643..94ffe6b9c940 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashSet; + use super::*; use dynamo_kv_router::protocols::RoutingConstraints as RsRoutingConstraints; use llm_rs::local_model::runtime_config::DisaggregatedEndpoint as RsDisaggregatedEndpoint; @@ -19,30 +21,36 @@ impl RoutingConstraints { fn new(required_taints: Option>, preferred_taints: Option>) -> Self { Self { inner: RsRoutingConstraints { - required_taints: required_taints.unwrap_or_default(), - preferred_taints: preferred_taints.unwrap_or_default(), + required_taints: required_taints + .unwrap_or_default() + .into_iter() + .collect::>(), + preferred_taints: preferred_taints + .unwrap_or_default() + .into_iter() + .collect::>(), }, } } #[getter] fn required_taints(&self) -> Vec { - self.inner.required_taints.clone() + self.inner.required_taints.iter().cloned().collect() } #[setter] fn set_required_taints(&mut self, required_taints: Vec) { - self.inner.required_taints = required_taints; + self.inner.required_taints = required_taints.into_iter().collect::>(); } #[getter] fn preferred_taints(&self) -> Vec { - self.inner.preferred_taints.clone() + self.inner.preferred_taints.iter().cloned().collect() } #[setter] fn set_preferred_taints(&mut self, preferred_taints: Vec) { - self.inner.preferred_taints = preferred_taints; + self.inner.preferred_taints = preferred_taints.into_iter().collect::>(); } } @@ -116,7 +124,7 @@ impl ModelRuntimeConfig { #[setter] fn set_taints(&mut self, taints: Vec) { - self.inner.taints = taints; + self.inner.taints = taints.into_iter().collect::>(); } fn set_engine_specific(&mut self, key: &str, value: String) -> PyResult<()> { @@ -231,6 +239,6 @@ impl ModelRuntimeConfig { #[getter] fn taints(&self) -> Vec { - self.inner.taints.clone() + self.inner.taints.iter().cloned().collect() } } diff --git a/lib/bindings/python/src/dynamo/llm/__init__.py b/lib/bindings/python/src/dynamo/llm/__init__.py index 1004f003b38f..757e2a6ff012 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -32,7 +32,6 @@ 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 9587df24a8ad..a3f82ff09f50 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashSet; use std::future::Future; use std::ops::Range; use std::sync::LazyLock; @@ -165,17 +166,17 @@ 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) -> &[String] { + fn taints(&self) -> &HashSet { &EMPTY_WORKER_TAINTS } } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] pub struct RoutingConstraints { - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub required_taints: Vec, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub preferred_taints: Vec, + #[serde(default, skip_serializing_if = "HashSet::is_empty")] + pub required_taints: HashSet, + #[serde(default, skip_serializing_if = "HashSet::is_empty")] + pub preferred_taints: HashSet, } impl RoutingConstraints { @@ -183,16 +184,14 @@ impl RoutingConstraints { self.required_taints.is_empty() && self.preferred_taints.is_empty() } - pub fn is_compatible_with_worker_taints(&self, worker_taints: &[String]) -> bool { - self.required_taints.iter().all(|taint| { - worker_taints - .iter() - .any(|worker_taint| worker_taint == taint) - }) + pub fn is_compatible_with_worker_taints(&self, worker_taints: &HashSet) -> bool { + self.required_taints + .iter() + .all(|taint| worker_taints.contains(taint)) } } -static EMPTY_WORKER_TAINTS: LazyLock> = LazyLock::new(Vec::new); +static EMPTY_WORKER_TAINTS: LazyLock> = LazyLock::new(HashSet::new); /// Transport abstraction for publishing batched router-visible KV cache events. pub trait RouterEventSink: Send + Sync { diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index bfc85ed652d2..fc1732e4890b 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -1102,7 +1102,7 @@ mod tests { 0_u64, SimpleWorkerConfig { max_num_batched_tokens: Some(256), - taints: vec!["mdc-a".to_string()], + taints: HashSet::from(["mdc-a".to_string()]), ..Default::default() }, ); @@ -1110,8 +1110,8 @@ mod tests { let (mut req, rx) = make_request("tainted", 256); req.routing_constraints = crate::protocols::RoutingConstraints { - required_taints: vec!["mdc-b".to_string()], - preferred_taints: Vec::new(), + required_taints: HashSet::from(["mdc-b".to_string()]), + preferred_taints: HashSet::new(), }; queue.enqueue(req).await; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 7c4627cb732e..d776d2b74e29 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -406,12 +406,14 @@ impl WorkerSelector for DefaultWorkerSelector { #[cfg(test)] mod tests { + use std::collections::HashSet; + use super::*; use crate::protocols::{SharedCacheHits, WorkerConfigLike}; #[derive(Clone, Default)] struct TaintedWorkerConfig { - taints: Vec, + taints: HashSet, } impl WorkerConfigLike for TaintedWorkerConfig { @@ -431,7 +433,7 @@ mod tests { None } - fn taints(&self) -> &[String] { + fn taints(&self) -> &HashSet { &self.taints } } @@ -605,7 +607,7 @@ mod tests { let workers = HashMap::from([( 10, TaintedWorkerConfig { - taints: vec!["mdc-a".to_string()], + taints: HashSet::from(["mdc-a".to_string()]), }, )]); let request = SchedulingRequest { @@ -626,8 +628,8 @@ mod tests { pinned_worker: None, allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { - required_taints: vec!["mdc-b".to_string()], - preferred_taints: Vec::new(), + required_taints: HashSet::from(["mdc-b".to_string()]), + preferred_taints: HashSet::new(), }, shared_cache_hits: None, resp_tx: None, @@ -644,13 +646,13 @@ mod tests { ( 10, TaintedWorkerConfig { - taints: vec!["mdc-a".to_string()], + taints: HashSet::from(["mdc-a".to_string()]), }, ), ( 20, TaintedWorkerConfig { - taints: vec!["mdc-b".to_string()], + taints: HashSet::from(["mdc-b".to_string()]), }, ), ]); @@ -672,8 +674,8 @@ mod tests { pinned_worker: None, allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { - required_taints: vec!["mdc-b".to_string()], - preferred_taints: Vec::new(), + required_taints: HashSet::from(["mdc-b".to_string()]), + preferred_taints: HashSet::new(), }, shared_cache_hits: None, resp_tx: None, @@ -690,13 +692,13 @@ mod tests { let name_b = "mdc-b".to_string(); let name_c = "mdc-c".to_string(); let taint_a = TaintedWorkerConfig { - taints: vec![name_a.clone()], + taints: HashSet::from([name_a.clone()]), }; let taint_b = TaintedWorkerConfig { - taints: vec![name_b.clone()], + taints: HashSet::from([name_b.clone()]), }; let taint_c = TaintedWorkerConfig { - taints: vec![name_c.clone()], + taints: HashSet::from([name_c.clone()]), }; let workers = HashMap::from([ (10, taint_a.clone()), @@ -734,8 +736,8 @@ mod tests { pinned_worker: None, allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { - required_taints: vec![required_taint.clone()], - preferred_taints: Vec::new(), + required_taints: HashSet::from([required_taint.clone()]), + preferred_taints: HashSet::new(), }, shared_cache_hits: None, resp_tx: None, diff --git a/lib/kv-router/src/test_utils.rs b/lib/kv-router/src/test_utils.rs index 48fd8f00d8fe..ce6817dd50c9 100644 --- a/lib/kv-router/src/test_utils.rs +++ b/lib/kv-router/src/test_utils.rs @@ -370,7 +370,7 @@ pub struct SimpleWorkerConfig { pub data_parallel_size: u32, pub max_num_batched_tokens: Option, pub total_kv_blocks: Option, - pub taints: Vec, + pub taints: HashSet, } impl Default for SimpleWorkerConfig { @@ -380,7 +380,7 @@ impl Default for SimpleWorkerConfig { data_parallel_size: 1, max_num_batched_tokens: None, total_kv_blocks: None, - taints: Vec::new(), + taints: HashSet::new(), } } } @@ -402,7 +402,8 @@ impl WorkerConfigLike for SimpleWorkerConfig { self.total_kv_blocks } - fn taints(&self) -> &[String] { + fn taints(&self) -> &HashSet { &self.taints } } +use std::collections::HashSet; diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 79d9debd7227..170e9fb8302d 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -1,7 +1,7 @@ // 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}; @@ -64,8 +64,8 @@ pub struct ModelRuntimeConfig { #[serde(default = "default_eagle")] pub enable_eagle: bool, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub taints: Vec, + #[serde(default, skip_serializing_if = "HashSet::is_empty")] + pub taints: HashSet, } const fn default_data_parallel_start_rank() -> u32 { @@ -104,7 +104,7 @@ impl Default for ModelRuntimeConfig { tensor_model_config: None, disaggregated_endpoint: None, enable_eagle: false, - taints: Vec::new(), + taints: HashSet::new(), } } } @@ -126,7 +126,7 @@ impl dynamo_kv_router::WorkerConfigLike for ModelRuntimeConfig { self.total_kv_blocks } - fn taints(&self) -> &[String] { + fn taints(&self) -> &HashSet { &self.taints } } From 5a4050d027fbb8a1d7fe07a0ecb9a5ed81442855 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 15:14:21 -0700 Subject: [PATCH 18/36] python api should be hashset Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/llm/kv.rs | 4 +- lib/bindings/python/rust/llm/local_model.rs | 57 +++++++-------------- lib/bindings/python/src/dynamo/_core.pyi | 10 ++-- 3 files changed, 26 insertions(+), 45 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 11ee1d5b10cb..cdf0965ea566 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -1033,7 +1033,7 @@ impl KvRouter { let routing = llm_rs::protocols::common::preprocessor::RoutingHints { backend_instance_id: worker_id, dp_rank, - routing_constraints: routing_constraints.map(|t| t.inner), + routing_constraints: routing_constraints.map(Into::into), ..Default::default() }; request_builder.routing(Some(routing)); @@ -1108,7 +1108,7 @@ impl KvRouter { 0.0, None, None, // allowed_worker_ids: pass via RoutingHints in PreprocessedRequest path - routing_constraints.map(|t| t.inner).unwrap_or_default(), + 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 94ffe6b9c940..60787492b218 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -11,46 +11,27 @@ use llm_rs::local_model::runtime_config::ModelRuntimeConfig as RsModelRuntimeCon #[pyclass] #[derive(Clone, Debug, Default)] pub struct RoutingConstraints { - pub(crate) inner: RsRoutingConstraints, + #[pyo3(get, set)] + pub required_taints: HashSet, + #[pyo3(get, set)] + pub preferred_taints: HashSet, } -#[pymethods] -impl RoutingConstraints { - #[new] - #[pyo3(signature = (required_taints=None, preferred_taints=None))] - fn new(required_taints: Option>, preferred_taints: Option>) -> Self { +impl From for RsRoutingConstraints { + fn from(value: RoutingConstraints) -> Self { Self { - inner: RsRoutingConstraints { - required_taints: required_taints - .unwrap_or_default() - .into_iter() - .collect::>(), - preferred_taints: preferred_taints - .unwrap_or_default() - .into_iter() - .collect::>(), - }, + required_taints: value.required_taints, + preferred_taints: value.preferred_taints, } } +} - #[getter] - fn required_taints(&self) -> Vec { - self.inner.required_taints.iter().cloned().collect() - } - - #[setter] - fn set_required_taints(&mut self, required_taints: Vec) { - self.inner.required_taints = required_taints.into_iter().collect::>(); - } - - #[getter] - fn preferred_taints(&self) -> Vec { - self.inner.preferred_taints.iter().cloned().collect() - } - - #[setter] - fn set_preferred_taints(&mut self, preferred_taints: Vec) { - self.inner.preferred_taints = preferred_taints.into_iter().collect::>(); +impl From for RoutingConstraints { + fn from(value: RsRoutingConstraints) -> Self { + Self { + required_taints: value.required_taints, + preferred_taints: value.preferred_taints, + } } } @@ -123,8 +104,8 @@ impl ModelRuntimeConfig { } #[setter] - fn set_taints(&mut self, taints: Vec) { - self.inner.taints = taints.into_iter().collect::>(); + fn set_taints(&mut self, taints: HashSet) { + self.inner.taints = taints; } fn set_engine_specific(&mut self, key: &str, value: String) -> PyResult<()> { @@ -238,7 +219,7 @@ impl ModelRuntimeConfig { } #[getter] - fn taints(&self) -> Vec { - self.inner.taints.iter().cloned().collect() + 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 e739e6f35ab9..b05b5cdbff91 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -514,7 +514,7 @@ class ModelRuntimeConfig: data_parallel_size: int enable_local_indexer: bool enable_eagle: bool - taints: List[str] + taints: Set[str] runtime_data: dict[str, Any] tensor_model_config: Any | None bootstrap_host: str | None @@ -547,13 +547,13 @@ class ModelRuntimeConfig: ... class RoutingConstraints: - required_taints: List[str] - preferred_taints: List[str] + required_taints: Set[str] + preferred_taints: Set[str] def __init__( self, - required_taints: Optional[List[str]] = None, - preferred_taints: Optional[List[str]] = None, + required_taints: Optional[Set[str]] = None, + preferred_taints: Optional[Set[str]] = None, ) -> None: ... class OverlapScores: From 250562b37f85c0e5af317b8cbd02afdcd5c61725 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 15:16:38 -0700 Subject: [PATCH 19/36] python api should be hashset -> now also in core.pyi Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/src/dynamo/_core.pyi | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index b05b5cdbff91..7715d52aad8f 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, ) From bffdef8fca75d8d51b1ec0c22bf279789daa03ea Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 15:19:34 -0700 Subject: [PATCH 20/36] pinned worker + required taints -> error instead of warning Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/llm/src/kv_router/push_router.rs | 30 ++++++++++++---------------- 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index f3d1c17b62bc..027993e7e06d 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -415,25 +415,21 @@ impl KvPushRouter { Some(config) if !routing_constraints.is_compatible_with_worker_taints(config.taints()) => { - tracing::warn!( - request_id = %context_id, - worker_id = pinned_worker_id, - dp_rank = ?resolved_dp_rank, - requested_taints = ?routing_constraints.required_taints, - worker_taints = ?config.taints(), - ?phase, - "Pinned worker fallback bypassed incompatible required taints" - ); + return Err(anyhow::anyhow!( + "Pinned worker {} does not satisfy required taints {:?}; worker taints: {:?}", + pinned_worker_id, + routing_constraints.required_taints, + config.taints() + ) + .into()); } None => { - tracing::warn!( - request_id = %context_id, - worker_id = pinned_worker_id, - dp_rank = ?resolved_dp_rank, - requested_taints = ?routing_constraints.required_taints, - ?phase, - "Pinned worker fallback could not validate required taints because worker config was unavailable" - ); + 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 + ) + .into()); } _ => {} } From ccd5d3b8dc2e4243684ccfd14beb043806bd5c3a Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 15:31:08 -0700 Subject: [PATCH 21/36] pinned worker + required taints -> error instead of warning FMT Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/llm/src/kv_router/push_router.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 027993e7e06d..6d4070c1ad05 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -420,16 +420,14 @@ impl KvPushRouter { pinned_worker_id, routing_constraints.required_taints, config.taints() - ) - .into()); + )); } 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 - ) - .into()); + )); } _ => {} } From 93f702e92733a738cbb44ce3d843ff9f069f2d61 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 16:08:41 -0700 Subject: [PATCH 22/36] add a prefered taint bias Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 7 +++ lib/kv-router/src/scheduling/config.rs | 15 +++++ lib/kv-router/src/scheduling/selector.rs | 75 +++++++++++++++++++++++- 3 files changed, 96 insertions(+), 1 deletion(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index a3f82ff09f50..ba7aa09bbee6 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -189,6 +189,13 @@ impl RoutingConstraints { .iter() .all(|taint| worker_taints.contains(taint)) } + + pub fn preferred_taint_matches(&self, worker_taints: &HashSet) -> usize { + self.preferred_taints + .iter() + .filter(|taint| worker_taints.contains(*taint)) + .count() + } } static EMPTY_WORKER_TAINTS: LazyLock> = LazyLock::new(HashSet::new); diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index e0723980b70b..f3d22e24a77e 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -47,6 +47,10 @@ const fn default_prefill_load_scale() -> f64 { 1.0 } +const fn default_preferred_taint_bias_pct() -> f64 { + 0.0 +} + pub const OVERLAP_SCORE_CREDIT_RANGE_ERROR: &str = "overlap_score_credit must be between 0.0 and 1.0"; pub const OVERLAP_SCORE_CREDIT_MIGRATION_ERROR: &str = concat!( @@ -292,6 +296,7 @@ struct KvRouterConfigSerde { host_cache_hit_weight: f64, disk_cache_hit_weight: f64, router_temperature: f64, + preferred_taint_bias_pct: f64, use_kv_events: bool, durable_kv_events: bool, router_replica_sync: bool, @@ -324,6 +329,7 @@ impl Default for KvRouterConfigSerde { host_cache_hit_weight: config.host_cache_hit_weight, disk_cache_hit_weight: config.disk_cache_hit_weight, router_temperature: config.router_temperature, + preferred_taint_bias_pct: config.preferred_taint_bias_pct, use_kv_events: config.use_kv_events, durable_kv_events: config.durable_kv_events, router_replica_sync: config.router_replica_sync, @@ -374,6 +380,13 @@ pub struct KvRouterConfig { #[validate(range(min = 0.0))] pub router_temperature: f64, + /// Fractional score reduction applied once per preferred taint match. + /// Higher routing scores are worse, so matching preferred taints lowers + /// the score multiplicatively. Default: 0.0 (disabled). + #[serde(default = "default_preferred_taint_bias_pct")] + #[validate(range(min = 0.0, max = 1.0))] + pub preferred_taint_bias_pct: f64, + pub use_kv_events: bool, /// **Deprecated:** Enable durable KV events using NATS JetStream instead of the default event plane. @@ -476,6 +489,7 @@ impl Default for KvRouterConfig { host_cache_hit_weight: default_host_cache_hit_weight(), disk_cache_hit_weight: default_disk_cache_hit_weight(), router_temperature: 0.0, + preferred_taint_bias_pct: default_preferred_taint_bias_pct(), use_kv_events: true, durable_kv_events: false, // default to NATS Core (local indexer mode) router_replica_sync: false, @@ -521,6 +535,7 @@ impl TryFrom for KvRouterConfig { host_cache_hit_weight: compat.host_cache_hit_weight, disk_cache_hit_weight: compat.disk_cache_hit_weight, router_temperature: compat.router_temperature, + preferred_taint_bias_pct: compat.preferred_taint_bias_pct, use_kv_events: compat.use_kv_events, durable_kv_events: compat.durable_kv_events, router_replica_sync: compat.router_replica_sync, diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index d776d2b74e29..5aed774392f0 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -285,9 +285,25 @@ impl WorkerSelector for DefaultWorkerSelector { .as_ref() .and_then(|cfg| cfg.router_temperature) .unwrap_or(self.kv_router_config.router_temperature); + let preferred_taint_bias_pct = self.kv_router_config.preferred_taint_bias_pct; 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"); + if preferred_taint_bias_pct == 0.0 { + return base_score; + } + let Some(config) = workers.get(&worker.worker_id) else { + return base_score; + }; + let preferred_matches = request + .routing_constraints + .preferred_taint_matches(config.taints()); + if preferred_matches == 0 { + return base_score; + } + let preferred_multiplier = + (1.0 - preferred_taint_bias_pct).powi(preferred_matches as i32); + base_score * preferred_multiplier }; let worker_iter = workers @@ -751,6 +767,63 @@ mod tests { } } + #[test] + fn test_preferred_taint_bias_pct_prefers_matching_worker() { + let selector = DefaultWorkerSelector::new( + Some(KvRouterConfig { + router_temperature: 0.0, + preferred_taint_bias_pct: 0.25, + ..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: HashSet::from(["mdc-a".to_string()]), + }, + shared_cache_hits: None, + resp_tx: None, + }; + + let result = selector.select_worker(&workers, &request, 16).unwrap(); + assert_eq!(result.worker.worker_id, 10); + } + /// Test the scoring formula with shared cache hits. /// /// Request [A, B, C, D], shared_cache_multiplier=0.5, block_size=1 From 34145569d0ab5c17f3871c6ae013c43771b56551 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 18:09:37 -0700 Subject: [PATCH 23/36] add a prefered taint bias as hashmap, with individual scores Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/llm/local_model.rs | 4 +- lib/bindings/python/src/dynamo/_core.pyi | 12 ++- lib/kv-router/src/protocols.rs | 20 +++-- lib/kv-router/src/scheduling/config.rs | 15 ---- lib/kv-router/src/scheduling/queue.rs | 2 +- lib/kv-router/src/scheduling/selector.rs | 85 ++++++++++++++++----- lib/llm/src/protocols/openai/nvext.rs | 23 ++++-- 7 files changed, 111 insertions(+), 50 deletions(-) diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index 60787492b218..73965f369d7d 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use super::*; use dynamo_kv_router::protocols::RoutingConstraints as RsRoutingConstraints; @@ -14,7 +14,7 @@ pub struct RoutingConstraints { #[pyo3(get, set)] pub required_taints: HashSet, #[pyo3(get, set)] - pub preferred_taints: HashSet, + pub preferred_taints: HashMap, } impl From for RsRoutingConstraints { diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 7715d52aad8f..68238df5938d 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -548,13 +548,21 @@ class ModelRuntimeConfig: ... class RoutingConstraints: + """ + Request-side routing constraints. + + ``required_taints`` is a hard eligibility filter. + ``preferred_taints`` maps taint -> weight in ``(-1.0, 1.0)``. + Positive weights prefer matching workers, negative weights avoid them, + and ``0.0`` is neutral. + """ required_taints: Set[str] - preferred_taints: Set[str] + preferred_taints: Dict[str, float] def __init__( self, required_taints: Optional[Set[str]] = None, - preferred_taints: Optional[Set[str]] = None, + preferred_taints: Optional[Dict[str, float]] = None, ) -> None: ... class OverlapScores: diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index ba7aa09bbee6..4a3f037ec57b 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::ops::Range; use std::sync::LazyLock; @@ -171,12 +171,12 @@ pub trait WorkerConfigLike { } } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +#[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 = "HashSet::is_empty")] - pub preferred_taints: HashSet, + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub preferred_taints: HashMap, } impl RoutingConstraints { @@ -192,10 +192,20 @@ impl RoutingConstraints { pub fn preferred_taint_matches(&self, worker_taints: &HashSet) -> usize { self.preferred_taints - .iter() + .keys() .filter(|taint| worker_taints.contains(*taint)) .count() } + + pub fn preferred_taint_multiplier(&self, worker_taints: &HashSet) -> f64 { + self.preferred_taints + .iter() + .filter(|(taint, _)| worker_taints.contains(*taint)) + .fold(1.0, |multiplier, (_, weight)| { + let weight = f64::from(*weight).clamp(-0.999_999, 0.999_999); + multiplier * (1.0 - weight) + }) + } } static EMPTY_WORKER_TAINTS: LazyLock> = LazyLock::new(HashSet::new); diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index f3d22e24a77e..e0723980b70b 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -47,10 +47,6 @@ const fn default_prefill_load_scale() -> f64 { 1.0 } -const fn default_preferred_taint_bias_pct() -> f64 { - 0.0 -} - pub const OVERLAP_SCORE_CREDIT_RANGE_ERROR: &str = "overlap_score_credit must be between 0.0 and 1.0"; pub const OVERLAP_SCORE_CREDIT_MIGRATION_ERROR: &str = concat!( @@ -296,7 +292,6 @@ struct KvRouterConfigSerde { host_cache_hit_weight: f64, disk_cache_hit_weight: f64, router_temperature: f64, - preferred_taint_bias_pct: f64, use_kv_events: bool, durable_kv_events: bool, router_replica_sync: bool, @@ -329,7 +324,6 @@ impl Default for KvRouterConfigSerde { host_cache_hit_weight: config.host_cache_hit_weight, disk_cache_hit_weight: config.disk_cache_hit_weight, router_temperature: config.router_temperature, - preferred_taint_bias_pct: config.preferred_taint_bias_pct, use_kv_events: config.use_kv_events, durable_kv_events: config.durable_kv_events, router_replica_sync: config.router_replica_sync, @@ -380,13 +374,6 @@ pub struct KvRouterConfig { #[validate(range(min = 0.0))] pub router_temperature: f64, - /// Fractional score reduction applied once per preferred taint match. - /// Higher routing scores are worse, so matching preferred taints lowers - /// the score multiplicatively. Default: 0.0 (disabled). - #[serde(default = "default_preferred_taint_bias_pct")] - #[validate(range(min = 0.0, max = 1.0))] - pub preferred_taint_bias_pct: f64, - pub use_kv_events: bool, /// **Deprecated:** Enable durable KV events using NATS JetStream instead of the default event plane. @@ -489,7 +476,6 @@ impl Default for KvRouterConfig { host_cache_hit_weight: default_host_cache_hit_weight(), disk_cache_hit_weight: default_disk_cache_hit_weight(), router_temperature: 0.0, - preferred_taint_bias_pct: default_preferred_taint_bias_pct(), use_kv_events: true, durable_kv_events: false, // default to NATS Core (local indexer mode) router_replica_sync: false, @@ -535,7 +521,6 @@ impl TryFrom for KvRouterConfig { host_cache_hit_weight: compat.host_cache_hit_weight, disk_cache_hit_weight: compat.disk_cache_hit_weight, router_temperature: compat.router_temperature, - preferred_taint_bias_pct: compat.preferred_taint_bias_pct, use_kv_events: compat.use_kv_events, durable_kv_events: compat.durable_kv_events, router_replica_sync: compat.router_replica_sync, diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index fc1732e4890b..a1452f3f2dc1 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -1111,7 +1111,7 @@ mod tests { let (mut req, rx) = make_request("tainted", 256); req.routing_constraints = crate::protocols::RoutingConstraints { required_taints: HashSet::from(["mdc-b".to_string()]), - preferred_taints: HashSet::new(), + preferred_taints: HashMap::new(), }; queue.enqueue(req).await; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 5aed774392f0..146c8c07d7c5 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -285,25 +285,15 @@ impl WorkerSelector for DefaultWorkerSelector { .as_ref() .and_then(|cfg| cfg.router_temperature) .unwrap_or(self.kv_router_config.router_temperature); - let preferred_taint_bias_pct = self.kv_router_config.preferred_taint_bias_pct; - let get_score = |worker: WorkerWithDpRank| -> f64 { let base_score = self.worker_logit(request, worker, block_size, weights, "Formula"); - if preferred_taint_bias_pct == 0.0 { - return base_score; - } let Some(config) = workers.get(&worker.worker_id) else { return base_score; }; - let preferred_matches = request - .routing_constraints - .preferred_taint_matches(config.taints()); - if preferred_matches == 0 { - return base_score; - } - let preferred_multiplier = - (1.0 - preferred_taint_bias_pct).powi(preferred_matches as i32); - base_score * preferred_multiplier + base_score + * request + .routing_constraints + .preferred_taint_multiplier(config.taints()) }; let worker_iter = workers @@ -645,7 +635,7 @@ mod tests { allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { required_taints: HashSet::from(["mdc-b".to_string()]), - preferred_taints: HashSet::new(), + preferred_taints: HashMap::new(), }, shared_cache_hits: None, resp_tx: None, @@ -691,7 +681,7 @@ mod tests { allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { required_taints: HashSet::from(["mdc-b".to_string()]), - preferred_taints: HashSet::new(), + preferred_taints: HashMap::new(), }, shared_cache_hits: None, resp_tx: None, @@ -753,7 +743,7 @@ mod tests { allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { required_taints: HashSet::from([required_taint.clone()]), - preferred_taints: HashSet::new(), + preferred_taints: HashMap::new(), }, shared_cache_hits: None, resp_tx: None, @@ -768,11 +758,10 @@ mod tests { } #[test] - fn test_preferred_taint_bias_pct_prefers_matching_worker() { + fn test_preferred_taints_prefer_matching_worker() { let selector = DefaultWorkerSelector::new( Some(KvRouterConfig { router_temperature: 0.0, - preferred_taint_bias_pct: 0.25, ..Default::default() }), "test", @@ -814,7 +803,7 @@ mod tests { allowed_worker_ids: None, routing_constraints: crate::protocols::RoutingConstraints { required_taints: HashSet::new(), - preferred_taints: HashSet::from(["mdc-a".to_string()]), + preferred_taints: HashMap::from([("mdc-a".to_string(), 0.85)]), }, shared_cache_hits: None, resp_tx: None, @@ -824,6 +813,62 @@ mod tests { 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 diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 631c73ee8a00..3865408f5684 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -290,15 +290,17 @@ impl NvExtResponseFieldSelection { /// 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, Eq, Default)] +#[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, - /// Reserved for future soft-preference routing. Currently not used by routing. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub preferred_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. + #[serde(default, skip_serializing_if = "std::collections::HashMap::is_empty")] + pub preferred_taints: std::collections::HashMap, } /// NVIDIA LLM extensions to the OpenAI API @@ -499,7 +501,18 @@ impl NvExt { } } -fn validate_nv_ext(_nv_ext: &NvExt) -> Result<(), ValidationError> { +fn validate_nv_ext(nv_ext: &NvExt) -> Result<(), ValidationError> { + if let Some(routing_constraints) = nv_ext.routing_constraints.as_ref() + && routing_constraints + .preferred_taints + .values() + .any(|weight| !(-1.0..1.0).contains(weight)) + { + let mut error = ValidationError::new("preferred_taint_weight_out_of_range"); + error.message = Some("preferred taint weights must be in the range (-1.0, 1.0)".into()); + return Err(error); + } + Ok(()) } From 5c122a9435364ea1f62dd919c9d9969a34d7bc26 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 20:52:51 -0700 Subject: [PATCH 24/36] address comments Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/llm/src/kv_router/push_router.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 6d4070c1ad05..34f8ab2b7cea 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -409,7 +409,7 @@ impl KvPushRouter { "Routing to specified worker" ); - if !routing_constraints.is_empty() { + if routing_constraints.has_hard_constraints() { let configs = self.chooser.workers_with_configs.borrow(); match configs.get(&pinned_worker_id) { Some(config) From 8b7487fa99ea52df56e9696905ed1e4cbcafb21f Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 20:53:03 -0700 Subject: [PATCH 25/36] address comments Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 4a3f037ec57b..0068f3d98098 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -184,6 +184,10 @@ impl RoutingConstraints { 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 { self.required_taints .iter() From bf2571b699216aa1f635fbe963ae8d176937c663 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 21:05:09 -0700 Subject: [PATCH 26/36] verify routing constraints with selector Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/rust/llm/local_model.rs | 15 +++++++++++ lib/kv-router/src/scheduling/queue.rs | 4 +-- lib/kv-router/src/scheduling/selector.rs | 29 ++++++++++----------- 3 files changed, 31 insertions(+), 17 deletions(-) diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index 73965f369d7d..e03cd9eefb6e 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -17,6 +17,21 @@ pub struct RoutingConstraints { 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 { diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index a1452f3f2dc1..02b824bb6324 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -366,7 +366,7 @@ impl< let Ok(config) = pinned_worker_config::(&*configs, worker) else { return false; }; - if !routing_constraints.is_empty() + if routing_constraints.has_hard_constraints() && !routing_constraints.is_compatible_with_worker_taints(config.taints()) { return false; @@ -386,7 +386,7 @@ impl< { continue; } - if !routing_constraints.is_empty() + if routing_constraints.has_hard_constraints() && !routing_constraints.is_compatible_with_worker_taints(config.taints()) { continue; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 146c8c07d7c5..a93ad5952941 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -207,26 +207,25 @@ impl WorkerSelector for DefaultWorkerSelector { let pinned_worker = request.pinned_worker; - if pinned_worker.is_none() - && !workers - .keys() - .any(|worker_id| request.is_worker_allowed(*worker_id)) - { - return Err(KvSchedulerError::NoEndpoints); - } - - if pinned_worker.is_none() - && !request.routing_constraints.is_empty() - && workers + if pinned_worker.is_none() { + let allowed_workers: Vec<_> = workers .iter() .filter(|(worker_id, _)| request.is_worker_allowed(**worker_id)) - .all(|(_, config)| { - !request + .collect(); + + if allowed_workers.is_empty() { + return Err(KvSchedulerError::NoEndpoints); + } + + if request.routing_constraints.has_hard_constraints() + && !allowed_workers.iter().any(|(_, config)| { + request .routing_constraints .is_compatible_with_worker_taints(config.taints()) }) - { - return Err(KvSchedulerError::NoEndpoints); + { + return Err(KvSchedulerError::NoEndpoints); + } } let request_blocks = request.request_blocks(block_size); From 7469e8985add4f8f64af6ca9e389a2591ff8fdca Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 21:32:17 -0700 Subject: [PATCH 27/36] add a comment Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 0068f3d98098..cc2e9e63addc 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -206,7 +206,9 @@ impl RoutingConstraints { .iter() .filter(|(taint, _)| worker_taints.contains(*taint)) .fold(1.0, |multiplier, (_, weight)| { - let weight = f64::from(*weight).clamp(-0.999_999, 0.999_999); + // Clamp weight to (-0.999, 0.999) to avoid multiplier <= 0 and numerical instabilities + // which would invert or zero out the score. + let weight = f64::from(*weight).clamp(-0.999, 0.999); multiplier * (1.0 - weight) }) } From 32b141ce88e499d9d5b8a9a408f4b49340793e55 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Thu, 14 May 2026 22:16:12 -0700 Subject: [PATCH 28/36] fmt Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index cc2e9e63addc..0fcc12625bd2 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -206,7 +206,7 @@ impl RoutingConstraints { .iter() .filter(|(taint, _)| worker_taints.contains(*taint)) .fold(1.0, |multiplier, (_, weight)| { - // Clamp weight to (-0.999, 0.999) to avoid multiplier <= 0 and numerical instabilities + // Clamp weight to (-0.999, 0.999) to avoid multiplier <= 0 and numerical instabilities // which would invert or zero out the score. let weight = f64::from(*weight).clamp(-0.999, 0.999); multiplier * (1.0 - weight) From 65596fe13db7434c7e7f57d6dc870d0592970b73 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 12:07:20 -0700 Subject: [PATCH 29/36] add improved heuristics support around preferred routing targets Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/src/dynamo/_core.pyi | 6 ++++-- lib/kv-router/src/protocols.rs | 14 +++++++------- lib/llm/src/protocols/openai/nvext.rs | 16 ++++------------ 3 files changed, 15 insertions(+), 21 deletions(-) diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 68238df5938d..83bfb24caebf 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -552,9 +552,11 @@ class RoutingConstraints: Request-side routing constraints. ``required_taints`` is a hard eligibility filter. - ``preferred_taints`` maps taint -> weight in ``(-1.0, 1.0)``. + ``preferred_taints`` maps taint -> signed weight. Positive weights prefer matching workers, negative weights avoid them, - and ``0.0`` is neutral. + 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] diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 0fcc12625bd2..7c1d3d67edcc 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -202,15 +202,15 @@ impl RoutingConstraints { } pub fn preferred_taint_multiplier(&self, worker_taints: &HashSet) -> f64 { - self.preferred_taints + let bias = self + .preferred_taints .iter() .filter(|(taint, _)| worker_taints.contains(*taint)) - .fold(1.0, |multiplier, (_, weight)| { - // Clamp weight to (-0.999, 0.999) to avoid multiplier <= 0 and numerical instabilities - // which would invert or zero out the score. - let weight = f64::from(*weight).clamp(-0.999, 0.999); - multiplier * (1.0 - weight) - }) + .map(|(_, weight)| f64::from(*weight)) + .sum::() + .tanh(); + + (-bias).exp() } } diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 3865408f5684..4aaffebeb2b5 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -299,6 +299,9 @@ pub struct RoutingConstraintsSchema { /// 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, } @@ -501,18 +504,7 @@ impl NvExt { } } -fn validate_nv_ext(nv_ext: &NvExt) -> Result<(), ValidationError> { - if let Some(routing_constraints) = nv_ext.routing_constraints.as_ref() - && routing_constraints - .preferred_taints - .values() - .any(|weight| !(-1.0..1.0).contains(weight)) - { - let mut error = ValidationError::new("preferred_taint_weight_out_of_range"); - error.message = Some("preferred taint weights must be in the range (-1.0, 1.0)".into()); - return Err(error); - } - +fn validate_nv_ext(_nv_ext: &NvExt) -> Result<(), ValidationError> { Ok(()) } From ce236a377434147bf0e9d9a6e5733214c49f44f4 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 13:19:49 -0700 Subject: [PATCH 30/36] add eligibility helper to bundle checks with usually inline Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 12 +++ lib/kv-router/src/scheduling/queue.rs | 38 +++----- lib/kv-router/src/scheduling/selector.rs | 47 +++------- lib/kv-router/src/scheduling/types.rs | 113 +++++++++++++++++------ 4 files changed, 127 insertions(+), 83 deletions(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 7c1d3d67edcc..1222b7adf17c 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -189,12 +189,20 @@ impl RoutingConstraints { } 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)) @@ -202,6 +210,10 @@ impl RoutingConstraints { } pub fn preferred_taint_multiplier(&self, worker_taints: &HashSet) -> f64 { + if self.preferred_taints.is_empty() { + return 1.0; + } + let bias = self .preferred_taints .iter() diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 02b824bb6324..7733e2c3bdf1 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,10 +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 crate::protocols::{ - PrefillLoadHint, RoutingConstraints, WorkerConfigLike, WorkerId, WorkerWithDpRank, +use super::types::{ + RoutingEligibility, SchedulingRequest, SchedulingResponse, pinned_worker_config, }; +use crate::protocols::{PrefillLoadHint, WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::{ActiveSequencesMultiWorker, SequencePublisher, SequenceRequest}; /// Large default for max_num_batched_tokens when not configured (effectively disables queueing for that worker) @@ -145,7 +145,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.pinned_worker) { request.respond(Err(error)); return; } @@ -158,16 +160,15 @@ impl< return; }; - if request.bypass_capacity_check() { + if eligibility.bypasses_capacity_check(request.pinned_worker) { self.admit_one(request, decay_now).await; return; } if self.all_workers_busy( threshold, - request.allowed_worker_ids.as_ref(), + request.eligibility(), request.pinned_worker, - &request.routing_constraints, decay_now, ) { tracing::debug!("all workers busy, queueing request"); @@ -218,9 +219,8 @@ impl< // 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.eligibility(), front.request.pinned_worker, - &front.request.routing_constraints, decay_now, ) { break; @@ -354,9 +354,8 @@ impl< fn all_workers_busy( &self, threshold: f64, - allowed: Option<&HashSet>, + eligibility: RoutingEligibility<'_>, pinned_worker: Option, - routing_constraints: &RoutingConstraints, decay_now: Instant, ) -> bool { let active_tokens = self.slots.active_tokens(decay_now); @@ -366,9 +365,7 @@ impl< let Ok(config) = pinned_worker_config::(&*configs, worker) else { return false; }; - if routing_constraints.has_hard_constraints() - && !routing_constraints.is_compatible_with_worker_taints(config.taints()) - { + if !eligibility.allows_worker(worker.worker_id, config) { return false; } @@ -381,14 +378,7 @@ impl< let mut checked_any = false; for (&worker_id, config) in configs.iter() { - if let Some(ids) = allowed - && !ids.contains(&worker_id) - { - continue; - } - if routing_constraints.has_hard_constraints() - && !routing_constraints.is_compatible_with_worker_taints(config.taints()) - { + if !eligibility.allows_worker(worker_id, config) { continue; } let dp_size = config.data_parallel_size(); @@ -412,7 +402,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; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index a93ad5952941..94bf75eba49a 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -203,29 +203,19 @@ 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(request.pinned_worker)?; let pinned_worker = request.pinned_worker; - if pinned_worker.is_none() { - let allowed_workers: Vec<_> = workers - .iter() - .filter(|(worker_id, _)| request.is_worker_allowed(**worker_id)) - .collect(); - - if allowed_workers.is_empty() { - return Err(KvSchedulerError::NoEndpoints); - } - - if request.routing_constraints.has_hard_constraints() - && !allowed_workers.iter().any(|(_, config)| { - request - .routing_constraints - .is_compatible_with_worker_taints(config.taints()) - }) - { - return Err(KvSchedulerError::NoEndpoints); - } + if pinned_worker.is_none() + && !eligibility.has_eligible_worker( + workers + .iter() + .map(|(&worker_id, config)| (worker_id, config)), + ) + { + return Err(KvSchedulerError::NoEndpoints); } let request_blocks = request.request_blocks(block_size); @@ -250,11 +240,10 @@ 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| { - !request - .routing_constraints - .is_compatible_with_worker_taints(config.taints()) - }) { + if workers + .get(&worker.worker_id) + .is_some_and(|config| !eligibility.allows_worker(worker.worker_id, config)) + { return Err(KvSchedulerError::NoEndpoints); } @@ -297,13 +286,7 @@ impl WorkerSelector for DefaultWorkerSelector { let worker_iter = workers .iter() - .filter(move |(worker_id, _)| request.is_worker_allowed(**worker_id)) - .filter(move |(_, config)| { - request.routing_constraints.is_empty() - || request - .routing_constraints - .is_compatible_with_worker_taints(config.taints()) - }) + .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(); diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index 4d5d9857c8c3..8d5f620e4813 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -84,7 +84,91 @@ pub struct SchedulingRequest { pub resp_tx: Option>>, } +#[derive(Clone, Copy)] +pub(crate) struct RoutingEligibility<'a> { + allowed_worker_ids: Option<&'a HashSet>, + routing_constraints: &'a RoutingConstraints, +} + +impl<'a> RoutingEligibility<'a> { + #[inline] + pub(crate) fn new( + allowed_worker_ids: Option<&'a HashSet>, + routing_constraints: &'a RoutingConstraints, + ) -> Self { + Self { + allowed_worker_ids, + routing_constraints, + } + } + + #[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, + pinned_worker: Option, + ) -> Result<(), KvSchedulerError> { + let Some(pinned_worker) = 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, pinned_worker: Option) -> bool { + pinned_worker.is_none() && self.allowed_worker_ids.is_some() + } +} + impl SchedulingRequest { + #[inline] + pub(crate) fn eligibility(&self) -> RoutingEligibility<'_> { + RoutingEligibility::new(self.allowed_worker_ids.as_ref(), &self.routing_constraints) + } + pub(crate) fn prefill_token_deltas(&self) -> PrefillTokenDeltas { if !self.track_prefill_tokens { return PrefillTokenDeltas::none(); @@ -113,12 +197,13 @@ impl SchedulingRequest { } pub(crate) fn best_effective_prefill_tokens(&self) -> usize { + let eligibility = self.eligibility(); 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)) + .filter(|(worker, _)| eligibility.allows_worker_id(worker.worker_id)) .map(|(_, cached_tokens)| *cached_tokens) .max() .unwrap_or(0), @@ -141,12 +226,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 { @@ -178,26 +257,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"); From d45ead4e7000c2ac93a2d8555199147ce5dc6d0e Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 13:40:16 -0700 Subject: [PATCH 31/36] add re-export and test utils import fix Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/src/dynamo/llm/__init__.py | 1 + lib/kv-router/src/test_utils.rs | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/lib/bindings/python/src/dynamo/llm/__init__.py b/lib/bindings/python/src/dynamo/llm/__init__.py index 757e2a6ff012..ac79da07a515 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -30,6 +30,7 @@ from dynamo._core import PythonAsyncEngine as PythonAsyncEngine from dynamo._core import RadixTree as RadixTree from dynamo._core import ReasoningConfig as ReasoningConfig +from dynamo._core import RoutingConstraints as RoutingConstraints from dynamo._core import RouterConfig as RouterConfig from dynamo._core import RouterMode as RouterMode from dynamo._core import SglangArgs as SglangArgs diff --git a/lib/kv-router/src/test_utils.rs b/lib/kv-router/src/test_utils.rs index ce6817dd50c9..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; @@ -406,4 +407,3 @@ impl WorkerConfigLike for SimpleWorkerConfig { &self.taints } } -use std::collections::HashSet; From b471affc4d6905d74b3ceaf49148e52fd6ff8730 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 14:44:14 -0700 Subject: [PATCH 32/36] make api more clear to only multiply if useful Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 6 +++--- lib/kv-router/src/scheduling/selector.rs | 11 +++++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 1222b7adf17c..9ca6c4b57bee 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -209,9 +209,9 @@ impl RoutingConstraints { .count() } - pub fn preferred_taint_multiplier(&self, worker_taints: &HashSet) -> f64 { + pub fn preferred_taint_multiplier(&self, worker_taints: &HashSet) -> Option { if self.preferred_taints.is_empty() { - return 1.0; + return None; } let bias = self @@ -222,7 +222,7 @@ impl RoutingConstraints { .sum::() .tanh(); - (-bias).exp() + Some((-bias).exp()) } } diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 94bf75eba49a..462cf6de0dab 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -278,10 +278,13 @@ impl WorkerSelector for DefaultWorkerSelector { let Some(config) = workers.get(&worker.worker_id) else { return base_score; }; - base_score - * request - .routing_constraints - .preferred_taint_multiplier(config.taints()) + match request + .routing_constraints + .preferred_taint_multiplier(config.taints()) + { + Some(multiplier) => base_score * multiplier, + None => base_score, + } }; let worker_iter = workers From 1922ae49984dbe8aa1a7fa1b139ee5b6a52ef712 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 14:45:03 -0700 Subject: [PATCH 33/36] make api more clear to only multiply if useful Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/bindings/python/src/dynamo/llm/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/bindings/python/src/dynamo/llm/__init__.py b/lib/bindings/python/src/dynamo/llm/__init__.py index e34899d7a5e1..0338e1c68288 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -32,9 +32,9 @@ from dynamo._core import PythonAsyncEngine as PythonAsyncEngine from dynamo._core import RadixTree as RadixTree from dynamo._core import ReasoningConfig as ReasoningConfig -from dynamo._core import RoutingConstraints as RoutingConstraints 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 From f802890b97576928c8b67d30adb35c1c57b2a31c Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 15:22:16 -0700 Subject: [PATCH 34/36] ove pinned worker in the api, scheduled context helper to make sure queueing computes the right thing Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/scheduling/policy.rs | 149 +++++++++++++----- lib/kv-router/src/scheduling/queue.rs | 37 +++-- lib/kv-router/src/scheduling/selector.rs | 4 +- lib/kv-router/src/scheduling/types.rs | 82 +++++++--- .../src/replay/offline/components/router.rs | 12 +- 5 files changed, 198 insertions(+), 86 deletions(-) diff --git a/lib/kv-router/src/scheduling/policy.rs b/lib/kv-router/src/scheduling/policy.rs index f1ac5e4cf518..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, @@ -171,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"); } @@ -183,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" @@ -199,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); } @@ -208,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"); } @@ -219,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" @@ -234,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 ---- @@ -249,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" ); } @@ -261,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)" @@ -275,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); } @@ -287,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" ); } @@ -302,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); } @@ -313,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); } @@ -324,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); } @@ -333,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); } @@ -343,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 7733e2c3bdf1..7ef6e5134c72 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -14,7 +14,8 @@ use super::policy::{FcfsPolicy, SchedulingPolicy}; use super::prefill_load::PrefillLoadEstimator; use super::selector::{DefaultWorkerSelector, WorkerSelector}; use super::types::{ - RoutingEligibility, SchedulingRequest, SchedulingResponse, pinned_worker_config, + RoutingEligibility, SchedulingContext, SchedulingRequest, SchedulingResponse, + pinned_worker_config, }; use crate::protocols::{PrefillLoadHint, WorkerConfigLike, WorkerId, WorkerWithDpRank}; use crate::sequences::{ActiveSequencesMultiWorker, SequencePublisher, SequenceRequest}; @@ -147,7 +148,7 @@ impl< pub async fn enqueue(&self, mut request: SchedulingRequest) { let eligibility = request.eligibility(); - if let Err(error) = eligibility.validate_pinned_worker(request.pinned_worker) { + if let Err(error) = eligibility.validate_pinned_worker() { request.respond(Err(error)); return; } @@ -160,20 +161,19 @@ impl< return; }; - if eligibility.bypasses_capacity_check(request.pinned_worker) { + if eligibility.bypasses_capacity_check() { self.admit_one(request, decay_now).await; return; } - if self.all_workers_busy( - threshold, - request.eligibility(), - 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); @@ -195,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(); @@ -217,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.eligibility(), - 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"); @@ -355,13 +355,12 @@ impl< &self, threshold: f64, eligibility: RoutingEligibility<'_>, - pinned_worker: Option, 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; }; diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 462cf6de0dab..e35c89b38835 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -204,9 +204,9 @@ impl WorkerSelector for DefaultWorkerSelector { ) -> Result { assert!(request.isl_tokens > 0); let eligibility = request.eligibility(); - eligibility.validate_pinned_worker(request.pinned_worker)?; + eligibility.validate_pinned_worker()?; - let pinned_worker = request.pinned_worker; + let pinned_worker = eligibility.pinned_worker(); if pinned_worker.is_none() && !eligibility.has_eligible_worker( diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index 8d5f620e4813..b1a08c87e953 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -87,6 +87,7 @@ pub struct SchedulingRequest { #[derive(Clone, Copy)] pub(crate) struct RoutingEligibility<'a> { allowed_worker_ids: Option<&'a HashSet>, + pinned_worker: Option, routing_constraints: &'a RoutingConstraints, } @@ -94,14 +95,21 @@ 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 @@ -140,11 +148,8 @@ impl<'a> RoutingEligibility<'a> { } #[inline] - pub(crate) fn validate_pinned_worker( - &self, - pinned_worker: Option, - ) -> Result<(), KvSchedulerError> { - let Some(pinned_worker) = pinned_worker else { + pub(crate) fn validate_pinned_worker(&self) -> Result<(), KvSchedulerError> { + let Some(pinned_worker) = self.pinned_worker else { return Ok(()); }; @@ -158,15 +163,60 @@ impl<'a> RoutingEligibility<'a> { } #[inline] - pub(crate) fn bypasses_capacity_check(&self, pinned_worker: Option) -> bool { - pinned_worker.is_none() && self.allowed_worker_ids.is_some() + 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.routing_constraints) + RoutingEligibility::new( + self.allowed_worker_ids.as_ref(), + self.pinned_worker, + &self.routing_constraints, + ) } pub(crate) fn prefill_token_deltas(&self) -> PrefillTokenDeltas { @@ -196,22 +246,6 @@ impl SchedulingRequest { PrefillTokenDeltas::new(self.isl_tokens, by_worker) } - pub(crate) fn best_effective_prefill_tokens(&self) -> usize { - let eligibility = self.eligibility(); - let cached_tokens = match self.pinned_worker { - Some(worker) => self.effective_cached_tokens_for(worker), - None => self - .effective_cached_tokens - .iter() - .filter(|(worker, _)| eligibility.allows_worker_id(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) diff --git a/lib/mocker/src/replay/offline/components/router.rs b/lib/mocker/src/replay/offline/components/router.rs index e2c0ed56a3ce..1c2ae9a87fa0 100644 --- a/lib/mocker/src/replay/offline/components/router.rs +++ b/lib/mocker/src/replay/offline/components/router.rs @@ -14,6 +14,7 @@ use dynamo_kv_router::protocols::{ 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, @@ -442,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), ) } From 64ccf6d7e5f6061b4807b71196065350a8de58bb Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 16:15:21 -0700 Subject: [PATCH 35/36] add comment why heuristic is choosen Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 9ca6c4b57bee..543a7cd88061 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -214,6 +214,10 @@ impl RoutingConstraints { return None; } + // Use exp(-tanh(sum)) so equal-magnitude positive and negative preferences + // have symmetric effect around the neutral multiplier 1.0, while keeping the + // multiplier strictly positive and bounded for numerically stable composition + // with the existing linear work score. let bias = self .preferred_taints .iter() From 6aca329d19996c45441012db4dd70f94c3f74319 Mon Sep 17 00:00:00 2001 From: michaelfeil <63565275+michaelfeil@users.noreply.github.com> Date: Fri, 15 May 2026 16:18:20 -0700 Subject: [PATCH 36/36] add comment why heuristic is choosen Signed-off-by: michaelfeil <63565275+michaelfeil@users.noreply.github.com> --- lib/kv-router/src/protocols.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 543a7cd88061..400ffc9a4fcf 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -215,9 +215,9 @@ impl RoutingConstraints { } // Use exp(-tanh(sum)) so equal-magnitude positive and negative preferences - // have symmetric effect around the neutral multiplier 1.0, while keeping the - // multiplier strictly positive and bounded for numerically stable composition - // with the existing linear work score. + // 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()