[router] Tokenize prompt once at ingress; forward input_ids to the engine (all policies) (#28744)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
c3bae61e16
commit
364bf976be
@@ -0,0 +1,210 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! End-to-end at the HTTP layer: the router tokenizes the prompt once at
|
||||
//! ingress and forwards the ids to the engine as `input_ids` (so the engine
|
||||
//! skips re-tokenizing the same prompt). Asserts the gating contract through
|
||||
//! the real chat handler + a MockWorker backend:
|
||||
//!
|
||||
//! * A plain text chat request on the engine-equivalent chat-encoder path →
|
||||
//! the forwarded body carries `input_ids` AND retains `messages`.
|
||||
//! * A request carrying `tools` → `input_ids` omitted (the router's encoder
|
||||
//! doesn't render tool schemas, so its ids would diverge from the engine).
|
||||
//! * A request with multimodal (array) content → `input_ids` omitted (a text
|
||||
//! tokenizer can't represent image content).
|
||||
//!
|
||||
//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches
|
||||
//! the built-in V4 chat encoder — the engine-equivalent path — without a
|
||||
//! template fixture.
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
|
||||
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry;
|
||||
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
|
||||
use sgl_router::proxy::Proxy;
|
||||
use sgl_router::server::app::build_router;
|
||||
use sgl_router::server::app_context::AppContext;
|
||||
use sgl_router::tokenizer::TokenizerRegistry;
|
||||
use sgl_router::workers::WorkerRegistry;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use crate::common::mock_worker::MockWorker;
|
||||
|
||||
const MODEL: &str = "deepseek-v4-tiny";
|
||||
|
||||
fn config() -> Config {
|
||||
Config {
|
||||
server: ServerConfig {
|
||||
host: "0".into(),
|
||||
port: 0,
|
||||
},
|
||||
observability: ObservabilityConfig::default(),
|
||||
model: ModelConfig {
|
||||
id: MODEL.into(),
|
||||
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
|
||||
policy: PolicyKind::CacheAwareZmq,
|
||||
circuit_breaker: None,
|
||||
cache_aware: Some(CacheAwareConfig::default()),
|
||||
sticky: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_ctx(url: String) -> Arc<AppContext> {
|
||||
let cfg = config();
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
assert!(
|
||||
tokenizers.has_chat_encoder(MODEL),
|
||||
"deepseek-v4 model id must auto-attach the built-in chat encoder"
|
||||
);
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
let _ = registry.add(WorkerSpec {
|
||||
id: WorkerId(url.clone()),
|
||||
url,
|
||||
mode: WorkerMode::Plain,
|
||||
model_ids: vec![ModelId(MODEL.into())],
|
||||
bootstrap_port: None,
|
||||
});
|
||||
// Use the real loaded tokenizers (not the empty-registry test default) so
|
||||
// the cache-aware policy can tokenize at ingress.
|
||||
let policies = Arc::new(
|
||||
build_registry(
|
||||
&cfg,
|
||||
Arc::new(HashTree::new()),
|
||||
Arc::clone(&tokenizers),
|
||||
BlockSizeOracle::new(),
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
|
||||
Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies))
|
||||
}
|
||||
|
||||
async fn send(ctx: Arc<AppContext>, body: Value) -> StatusCode {
|
||||
let app = build_router(ctx);
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/chat/completions")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap();
|
||||
app.oneshot(req).await.unwrap().status()
|
||||
}
|
||||
|
||||
fn captured(mock: &MockWorker) -> Value {
|
||||
let b = mock
|
||||
.captured
|
||||
.lock()
|
||||
.unwrap()
|
||||
.last_body
|
||||
.clone()
|
||||
.expect("worker captured a request body");
|
||||
serde_json::from_slice(&b).expect("captured body is valid JSON")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn plain_chat_forwards_input_ids_and_keeps_messages() {
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(mock.url.clone());
|
||||
let status = send(
|
||||
ctx,
|
||||
json!({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "hello there friend"}],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let body = captured(&mock);
|
||||
let ids = body.get("input_ids").and_then(|v| v.as_array());
|
||||
assert!(
|
||||
ids.is_some_and(|a| !a.is_empty()),
|
||||
"engine must receive non-empty input_ids; got {body}"
|
||||
);
|
||||
assert!(
|
||||
body.get("messages").is_some(),
|
||||
"messages must be retained alongside input_ids; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_request_omits_input_ids() {
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(mock.url.clone());
|
||||
let status = send(
|
||||
ctx,
|
||||
json!({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "f"}}],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let body = captured(&mock);
|
||||
assert!(
|
||||
body.get("input_ids").is_none(),
|
||||
"tool requests must not forward input_ids; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn thinking_request_omits_input_ids() {
|
||||
// `chat_template_kwargs` steers engine-side thinking mode, which the
|
||||
// router's encoder renders in the default mode only — forwarding ids would
|
||||
// silently run the wrong mode, so the handler must omit them.
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(mock.url.clone());
|
||||
let status = send(
|
||||
ctx,
|
||||
json!({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"chat_template_kwargs": {"enable_thinking": true},
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let body = captured(&mock);
|
||||
assert!(
|
||||
body.get("input_ids").is_none(),
|
||||
"thinking-mode requests must not forward input_ids; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multimodal_request_omits_input_ids() {
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(mock.url.clone());
|
||||
let status = send(
|
||||
ctx,
|
||||
json!({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": "x"}]}],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let body = captured(&mock);
|
||||
assert!(
|
||||
body.get("input_ids").is_none(),
|
||||
"multimodal requests must not forward input_ids; got {body}"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user