[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}"
|
||||
);
|
||||
}
|
||||
@@ -10,11 +10,14 @@
|
||||
|
||||
mod common;
|
||||
|
||||
mod cache_aware_input_ids;
|
||||
mod chat_routing;
|
||||
mod failover;
|
||||
mod graceful_shutdown;
|
||||
mod header_forwarding;
|
||||
mod pd_bootstrap_injection;
|
||||
mod pd_pool_isolation;
|
||||
mod roundrobin_input_ids;
|
||||
mod sticky_input_ids;
|
||||
mod sticky_routing;
|
||||
mod timeout;
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! `input_ids` forwarding is policy-independent: a load-only **round-robin**
|
||||
//! policy on a chat-encoder model still forwards `input_ids` to the engine
|
||||
//! (the engine-tokenization offload), even though it picks workers round-robin
|
||||
//! and ignores the tokens for routing. Tokenization is gated on the model's
|
||||
//! chat encoder at ingress, not on the policy.
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry_with_defaults;
|
||||
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;
|
||||
|
||||
// deepseek-v4 id → the tokenizer registry auto-attaches the built-in V4 chat
|
||||
// encoder, so the model has an engine-equivalent encode path.
|
||||
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::RoundRobin,
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
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();
|
||||
// The handler tokenizes via the AppContext's registry (which carries the V4
|
||||
// encoder); the RoundRobin policy itself needs no tokenizer.
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
assert!(tokenizers.has_chat_encoder(MODEL));
|
||||
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,
|
||||
});
|
||||
let policies = Arc::new(build_registry_with_defaults(&cfg).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")
|
||||
}
|
||||
|
||||
/// A round-robin (load-only) policy still forwards `input_ids` on a
|
||||
/// chat-encoder model — the offload is decoupled from routing.
|
||||
#[tokio::test]
|
||||
async fn round_robin_plain_chat_forwards_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": "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()),
|
||||
"round-robin must forward input_ids on a chat-encoder model; got {body}"
|
||||
);
|
||||
assert!(
|
||||
body.get("messages").is_some(),
|
||||
"messages must be retained alongside input_ids; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Even under round-robin, a tool request omits `input_ids` (the safe predicate
|
||||
/// is policy-independent too).
|
||||
#[tokio::test]
|
||||
async fn round_robin_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 under any policy; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
/// A successful plain-chat forward on a chat-encoder model must NOT emit
|
||||
/// `sgl_router_ingress_tokenize_errors_total` — that counter fires only when the
|
||||
/// offload was expected but the encoder failed. A tool request on the same model
|
||||
/// is an *expected* omission (its ids are still engine-equivalent; the
|
||||
/// safe-predicate withholds forwarding for other reasons), so it must not emit
|
||||
/// the error counter either.
|
||||
#[tokio::test]
|
||||
async fn successful_forward_does_not_emit_ingress_tokenize_error() {
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(mock.url.clone());
|
||||
|
||||
let status = send(
|
||||
Arc::clone(&ctx),
|
||||
json!({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "hello there friend"}],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let status = send(
|
||||
Arc::clone(&ctx),
|
||||
json!({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "f"}}],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
let m = ctx.metrics.render();
|
||||
assert!(
|
||||
m.contains("# TYPE sgl_router_ingress_tokenize_errors_total counter"),
|
||||
"the error counter family must be exposed; got:\n{m}",
|
||||
);
|
||||
assert!(
|
||||
!m.contains("sgl_router_ingress_tokenize_errors_total{"),
|
||||
"healthy forwards (and expected omissions) must not emit the error counter; got:\n{m}",
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! Tokenize-once at ingress under the STICKY policy. The engine-tokenization
|
||||
//! offload (`input_ids` forwarding) is a property of the MODEL — does it have a
|
||||
//! chat encoder? — not of the routing policy, so a sticky-routed request on a
|
||||
//! chat-encoder model must forward `input_ids` exactly like cache-aware does,
|
||||
//! while still pinning sessions O(1) by header.
|
||||
//!
|
||||
//! Asserts through the real chat handler + `MockWorker` backends:
|
||||
//!
|
||||
//! * A plain text chat request forwards `input_ids` AND retains `messages`,
|
||||
//! even though sticky never consults the tokens for routing.
|
||||
//! * A request carrying `tools` / multimodal content omits `input_ids` — the
|
||||
//! same safe-to-forward predicate applies regardless of policy.
|
||||
//! * Same-session-header requests still pin to a single worker (O(1) sticky
|
||||
//! routing is unchanged by the added tokenization).
|
||||
//!
|
||||
//! 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, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
|
||||
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";
|
||||
const HEADER: &str = "x-sgl-routing-key";
|
||||
|
||||
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::Sticky,
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
// Push eviction far out so the background sweeper never fires
|
||||
// mid-test; round-robin fallback for the initial pin of a key.
|
||||
sticky: Some(StickyConfig {
|
||||
header_name: HEADER.to_string(),
|
||||
fallback_policy: PolicyKind::RoundRobin,
|
||||
idle_secs: 3600,
|
||||
eviction_interval_secs: 3600,
|
||||
}),
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an `AppContext` running the sticky policy over the given workers.
|
||||
/// The tokenizer registry is loaded from config (real tiny tokenizer + the
|
||||
/// auto-attached V4 chat encoder) so the ingress can tokenize — the sticky
|
||||
/// policy itself holds no tokenizer.
|
||||
fn build_ctx(worker_urls: &[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());
|
||||
for (i, url) in worker_urls.iter().enumerate() {
|
||||
let _ = registry.add(WorkerSpec {
|
||||
id: WorkerId(format!("w{i}")),
|
||||
url: url.clone(),
|
||||
mode: WorkerMode::Plain,
|
||||
model_ids: vec![ModelId(MODEL.into())],
|
||||
bootstrap_port: None,
|
||||
});
|
||||
}
|
||||
// Sticky needs no cache-aware deps, so the defaults registry is fine — the
|
||||
// ingress tokenizes via `ctx.tokenizers`, not the policy.
|
||||
let policies = Arc::new(build_policy_registry(&cfg).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>, routing_key: &str, body: Value) -> StatusCode {
|
||||
let app = build_router(ctx);
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/chat/completions")
|
||||
.header("content-type", "application/json")
|
||||
.header(HEADER, routing_key)
|
||||
.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 sticky_plain_chat_forwards_input_ids_and_keeps_messages() {
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(std::slice::from_ref(&mock.url));
|
||||
let status = send(
|
||||
ctx,
|
||||
"alice",
|
||||
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()),
|
||||
"sticky-routed chat must still forward 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 sticky_tool_request_omits_input_ids() {
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(std::slice::from_ref(&mock.url));
|
||||
let status = send(
|
||||
ctx,
|
||||
"alice",
|
||||
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 even under sticky; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sticky_thinking_request_omits_input_ids() {
|
||||
// `chat_template_kwargs` steers engine-side thinking mode the router's
|
||||
// encoder renders in the default mode only — the safe-to-forward predicate
|
||||
// is policy-independent, so sticky must omit ids here too.
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(std::slice::from_ref(&mock.url));
|
||||
let status = send(
|
||||
ctx,
|
||||
"alice",
|
||||
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 under sticky; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sticky_multimodal_request_omits_input_ids() {
|
||||
let mock = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(std::slice::from_ref(&mock.url));
|
||||
let status = send(
|
||||
ctx,
|
||||
"alice",
|
||||
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 under sticky; got {body}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Routing is unchanged: same session header pins every request to one worker
|
||||
/// (O(1) sticky), even though the ingress now also tokenizes. With two
|
||||
/// backends, all same-key requests must land on exactly one of them.
|
||||
#[tokio::test]
|
||||
async fn sticky_pins_session_by_header_with_tokenization_on() {
|
||||
let w0 = MockWorker::start(vec![]).await;
|
||||
let w1 = MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(&[w0.url.clone(), w1.url.clone()]);
|
||||
let app = build_router(ctx.clone());
|
||||
|
||||
const N: usize = 5;
|
||||
for _ in 0..N {
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/chat/completions")
|
||||
.header("content-type", "application/json")
|
||||
.header(HEADER, "alice")
|
||||
.body(Body::from(
|
||||
serde_json::to_vec(&json!({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "hello there friend"}],
|
||||
}))
|
||||
.unwrap(),
|
||||
))
|
||||
.unwrap();
|
||||
let res = app.clone().oneshot(req).await.unwrap();
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
// Exactly one worker captured a body — all same-key requests pinned to it.
|
||||
let w0_hit = w0.captured.lock().unwrap().last_body.is_some();
|
||||
let w1_hit = w1.captured.lock().unwrap().last_body.is_some();
|
||||
assert!(
|
||||
w0_hit ^ w1_hit,
|
||||
"same routing key must pin to exactly one worker (w0_hit={w0_hit}, w1_hit={w1_hit})"
|
||||
);
|
||||
|
||||
// And the pinned worker still received forwarded input_ids — the offload
|
||||
// and the pin coexist.
|
||||
let pinned = if w0_hit { &w0 } else { &w1 };
|
||||
let body = captured(pinned);
|
||||
assert!(
|
||||
body.get("input_ids")
|
||||
.and_then(|v| v.as_array())
|
||||
.is_some_and(|a| !a.is_empty()),
|
||||
"the pinned worker must receive forwarded input_ids; got {body}"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user