fix(gateway): reject over-long prompts with 400 context_length_exceeded

e2e_test/router/test_context_length.py

@@ -0,0 +1,130 @@
+"""An input longer than the model's context window is a 400 from the gateway.
+
+Before this check, an over-long prompt on the disaggregated gRPC path reached
+the prefill engine, whose rejection came back as::
+
+    HTTP 500 {"error": {"code": "prefill_worker_failed_to_start", ...}}
+
+a server error for a client mistake, blamed on a healthy worker (#2380).
+The gateway learns the engine's window at discovery (``max_context_length``
+on the model card), tokenizes the prompt itself on the gRPC path, and can
+therefore answer ``400 context_length_exceeded`` without dispatching.
+
+The engines run with a deliberately small window so a modest prompt overruns
+it. A prompt that fits still goes through; the engine stays the arbiter of
+whether ``input + max_tokens`` fits.
+
+Usage:
+    E2E_RUNTIME=vllm pytest e2e_test/router/test_context_length.py -v
+"""
+
+from __future__ import annotations
+
+import logging
+
+import httpx
+import pytest
+from infra.constants import get_runtime
+
+logger = logging.getLogger(__name__)
+
+_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
+_MODEL_BY_ENGINE = {"tokenspeed": "Qwen/Qwen3.5-9B"}
+# Small enough that a few hundred words overrun it, large enough for the
+# chat template plus a short answer.
+_WINDOW = 2048
+_WINDOW_ARGS = {
+    "sglang": ["--context-length", str(_WINDOW)],
+    "vllm": ["--max-model-len", str(_WINDOW)],
+    "tokenspeed": ["--max-model-len", str(_WINDOW)],
+}.get(get_runtime(), [])
+# Common English words are one token each on these tokenizers; twice the
+# window's worth of them is over the window on any of them.
+_OVER_LONG_WORDS = _WINDOW * 2
+_FITS_WORDS = 64
+
+
+def _chat(gateway, model: str, words: int, *, stream: bool = False) -> httpx.Response:
+    prompt = " ".join(["hello"] * words)
+    return httpx.post(
+        f"{gateway.base_url}/v1/chat/completions",
+        json={
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": 8,
+            "stream": stream,
+        },
+        timeout=120.0,
+    )
+
+
+def _assert_context_length_rejection(resp: httpx.Response) -> None:
+    logger.info(
+        "status=%s content-type=%s body=%s",
+        resp.status_code,
+        resp.headers.get("content-type"),
+        resp.text[:300],
+    )
+    assert resp.status_code == 400, f"expected 400, got {resp.status_code}: {resp.text[:300]}"
+    content_type = resp.headers.get("content-type", "")
+    assert content_type.startswith("application/json"), (
+        f"rejection must be a JSON error, got {content_type!r}: {resp.text[:300]}"
+    )
+    error = resp.json()["error"]
+    assert error["code"] == "context_length_exceeded", resp.text
+    assert "prefill_worker_failed_to_start" not in resp.text, (
+        f"a client mistake must not be reported as a worker failure: {resp.text[:300]}"
+    )
+    assert f"{_WINDOW} tokens" in error["message"], error["message"]
+
+
+class _ContextLengthContract:
+    """Shared bodies; subclasses pin the topology via the setup marker."""
+
+    def test_over_long_prompt_is_a_400_context_length_exceeded(self, setup_backend):
+        _, model, _, gateway = setup_backend
+        _assert_context_length_rejection(_chat(gateway, model, _OVER_LONG_WORDS))
+
+    def test_over_long_streaming_prompt_is_a_json_400(self, setup_backend):
+        _, model, _, gateway = setup_backend
+        _assert_context_length_rejection(_chat(gateway, model, _OVER_LONG_WORDS, stream=True))
+
+    def test_prompt_within_the_window_is_served(self, setup_backend):
+        _, model, _, gateway = setup_backend
+        resp = _chat(gateway, model, _FITS_WORDS)
+        assert resp.status_code == 200, f"{resp.status_code}: {resp.text[:300]}"
+        assert resp.json()["choices"], resp.text
+
+    def test_workers_still_healthy_after_rejection(self, setup_backend):
+        """The rejection never reached an engine: every worker stays healthy
+        and serves the next request."""
+        _, model, _, gateway = setup_backend
+        _assert_context_length_rejection(_chat(gateway, model, _OVER_LONG_WORDS))
+        workers = gateway.list_workers()
+        assert workers, "no workers registered"
+        unhealthy = [(w.url, w.status) for w in workers if w.status.lower() != "healthy"]
+        assert not unhealthy, f"workers unhealthy after a rejected request: {unhealthy}"
+        resp = _chat(gateway, model, _FITS_WORDS)
+        assert resp.status_code == 200, f"{resp.status_code}: {resp.text[:300]}"
+
+
+@pytest.mark.engine("sglang", "vllm", "tokenspeed")
+@pytest.mark.gpu(2)
+@pytest.mark.e2e
+@pytest.mark.model(_MODEL, tokenspeed=_MODEL_BY_ENGINE["tokenspeed"])
+@pytest.mark.workers(extra_engine_args=_WINDOW_ARGS)
+@pytest.mark.parametrize(
+    "setup_backend", [pytest.param(("pd_grpc", (1, 1)), id="1p1d")], indirect=True
+)
+class TestPDContextLength(_ContextLengthContract):
+    """The disaggregated gRPC path: the case #2380 reported."""
+
+
+@pytest.mark.engine("sglang", "vllm", "tokenspeed")
+@pytest.mark.gpu(1)
+@pytest.mark.e2e
+@pytest.mark.model(_MODEL, tokenspeed=_MODEL_BY_ENGINE["tokenspeed"])
+@pytest.mark.workers(extra_engine_args=_WINDOW_ARGS)
+@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
+class TestRegularContextLength(_ContextLengthContract):
+    """The single-worker gRPC path is bounded the same way."""

