fix: unblock long-context requests in the rust router, and kvd bug

infera/engine/sglang/kvd_wiring.py

@@ -122,12 +122,38 @@ def _finish_wiring(args: Any, socket_path: str) -> None:
     # subprocess re-parses argv, so this does not affect the running engine;
     # `_append_sglang_hicache_argv` is what actually selects the backend.
     sa = args.server_args
+
+    def _set_metadata(field: str, value: Any) -> bool:
+        """Best-effort ``setattr`` on ServerArgs. True if the value landed.
+
+        ServerArgs.__setattr__ raises AttributeError for any public field
+        assigned after resolution ("server_args is read-only -- use
+        get_context().override(source, ...)"), and on the v0.5.17 base that
+        guard is unconditional where it used to be gated on
+        SGLANG_STRICT_CONFIG_MUTATION. Every write in this function is
+        metadata-only, so being refused must not be fatal: it would take down a
+        leg over a field the engine subprocess never reads (observed on
+        lmsysorg/sglang:v0.5.17-rocm720-mi35x, where a kvd prefill leg died at
+        startup on enable_hierarchical_cache).
+        """
+        try:
+            setattr(sa, field, value)
+            return True
+        except AttributeError as exc:
+            logger.debug(
+                "ServerArgs.%s is read-only on this SGLang (%s); skipping the "
+                "metadata sync. The engine reads these off the forwarded argv.",
+                field,
+                exc,
+            )
+            return False
+
     if hasattr(sa, "enable_hierarchical_cache") and not sa.enable_hierarchical_cache:
-        sa.enable_hierarchical_cache = True
-        logger.info("--infera-kvd-socket implies --enable-hierarchical-cache")
+        if _set_metadata("enable_hierarchical_cache", True):
+            logger.info("--infera-kvd-socket implies --enable-hierarchical-cache")
     if hasattr(sa, "hicache_storage_backend") and not sa.hicache_storage_backend:
-        sa.hicache_storage_backend = "infera-kvd"
-        logger.info("--infera-kvd-socket implies --hicache-storage-backend infera-kvd")
+        if _set_metadata("hicache_storage_backend", "infera-kvd"):
+            logger.info("--infera-kvd-socket implies --hicache-storage-backend infera-kvd")
     # PR #9 review fix P1 (prefetch_threshold silent perf failure):
     # SGLang's default prefetch_threshold is 256 tokens. The runbook on
     # MI355X documents that 64 is needed for cache_control workloads
@@ -150,12 +176,12 @@ def _finish_wiring(args: Any, socket_path: str) -> None:
         if hasattr(sa, field):
             current = getattr(sa, field)
             if current is None or current == 256:
-                setattr(sa, field, 64)
-                logger.info(
-                    "--infera-kvd-socket lowered SGLang.%s to 64 "
-                    "(cache_control workloads on short prompts)",
-                    field,
-                )
+                if _set_metadata(field, 64):
+                    logger.info(
+                        "--infera-kvd-socket lowered SGLang.%s to 64 "
+                        "(cache_control workloads on short prompts)",
+                        field,
+                    )
             else:
                 logger.info(
                     "SGLang.%s already set to %s — leaving operator value in place",

rust/router/src/handlers.rs

@@ -9,7 +9,7 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use axum::body::Bytes;
-use axum::extract::State;
+use axum::extract::{DefaultBodyLimit, State};
 use axum::response::{IntoResponse, Response};
 use axum::routing::{get, post};
 use axum::{Json, Router};
@@ -45,6 +45,10 @@ pub fn app(state: AppState) -> Router {
         .route("/v1/workers", get(workers))
         .route("/v1/models", get(models))
         .route("/metrics", get(metrics))
+        // axum's default 2 MiB cap on `Bytes` would 413 long-context prompts, but
+        // fully disabling the limit allows unbounded buffering into memory (DoS).
+        // Raise the limit enough for expected prompts; the engine still enforces `--context-length`.
+        .layer(DefaultBodyLimit::max(8 * 1024 * 1024))
         .with_state(state)
 }
 

rust/router/tests/functional.rs

@@ -12,7 +12,7 @@ use std::time::{Duration, Instant};
 
 use arc_swap::ArcSwap;
 use axum::body::{Body, Bytes};
-use axum::extract::State;
+use axum::extract::{DefaultBodyLimit, State};
 use axum::http::header::CONTENT_TYPE;
 use axum::http::{HeaderMap, StatusCode};
 use axum::response::{IntoResponse, Response};
@@ -83,6 +83,9 @@ async fn spawn_mock(status: u16, sse: bool, reply: Value) -> (String, Arc<MockSt
     let router = Router::new()
         .route("/v1/chat/completions", post(mock_handle))
         .route("/v1/completions", post(mock_handle))
+        // Stand in for a real engine, which caps a prompt by context length
+        // rather than by request bytes.
+        .layer(DefaultBodyLimit::disable())
         .with_state(state.clone());
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
     let port = listener.local_addr().unwrap().port();
@@ -152,6 +155,32 @@ async fn mixed_unary_ok() {
     assert_eq!(mock.hit_count(), 1);
 }
 
+/// A long-context prompt is a normal request, not an oversized one: axum's
+/// default 2 MiB body cap used to 413 it before it reached a worker.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn mixed_forwards_body_over_axum_default_limit() {
+    let (url, mock) = spawn_mock(200, false, json!({"answer": 42})).await;
+    let state = make_state(
+        vec![worker(json!({
+            "worker_id": "w1", "url": url, "model_name": "m", "disagg_mode": "mixed"
+        }))],
+        0,
+    );
+    let router = spawn_router(state).await;
+    let prompt = "x".repeat((2 << 20) + 4096);
+
+    let resp = client()
+        .post(format!("{router}/v1/chat/completions"))
+        .json(&json!({"model": "m", "prompt": prompt}))
+        .send()
+        .await
+        .unwrap();
+
+    assert_eq!(resp.status(), 200);
+    let hits = mock.hits.lock().unwrap();
+    assert_eq!(hits[0].body["prompt"].as_str().unwrap().len(), prompt.len());
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn mixed_round_robin_spreads_load() {
     let (url_a, a) = spawn_mock(200, false, json!({"w": "a"})).await;