perf(kv-router): parse selection requests without serde flatten or path tracking

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<T: serde::de::DeserializeOwned>(
+    request: &Bound<'_, PyAny>,
+    prompt: impl FnOnce(&mut T) -> &mut PromptRequest,
+) -> PyResult<T> {
+    let invalid = |e: pythonize::PythonizeError| PyValueError::new_err(e.to_string());
+    let Ok(dict) = request.downcast::<pyo3::types::PyDict>() 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::<Vec<u32>>()
+        .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<Bound<'p, PyAny>> {
-        let req: OverlapScoresRequest =
-            depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
+        let req = depythonize_prompt_request::<OverlapScoresRequest>(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<Bound<'p, PyAny>> {
-        let req: SelectRequest =
-            depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
+        let req =
+            depythonize_prompt_request::<SelectRequest>(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<Bound<'p, PyAny>> {
-        let req: SelectAndReserveRequest =
-            depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
+        let req = depythonize_prompt_request::<SelectAndReserveRequest>(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<Bound<'p, PyAny>> {
-        let req: ReservationRequest =
-            depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
+        let req = depythonize_prompt_request::<ReservationRequest>(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<Bound<'p, PyAny>> {
-        let req: PotentialLoadsRequest =
-            depythonize(request.bind(py)).map_err(|e| PyValueError::new_err(e.to_string()))?;
+        let req = depythonize_prompt_request::<PotentialLoadsRequest>(request.bind(py), |req| {
+            &mut req.prompt
+        })?;
         let core = self.inner.clone();
         crate::future_into_py(py, async move {
             let resp = core

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()
 

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<Arc<AppState>>,
     headers: HeaderMap,
-    payload: Result<Json<SelectRequest>, JsonRejection>,
+    PromptJson(req): PromptJson<SelectRequest>,
 ) -> 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<Arc<AppState>>,
     headers: HeaderMap,
-    payload: Result<Json<SelectAndReserveRequest>, JsonRejection>,
+    PromptJson(req): PromptJson<SelectAndReserveRequest>,
 ) -> 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<Arc<AppState>>,
-    payload: Result<Json<ReservationRequest>, JsonRejection>,
+    PromptJson(req): PromptJson<ReservationRequest>,
 ) -> 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<Arc<AppState>>, Query(params): Query<FilterQu
 
 async fn potential_loads(
     State(state): State<Arc<AppState>>,
-    payload: Result<Json<PotentialLoadsRequest>, JsonRejection>,
+    PromptJson(req): PromptJson<PotentialLoadsRequest>,
 ) -> 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<Arc<AppState>>,
-    payload: Result<Json<OverlapScoresRequest>, JsonRejection>,
+    PromptJson(req): PromptJson<OverlapScoresRequest>,
 ) -> 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<Arc<AppState>>) -> 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>(T);
+
+impl<S: Send + Sync, T: DeserializeOwned> FromRequest<S> for PromptJson<T> {
+    type Rejection = Response;
+
+    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
+        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")
 }

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();