model_gateway/src/routers/grpc/common/stages/context_length.rs

@@ -0,0 +1,268 @@
+//! Gateway-side context-window check: reject an input the selected engine
+//! cannot hold before it is dispatched.
+//!
+//! The gRPC path tokenizes at the gateway and ships token ids, so the count
+//! here is exactly what the engine would see. Without this check an over-long
+//! prompt travels to the engine, which rejects it, and on the disaggregated
+//! path that rejection comes back as `500 prefill_worker_failed_to_start`: a
+//! server error for a client mistake, blamed on a healthy worker (#2380).
+//!
+//! Runs once at ingress, after worker selection (the window is a property of
+//! the chosen worker's model card) and before any client is acquired.
+
+use axum::response::Response;
+use tracing::debug;
+
+use crate::routers::{
+    error,
+    grpc::context::{PreparationOutput, WorkerSelection},
+};
+
+/// Error code for an input that does not fit the model's window; matches the
+/// OpenAI API's code for the same condition.
+pub(crate) const CONTEXT_LENGTH_EXCEEDED: &str = "context_length_exceeded";
+
+/// Reject the request when its longest input exceeds the context window the
+/// selected worker(s) advertise for `model_id`.
+///
+/// Only a *known* window is enforced: a worker that never advertised one
+/// leaves the check to the engine, as before. The check is on the input
+/// alone (strictly longer than the window), the one condition every engine
+/// rejects; whether `input + max_tokens` must also fit is engine-specific
+/// (vLLM rejects, SGLang clamps) and stays with the engine.
+pub(crate) fn enforce_context_length(
+    prep: &PreparationOutput,
+    workers: &WorkerSelection,
+    model_id: &str,
+) -> Result<(), Response> {
+    let Some(limit) = selection_context_length(workers, model_id) else {
+        return Ok(());
+    };
+    let input_tokens = prep.max_input_token_count();
+    if input_tokens <= limit as usize {
+        return Ok(());
+    }
+    debug!(
+        function = "enforce_context_length",
+        input_tokens, limit, model_id, "Rejecting input longer than the model's context window"
+    );
+    Err(error::bad_request(
+        CONTEXT_LENGTH_EXCEEDED,
+        format!(
+            "This model's maximum context length is {limit} tokens. However, your request has \
+             {input_tokens} input tokens. Please reduce the length of the input."
+        ),
+    ))
+}
+
+/// Tightest context window across the selected legs, `None` if no leg
+/// advertises one. A disaggregated prompt must fit both the prefill and the
+/// decode engine, so the smaller window governs. Encode workers process media,
+/// not the token prompt, and are not consulted.
+fn selection_context_length(workers: &WorkerSelection, model_id: &str) -> Option<u32> {
+    match workers {
+        WorkerSelection::Single { worker } => worker.context_length(model_id),
+        WorkerSelection::Disaggregated {
+            prefill, decode, ..
+        } => match (
+            prefill.context_length(model_id),
+            decode.context_length(model_id),
+        ) {
+            (Some(prefill), Some(decode)) => Some(prefill.min(decode)),
+            (prefill, decode) => prefill.or(decode),
+        },
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use axum::http::StatusCode;
+    use openai_protocol::{model_card::ModelCard, worker::HealthCheckConfig};
+
+    use super::*;
+    use crate::{
+        routers::{error::HEADER_X_SMG_ERROR_CODE, grpc::context::CompletionItem},
+        worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, Worker, WorkerType},
+    };
+
+    const MODEL: &str = "ctx-test-model";
+
+    fn worker(worker_type: WorkerType, context_length: Option<u32>) -> Arc<dyn Worker> {
+        let mut card = ModelCard::new(MODEL);
+        if let Some(len) = context_length {
+            card = card.with_context_length(len);
+        }
+        Arc::new(
+            BasicWorkerBuilder::new(format!("grpc://127.0.0.1:1/{worker_type:?}"))
+                .worker_type(worker_type)
+                .connection_mode(ConnectionMode::Grpc)
+                .runtime_type(RuntimeType::Sglang)
+                .model(card)
+                .health_config(HealthCheckConfig {
+                    disable_health_check: true,
+                    ..Default::default()
+                })
+                .build(),
+        )
+    }
+
+    fn single(context_length: Option<u32>) -> WorkerSelection {
+        WorkerSelection::Single {
+            worker: worker(WorkerType::Regular, context_length),
+        }
+    }
+
+    fn pd(prefill: Option<u32>, decode: Option<u32>) -> WorkerSelection {
+        WorkerSelection::Disaggregated {
+            encode_assignments: None,
+            prefill: worker(WorkerType::Prefill, prefill),
+            decode: worker(WorkerType::Decode, decode),
+            runtime_type: RuntimeType::Sglang,
+        }
+    }
+
+    fn input(tokens: usize) -> PreparationOutput {
+        PreparationOutput::Generate {
+            original_text: None,
+            token_ids: vec![7; tokens],
+        }
+    }
+
+    fn assert_rejected(result: Result<(), Response>) {
+        let response = result.expect_err("over-long input must be rejected");
+        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+        assert_eq!(
+            response
+                .headers()
+                .get(HEADER_X_SMG_ERROR_CODE)
+                .and_then(|v| v.to_str().ok()),
+            Some(CONTEXT_LENGTH_EXCEEDED)
+        );
+    }
+
+    #[test]
+    fn input_within_the_window_passes() {
+        assert!(enforce_context_length(&input(8), &single(Some(16)), MODEL).is_ok());
+    }
+
+    /// The window is a capacity, not a strict bound: a prompt that exactly
+    /// fills it is not "longer than" it, and whether the engine can still
+    /// generate from it is the engine's call (embeddings, for one, can).
+    #[test]
+    fn input_exactly_at_the_window_passes() {
+        assert!(enforce_context_length(&input(16), &single(Some(16)), MODEL).is_ok());
+    }
+
+    #[test]
+    fn input_over_the_window_is_a_400_context_length_exceeded() {
+        assert_rejected(enforce_context_length(&input(17), &single(Some(16)), MODEL));
+    }
+
+    #[tokio::test]
+    async fn rejection_body_names_the_limit_and_the_actual_count() {
+        let response = enforce_context_length(&input(40), &single(Some(16)), MODEL)
+            .expect_err("over-long input must be rejected");
+        let bytes = axum::body::to_bytes(response.into_body(), 1 << 16)
+            .await
+            .expect("body");
+        let json: serde_json::Value = serde_json::from_slice(&bytes).expect("json body");
+        assert_eq!(json["error"]["code"], CONTEXT_LENGTH_EXCEEDED);
+        let message = json["error"]["message"].as_str().expect("message");
+        assert!(message.contains("16 tokens"), "{message}");
+        assert!(message.contains("40 input tokens"), "{message}");
+        assert!(
+            !message.contains("worker"),
+            "the client mistake must not be blamed on a worker: {message}"
+        );
+    }
+
+    /// No advertised window means no gateway opinion: the engine keeps
+    /// enforcing its own limit exactly as before this check existed.
+    #[test]
+    fn unknown_window_is_not_enforced() {
+        assert!(enforce_context_length(&input(1_000_000), &single(None), MODEL).is_ok());
+        assert!(enforce_context_length(&input(1_000_000), &pd(None, None), MODEL).is_ok());
+    }
+
+    /// A disaggregated prompt has to fit both engines; the tighter one wins.
+    #[test]
+    fn disaggregated_uses_the_tighter_leg() {
+        assert!(enforce_context_length(&input(10), &pd(Some(16), Some(16)), MODEL).is_ok());
+        assert_rejected(enforce_context_length(
+            &input(10),
+            &pd(Some(16), Some(8)),
+            MODEL,
+        ));
+        assert_rejected(enforce_context_length(
+            &input(10),
+            &pd(Some(8), Some(16)),
+            MODEL,
+        ));
+    }
+
+    #[test]
+    fn disaggregated_with_one_unknown_leg_uses_the_known_one() {
+        assert!(enforce_context_length(&input(10), &pd(None, Some(16)), MODEL).is_ok());
+        assert_rejected(enforce_context_length(
+            &input(10),
+            &pd(Some(8), None),
+            MODEL,
+        ));
+        assert_rejected(enforce_context_length(
+            &input(10),
+            &pd(None, Some(8)),
+            MODEL,
+        ));
+    }
+
+    /// Each batched prompt is its own engine request: a batch of short prompts
+    /// whose *sum* exceeds the window is fine, one long prompt anywhere in the
+    /// batch is not.
+    #[test]
+    fn batched_completion_is_bounded_per_item() {
+        let batch = |lens: &[usize]| PreparationOutput::Completion {
+            items: lens
+                .iter()
+                .map(|&len| CompletionItem {
+                    text: String::new(),
+                    token_ids: vec![7; len],
+                })
+                .collect(),
+            joined_routing_text: None,
+        };
+        assert!(enforce_context_length(&batch(&[10, 10, 10]), &single(Some(16)), MODEL).is_ok());
+        assert_rejected(enforce_context_length(
+            &batch(&[4, 17]),
+            &single(Some(16)),
+            MODEL,
+        ));
+    }
+
+    /// The window is looked up for the model actually requested, so a worker
+    /// serving several cards with different windows is bounded by the right one.
+    #[test]
+    fn window_is_resolved_for_the_requested_model_card() {
+        let worker: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("grpc://127.0.0.1:1/multi")
+                .worker_type(WorkerType::Regular)
+                .connection_mode(ConnectionMode::Grpc)
+                .runtime_type(RuntimeType::Sglang)
+                .models(vec![
+                    ModelCard::new("wide").with_context_length(64),
+                    ModelCard::new("narrow").with_context_length(8),
+                ])
+                .health_config(HealthCheckConfig {
+                    disable_health_check: true,
+                    ..Default::default()
+                })
+                .build(),
+        );
+        let selection = WorkerSelection::Single { worker };
+        assert!(enforce_context_length(&input(10), &selection, "wide").is_ok());
+        assert_rejected(enforce_context_length(&input(10), &selection, "narrow"));
+        // An id no card matches falls back to the primary (first) card.
+        assert!(enforce_context_length(&input(10), &selection, "unlisted").is_ok());
+    }
+}