diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index f987e209abf0..4fea01a3b6ba 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -32,10 +32,11 @@ use dynamo_kv_router::protocols::*; use dynamo_kv_router::services::indexer::{self, IndexerConfig}; #[cfg(feature = "select-service")] use dynamo_kv_router::services::selection::{ - self, OverlapScoresRequest, PotentialLoadsRequest, ReservationRequest, SelectAndReserveRequest, - SelectRequest, SelectionCacheConfig as RsSelectionCacheConfig, SelectionError, - SelectionService as RustSelectionService, SelectionServiceBuilder, SelectionServiceConfig, - WorkerPatchRequest, WorkerRequest, warn_for_unserved_worker_selection_policies, + self, OverlapScoresRequest, PotentialLoadsRequest, PromptRequest, ReservationRequest, + SelectAndReserveRequest, SelectRequest, SelectionCacheConfig as RsSelectionCacheConfig, + SelectionError, SelectionService as RustSelectionService, SelectionServiceBuilder, + SelectionServiceConfig, WorkerPatchRequest, WorkerRequest, + warn_for_unserved_worker_selection_policies, }; #[cfg(feature = "slot-tracker")] use dynamo_kv_router::services::slot_tracker::{self, SlotTrackerConfig}; @@ -571,6 +572,34 @@ where rt.block_on(selection::run_server(config.port, service)) } +/// Depythonize a prompt-carrying request dict. `token_ids` is extracted with pyo3's typed +/// list conversion: pythonize's per-element serde path costs several times more per token, +/// all of it on the calling thread under the GIL. +#[cfg(feature = "select-service")] +fn depythonize_prompt_request( + request: &Bound<'_, PyAny>, + prompt: impl FnOnce(&mut T) -> &mut PromptRequest, +) -> PyResult { + let invalid = |e: pythonize::PythonizeError| PyValueError::new_err(e.to_string()); + let Ok(dict) = request.downcast::() else { + return depythonize(request).map_err(invalid); + }; + let Some(tokens) = dict + .get_item("token_ids")? + .filter(|tokens| !tokens.is_none()) + else { + return depythonize(request).map_err(invalid); + }; + let token_ids = tokens + .extract::>() + .map_err(|e| PyValueError::new_err(format!("token_ids: {e}")))?; + let rest = dict.copy()?; + rest.del_item("token_ids")?; + let mut req: T = depythonize(rest.as_any()).map_err(invalid)?; + prompt(&mut req).token_ids = Some(token_ids); + Ok(req) +} + /// Map a [`SelectionError`] to a Python exception: invalid input becomes a /// `ValueError`, anything else a `SelectionServiceError` carrying `kind` and /// `status_code`. @@ -793,8 +822,9 @@ impl SelectionService { /// Per-worker KV-overlap scores for a prompt (dict, see `OverlapScoresRequest`). fn overlap_scores<'p>(&self, py: Python<'p>, request: PyObject) -> PyResult> { - let req: OverlapScoresRequest = - depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; + let req = depythonize_prompt_request::(request.bind(py), |req| { + &mut req.prompt + })?; let core = self.inner.clone(); crate::future_into_py(py, async move { let resp = core.overlap_scores(req).await.map_err(selection_to_pyerr)?; @@ -804,8 +834,8 @@ impl SelectionService { /// Select the best worker by KV-overlap + load, no booking (dict, see `SelectRequest`). fn select<'p>(&self, py: Python<'p>, request: PyObject) -> PyResult> { - let req: SelectRequest = - depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; + let req = + depythonize_prompt_request::(request.bind(py), |req| &mut req.prompt)?; let core = self.inner.clone(); crate::future_into_py(py, async move { let resp = core.select(req).await.map_err(selection_to_pyerr)?; @@ -819,8 +849,9 @@ impl SelectionService { py: Python<'p>, request: PyObject, ) -> PyResult> { - let req: SelectAndReserveRequest = - depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; + let req = depythonize_prompt_request::(request.bind(py), |req| { + &mut req.prompt + })?; let core = self.inner.clone(); crate::future_into_py(py, async move { let resp = core @@ -838,8 +869,9 @@ impl SelectionService { py: Python<'p>, request: PyObject, ) -> PyResult> { - let req: ReservationRequest = - depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; + let req = depythonize_prompt_request::(request.bind(py), |req| { + &mut req.prompt + })?; let core = self.inner.clone(); crate::future_into_py(py, async move { let resp = core @@ -904,8 +936,9 @@ impl SelectionService { /// Per-worker potential loads for a prompt, without booking (dict, see `PotentialLoadsRequest`). fn potential_loads<'p>(&self, py: Python<'p>, request: PyObject) -> PyResult> { - let req: PotentialLoadsRequest = - depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?; + let req = depythonize_prompt_request::(request.bind(py), |req| { + &mut req.prompt + })?; let core = self.inner.clone(); crate::future_into_py(py, async move { let resp = core diff --git a/lib/bindings/python/tests/test_kv_bindings.py b/lib/bindings/python/tests/test_kv_bindings.py index e7643d8baf61..9d904a2f98f6 100644 --- a/lib/bindings/python/tests/test_kv_bindings.py +++ b/lib/bindings/python/tests/test_kv_bindings.py @@ -305,6 +305,9 @@ async def test_selection_service_malformed_payload_raises_value_error(monkeypatc # Malformed input surfaces as a plain ValueError, not a SelectionServiceError. with pytest.raises(ValueError): await service.upsert_worker({"model_name": "model"}) + # Prompt booleans are not coerced by truthiness. + with pytest.raises(ValueError, match="boolean"): + await service.select({"token_ids": [1, 2, 3, 4], "is_eagle": "false"}) finally: service.shutdown() diff --git a/lib/kv-router/src/services/selection/server.rs b/lib/kv-router/src/services/selection/server.rs index cfb6af7ab39f..e1d3093be759 100644 --- a/lib/kv-router/src/services/selection/server.rs +++ b/lib/kv-router/src/services/selection/server.rs @@ -3,13 +3,15 @@ use std::sync::Arc; +use axum::body::Bytes; use axum::extract::rejection::JsonRejection; -use axum::extract::{Path, Query, State}; -use axum::http::{HeaderMap, StatusCode}; +use axum::extract::{FromRequest, Path, Query, Request, State}; +use axum::http::{HeaderMap, StatusCode, header}; use axum::response::{IntoResponse, Response}; use axum::routing::{delete, get, patch, post}; use axum::{Json, Router}; use serde::Deserialize; +use serde::de::DeserializeOwned; use tokio::net::TcpListener; use crate::protocols::WorkerId; @@ -85,12 +87,8 @@ async fn list_workers( async fn select( State(state): State>, headers: HeaderMap, - payload: Result, JsonRejection>, + PromptJson(req): PromptJson, ) -> Response { - let Json(req) = match payload { - Ok(payload) => payload, - Err(error) => return json_rejection(error), - }; match state .service .select_with_policy_class(req, policy_class_from_headers(&headers)) @@ -104,12 +102,8 @@ async fn select( async fn select_and_reserve( State(state): State>, headers: HeaderMap, - payload: Result, JsonRejection>, + PromptJson(req): PromptJson, ) -> Response { - let Json(req) = match payload { - Ok(payload) => payload, - Err(error) => return json_rejection(error), - }; match state .service .select_and_reserve_with_policy_class(req, policy_class_from_headers(&headers)) @@ -122,12 +116,8 @@ async fn select_and_reserve( async fn create_reservation( State(state): State>, - payload: Result, JsonRejection>, + PromptJson(req): PromptJson, ) -> Response { - let Json(req) = match payload { - Ok(payload) => payload, - Err(error) => return json_rejection(error), - }; match state.service.create_reservation(req).await { Ok(response) => (StatusCode::CREATED, Json(response)).into_response(), Err(error) => error.into_response(), @@ -195,12 +185,8 @@ async fn loads(State(state): State>, Query(params): Query>, - payload: Result, JsonRejection>, + PromptJson(req): PromptJson, ) -> Response { - let Json(req) = match payload { - Ok(payload) => payload, - Err(error) => return json_rejection(error), - }; match state.service.potential_loads(req).await { Ok(response) => Json(response).into_response(), Err(error) => error.into_response(), @@ -209,12 +195,8 @@ async fn potential_loads( async fn overlap_scores( State(state): State>, - payload: Result, JsonRejection>, + PromptJson(req): PromptJson, ) -> Response { - let Json(req) = match payload { - Ok(payload) => payload, - Err(error) => return json_rejection(error), - }; match state.service.overlap_scores(req).await { Ok(response) => Json(response).into_response(), Err(error) => error.into_response(), @@ -225,6 +207,58 @@ async fn dump_events(State(state): State>) -> Response { Json(state.service.indexer_snapshot().await).into_response() } +/// axum's `Json` for prompt-carrying bodies, without `serde_path_to_error`: `Json` +/// tracks the path of every element it deserializes, which dominates parsing a large +/// `token_ids` array. Rejections keep `Json`'s statuses: 415 for a missing or non-JSON +/// `Content-Type`, 400 for malformed JSON, 422 for a body of the wrong shape. +struct PromptJson(T); + +impl FromRequest for PromptJson { + type Rejection = Response; + + async fn from_request(req: Request, state: &S) -> Result { + if !has_json_content_type(req.headers()) { + return Err(json_error( + StatusCode::UNSUPPORTED_MEDIA_TYPE, + "Expected request with `Content-Type: application/json`", + )); + } + let body = Bytes::from_request(req, state) + .await + .map_err(|rejection| json_error(rejection.status(), rejection.body_text()))?; + serde_json::from_slice(&body).map(Self).map_err(|error| { + if error.is_data() { + json_error( + StatusCode::UNPROCESSABLE_ENTITY, + format!("Failed to deserialize the JSON body into the target type: {error}"), + ) + } else { + json_error( + StatusCode::BAD_REQUEST, + format!("Failed to parse the request body as JSON: {error}"), + ) + } + }) + } +} + +/// `application/json` or an `application/*+json` subtype, as axum's `Json` accepts. +fn has_json_content_type(headers: &HeaderMap) -> bool { + let essence = headers + .get(header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(';').next()) + .unwrap_or_default() + .trim(); + let Some((kind, subtype)) = essence.split_once('/') else { + return false; + }; + let suffix = subtype + .rsplit_once('+') + .map_or(subtype, |(_, suffix)| suffix); + kind.eq_ignore_ascii_case("application") && suffix.eq_ignore_ascii_case("json") +} + async fn not_found() -> Response { json_error(StatusCode::NOT_FOUND, "route not found") } diff --git a/lib/kv-router/src/services/selection/tests.rs b/lib/kv-router/src/services/selection/tests.rs index ce1d625918e2..a30e40d514ae 100644 --- a/lib/kv-router/src/services/selection/tests.rs +++ b/lib/kv-router/src/services/selection/tests.rs @@ -576,6 +576,58 @@ fn prompt_request_cache_salt_changes_normalized_hashes() { assert_ne!(salted.sequence_hashes, unsalted.sequence_hashes); } +#[test] +fn prompt_carrying_requests_read_top_level_prompt_fields() { + let body = serde_json::json!({ + "selection_id": "sel-1", + "worker_id": 7, + "token_ids": [1, 2, 3], + "cache_salt": "tenant-a", + "is_eagle": true, + }); + let assert_prompt = |prompt: &PromptRequest| { + assert_eq!(prompt.token_ids.as_deref(), Some(&[1, 2, 3][..])); + assert_eq!(prompt.cache_namespace.as_deref(), Some("tenant-a")); + assert_eq!(prompt.is_eagle, Some(true)); + }; + + let select: SelectRequest = serde_json::from_value(body.clone()).unwrap(); + assert_prompt(&select.prompt); + assert_eq!(select.selection_id.as_deref(), Some("sel-1")); + assert_eq!(select.model_name, DEFAULT_MODEL_NAME); + let reserve: SelectAndReserveRequest = serde_json::from_value(body.clone()).unwrap(); + assert_prompt(&reserve.prompt); + let booking: ReservationRequest = serde_json::from_value(body.clone()).unwrap(); + assert_prompt(&booking.prompt); + assert_eq!(booking.worker_id, Some(7)); + let loads: PotentialLoadsRequest = serde_json::from_value(body.clone()).unwrap(); + assert_prompt(&loads.prompt); + let scores: OverlapScoresRequest = serde_json::from_value(body).unwrap(); + assert_prompt(&scores.prompt); +} + +#[tokio::test] +async fn prompt_routes_keep_json_rejection_statuses() { + let untyped = app() + .oneshot( + Request::builder() + .method("POST") + .uri("/select") + .body(Body::from(r#"{"token_ids":[1]}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(untyped.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE); + + let malformed = post(app(), "/select", "{").await; + assert_eq!(malformed.status(), StatusCode::BAD_REQUEST); + + let wrong_shape = post(app(), "/select", r#"{"token_ids":"oops"}"#).await; + assert_eq!(wrong_shape.status(), StatusCode::UNPROCESSABLE_ENTITY); + assert!(response_json(wrong_shape).await["error"].is_string()); +} + #[test] fn keyed_prompt_tracking_leaves_indexer_hashes_public() { let mut key_file = NamedTempFile::new().unwrap(); diff --git a/lib/kv-router/src/services/selection/types.rs b/lib/kv-router/src/services/selection/types.rs index 26a51a93e225..212a60aec6b6 100644 --- a/lib/kv-router/src/services/selection/types.rs +++ b/lib/kv-router/src/services/selection/types.rs @@ -2,24 +2,27 @@ // SPDX-License-Identifier: Apache-2.0 use std::collections::{BTreeMap, HashMap, HashSet}; +use std::fmt; use std::sync::Arc; -use serde::{Deserialize, Serialize}; +use serde::de::{self, Visitor}; +use serde::{Deserialize, Deserializer, Serialize}; use crate::identity::{RoutingPartitionId, default_routing_group}; use crate::kv_hints::KvHint; use crate::protocols::{ - DpRank, KvHintTransferWorkerMetadata, KvTransferEnforcement, RoutingConstraints, - WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, + BlockExtraInfo, DpRank, KvHintTransferWorkerMetadata, KvTransferEnforcement, + RoutingConstraints, WorkerAffinityTarget, WorkerConfigLike, WorkerId, WorkerWithDpRank, }; use crate::scheduling::config::RouterConfigOverride; pub use crate::scheduling::{OverlapScoresResponse, SharedCacheOverlapScore, WorkerOverlapScore}; use crate::scheduling::{PotentialLoad, SessionContext, WorkerSelectionInputTrigger}; use crate::services::overlap::MooncakeOverlapSummary; -use super::input::PromptRequest; +use super::input::{MmRoutingInfoRequest, PromptRequest}; pub const DEFAULT_MODEL_NAME: &str = "default"; + pub(super) const REQUEST_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; fn default_model_name() -> String { @@ -417,14 +420,87 @@ impl WorkerCatalogRecord { } } +/// Request bodies carry `PromptRequest`'s fields at the top level. `#[serde(flatten)]` +/// would buffer the whole body through serde's `Content` first: every token id becomes a +/// boxed `Content` value, and self-describing-only inputs (pythonize) fall back to untyped +/// integer extraction. Each request type instead deserializes through a flat wire struct +/// with the prompt fields inlined. +macro_rules! inline_prompt_wire { + ($wire:ident => $target:ident { + $($(#[$meta:meta])* $field:ident: $ty:ty,)* + }) => { + #[derive(Deserialize)] + struct $wire { + $($(#[$meta])* $field: $ty,)* + token_ids: Option>, + mm_routing_info: Option, + block_mm_infos: Option>>, + block_hashes: Option>, + sequence_hashes: Option>, + isl_tokens: Option, + lora_name: Option, + #[serde(default, rename = "cache_salt")] + cache_namespace: Option, + #[serde(default, deserialize_with = "strict_bool")] + is_eagle: Option, + } + + impl From<$wire> for $target { + fn from(wire: $wire) -> Self { + Self { + $($field: wire.$field,)* + prompt: PromptRequest { + token_ids: wire.token_ids, + mm_routing_info: wire.mm_routing_info, + block_mm_infos: wire.block_mm_infos, + block_hashes: wire.block_hashes, + sequence_hashes: wire.sequence_hashes, + isl_tokens: wire.isl_tokens, + lora_name: wire.lora_name, + cache_namespace: wire.cache_namespace, + is_eagle: wire.is_eagle, + }, + } + } + } + }; +} + +/// A boolean or null, nothing else. A typed `bool` field would let pythonize coerce any +/// truthy object, so `"is_eagle": "false"` would enable Eagle hashing. +fn strict_bool<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + struct StrictBool; + + impl Visitor<'_> for StrictBool { + type Value = Option; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a boolean or null") + } + + fn visit_bool(self, value: bool) -> Result { + Ok(Some(value)) + } + + fn visit_unit(self) -> Result { + Ok(None) + } + + fn visit_none(self) -> Result { + Ok(None) + } + } + + deserializer.deserialize_any(StrictBool) +} + #[derive(Debug, Deserialize)] +#[serde(from = "SelectRequestWire")] pub struct SelectRequest { - #[serde(default = "default_model_name")] pub model_name: String, - #[serde(default = "default_routing_group")] pub routing_group: String, pub selection_id: Option, - #[serde(flatten)] + /// Inlined at the top level of the body. pub prompt: PromptRequest, pub router_config_override: Option, pub expected_output_tokens: Option, @@ -436,7 +512,6 @@ pub struct SelectRequest { pub affinity_target: Option, pub pinned_worker: Option, pub allowed_worker_ids: Option>, - #[serde(default)] pub routing_constraints: RoutingConstraints, /// Select from current scheduler state without queue admission. /// @@ -447,10 +522,30 @@ pub struct SelectRequest { /// A `selection_id` still caches the booking inputs for a follow-up /// `create_reservation`. Ignored on `select_and_reserve`, which always /// books. - #[serde(default)] pub advisory: bool, } +inline_prompt_wire!(SelectRequestWire => SelectRequest { + #[serde(default = "default_model_name")] + model_name: String, + #[serde(default = "default_routing_group")] + routing_group: String, + selection_id: Option, + router_config_override: Option, + expected_output_tokens: Option, + priority_jump: Option, + strict_priority: Option, + session_id: Option, + session_context: Option, + affinity_target: Option, + pinned_worker: Option, + allowed_worker_ids: Option>, + #[serde(default)] + routing_constraints: RoutingConstraints, + #[serde(default)] + advisory: bool, +}); + impl SelectRequest { pub(super) fn take_session_context(&mut self) -> Option { resolve_session_context(self.session_context.take(), self.session_id.take()) @@ -458,13 +553,12 @@ impl SelectRequest { } #[derive(Debug, Deserialize)] +#[serde(from = "SelectAndReserveRequestWire")] pub struct SelectAndReserveRequest { - #[serde(default = "default_model_name")] pub model_name: String, - #[serde(default = "default_routing_group")] pub routing_group: String, pub selection_id: Option, - #[serde(flatten)] + /// Inlined at the top level of the body. pub prompt: PromptRequest, pub router_config_override: Option, pub expected_output_tokens: Option, @@ -476,10 +570,28 @@ pub struct SelectAndReserveRequest { pub affinity_target: Option, pub pinned_worker: Option, pub allowed_worker_ids: Option>, - #[serde(default)] pub routing_constraints: RoutingConstraints, } +inline_prompt_wire!(SelectAndReserveRequestWire => SelectAndReserveRequest { + #[serde(default = "default_model_name")] + model_name: String, + #[serde(default = "default_routing_group")] + routing_group: String, + selection_id: Option, + router_config_override: Option, + expected_output_tokens: Option, + priority_jump: Option, + strict_priority: Option, + session_id: Option, + session_context: Option, + affinity_target: Option, + pinned_worker: Option, + allowed_worker_ids: Option>, + #[serde(default)] + routing_constraints: RoutingConstraints, +}); + impl SelectAndReserveRequest { pub(super) fn take_session_context(&mut self) -> Option { resolve_session_context(self.session_context.take(), self.session_id.take()) @@ -553,10 +665,9 @@ fn resolve_session_context( /// self-contained with `worker_id`. The replay books exactly what `select` captured; /// request fields other than the ids and model/routing-group are ignored. #[derive(Debug, Deserialize)] +#[serde(from = "ReservationRequestWire")] pub struct ReservationRequest { - #[serde(default = "default_model_name")] pub model_name: String, - #[serde(default = "default_routing_group")] pub routing_group: String, /// The single booking id: the cache key to replay and the scheduler request /// id the booking lands under (the `selection_id` from the matching `select`). @@ -565,7 +676,7 @@ pub struct ReservationRequest { /// without a cached select. Omit to replay the cached `selection_id`. pub worker_id: Option, pub dp_rank: Option, - #[serde(flatten)] + /// Inlined at the top level of the body. pub prompt: PromptRequest, pub router_config_override: Option, pub expected_output_tokens: Option, @@ -573,33 +684,61 @@ pub struct ReservationRequest { pub track_prefill_tokens: Option, } +inline_prompt_wire!(ReservationRequestWire => ReservationRequest { + #[serde(default = "default_model_name")] + model_name: String, + #[serde(default = "default_routing_group")] + routing_group: String, + selection_id: String, + worker_id: Option, + dp_rank: Option, + router_config_override: Option, + expected_output_tokens: Option, + effective_prefill_tokens: Option, + track_prefill_tokens: Option, +}); + #[derive(Debug, Deserialize)] pub struct OutputBlockRequest { pub decay_fraction: Option, } #[derive(Debug, Deserialize)] +#[serde(from = "PotentialLoadsRequestWire")] pub struct PotentialLoadsRequest { - #[serde(default = "default_model_name")] pub model_name: String, - #[serde(default = "default_routing_group")] pub routing_group: String, - #[serde(flatten)] + /// Inlined at the top level of the body. pub prompt: PromptRequest, pub router_config_override: Option, } +inline_prompt_wire!(PotentialLoadsRequestWire => PotentialLoadsRequest { + #[serde(default = "default_model_name")] + model_name: String, + #[serde(default = "default_routing_group")] + routing_group: String, + router_config_override: Option, +}); + #[derive(Debug, Deserialize)] +#[serde(from = "OverlapScoresRequestWire")] pub struct OverlapScoresRequest { - #[serde(default = "default_model_name")] pub model_name: String, - #[serde(default = "default_routing_group")] pub routing_group: String, - #[serde(flatten)] + /// Inlined at the top level of the body. pub prompt: PromptRequest, pub router_config_override: Option, } +inline_prompt_wire!(OverlapScoresRequestWire => OverlapScoresRequest { + #[serde(default = "default_model_name")] + model_name: String, + #[serde(default = "default_routing_group")] + routing_group: String, + router_config_override: Option, +}); + #[derive(Debug, Serialize)] pub struct SelectResponse { #[serde(skip_serializing_if = "Option::is_none")]