From 2a491dfe10f23735ef832a526e6e39c2c421043e Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 11 Jun 2026 19:42:35 -0700 Subject: [PATCH 1/4] feat(kv-router): add runtime-free selection service Signed-off-by: PeaBrane --- lib/bindings/python/Cargo.toml | 1 + lib/bindings/python/rust/lib.rs | 9 + lib/bindings/python/rust/llm/kv.rs | 65 ++- .../python/src/dynamo/llm/__init__.py | 1 + .../src/dynamo/select_service/__init__.py | 2 + .../src/dynamo/select_service/__main__.py | 7 + .../python/src/dynamo/select_service/main.py | 19 + lib/kv-router/Cargo.toml | 1 + lib/kv-router/src/services/mod.rs | 9 +- .../src/services/selection/README.md | 99 ++++ .../src/services/selection/catalog.rs | 144 +++++ lib/kv-router/src/services/selection/error.rs | 77 +++ lib/kv-router/src/services/selection/input.rs | 137 +++++ lib/kv-router/src/services/selection/mod.rs | 27 + .../src/services/selection/scoring.rs | 165 ++++++ .../src/services/selection/server.rs | 278 +++++++++ lib/kv-router/src/services/selection/tests.rs | 325 +++++++++++ lib/kv-router/src/services/selection/types.rs | 532 ++++++++++++++++++ tests/router/common.py | 44 ++ tests/router/helper.py | 57 ++ tests/router/mocker_process.py | 90 ++- tests/router/test_router_e2e_with_mockers.py | 2 + 22 files changed, 2082 insertions(+), 9 deletions(-) create mode 100644 lib/bindings/python/src/dynamo/select_service/__init__.py create mode 100644 lib/bindings/python/src/dynamo/select_service/__main__.py create mode 100644 lib/bindings/python/src/dynamo/select_service/main.py create mode 100644 lib/kv-router/src/services/selection/README.md create mode 100644 lib/kv-router/src/services/selection/catalog.rs create mode 100644 lib/kv-router/src/services/selection/error.rs create mode 100644 lib/kv-router/src/services/selection/input.rs create mode 100644 lib/kv-router/src/services/selection/mod.rs create mode 100644 lib/kv-router/src/services/selection/scoring.rs create mode 100644 lib/kv-router/src/services/selection/server.rs create mode 100644 lib/kv-router/src/services/selection/tests.rs create mode 100644 lib/kv-router/src/services/selection/types.rs diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index d8ea4a874ab1..6a76e8ee536d 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -27,6 +27,7 @@ media-ffmpeg = ["dynamo-llm/media-ffmpeg"] kv-indexer = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-indexer"] kv-indexer-metrics = ["kv-indexer", "dynamo-kv-router/metrics"] slot-tracker = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-slot-tracker"] +select-service = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-selection"] mocker-kvbm-offload = ["dynamo-mocker/kvbm-offload"] aic-forward-pass = ["dynamo-mocker/aic-forward-pass", "dep:aiconfigurator-core"] nvtx = ["dynamo-runtime/nvtx"] diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index e90b40a2757c..fb14450ecb53 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -168,6 +168,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(fetch_model, m)?)?; m.add_function(wrap_pyfunction!(run_kv_indexer, m)?)?; m.add_function(wrap_pyfunction!(run_slot_tracker, m)?)?; + m.add_function(wrap_pyfunction!(run_select_service, m)?)?; m.add_function(wrap_pyfunction!(llm::entrypoint::make_engine, m)?)?; m.add_function(wrap_pyfunction!(llm::replay::run_mocker_trace_replay, m)?)?; m.add_function(wrap_pyfunction!( @@ -290,6 +291,14 @@ fn run_slot_tracker(py: Python<'_>, argv: Option>) -> PyResult<()> { .map_err(standalone_to_pyerr) } +#[pyfunction(name = "run_select_service")] +#[pyo3(signature = (argv=None))] +fn run_select_service(py: Python<'_>, argv: Option>) -> PyResult<()> { + let argv = argv.unwrap_or_default(); + py.allow_threads(move || llm::kv::run_select_service_cli(argv)) + .map_err(standalone_to_pyerr) +} + /// Log a message from Python with file and line info #[pyfunction] #[pyo3(text_signature = "(level, message, module, file, line)")] diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 0c262405cc36..7779db583a33 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -12,13 +12,21 @@ use tokio_stream::StreamExt; use super::local_model::RoutingConstraints; use super::*; use crate::Endpoint; -#[cfg(any(feature = "kv-indexer", feature = "slot-tracker"))] +#[cfg(any( + feature = "kv-indexer", + feature = "slot-tracker", + feature = "select-service" +))] use clap::Parser; -use dynamo_kv_router::config::{KvRouterConfig, RouterConfigOverride}; +use dynamo_kv_router::config::{ + KvRouterConfig, RouterConfigOverride, kv_router_config_from_dynamo_env, +}; use dynamo_kv_router::protocols::compute_block_hash_for_seq; use dynamo_kv_router::protocols::*; #[cfg(feature = "kv-indexer")] use dynamo_kv_router::services::indexer::{self, IndexerConfig}; +#[cfg(feature = "select-service")] +use dynamo_kv_router::services::selection::{self, SelectionServiceConfig}; #[cfg(feature = "slot-tracker")] use dynamo_kv_router::services::slot_tracker::{self, SlotTrackerConfig}; use rs::pipeline::{AsyncEngine, SingleIn}; @@ -191,7 +199,58 @@ where } } -#[cfg(any(feature = "kv-indexer", feature = "slot-tracker"))] +#[cfg(feature = "select-service")] +#[derive(Parser)] +#[command( + name = "python -m dynamo.select_service", + about = "Runtime-free Dynamo worker selection service" +)] +struct SelectServiceCli { + /// HTTP server port + #[arg(long, default_value_t = 8092)] + port: u16, + + /// Number of KV indexer worker threads + #[arg(long, default_value_t = 4)] + threads: usize, +} + +pub fn run_select_service_cli(args: I) -> anyhow::Result<()> +where + I: IntoIterator, + T: Into, +{ + #[cfg(feature = "select-service")] + { + let cli = SelectServiceCli::try_parse_from( + std::iter::once(OsString::from("python -m dynamo.select_service")) + .chain(args.into_iter().map(Into::into)), + )?; + + init_standalone_logging(); + + let rt = tokio::runtime::Runtime::new()?; + rt.block_on(selection::run_server(SelectionServiceConfig { + port: cli.port, + threads: cli.threads, + kv_router_config: kv_router_config_from_dynamo_env(), + })) + } + + #[cfg(not(feature = "select-service"))] + { + let _ = args; + anyhow::bail!( + "dynamo.select_service is not available in this build; reinstall with --features select-service" + ) + } +} + +#[cfg(any( + feature = "kv-indexer", + feature = "slot-tracker", + feature = "select-service" +))] fn init_standalone_logging() { let _ = tracing_subscriber::fmt() .with_env_filter( diff --git a/lib/bindings/python/src/dynamo/llm/__init__.py b/lib/bindings/python/src/dynamo/llm/__init__.py index ebd096568819..3ff854436ee2 100644 --- a/lib/bindings/python/src/dynamo/llm/__init__.py +++ b/lib/bindings/python/src/dynamo/llm/__init__.py @@ -41,6 +41,7 @@ from dynamo._core import register_model as register_model from dynamo._core import run_input from dynamo._core import run_kv_indexer as run_kv_indexer +from dynamo._core import run_select_service as run_select_service from dynamo._core import run_slot_tracker as run_slot_tracker from dynamo._core import unregister_model as unregister_model diff --git a/lib/bindings/python/src/dynamo/select_service/__init__.py b/lib/bindings/python/src/dynamo/select_service/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/lib/bindings/python/src/dynamo/select_service/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/lib/bindings/python/src/dynamo/select_service/__main__.py b/lib/bindings/python/src/dynamo/select_service/__main__.py new file mode 100644 index 000000000000..09b36ca221bf --- /dev/null +++ b/lib/bindings/python/src/dynamo/select_service/__main__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dynamo.select_service.main import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/lib/bindings/python/src/dynamo/select_service/main.py b/lib/bindings/python/src/dynamo/select_service/main.py new file mode 100644 index 000000000000..bf05f21ca839 --- /dev/null +++ b/lib/bindings/python/src/dynamo/select_service/main.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sys +from collections.abc import Sequence + +from dynamo.llm import run_select_service + + +def main(argv: Sequence[str] | None = None) -> int: + args = list(sys.argv[1:] if argv is None else argv) + try: + run_select_service(args) + except Exception as exc: + if "-h" in args or "--help" in args: + print(exc) + return 0 + raise + return 0 diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index 72a94450c133..e335a2151cb1 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -21,6 +21,7 @@ profile = [] shard-metrics = [] standalone-indexer = ["dep:axum", "dep:serde_json", "dep:reqwest", "dep:zmq"] standalone-slot-tracker = ["dep:axum", "dep:serde_json", "dep:zmq"] +standalone-selection = ["standalone-indexer"] [dependencies] # repo diff --git a/lib/kv-router/src/services/mod.rs b/lib/kv-router/src/services/mod.rs index 6b02ba3e0524..66c223168a9f 100644 --- a/lib/kv-router/src/services/mod.rs +++ b/lib/kv-router/src/services/mod.rs @@ -3,7 +3,11 @@ //! Standalone services built from brokerless transport primitives. -#[cfg(any(feature = "standalone-indexer", feature = "standalone-slot-tracker"))] +#[cfg(any( + feature = "standalone-indexer", + feature = "standalone-slot-tracker", + feature = "standalone-selection" +))] pub(crate) mod zmq; #[cfg(feature = "standalone-indexer")] @@ -12,5 +16,8 @@ pub mod indexer; #[cfg(feature = "standalone-indexer")] pub mod shared_cache; +#[cfg(feature = "standalone-selection")] +pub mod selection; + #[cfg(feature = "standalone-slot-tracker")] pub mod slot_tracker; diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md new file mode 100644 index 000000000000..d68dcdb7249d --- /dev/null +++ b/lib/kv-router/src/services/selection/README.md @@ -0,0 +1,99 @@ +# Dynamo Selection Service + +`dynamo-select-service` is a runtime-free HTTP service built on `dynamo-kv-router`. +It selects workers and tracks reservations, but it never forwards model requests +and never owns model responses. + +## Readiness + +- `GET /health` returns process liveness. +- `GET /ready` returns `200` only after at least one worker is schedulable. +- `GET /ready` returns `503` with worker lifecycle details while no workers are schedulable. +- `/select`, `/select_and_reserve`, `/reservations`, `/potential_loads`, and + `/overlap_scores` return `503` while the target worker pool is not ready. + +## Worker Lifecycle + +Workers are accepted into a single `WorkerCatalog` and become schedulable only +after the reconciler has published consistent derived state to the scheduler, +slot tracker, and indexer/listener view. + +Lifecycle values: + +- `incomplete`: accepted, but missing metadata required for valid booking. +- `schedulable`: metadata is complete and the scheduler can safely book it. +- `draining`: removed from the scheduler view before cleanup. +- `unschedulable`: cleanup has completed and the worker is no longer selectable. + +When queueing is enabled, `max_num_batched_tokens` is required before a worker +can become schedulable. When KV events are enabled, KV event listener endpoints +are also required. + +## Selection APIs + +### `POST /select` + +Selects a worker through the existing `SchedulerQueue` admission path with +`update_states=false`. + +This endpoint is query-only. It does not book active-load state. If +`selection_id` is provided, the service accepts and echoes it for observability, +but it does not dedupe, replay, or otherwise make retries idempotent. Retrying +`/select` may produce a fresh selection. + +### `POST /select_and_reserve` + +Selects a worker through the same `SchedulerQueue` admission path and books +active-load state before returning. + +`reservation_id` is used as the scheduler and slot-tracker request ID. If it is +omitted, the service generates one. Duplicate IDs use the existing slot-tracker +behavior and may return `409 Conflict`; V1 does not provide retry idempotency, +distributed atomicity, or exactly-once booking semantics. + +### `POST /reservations` + +Books an explicit post-select reservation on the supplied worker and dp-rank. +`reservation_id` is the booking request ID. Duplicate IDs use existing +slot-tracker behavior and may return `409 Conflict`. + +Clients must treat reservation creation as at-most-once unless they can tolerate +duplicate/conflict handling. + +### Lifecycle + +- `POST /reservations/{reservation_id}/prefill_complete` +- `DELETE /reservations/{reservation_id}` + +Lifecycle endpoints operate on `reservation_id`. + +## Request Inputs + +For `/select`, `/select_and_reserve`, `/potential_loads`, and `/overlap_scores`: + +- If `token_ids` is present, the service computes block hashes and sequence + hashes internally and ignores supplied hashes. `isl_tokens` is derived from + `token_ids.len()`. +- Otherwise, callers must provide `block_hashes`, `sequence_hashes`, and + `isl_tokens`. + +For `POST /reservations`: + +- If `token_ids` is present, the service computes sequence hashes internally. +- Otherwise, callers must provide `sequence_hashes` and `isl_tokens`. + +Signed hash values are reinterpreted bit-for-bit as unsigned router hashes, +matching the standalone indexer and slot-tracker APIs. + +## Retry Semantics + +V1 intentionally has no internal idempotency ledger. + +- `/select` retries may return a different worker. +- `/select_and_reserve` retries with the same `reservation_id` may return + `409 Conflict` if the first booking succeeded. +- `POST /reservations` retries with the same `reservation_id` may return + `409 Conflict` if the first booking succeeded. +- Clients that need replayable retry semantics should keep their own request + ledger. The service core is structured so a core idempotency ledger can be + added later without changing the scheduler queue contract. diff --git a/lib/kv-router/src/services/selection/catalog.rs b/lib/kv-router/src/services/selection/catalog.rs new file mode 100644 index 000000000000..205e74a9f359 --- /dev/null +++ b/lib/kv-router/src/services/selection/catalog.rs @@ -0,0 +1,144 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashMap; + +use parking_lot::RwLock; + +use crate::protocols::WorkerId; + +use super::error::SelectionError; +use super::types::{ + SelectionKey, SelectionWorkerConfig, WorkerCatalogRecord, WorkerLifecycle, WorkerPatchRequest, + WorkerRequest, +}; + +#[derive(Debug, Default)] +pub(super) struct WorkerCatalog { + workers: RwLock>, +} + +impl WorkerCatalog { + pub(super) fn upsert( + &self, + req: WorkerRequest, + ) -> (Option, WorkerCatalogRecord) { + let mut workers = self.workers.write(); + let previous = workers.get(&req.worker_id).cloned(); + let record = WorkerCatalogRecord::new(req); + workers.insert(record.worker_id, record.clone()); + (previous, record) + } + + pub(super) fn patch( + &self, + worker_id: WorkerId, + patch: WorkerPatchRequest, + ) -> Result<(WorkerCatalogRecord, WorkerCatalogRecord), SelectionError> { + let mut workers = self.workers.write(); + let Some(record) = workers.get_mut(&worker_id) else { + return Err(SelectionError::NotFound(format!( + "worker {worker_id} not found" + ))); + }; + let previous = record.clone(); + record.apply_patch(patch); + record.lifecycle = WorkerLifecycle::Incomplete; + record.not_schedulable_reasons.clear(); + Ok((previous, record.clone())) + } + + pub(super) fn get(&self, worker_id: WorkerId) -> Option { + self.workers.read().get(&worker_id).cloned() + } + + pub(super) fn set_lifecycle( + &self, + worker_id: WorkerId, + lifecycle: WorkerLifecycle, + reasons: Vec, + ) -> Option { + let mut workers = self.workers.write(); + let record = workers.get_mut(&worker_id)?; + record.lifecycle = lifecycle; + record.not_schedulable_reasons = reasons; + Some(record.clone()) + } + + pub(super) fn list( + &self, + model_name: Option<&str>, + tenant_id: Option<&str>, + ) -> Vec { + let mut records: Vec<_> = self + .workers + .read() + .values() + .filter(|record| { + model_name.is_none_or(|model_name| record.model_name == model_name) + && tenant_id.is_none_or(|tenant_id| record.tenant_id == tenant_id) + }) + .cloned() + .collect(); + records.sort_by_key(|record| { + ( + record.model_name.clone(), + record.tenant_id.clone(), + record.worker_id, + ) + }); + records + } + + pub(super) fn has_schedulable_for_key(&self, key: &SelectionKey) -> bool { + self.workers.read().values().any(|record| { + record.lifecycle == WorkerLifecycle::Schedulable + && record.model_name == key.model_name + && record.tenant_id == key.tenant_id + }) + } + + pub(super) fn scheduler_configs_for_key( + &self, + key: &SelectionKey, + ) -> HashMap { + self.workers + .read() + .values() + .filter(|record| { + record.lifecycle == WorkerLifecycle::Schedulable + && record.model_name == key.model_name + && record.tenant_id == key.tenant_id + }) + .filter_map(|record| { + record + .scheduler_config() + .map(|config| (record.worker_id, config)) + }) + .collect() + } + + pub(super) fn schedulable_count(&self) -> usize { + self.workers + .read() + .values() + .filter(|record| record.lifecycle == WorkerLifecycle::Schedulable) + .count() + } + + pub(super) fn schedulable_endpoint( + &self, + worker_id: WorkerId, + key: &SelectionKey, + ) -> Option { + let workers = self.workers.read(); + let record = workers.get(&worker_id)?; + if record.lifecycle != WorkerLifecycle::Schedulable + || record.model_name != key.model_name + || record.tenant_id != key.tenant_id + { + return None; + } + record.endpoint.clone() + } +} diff --git a/lib/kv-router/src/services/selection/error.rs b/lib/kv-router/src/services/selection/error.rs new file mode 100644 index 000000000000..7b4b2f154f04 --- /dev/null +++ b/lib/kv-router/src/services/selection/error.rs @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use axum::Json; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; + +use crate::scheduling::KvSchedulerError; +use crate::sequences::SequenceError; + +#[derive(Debug, thiserror::Error)] +pub enum SelectionError { + #[error("{0}")] + BadRequest(String), + #[error("{0}")] + NotReady(String), + #[error("{0}")] + NotFound(String), + #[error("{0}")] + Conflict(String), + #[error("{0}")] + Internal(String), + #[error(transparent)] + Scheduler(#[from] KvSchedulerError), + #[error(transparent)] + Sequence(#[from] SequenceError), +} + +impl SelectionError { + fn status(&self) -> StatusCode { + match self { + Self::BadRequest(_) => StatusCode::BAD_REQUEST, + Self::NotReady(_) => StatusCode::SERVICE_UNAVAILABLE, + Self::NotFound(_) => StatusCode::NOT_FOUND, + Self::Conflict(_) => StatusCode::CONFLICT, + Self::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, + Self::Scheduler(error) => scheduler_error_status(error), + Self::Sequence(error) => sequence_error_status(error), + } + } +} + +fn scheduler_error_status(error: &KvSchedulerError) -> StatusCode { + match error { + KvSchedulerError::NoEndpoints + | KvSchedulerError::SubscriberShutdown + | KvSchedulerError::InitFailed(_) => StatusCode::SERVICE_UNAVAILABLE, + KvSchedulerError::Backpressure { .. } + | KvSchedulerError::AllEligibleWorkersOverloaded + | KvSchedulerError::PinnedWorkerOverloaded { .. } => StatusCode::TOO_MANY_REQUESTS, + KvSchedulerError::PinnedWorkerNotAllowed { .. } => StatusCode::BAD_REQUEST, + KvSchedulerError::BookingFailed(message) if message.contains("already exists") => { + StatusCode::CONFLICT + } + KvSchedulerError::BookingFailed(_) => StatusCode::CONFLICT, + } +} + +fn sequence_error_status(error: &SequenceError) -> StatusCode { + match error { + SequenceError::WorkerNotFound { .. } | SequenceError::RequestNotFound { .. } => { + StatusCode::NOT_FOUND + } + SequenceError::DuplicateRequest { .. } => StatusCode::CONFLICT, + SequenceError::ReplicaSyncPublishFailed(_) => StatusCode::INTERNAL_SERVER_ERROR, + } +} + +impl IntoResponse for SelectionError { + fn into_response(self) -> Response { + ( + self.status(), + Json(serde_json::json!({"error": self.to_string()})), + ) + .into_response() + } +} diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs new file mode 100644 index 000000000000..a15296a6fc81 --- /dev/null +++ b/lib/kv-router/src/services/selection/input.rs @@ -0,0 +1,137 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dynamo_tokens::SequenceHash; +use serde::Deserialize; + +use crate::protocols::{ + BlockHashOptions, LocalBlockHash, compute_block_hash_for_seq, compute_seq_hash_for_block, +}; + +use super::error::SelectionError; + +#[derive(Debug, Clone, Deserialize)] +pub struct PromptRequest { + #[serde(default)] + pub token_ids: Option>, + #[serde(default)] + pub block_hashes: Option>, + #[serde(default)] + pub sequence_hashes: Option>, + #[serde(default)] + pub isl_tokens: Option, + #[serde(default)] + pub lora_name: Option, +} + +impl PromptRequest { + pub(super) fn normalize_for_selection( + &self, + block_size: u32, + ) -> Result { + if let Some(token_ids) = &self.token_ids { + return Ok(normalize_tokens( + token_ids, + block_size, + self.lora_name.as_deref(), + )); + } + + let block_hashes = self.block_hashes.as_ref().ok_or_else(|| { + SelectionError::BadRequest("block_hashes is required without token_ids".to_string()) + })?; + let sequence_hashes = self.sequence_hashes.as_ref().ok_or_else(|| { + SelectionError::BadRequest("sequence_hashes is required without token_ids".to_string()) + })?; + let isl_tokens = self.isl_tokens.ok_or_else(|| { + SelectionError::BadRequest("isl_tokens is required without token_ids".to_string()) + })?; + normalize_hashes(block_hashes, sequence_hashes, isl_tokens) + } + + pub(super) fn normalize_for_reservation( + &self, + block_size: u32, + ) -> Result { + if let Some(token_ids) = &self.token_ids { + let normalized = normalize_tokens(token_ids, block_size, self.lora_name.as_deref()); + return Ok(NormalizedReservation { + sequence_hashes: normalized.sequence_hashes, + }); + } + + let sequence_hashes = self.sequence_hashes.as_ref().ok_or_else(|| { + SelectionError::BadRequest("sequence_hashes is required without token_ids".to_string()) + })?; + if self.isl_tokens.is_none() { + return Err(SelectionError::BadRequest( + "isl_tokens is required without token_ids".to_string(), + )); + } + Ok(NormalizedReservation { + sequence_hashes: signed_sequence_hashes(sequence_hashes), + }) + } +} + +fn normalize_tokens( + token_ids: &[u32], + block_size: u32, + lora_name: Option<&str>, +) -> NormalizedPrompt { + let block_hashes = compute_block_hash_for_seq( + token_ids, + block_size, + BlockHashOptions { + lora_name, + ..Default::default() + }, + ); + let sequence_hashes = compute_seq_hash_for_block(&block_hashes); + NormalizedPrompt { + block_hashes, + sequence_hashes, + isl_tokens: token_ids.len(), + } +} + +fn normalize_hashes( + block_hashes: &[i64], + sequence_hashes: &[i64], + isl_tokens: usize, +) -> Result { + if isl_tokens == 0 { + return Err(SelectionError::BadRequest( + "isl_tokens must be greater than 0".to_string(), + )); + } + if block_hashes.len() != sequence_hashes.len() { + return Err(SelectionError::BadRequest(format!( + "block_hashes length {} must match sequence_hashes length {}", + block_hashes.len(), + sequence_hashes.len() + ))); + } + Ok(NormalizedPrompt { + block_hashes: block_hashes + .iter() + .map(|hash| LocalBlockHash(*hash as u64)) + .collect(), + sequence_hashes: signed_sequence_hashes(sequence_hashes), + isl_tokens, + }) +} + +fn signed_sequence_hashes(sequence_hashes: &[i64]) -> Vec { + sequence_hashes.iter().map(|hash| *hash as u64).collect() +} + +pub(super) struct NormalizedPrompt { + pub(super) block_hashes: Vec, + pub(super) sequence_hashes: Vec, + pub(super) isl_tokens: usize, +} + +pub(super) struct NormalizedReservation { + pub(super) sequence_hashes: Vec, +} diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs new file mode 100644 index 000000000000..a13fbcbf261c --- /dev/null +++ b/lib/kv-router/src/services/selection/mod.rs @@ -0,0 +1,27 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Runtime-free worker selection service. +//! +//! The service owns worker selection and reservation state, but never forwards +//! model requests and never owns model responses. + +mod catalog; +mod core; +mod error; +mod input; +mod scoring; +mod server; +mod types; + +#[cfg(test)] +mod tests; + +pub use core::{SelectionCore, SelectionServiceConfig}; +pub use server::{AppState, create_router, run_server}; +pub use types::{ + ModelLoadResponse, OverlapScoresRequest, OverlapScoresResponse, PotentialLoadsRequest, + ReadyResponse, ReservationRequest, ReservationResponse, SelectAndReserveRequest, SelectRequest, + SelectResponse, SelectionKey, SelectionWorkerConfig, WorkerCatalogRecord, WorkerLifecycle, + WorkerOverlapScore, WorkerPatchRequest, WorkerRequest, +}; diff --git a/lib/kv-router/src/services/selection/scoring.rs b/lib/kv-router/src/services/selection/scoring.rs new file mode 100644 index 000000000000..664e7b0423e2 --- /dev/null +++ b/lib/kv-router/src/services/selection/scoring.rs @@ -0,0 +1,165 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{HashMap, HashSet}; + +use crate::indexer::TieredMatchDetails; +use crate::protocols::{StorageTier, WorkerWithDpRank}; +use crate::scheduling::TierOverlapBlocks; + +use super::types::{OverlapScoresResponse, WorkerOverlapScore}; + +#[derive(Default)] +pub(super) struct OverlapInputs { + pub(super) tier_overlap_blocks: TierOverlapBlocks, + pub(super) effective_overlap_blocks: HashMap, + pub(super) effective_cached_tokens: HashMap, +} + +#[derive(Default)] +pub(super) struct CacheHitEstimates { + pub(super) effective_overlap_blocks: HashMap, + pub(super) cached_tokens: HashMap, +} + +fn cache_hit_weight_for_tier(config: &crate::config::KvRouterConfig, tier: StorageTier) -> f64 { + match tier { + StorageTier::Device => 1.0, + StorageTier::HostPinned => config.host_cache_hit_weight, + StorageTier::Disk | StorageTier::External => config.disk_cache_hit_weight, + } +} + +pub(super) fn cache_hit_estimates_from_tiered_matches( + config: &crate::config::KvRouterConfig, + block_size: u32, + tiered_matches: &TieredMatchDetails, +) -> CacheHitEstimates { + let mut effective_overlap_blocks = HashMap::new(); + + for (worker, overlap) in &tiered_matches.device.overlap_scores.scores { + effective_overlap_blocks.insert(*worker, *overlap as f64); + } + + for (storage_tier, tier_matches) in &tiered_matches.lower_tier { + let weight = cache_hit_weight_for_tier(config, *storage_tier); + if weight == 0.0 { + continue; + } + for (worker, hits) in &tier_matches.hits { + *effective_overlap_blocks.entry(*worker).or_insert(0.0) += *hits as f64 * weight; + } + } + + let cached_tokens = effective_overlap_blocks + .iter() + .map(|(worker, overlap)| { + ( + *worker, + (*overlap * block_size as f64).round().max(0.0) as usize, + ) + }) + .collect(); + + CacheHitEstimates { + effective_overlap_blocks, + cached_tokens, + } +} + +pub(super) fn tier_overlap_blocks_from_tiered_matches( + tiered_matches: &TieredMatchDetails, +) -> TierOverlapBlocks { + let mut tier_overlap_blocks = TierOverlapBlocks::default(); + tier_overlap_blocks.device.extend( + tiered_matches + .device + .overlap_scores + .scores + .iter() + .map(|(worker, hits)| (*worker, *hits as usize)), + ); + + if let Some(host_matches) = tiered_matches.lower_tier.get(&StorageTier::HostPinned) { + tier_overlap_blocks.host_pinned.extend( + host_matches + .hits + .iter() + .map(|(worker, hits)| (*worker, *hits)), + ); + } + + for tier in [StorageTier::Disk, StorageTier::External] { + if let Some(matches) = tiered_matches.lower_tier.get(&tier) { + for (worker, hits) in &matches.hits { + *tier_overlap_blocks.disk.entry(*worker).or_default() += *hits; + } + } + } + + tier_overlap_blocks +} + +pub(super) fn build_overlap_scores_response( + config: &crate::config::KvRouterConfig, + tiered: &TieredMatchDetails, + block_size: u32, +) -> OverlapScoresResponse { + let mut all_workers = HashSet::new(); + for worker in tiered.device.overlap_scores.scores.keys() { + all_workers.insert(*worker); + } + for matches in tiered.lower_tier.values() { + for worker in matches.hits.keys() { + all_workers.insert(*worker); + } + } + + let host = tiered.lower_tier.get(&StorageTier::HostPinned); + let disk = tiered.lower_tier.get(&StorageTier::Disk); + let external = tiered.lower_tier.get(&StorageTier::External); + + let mut workers: Vec<_> = all_workers + .into_iter() + .map(|worker| { + let device_blocks = tiered + .device + .overlap_scores + .scores + .get(&worker) + .copied() + .unwrap_or(0) as usize; + let host_pinned_extension_blocks = host + .and_then(|matches| matches.hits.get(&worker)) + .copied() + .unwrap_or(0); + let disk_extension_blocks = disk + .and_then(|matches| matches.hits.get(&worker)) + .copied() + .unwrap_or(0) + + external + .and_then(|matches| matches.hits.get(&worker)) + .copied() + .unwrap_or(0); + let router_credit_blocks = device_blocks as f64 + + config.host_cache_hit_weight * host_pinned_extension_blocks as f64 + + config.disk_cache_hit_weight * disk_extension_blocks as f64; + + WorkerOverlapScore { + worker_id: worker.worker_id, + dp_rank: worker.dp_rank, + device_blocks, + host_pinned_extension_blocks, + disk_extension_blocks, + router_credit_blocks, + } + }) + .collect(); + workers.sort_by_key(|worker| (worker.worker_id, worker.dp_rank)); + + OverlapScoresResponse { + block_size, + num_blocks: tiered.device.overlap_scores.frequencies.len(), + workers, + } +} diff --git a/lib/kv-router/src/services/selection/server.rs b/lib/kv-router/src/services/selection/server.rs new file mode 100644 index 000000000000..e9479221bfd4 --- /dev/null +++ b/lib/kv-router/src/services/selection/server.rs @@ -0,0 +1,278 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::fmt; +use std::sync::Arc; + +use axum::extract::rejection::JsonRejection; +use axum::extract::{Path, Query, State}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::{delete, get, patch, post}; +use axum::{Json, Router}; +use serde::Deserialize; +use tokio::net::TcpListener; +use tokio_util::sync::CancellationToken; + +use crate::protocols::WorkerId; + +use super::core::{SelectionCore, SelectionServiceConfig}; +use super::types::{ + OverlapScoresRequest, PotentialLoadsRequest, REQUEST_BODY_LIMIT_BYTES, ReservationRequest, + SelectAndReserveRequest, SelectRequest, WorkerPatchRequest, WorkerRequest, +}; + +#[derive(Debug, Deserialize)] +struct FilterQuery { + model_name: Option, + tenant_id: Option, +} + +pub struct AppState { + pub core: Arc, +} + +async fn create_worker( + State(state): State>, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state.core.upsert_worker(req).await { + Ok(worker) => (StatusCode::CREATED, Json(worker)).into_response(), + Err(error) => error.into_response(), + } +} + +async fn patch_worker( + State(state): State>, + Path(worker_id): Path, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state.core.patch_worker(worker_id, req).await { + Ok(worker) => Json(worker).into_response(), + Err(error) => error.into_response(), + } +} + +async fn delete_worker( + State(state): State>, + Path(worker_id): Path, +) -> Response { + match state.core.delete_worker(worker_id).await { + Ok(worker) => Json(worker).into_response(), + Err(error) => error.into_response(), + } +} + +async fn list_workers( + State(state): State>, + Query(params): Query, +) -> Response { + Json( + state + .core + .list_workers(params.model_name.as_deref(), params.tenant_id.as_deref()), + ) + .into_response() +} + +async fn select( + State(state): State>, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state.core.select(req).await { + Ok(response) => Json(response).into_response(), + Err(error) => error.into_response(), + } +} + +async fn select_and_reserve( + State(state): State>, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state.core.select_and_reserve(req).await { + Ok(response) => Json(response).into_response(), + Err(error) => error.into_response(), + } +} + +async fn create_reservation( + State(state): State>, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state.core.create_reservation(req).await { + Ok(response) => (StatusCode::CREATED, Json(response)).into_response(), + Err(error) => error.into_response(), + } +} + +async fn prefill_complete( + State(state): State>, + Path(reservation_id): Path, +) -> Response { + match state.core.prefill_complete(&reservation_id).await { + Ok(()) => json_ok(StatusCode::OK), + Err(error) => error.into_response(), + } +} + +async fn delete_reservation( + State(state): State>, + Path(reservation_id): Path, +) -> Response { + match state.core.free_reservation(&reservation_id).await { + Ok(()) => json_ok(StatusCode::OK), + Err(error) => error.into_response(), + } +} + +async fn health() -> Response { + json_ok(StatusCode::OK) +} + +async fn ready(State(state): State>) -> Response { + let response = state.core.ready(); + if response.ready { + Json(response).into_response() + } else { + (StatusCode::SERVICE_UNAVAILABLE, Json(response)).into_response() + } +} + +async fn loads(State(state): State>, Query(params): Query) -> Response { + Json( + state + .core + .loads(params.model_name.as_deref(), params.tenant_id.as_deref()), + ) + .into_response() +} + +async fn potential_loads( + State(state): State>, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state.core.potential_loads(req).await { + Ok(response) => Json(response).into_response(), + Err(error) => error.into_response(), + } +} + +async fn overlap_scores( + State(state): State>, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state.core.overlap_scores(req).await { + Ok(response) => Json(response).into_response(), + Err(error) => error.into_response(), + } +} + +async fn not_found() -> Response { + json_error(StatusCode::NOT_FOUND, "route not found") +} + +async fn method_not_allowed() -> Response { + json_error(StatusCode::METHOD_NOT_ALLOWED, "method not allowed") +} + +fn json_ok(status: StatusCode) -> Response { + (status, Json(serde_json::json!({"status": "ok"}))).into_response() +} + +fn json_error(status: StatusCode, error: impl fmt::Display) -> Response { + ( + status, + Json(serde_json::json!({"error": error.to_string()})), + ) + .into_response() +} + +fn json_rejection(error: JsonRejection) -> Response { + json_error(error.status(), error.body_text()) +} + +pub fn create_router(state: Arc) -> Router { + Router::new() + .route("/select", post(select)) + .route("/select_and_reserve", post(select_and_reserve)) + .route("/reservations", post(create_reservation)) + .route( + "/reservations/{reservation_id}/prefill_complete", + post(prefill_complete), + ) + .route("/reservations/{reservation_id}", delete(delete_reservation)) + .route("/workers", post(create_worker).get(list_workers)) + .route( + "/workers/{worker_id}", + patch(patch_worker).delete(delete_worker), + ) + .route("/health", get(health)) + .route("/ready", get(ready)) + .route("/loads", get(loads)) + .route("/potential_loads", post(potential_loads)) + .route("/overlap_scores", post(overlap_scores)) + .fallback(not_found) + .method_not_allowed_fallback(method_not_allowed) + .layer(axum::extract::DefaultBodyLimit::max( + REQUEST_BODY_LIMIT_BYTES, + )) + .with_state(state) +} + +pub async fn run_server(config: SelectionServiceConfig) -> anyhow::Result<()> { + let cancel_token = CancellationToken::new(); + let shutdown_token = cancel_token.clone(); + tokio::spawn(async move { + tokio::signal::ctrl_c().await.ok(); + tracing::info!("received shutdown signal"); + shutdown_token.cancel(); + }); + + tracing::info!( + port = config.port, + threads = config.threads, + "Starting Dynamo selection service" + ); + + let core = Arc::new(SelectionCore::new( + config.kv_router_config, + config.threads, + cancel_token.clone(), + )); + let app = create_router(Arc::new(AppState { core })); + let listener = TcpListener::bind(("0.0.0.0", config.port)).await?; + axum::serve(listener, app) + .with_graceful_shutdown(async move { + cancel_token.cancelled().await; + }) + .await?; + Ok(()) +} diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs new file mode 100644 index 000000000000..b383389645fb --- /dev/null +++ b/lib/kv-router/src/services/selection/tests.rs @@ -0,0 +1,325 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use axum::Router; +use axum::body::{Body, to_bytes}; +use axum::http::{Request, StatusCode, header}; +use axum::response::Response; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; +use tower::ServiceExt; + +use super::*; + +fn test_config() -> crate::config::KvRouterConfig { + let mut config = crate::config::KvRouterConfig::default(); + config.use_kv_events = false; + config.router_queue_threshold = None; + config +} + +fn app() -> Router { + let core = Arc::new(SelectionCore::new( + test_config(), + 1, + CancellationToken::new(), + )); + create_router(Arc::new(AppState { core })) +} + +async fn response_json(response: Response) -> serde_json::Value { + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("read response body"); + serde_json::from_slice(&body).expect("response JSON") +} + +async fn post(app: Router, uri: &str, body: &str) -> Response { + app.oneshot( + Request::builder() + .method("POST") + .uri(uri) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap() +} + +async fn register_worker(app: Router, max_tokens: Option) -> Response { + register_worker_id(app, 1, max_tokens).await +} + +async fn register_worker_id(app: Router, worker_id: u64, max_tokens: Option) -> Response { + let mut body = serde_json::json!({ + "worker_id": worker_id, + "model_name": "model", + "endpoint": format!("http://worker-{worker_id}:8000"), + "block_size": 4 + }); + if let Some(max_tokens) = max_tokens { + body["max_num_batched_tokens"] = serde_json::json!(max_tokens); + } + post(app, "/workers", &body.to_string()).await +} + +#[tokio::test] +async fn ready_and_select_report_not_ready_without_schedulable_workers() { + let app = app(); + let ready_response = app + .clone() + .oneshot( + Request::builder() + .uri("/ready") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(ready_response.status(), StatusCode::SERVICE_UNAVAILABLE); + + let select_response = post( + app, + "/select", + r#"{"model_name":"model","token_ids":[1,2,3,4],"selection_id":"s1"}"#, + ) + .await; + assert_eq!(select_response.status(), StatusCode::SERVICE_UNAVAILABLE); +} + +#[tokio::test] +async fn incomplete_worker_is_accepted_but_not_schedulable() { + let mut config = test_config(); + config.router_queue_threshold = Some(1.0); + let core = Arc::new(SelectionCore::new(config, 1, CancellationToken::new())); + let app = create_router(Arc::new(AppState { core })); + + let response = post( + app.clone(), + "/workers", + r#"{"worker_id":1,"model_name":"model","endpoint":"http://worker-1:8000","block_size":4}"#, + ) + .await; + assert_eq!(response.status(), StatusCode::CREATED); + let body = response_json(response).await; + assert_eq!(body["lifecycle"], "incomplete"); + assert!( + body["not_schedulable_reasons"][0] + .as_str() + .unwrap() + .contains("max_num_batched_tokens") + ); + + let select_response = post( + app, + "/select", + r#"{"model_name":"model","token_ids":[1,2,3,4]}"#, + ) + .await; + assert_eq!(select_response.status(), StatusCode::SERVICE_UNAVAILABLE); +} + +#[tokio::test] +async fn select_echoes_selection_id_and_does_not_book_load() { + let app = app(); + assert_eq!( + register_worker(app.clone(), None).await.status(), + StatusCode::CREATED + ); + + let response = post( + app.clone(), + "/select", + r#"{"model_name":"model","token_ids":[1,2,3,4],"selection_id":"sel-a"}"#, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["selection_id"], "sel-a"); + assert_eq!(body["worker_id"], 1); + + let loads_response = app + .oneshot( + Request::builder() + .uri("/loads") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(loads_response.status(), StatusCode::OK); + let loads = response_json(loads_response).await; + assert_eq!(loads[0]["loads"][0]["active_requests"], 0); +} + +#[tokio::test] +async fn select_and_reserve_books_and_duplicate_reservation_conflicts() { + let app = app(); + assert_eq!( + register_worker(app.clone(), None).await.status(), + StatusCode::CREATED + ); + + let response = post( + app.clone(), + "/select_and_reserve", + r#"{"model_name":"model","token_ids":[1,2,3,4],"reservation_id":"res-a"}"#, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["reservation_id"], "res-a"); + + let duplicate = post( + app.clone(), + "/select_and_reserve", + r#"{"model_name":"model","token_ids":[1,2,3,4],"reservation_id":"res-a"}"#, + ) + .await; + assert_eq!(duplicate.status(), StatusCode::CONFLICT); + + let free = app + .oneshot( + Request::builder() + .method("DELETE") + .uri("/reservations/res-a") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(free.status(), StatusCode::OK); +} + +#[tokio::test] +async fn explicit_reservation_books_after_select() { + let app = app(); + assert_eq!( + register_worker(app.clone(), None).await.status(), + StatusCode::CREATED + ); + + let select_response = post( + app.clone(), + "/select", + r#"{"model_name":"model","token_ids":[1,2,3,4]}"#, + ) + .await; + assert_eq!(select_response.status(), StatusCode::OK); + let selected = response_json(select_response).await; + + let reservation = serde_json::json!({ + "model_name": "model", + "reservation_id": "res-b", + "worker_id": selected["worker_id"], + "dp_rank": selected["dp_rank"], + "sequence_hashes": [1], + "isl_tokens": 4 + }); + let reserve_response = post(app.clone(), "/reservations", &reservation.to_string()).await; + assert_eq!(reserve_response.status(), StatusCode::CREATED); + + let prefill_response = post(app, "/reservations/res-b/prefill_complete", "{}").await; + assert_eq!(prefill_response.status(), StatusCode::OK); +} + +#[tokio::test] +async fn explicit_reservation_rejects_unschedulable_worker() { + let app = app(); + let incomplete = post( + app.clone(), + "/workers", + r#"{"worker_id":1,"model_name":"model","block_size":4}"#, + ) + .await; + assert_eq!(incomplete.status(), StatusCode::CREATED); + assert_eq!(response_json(incomplete).await["lifecycle"], "incomplete"); + assert_eq!( + register_worker_id(app.clone(), 2, None).await.status(), + StatusCode::CREATED + ); + + let reservation = serde_json::json!({ + "model_name": "model", + "reservation_id": "res-unschedulable", + "worker_id": 1, + "sequence_hashes": [1], + "isl_tokens": 4 + }); + let reserve_response = post(app, "/reservations", &reservation.to_string()).await; + assert_eq!(reserve_response.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn reconcile_rolls_back_partial_listener_registration() { + let mut config = test_config(); + config.use_kv_events = true; + let core = Arc::new(SelectionCore::new(config, 1, CancellationToken::new())); + let app = create_router(Arc::new(AppState { core })); + + let response = post( + app.clone(), + "/workers", + r#"{ + "worker_id": 7, + "model_name": "model", + "endpoint": "http://worker-7:8000", + "block_size": 4, + "data_parallel_size": 2, + "kv_events_endpoints": { + "0": "tcp://127.0.0.1:5557", + "1": "not-a-zmq-endpoint" + } + }"#, + ) + .await; + assert_eq!(response.status(), StatusCode::CREATED); + let body = response_json(response).await; + assert_eq!(body["lifecycle"], "incomplete"); + assert!( + body["not_schedulable_reasons"][0] + .as_str() + .unwrap() + .contains("reconciliation failed") + ); + + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/workers/7") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + r#"{ + "kv_events_endpoints": { + "0": "tcp://127.0.0.1:5557", + "1": "tcp://127.0.0.1:5558" + } + }"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["lifecycle"], "schedulable"); +} + +#[tokio::test] +async fn hash_path_validation_returns_bad_request() { + let app = app(); + assert_eq!( + register_worker(app.clone(), None).await.status(), + StatusCode::CREATED + ); + + let response = post( + app, + "/select", + r#"{"model_name":"model","block_hashes":[1],"sequence_hashes":[1,2],"isl_tokens":4}"#, + ) + .await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); +} diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs new file mode 100644 index 000000000000..7d49a08fcaaa --- /dev/null +++ b/lib/kv-router/src/services/selection/types.rs @@ -0,0 +1,532 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{HashMap, HashSet}; +use std::fmt; + +use serde::{Deserialize, Serialize}; + +use crate::protocols::{ + DpRank, KvTransferEnforcement, RoutingConstraints, WorkerConfigLike, WorkerId, WorkerWithDpRank, +}; +use crate::scheduling::PotentialLoad; +use crate::scheduling::config::RouterConfigOverride; +use crate::services::indexer::registry::IndexerKey; + +use super::input::PromptRequest; + +const DEFAULT_MODEL_NAME: &str = "default"; +const DEFAULT_TENANT_ID: &str = "default"; +pub(super) const WORKER_TYPE: &str = "select"; +pub(super) const REQUEST_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; + +fn default_model_name() -> String { + DEFAULT_MODEL_NAME.to_string() +} + +fn default_tenant_id() -> String { + DEFAULT_TENANT_ID.to_string() +} + +#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize)] +pub struct SelectionKey { + pub model_name: String, + pub tenant_id: String, +} + +impl SelectionKey { + pub(super) fn new(model_name: impl Into, tenant_id: impl Into) -> Self { + Self { + model_name: model_name.into(), + tenant_id: tenant_id.into(), + } + } + + pub(super) fn indexer_key(&self) -> IndexerKey { + IndexerKey { + model_name: self.model_name.clone(), + tenant_id: self.tenant_id.clone(), + } + } +} + +impl fmt::Display for SelectionKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "model={} tenant={}", self.model_name, self.tenant_id) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WorkerLifecycle { + Incomplete, + Schedulable, + Draining, + Unschedulable, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct SelectionWorkerConfig { + pub endpoint: String, + pub data_parallel_start_rank: u32, + pub data_parallel_size: u32, + pub max_num_batched_tokens: Option, + pub total_kv_blocks: Option, + pub stable_routing_id: Option, + #[serde(default)] + pub taints: HashSet, + #[serde(default)] + pub topology_domains: HashMap, + pub kv_transfer_domain: Option, + pub kv_transfer_enforcement: Option, + pub kv_transfer_preferred_weight: Option, +} + +impl WorkerConfigLike for SelectionWorkerConfig { + fn data_parallel_start_rank(&self) -> u32 { + self.data_parallel_start_rank + } + + fn data_parallel_size(&self) -> u32 { + self.data_parallel_size + } + + fn max_num_batched_tokens(&self) -> Option { + self.max_num_batched_tokens + } + + fn total_kv_blocks(&self) -> Option { + self.total_kv_blocks + } + + fn taints(&self) -> &HashSet { + &self.taints + } + + fn stable_routing_id(&self) -> Option<&str> { + self.stable_routing_id.as_deref() + } + + fn topology_domains(&self) -> Option<&HashMap> { + Some(&self.topology_domains) + } + + fn kv_transfer_domain(&self) -> Option<&str> { + self.kv_transfer_domain.as_deref() + } + + fn kv_transfer_enforcement(&self) -> Option { + self.kv_transfer_enforcement + } + + fn kv_transfer_preferred_weight(&self) -> Option { + self.kv_transfer_preferred_weight + } +} + +#[derive(Debug, Clone, Serialize)] +pub struct WorkerCatalogRecord { + pub worker_id: WorkerId, + pub model_name: String, + pub tenant_id: String, + pub lifecycle: WorkerLifecycle, + pub endpoint: Option, + pub kv_events_endpoint: Option, + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub kv_events_endpoints: HashMap, + pub replay_endpoint: Option, + pub block_size: Option, + pub data_parallel_start_rank: Option, + pub data_parallel_size: Option, + pub max_num_batched_tokens: Option, + pub total_kv_blocks: Option, + pub stable_routing_id: Option, + #[serde(default)] + pub taints: HashSet, + #[serde(default)] + pub topology_domains: HashMap, + pub kv_transfer_domain: Option, + pub kv_transfer_enforcement: Option, + pub kv_transfer_preferred_weight: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub not_schedulable_reasons: Vec, +} + +impl WorkerCatalogRecord { + pub(super) fn new(req: WorkerRequest) -> Self { + Self { + worker_id: req.worker_id, + model_name: req.model_name, + tenant_id: req.tenant_id, + lifecycle: WorkerLifecycle::Incomplete, + endpoint: req.endpoint, + kv_events_endpoint: req.kv_events_endpoint, + kv_events_endpoints: req.kv_events_endpoints, + replay_endpoint: req.replay_endpoint, + block_size: req.block_size, + data_parallel_start_rank: req.data_parallel_start_rank, + data_parallel_size: req.data_parallel_size, + max_num_batched_tokens: req.max_num_batched_tokens, + total_kv_blocks: req.total_kv_blocks, + stable_routing_id: req.stable_routing_id, + taints: req.taints, + topology_domains: req.topology_domains, + kv_transfer_domain: req.kv_transfer_domain, + kv_transfer_enforcement: req.kv_transfer_enforcement, + kv_transfer_preferred_weight: req.kv_transfer_preferred_weight, + not_schedulable_reasons: Vec::new(), + } + } + + pub(super) fn key(&self) -> SelectionKey { + SelectionKey::new(self.model_name.clone(), self.tenant_id.clone()) + } + + pub(super) fn dp_start(&self) -> u32 { + self.data_parallel_start_rank.unwrap_or(0) + } + + pub(super) fn dp_size(&self) -> u32 { + self.data_parallel_size.unwrap_or(1) + } + + pub(super) fn dp_ranks(&self) -> impl Iterator { + let start = self.dp_start(); + let size = self.dp_size(); + start..start.saturating_add(size) + } + + pub(super) fn scheduler_config(&self) -> Option { + Some(SelectionWorkerConfig { + endpoint: self.endpoint.clone()?, + data_parallel_start_rank: self.dp_start(), + data_parallel_size: self.dp_size(), + max_num_batched_tokens: self.max_num_batched_tokens, + total_kv_blocks: self.total_kv_blocks, + stable_routing_id: self.stable_routing_id.clone(), + taints: self.taints.clone(), + topology_domains: self.topology_domains.clone(), + kv_transfer_domain: self.kv_transfer_domain.clone(), + kv_transfer_enforcement: self.kv_transfer_enforcement, + kv_transfer_preferred_weight: self.kv_transfer_preferred_weight, + }) + } + + pub(super) fn listener_endpoints(&self) -> HashMap { + if !self.kv_events_endpoints.is_empty() { + return self.kv_events_endpoints.clone(); + } + + match (self.dp_size(), self.kv_events_endpoint.clone()) { + (1, Some(endpoint)) => HashMap::from([(self.dp_start(), endpoint)]), + _ => HashMap::new(), + } + } + + pub(super) fn missing_schedulable_metadata( + &self, + queueing_enabled: bool, + kv_events_enabled: bool, + ) -> Vec { + let mut missing = Vec::new(); + + if self.endpoint.as_deref().is_none_or(str::is_empty) { + missing.push("endpoint is required".to_string()); + } + if self.block_size.is_none_or(|block_size| block_size == 0) { + missing.push("block_size must be greater than 0".to_string()); + } + if self.dp_size() == 0 { + missing.push("data_parallel_size must be greater than 0".to_string()); + } + if queueing_enabled && self.max_num_batched_tokens.is_none() { + missing + .push("max_num_batched_tokens is required while queueing is enabled".to_string()); + } + if kv_events_enabled { + let endpoints = self.listener_endpoints(); + for rank in self.dp_ranks() { + if endpoints + .get(&rank) + .is_none_or(|endpoint| endpoint.is_empty()) + { + missing.push(format!("kv_events endpoint is required for dp_rank {rank}")); + } + } + } + + missing + } +} + +#[derive(Debug, Deserialize)] +pub struct WorkerRequest { + pub worker_id: WorkerId, + #[serde(default = "default_model_name")] + pub model_name: String, + #[serde(default = "default_tenant_id")] + pub tenant_id: String, + #[serde(default)] + pub endpoint: Option, + #[serde(default)] + pub kv_events_endpoint: Option, + #[serde(default)] + pub kv_events_endpoints: HashMap, + #[serde(default)] + pub replay_endpoint: Option, + #[serde(default)] + pub block_size: Option, + #[serde(default)] + pub data_parallel_start_rank: Option, + #[serde(default)] + pub data_parallel_size: Option, + #[serde(default)] + pub max_num_batched_tokens: Option, + #[serde(default)] + pub total_kv_blocks: Option, + #[serde(default)] + pub stable_routing_id: Option, + #[serde(default)] + pub taints: HashSet, + #[serde(default)] + pub topology_domains: HashMap, + #[serde(default)] + pub kv_transfer_domain: Option, + #[serde(default)] + pub kv_transfer_enforcement: Option, + #[serde(default)] + pub kv_transfer_preferred_weight: Option, +} + +#[derive(Debug, Deserialize)] +pub struct WorkerPatchRequest { + #[serde(default)] + pub endpoint: Option, + #[serde(default)] + pub kv_events_endpoint: Option, + #[serde(default)] + pub kv_events_endpoints: Option>, + #[serde(default)] + pub replay_endpoint: Option, + #[serde(default)] + pub block_size: Option, + #[serde(default)] + pub data_parallel_start_rank: Option, + #[serde(default)] + pub data_parallel_size: Option, + #[serde(default)] + pub max_num_batched_tokens: Option, + #[serde(default)] + pub total_kv_blocks: Option, + #[serde(default)] + pub stable_routing_id: Option, + #[serde(default)] + pub taints: Option>, + #[serde(default)] + pub topology_domains: Option>, + #[serde(default)] + pub kv_transfer_domain: Option, + #[serde(default)] + pub kv_transfer_enforcement: Option, + #[serde(default)] + pub kv_transfer_preferred_weight: Option, +} + +impl WorkerCatalogRecord { + pub(super) fn apply_patch(&mut self, patch: WorkerPatchRequest) { + if patch.endpoint.is_some() { + self.endpoint = patch.endpoint; + } + if patch.kv_events_endpoint.is_some() { + self.kv_events_endpoint = patch.kv_events_endpoint; + } + if let Some(endpoints) = patch.kv_events_endpoints { + self.kv_events_endpoints = endpoints; + } + if patch.replay_endpoint.is_some() { + self.replay_endpoint = patch.replay_endpoint; + } + if patch.block_size.is_some() { + self.block_size = patch.block_size; + } + if patch.data_parallel_start_rank.is_some() { + self.data_parallel_start_rank = patch.data_parallel_start_rank; + } + if patch.data_parallel_size.is_some() { + self.data_parallel_size = patch.data_parallel_size; + } + if patch.max_num_batched_tokens.is_some() { + self.max_num_batched_tokens = patch.max_num_batched_tokens; + } + if patch.total_kv_blocks.is_some() { + self.total_kv_blocks = patch.total_kv_blocks; + } + if patch.stable_routing_id.is_some() { + self.stable_routing_id = patch.stable_routing_id; + } + if let Some(taints) = patch.taints { + self.taints = taints; + } + if let Some(topology_domains) = patch.topology_domains { + self.topology_domains = topology_domains; + } + if patch.kv_transfer_domain.is_some() { + self.kv_transfer_domain = patch.kv_transfer_domain; + } + if patch.kv_transfer_enforcement.is_some() { + self.kv_transfer_enforcement = patch.kv_transfer_enforcement; + } + if patch.kv_transfer_preferred_weight.is_some() { + self.kv_transfer_preferred_weight = patch.kv_transfer_preferred_weight; + } + } +} + +#[derive(Debug, Deserialize)] +pub struct SelectRequest { + #[serde(default = "default_model_name")] + pub model_name: String, + #[serde(default = "default_tenant_id")] + pub tenant_id: String, + #[serde(default)] + pub selection_id: Option, + #[serde(flatten)] + pub prompt: PromptRequest, + #[serde(default)] + pub router_config_override: Option, + #[serde(default)] + pub expected_output_tokens: Option, + #[serde(default)] + pub priority_jump: Option, + #[serde(default)] + pub pinned_worker: Option, + #[serde(default)] + pub allowed_worker_ids: Option>, + #[serde(default)] + pub routing_constraints: RoutingConstraints, +} + +#[derive(Debug, Deserialize)] +pub struct SelectAndReserveRequest { + #[serde(default = "default_model_name")] + pub model_name: String, + #[serde(default = "default_tenant_id")] + pub tenant_id: String, + #[serde(default)] + pub selection_id: Option, + #[serde(default)] + pub reservation_id: Option, + #[serde(flatten)] + pub prompt: PromptRequest, + #[serde(default)] + pub router_config_override: Option, + #[serde(default)] + pub expected_output_tokens: Option, + #[serde(default)] + pub priority_jump: Option, + #[serde(default)] + pub pinned_worker: Option, + #[serde(default)] + pub allowed_worker_ids: Option>, + #[serde(default)] + pub routing_constraints: RoutingConstraints, +} + +#[derive(Debug, Deserialize)] +pub struct ReservationRequest { + #[serde(default = "default_model_name")] + pub model_name: String, + #[serde(default = "default_tenant_id")] + pub tenant_id: String, + pub reservation_id: String, + pub worker_id: WorkerId, + #[serde(default)] + pub dp_rank: Option, + #[serde(flatten)] + pub prompt: PromptRequest, + #[serde(default)] + pub router_config_override: Option, + #[serde(default)] + pub expected_output_tokens: Option, +} + +#[derive(Debug, Deserialize)] +pub struct PotentialLoadsRequest { + #[serde(default = "default_model_name")] + pub model_name: String, + #[serde(default = "default_tenant_id")] + pub tenant_id: String, + #[serde(flatten)] + pub prompt: PromptRequest, + #[serde(default)] + pub router_config_override: Option, +} + +#[derive(Debug, Deserialize)] +pub struct OverlapScoresRequest { + #[serde(default = "default_model_name")] + pub model_name: String, + #[serde(default = "default_tenant_id")] + pub tenant_id: String, + #[serde(flatten)] + pub prompt: PromptRequest, +} + +#[derive(Debug, Serialize)] +pub struct SelectResponse { + #[serde(skip_serializing_if = "Option::is_none")] + pub selection_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub reservation_id: Option, + pub model_name: String, + pub tenant_id: String, + pub worker_id: WorkerId, + pub dp_rank: DpRank, + pub endpoint: String, + pub effective_overlap_blocks: f64, + pub cached_tokens: usize, + pub block_size: u32, +} + +#[derive(Debug, Serialize)] +pub struct ReservationResponse { + pub reservation_id: String, + pub model_name: String, + pub tenant_id: String, + pub worker_id: WorkerId, + pub dp_rank: DpRank, + pub endpoint: String, +} + +#[derive(Debug, Serialize)] +pub struct ReadyResponse { + pub ready: bool, + pub schedulable_workers: usize, + pub workers: Vec, +} + +#[derive(Debug, Serialize)] +pub struct ModelLoadResponse { + pub model_name: String, + pub tenant_id: String, + pub loads: Vec, + pub pending_count: usize, + pub pending_isl_tokens: usize, +} + +#[derive(Debug, Serialize)] +pub struct WorkerOverlapScore { + pub worker_id: WorkerId, + pub dp_rank: DpRank, + pub device_blocks: usize, + pub host_pinned_extension_blocks: usize, + pub disk_extension_blocks: usize, + pub router_credit_blocks: f64, +} + +#[derive(Debug, Serialize)] +pub struct OverlapScoresResponse { + pub block_size: u32, + pub num_blocks: usize, + pub workers: Vec, +} diff --git a/tests/router/common.py b/tests/router/common.py index 4088259fa856..13653339dad6 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -2723,6 +2723,7 @@ def _test_router_decisions( durable_kv_events: bool = False, router_event_threads: int = 4, standalone_indexer_url: Optional[str] = None, + standalone_selector_url: Optional[str] = None, router_aic_config: Optional[dict[str, Any]] = None, router_predicted_ttl_secs: Optional[float] = None, initial_wait: float = 0.25, @@ -3046,6 +3047,49 @@ async def _verify_scores(): asyncio.run(_verify_scores()) + # Verify standalone selection service scores via HTTP POST /overlap_scores. + if standalone_selector_url: + _dp_a = dp_rank_a if dp_rank_a is not None else 0 + _dp_b = dp_rank_b if dp_rank_b is not None else 0 + + async def _verify_selection_service_scores(): + # The sidecar readiness gate confirms registration; allow KV events to settle. + await asyncio.sleep(3) + + async with aiohttp.ClientSession() as session: + async with session.post( + f"{standalone_selector_url}/overlap_scores", + json={"token_ids": req4_tokens, "model_name": model_name}, + ) as resp: + assert ( + resp.status == 200 + ), f"POST /overlap_scores failed: {resp.status} {await resp.text()}" + body = await resp.json() + + scores = { + (worker["worker_id"], worker["dp_rank"]): worker + for worker in body["workers"] + } + score_a = scores[(worker_a_id, _dp_a)]["device_blocks"] + score_b = scores[(worker_b_id, _dp_b)]["device_blocks"] + + logger.info( + "Standalone selection /overlap_scores: %s[%s]=%s, %s[%s]=%s", + worker_a_id, + _dp_a, + score_a, + worker_b_id, + _dp_b, + score_b, + ) + assert score_a > score_b, ( + f"Expected selection service worker {worker_a_id} dp_rank {_dp_a} " + f"score {score_a} > worker {worker_b_id} dp_rank {_dp_b} " + f"score {score_b} for req4 tokens" + ) + + asyncio.run(_verify_selection_service_scores()) + def _test_busy_threshold_endpoint( engine_workers, diff --git a/tests/router/helper.py b/tests/router/helper.py index 07fa1dcf0ac9..395357087951 100644 --- a/tests/router/helper.py +++ b/tests/router/helper.py @@ -38,6 +38,11 @@ def get_kv_indexer_command() -> list[str]: return [sys.executable, "-m", "dynamo.indexer"] +def get_select_service_command() -> list[str]: + """Return the preferred standalone selection service command.""" + return [sys.executable, "-m", "dynamo.select_service"] + + def get_kv_indexer_test_env() -> Dict[str, str]: """Indexer launch env that enables the listener-control test endpoints (gated off by default; used by the ZMQ replay scenario).""" @@ -430,6 +435,58 @@ async def wait_for_indexer_workers_active( ) +async def wait_for_selection_service_ready( + selector_url: str, + expected_worker_ids: set[int], + timeout_s: float = 30.0, +) -> None: + """Wait until the standalone selection service reports expected workers ready.""" + if not expected_worker_ids: + return + + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + ready_url = f"{selector_url}/ready" + + async with aiohttp.ClientSession() as session: + while loop.time() < deadline: + remaining_s = deadline - loop.time() + if remaining_s <= 0: + break + + try: + request_timeout = aiohttp.ClientTimeout(total=min(2.0, remaining_s)) + async with session.get(ready_url, timeout=request_timeout) as resp: + if resp.status not in (200, 503): + await asyncio.sleep(0.5) + continue + body = await resp.json() + except (aiohttp.ClientError, asyncio.TimeoutError): + await asyncio.sleep(0.5) + continue + + workers_by_id = { + worker["worker_id"]: worker for worker in body.get("workers", []) + } + all_schedulable = all( + workers_by_id.get(worker_id, {}).get("lifecycle") == "schedulable" + for worker_id in expected_worker_ids + ) + if ( + resp.status == 200 + and body.get("ready") is True + and body.get("schedulable_workers", 0) >= len(expected_worker_ids) + and all_schedulable + ): + return + + await asyncio.sleep(0.5) + + raise RuntimeError( + f"Timed out waiting for selection service workers to become ready at {ready_url}" + ) + + async def send_request_with_retry(url: str, payload: dict, max_retries: int = 8): """Send a single request with exponential backoff retry""" wait_time = 1 # Start with 1 second diff --git a/tests/router/mocker_process.py b/tests/router/mocker_process.py index 9d8c09f2346a..8e27ee171305 100644 --- a/tests/router/mocker_process.py +++ b/tests/router/mocker_process.py @@ -14,7 +14,9 @@ generate_random_suffix, get_kv_indexer_command, get_kv_indexer_test_env, + get_select_service_command, wait_for_indexer_workers_active, + wait_for_selection_service_ready, ) from tests.utils.constants import ROUTER_MODEL_NAME from tests.utils.managed_process import ManagedProcess @@ -126,9 +128,13 @@ def __init__( request_plane: str = "nats", zmq_kv_events: bool = False, standalone_indexer: bool = False, + standalone_selector: bool = False, model_name: str = "mocker", zmq_replay: bool = False, ): + if standalone_selector and not standalone_indexer: + raise ValueError("standalone_selector requires standalone_indexer=True") + namespace_suffix = generate_random_suffix() self.namespace = f"test-namespace-{namespace_suffix}" self.component_name = "mocker" @@ -138,10 +144,13 @@ def __init__( self._zmq_kv_events_ports: list[int] = [] self._zmq_replay_ports: list[int] = [] self._standalone_indexer = standalone_indexer + self._standalone_selector = standalone_selector self._standalone_indexer_port: Optional[int] = None self._standalone_indexer_b_port: Optional[int] = None + self._standalone_selector_port: Optional[int] = None self._indexer_process: Optional[ManagedProcess] = None self._indexer_b_process: Optional[ManagedProcess] = None + self._selector_process: Optional[ManagedProcess] = None self._mocker_processes: list[ManagedProcess] = [] self._request = request self._store_backend = store_backend @@ -186,10 +195,12 @@ def __init__( ) if standalone_indexer: - indexer_ports = allocate_ports(2, BASE_PORT) - self._standalone_indexer_port = indexer_ports[0] - self._standalone_indexer_b_port = indexer_ports[1] - request.addfinalizer(lambda: deallocate_ports(indexer_ports)) + sidecar_ports = allocate_ports(3 if standalone_selector else 2, BASE_PORT) + self._standalone_indexer_port = sidecar_ports[0] + self._standalone_indexer_b_port = sidecar_ports[1] + if standalone_selector: + self._standalone_selector_port = sidecar_ports[2] + request.addfinalizer(lambda: deallocate_ports(sidecar_ports)) self._process = None else: command = _build_mocker_command( @@ -213,10 +224,11 @@ def __init__( ) logger.info( - "Created mocker process with %s worker(s), endpoint: %s%s", + "Created mocker process with %s worker(s), endpoint: %s%s%s", num_mockers, self.endpoint, ", standalone_indexer=True" if standalone_indexer else "", + ", standalone_selector=True" if standalone_selector else "", ) @property @@ -231,6 +243,12 @@ def standalone_indexer_b_url(self) -> Optional[str]: return f"http://localhost:{self._standalone_indexer_b_port}" return None + @property + def standalone_selector_url(self) -> Optional[str]: + if self._standalone_selector_port is not None: + return f"http://localhost:{self._standalone_selector_port}" + return None + def __enter__(self): if self._standalone_indexer: block_size = self._mocker_args_orig.get("block_size", BLOCK_SIZE) @@ -257,6 +275,29 @@ def __enter__(self): self._standalone_indexer_port, ) self._indexer_process.__enter__() + + if self._standalone_selector: + selector_cmd = [ + *get_select_service_command(), + "--port", + str(self._standalone_selector_port), + ] + self._selector_process = ManagedProcess( + command=selector_cmd, + timeout=120, + display_output=True, + health_check_ports=[self._standalone_selector_port], + health_check_urls=[], + log_dir=self._request.node.name, + terminate_all_matching_process_names=False, + display_name="dynamo-select-service", + env=os.environ.copy(), + ) + logger.info( + "Starting standalone selection service on port %s", + self._standalone_selector_port, + ) + self._selector_process.__enter__() else: logger.info("Starting mocker process with %s worker(s)", self.num_workers) self._process.__enter__() @@ -346,6 +387,33 @@ async def launch_workers_with_indexer(self, endpoint): f"dp_rank {dp_rank}: {response.status} {body}" ) + if self.standalone_selector_url: + select_payload = { + "worker_id": new_worker_id, + "model_name": self.model_name, + "endpoint": self.endpoint, + "kv_events_endpoints": zmq_addresses, + "block_size": self._mocker_args_orig.get( + "block_size", BLOCK_SIZE + ), + "data_parallel_start_rank": 0, + "data_parallel_size": dp_size, + "max_num_batched_tokens": self._mocker_args_orig.get( + "max_num_batched_tokens", + 8192, + ), + } + async with session.post( + f"{self.standalone_selector_url}/workers", + json=select_payload, + ) as response: + if response.status != 201: + body = await response.text() + raise RuntimeError( + f"Failed to register selection service worker " + f"{new_worker_id}: {response.status} {body}" + ) + self.worker_id_to_zmq_ports[new_worker_id] = zmq_addresses logger.info( "Mocker %s: worker_id=%s, zmq_addresses=%s", @@ -357,6 +425,11 @@ async def launch_workers_with_indexer(self, endpoint): await wait_for_indexer_workers_active( self.standalone_indexer_url, self.worker_id_to_zmq_ports ) + if self.standalone_selector_url: + await wait_for_selection_service_ready( + self.standalone_selector_url, + set(self.worker_id_to_zmq_ports), + ) logger.info( "All %s mockers launched and registered with indexer", self.num_workers, @@ -423,6 +496,13 @@ def __exit__(self, exc_type, exc_val, exc_tb): logger.warning("Error stopping indexer B process: %s", error) self._indexer_b_process = None + if self._selector_process is not None: + try: + self._selector_process.__exit__(exc_type, exc_val, exc_tb) + except Exception as error: + logger.warning("Error stopping selection service process: %s", error) + self._selector_process = None + if self._indexer_process is not None: try: self._indexer_process.__exit__(exc_type, exc_val, exc_tb) diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index ca33fa1c90db..3b25e9a9e87c 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -801,6 +801,7 @@ def test_router_decisions( request_plane=request_plane, zmq_kv_events=zmq_kv_events, standalone_indexer=zmq_kv_events, + standalone_selector=zmq_kv_events, model_name=MODEL_NAME, ) as mockers: logger.info(f"All mockers using endpoint: {mockers.endpoint}") @@ -829,6 +830,7 @@ def test_router_decisions( use_kv_events=use_kv_events, durable_kv_events=durable_kv_events, standalone_indexer_url=mockers.standalone_indexer_url, + standalone_selector_url=mockers.standalone_selector_url, router_predicted_ttl_secs=router_predicted_ttl_secs, ) From 9409c1e75891df5c406e5576931251a0c71ee93e Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 11 Jun 2026 20:11:45 -0700 Subject: [PATCH 2/4] fix(kv-router): close select service CI gaps Signed-off-by: PeaBrane --- lib/bindings/python/rust/lib.rs | 6 +- lib/bindings/python/rust/llm/kv.rs | 4 +- lib/bindings/python/src/dynamo/_core.pyi | 4 + .../python/src/dynamo/select_service/main.py | 8 +- .../src/services/selection/README.md | 16 + .../src/services/selection/core/mod.rs | 750 ++++++++++++++++++ lib/kv-router/src/services/selection/error.rs | 3 - lib/kv-router/src/services/selection/input.rs | 54 +- lib/kv-router/src/services/selection/mod.rs | 8 +- .../src/services/selection/scoring.rs | 22 +- .../src/services/selection/server.rs | 26 +- lib/kv-router/src/services/selection/tests.rs | 216 ++++- lib/kv-router/src/services/selection/types.rs | 31 + lib/kv-router/src/services/zmq.rs | 2 +- 14 files changed, 1118 insertions(+), 32 deletions(-) create mode 100644 lib/kv-router/src/services/selection/core/mod.rs diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index fb14450ecb53..162faee7d816 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -252,7 +252,11 @@ where } fn standalone_to_pyerr(err: anyhow::Error) -> PyErr { - #[cfg(any(feature = "kv-indexer", feature = "slot-tracker"))] + #[cfg(any( + feature = "kv-indexer", + feature = "slot-tracker", + feature = "select-service" + ))] if let Some(clap_error) = err.downcast_ref::() { let _ = clap_error.print(); return pyo3::exceptions::PySystemExit::new_err(clap_error.exit_code()); diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 7779db583a33..a98623d64d0a 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -18,9 +18,7 @@ use crate::Endpoint; feature = "select-service" ))] use clap::Parser; -use dynamo_kv_router::config::{ - KvRouterConfig, RouterConfigOverride, kv_router_config_from_dynamo_env, -}; +use dynamo_kv_router::config::{KvRouterConfig, RouterConfigOverride}; use dynamo_kv_router::protocols::compute_block_hash_for_seq; use dynamo_kv_router::protocols::*; #[cfg(feature = "kv-indexer")] diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index f485dc3049e1..f5f7d88d2c6a 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -112,6 +112,10 @@ def run_slot_tracker(args: List[str]) -> None: """Run the KV router slot tracker with the given arguments.""" ... +def run_select_service(args: List[str]) -> None: + """Run the Dynamo selection service with the given arguments.""" + ... + # Any Python object that can be serialized to JSON (dict, list, str, int, etc.) JsonLike = Any diff --git a/lib/bindings/python/src/dynamo/select_service/main.py b/lib/bindings/python/src/dynamo/select_service/main.py index bf05f21ca839..23fbb3e34547 100644 --- a/lib/bindings/python/src/dynamo/select_service/main.py +++ b/lib/bindings/python/src/dynamo/select_service/main.py @@ -9,11 +9,5 @@ def main(argv: Sequence[str] | None = None) -> int: args = list(sys.argv[1:] if argv is None else argv) - try: - run_select_service(args) - except Exception as exc: - if "-h" in args or "--help" in args: - print(exc) - return 0 - raise + run_select_service(args) return 0 diff --git a/lib/kv-router/src/services/selection/README.md b/lib/kv-router/src/services/selection/README.md index d68dcdb7249d..52f09f4d2840 100644 --- a/lib/kv-router/src/services/selection/README.md +++ b/lib/kv-router/src/services/selection/README.md @@ -63,9 +63,13 @@ duplicate/conflict handling. ### Lifecycle - `POST /reservations/{reservation_id}/prefill_complete` +- `POST /reservations/{reservation_id}/output_block` - `DELETE /reservations/{reservation_id}` Lifecycle endpoints operate on `reservation_id`. +`/output_block` accepts an optional `decay_fraction` in `[0.0, 1.0]` and is +intended for clients that want parity with push-router output-block load +tracking while they own the response stream. ## Request Inputs @@ -74,6 +78,13 @@ For `/select`, `/select_and_reserve`, `/potential_loads`, and `/overlap_scores`: - If `token_ids` is present, the service computes block hashes and sequence hashes internally and ignores supplied hashes. `isl_tokens` is derived from `token_ids.len()`. +- If `mm_routing_info.routing_token_ids` is present and non-empty, those tokens + are used for routing hash computation instead of `token_ids`, with + `mm_routing_info.block_mm_infos` aligned to routing blocks. Top-level + `block_mm_infos` may be supplied with `token_ids` for direct block-level MM + hashing. +- `is_eagle` may be supplied per request; otherwise the worker catalog's + model-level `is_eagle` value is used, defaulting to `false`. - Otherwise, callers must provide `block_hashes`, `sequence_hashes`, and `isl_tokens`. @@ -85,6 +96,11 @@ For `POST /reservations`: Signed hash values are reinterpreted bit-for-bit as unsigned router hashes, matching the standalone indexer and slot-tracker APIs. +`/overlap_scores` returns all currently schedulable worker ranks for the model +and tenant, including zero-hit rows, using the same host/disk cumulative fields +as the Python-bound KV router. Shared-cache fields are present but disabled in +this runtime-free service unless a shared-cache integration is added later. + ## Retry Semantics V1 intentionally has no internal idempotency ledger. diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs new file mode 100644 index 000000000000..3f49ce82d54a --- /dev/null +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -0,0 +1,750 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashMap; +use std::sync::Arc; + +use dynamo_tokens::SequenceHash; +use parking_lot::RwLock; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; + +use crate::protocols::{ + ActiveLoad, ActiveSequenceEvent, LocalBlockHash, WorkerId, WorkerWithDpRank, +}; +use crate::scheduling::policy::RouterSchedulingPolicy; +use crate::scheduling::selector::DefaultWorkerSelector; +use crate::scheduling::{LocalScheduler, PotentialLoad}; +use crate::sequences::{ + ActiveSequencesMultiWorker, ReplicaWorkerPolicy, SequenceError, SequencePublisher, + SequenceRequest, +}; +use crate::services::indexer::backend::Indexer; +use crate::services::indexer::registry::WorkerRegistry; + +use super::catalog::WorkerCatalog; +use super::error::SelectionError; +use super::input::PromptRequest; +use super::scoring::{ + OverlapInputs, build_overlap_scores_response, cache_hit_estimates_from_tiered_matches, + tier_overlap_blocks_from_tiered_matches, +}; +use super::types::{ + ModelLoadResponse, OverlapScoresRequest, OverlapScoresResponse, PotentialLoadsRequest, + ReadyResponse, ReservationRequest, ReservationResponse, SelectAndReserveRequest, SelectRequest, + SelectResponse, SelectionKey, SelectionWorkerConfig, WORKER_TYPE, WorkerCatalogRecord, + WorkerLifecycle, WorkerPatchRequest, WorkerRequest, +}; + +type SelectionScheduler = + LocalScheduler; + +struct SelectionEntry { + key: SelectionKey, + block_size: u32, + is_eagle: bool, + indexer: Indexer, + workers_tx: watch::Sender>, + scheduler: SelectionScheduler, +} + +struct PreparedSelectionInputs { + sequence_hashes: Vec, + isl_tokens: usize, + overlap: OverlapInputs, +} + +#[derive(Clone)] +struct SelectionSequencePublisher; + +impl SequencePublisher for SelectionSequencePublisher { + async fn publish_event(&self, _event: &ActiveSequenceEvent) -> anyhow::Result<()> { + Ok(()) + } + + fn publish_load(&self, _load: ActiveLoad) {} + + fn observe_load( + &self, + _worker: &WorkerWithDpRank, + _worker_type: &str, + _blocks: usize, + _tokens: usize, + ) { + } +} + +#[derive(Debug, Clone)] +pub struct SelectionServiceConfig { + pub port: u16, + pub threads: usize, + pub kv_router_config: crate::config::KvRouterConfig, +} + +pub struct SelectionCore { + catalog: WorkerCatalog, + entries: RwLock>>, + indexer_registry: Arc, + kv_router_config: crate::config::KvRouterConfig, + cancel_token: CancellationToken, +} + +impl SelectionCore { + pub fn new( + kv_router_config: crate::config::KvRouterConfig, + indexer_threads: usize, + cancel_token: CancellationToken, + ) -> Self { + let indexer_registry = Arc::new(WorkerRegistry::new(indexer_threads)); + indexer_registry.signal_ready(); + Self { + catalog: WorkerCatalog::default(), + entries: RwLock::new(HashMap::new()), + indexer_registry, + kv_router_config, + cancel_token, + } + } + + pub async fn upsert_worker( + &self, + req: WorkerRequest, + ) -> Result { + let (previous, record) = self.catalog.upsert(req); + self.reconcile_worker(record.worker_id, previous).await + } + + pub async fn patch_worker( + &self, + worker_id: WorkerId, + patch: WorkerPatchRequest, + ) -> Result { + let (previous, record) = self.catalog.patch(worker_id, patch)?; + self.reconcile_worker(record.worker_id, Some(previous)) + .await + } + + pub async fn delete_worker( + &self, + worker_id: WorkerId, + ) -> Result { + let Some(previous) = self.catalog.get(worker_id) else { + return Err(SelectionError::NotFound(format!( + "worker {worker_id} not found" + ))); + }; + let key = previous.key(); + self.catalog + .set_lifecycle(worker_id, WorkerLifecycle::Draining, Vec::new()); + self.publish_scheduler_config(&key)?; + self.cleanup_indexer_registration(&previous).await; + let record = self + .catalog + .set_lifecycle(worker_id, WorkerLifecycle::Unschedulable, Vec::new()) + .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; + self.publish_scheduler_config(&key)?; + Ok(record) + } + + pub fn list_workers( + &self, + model_name: Option<&str>, + tenant_id: Option<&str>, + ) -> Vec { + self.catalog.list(model_name, tenant_id) + } + + pub fn ready(&self) -> ReadyResponse { + let schedulable_workers = self.catalog.schedulable_count(); + let workers = self.catalog.list(None, None); + ReadyResponse { + ready: schedulable_workers > 0, + schedulable_workers, + workers, + } + } + + async fn reconcile_worker( + &self, + worker_id: WorkerId, + previous: Option, + ) -> Result { + let Some(record) = self.catalog.get(worker_id) else { + return Err(SelectionError::NotFound(format!( + "worker {worker_id} not found" + ))); + }; + + if previous + .as_ref() + .is_some_and(|record| record.lifecycle == WorkerLifecycle::Schedulable) + { + self.catalog + .set_lifecycle(worker_id, WorkerLifecycle::Draining, Vec::new()); + self.publish_scheduler_config(&previous.as_ref().expect("checked").key())?; + self.cleanup_indexer_registration(previous.as_ref().expect("checked")) + .await; + } + + let reasons = record.missing_schedulable_metadata( + self.kv_router_config.router_queue_threshold.is_some(), + self.kv_router_config.use_kv_events, + ); + if !reasons.is_empty() { + let updated = self + .catalog + .set_lifecycle(worker_id, WorkerLifecycle::Incomplete, reasons) + .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; + self.publish_scheduler_config(&updated.key())?; + return Ok(updated); + } + + if let Err(error) = self.ensure_entry(&record) { + return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); + } + if self.kv_router_config.use_kv_events + && let Err(error) = self.register_indexer_listeners(&record).await + { + self.cleanup_indexer_registration(&record).await; + return self.mark_incomplete_after_reconcile_error(worker_id, record.key(), error); + } + + let updated = self + .catalog + .set_lifecycle(worker_id, WorkerLifecycle::Schedulable, Vec::new()) + .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; + self.publish_scheduler_config(&updated.key())?; + Ok(updated) + } + + fn mark_incomplete_after_reconcile_error( + &self, + worker_id: WorkerId, + key: SelectionKey, + error: SelectionError, + ) -> Result { + let updated = self + .catalog + .set_lifecycle( + worker_id, + WorkerLifecycle::Incomplete, + vec![format!("reconciliation failed: {error}")], + ) + .ok_or_else(|| SelectionError::NotFound(format!("worker {worker_id} not found")))?; + self.publish_scheduler_config(&key)?; + Ok(updated) + } + + fn ensure_entry( + &self, + record: &WorkerCatalogRecord, + ) -> Result, SelectionError> { + let block_size = record + .block_size + .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; + let is_eagle = record.is_eagle.unwrap_or(false); + let key = record.key(); + + if let Some(entry) = self.entries.read().get(&key).cloned() { + if entry.block_size != block_size { + return Err(SelectionError::Conflict(format!( + "block_size mismatch for {key}: existing={} requested={block_size}", + entry.block_size + ))); + } + if entry.is_eagle != is_eagle { + return Err(SelectionError::Conflict(format!( + "is_eagle mismatch for {key}: existing={} requested={is_eagle}", + entry.is_eagle + ))); + } + return Ok(entry); + } + + let mut entries = self.entries.write(); + if let Some(entry) = entries.get(&key).cloned() { + if entry.block_size != block_size { + return Err(SelectionError::Conflict(format!( + "block_size mismatch for {key}: existing={} requested={block_size}", + entry.block_size + ))); + } + if entry.is_eagle != is_eagle { + return Err(SelectionError::Conflict(format!( + "is_eagle mismatch for {key}: existing={} requested={is_eagle}", + entry.is_eagle + ))); + } + return Ok(entry); + } + + let (workers_tx, workers_rx) = watch::channel(HashMap::new()); + let slots = Arc::new(ActiveSequencesMultiWorker::new_with_replica_worker_policy( + SelectionSequencePublisher, + block_size as usize, + HashMap::new(), + self.kv_router_config.router_replica_sync, + uuid::Uuid::new_v4().as_u64_pair().0, + WORKER_TYPE, + ReplicaWorkerPolicy::RequireRegistered, + )); + slots.start_periodic_force_expiry_across_all_workers(self.cancel_token.child_token()); + + let selector = DefaultWorkerSelector::new(Some(self.kv_router_config.clone()), WORKER_TYPE); + let scheduler = LocalScheduler::new_without_overlap_refresh( + slots, + workers_rx, + self.kv_router_config.router_queue_threshold, + self.kv_router_config + .router_queue_by_incoming_missing_isl + .clone(), + block_size, + selector, + RouterSchedulingPolicy::new(self.kv_router_config.router_queue_policy), + None, + self.kv_router_config.router_queue_recheck_interval(), + self.kv_router_config.router_track_prefill_tokens, + self.cancel_token.child_token(), + WORKER_TYPE, + true, + ); + + let indexer = self + .indexer_registry + .get_or_create_indexer(key.indexer_key(), block_size); + let entry = Arc::new(SelectionEntry { + key: key.clone(), + block_size, + is_eagle, + indexer, + workers_tx, + scheduler, + }); + entries.insert(key, entry.clone()); + Ok(entry) + } + + async fn register_indexer_listeners( + &self, + record: &WorkerCatalogRecord, + ) -> Result<(), SelectionError> { + let block_size = record + .block_size + .ok_or_else(|| SelectionError::BadRequest("block_size is required".to_string()))?; + let mut endpoints: Vec<_> = record.listener_endpoints().into_iter().collect(); + endpoints.sort_by_key(|(dp_rank, _)| *dp_rank); + for (dp_rank, endpoint) in endpoints { + crate::services::zmq::validate_endpoint(&endpoint).map_err(|error| { + SelectionError::BadRequest(format!( + "invalid kv_events endpoint for worker {} dp_rank {dp_rank}: {error}", + record.worker_id + )) + })?; + if let Some(replay_endpoint) = record.replay_endpoint.as_deref() { + crate::services::zmq::validate_endpoint(replay_endpoint).map_err(|error| { + SelectionError::BadRequest(format!( + "invalid replay endpoint for worker {} dp_rank {dp_rank}: {error}", + record.worker_id + )) + })?; + } + self.indexer_registry + .register( + record.worker_id, + endpoint, + dp_rank, + record.model_name.clone(), + record.tenant_id.clone(), + block_size, + record.replay_endpoint.clone(), + ) + .await + .map_err(|error| SelectionError::BadRequest(error.to_string()))?; + } + Ok(()) + } + + async fn cleanup_indexer_registration(&self, record: &WorkerCatalogRecord) { + if self.kv_router_config.use_kv_events { + if let Err(error) = self + .indexer_registry + .deregister(record.worker_id, &record.model_name, &record.tenant_id) + .await + { + tracing::debug!( + worker_id = record.worker_id, + error = %error, + "indexer deregistration skipped or failed" + ); + } + return; + } + + let key = record.key().indexer_key(); + let indexer = self + .indexer_registry + .get_indexer(&key) + .map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(record.worker_id).await; + } + } + + fn publish_scheduler_config(&self, key: &SelectionKey) -> Result<(), SelectionError> { + let Some(entry) = self.entries.read().get(key).cloned() else { + return Ok(()); + }; + let workers = self.catalog.scheduler_configs_for_key(key); + entry.workers_tx.send(workers).map_err(|_| { + SelectionError::Internal(format!("scheduler worker watch closed for {key}")) + }) + } + + fn ready_entry(&self, key: &SelectionKey) -> Result, SelectionError> { + if self.catalog.schedulable_count() == 0 { + return Err(SelectionError::NotReady( + "no schedulable workers are available".to_string(), + )); + } + + let Some(entry) = self.entries.read().get(key).cloned() else { + return Err(SelectionError::NotReady(format!( + "no schedulable workers for {key}" + ))); + }; + if !self.catalog.has_schedulable_for_key(key) { + return Err(SelectionError::NotReady(format!( + "no schedulable workers for {key}" + ))); + } + Ok(entry) + } + + pub async fn select(&self, req: SelectRequest) -> Result { + let key = SelectionKey::new(req.model_name.clone(), req.tenant_id.clone()); + let entry = self.ready_entry(&key)?; + let prepared = self.prepare_selection_inputs(&entry, &req.prompt).await?; + let response = entry + .scheduler + .schedule_with_block_hashes( + None, + prepared.isl_tokens, + Some(prepared.sequence_hashes), + None, + prepared.overlap.tier_overlap_blocks, + prepared.overlap.effective_overlap_blocks, + prepared.overlap.effective_cached_tokens, + req.router_config_override.as_ref(), + false, + req.prompt.lora_name.clone(), + req.priority_jump.unwrap_or_default(), + req.expected_output_tokens, + req.pinned_worker, + req.allowed_worker_ids, + req.routing_constraints, + None, + ) + .await?; + self.selection_response(key, entry.block_size, req.selection_id, None, response) + } + + pub async fn select_and_reserve( + &self, + mut req: SelectAndReserveRequest, + ) -> Result { + let reservation_id = req + .reservation_id + .take() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let key = SelectionKey::new(req.model_name.clone(), req.tenant_id.clone()); + let entry = self.ready_entry(&key)?; + let prepared = self.prepare_selection_inputs(&entry, &req.prompt).await?; + let response = entry + .scheduler + .schedule_with_block_hashes( + Some(reservation_id.clone()), + prepared.isl_tokens, + Some(prepared.sequence_hashes), + None, + prepared.overlap.tier_overlap_blocks, + prepared.overlap.effective_overlap_blocks, + prepared.overlap.effective_cached_tokens, + req.router_config_override.as_ref(), + true, + req.prompt.lora_name.clone(), + req.priority_jump.unwrap_or_default(), + req.expected_output_tokens, + req.pinned_worker, + req.allowed_worker_ids, + req.routing_constraints, + None, + ) + .await?; + self.selection_response( + key, + entry.block_size, + req.selection_id, + Some(reservation_id), + response, + ) + } + + pub async fn create_reservation( + &self, + req: ReservationRequest, + ) -> Result { + let key = SelectionKey::new(req.model_name.clone(), req.tenant_id.clone()); + let entry = self.ready_entry(&key)?; + let normalized = req + .prompt + .normalize_for_reservation(entry.block_size, entry.is_eagle)?; + let worker = WorkerWithDpRank::new(req.worker_id, req.dp_rank.unwrap_or(0)); + let endpoint = self + .catalog + .schedulable_endpoint(req.worker_id, &key) + .ok_or_else(|| { + SelectionError::NotFound(format!( + "schedulable worker {} not found for {key}", + req.worker_id + )) + })?; + let track_prefill_tokens = req + .router_config_override + .as_ref() + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.kv_router_config.router_track_prefill_tokens); + + entry + .scheduler + .add_request(SequenceRequest { + request_id: req.reservation_id.clone(), + token_sequence: Some(normalized.sequence_hashes), + track_prefill_tokens, + expected_output_tokens: req.expected_output_tokens, + prefill_load_hint: None, + worker, + lora_name: req.prompt.lora_name.clone(), + }) + .await?; + + Ok(ReservationResponse { + reservation_id: req.reservation_id, + model_name: key.model_name, + tenant_id: key.tenant_id, + worker_id: worker.worker_id, + dp_rank: worker.dp_rank, + endpoint, + }) + } + + pub async fn prefill_complete(&self, reservation_id: &str) -> Result<(), SelectionError> { + let entries = { self.entries.read().values().cloned().collect::>() }; + for entry in entries { + match entry.scheduler.mark_prefill_completed(reservation_id).await { + Ok(()) => return Ok(()), + Err(SequenceError::RequestNotFound { .. }) => continue, + Err(error) => return Err(error.into()), + } + } + Err(SelectionError::NotFound(format!( + "reservation {reservation_id} not found" + ))) + } + + pub async fn free_reservation(&self, reservation_id: &str) -> Result<(), SelectionError> { + let entries = { self.entries.read().values().cloned().collect::>() }; + for entry in entries { + match entry.scheduler.free(reservation_id).await { + Ok(()) => return Ok(()), + Err(SequenceError::RequestNotFound { .. }) => continue, + Err(error) => return Err(error.into()), + } + } + Err(SelectionError::NotFound(format!( + "reservation {reservation_id} not found" + ))) + } + + pub fn add_output_block( + &self, + reservation_id: &str, + decay_fraction: Option, + ) -> Result<(), SelectionError> { + if let Some(frac) = decay_fraction + && !(0.0..=1.0).contains(&frac) + { + return Err(SelectionError::BadRequest( + "decay_fraction must be between 0.0 and 1.0".to_string(), + )); + } + + let entries = { self.entries.read().values().cloned().collect::>() }; + for entry in entries { + match entry + .scheduler + .add_output_block(reservation_id, decay_fraction) + { + Ok(()) => return Ok(()), + Err(SequenceError::RequestNotFound { .. }) => continue, + Err(error) => return Err(error.into()), + } + } + Err(SelectionError::NotFound(format!( + "reservation {reservation_id} not found" + ))) + } + + pub fn loads( + &self, + model_name: Option<&str>, + tenant_id: Option<&str>, + ) -> Vec { + let entries: Vec<_> = self.entries.read().values().cloned().collect(); + let mut loads = Vec::new(); + for entry in entries { + if model_name.is_some_and(|model_name| entry.key.model_name != model_name) + || tenant_id.is_some_and(|tenant_id| entry.key.tenant_id != tenant_id) + { + continue; + } + loads.push(ModelLoadResponse { + model_name: entry.key.model_name.clone(), + tenant_id: entry.key.tenant_id.clone(), + loads: entry + .scheduler + .get_potential_loads(None, 0, HashMap::new(), false), + pending_count: entry.scheduler.pending_count(), + pending_isl_tokens: entry.scheduler.pending_isl_tokens(), + }); + } + loads.sort_by(|a, b| (&a.model_name, &a.tenant_id).cmp(&(&b.model_name, &b.tenant_id))); + loads + } + + pub async fn potential_loads( + &self, + req: PotentialLoadsRequest, + ) -> Result, SelectionError> { + let key = SelectionKey::new(req.model_name.clone(), req.tenant_id.clone()); + let entry = self.ready_entry(&key)?; + let prepared = self.prepare_selection_inputs(&entry, &req.prompt).await?; + let track_prefill_tokens = req + .router_config_override + .as_ref() + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.kv_router_config.router_track_prefill_tokens); + Ok(entry.scheduler.get_potential_loads( + Some(prepared.sequence_hashes), + prepared.isl_tokens, + prepared.overlap.effective_cached_tokens, + track_prefill_tokens, + )) + } + + pub async fn overlap_scores( + &self, + req: OverlapScoresRequest, + ) -> Result { + let key = SelectionKey::new(req.model_name.clone(), req.tenant_id.clone()); + let entry = self.ready_entry(&key)?; + let normalized = req + .prompt + .normalize_for_selection(entry.block_size, entry.is_eagle)?; + let tiered = entry + .indexer + .find_tiered_matches(normalized.block_hashes) + .await + .map_err(|error| SelectionError::Internal(error.to_string()))?; + let schedulable_workers = self.schedulable_worker_ranks(&key); + Ok(build_overlap_scores_response( + &self.kv_router_config, + req.router_config_override.as_ref(), + &tiered, + entry.block_size, + schedulable_workers, + )) + } + + async fn prepare_selection_inputs( + &self, + entry: &SelectionEntry, + prompt: &PromptRequest, + ) -> Result { + let normalized = prompt.normalize_for_selection(entry.block_size, entry.is_eagle)?; + let overlap = self.overlap_inputs(entry, &normalized.block_hashes).await?; + Ok(PreparedSelectionInputs { + sequence_hashes: normalized.sequence_hashes, + isl_tokens: normalized.isl_tokens, + overlap, + }) + } + + async fn overlap_inputs( + &self, + entry: &SelectionEntry, + block_hashes: &[LocalBlockHash], + ) -> Result { + if block_hashes.is_empty() { + return Ok(OverlapInputs::default()); + } + let tiered = entry + .indexer + .find_tiered_matches(block_hashes.to_vec()) + .await + .map_err(|error| SelectionError::Internal(error.to_string()))?; + let estimates = cache_hit_estimates_from_tiered_matches( + &self.kv_router_config, + entry.block_size, + &tiered, + ); + Ok(OverlapInputs { + tier_overlap_blocks: tier_overlap_blocks_from_tiered_matches(&tiered), + effective_overlap_blocks: estimates.effective_overlap_blocks, + effective_cached_tokens: estimates.cached_tokens, + }) + } + + fn selection_response( + &self, + key: SelectionKey, + block_size: u32, + selection_id: Option, + reservation_id: Option, + response: crate::scheduling::SchedulingResponse, + ) -> Result { + let endpoint = self + .catalog + .schedulable_endpoint(response.best_worker.worker_id, &key) + .ok_or_else(|| { + SelectionError::Internal(format!( + "selected worker {} is no longer schedulable", + response.best_worker.worker_id + )) + })?; + Ok(SelectResponse { + selection_id, + reservation_id, + model_name: key.model_name, + tenant_id: key.tenant_id, + worker_id: response.best_worker.worker_id, + dp_rank: response.best_worker.dp_rank, + endpoint, + effective_overlap_blocks: response.effective_overlap_blocks, + cached_tokens: response.cached_tokens, + block_size, + }) + } + + fn schedulable_worker_ranks(&self, key: &SelectionKey) -> Vec { + let configs = self.catalog.scheduler_configs_for_key(key); + let mut workers = Vec::new(); + for (worker_id, config) in configs { + let start = config.data_parallel_start_rank; + let end = start.saturating_add(config.data_parallel_size); + for dp_rank in start..end { + workers.push(WorkerWithDpRank::new(worker_id, dp_rank)); + } + } + workers + } +} diff --git a/lib/kv-router/src/services/selection/error.rs b/lib/kv-router/src/services/selection/error.rs index 7b4b2f154f04..40e67f76f1e3 100644 --- a/lib/kv-router/src/services/selection/error.rs +++ b/lib/kv-router/src/services/selection/error.rs @@ -49,9 +49,6 @@ fn scheduler_error_status(error: &KvSchedulerError) -> StatusCode { | KvSchedulerError::AllEligibleWorkersOverloaded | KvSchedulerError::PinnedWorkerOverloaded { .. } => StatusCode::TOO_MANY_REQUESTS, KvSchedulerError::PinnedWorkerNotAllowed { .. } => StatusCode::BAD_REQUEST, - KvSchedulerError::BookingFailed(message) if message.contains("already exists") => { - StatusCode::CONFLICT - } KvSchedulerError::BookingFailed(_) => StatusCode::CONFLICT, } } diff --git a/lib/kv-router/src/services/selection/input.rs b/lib/kv-router/src/services/selection/input.rs index a15296a6fc81..f189eccb1add 100644 --- a/lib/kv-router/src/services/selection/input.rs +++ b/lib/kv-router/src/services/selection/input.rs @@ -5,16 +5,30 @@ use dynamo_tokens::SequenceHash; use serde::Deserialize; use crate::protocols::{ - BlockHashOptions, LocalBlockHash, compute_block_hash_for_seq, compute_seq_hash_for_block, + BlockExtraInfo, BlockHashOptions, LocalBlockHash, compute_block_hash_for_seq, + compute_seq_hash_for_block, }; use super::error::SelectionError; +type RoutingTokensAndMmInfos<'a> = (&'a [u32], Option<&'a [Option]>); + +#[derive(Debug, Clone, Deserialize)] +pub struct MmRoutingInfoRequest { + pub routing_token_ids: Vec, + #[serde(default)] + pub block_mm_infos: Vec>, +} + #[derive(Debug, Clone, Deserialize)] pub struct PromptRequest { #[serde(default)] pub token_ids: Option>, #[serde(default)] + pub mm_routing_info: Option, + #[serde(default)] + pub block_mm_infos: Option>>, + #[serde(default)] pub block_hashes: Option>, #[serde(default)] pub sequence_hashes: Option>, @@ -22,18 +36,23 @@ pub struct PromptRequest { pub isl_tokens: Option, #[serde(default)] pub lora_name: Option, + #[serde(default)] + pub is_eagle: Option, } impl PromptRequest { pub(super) fn normalize_for_selection( &self, block_size: u32, + default_is_eagle: bool, ) -> Result { - if let Some(token_ids) = &self.token_ids { + if let Some((token_ids, block_mm_infos)) = self.routing_tokens_and_mm_infos() { return Ok(normalize_tokens( token_ids, block_size, self.lora_name.as_deref(), + block_mm_infos, + self.is_eagle.unwrap_or(default_is_eagle), )); } @@ -52,9 +71,16 @@ impl PromptRequest { pub(super) fn normalize_for_reservation( &self, block_size: u32, + default_is_eagle: bool, ) -> Result { - if let Some(token_ids) = &self.token_ids { - let normalized = normalize_tokens(token_ids, block_size, self.lora_name.as_deref()); + if let Some((token_ids, block_mm_infos)) = self.routing_tokens_and_mm_infos() { + let normalized = normalize_tokens( + token_ids, + block_size, + self.lora_name.as_deref(), + block_mm_infos, + self.is_eagle.unwrap_or(default_is_eagle), + ); return Ok(NormalizedReservation { sequence_hashes: normalized.sequence_hashes, }); @@ -72,19 +98,37 @@ impl PromptRequest { sequence_hashes: signed_sequence_hashes(sequence_hashes), }) } + + fn routing_tokens_and_mm_infos(&self) -> Option> { + if let Some(mm_routing_info) = &self.mm_routing_info + && !mm_routing_info.routing_token_ids.is_empty() + { + return Some(( + &mm_routing_info.routing_token_ids, + Some(mm_routing_info.block_mm_infos.as_slice()), + )); + } + + self.token_ids + .as_deref() + .map(|token_ids| (token_ids, self.block_mm_infos.as_deref())) + } } fn normalize_tokens( token_ids: &[u32], block_size: u32, lora_name: Option<&str>, + block_mm_infos: Option<&[Option]>, + is_eagle: bool, ) -> NormalizedPrompt { let block_hashes = compute_block_hash_for_seq( token_ids, block_size, BlockHashOptions { + block_mm_infos, lora_name, - ..Default::default() + is_eagle: Some(is_eagle), }, ); let sequence_hashes = compute_seq_hash_for_block(&block_hashes); diff --git a/lib/kv-router/src/services/selection/mod.rs b/lib/kv-router/src/services/selection/mod.rs index a13fbcbf261c..bd4a37da8c36 100644 --- a/lib/kv-router/src/services/selection/mod.rs +++ b/lib/kv-router/src/services/selection/mod.rs @@ -20,8 +20,8 @@ mod tests; pub use core::{SelectionCore, SelectionServiceConfig}; pub use server::{AppState, create_router, run_server}; pub use types::{ - ModelLoadResponse, OverlapScoresRequest, OverlapScoresResponse, PotentialLoadsRequest, - ReadyResponse, ReservationRequest, ReservationResponse, SelectAndReserveRequest, SelectRequest, - SelectResponse, SelectionKey, SelectionWorkerConfig, WorkerCatalogRecord, WorkerLifecycle, - WorkerOverlapScore, WorkerPatchRequest, WorkerRequest, + ModelLoadResponse, OutputBlockRequest, OverlapScoresRequest, OverlapScoresResponse, + PotentialLoadsRequest, ReadyResponse, ReservationRequest, ReservationResponse, + SelectAndReserveRequest, SelectRequest, SelectResponse, SelectionKey, SelectionWorkerConfig, + WorkerCatalogRecord, WorkerLifecycle, WorkerOverlapScore, WorkerPatchRequest, WorkerRequest, }; diff --git a/lib/kv-router/src/services/selection/scoring.rs b/lib/kv-router/src/services/selection/scoring.rs index 664e7b0423e2..a019dec2c125 100644 --- a/lib/kv-router/src/services/selection/scoring.rs +++ b/lib/kv-router/src/services/selection/scoring.rs @@ -6,8 +6,9 @@ use std::collections::{HashMap, HashSet}; use crate::indexer::TieredMatchDetails; use crate::protocols::{StorageTier, WorkerWithDpRank}; use crate::scheduling::TierOverlapBlocks; +use crate::scheduling::config::RouterConfigOverride; -use super::types::{OverlapScoresResponse, WorkerOverlapScore}; +use super::types::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; #[derive(Default)] pub(super) struct OverlapInputs { @@ -102,10 +103,13 @@ pub(super) fn tier_overlap_blocks_from_tiered_matches( pub(super) fn build_overlap_scores_response( config: &crate::config::KvRouterConfig, + config_override: Option<&RouterConfigOverride>, tiered: &TieredMatchDetails, block_size: u32, + schedulable_workers: impl IntoIterator, ) -> OverlapScoresResponse { let mut all_workers = HashSet::new(); + all_workers.extend(schedulable_workers); for worker in tiered.device.overlap_scores.scores.keys() { all_workers.insert(*worker); } @@ -118,6 +122,9 @@ pub(super) fn build_overlap_scores_response( let host = tiered.lower_tier.get(&StorageTier::HostPinned); let disk = tiered.lower_tier.get(&StorageTier::Disk); let external = tiered.lower_tier.get(&StorageTier::External); + let overlap_score_credit = config_override + .and_then(|cfg| cfg.overlap_score_credit) + .unwrap_or(config.overlap_score_credit); let mut workers: Vec<_> = all_workers .into_iter() @@ -141,7 +148,9 @@ pub(super) fn build_overlap_scores_response( .and_then(|matches| matches.hits.get(&worker)) .copied() .unwrap_or(0); - let router_credit_blocks = device_blocks as f64 + let host_pinned_blocks = device_blocks + host_pinned_extension_blocks; + let disk_blocks = host_pinned_blocks + disk_extension_blocks; + let router_credit_blocks = overlap_score_credit * device_blocks as f64 + config.host_cache_hit_weight * host_pinned_extension_blocks as f64 + config.disk_cache_hit_weight * disk_extension_blocks as f64; @@ -149,8 +158,11 @@ pub(super) fn build_overlap_scores_response( worker_id: worker.worker_id, dp_rank: worker.dp_rank, device_blocks, + host_pinned_blocks, + disk_blocks, host_pinned_extension_blocks, disk_extension_blocks, + shared_beyond_device_blocks: None, router_credit_blocks, } }) @@ -161,5 +173,11 @@ pub(super) fn build_overlap_scores_response( block_size, num_blocks: tiered.device.overlap_scores.frequencies.len(), workers, + shared_cache: SharedCacheOverlapScore { + enabled: false, + total_hit_blocks: 0, + ranges: Vec::new(), + error: None, + }, } } diff --git a/lib/kv-router/src/services/selection/server.rs b/lib/kv-router/src/services/selection/server.rs index e9479221bfd4..a79b7a5e35bf 100644 --- a/lib/kv-router/src/services/selection/server.rs +++ b/lib/kv-router/src/services/selection/server.rs @@ -18,8 +18,8 @@ use crate::protocols::WorkerId; use super::core::{SelectionCore, SelectionServiceConfig}; use super::types::{ - OverlapScoresRequest, PotentialLoadsRequest, REQUEST_BODY_LIMIT_BYTES, ReservationRequest, - SelectAndReserveRequest, SelectRequest, WorkerPatchRequest, WorkerRequest, + OutputBlockRequest, OverlapScoresRequest, PotentialLoadsRequest, REQUEST_BODY_LIMIT_BYTES, + ReservationRequest, SelectAndReserveRequest, SelectRequest, WorkerPatchRequest, WorkerRequest, }; #[derive(Debug, Deserialize)] @@ -145,6 +145,24 @@ async fn delete_reservation( } } +async fn add_output_block( + State(state): State>, + Path(reservation_id): Path, + payload: Result, JsonRejection>, +) -> Response { + let Json(req) = match payload { + Ok(payload) => payload, + Err(error) => return json_rejection(error), + }; + match state + .core + .add_output_block(&reservation_id, req.decay_fraction) + { + Ok(()) => json_ok(StatusCode::OK), + Err(error) => error.into_response(), + } +} + async fn health() -> Response { json_ok(StatusCode::OK) } @@ -228,6 +246,10 @@ pub fn create_router(state: Arc) -> Router { "/reservations/{reservation_id}/prefill_complete", post(prefill_complete), ) + .route( + "/reservations/{reservation_id}/output_block", + post(add_output_block), + ) .route("/reservations/{reservation_id}", delete(delete_reservation)) .route("/workers", post(create_worker).get(list_workers)) .route( diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index b383389645fb..8cd616d7271c 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -9,13 +9,23 @@ use std::sync::Arc; use tokio_util::sync::CancellationToken; use tower::ServiceExt; +use crate::indexer::{LowerTierMatchDetails, MatchDetails, TieredMatchDetails}; +use crate::protocols::{ + BlockExtraInfo, BlockHashOptions, BlockMmObjectInfo, OverlapScores, StorageTier, + WorkerWithDpRank, compute_block_hash_for_seq, compute_seq_hash_for_block, +}; +use crate::scheduling::config::RouterConfigOverride; + +use super::input::{MmRoutingInfoRequest, PromptRequest}; +use super::scoring::build_overlap_scores_response; use super::*; fn test_config() -> crate::config::KvRouterConfig { - let mut config = crate::config::KvRouterConfig::default(); - config.use_kv_events = false; - config.router_queue_threshold = None; - config + crate::config::KvRouterConfig { + use_kv_events: false, + router_queue_threshold: None, + ..Default::default() + } } fn app() -> Router { @@ -64,6 +74,95 @@ async fn register_worker_id(app: Router, worker_id: u64, max_tokens: Option post(app, "/workers", &body.to_string()).await } +#[test] +fn prompt_normalization_uses_mm_routing_info_and_eagle_hashing() { + let mm_infos = vec![Some(BlockExtraInfo { + mm_objects: vec![BlockMmObjectInfo { + mm_hash: 42, + offsets: vec![(0, 2)], + }], + })]; + let request = PromptRequest { + token_ids: Some(vec![1, 2, 3, 4]), + mm_routing_info: Some(MmRoutingInfoRequest { + routing_token_ids: vec![10, 11, 12, 13, 14, 15, 16, 17], + block_mm_infos: mm_infos.clone(), + }), + block_mm_infos: None, + block_hashes: None, + sequence_hashes: None, + isl_tokens: None, + lora_name: Some("adapter".to_string()), + is_eagle: Some(true), + }; + + let normalized = request + .normalize_for_selection(4, false) + .expect("normalize prompt"); + let expected_block_hashes = compute_block_hash_for_seq( + &[10, 11, 12, 13, 14, 15, 16, 17], + 4, + BlockHashOptions { + block_mm_infos: Some(&mm_infos), + lora_name: Some("adapter"), + is_eagle: Some(true), + }, + ); + assert_eq!(normalized.block_hashes, expected_block_hashes); + assert_eq!( + normalized.sequence_hashes, + compute_seq_hash_for_block(&expected_block_hashes) + ); + assert_eq!(normalized.isl_tokens, 8); +} + +#[test] +fn overlap_scores_response_honors_override_and_includes_python_shape_fields() { + let worker = WorkerWithDpRank::new(1, 0); + let idle_worker = WorkerWithDpRank::new(2, 0); + let mut device_scores = OverlapScores::new(); + device_scores.scores.insert(worker, 2); + device_scores.frequencies = vec![1, 1]; + let mut host = LowerTierMatchDetails::default(); + host.hits.insert(worker, 1); + let mut tiered = TieredMatchDetails { + device: MatchDetails { + overlap_scores: device_scores, + last_matched_hashes: Default::default(), + }, + lower_tier: Default::default(), + }; + tiered.lower_tier.insert(StorageTier::HostPinned, host); + + let mut config = test_config(); + config.host_cache_hit_weight = 0.75; + let override_config = RouterConfigOverride { + overlap_score_credit: Some(0.5), + ..Default::default() + }; + let response = build_overlap_scores_response( + &config, + Some(&override_config), + &tiered, + 4, + [worker, idle_worker], + ); + + assert_eq!(response.workers.len(), 2); + let selected = response + .workers + .iter() + .find(|row| row.worker_id == 1) + .expect("worker row"); + assert_eq!(selected.device_blocks, 2); + assert_eq!(selected.host_pinned_blocks, 3); + assert_eq!(selected.disk_blocks, 3); + assert_eq!(selected.host_pinned_extension_blocks, 1); + assert_eq!(selected.shared_beyond_device_blocks, None); + assert_eq!(selected.router_credit_blocks, 1.75); + assert!(!response.shared_cache.enabled); +} + #[tokio::test] async fn ready_and_select_report_not_ready_without_schedulable_workers() { let app = app(); @@ -153,6 +252,52 @@ async fn select_echoes_selection_id_and_does_not_book_load() { assert_eq!(loads[0]["loads"][0]["active_requests"], 0); } +#[tokio::test] +async fn overlap_scores_returns_all_schedulable_worker_ranks() { + let app = app(); + assert_eq!( + register_worker(app.clone(), None).await.status(), + StatusCode::CREATED + ); + let response = post( + app.clone(), + "/workers", + r#"{ + "worker_id": 2, + "model_name": "model", + "endpoint": "http://worker-2:8000", + "block_size": 4, + "data_parallel_start_rank": 2, + "data_parallel_size": 2 + }"#, + ) + .await; + assert_eq!(response.status(), StatusCode::CREATED); + + let response = post( + app, + "/overlap_scores", + r#"{"model_name":"model","token_ids":[1,2,3,4]}"#, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + let workers = body["workers"].as_array().expect("workers array"); + let ids: Vec<_> = workers + .iter() + .map(|row| { + ( + row["worker_id"].as_u64().unwrap(), + row["dp_rank"].as_u64().unwrap(), + ) + }) + .collect(); + assert_eq!(ids, vec![(1, 0), (2, 2), (2, 3)]); + assert_eq!(workers[0]["host_pinned_blocks"], 0); + assert_eq!(workers[0]["disk_blocks"], 0); + assert_eq!(body["shared_cache"]["enabled"], false); +} + #[tokio::test] async fn select_and_reserve_books_and_duplicate_reservation_conflicts() { let app = app(); @@ -192,6 +337,69 @@ async fn select_and_reserve_books_and_duplicate_reservation_conflicts() { assert_eq!(free.status(), StatusCode::OK); } +#[tokio::test] +async fn output_block_endpoint_updates_reserved_load() { + let app = app(); + assert_eq!( + register_worker(app.clone(), None).await.status(), + StatusCode::CREATED + ); + + let response = post( + app.clone(), + "/select_and_reserve", + r#"{"model_name":"model","token_ids":[1,2,3,4],"reservation_id":"res-output"}"#, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + + let loads_before = app + .clone() + .oneshot( + Request::builder() + .uri("/loads") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let before = response_json(loads_before).await; + let before_blocks = before[0]["loads"][0]["potential_decode_blocks"] + .as_u64() + .unwrap(); + + let response = post( + app.clone(), + "/reservations/res-output/output_block", + r#"{}"#, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + + let invalid = post( + app.clone(), + "/reservations/res-output/output_block", + r#"{"decay_fraction":1.5}"#, + ) + .await; + assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); + + let loads_after = app + .oneshot( + Request::builder() + .uri("/loads") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let after = response_json(loads_after).await; + let after_blocks = after[0]["loads"][0]["potential_decode_blocks"] + .as_u64() + .unwrap(); + assert_eq!(after_blocks, before_blocks + 1); +} + #[tokio::test] async fn explicit_reservation_books_after_select() { let app = app(); diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 7d49a08fcaaa..1701bf872ea3 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -73,6 +73,7 @@ pub struct SelectionWorkerConfig { pub max_num_batched_tokens: Option, pub total_kv_blocks: Option, pub stable_routing_id: Option, + pub is_eagle: Option, #[serde(default)] pub taints: HashSet, #[serde(default)] @@ -141,6 +142,7 @@ pub struct WorkerCatalogRecord { pub max_num_batched_tokens: Option, pub total_kv_blocks: Option, pub stable_routing_id: Option, + pub is_eagle: Option, #[serde(default)] pub taints: HashSet, #[serde(default)] @@ -169,6 +171,7 @@ impl WorkerCatalogRecord { max_num_batched_tokens: req.max_num_batched_tokens, total_kv_blocks: req.total_kv_blocks, stable_routing_id: req.stable_routing_id, + is_eagle: req.is_eagle, taints: req.taints, topology_domains: req.topology_domains, kv_transfer_domain: req.kv_transfer_domain, @@ -204,6 +207,7 @@ impl WorkerCatalogRecord { max_num_batched_tokens: self.max_num_batched_tokens, total_kv_blocks: self.total_kv_blocks, stable_routing_id: self.stable_routing_id.clone(), + is_eagle: self.is_eagle, taints: self.taints.clone(), topology_domains: self.topology_domains.clone(), kv_transfer_domain: self.kv_transfer_domain.clone(), @@ -287,6 +291,8 @@ pub struct WorkerRequest { #[serde(default)] pub stable_routing_id: Option, #[serde(default)] + pub is_eagle: Option, + #[serde(default)] pub taints: HashSet, #[serde(default)] pub topology_domains: HashMap, @@ -321,6 +327,8 @@ pub struct WorkerPatchRequest { #[serde(default)] pub stable_routing_id: Option, #[serde(default)] + pub is_eagle: Option, + #[serde(default)] pub taints: Option>, #[serde(default)] pub topology_domains: Option>, @@ -364,6 +372,9 @@ impl WorkerCatalogRecord { if patch.stable_routing_id.is_some() { self.stable_routing_id = patch.stable_routing_id; } + if patch.is_eagle.is_some() { + self.is_eagle = patch.is_eagle; + } if let Some(taints) = patch.taints { self.taints = taints; } @@ -450,6 +461,12 @@ pub struct ReservationRequest { pub expected_output_tokens: Option, } +#[derive(Debug, Deserialize)] +pub struct OutputBlockRequest { + #[serde(default)] + pub decay_fraction: Option, +} + #[derive(Debug, Deserialize)] pub struct PotentialLoadsRequest { #[serde(default = "default_model_name")] @@ -470,6 +487,8 @@ pub struct OverlapScoresRequest { pub tenant_id: String, #[serde(flatten)] pub prompt: PromptRequest, + #[serde(default)] + pub router_config_override: Option, } #[derive(Debug, Serialize)] @@ -519,14 +538,26 @@ pub struct WorkerOverlapScore { pub worker_id: WorkerId, pub dp_rank: DpRank, pub device_blocks: usize, + pub host_pinned_blocks: usize, + pub disk_blocks: usize, pub host_pinned_extension_blocks: usize, pub disk_extension_blocks: usize, + pub shared_beyond_device_blocks: Option, pub router_credit_blocks: f64, } +#[derive(Debug, Serialize)] +pub struct SharedCacheOverlapScore { + pub enabled: bool, + pub total_hit_blocks: u32, + pub ranges: Vec<(u32, u32)>, + pub error: Option, +} + #[derive(Debug, Serialize)] pub struct OverlapScoresResponse { pub block_size: u32, pub num_blocks: usize, pub workers: Vec, + pub shared_cache: SharedCacheOverlapScore, } diff --git a/lib/kv-router/src/services/zmq.rs b/lib/kv-router/src/services/zmq.rs index dc75789645e6..05cb49ed384b 100644 --- a/lib/kv-router/src/services/zmq.rs +++ b/lib/kv-router/src/services/zmq.rs @@ -63,7 +63,7 @@ impl ZmqSocket { Ok(()) } - #[cfg(any(feature = "standalone-slot-tracker", test))] + #[cfg(feature = "standalone-slot-tracker")] pub(crate) fn disconnect(&self, endpoint: &str) -> Result<()> { self.socket().disconnect(endpoint)?; Ok(()) From f5bfec23563246796f98c857aa10354bf8443b41 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 11 Jun 2026 23:09:48 -0700 Subject: [PATCH 3/4] fix(ci): include selection service in runtime wheel Signed-off-by: PeaBrane --- container/templates/wheel_builder.Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/container/templates/wheel_builder.Dockerfile b/container/templates/wheel_builder.Dockerfile index 7d04072feb96..2e07cb93eb18 100644 --- a/container/templates/wheel_builder.Dockerfile +++ b/container/templates/wheel_builder.Dockerfile @@ -488,9 +488,9 @@ RUN --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token uv build --wheel --out-dir /opt/dynamo/dist && \ cd /opt/dynamo/lib/bindings/python && \ if [ "$ENABLE_MEDIA_FFMPEG" = "true" ]; then \ - maturin build --release --features "media-ffmpeg,kv-indexer,slot-tracker,mm-routing,aic-forward-pass" --out /opt/dynamo/dist; \ + maturin build --release --features "media-ffmpeg,kv-indexer,slot-tracker,select-service,mm-routing,aic-forward-pass" --out /opt/dynamo/dist; \ else \ - maturin build --release --features "kv-indexer,slot-tracker,mm-routing,aic-forward-pass" --out /opt/dynamo/dist; \ + maturin build --release --features "kv-indexer,slot-tracker,select-service,mm-routing,aic-forward-pass" --out /opt/dynamo/dist; \ fi && \ /tmp/use-sccache.sh show-stats "Dynamo Runtime" From ba0ed3f24f65d2554522340789005a45c7083348 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Thu, 11 Jun 2026 23:40:41 -0700 Subject: [PATCH 4/4] fix(bindings): import select service router config Signed-off-by: PeaBrane --- lib/bindings/python/rust/llm/kv.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index a98623d64d0a..3a6c1a6094ff 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -18,6 +18,8 @@ use crate::Endpoint; feature = "select-service" ))] use clap::Parser; +#[cfg(feature = "select-service")] +use dynamo_kv_router::config::kv_router_config_from_dynamo_env; use dynamo_kv_router::config::{KvRouterConfig, RouterConfigOverride}; use dynamo_kv_router::protocols::compute_block_hash_for_seq; use dynamo_kv_router::protocols::*;