use std::collections::HashSet;
use dynamo_kv_router::{
RouterConfigOverride,
indexer::RoutingDecisionHashes,
protocols::{
BlockExtraInfo, RoutingConstraints, WorkerAffinityTarget, WorkerId, WorkerWithDpRank,
},
router_hint::RouterHint,
scheduling::{AdmissionAttempt, AdvisoryWorkerLoad, QueueRejection, RoutingEligibility},
selector::WorkerSelector,
};
use dynamo_runtime::{dynamo_nvtx_range, pipeline::Error};
use crate::{
kv_router::{
FindBestMatchAdmission, FindBestMatchInnerOutcome, FindBestMatchOutcome,
routing_host::RoutingHost,
},
local_model::runtime_config::ModelRuntimeConfig,
preprocessor::PreprocessedRequest,
protocols::{
TokenIdType,
common::{preprocessor::RoutingHints, timing::RequestPhase},
},
session_affinity::AffinityTarget,
};
pub(super) struct WorkerSelection {
pub(super) worker: WorkerWithDpRank,
pub(super) attempt: AdmissionAttempt,
pub(super) overlap_amount: u32,
pub(super) effective_overlap_blocks: f64,
pub(super) cached_tokens: usize,
pub(super) potential_decode_blocks: u64,
pub(super) selected_worker_load: Option<AdvisoryWorkerLoad>,
pub(super) routing_hashes: Option<RoutingDecisionHashes>,
pub(super) router_hint: Option<RouterHint>,
}
pub(super) enum SelectionOutcome {
Routed(WorkerSelection),
QueueRejected(QueueRejection),
}
impl SelectionOutcome {
pub(super) fn into_result(self) -> Result<WorkerSelection, Error> {
match self {
Self::Routed(selection) => Ok(selection),
Self::QueueRejected(rejection) => Err(rejection.into()),
}
}
}
#[derive(Clone, Copy)]
pub(super) struct RoutingRequestParts<'a> {
pub(super) token_ids: &'a [TokenIdType],
pub(super) block_mm_infos: Option<&'a [Option<BlockExtraInfo>]>,
}
impl<'a> RoutingRequestParts<'a> {
pub(super) fn new(request: &'a PreprocessedRequest) -> Self {
let (token_ids, block_mm_infos) = request.block_mm_routing_info();
Self {
token_ids,
block_mm_infos,
}
}
}
pub(super) struct SelectionOptions {
pub(super) pinned_target: Option<AffinityTarget>,
pub(super) affinity_target: Option<AffinityTarget>,
pub(super) planned_worker: Option<WorkerWithDpRank>,
pub(super) policy_class: Option<String>,
pub(super) session_context: Option<dynamo_kv_router::SessionContext>,
pub(super) admission: FindBestMatchAdmission,
}
struct BestMatchArgs<'a> {
context_id: &'a str,
routing_parts: RoutingRequestParts<'a>,
router_config_override: Option<&'a RouterConfigOverride>,
update_states: bool,
return_routing_hashes: bool,
lora_name: Option<String>,
cache_namespace: Option<String>,
priority_jump: f64,
strict_priority: u32,
policy_class: Option<String>,
session_context: Option<dynamo_kv_router::SessionContext>,
expected_output_tokens: Option<u32>,
affinity_target: Option<WorkerAffinityTarget>,
pinned_worker: Option<WorkerWithDpRank>,
allowed_worker_ids: Option<HashSet<WorkerId>>,
routing_constraints: RoutingConstraints,
admission: FindBestMatchAdmission,
}
impl<Sel> RoutingHost<Sel>
where
Sel: WorkerSelector<ModelRuntimeConfig> + Send + 'static,
{
async fn select_best_match(&self, args: BestMatchArgs<'_>) -> Result<SelectionOutcome, Error> {
let outcome = self
.kv_router()
.find_best_match_details_with_policy_class_inner(
Some(args.context_id),
args.routing_parts.token_ids,
args.routing_parts.block_mm_infos,
args.router_config_override,
args.update_states,
args.return_routing_hashes,
args.lora_name,
args.cache_namespace,
args.priority_jump,
args.strict_priority,
args.policy_class,
args.session_context,
args.expected_output_tokens,
args.affinity_target,
args.pinned_worker,
args.allowed_worker_ids,
args.routing_constraints,
args.admission,
)
.await?;
match outcome {
FindBestMatchInnerOutcome::WithAdmission(admitted) => match admitted.outcome {
FindBestMatchOutcome::Routed {
worker,
overlap_blocks,
effective_overlap_blocks,
cached_tokens,
potential_decode_blocks,
routing_hashes,
router_hint,
} => Ok(SelectionOutcome::Routed(WorkerSelection {
worker,
attempt: admitted.attempt,
overlap_amount: overlap_blocks,
effective_overlap_blocks,
cached_tokens,
potential_decode_blocks,
selected_worker_load: None,
routing_hashes,
router_hint,
})),
FindBestMatchOutcome::QueueRejected { rejection } => {
Ok(SelectionOutcome::QueueRejected(rejection))
}
},
FindBestMatchInnerOutcome::WithoutAdmission(outcome) => match outcome {
crate::kv_router::FindBestMatchAdvisoryOutcome::Routed {
worker,
overlap_blocks,
effective_overlap_blocks,
cached_tokens,
potential_decode_blocks,
selected_worker_load,
routing_hashes,
} => Ok(SelectionOutcome::Routed(WorkerSelection {
worker,
attempt: AdmissionAttempt::Untracked,
overlap_amount: overlap_blocks,
effective_overlap_blocks,
cached_tokens,
potential_decode_blocks,
selected_worker_load: Some(selected_worker_load),
routing_hashes,
router_hint: None,
})),
crate::kv_router::FindBestMatchAdvisoryOutcome::QueueRejected { rejection } => {
Ok(SelectionOutcome::QueueRejected(rejection))
}
},
}
}
pub(super) async fn select_worker_outcome(
&self,
context_id: &str,
request: &PreprocessedRequest,
routing_parts: RoutingRequestParts<'_>,
phase: RequestPhase,
is_query_only: bool,
options: SelectionOptions,
) -> Result<SelectionOutcome, Error> {
let _nvtx_select = dynamo_nvtx_range!("route.select_worker");
let routing = request.routing.as_ref();
let explicit_pin = pinned_worker_hint(phase, routing);
let lora_name = routing.and_then(|routing| routing.lora_name.clone());
let cache_namespace = routing.and_then(|routing| routing.cache_namespace.clone());
let priority_jump = routing
.and_then(|routing| routing.priority_jump)
.unwrap_or(0.0);
let strict_priority = routing
.and_then(|routing| routing.strict_priority)
.unwrap_or(0);
let expected_output_tokens = routing.and_then(|routing| routing.expected_output_tokens);
let routing_constraints = routing
.and_then(|routing| routing.routing_constraints.clone())
.unwrap_or_default();
let mut allowed_worker_ids = routing.and_then(|routing| routing.allowed_worker_ids.clone());
let migration_excluded_worker_ids = request
.migration_state
.as_ref()
.map(|state| state.excluded_worker_ids())
.unwrap_or_default();
if explicit_pin.is_none() && !migration_excluded_worker_ids.is_empty() {
let workers = self.kv_router().workers_with_configs.borrow();
let eligible =
allowed_worker_ids.get_or_insert_with(|| workers.keys().copied().collect());
eligible.retain(|worker_id| {
workers.get(worker_id).is_some_and(|config| {
routing_constraints.is_compatible_with_worker_taints(&config.taints)
}) && !migration_excluded_worker_ids.contains(worker_id)
});
if eligible.is_empty()
&& let Some(error) = request
.migration_state
.as_ref()
.and_then(|state| state.exhausted_error())
{
return Err(anyhow::anyhow!(error));
}
}
let return_routing_hashes =
!is_query_only && self.kv_router().indexer().records_routing_decisions();
let SelectionOptions {
pinned_target,
affinity_target,
planned_worker,
policy_class,
session_context,
admission,
} = options;
let worker_only_affinity = pinned_target.filter(|target| target.dp_rank.is_none());
if let Some(target) = worker_only_affinity {
match &mut allowed_worker_ids {
Some(allowed_workers) => {
allowed_workers.retain(|worker_id| *worker_id == target.worker_id);
}
None => {
allowed_worker_ids = Some(HashSet::from([target.worker_id]));
}
}
}
let explicit_pin = match (explicit_pin, worker_only_affinity) {
(Some((worker_id, None)), Some(affinity_target))
if worker_id == affinity_target.worker_id =>
{
None
}
(explicit_pin, _) => explicit_pin,
};
let affinity_pin = pinned_target.and_then(|target| {
target
.dp_rank
.map(|dp_rank| (target.worker_id, Some(dp_rank)))
});
let requested_pin = merge_affinity_pin(explicit_pin, affinity_pin);
let pinned_worker = match planned_worker {
Some(planned_worker) => {
if let Some((worker_id, dp_rank)) = requested_pin
&& (worker_id != planned_worker.worker_id
|| dp_rank.is_some_and(|dp_rank| dp_rank != planned_worker.dp_rank))
{
return Err(anyhow::anyhow!(
"Previewed worker {} dp_rank {} conflicts with requested worker {} dp_rank {:?}",
planned_worker.worker_id,
planned_worker.dp_rank,
worker_id,
dp_rank,
));
}
Some(planned_worker)
}
None => match requested_pin {
Some((worker_id, requested_dp_rank)) => Some(resolve_pinned_worker_rank(
worker_id,
requested_dp_rank,
self.kv_router().unique_dp_rank_for_worker(worker_id),
)?),
None => None,
},
};
let Some(pinned_worker) = pinned_worker else {
let _nvtx_kv = dynamo_nvtx_range!("route.kv_match");
let selection = self
.select_best_match(BestMatchArgs {
context_id,
routing_parts,
router_config_override: request.router_config_override.as_ref(),
update_states: !is_query_only,
return_routing_hashes,
lora_name,
cache_namespace,
priority_jump,
strict_priority,
policy_class,
session_context,
expected_output_tokens,
affinity_target: affinity_target
.map(|target| WorkerAffinityTarget::new(target.worker_id, target.dp_rank)),
pinned_worker: None,
allowed_worker_ids,
routing_constraints: routing_constraints.clone(),
admission,
})
.await?;
if !is_query_only && let SelectionOutcome::Routed(selection) = &selection {
let total_blocks = routing_parts
.token_ids
.len()
.div_ceil(self.kv_router().block_size() as usize);
tracing::debug!(
request_id = %context_id,
worker_id = selection.worker.worker_id,
dp_rank = selection.worker.dp_rank,
overlap_blocks = selection.overlap_amount,
total_blocks,
"[ROUTING] Best: worker_{} dp_rank={} with {}/{} blocks overlap",
selection.worker.worker_id,
selection.worker.dp_rank,
selection.overlap_amount,
total_blocks,
);
}
return Ok(selection);
};
{
let configs = self.kv_router().workers_with_configs.borrow();
let eligibility = RoutingEligibility::new(
allowed_worker_ids.as_ref(),
None,
Some(pinned_worker),
&routing_constraints,
);
if let Err(error) = eligibility.validate_worker_rank(&configs, pinned_worker) {
return Err(anyhow::anyhow!(
"Pinned worker {} dp_rank {} is not eligible: {error}",
pinned_worker.worker_id,
pinned_worker.dp_rank
));
}
}
tracing::debug!(
worker_id = pinned_worker.worker_id,
dp_rank = pinned_worker.dp_rank,
?phase,
"Routing to specified worker"
);
self.select_best_match(BestMatchArgs {
context_id,
routing_parts,
router_config_override: request.router_config_override.as_ref(),
update_states: !is_query_only,
return_routing_hashes,
lora_name,
cache_namespace,
priority_jump,
strict_priority,
policy_class,
session_context,
expected_output_tokens,
affinity_target: None,
pinned_worker: Some(pinned_worker),
allowed_worker_ids,
routing_constraints,
admission,
})
.await
}
}
fn merge_affinity_pin(
explicit: Option<(u64, Option<u32>)>,
affinity: Option<(u64, Option<u32>)>,
) -> Option<(u64, Option<u32>)> {
match (explicit, affinity) {
(Some((worker_id, None)), Some((affinity_worker_id, affinity_rank)))
if worker_id == affinity_worker_id =>
{
Some((worker_id, affinity_rank))
}
(Some(explicit), _) => Some(explicit),
(None, affinity) => affinity,
}
}
fn resolve_pinned_worker_rank(
worker_id: WorkerId,
requested_dp_rank: Option<u32>,
unique_dp_rank: Option<u32>,
) -> Result<WorkerWithDpRank, Error> {
let Some(dp_rank) = requested_dp_rank.or(unique_dp_rank) else {
return Err(anyhow::anyhow!(
"Pinned worker {worker_id} requires an explicit dp_rank because it has multiple or unknown DP ranks"
));
};
Ok(WorkerWithDpRank::new(worker_id, dp_rank))
}
fn pinned_worker_hint(
phase: RequestPhase,
routing: Option<&RoutingHints>,
) -> Option<(u64, Option<u32>)> {
let routing = routing?;
match phase {
RequestPhase::Prefill => {
let worker_id = routing.prefill_worker_id.or(routing.backend_instance_id)?;
let dp_rank = routing.prefill_dp_rank.or(routing.dp_rank);
Some((worker_id, dp_rank))
}
RequestPhase::Decode => {
let worker_id = routing.decode_worker_id.or(routing.backend_instance_id)?;
Some((worker_id, routing.dp_rank))
}
RequestPhase::Aggregated => {
let worker_id = routing.decode_worker_id.or(routing.backend_instance_id)?;
Some((worker_id, routing.dp_rank))
}
}
}
#[cfg(test)]
mod tests {
use std::collections::{HashMap, HashSet};
use dynamo_kv_router::{
protocols::{RoutingConstraints, WorkerWithDpRank},
scheduling::{RoutingEligibility, WorkerEligibilityError},
};
use super::{merge_affinity_pin, pinned_worker_hint, resolve_pinned_worker_rank};
use crate::{
local_model::runtime_config::ModelRuntimeConfig,
protocols::common::{preprocessor::RoutingHints, timing::RequestPhase},
};
#[test]
fn resolve_pinned_worker_rank_uses_explicit_rank_including_zero() {
let worker = resolve_pinned_worker_rank(7, Some(0), Some(3)).unwrap();
assert_eq!(worker.worker_id, 7);
assert_eq!(worker.dp_rank, 0);
}
#[test]
fn resolve_pinned_worker_rank_uses_unique_rank_when_unset() {
let worker = resolve_pinned_worker_rank(7, None, Some(3)).unwrap();
assert_eq!(worker.worker_id, 7);
assert_eq!(worker.dp_rank, 3);
}
#[test]
fn resolve_pinned_worker_rank_rejects_unresolved_rank() {
let error = resolve_pinned_worker_rank(7, None, None)
.unwrap_err()
.to_string();
assert!(error.contains("requires an explicit dp_rank"));
}
#[test]
fn affinity_pin_supplies_rank_for_matching_explicit_worker() {
assert_eq!(
merge_affinity_pin(Some((7, None)), Some((7, Some(0)))),
Some((7, Some(0)))
);
assert_eq!(
merge_affinity_pin(Some((7, Some(2))), Some((7, Some(3)))),
Some((7, Some(2)))
);
}
#[test]
fn pinned_worker_hint_prefill_uses_prefill_worker_before_backend() {
let routing = RoutingHints {
backend_instance_id: Some(1),
prefill_worker_id: Some(2),
dp_rank: Some(3),
prefill_dp_rank: Some(4),
..Default::default()
};
assert_eq!(
pinned_worker_hint(RequestPhase::Prefill, Some(&routing)),
Some((2, Some(4)))
);
}
#[test]
fn pinned_worker_hint_decode_uses_decode_worker_before_backend() {
let routing = RoutingHints {
backend_instance_id: Some(1),
decode_worker_id: Some(5),
dp_rank: Some(6),
..Default::default()
};
assert_eq!(
pinned_worker_hint(RequestPhase::Decode, Some(&routing)),
Some((5, Some(6)))
);
}
#[test]
fn pinned_worker_hint_aggregated_uses_decode_worker_before_backend() {
let routing = RoutingHints {
backend_instance_id: Some(9),
decode_worker_id: Some(5),
dp_rank: Some(7),
..Default::default()
};
assert_eq!(
pinned_worker_hint(RequestPhase::Aggregated, Some(&routing)),
Some((5, Some(7)))
);
}
#[test]
fn affinity_validation_ignores_transient_overload() {
let worker = WorkerWithDpRank::new(7, 0);
let configs = HashMap::from([(7, ModelRuntimeConfig::default())]);
let constraints = RoutingConstraints::default();
let overloaded = HashSet::from([7]);
let scheduling_eligibility =
RoutingEligibility::new(None, Some(&overloaded), Some(worker), &constraints);
let affinity_eligibility = RoutingEligibility::new(None, None, Some(worker), &constraints);
assert_eq!(
scheduling_eligibility
.validate_worker_rank(&configs, worker)
.err(),
Some(WorkerEligibilityError::WorkerOverloaded { worker_id: 7 })
);
assert!(
affinity_eligibility
.validate_worker_rank(&configs, worker)
.is_ok()
);
}
}