feat: rust sglang server openai apis (#33103)
Co-authored-by: Rain Jiang <rain-jiang@outlook.com>
This commit is contained in:
Chengyu Lin
co-authored by
Rain Jiang
Rain Jiang
parent
0bf0640b9d
commit
e00f32ed4f
Generated
+1414
-146
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,4 @@
|
||||
[toolchain]
|
||||
channel = "1.90"
|
||||
channel = "1.92"
|
||||
profile = "minimal"
|
||||
components = ["clippy", "rustfmt"]
|
||||
|
||||
@@ -36,6 +36,12 @@ axum = { version = "0.8.9", features = ["json", "tokio"] }
|
||||
core_affinity = "0.8"
|
||||
# the dynamo-tokenizers is deps on hf-hub, should bump version together.
|
||||
dynamo-tokenizers = "1.7.0"
|
||||
# Keep the HTTP adapter on Dynamo's public OpenAI types, renderer, and parsers.
|
||||
# `dynamo-llm` stays excluded because it pulls the full runtime just to reuse
|
||||
# small service helpers.
|
||||
dynamo-parsers = "7.0.1"
|
||||
dynamo-protocols = "5.1.0"
|
||||
dynamo-renderer = "5.0.0"
|
||||
flume = "0.12.0"
|
||||
itertools = "0.14"
|
||||
hf-hub = { version = "0.4", default-features = false }
|
||||
@@ -45,3 +51,7 @@ rmpv = { version = "1", features = ["with-serde"] }
|
||||
# "anything Rust admits, Python can compile" invariant in `message::sampling`.
|
||||
# A minor bump can widen it and silently reopen a scheduler-killing hole.
|
||||
regex-syntax = "=0.8.11"
|
||||
|
||||
[dev-dependencies]
|
||||
# `Router::oneshot` for handler-level router tests.
|
||||
tower = { version = "0.5", features = ["util"] }
|
||||
|
||||
@@ -20,13 +20,14 @@ use crate::runtime::ServerArgs;
|
||||
use crate::tokenizer_manager::ActivityCounter;
|
||||
use crate::tokenizer_manager::Senders;
|
||||
|
||||
/// Shared handler state: the submit machinery (`senders`, `egress_buf`)
|
||||
/// + shared tokenizer.
|
||||
/// Shared handler state: submission handles, immutable server configuration,
|
||||
/// and the API-owned chat formatter.
|
||||
#[derive(Clone)]
|
||||
struct AppState {
|
||||
senders: Senders,
|
||||
egress_buf: usize,
|
||||
server_args: Arc<ServerArgs>,
|
||||
chat_formatter: Option<openai::ChatFormatter>,
|
||||
/// Egress heartbeat (bumped per drained ring frame).
|
||||
egress_activity: ActivityCounter,
|
||||
}
|
||||
@@ -42,10 +43,12 @@ pub async fn serve(
|
||||
// releases.
|
||||
shutdown: flume::Receiver<()>,
|
||||
) {
|
||||
let chat_formatter = openai::load_chat_support(&server_args);
|
||||
let state = AppState {
|
||||
senders,
|
||||
egress_buf,
|
||||
server_args: server_args.clone(),
|
||||
chat_formatter,
|
||||
egress_activity,
|
||||
};
|
||||
// Each endpoint module registers its own routes and merges here.
|
||||
|
||||
@@ -56,12 +56,14 @@ async fn await_control_result(
|
||||
StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
Err((code, e.to_string()).into_response())
|
||||
}
|
||||
// A control request never receives generation frames.
|
||||
Some(EgressItem::Frame(_)) | Some(EgressItem::Done(_)) => Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"unexpected generation output for control request",
|
||||
)
|
||||
.into_response()),
|
||||
// A control request never receives generation frames or service-call data.
|
||||
Some(EgressItem::Frame(_)) | Some(EgressItem::Done(_)) | Some(EgressItem::Data(_)) => {
|
||||
Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"unexpected generation output for control request",
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
None => Err((StatusCode::from_u16(499).unwrap(), "request aborted").into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -243,7 +243,7 @@ async fn drain_unary(
|
||||
StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
return (status, error_value(code, &e.to_string()), true);
|
||||
}
|
||||
EgressItem::Control(_) => continue, // never on `/generate`
|
||||
EgressItem::Control(_) | EgressItem::Data(_) => continue, // never on `/generate`
|
||||
}
|
||||
}
|
||||
// Sender dropped without a terminal item: the shard dropped this request (a
|
||||
@@ -376,7 +376,7 @@ fn generation_event_stream(
|
||||
terminal = Some(out);
|
||||
}
|
||||
EgressItem::Error(e) => failed = Some(e),
|
||||
EgressItem::Control(_) => {} // never on /generate
|
||||
EgressItem::Control(_) | EgressItem::Data(_) => {} // never on /generate
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,55 +1,240 @@
|
||||
//! OpenAI-compatible endpoints: `/v1/completions`, `/v1/chat/completions`, and
|
||||
//! `/v1/models`. Each runs the same tokenize→generate→detok pipeline as
|
||||
//! `/generate` and shapes the neutral [`ChunkEvent`] delta into OpenAI types
|
||||
//! (`dynamo-protocols`), with chat-template rendering (`dynamo-renderer`) and
|
||||
//! reasoning / tool-call parsing (`dynamo-parsers`).
|
||||
//! OpenAI-compatible generation endpoints.
|
||||
//!
|
||||
//! Mounted on the shared [`AppState`](super::AppState) by the parent
|
||||
//! `api_server` module; the submit machinery and control plane live there.
|
||||
//! The HTTP adapter stays deliberately thin: Dynamo owns the standard OpenAI
|
||||
//! request and response primitives. Native [`ChunkEvent`] values remain the one
|
||||
//! backend output type for both unary and streaming responses.
|
||||
|
||||
use axum::{
|
||||
extract::State,
|
||||
Json, Router,
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use futures::StreamExt;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use axum::{Router, routing::get};
|
||||
mod chat;
|
||||
mod completions;
|
||||
mod models;
|
||||
mod reasoning;
|
||||
mod template;
|
||||
mod tools;
|
||||
|
||||
pub(super) use template::ChatFormatter;
|
||||
|
||||
use super::AppState;
|
||||
use super::frame::OutputAccumulator;
|
||||
use super::guard::AbortGuard;
|
||||
use super::submit::submit;
|
||||
use crate::ids::Rid;
|
||||
use crate::message::{ChunkEvent, EgressItem, GenerateRequest, RequestKind};
|
||||
use crate::runtime::ServerArgs;
|
||||
|
||||
const MAX_OPENAI_CHOICES: usize = 4096;
|
||||
|
||||
/// The routes this module owns, mounted by `api_server::serve`.
|
||||
pub(super) fn routes() -> Router<AppState> {
|
||||
// `/v1/models` is OpenAI-compatible; completions/chat land here too.
|
||||
Router::new().route("/v1/models", get(available_models))
|
||||
Router::new()
|
||||
.merge(models::routes())
|
||||
.merge(completions::routes())
|
||||
.merge(chat::routes())
|
||||
}
|
||||
|
||||
/// `GET /v1/models` — OpenAI-compatible model list. Served from `server_args`;
|
||||
/// no scheduler round-trip. Mirrors `http_server.available_models`.
|
||||
///
|
||||
/// TODO(v1/models): when `--enable-lora`, append a `ModelCard` per loaded LoRA
|
||||
/// adapter (`id=lora_name, root=lora_path, parent=served_model_name,
|
||||
/// max_model_len=None`). Adapters load/unload at runtime, so that part needs a
|
||||
/// control-request query to the scheduler's LoRA registry.
|
||||
async fn available_models(State(state): State<AppState>) -> Response {
|
||||
let created = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
let name = &state.server_args.served_model_name;
|
||||
let base = serde_json::json!({
|
||||
"id": name,
|
||||
"object": "model",
|
||||
"created": created,
|
||||
"owned_by": "sglang",
|
||||
"root": name,
|
||||
"parent": serde_json::Value::Null,
|
||||
"max_model_len": state.server_args.model_config.context_len,
|
||||
});
|
||||
let list = serde_json::json!({ "object": "list", "data": [base] });
|
||||
(
|
||||
StatusCode::OK,
|
||||
[("content-type", "application/json")],
|
||||
serde_json::to_vec(&list).unwrap_or_default(),
|
||||
)
|
||||
.into_response()
|
||||
/// Resolve the chat formatter, or `None` to disable the OpenAI chat-completions
|
||||
/// endpoint. Tokenization is the tokenizer pool's job (the api server never
|
||||
/// encodes); the formatter needs at most `tokenizer_config.json` — a built-in
|
||||
/// `--chat-template` name or a model-path-inferred legacy template resolve
|
||||
/// without it, so its absence must not disable chat.
|
||||
pub(super) fn load_chat_support(server_args: &ServerArgs) -> Option<ChatFormatter> {
|
||||
// Chat needs the tokenizer pool behind it: under `skip_tokenizer_init`
|
||||
// there is none (text cannot be submitted), so chat is disabled.
|
||||
if server_args.skip_tokenizer_init || server_args.tokenizer_path.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let config_file = crate::tokenizer::resolve_model_file(
|
||||
&server_args.tokenizer_path,
|
||||
server_args.revision.as_deref(),
|
||||
"tokenizer_config.json",
|
||||
);
|
||||
|
||||
match template::load_chat_formatter(
|
||||
config_file.as_deref(),
|
||||
(!server_args.model_path.is_empty()).then_some(server_args.model_path.as_str()),
|
||||
server_args.chat_template.as_deref(),
|
||||
) {
|
||||
Ok(formatter) => {
|
||||
tracing::info!(
|
||||
config = ?config_file.as_deref().unwrap_or("<built-in / inferred>"),
|
||||
"loaded OpenAI chat template"
|
||||
);
|
||||
Some(formatter)
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "OpenAI chat completions disabled");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn unix_seconds() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn unix_seconds_u32() -> u32 {
|
||||
u32::try_from(unix_seconds()).unwrap_or(u32::MAX)
|
||||
}
|
||||
|
||||
/// The OpenAI `{"error": {...}}` payload: `type` is the SDK-facing error kind
|
||||
/// (`AuthenticationError` / `InternalServerError` / `BadRequestError`), and
|
||||
/// `code` carries the HTTP status — the shape Python's OpenAI frontend emits.
|
||||
fn error_payload(code: StatusCode, message: String) -> serde_json::Value {
|
||||
let error_type = if code == StatusCode::UNAUTHORIZED {
|
||||
"AuthenticationError"
|
||||
} else if code.is_server_error() {
|
||||
"InternalServerError"
|
||||
} else {
|
||||
"BadRequestError"
|
||||
};
|
||||
serde_json::json!({
|
||||
"error": {
|
||||
"object": "error",
|
||||
"message": message,
|
||||
"type": error_type,
|
||||
"param": null,
|
||||
"code": code.as_u16(),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Shape a `StatusCode` + message into an OpenAI error response, mirroring
|
||||
/// `pre_submit_error`'s rule: unary requests get the JSON error with its
|
||||
/// status; a request whose stream is already committed gets 200 + one SSE
|
||||
/// error frame + `[DONE]`.
|
||||
pub(super) fn openai_error_response(
|
||||
code: StatusCode,
|
||||
message: impl Into<String>,
|
||||
stream: bool,
|
||||
) -> Response {
|
||||
let body = error_payload(code, message.into());
|
||||
if !stream {
|
||||
return (code, Json(body)).into_response();
|
||||
}
|
||||
super::submit::sse_error_response(body)
|
||||
}
|
||||
|
||||
/// Unary OpenAI error — the common pre-submit case (Python validates before
|
||||
/// the stream starts and answers 4xx in JSON even for `stream=true`).
|
||||
fn openai_error(code: StatusCode, message: impl Into<String>) -> Response {
|
||||
openai_error_response(code, message, false)
|
||||
}
|
||||
|
||||
/// The OpenAI error frame payload for errors raised *inside* a committed
|
||||
/// stream, where only a `data:` frame can be emitted (the status is folded
|
||||
/// into the body, since the response status is already 200).
|
||||
pub(super) fn streaming_error(code: u16, message: impl Into<String>) -> String {
|
||||
let status = StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
error_payload(status, message.into()).to_string()
|
||||
}
|
||||
|
||||
async fn collect_output(
|
||||
mut rx: mpsc::Receiver<EgressItem>,
|
||||
guard: &mut AbortGuard,
|
||||
rid: &Rid,
|
||||
) -> Result<ChunkEvent, (StatusCode, String)> {
|
||||
let mut accumulator = OutputAccumulator::default();
|
||||
let output = loop {
|
||||
match rx.recv().await {
|
||||
Some(EgressItem::Frame(output)) => accumulator.fold(&output),
|
||||
Some(EgressItem::Done(output)) => {
|
||||
accumulator.fold(&output);
|
||||
break accumulator.into_output();
|
||||
}
|
||||
Some(EgressItem::Error(error)) => {
|
||||
guard.disarm(rid);
|
||||
let status = StatusCode::from_u16(error.http_status())
|
||||
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
return Err((status, error.to_string()));
|
||||
}
|
||||
Some(EgressItem::Control(_)) | Some(EgressItem::Data(_)) => {}
|
||||
None => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"response truncated before completion".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
guard.disarm(rid);
|
||||
if let Some((code, message)) = output
|
||||
.finish_reason
|
||||
.as_ref()
|
||||
.and_then(|reason| reason.abort_status())
|
||||
{
|
||||
return Err((
|
||||
StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
|
||||
message.to_owned(),
|
||||
));
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
async fn submit_generation(
|
||||
state: &AppState,
|
||||
request: GenerateRequest,
|
||||
stream: bool,
|
||||
guard: &mut AbortGuard,
|
||||
) -> Result<mpsc::Receiver<EgressItem>, Response> {
|
||||
match submit(state, RequestKind::Generate(Box::new(request)), stream).await {
|
||||
Ok((rid, rx)) => {
|
||||
guard.arm(rid);
|
||||
Ok(rx)
|
||||
}
|
||||
// Same rule as `pre_submit_error`: a committed stream gets 200 plus an
|
||||
// SSE error frame + `[DONE]`, not a unary 503 — but with the OpenAI
|
||||
// error shape, since this is the OpenAI frontend.
|
||||
Err(_) => Err(openai_error_response(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"service unavailable",
|
||||
stream,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn indexed_egress_stream(
|
||||
index: usize,
|
||||
rx: mpsc::Receiver<EgressItem>,
|
||||
) -> futures::stream::BoxStream<'static, (usize, Option<EgressItem>)> {
|
||||
futures::stream::unfold((rx, false), move |(mut rx, finished)| async move {
|
||||
if finished {
|
||||
return None;
|
||||
}
|
||||
match rx.recv().await {
|
||||
Some(item) => {
|
||||
let finished = matches!(item, EgressItem::Done(_) | EgressItem::Error(_));
|
||||
Some(((index, Some(item)), (rx, finished)))
|
||||
}
|
||||
None => Some(((index, None), (rx, true))),
|
||||
}
|
||||
})
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn contains_media(value: &serde_json::Value) -> bool {
|
||||
match value {
|
||||
serde_json::Value::Array(values) => values.iter().any(contains_media),
|
||||
serde_json::Value::Object(object) => {
|
||||
object.keys().any(|key| {
|
||||
matches!(
|
||||
key.as_str(),
|
||||
"image_url" | "video_url" | "input_audio" | "audio_url" | "file"
|
||||
)
|
||||
}) || object.values().any(contains_media)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_utils;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,864 @@
|
||||
//! OpenAI legacy text-completion endpoint and wire shaping.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::convert::Infallible;
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{State, rejection::JsonRejection},
|
||||
http::StatusCode,
|
||||
response::{
|
||||
IntoResponse, Response,
|
||||
sse::{Event, Sse},
|
||||
},
|
||||
routing::post,
|
||||
};
|
||||
use dynamo_protocols::types::{
|
||||
Choice, CompletionFinishReason, CompletionUsage, CreateCompletionRequest,
|
||||
CreateCompletionResponse, Logprobs, Prompt, Stop,
|
||||
};
|
||||
use futures::StreamExt;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::super::guard::AbortGuard;
|
||||
use super::super::submit::submit;
|
||||
use super::{
|
||||
AppState, MAX_OPENAI_CHOICES, collect_output, indexed_egress_stream, openai_error,
|
||||
streaming_error, submit_generation, unix_seconds_u32,
|
||||
};
|
||||
use crate::ids::Rid;
|
||||
use crate::message::{
|
||||
ChunkEvent, ChunkExtras, EgressItem, GenerateRequest, Matched, OneOrMany, RequestKind,
|
||||
SamplingParams, TokenIds,
|
||||
};
|
||||
|
||||
pub(super) fn routes() -> Router<AppState> {
|
||||
Router::new().route("/v1/completions", post(completions))
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
enum PromptSpec {
|
||||
Text(String),
|
||||
TokenIds(TokenIds),
|
||||
}
|
||||
|
||||
pub(super) struct SubmittedChoice {
|
||||
pub(super) index: usize,
|
||||
pub(super) prompt_index: usize,
|
||||
pub(super) rid: Rid,
|
||||
pub(super) echo: String,
|
||||
pub(super) rx: mpsc::Receiver<EgressItem>,
|
||||
}
|
||||
#[derive(Debug, Default)]
|
||||
pub(super) struct ChoiceExtensions {
|
||||
matched_stop: Option<serde_json::Value>,
|
||||
/// Dynamo's enum covers the standard values. Python additionally exposes
|
||||
/// `abort`, and native unknown finish types are preserved rather than lost.
|
||||
finish_reason_override: Option<String>,
|
||||
}
|
||||
|
||||
async fn completions(
|
||||
State(state): State<AppState>,
|
||||
body: Result<Json<CreateCompletionRequest>, JsonRejection>,
|
||||
) -> Response {
|
||||
let request = match body {
|
||||
Ok(Json(request)) => request,
|
||||
Err(rejection) => {
|
||||
return openai_error(StatusCode::BAD_REQUEST, rejection.body_text());
|
||||
}
|
||||
};
|
||||
let stream = request.stream.unwrap_or(false);
|
||||
let echo = request.echo.unwrap_or(false);
|
||||
let model = request.model.clone();
|
||||
if model != state.server_args.served_model_name {
|
||||
return openai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("The model `{model}` does not exist"),
|
||||
);
|
||||
}
|
||||
|
||||
if request.prompt_embeds.is_some() {
|
||||
return openai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"prompt_embeds is not supported by the Rust frontend",
|
||||
);
|
||||
}
|
||||
if request.suffix.is_some() {
|
||||
return openai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"suffix is not supported by this model",
|
||||
);
|
||||
}
|
||||
if request.best_of.is_some_and(|best_of| best_of != 1) {
|
||||
return openai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"best_of values greater than 1 are not supported",
|
||||
);
|
||||
}
|
||||
if request.max_tokens == Some(0) {
|
||||
return openai_error(StatusCode::BAD_REQUEST, "max_tokens must be positive");
|
||||
}
|
||||
if request.n == Some(0) {
|
||||
return openai_error(StatusCode::BAD_REQUEST, "n must be at least 1");
|
||||
}
|
||||
let prompts = match completion_prompt_specs(&request.prompt) {
|
||||
Ok(prompts) => prompts,
|
||||
Err(message) => return openai_error(StatusCode::BAD_REQUEST, message),
|
||||
};
|
||||
let mut sampling = match completion_sampling_params(&request) {
|
||||
Ok(sampling) => sampling,
|
||||
Err(message) => return openai_error(StatusCode::BAD_REQUEST, message),
|
||||
};
|
||||
if let Err(error) = sampling.normalize(
|
||||
state.server_args.skip_tokenizer_init,
|
||||
state
|
||||
.server_args
|
||||
.model_config
|
||||
.vocab_size
|
||||
.unwrap_or(u64::MAX),
|
||||
) {
|
||||
return openai_error(StatusCode::BAD_REQUEST, error.to_string());
|
||||
}
|
||||
|
||||
let n = request.n.unwrap_or(1) as usize;
|
||||
let choice_count = match prompts.len().checked_mul(n) {
|
||||
Some(count) if count <= MAX_OPENAI_CHOICES => count,
|
||||
_ => {
|
||||
return openai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("prompt count times n exceeds the maximum of {MAX_OPENAI_CHOICES}"),
|
||||
);
|
||||
}
|
||||
};
|
||||
let response_id = format!("cmpl-{}", uuid::Uuid::new_v4().simple());
|
||||
let created = unix_seconds_u32();
|
||||
let mut guard = AbortGuard::new_empty(state.senders.clone());
|
||||
let mut submitted = Vec::with_capacity(choice_count);
|
||||
|
||||
for (prompt_index, prompt) in prompts.into_iter().enumerate() {
|
||||
let (text, input_ids, mut prompt_echo) = match prompt {
|
||||
PromptSpec::Text(text) => {
|
||||
let prompt_echo = if echo { text.clone() } else { String::new() };
|
||||
(Some(text), None, prompt_echo)
|
||||
}
|
||||
PromptSpec::TokenIds(input_ids) => (None, Some(input_ids), String::new()),
|
||||
};
|
||||
for sample_index in 0..n {
|
||||
let index = prompt_index * n + sample_index;
|
||||
let rid = Rid::from_client(&format!("{response_id}-{index}"));
|
||||
if echo
|
||||
&& sample_index == 0
|
||||
&& let Some(token_ids) = &input_ids
|
||||
{
|
||||
prompt_echo = match decode_prompt_echo(&state, token_ids.clone()).await {
|
||||
Ok(echo) => echo,
|
||||
Err(response) => return response,
|
||||
};
|
||||
}
|
||||
let native = GenerateRequest {
|
||||
rid: rid.clone(),
|
||||
text: text.clone(),
|
||||
input_ids: input_ids.clone(),
|
||||
sampling_params: sampling.clone(),
|
||||
stream,
|
||||
return_logprob: request.logprobs.is_some(),
|
||||
logprob_start_len: if echo && request.logprobs.is_some() {
|
||||
0
|
||||
} else {
|
||||
-1
|
||||
},
|
||||
top_logprobs_num: request.logprobs.unwrap_or(0) as i64,
|
||||
return_text_in_logprobs: request.logprobs.map(|_| true),
|
||||
..Default::default()
|
||||
};
|
||||
let rx = match submit_generation(&state, native, stream, &mut guard).await {
|
||||
Ok(rx) => rx,
|
||||
Err(response) => return response,
|
||||
};
|
||||
submitted.push(SubmittedChoice {
|
||||
index,
|
||||
prompt_index,
|
||||
rid,
|
||||
echo: prompt_echo.clone(),
|
||||
rx,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if stream {
|
||||
let include_usage = request
|
||||
.stream_options
|
||||
.map(|o| o.include_usage)
|
||||
.unwrap_or(false)
|
||||
|| state.server_args.stream_response_default_include_usage;
|
||||
let continuous_usage = request
|
||||
.stream_options
|
||||
.map(|o| o.continuous_usage_stats)
|
||||
.unwrap_or(false);
|
||||
let want_logprobs = request.logprobs.is_some();
|
||||
let s = completion_event_stream(
|
||||
submitted,
|
||||
guard,
|
||||
response_id,
|
||||
model,
|
||||
created,
|
||||
echo,
|
||||
want_logprobs,
|
||||
include_usage,
|
||||
continuous_usage,
|
||||
)
|
||||
.map(|data| Ok::<_, Infallible>(Event::default().data(data)));
|
||||
Sse::new(s).into_response()
|
||||
} else {
|
||||
unary_completion(
|
||||
submitted,
|
||||
guard,
|
||||
response_id,
|
||||
model,
|
||||
created,
|
||||
echo,
|
||||
request.logprobs.is_some(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode a token-id prompt back to text for `echo=true`, via a
|
||||
/// `RequestKind::Detokenize` request through the regular submit path — the
|
||||
/// detok stage answers it with a single `Data` payload (the raw UTF-8 text),
|
||||
/// or an `Error` (e.g. out-of-range ids → `Validation` → 400).
|
||||
async fn decode_prompt_echo(state: &AppState, token_ids: TokenIds) -> Result<String, Response> {
|
||||
let Ok((_rid, mut rx)) = submit(state, RequestKind::Detokenize { token_ids }, false).await
|
||||
else {
|
||||
// Same rule as `submit_generation`: rebuild the refusal in the OpenAI
|
||||
// error shape rather than forwarding the native-shaped response.
|
||||
return Err(openai_error(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"service unavailable",
|
||||
));
|
||||
};
|
||||
match rx.recv().await {
|
||||
Some(EgressItem::Data(payload)) => String::from_utf8(payload.to_vec()).map_err(|_| {
|
||||
openai_error(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"detokenized prompt is not valid UTF-8",
|
||||
)
|
||||
}),
|
||||
Some(EgressItem::Error(crate::error::Error::Validation(message))) => {
|
||||
Err(openai_error(StatusCode::BAD_REQUEST, message))
|
||||
}
|
||||
Some(EgressItem::Error(error)) => {
|
||||
let status = StatusCode::from_u16(error.http_status())
|
||||
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
Err(openai_error(
|
||||
status,
|
||||
format!("failed to decode prompt for echo: {error}"),
|
||||
))
|
||||
}
|
||||
Some(_) | None => Err(openai_error(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"failed to decode prompt for echo: reply channel closed",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn completion_prompt_specs(prompt: &Prompt) -> Result<Vec<PromptSpec>, String> {
|
||||
match prompt {
|
||||
Prompt::String(text) => {
|
||||
if text.is_empty() {
|
||||
return Err("Prompt cannot be empty".into());
|
||||
}
|
||||
Ok(vec![PromptSpec::Text(text.clone())])
|
||||
}
|
||||
Prompt::StringArray(texts) => {
|
||||
if texts.is_empty() || texts.iter().any(String::is_empty) {
|
||||
return Err("Prompt cannot be empty".into());
|
||||
}
|
||||
Ok(texts.iter().cloned().map(PromptSpec::Text).collect())
|
||||
}
|
||||
Prompt::IntegerArray(ids) => Ok(vec![token_prompt_spec(ids)?]),
|
||||
Prompt::ArrayOfIntegerArray(prompts) => {
|
||||
if prompts.is_empty() {
|
||||
return Err("Prompt cannot be empty".into());
|
||||
}
|
||||
prompts.iter().map(|ids| token_prompt_spec(ids)).collect()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn token_prompt_spec(ids: &[u32]) -> Result<PromptSpec, String> {
|
||||
if ids.is_empty() {
|
||||
return Err("Prompt cannot be empty".into());
|
||||
}
|
||||
let input_ids = ids
|
||||
.iter()
|
||||
.map(|&id| i32::try_from(id).map_err(|_| format!("Token ID {id} is out of range")))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(PromptSpec::TokenIds(input_ids))
|
||||
}
|
||||
|
||||
fn completion_sampling_params(request: &CreateCompletionRequest) -> Result<SamplingParams, String> {
|
||||
let mut stop = None;
|
||||
let mut stop_token_ids = None;
|
||||
match request.stop.as_ref() {
|
||||
Some(Stop::String(value)) => stop = Some(OneOrMany::One(value.clone())),
|
||||
Some(Stop::StringArray(values)) => stop = Some(OneOrMany::Many(values.clone())),
|
||||
Some(Stop::TokenIdArray(values)) => {
|
||||
stop_token_ids
|
||||
.get_or_insert_with(Vec::new)
|
||||
.extend(values.iter().map(|&id| id as i64));
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
|
||||
let mut logit_bias = BTreeMap::new();
|
||||
if let Some(values) = request.logit_bias.as_ref() {
|
||||
for (token, bias) in values {
|
||||
let bias = bias
|
||||
.as_f64()
|
||||
.ok_or_else(|| format!("logit_bias[{token:?}] must be a number"))?;
|
||||
logit_bias.insert(token.clone(), bias);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(SamplingParams {
|
||||
max_new_tokens: Some(request.max_tokens.unwrap_or(16) as i64),
|
||||
stop,
|
||||
stop_token_ids,
|
||||
temperature: request.temperature.unwrap_or(1.0) as f64,
|
||||
top_p: request.top_p.unwrap_or(1.0) as f64,
|
||||
frequency_penalty: request.frequency_penalty.unwrap_or(0.0) as f64,
|
||||
presence_penalty: request.presence_penalty.unwrap_or(0.0) as f64,
|
||||
// OpenAI `n` is implemented by fan-out: every native request has one
|
||||
// output, avoiding the native path's intentional `n > 1` rejection.
|
||||
n: 1,
|
||||
logit_bias: (!logit_bias.is_empty()).then_some(logit_bias),
|
||||
sampling_seed: request.seed,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn unary_completion(
|
||||
submitted: Vec<SubmittedChoice>,
|
||||
mut guard: AbortGuard,
|
||||
response_id: String,
|
||||
model: String,
|
||||
created: u32,
|
||||
echo: bool,
|
||||
want_logprobs: bool,
|
||||
) -> Response {
|
||||
// Every request is already submitted, so draining in choice order does not
|
||||
// serialize generation. The non-streaming native path sends one terminal
|
||||
// result, and the accumulator also tolerates intermediate frames.
|
||||
let mut choices = Vec::with_capacity(submitted.len());
|
||||
let mut extensions = Vec::with_capacity(submitted.len());
|
||||
let mut prompt_tokens = BTreeMap::<usize, u32>::new();
|
||||
let mut completion_tokens = 0u64;
|
||||
|
||||
for choice in submitted {
|
||||
let output = match collect_output(choice.rx, &mut guard, &choice.rid).await {
|
||||
Ok(output) => output,
|
||||
Err((status, message)) => return openai_error(status, message),
|
||||
};
|
||||
|
||||
prompt_tokens
|
||||
.entry(choice.prompt_index)
|
||||
.or_insert(output.prompt_tokens);
|
||||
completion_tokens = completion_tokens.saturating_add(output.completion_tokens);
|
||||
let (response_choice, extension) = completion_choice(
|
||||
choice.index,
|
||||
if echo {
|
||||
choice.echo + &output.text
|
||||
} else {
|
||||
output.text.clone()
|
||||
},
|
||||
&output,
|
||||
want_logprobs,
|
||||
echo,
|
||||
);
|
||||
choices.push(response_choice);
|
||||
extensions.push(extension);
|
||||
}
|
||||
|
||||
let prompt_tokens = prompt_tokens
|
||||
.values()
|
||||
.copied()
|
||||
.fold(0u32, u32::saturating_add);
|
||||
let usage = completion_usage(
|
||||
prompt_tokens,
|
||||
u32::try_from(completion_tokens).unwrap_or(u32::MAX),
|
||||
);
|
||||
|
||||
Json(completion_response_value(
|
||||
CreateCompletionResponse {
|
||||
id: response_id,
|
||||
choices,
|
||||
created,
|
||||
model,
|
||||
system_fingerprint: None,
|
||||
object: "text_completion".into(),
|
||||
usage: Some(usage),
|
||||
},
|
||||
&extensions,
|
||||
))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn completion_choice(
|
||||
index: usize,
|
||||
text: String,
|
||||
output: &ChunkEvent,
|
||||
want_logprobs: bool,
|
||||
include_input_logprobs: bool,
|
||||
) -> (Choice, ChoiceExtensions) {
|
||||
let reason = output.finish_reason.as_ref();
|
||||
let (finish_reason, finish_reason_override) = {
|
||||
match reason.and_then(|reason| reason.kind_name()) {
|
||||
Some("stop") => (Some(CompletionFinishReason::Stop), None),
|
||||
Some("length") => (Some(CompletionFinishReason::Length), None),
|
||||
Some("content_filter") => (Some(CompletionFinishReason::ContentFilter), None),
|
||||
Some(other) => (None, Some(other.into())),
|
||||
None => (None, None),
|
||||
}
|
||||
};
|
||||
let matched_stop = reason
|
||||
.and_then(|reason| reason.matched())
|
||||
.map(|matched| match matched {
|
||||
Matched::Token(id) => serde_json::json!(id),
|
||||
Matched::Str(value) => serde_json::json!(value),
|
||||
// Python's OpenAI schema supports an integer or string here, not a
|
||||
// multi-token list. Preserve the native value rather than dropping it.
|
||||
Matched::Tokens(ids) => serde_json::json!(ids),
|
||||
});
|
||||
(
|
||||
Choice {
|
||||
text,
|
||||
index: u32::try_from(index).unwrap_or(u32::MAX),
|
||||
logprobs: want_logprobs
|
||||
.then(|| completion_logprobs(output.extras.as_deref(), include_input_logprobs)),
|
||||
finish_reason,
|
||||
},
|
||||
ChoiceExtensions {
|
||||
matched_stop,
|
||||
finish_reason_override,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Serialize Dynamo's standard response and add only SGLang/Python fields that
|
||||
/// its schema cannot represent. `text_offset` is corrected here because Dynamo
|
||||
/// types it as `u32`, while Python deliberately emits `-1`.
|
||||
pub(super) fn completion_response_value(
|
||||
response: CreateCompletionResponse,
|
||||
extensions: &[ChoiceExtensions],
|
||||
) -> serde_json::Value {
|
||||
let mut value = serde_json::to_value(response).expect("OpenAI response must serialize");
|
||||
let Some(root) = value.as_object_mut() else {
|
||||
return value;
|
||||
};
|
||||
// Python's Completion response does not expose this OpenAI field.
|
||||
root.remove("system_fingerprint");
|
||||
let Some(choices) = root
|
||||
.get_mut("choices")
|
||||
.and_then(serde_json::Value::as_array_mut)
|
||||
else {
|
||||
return value;
|
||||
};
|
||||
for (choice, extension) in choices.iter_mut().zip(extensions) {
|
||||
let Some(choice) = choice.as_object_mut() else {
|
||||
continue;
|
||||
};
|
||||
if let Some(reason) = &extension.finish_reason_override {
|
||||
choice.insert("finish_reason".into(), serde_json::json!(reason));
|
||||
}
|
||||
choice.insert(
|
||||
"matched_stop".into(),
|
||||
extension
|
||||
.matched_stop
|
||||
.clone()
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
);
|
||||
if let Some(logprobs) = choice
|
||||
.get_mut("logprobs")
|
||||
.and_then(serde_json::Value::as_object_mut)
|
||||
{
|
||||
let count = logprobs
|
||||
.get("tokens")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map_or(0, Vec::len);
|
||||
logprobs.insert("text_offset".into(), serde_json::json!(vec![-1; count]));
|
||||
}
|
||||
}
|
||||
value
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn completion_event_stream(
|
||||
submitted: Vec<SubmittedChoice>,
|
||||
mut guard: AbortGuard,
|
||||
response_id: String,
|
||||
model: String,
|
||||
created: u32,
|
||||
echo: bool,
|
||||
want_logprobs: bool,
|
||||
include_usage: bool,
|
||||
continuous_usage: bool,
|
||||
) -> impl futures::Stream<Item = String> {
|
||||
async_stream::stream! {
|
||||
let count = submitted.len();
|
||||
let mut rids = Vec::with_capacity(count);
|
||||
let mut prompt_indexes = Vec::with_capacity(count);
|
||||
let mut echoes = Vec::with_capacity(count);
|
||||
let mut first_chunks = vec![true; count];
|
||||
let mut prompt_tokens_by_prompt = BTreeMap::<usize, u32>::new();
|
||||
let mut completion_tokens_by_choice = vec![0u64; count];
|
||||
let mut streams = Vec::with_capacity(count);
|
||||
|
||||
for choice in submitted {
|
||||
let index = choice.index;
|
||||
rids.push(choice.rid);
|
||||
prompt_indexes.push(choice.prompt_index);
|
||||
echoes.push(choice.echo);
|
||||
streams.push(indexed_egress_stream(index, choice.rx));
|
||||
}
|
||||
let mut events = futures::stream::select_all(streams);
|
||||
|
||||
while let Some((index, item)) = events.next().await {
|
||||
let Some(item) = item else {
|
||||
yield streaming_error(500, "response truncated before completion");
|
||||
continue;
|
||||
};
|
||||
let output = match item {
|
||||
EgressItem::Frame(output) => output,
|
||||
EgressItem::Done(output) => {
|
||||
guard.disarm(&rids[index]);
|
||||
output
|
||||
}
|
||||
EgressItem::Error(error) => {
|
||||
guard.disarm(&rids[index]);
|
||||
yield streaming_error(error.http_status(), error.to_string());
|
||||
continue;
|
||||
}
|
||||
EgressItem::Control(_) | EgressItem::Data(_) => continue,
|
||||
};
|
||||
|
||||
if let Some((code, message)) = output
|
||||
.finish_reason
|
||||
.as_ref()
|
||||
.and_then(|reason| reason.abort_status())
|
||||
{
|
||||
yield streaming_error(code, message);
|
||||
continue;
|
||||
}
|
||||
|
||||
prompt_tokens_by_prompt
|
||||
.entry(prompt_indexes[index])
|
||||
.or_insert(output.prompt_tokens);
|
||||
completion_tokens_by_choice[index] = completion_tokens_by_choice[index]
|
||||
.saturating_add(output.completion_tokens);
|
||||
let first = std::mem::replace(&mut first_chunks[index], false);
|
||||
let text = if echo && first {
|
||||
echoes[index].clone() + &output.text
|
||||
} else {
|
||||
output.text.clone()
|
||||
};
|
||||
let chunk_usage = continuous_usage.then(|| {
|
||||
completion_usage(
|
||||
output.prompt_tokens,
|
||||
u32::try_from(completion_tokens_by_choice[index]).unwrap_or(u32::MAX),
|
||||
)
|
||||
});
|
||||
let (choice, extension) = completion_choice(
|
||||
index,
|
||||
text,
|
||||
&output,
|
||||
want_logprobs,
|
||||
echo && first,
|
||||
);
|
||||
let chunk = CreateCompletionResponse {
|
||||
id: response_id.clone(),
|
||||
choices: vec![choice],
|
||||
created,
|
||||
model: model.clone(),
|
||||
system_fingerprint: None,
|
||||
object: "text_completion".into(),
|
||||
usage: chunk_usage,
|
||||
};
|
||||
yield completion_response_value(chunk, &[extension]).to_string();
|
||||
}
|
||||
|
||||
if include_usage {
|
||||
let prompt_tokens = prompt_tokens_by_prompt
|
||||
.values()
|
||||
.copied()
|
||||
.fold(0u32, u32::saturating_add);
|
||||
let completion_tokens = completion_tokens_by_choice
|
||||
.into_iter()
|
||||
.fold(0u64, u64::saturating_add);
|
||||
let final_chunk = CreateCompletionResponse {
|
||||
id: response_id,
|
||||
choices: vec![],
|
||||
created,
|
||||
model,
|
||||
system_fingerprint: None,
|
||||
object: "text_completion".into(),
|
||||
usage: Some(completion_usage(
|
||||
prompt_tokens,
|
||||
u32::try_from(completion_tokens).unwrap_or(u32::MAX),
|
||||
)),
|
||||
};
|
||||
yield completion_response_value(final_chunk, &[]).to_string();
|
||||
}
|
||||
yield "[DONE]".to_string();
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn completion_usage(prompt_tokens: u32, completion_tokens: u32) -> CompletionUsage {
|
||||
CompletionUsage {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens: prompt_tokens.saturating_add(completion_tokens),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn completion_logprobs(extras: Option<&ChunkExtras>, include_input: bool) -> Logprobs {
|
||||
let mut result = Logprobs {
|
||||
tokens: Vec::new(),
|
||||
token_logprobs: Vec::new(),
|
||||
top_logprobs: Vec::new(),
|
||||
text_offset: Vec::new(),
|
||||
};
|
||||
let Some(extras) = extras else {
|
||||
return result;
|
||||
};
|
||||
if include_input {
|
||||
append_selected_logprobs(
|
||||
&mut result,
|
||||
&extras.in_lp_val,
|
||||
&extras.in_lp_idx,
|
||||
&extras.in_lp_txt,
|
||||
);
|
||||
append_top_logprobs(
|
||||
&mut result,
|
||||
&extras.in_top_val,
|
||||
&extras.in_top_idx,
|
||||
&extras.in_top_lens,
|
||||
&extras.in_top_txt,
|
||||
);
|
||||
}
|
||||
append_selected_logprobs(
|
||||
&mut result,
|
||||
&extras.out_lp_val,
|
||||
&extras.out_lp_idx,
|
||||
&extras.out_lp_txt,
|
||||
);
|
||||
append_top_logprobs(
|
||||
&mut result,
|
||||
&extras.out_top_val,
|
||||
&extras.out_top_idx,
|
||||
&extras.out_top_lens,
|
||||
&extras.out_top_txt,
|
||||
);
|
||||
result
|
||||
}
|
||||
|
||||
fn append_selected_logprobs(result: &mut Logprobs, values: &[f32], ids: &[i32], texts: &[String]) {
|
||||
for (index, (&value, &id)) in values.iter().zip(ids).enumerate() {
|
||||
result.tokens.push(
|
||||
texts
|
||||
.get(index)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| format!("token_id:{id}")),
|
||||
);
|
||||
result
|
||||
.token_logprobs
|
||||
.push((!value.is_nan()).then_some(value));
|
||||
// Dynamo's field is `u32`; Python's `-1` sentinel is applied once at
|
||||
// final wire shaping in `completion_response_value`.
|
||||
result.text_offset.push(0);
|
||||
}
|
||||
}
|
||||
|
||||
fn append_top_logprobs(
|
||||
result: &mut Logprobs,
|
||||
values: &[f32],
|
||||
ids: &[i32],
|
||||
lens: &[u32],
|
||||
texts: &[String],
|
||||
) {
|
||||
let mut offset = 0usize;
|
||||
for &len in lens {
|
||||
let len = len as usize;
|
||||
if len == 0 {
|
||||
result.top_logprobs.push(serde_json::Value::Null);
|
||||
continue;
|
||||
}
|
||||
let mut top = BTreeMap::new();
|
||||
for index in offset..offset.saturating_add(len) {
|
||||
let (Some(&value), Some(&id)) = (values.get(index), ids.get(index)) else {
|
||||
continue;
|
||||
};
|
||||
top.insert(
|
||||
texts
|
||||
.get(index)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| format!("token_id:{id}")),
|
||||
value,
|
||||
);
|
||||
}
|
||||
result.top_logprobs.push(serde_json::json!(top));
|
||||
offset = offset.saturating_add(len);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::test_utils::{chunk, senders, submitted};
|
||||
use super::{
|
||||
ChoiceExtensions, PromptSpec, completion_event_stream, completion_logprobs,
|
||||
completion_prompt_specs, completion_response_value, unary_completion,
|
||||
};
|
||||
use crate::api_server::guard::AbortGuard;
|
||||
use crate::message::ChunkExtras;
|
||||
use axum::http::StatusCode;
|
||||
use dynamo_protocols::types::{
|
||||
Choice, CreateCompletionRequest, CreateCompletionResponse, Prompt,
|
||||
};
|
||||
use futures::StreamExt;
|
||||
|
||||
#[test]
|
||||
fn dynamo_completion_request_deserializes_directly() {
|
||||
let request: CreateCompletionRequest = serde_json::from_value(serde_json::json!({
|
||||
"model": "m",
|
||||
"prompt": ["a", "b"],
|
||||
"max_tokens": 8,
|
||||
"n": 2,
|
||||
"stream_options": {
|
||||
"include_usage": true,
|
||||
"continuous_usage_stats": true
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
assert!(matches!(request.prompt, Prompt::StringArray(_)));
|
||||
assert_eq!(request.n, Some(2));
|
||||
assert!(request.stream_options.unwrap().continuous_usage_stats);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_tokens_zero_is_rejected_before_submission() {
|
||||
let request: CreateCompletionRequest = serde_json::from_value(serde_json::json!({
|
||||
"model": "m",
|
||||
"prompt": "hello",
|
||||
"max_tokens": 0
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(request.max_tokens, Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_prompt_is_normalized_without_echo_state() {
|
||||
let specs = completion_prompt_specs(&Prompt::IntegerArray(vec![1, 2])).unwrap();
|
||||
assert_eq!(specs, [PromptSpec::TokenIds(vec![1, 2])]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_top_logprobs_keeps_selected_token_and_empty_top_map() {
|
||||
let extras = ChunkExtras {
|
||||
out_lp_val: vec![-0.25],
|
||||
out_lp_idx: vec![7],
|
||||
out_lp_txt: vec!["x".into()],
|
||||
out_top_lens: vec![0],
|
||||
..Default::default()
|
||||
};
|
||||
let logprobs = completion_logprobs(Some(&extras), false);
|
||||
assert_eq!(logprobs.tokens, ["x"]);
|
||||
assert_eq!(logprobs.token_logprobs, [Some(-0.25)]);
|
||||
assert_eq!(logprobs.top_logprobs, [serde_json::Value::Null]);
|
||||
|
||||
let value = completion_response_value(
|
||||
CreateCompletionResponse {
|
||||
id: "cmpl-test".into(),
|
||||
choices: vec![Choice {
|
||||
text: "x".into(),
|
||||
index: 0,
|
||||
logprobs: Some(logprobs),
|
||||
finish_reason: None,
|
||||
}],
|
||||
created: 1,
|
||||
model: "model".into(),
|
||||
system_fingerprint: None,
|
||||
object: "text_completion".into(),
|
||||
usage: None,
|
||||
},
|
||||
&[ChoiceExtensions::default()],
|
||||
);
|
||||
assert_eq!(
|
||||
value["choices"][0]["logprobs"]["text_offset"],
|
||||
serde_json::json!([-1])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unary_fold_orders_choices_and_counts_each_prompt_once() {
|
||||
let (choice0, tx0) = submitted(0, 0, "r0");
|
||||
let (choice1, tx1) = submitted(1, 0, "r1");
|
||||
tx0.send(chunk("r0", "a", false)).await.unwrap();
|
||||
tx0.send(chunk("r0", "b", true)).await.unwrap();
|
||||
tx1.send(chunk("r1", "x", false)).await.unwrap();
|
||||
tx1.send(chunk("r1", "y", true)).await.unwrap();
|
||||
|
||||
let response = unary_completion(
|
||||
vec![choice0, choice1],
|
||||
AbortGuard::new_empty(senders()),
|
||||
"cmpl-test".into(),
|
||||
"model".into(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
assert_eq!(value["choices"][0]["text"], "ab");
|
||||
assert_eq!(value["choices"][1]["text"], "xy");
|
||||
assert_eq!(value["choices"][0]["matched_stop"], "</s>");
|
||||
assert_eq!(value["usage"]["prompt_tokens"], 5);
|
||||
assert_eq!(value["usage"]["completion_tokens"], 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_uses_deltas_then_usage_and_done() {
|
||||
let (choice, tx) = submitted(0, 0, "r0");
|
||||
tx.send(chunk("r0", "a", false)).await.unwrap();
|
||||
tx.send(chunk("r0", "b", true)).await.unwrap();
|
||||
|
||||
let stream = completion_event_stream(
|
||||
vec![choice],
|
||||
AbortGuard::new_empty(senders()),
|
||||
"cmpl-test".into(),
|
||||
"model".into(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
);
|
||||
futures::pin_mut!(stream);
|
||||
let frames: Vec<String> = stream.collect().await;
|
||||
assert_eq!(frames.len(), 4);
|
||||
let first: serde_json::Value = serde_json::from_str(&frames[0]).unwrap();
|
||||
let terminal: serde_json::Value = serde_json::from_str(&frames[1]).unwrap();
|
||||
let usage: serde_json::Value = serde_json::from_str(&frames[2]).unwrap();
|
||||
assert_eq!(first["choices"][0]["text"], "a");
|
||||
assert_eq!(terminal["choices"][0]["text"], "b");
|
||||
assert_eq!(terminal["choices"][0]["finish_reason"], "stop");
|
||||
assert!(usage["choices"].as_array().unwrap().is_empty());
|
||||
assert_eq!(usage["usage"]["prompt_tokens"], 5);
|
||||
assert_eq!(usage["usage"]["completion_tokens"], 2);
|
||||
assert_eq!(frames[3], "[DONE]");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
//! OpenAI model discovery endpoints.
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
routing::get,
|
||||
};
|
||||
|
||||
use super::{AppState, openai_error, unix_seconds_u32};
|
||||
|
||||
pub(super) fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/v1/models", get(available_models))
|
||||
.route("/v1/models/{model}", get(retrieve_model))
|
||||
}
|
||||
|
||||
/// `GET /v1/models` — OpenAI-compatible model list. Served from `server_args`;
|
||||
/// no scheduler round-trip.
|
||||
async fn available_models(State(state): State<AppState>) -> Response {
|
||||
let base = model_card(&state);
|
||||
Json(serde_json::json!({ "object": "list", "data": [base] })).into_response()
|
||||
}
|
||||
|
||||
async fn retrieve_model(State(state): State<AppState>, Path(model): Path<String>) -> Response {
|
||||
if model != state.server_args.served_model_name {
|
||||
return openai_error(
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("The model `{model}` does not exist"),
|
||||
);
|
||||
}
|
||||
Json(model_card(&state)).into_response()
|
||||
}
|
||||
|
||||
fn model_card(state: &AppState) -> serde_json::Value {
|
||||
let name = &state.server_args.served_model_name;
|
||||
serde_json::json!({
|
||||
"id": name,
|
||||
"object": "model",
|
||||
"created": unix_seconds_u32(),
|
||||
"owned_by": "sglang",
|
||||
"root": name,
|
||||
"parent": serde_json::Value::Null,
|
||||
"max_model_len": state.server_args.model_config.context_len,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
//! Reasoning-content splitting for Chat Completions (`--reasoning-parser`).
|
||||
//!
|
||||
//! Mirrors the Python frontend (`sglang.srt.parser.reasoning_parser` +
|
||||
//! `serving_chat._process_reasoning_stream`): when the
|
||||
//! server was launched with `--reasoning-parser <name>` (and the request keeps
|
||||
//! the default `separate_reasoning=true`, which the Dynamo request type cannot
|
||||
//! express), the model's `<think>`-style markers are stripped out of `content`
|
||||
//! into `reasoning_content` — for unary responses and streaming deltas alike.
|
||||
//!
|
||||
//! The parser lifecycle (lazy build, per-frame incremental split, terminal
|
||||
//! flush of *both* buffered columns) lives here so the endpoint cannot drop
|
||||
//! the tail half.
|
||||
|
||||
use dynamo_parsers::reasoning::{
|
||||
ReasoningParser as _, ReasoningParserType, ReasoningParserWrapper,
|
||||
};
|
||||
|
||||
/// Build the parser the Python `--reasoning-parser` name selects.
|
||||
///
|
||||
/// The names come from Python's `ReasoningParser.DetectorMap`, which differs
|
||||
/// from the dynamo-parsers registry keys in a few spellings (deepseek-r1 vs
|
||||
/// deepseek_r1, kimi_k2 vs kimi_k25, …) and has a few entries that Python maps
|
||||
/// onto a forced-reasoning `<think>` parser (qwen3-thinking, minimax). Names
|
||||
/// dynamo does not know (hunyuan, inkling, apertus2509, mimo, poolside_v1,
|
||||
/// cohere_command4 — all tokenizer-driven parsers) fall through to the
|
||||
/// registry, which warns and falls back to the non-forced Basic parser.
|
||||
pub(super) fn build_reasoning_parser(server_name: &str) -> ReasoningParserWrapper {
|
||||
let name = match server_name {
|
||||
// Python DetectorMap spellings that differ from the dynamo registry keys.
|
||||
"deepseek-r1" | "step3p5" => "deepseek_r1",
|
||||
"kimi_k2" => "kimi_k25",
|
||||
"gpt-oss" => "gpt_oss",
|
||||
"nemotron_3" => "nemotron3",
|
||||
"interns1" => "qwen3",
|
||||
// Python forces reasoning for these; the R1 parser is the same
|
||||
// `<think>` / `</think>` configuration with `force_reasoning=true`.
|
||||
"qwen3-thinking" | "minimax" => "deepseek_r1",
|
||||
_ => server_name,
|
||||
};
|
||||
ReasoningParserType::get_reasoning_parser_from_name(name)
|
||||
}
|
||||
|
||||
/// Split a completed generation's text into `(reasoning_text, normal_text)`
|
||||
/// when `--reasoning-parser` selects a parser; otherwise the text passes
|
||||
/// through untouched as normal text. Chat splits before tool-call parsing.
|
||||
pub(super) fn split_reasoning_unary(
|
||||
name: Option<&str>,
|
||||
text: &str,
|
||||
token_ids: &[i32],
|
||||
) -> (String, String) {
|
||||
let Some(name) = name else {
|
||||
return (String::new(), text.to_owned());
|
||||
};
|
||||
let mut parser = build_reasoning_parser(name);
|
||||
let token_ids = token_ids
|
||||
.iter()
|
||||
.filter_map(|&id| u32::try_from(id).ok())
|
||||
.collect::<Vec<_>>();
|
||||
let split = parser.detect_and_parse_reasoning(text, &token_ids);
|
||||
(split.reasoning_text, split.normal_text)
|
||||
}
|
||||
|
||||
/// Stateful reasoning split for one streaming response. Mirrors Python's
|
||||
/// `reasoning_parser_dict` entries: the parser is built lazily on the first
|
||||
/// content delta, each frame is split into `(reasoning, normal)` deltas, and
|
||||
/// [`finish`](Self::finish) flushes the parser-buffered tail — *both* columns,
|
||||
/// since the buffered text can sit in either one (e.g. MiniMax M3's
|
||||
/// implicit-tool-start recovery holds the leading answer text until the think
|
||||
/// opener or a tool marker establishes the mode, and releases it as normal
|
||||
/// text at EOF).
|
||||
#[derive(Default)]
|
||||
pub(super) struct ReasoningStreamSplitter {
|
||||
name: Option<String>,
|
||||
parser: Option<ReasoningParserWrapper>,
|
||||
}
|
||||
|
||||
impl ReasoningStreamSplitter {
|
||||
pub(super) fn new(name: Option<&str>) -> Self {
|
||||
Self {
|
||||
name: name.map(str::to_owned),
|
||||
parser: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Split one frame's text into `(reasoning_text, normal_text)` deltas.
|
||||
pub(super) fn split(&mut self, text: &str, token_ids: &[i32]) -> (String, String) {
|
||||
let Some(name) = self.name.as_deref() else {
|
||||
return (String::new(), text.to_owned());
|
||||
};
|
||||
let parser = self
|
||||
.parser
|
||||
.get_or_insert_with(|| build_reasoning_parser(name));
|
||||
let token_ids = token_ids
|
||||
.iter()
|
||||
.filter_map(|&id| u32::try_from(id).ok())
|
||||
.collect::<Vec<_>>();
|
||||
let split = parser.parse_reasoning_streaming_incremental(text, &token_ids);
|
||||
(split.reasoning_text, split.normal_text)
|
||||
}
|
||||
|
||||
/// Flush the parser-buffered tail at stream end, releasing both columns.
|
||||
pub(super) fn finish(&mut self) -> (String, String) {
|
||||
let Some(parser) = self.parser.as_mut() else {
|
||||
return (String::new(), String::new());
|
||||
};
|
||||
let tail = parser.finish_reasoning_stream();
|
||||
(tail.reasoning_text, tail.normal_text)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{ReasoningStreamSplitter, build_reasoning_parser, split_reasoning_unary};
|
||||
use dynamo_parsers::reasoning::ReasoningParser;
|
||||
|
||||
#[test]
|
||||
fn python_deepseek_r1_name_splits_forced_reasoning() {
|
||||
let mut parser = build_reasoning_parser("deepseek-r1");
|
||||
// Forced: text before any marker is reasoning.
|
||||
let split = parser.detect_and_parse_reasoning("think hard</think>Paris", &[]);
|
||||
assert_eq!(split.reasoning_text, "think hard");
|
||||
assert_eq!(split.normal_text, "Paris");
|
||||
let split = parser.detect_and_parse_reasoning("<think>yes</think>answer", &[]);
|
||||
assert_eq!(split.reasoning_text, "yes");
|
||||
assert_eq!(split.normal_text, "answer");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn python_kimi_k2_name_maps_to_kimi_k25() {
|
||||
let mut parser = build_reasoning_parser("kimi_k2");
|
||||
let split = parser.detect_and_parse_reasoning("<think>k</think>out", &[]);
|
||||
assert_eq!(split.reasoning_text, "k");
|
||||
assert_eq!(split.normal_text, "out");
|
||||
// Kimi-K2.5 interrupts reasoning at the tool-call section marker.
|
||||
let mut parser = build_reasoning_parser("kimi_k2");
|
||||
let split =
|
||||
parser.detect_and_parse_reasoning("reasons<|tool_calls_section_begin|>calls", &[]);
|
||||
assert_eq!(split.reasoning_text, "reasons");
|
||||
assert_eq!(split.normal_text, "<|tool_calls_section_begin|>calls");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen3_thinking_forces_reasoning_like_python() {
|
||||
let mut parser = build_reasoning_parser("qwen3-thinking");
|
||||
let split = parser.detect_and_parse_reasoning("plain text", &[]);
|
||||
assert_eq!(split.reasoning_text, "plain text");
|
||||
assert_eq!(split.normal_text, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streaming_split_keeps_markers_out_of_both_columns() {
|
||||
let mut parser = build_reasoning_parser("deepseek-r1");
|
||||
let mut reasoning = String::new();
|
||||
let mut normal = String::new();
|
||||
for chunk in ["<think>rea", "son</think>an", "swer"] {
|
||||
let split = parser.parse_reasoning_streaming_incremental(chunk, &[]);
|
||||
reasoning.push_str(&split.reasoning_text);
|
||||
normal.push_str(&split.normal_text);
|
||||
}
|
||||
let tail = parser.finish_reasoning_stream();
|
||||
reasoning.push_str(&tail.reasoning_text);
|
||||
normal.push_str(&tail.normal_text);
|
||||
assert_eq!(reasoning, "reason");
|
||||
assert_eq!(normal, "answer");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unary_split_passes_text_through_without_a_parser() {
|
||||
let (reasoning, normal) = split_reasoning_unary(None, "<think>kept as text</think>", &[1]);
|
||||
assert_eq!(reasoning, "");
|
||||
assert_eq!(normal, "<think>kept as text</think>");
|
||||
}
|
||||
|
||||
/// REASONING_P1: MiniMax M3's implicit-tool-start recovery buffers the
|
||||
/// answer text until a boundary establishes the mode; with no opener the
|
||||
/// whole buffer is released as normal text only at `finish`. The chat
|
||||
/// terminal flush must emit the normal half of the tail.
|
||||
#[test]
|
||||
fn streaming_tail_releases_normal_text_only_at_finish() {
|
||||
let mut splitter = ReasoningStreamSplitter::new(Some("minimax_m3"));
|
||||
let (reasoning, normal) = splitter.split("The answer is", &[]);
|
||||
assert_eq!(reasoning, "");
|
||||
assert_eq!(normal, "", "M3 holds the ambiguous prefix until a boundary");
|
||||
let (reasoning, normal) = splitter.split(" 42", &[]);
|
||||
assert_eq!(reasoning, "");
|
||||
assert_eq!(normal, "");
|
||||
let (reasoning, normal) = splitter.finish();
|
||||
assert_eq!(reasoning, "");
|
||||
assert_eq!(normal, "The answer is 42");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streaming_tail_releases_reasoning_after_marker_boundary() {
|
||||
let mut splitter = ReasoningStreamSplitter::new(Some("minimax_m3"));
|
||||
let (reasoning, normal) = splitter.split("<mm:think>think", &[]);
|
||||
assert_eq!(reasoning, "think");
|
||||
assert_eq!(normal, "");
|
||||
let (reasoning, normal) = splitter.split(" hard</mm:think>", &[]);
|
||||
assert_eq!(reasoning, " hard");
|
||||
assert_eq!(normal, "");
|
||||
let (reasoning, normal) = splitter.finish();
|
||||
assert_eq!(reasoning, "");
|
||||
assert_eq!(normal, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finish_without_a_parser_is_empty() {
|
||||
let mut splitter = ReasoningStreamSplitter::new(None);
|
||||
let (reasoning, normal) = splitter.split("plain", &[]);
|
||||
assert_eq!(reasoning, "");
|
||||
assert_eq!(normal, "plain");
|
||||
let (reasoning, normal) = splitter.finish();
|
||||
assert_eq!(reasoning, "");
|
||||
assert_eq!(normal, "");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,319 @@
|
||||
//! Shared HTTP test harness and `openai.rs`-level handler tests.
|
||||
//!
|
||||
//! Submodule tests live next to the code they cover: `chat`, `completions`,
|
||||
//! `tools`, and `reasoning` each carry their own
|
||||
//! `#[cfg(test)] mod tests`. This module keeps the fixtures they all share —
|
||||
//! channel fixtures (`senders`, `chunk`, `submitted`, `chat_submitted`) and the
|
||||
//! full-router harness (`server_args`, `app_state`,
|
||||
//! `oneshot`, `post_json`, `body_json`) — plus the handler-level tests that
|
||||
//! exercise [`routes`] end to end. The helpers are `pub(super)` so sibling
|
||||
//! test modules can import them via `super::super::test_utils::*`.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use axum::response::Response;
|
||||
use serde_json::json;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use super::routes;
|
||||
use crate::ids::Rid;
|
||||
use crate::message::{ChunkEvent, EgressItem};
|
||||
use crate::runtime::ServerArgs;
|
||||
use crate::tokenizer_manager::Senders;
|
||||
|
||||
pub(super) fn senders() -> Senders {
|
||||
Senders {
|
||||
tm: flume::unbounded().0,
|
||||
abort: flume::unbounded().0,
|
||||
tok: flume::unbounded().0,
|
||||
detok: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn chunk(rid: &str, text: &str, done: bool) -> EgressItem {
|
||||
let output = ChunkEvent {
|
||||
rid: rid.into(),
|
||||
text: text.into(),
|
||||
token_ids: vec![1],
|
||||
prompt_tokens: 5,
|
||||
completion_tokens: 1,
|
||||
finish_reason: done.then(|| {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"type": "stop",
|
||||
"matched": "</s>"
|
||||
}))
|
||||
.unwrap()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
if done {
|
||||
EgressItem::Done(output)
|
||||
} else {
|
||||
EgressItem::Frame(output)
|
||||
}
|
||||
}
|
||||
|
||||
/// A submitted legacy completion choice with its egress channel.
|
||||
pub(super) fn submitted(
|
||||
index: usize,
|
||||
prompt_index: usize,
|
||||
rid: &str,
|
||||
) -> (
|
||||
super::completions::SubmittedChoice,
|
||||
tokio::sync::mpsc::Sender<EgressItem>,
|
||||
) {
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(8);
|
||||
(
|
||||
super::completions::SubmittedChoice {
|
||||
index,
|
||||
prompt_index,
|
||||
rid: rid.into(),
|
||||
echo: String::new(),
|
||||
rx,
|
||||
},
|
||||
tx,
|
||||
)
|
||||
}
|
||||
|
||||
/// A submitted chat choice (the tuple `chat_event_stream` consumes) with its
|
||||
/// egress channel.
|
||||
pub(super) fn chat_submitted(
|
||||
index: usize,
|
||||
rid: &str,
|
||||
) -> (
|
||||
(usize, Rid, tokio::sync::mpsc::Receiver<EgressItem>),
|
||||
tokio::sync::mpsc::Sender<EgressItem>,
|
||||
) {
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(8);
|
||||
((index, rid.into(), rx), tx)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------
|
||||
// Handler-level tests: full router, real extractors, no scheduler. A
|
||||
// request that reaches `submit` with an OPEN tm lane would wait on the
|
||||
// egress receiver forever, so submission-reaching cases use `senders_closed`
|
||||
// (503) and everything else fails validation before submit.
|
||||
// ---------------------------------------------------------------------
|
||||
|
||||
pub(super) fn server_args() -> Arc<ServerArgs> {
|
||||
Arc::new(
|
||||
serde_json::from_value(serde_json::json!({ "served_model_name": "model" }))
|
||||
.expect("ServerArgs must deserialize"),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn app_state(senders: Senders) -> super::AppState {
|
||||
super::AppState {
|
||||
senders,
|
||||
egress_buf: 8,
|
||||
server_args: server_args(),
|
||||
chat_formatter: None,
|
||||
egress_activity: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn senders_closed() -> Senders {
|
||||
// Dropping the receivers disconnects the channels; the senders stay
|
||||
// valid (moveable) but every send reports `Err`, the shutdown state
|
||||
// `submit` surfaces as a 503.
|
||||
let (tm_tx, tm_rx) = flume::unbounded();
|
||||
drop(tm_rx);
|
||||
let (abort_tx, abort_rx) = flume::unbounded();
|
||||
drop(abort_rx);
|
||||
let (tok_tx, tok_rx) = flume::unbounded();
|
||||
drop(tok_rx);
|
||||
Senders {
|
||||
tm: tm_tx,
|
||||
abort: abort_tx,
|
||||
tok: tok_tx,
|
||||
detok: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
/// Serve one request through the full router (extractors, auth, routing).
|
||||
/// `with_state` consumes the state into a `Router<()>`, which is what
|
||||
/// implements `tower::Service`.
|
||||
pub(super) async fn oneshot(app: Router<()>, req: Request<Body>) -> Response {
|
||||
app.oneshot(req).await.unwrap()
|
||||
}
|
||||
|
||||
pub(super) async fn post_json(app: Router<()>, path: &str, body: serde_json::Value) -> Response {
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri(path)
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap();
|
||||
oneshot(app, req).await
|
||||
}
|
||||
|
||||
pub(super) async fn body_json(response: Response) -> serde_json::Value {
|
||||
let bytes = axum::body::to_bytes(response.into_body(), 64 * 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
/// The common StatusCode→error helper follows `pre_submit_error`'s shape:
|
||||
/// unary requests get the JSON error with its status; a committed stream gets
|
||||
/// 200 + one SSE error frame + `[DONE]`, and the frame carries the OpenAI
|
||||
/// error fields (`type`, `param`, `code`) that the SDKs dispatch on.
|
||||
#[tokio::test]
|
||||
async fn openai_error_response_covers_unary_and_sse() {
|
||||
let unary = super::openai_error_response(StatusCode::BAD_REQUEST, "bad input", false);
|
||||
assert_eq!(unary.status(), StatusCode::BAD_REQUEST);
|
||||
let value = body_json(unary).await;
|
||||
assert_eq!(value["error"]["message"], "bad input");
|
||||
assert_eq!(value["error"]["type"], "BadRequestError");
|
||||
assert_eq!(value["error"]["code"], 400);
|
||||
assert!(value["error"]["param"].is_null());
|
||||
|
||||
let streamed = super::openai_error_response(StatusCode::BAD_REQUEST, "bad input", true);
|
||||
assert_eq!(streamed.status(), StatusCode::OK);
|
||||
let bytes = axum::body::to_bytes(streamed.into_body(), 64 * 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
let text = String::from_utf8(bytes.to_vec()).unwrap();
|
||||
let frame = text
|
||||
.split("\n\n")
|
||||
.next()
|
||||
.unwrap()
|
||||
.strip_prefix("data: ")
|
||||
.unwrap();
|
||||
let frame: serde_json::Value = serde_json::from_str(frame).unwrap();
|
||||
assert_eq!(frame["error"]["message"], "bad input");
|
||||
assert_eq!(frame["error"]["type"], "BadRequestError");
|
||||
assert!(text.contains("[DONE]"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn completions_handler_validates_before_submit() {
|
||||
let app = routes().with_state(app_state(senders()));
|
||||
let cases = [
|
||||
(json!({"model": "other", "prompt": "hi"}), "unknown model"),
|
||||
(json!({"model": "model", "prompt": "hi", "n": 0}), "n=0"),
|
||||
(
|
||||
json!({"model": "model", "prompt": "hi", "max_tokens": 0}),
|
||||
"max_tokens=0",
|
||||
),
|
||||
(json!({"model": "model", "prompt": ""}), "empty prompt"),
|
||||
(
|
||||
json!({"model": "model", "prompt": "hi", "best_of": 2}),
|
||||
"best_of>1",
|
||||
),
|
||||
(
|
||||
json!({"model": "model", "prompt": "hi", "suffix": "x"}),
|
||||
"suffix",
|
||||
),
|
||||
(
|
||||
json!({"model": "model", "prompt": "hi", "prompt_embeds": [[1.0]]}),
|
||||
"prompt_embeds",
|
||||
),
|
||||
];
|
||||
for (body, label) in cases {
|
||||
let response = post_json(app.clone(), "/v1/completions", body).await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{label}");
|
||||
}
|
||||
// Malformed JSON → 400 (JsonRejection path).
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/completions")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from("not json"))
|
||||
.unwrap();
|
||||
let response = oneshot(app.clone(), req).await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
// A closed tm inbox (shutdown) surfaces as 503.
|
||||
let app = routes().with_state(app_state(senders_closed()));
|
||||
let response = post_json(
|
||||
app.clone(),
|
||||
"/v1/completions",
|
||||
json!({"model": "model", "prompt": "hi"}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chat_handler_validates_before_submit() {
|
||||
let app = routes().with_state(app_state(senders()));
|
||||
let cases = [
|
||||
(
|
||||
json!({"model": "other", "messages": [{"role": "user", "content": "hi"}]}),
|
||||
"unknown model",
|
||||
),
|
||||
(json!({"model": "model", "messages": []}), "empty messages"),
|
||||
(
|
||||
json!({"model": "model", "messages": [{"role": "user", "content": "hi"}], "n": 0}),
|
||||
"n=0",
|
||||
),
|
||||
(
|
||||
json!({"model": "model", "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "http://example.com/x.png"}}]}]}),
|
||||
"media content",
|
||||
),
|
||||
(
|
||||
json!({"model": "model", "messages": [{"role": "user", "content": "hi"}], "function_call": "auto"}),
|
||||
"deprecated function_call",
|
||||
),
|
||||
(
|
||||
json!({"model": "model", "messages": [{"role": "user", "content": "hi"}], "audio": {"input_audio": {"data": "x", "format": "wav"}}}),
|
||||
"audio",
|
||||
),
|
||||
(
|
||||
json!({"model": "model", "messages": [{"role": "user", "content": "hi"}], "max_completion_tokens": 0}),
|
||||
"max_completion_tokens=0",
|
||||
),
|
||||
];
|
||||
for (body, label) in cases {
|
||||
let response = post_json(app.clone(), "/v1/chat/completions", body).await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{label}");
|
||||
}
|
||||
// A valid request with no loaded chat template → 400 (template gate).
|
||||
let response = post_json(
|
||||
app.clone(),
|
||||
"/v1/chat/completions",
|
||||
json!({"model": "model", "messages": [{"role": "user", "content": "hi"}]}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn basic_openai_router_excludes_responses_api() {
|
||||
let app = routes().with_state(app_state(senders()));
|
||||
let response = post_json(app, "/v1/responses", json!({"input": "hi"})).await;
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
/// A closed tm inbox with a *streaming* request must answer inside the
|
||||
/// committed stream: 200 + one OpenAI-shaped SSE error frame + `[DONE]` (the
|
||||
/// same rule `pre_submit_error` applies to the native API), not a unary 503.
|
||||
#[tokio::test]
|
||||
async fn streaming_submit_failure_answers_inside_the_stream() {
|
||||
let app = routes().with_state(app_state(senders_closed()));
|
||||
let response = post_json(
|
||||
app,
|
||||
"/v1/completions",
|
||||
json!({"model": "model", "prompt": "hi", "stream": true}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let bytes = axum::body::to_bytes(response.into_body(), 64 * 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
let text = String::from_utf8(bytes.to_vec()).unwrap();
|
||||
let frame = text
|
||||
.split("\n\n")
|
||||
.next()
|
||||
.unwrap()
|
||||
.strip_prefix("data: ")
|
||||
.unwrap();
|
||||
let frame: serde_json::Value = serde_json::from_str(frame).unwrap();
|
||||
assert_eq!(frame["error"]["message"], "service unavailable");
|
||||
assert_eq!(frame["error"]["type"], "InternalServerError");
|
||||
assert_eq!(frame["error"]["code"], 503);
|
||||
assert!(text.contains("[DONE]"));
|
||||
}
|
||||
@@ -0,0 +1,836 @@
|
||||
//! Tool-choice constraints and unary tool-call parsing.
|
||||
//!
|
||||
//! Two complementary mechanisms, mirroring the Python frontend
|
||||
//! (`serving_chat.py` + dynamo-parsers):
|
||||
//!
|
||||
//! - [`apply_tool_constraint`] turns `tool_choice` into a sampling constraint
|
||||
//! *before* submission: a `structural_tag` when the parser supports one (or
|
||||
//! when strict tools need the llama3 triggered-tag format), otherwise a
|
||||
//! `json_schema` array restricting the output to tool calls. It validates
|
||||
//! `tool_choice`/`tools` agreement first.
|
||||
//! - [`parse_chat_tool_calls`] strips the model's tool-call markers out of a
|
||||
//! finished response. Streaming responses use Dynamo's
|
||||
//! `apply_tool_calling_jail` directly in `chat.rs`.
|
||||
//!
|
||||
//! [`dynamo_parser_name`] canonicalizes the SGLang CLI parser names onto the
|
||||
//! dynamo-parsers registry keys, [`chat_delta`] builds the stream deltas these
|
||||
//! paths emit, and [`chat_finish_reason`] maps the scheduler's finish reason
|
||||
//! onto the OpenAI wire values.
|
||||
//!
|
||||
//! # Test coverage
|
||||
//!
|
||||
//! The tests cover parser-name canonicalization, every `tool_choice` branch of
|
||||
//! [`apply_tool_constraint`], Dynamo's streaming jail integration, unary
|
||||
//! parsing, and finish-reason mapping.
|
||||
|
||||
use dynamo_parsers::parsers::get_tool_parser_map;
|
||||
use dynamo_parsers::{
|
||||
StructuralTagBuilder, StructuralTagSchemaMode, ToolCallFormatBuildContext,
|
||||
ToolChoice as DynamoToolChoice, ToolDefinition, TriggeredTagsConfig,
|
||||
try_tool_call_parse_aggregate_finalize,
|
||||
};
|
||||
use dynamo_protocols::types::{
|
||||
ChatCompletionMessageContent, ChatCompletionMessageToolCall,
|
||||
ChatCompletionMessageToolCallChunk, ChatCompletionStreamResponseDelta,
|
||||
ChatCompletionToolChoiceOption, FinishReason as OpenAIFinishReason, FunctionCall, FunctionType,
|
||||
Role,
|
||||
};
|
||||
|
||||
use crate::message::{ChunkEvent, SamplingParams};
|
||||
|
||||
/// Canonicalize a tool-call parser name onto the dynamo-parsers registry keys.
|
||||
///
|
||||
/// SGLang canonicalizes these legacy CLI names in the opposite direction from
|
||||
/// the current Dynamo parser registry.
|
||||
pub(super) fn dynamo_parser_name(parser: &str) -> &str {
|
||||
match parser {
|
||||
"llama3" => "llama3_json",
|
||||
"qwen" => "qwen25",
|
||||
"glm" | "glm45" => "glm47",
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
/// Map the OpenAI wire `tool_choice` onto the Dynamo choice. A missing/auto
|
||||
/// choice reads as `Auto`.
|
||||
pub(super) fn dynamo_tool_choice(
|
||||
choice: &Option<ChatCompletionToolChoiceOption>,
|
||||
) -> DynamoToolChoice {
|
||||
match choice {
|
||||
Some(ChatCompletionToolChoiceOption::None) => DynamoToolChoice::None,
|
||||
Some(ChatCompletionToolChoiceOption::Required) => DynamoToolChoice::Required,
|
||||
Some(ChatCompletionToolChoiceOption::Named(choice)) => {
|
||||
DynamoToolChoice::Named(choice.function.name.clone())
|
||||
}
|
||||
Some(ChatCompletionToolChoiceOption::Auto) | None => DynamoToolChoice::Auto,
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate `tool_choice` against `tools`, then — when a tool-call `parser`
|
||||
/// is configured — turn it into a sampling constraint, mirroring Python's
|
||||
/// `serving_chat` logic. Validation runs even without a parser, so an invalid
|
||||
/// choice (required/named with nothing to select) is rejected before
|
||||
/// submission in every mode.
|
||||
///
|
||||
/// Prefers a structural-tag constraint: the parser's own registered builder,
|
||||
/// or — for llama3 with strict tools under `auto` — a triggered-tag builder
|
||||
/// so the model emits calls in the exact `<|python_tag|>` format. Otherwise
|
||||
/// `required`/`named` choices fall back to a JSON-schema array constraining
|
||||
/// the output to `{"name", "parameters"}` objects (`maxItems: 1` when
|
||||
/// `parallel_tool_calls` is false).
|
||||
pub(super) fn apply_tool_constraint(
|
||||
sampling: &mut SamplingParams,
|
||||
parser: Option<&str>,
|
||||
tool_choice: &DynamoToolChoice,
|
||||
tools: &[ToolDefinition],
|
||||
parallel_tool_calls: Option<bool>,
|
||||
) -> Result<(), String> {
|
||||
if *tool_choice == DynamoToolChoice::None {
|
||||
return Ok(());
|
||||
}
|
||||
if *tool_choice == DynamoToolChoice::Required && tools.is_empty() {
|
||||
return Err("tool_choice is \"required\" but tools is empty".into());
|
||||
}
|
||||
if let DynamoToolChoice::Named(name) = tool_choice
|
||||
&& !tools.iter().any(|tool| &tool.name == name)
|
||||
{
|
||||
return Err(format!(
|
||||
"tool named \"{name}\" in tool_choice is not present in tools"
|
||||
));
|
||||
}
|
||||
|
||||
let Some(parser) = parser else {
|
||||
return Ok(()); // validation only
|
||||
};
|
||||
let parser = dynamo_parser_name(parser);
|
||||
let config = get_tool_parser_map()
|
||||
.get(parser)
|
||||
.ok_or_else(|| format!("tool-call parser `{parser}` is not supported by Dynamo"))?;
|
||||
let builder = config.structural_tag_builder.clone().or_else(|| {
|
||||
(parser == "llama3_json"
|
||||
&& *tool_choice == DynamoToolChoice::Auto
|
||||
&& tools.iter().any(|tool| tool.strict.unwrap_or(false)))
|
||||
.then(|| {
|
||||
StructuralTagBuilder::TriggeredTags(TriggeredTagsConfig {
|
||||
begin_template: r#"<|python_tag|>{"name":"{name}", "arguments":"#.to_string(),
|
||||
end_template: "}".to_string(),
|
||||
triggers: vec!["<|python_tag|>".to_string()],
|
||||
content_style: Default::default(),
|
||||
tool_call_ban_tokens: Vec::new(),
|
||||
reasoning_end: None,
|
||||
})
|
||||
})
|
||||
});
|
||||
if let Some(builder) = builder
|
||||
&& let Some(tag) = builder
|
||||
.build_tool_call_format(&ToolCallFormatBuildContext {
|
||||
tool_choice,
|
||||
tools,
|
||||
parallel_tool_calls,
|
||||
schema_mode: StructuralTagSchemaMode::Auto,
|
||||
starts_in_reasoning: false,
|
||||
})
|
||||
.map_err(|error| error.to_string())?
|
||||
{
|
||||
sampling.structural_tag = Some(tag.to_string());
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if matches!(
|
||||
tool_choice,
|
||||
DynamoToolChoice::Required | DynamoToolChoice::Named(_)
|
||||
) {
|
||||
let selected = match tool_choice {
|
||||
DynamoToolChoice::Named(name) => tools
|
||||
.iter()
|
||||
.filter(|tool| tool.name == *name)
|
||||
.collect::<Vec<_>>(),
|
||||
_ => tools.iter().collect(),
|
||||
};
|
||||
let schemas = selected
|
||||
.into_iter()
|
||||
.map(|tool| {
|
||||
serde_json::json!({
|
||||
"properties": {
|
||||
"name": {"type": "string", "enum": [tool.name]},
|
||||
"parameters": tool.parameters.clone().unwrap_or_else(|| {
|
||||
serde_json::json!({"type": "object", "properties": {}})
|
||||
}),
|
||||
},
|
||||
"required": ["name", "parameters"],
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let items = if schemas.len() == 1 {
|
||||
schemas.into_iter().next().expect("one schema")
|
||||
} else {
|
||||
serde_json::json!({"type": "object", "anyOf": schemas})
|
||||
};
|
||||
let mut schema = serde_json::json!({
|
||||
"type": "array",
|
||||
"minItems": 1,
|
||||
"items": items,
|
||||
});
|
||||
if parallel_tool_calls == Some(false) {
|
||||
schema["maxItems"] = serde_json::json!(1);
|
||||
}
|
||||
sampling.json_schema = Some(schema.to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build a chat-streaming delta carrying any of the optional columns.
|
||||
///
|
||||
/// The deprecated `function_call` field stays `None` — tool calls go through
|
||||
/// the `tool_calls` array.
|
||||
#[allow(deprecated)]
|
||||
pub(super) fn chat_delta(
|
||||
content: Option<String>,
|
||||
role: Option<Role>,
|
||||
tool_calls: Option<Vec<ChatCompletionMessageToolCallChunk>>,
|
||||
reasoning_content: Option<String>,
|
||||
) -> ChatCompletionStreamResponseDelta {
|
||||
ChatCompletionStreamResponseDelta {
|
||||
content: content.map(ChatCompletionMessageContent::Text),
|
||||
function_call: None,
|
||||
tool_calls,
|
||||
role,
|
||||
refusal: None,
|
||||
reasoning_content,
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse tool calls out of a completed (unary) generation's content.
|
||||
///
|
||||
/// Returns `(content, None)` when no parser is configured or no call parses —
|
||||
/// the content passes through untouched. With a parser, a successful parse
|
||||
/// returns the leftover non-tool text and the calls; `parallel_tool_calls`
|
||||
/// false truncates the batch to the first call, mirroring Python.
|
||||
pub(super) async fn parse_chat_tool_calls(
|
||||
content: String,
|
||||
parser: Option<&str>,
|
||||
tools: Option<&[ToolDefinition]>,
|
||||
parallel_tool_calls: bool,
|
||||
) -> (String, Option<Vec<ChatCompletionMessageToolCall>>) {
|
||||
let Some(parser) = parser else {
|
||||
return (content, None);
|
||||
};
|
||||
let parser = dynamo_parser_name(parser);
|
||||
match try_tool_call_parse_aggregate_finalize(&content, Some(parser), tools).await {
|
||||
Ok((mut calls, normal)) if !calls.is_empty() => {
|
||||
if !parallel_tool_calls {
|
||||
calls.truncate(1);
|
||||
}
|
||||
(
|
||||
normal.unwrap_or_default(),
|
||||
Some(
|
||||
calls
|
||||
.into_iter()
|
||||
.map(|call| ChatCompletionMessageToolCall {
|
||||
id: call.id,
|
||||
r#type: FunctionType::Function,
|
||||
function: FunctionCall {
|
||||
name: call.function.name,
|
||||
arguments: call.function.arguments,
|
||||
},
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
)
|
||||
}
|
||||
_ => (content, None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Map the scheduler's finish kind onto the OpenAI wire values. Length and
|
||||
/// content-filter keep their names; everything else (including a bare abort)
|
||||
/// reports as `stop`, matching Python's fallback.
|
||||
pub(super) fn chat_finish_reason(output: &ChunkEvent) -> Option<OpenAIFinishReason> {
|
||||
let kind = output
|
||||
.finish_reason
|
||||
.as_ref()
|
||||
.and_then(|reason| reason.kind_name());
|
||||
kind.map(|kind| match kind {
|
||||
"length" => OpenAIFinishReason::Length,
|
||||
"content_filter" => OpenAIFinishReason::ContentFilter,
|
||||
_ => OpenAIFinishReason::Stop,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
apply_tool_constraint, chat_delta, chat_finish_reason, dynamo_parser_name,
|
||||
dynamo_tool_choice, parse_chat_tool_calls,
|
||||
};
|
||||
use crate::message::{ChunkEvent, SamplingParams};
|
||||
use dynamo_parsers::tool_calling::jail::{Annotated, apply_tool_calling_jail};
|
||||
use dynamo_parsers::{ToolChoice as DynamoToolChoice, ToolDefinition};
|
||||
use dynamo_protocols::types::CreateChatCompletionStreamResponse as StreamResponse;
|
||||
use dynamo_protocols::types::{
|
||||
ChatChoiceStream, ChatCompletionMessageContent, ChatCompletionMessageToolCallChunk,
|
||||
ChatCompletionNamedToolChoice, ChatCompletionToolChoiceOption, ChatCompletionToolType,
|
||||
FinishReason as OpenAIFinishReason, FunctionCallStream, FunctionName, FunctionType, Role,
|
||||
};
|
||||
use futures::{StreamExt, stream};
|
||||
|
||||
fn tool(name: &str, strict: bool) -> ToolDefinition {
|
||||
ToolDefinition {
|
||||
name: name.into(),
|
||||
parameters: Some(serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"]
|
||||
})),
|
||||
strict: Some(strict),
|
||||
}
|
||||
}
|
||||
|
||||
fn stream_item(text: &str, finish: Option<OpenAIFinishReason>) -> Annotated<StreamResponse> {
|
||||
Annotated {
|
||||
data: Some(StreamResponse {
|
||||
id: "chatcmpl-test".into(),
|
||||
choices: vec![ChatChoiceStream {
|
||||
index: 0,
|
||||
delta: chat_delta(Some(text.into()), Some(Role::Assistant), None, None),
|
||||
finish_reason: finish,
|
||||
logprobs: None,
|
||||
}],
|
||||
created: 1,
|
||||
model: "model".into(),
|
||||
service_tier: None,
|
||||
system_fingerprint: None,
|
||||
object: "chat.completion.chunk".into(),
|
||||
usage: None,
|
||||
}),
|
||||
id: None,
|
||||
event: None,
|
||||
comment: None,
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// A terminal chunk with no text — what the upstream path emits for an
|
||||
/// empty `Done` frame (content `None`, not an empty string).
|
||||
fn stream_done(finish: OpenAIFinishReason) -> Annotated<StreamResponse> {
|
||||
Annotated {
|
||||
data: Some(StreamResponse {
|
||||
id: "chatcmpl-test".into(),
|
||||
choices: vec![ChatChoiceStream {
|
||||
index: 0,
|
||||
delta: chat_delta(None, Some(Role::Assistant), None, None),
|
||||
finish_reason: Some(finish),
|
||||
logprobs: None,
|
||||
}],
|
||||
created: 1,
|
||||
model: "model".into(),
|
||||
service_tier: None,
|
||||
system_fingerprint: None,
|
||||
object: "chat.completion.chunk".into(),
|
||||
usage: None,
|
||||
}),
|
||||
id: None,
|
||||
event: None,
|
||||
comment: None,
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn choice(item: &Annotated<StreamResponse>) -> &ChatChoiceStream {
|
||||
item.data.as_ref().unwrap().choices.first().unwrap()
|
||||
}
|
||||
|
||||
fn delta_text(item: &Annotated<StreamResponse>) -> String {
|
||||
match choice(item).delta.content.as_ref().unwrap() {
|
||||
ChatCompletionMessageContent::Text(text) => text.clone(),
|
||||
_ => panic!("expected a text delta"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn apply_jail(
|
||||
items: Vec<Annotated<StreamResponse>>,
|
||||
parser: &str,
|
||||
) -> Vec<Annotated<StreamResponse>> {
|
||||
apply_tool_calling_jail(
|
||||
Some(parser.into()),
|
||||
Some(ChatCompletionToolChoiceOption::Auto),
|
||||
None,
|
||||
false,
|
||||
stream::iter(items),
|
||||
)
|
||||
.collect()
|
||||
.await
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dynamo_parser_name_canonicalizes_cli_names() {
|
||||
assert_eq!(dynamo_parser_name("llama3"), "llama3_json");
|
||||
assert_eq!(dynamo_parser_name("qwen"), "qwen25");
|
||||
assert_eq!(dynamo_parser_name("glm"), "glm47");
|
||||
assert_eq!(dynamo_parser_name("glm45"), "glm47");
|
||||
assert_eq!(dynamo_parser_name("qwen25"), "qwen25");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_tool_choice_builds_python_compatible_constraint() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::Required,
|
||||
&[tool("get_weather", true)],
|
||||
Some(false),
|
||||
)
|
||||
.unwrap();
|
||||
let schema: serde_json::Value =
|
||||
serde_json::from_str(sampling.json_schema.as_deref().unwrap()).unwrap();
|
||||
assert_eq!(schema["type"], "array");
|
||||
assert_eq!(schema["minItems"], 1);
|
||||
assert_eq!(schema["maxItems"], 1);
|
||||
assert_eq!(
|
||||
schema["items"]["properties"]["name"]["enum"][0],
|
||||
"get_weather"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_choice_without_parallel_flag_has_no_max_items() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::Required,
|
||||
&[tool("get_weather", false)],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let schema: serde_json::Value =
|
||||
serde_json::from_str(sampling.json_schema.as_deref().unwrap()).unwrap();
|
||||
assert_eq!(schema["type"], "array");
|
||||
assert!(schema.get("maxItems").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn named_tool_choice_restricts_the_schema_enum() {
|
||||
let tools = [tool("get_weather", false), tool("get_time", false)];
|
||||
let mut sampling = SamplingParams::default();
|
||||
apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::Named("get_time".into()),
|
||||
&tools,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let schema: serde_json::Value =
|
||||
serde_json::from_str(sampling.json_schema.as_deref().unwrap()).unwrap();
|
||||
// One candidate → a single schema, not an anyOf.
|
||||
assert_eq!(schema["items"]["properties"]["name"]["enum"][0], "get_time");
|
||||
assert!(schema["items"].get("anyOf").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strict_auto_llama_tool_uses_python_compatible_constraint() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::Auto,
|
||||
&[tool("get_weather", true)],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let schema: serde_json::Value =
|
||||
serde_json::from_str(sampling.structural_tag.as_deref().unwrap()).unwrap();
|
||||
assert_eq!(schema["type"], "structural_tag");
|
||||
assert_eq!(schema["format"]["type"], "triggered_tags");
|
||||
assert_eq!(schema["format"]["at_least_one"], false);
|
||||
assert_eq!(
|
||||
schema["format"]["tags"][0]["content"]["json_schema"]["required"][0],
|
||||
"city"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_without_strict_tools_stays_unconstrained() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::Auto,
|
||||
&[tool("get_weather", false)],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(sampling.json_schema.is_none());
|
||||
assert!(sampling.structural_tag.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_choice_none_is_a_no_op() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::None,
|
||||
&[],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(sampling.json_schema.is_none());
|
||||
assert!(sampling.structural_tag.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_tool_choices_are_rejected_before_submission() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
let error = apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::Required,
|
||||
&[],
|
||||
None,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(error.contains("required"));
|
||||
|
||||
let error = apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("llama3"),
|
||||
&DynamoToolChoice::Named("missing".into()),
|
||||
&[tool("get_weather", false)],
|
||||
None,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(error.contains("missing"));
|
||||
}
|
||||
|
||||
/// Validation runs even without a parser (the handler calls this in every
|
||||
/// mode), so an invalid choice is rejected before submission there too.
|
||||
#[test]
|
||||
fn missing_parser_still_validates_the_tool_choice() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
let error =
|
||||
apply_tool_constraint(&mut sampling, None, &DynamoToolChoice::Required, &[], None)
|
||||
.unwrap_err();
|
||||
assert!(error.contains("required"));
|
||||
let error = apply_tool_constraint(
|
||||
&mut sampling,
|
||||
None,
|
||||
&DynamoToolChoice::Named("missing".into()),
|
||||
&[tool("get_weather", false)],
|
||||
None,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(error.contains("missing"));
|
||||
// A valid choice with no parser stays unconstrained.
|
||||
apply_tool_constraint(
|
||||
&mut sampling,
|
||||
None,
|
||||
&DynamoToolChoice::Auto,
|
||||
&[tool("get_weather", false)],
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(sampling.json_schema.is_none());
|
||||
assert!(sampling.structural_tag.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dynamo_tool_choice_maps_the_openai_wire_values() {
|
||||
let named = |name: &str| {
|
||||
Some(ChatCompletionToolChoiceOption::Named(
|
||||
ChatCompletionNamedToolChoice {
|
||||
r#type: ChatCompletionToolType::Function,
|
||||
function: FunctionName { name: name.into() },
|
||||
},
|
||||
))
|
||||
};
|
||||
assert!(matches!(dynamo_tool_choice(&None), DynamoToolChoice::Auto));
|
||||
assert!(matches!(
|
||||
dynamo_tool_choice(&Some(ChatCompletionToolChoiceOption::Auto)),
|
||||
DynamoToolChoice::Auto
|
||||
));
|
||||
assert!(matches!(
|
||||
dynamo_tool_choice(&Some(ChatCompletionToolChoiceOption::Required)),
|
||||
DynamoToolChoice::Required
|
||||
));
|
||||
assert!(matches!(
|
||||
dynamo_tool_choice(&Some(ChatCompletionToolChoiceOption::None)),
|
||||
DynamoToolChoice::None
|
||||
));
|
||||
assert!(matches!(
|
||||
dynamo_tool_choice(&named("get_weather")),
|
||||
DynamoToolChoice::Named(name) if name == "get_weather"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_parser_is_rejected() {
|
||||
let mut sampling = SamplingParams::default();
|
||||
let error = apply_tool_constraint(
|
||||
&mut sampling,
|
||||
Some("not-a-parser"),
|
||||
&DynamoToolChoice::Auto,
|
||||
&[tool("get_weather", false)],
|
||||
None,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(error.contains("not supported"));
|
||||
assert!(sampling.json_schema.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_jail_emits_plain_text_without_buffering() {
|
||||
let items = apply_jail(
|
||||
vec![
|
||||
stream_item("Par", None),
|
||||
stream_item("is", Some(OpenAIFinishReason::Stop)),
|
||||
],
|
||||
"llama3_json",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(items.len(), 2);
|
||||
assert_eq!(delta_text(&items[0]), "Par");
|
||||
assert_eq!(choice(&items[0]).finish_reason, None);
|
||||
assert_eq!(delta_text(&items[1]), "is");
|
||||
assert_eq!(
|
||||
choice(&items[1]).finish_reason,
|
||||
Some(OpenAIFinishReason::Stop)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_jail_buffers_a_whole_call_until_done() {
|
||||
let items = apply_jail(
|
||||
vec![stream_item(
|
||||
r#"<|python_tag|>{"name":"get_weather","parameters":{"city":"Paris"}}"#,
|
||||
Some(OpenAIFinishReason::Stop),
|
||||
)],
|
||||
"llama3_json",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(items.len(), 1);
|
||||
let terminal = choice(&items[0]);
|
||||
assert!(matches!(
|
||||
terminal.delta.content.as_ref(),
|
||||
Some(ChatCompletionMessageContent::Text(text)) if text.is_empty()
|
||||
));
|
||||
assert_eq!(
|
||||
terminal.delta.tool_calls.as_ref().unwrap()[0]
|
||||
.function
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.name,
|
||||
Some("get_weather".into())
|
||||
);
|
||||
// The terminal reason is rewritten: calls were emitted.
|
||||
assert_eq!(terminal.finish_reason, Some(OpenAIFinishReason::ToolCalls));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_jail_detects_bare_json_without_a_start_marker() {
|
||||
let items = apply_jail(
|
||||
vec![
|
||||
stream_item(r#"{"name":"get_weather","parameters":{"#, None),
|
||||
stream_item(r#""city":"Paris"}}"#, Some(OpenAIFinishReason::Stop)),
|
||||
],
|
||||
"llama3_json",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(items.len(), 1);
|
||||
let terminal = choice(&items[0]);
|
||||
assert!(matches!(
|
||||
terminal.delta.content.as_ref(),
|
||||
Some(ChatCompletionMessageContent::Text(text)) if text.is_empty()
|
||||
));
|
||||
assert_eq!(
|
||||
terminal.delta.tool_calls.as_ref().unwrap()[0]
|
||||
.function
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.name,
|
||||
Some("get_weather".into())
|
||||
);
|
||||
assert_eq!(terminal.finish_reason, Some(OpenAIFinishReason::ToolCalls));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_jail_holds_only_a_split_marker() {
|
||||
let items = apply_jail(
|
||||
vec![
|
||||
stream_item("Before <|python_", None),
|
||||
stream_item(
|
||||
r#"tag|>{"name":"get_weather","parameters":{"city":"Paris"}}"#,
|
||||
Some(OpenAIFinishReason::Stop),
|
||||
),
|
||||
],
|
||||
"llama3_json",
|
||||
)
|
||||
.await;
|
||||
// The safe prefix streams immediately; the held marker suffix joins
|
||||
// the next chunk, which parses into a tool call.
|
||||
assert_eq!(delta_text(&items[0]), "Before ");
|
||||
let tool_call = choice(&items[1]);
|
||||
assert!(matches!(
|
||||
tool_call.delta.content.as_ref(),
|
||||
Some(ChatCompletionMessageContent::Text(text)) if text.is_empty()
|
||||
));
|
||||
assert_eq!(
|
||||
tool_call.delta.tool_calls.as_ref().unwrap()[0]
|
||||
.function
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.name,
|
||||
Some("get_weather".into())
|
||||
);
|
||||
assert_eq!(tool_call.finish_reason, Some(OpenAIFinishReason::ToolCalls));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_jail_releases_an_incomplete_marker_at_done() {
|
||||
let items = apply_jail(
|
||||
vec![
|
||||
stream_item("Before <|python_", None),
|
||||
stream_done(OpenAIFinishReason::Stop),
|
||||
],
|
||||
"llama3_json",
|
||||
)
|
||||
.await;
|
||||
let text = items
|
||||
.iter()
|
||||
.filter_map(|item| choice(item).delta.content.as_ref())
|
||||
.filter_map(|content| match content {
|
||||
ChatCompletionMessageContent::Text(text) => Some(text.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<String>();
|
||||
assert_eq!(text, "Before <|python_");
|
||||
assert_eq!(
|
||||
choice(&items[1]).finish_reason,
|
||||
Some(OpenAIFinishReason::Stop)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_jail_emits_a_complete_tool_call_before_done() {
|
||||
let items = apply_jail(
|
||||
vec![
|
||||
stream_item(
|
||||
r#"<|python_tag|>{"name":"get_weather","parameters":{"city":"Paris"}}"#,
|
||||
None,
|
||||
),
|
||||
stream_done(OpenAIFinishReason::Stop),
|
||||
],
|
||||
"llama3_json",
|
||||
)
|
||||
.await;
|
||||
let tool_position = items
|
||||
.iter()
|
||||
.position(|item| choice(item).delta.tool_calls.is_some())
|
||||
.expect("tool call chunk");
|
||||
let terminal_position = items
|
||||
.iter()
|
||||
.position(|item| choice(item).finish_reason.is_some())
|
||||
.expect("terminal chunk");
|
||||
assert!(tool_position < terminal_position);
|
||||
// Calls were emitted, so the terminal reason is rewritten.
|
||||
assert_eq!(
|
||||
choice(&items[terminal_position]).finish_reason,
|
||||
Some(OpenAIFinishReason::ToolCalls)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn canonical_qwen_parser_name_uses_dynamo_qwen25() {
|
||||
let (content, calls) = parse_chat_tool_calls(
|
||||
r#"<tool_call>{"name":"get_weather","arguments":{"city":"Paris"}}</tool_call>"#.into(),
|
||||
Some("qwen"),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert!(content.is_empty());
|
||||
assert_eq!(calls.unwrap()[0].function.name, "get_weather");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unary_parse_without_a_parser_passes_content_through() {
|
||||
let (content, calls) =
|
||||
parse_chat_tool_calls("<|python_tag|>call".into(), None, None, true).await;
|
||||
assert_eq!(content, "<|python_tag|>call");
|
||||
assert!(calls.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_finish_reason_maps_scheduler_kinds() {
|
||||
let output = |finish: serde_json::Value| ChunkEvent {
|
||||
rid: "r".into(),
|
||||
text: "x".into(),
|
||||
token_ids: vec![1],
|
||||
prompt_tokens: 1,
|
||||
completion_tokens: 1,
|
||||
finish_reason: Some(serde_json::from_value(finish).unwrap()),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
chat_finish_reason(&output(
|
||||
serde_json::json!({"type": "stop", "matched": "</s>"})
|
||||
)),
|
||||
Some(OpenAIFinishReason::Stop)
|
||||
);
|
||||
assert_eq!(
|
||||
chat_finish_reason(&output(serde_json::json!({"type": "length", "length": 8}))),
|
||||
Some(OpenAIFinishReason::Length)
|
||||
);
|
||||
assert_eq!(
|
||||
chat_finish_reason(&output(serde_json::json!({"type": "content_filter"}))),
|
||||
Some(OpenAIFinishReason::ContentFilter)
|
||||
);
|
||||
// Unknown kinds (including a bare abort) fall back to `stop`.
|
||||
assert_eq!(
|
||||
chat_finish_reason(&output(serde_json::json!({"type": "abort"}))),
|
||||
Some(OpenAIFinishReason::Stop)
|
||||
);
|
||||
assert_eq!(
|
||||
chat_finish_reason(&ChunkEvent {
|
||||
finish_reason: None,
|
||||
..Default::default()
|
||||
}),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_delta_carries_the_optional_columns() {
|
||||
let delta = chat_delta(
|
||||
Some("hi".into()),
|
||||
Some(Role::Assistant),
|
||||
Some(vec![ChatCompletionMessageToolCallChunk {
|
||||
index: 0,
|
||||
id: Some("call_1".into()),
|
||||
r#type: Some(FunctionType::Function),
|
||||
function: Some(FunctionCallStream {
|
||||
name: Some("get_weather".into()),
|
||||
arguments: Some("{}".into()),
|
||||
}),
|
||||
}]),
|
||||
Some("thinking".into()),
|
||||
);
|
||||
assert_eq!(
|
||||
delta.content,
|
||||
Some(ChatCompletionMessageContent::Text("hi".into()))
|
||||
);
|
||||
assert_eq!(delta.role, Some(Role::Assistant));
|
||||
assert_eq!(delta.reasoning_content, Some("thinking".into()));
|
||||
assert_eq!(
|
||||
delta.tool_calls.as_ref().unwrap()[0]
|
||||
.function
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.name,
|
||||
Some("get_weather".into())
|
||||
);
|
||||
assert!(chat_delta(None, None, None, None).content.is_none());
|
||||
}
|
||||
}
|
||||
@@ -38,6 +38,8 @@ pub(super) async fn submit(
|
||||
// client's, or minted one. Control requests have no client-facing rid.
|
||||
RequestKind::Generate(g) => g.rid.clone(),
|
||||
RequestKind::Control(c) => c.rid().into(),
|
||||
// Internal service call — no client-facing rid; mint a fresh one.
|
||||
RequestKind::Detokenize { .. } => Rid::new(),
|
||||
};
|
||||
// Two in-flight requests can name the same client rid, but they cannot share a
|
||||
// `Rid`: `into_requests` built each through `Rid::from_client`, which appends a
|
||||
@@ -79,6 +81,14 @@ pub(super) fn pre_submit_error(code: StatusCode, message: &str, stream: bool) ->
|
||||
if !stream {
|
||||
return (code, Json(body)).into_response();
|
||||
}
|
||||
sse_error_response(body)
|
||||
}
|
||||
|
||||
/// A 200 SSE response carrying one error frame + `[DONE]` — how a stream the
|
||||
/// client is already committed to reading reports a failure. Shared by every
|
||||
/// endpoint family: the native API via [`pre_submit_error`] and the OpenAI
|
||||
/// frontend's `openai_error_response`.
|
||||
pub(super) fn sse_error_response(body: serde_json::Value) -> Response {
|
||||
let frames = [body.to_string(), "[DONE]".to_string()];
|
||||
Sse::new(futures::stream::iter(
|
||||
frames.map(|data| Ok::<_, Infallible>(Event::default().data(data))),
|
||||
|
||||
@@ -91,6 +91,21 @@ impl DetokenizerBackend {
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode one complete sequence without creating request-scoped streaming
|
||||
/// state. This runs on a pinned detokenizer worker, never on an API runtime
|
||||
/// thread.
|
||||
fn decode_once(&self, token_ids: &[u32]) -> Result<String, Error> {
|
||||
match self {
|
||||
DetokenizerBackend::Dynamo(tokenizer) => tokenizer
|
||||
.decode(token_ids, true)
|
||||
.map(String::from)
|
||||
.map_err(|error| Error::Detokenize(error.to_string())),
|
||||
DetokenizerBackend::Skip => Err(Error::Validation(
|
||||
"echo for token-ID prompts is unavailable when skip_tokenizer_init=True".into(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode each logprob token id to its own text (one id at a time, matching
|
||||
/// Python's `batch_decode([[id] for id in ids])`). Runs on this CPU-bound
|
||||
/// shard, not the api-server I/O threads. `Skip` mode (no tokenizer) yields
|
||||
@@ -194,6 +209,9 @@ impl Runnable for DetokenizerWorker {
|
||||
handle_chunk(&mut table, ev, &self.backend, &self.abort);
|
||||
}
|
||||
}
|
||||
DetokMsg::Decode { rid, token_ids } => {
|
||||
handle_decode(&mut table, &rid, &token_ids, &self.backend)
|
||||
}
|
||||
DetokMsg::Result { rid, payload } => handle_result(&mut table, &rid, payload),
|
||||
DetokMsg::Fail { rid, message } => {
|
||||
handle_fail(&mut table, &rid, message, &self.abort)
|
||||
@@ -206,6 +224,27 @@ impl Runnable for DetokenizerWorker {
|
||||
}
|
||||
}
|
||||
|
||||
/// The `RequestKind::Detokenize` backend stage: tm-ingress queued this rid's
|
||||
/// `Register` just before on this same channel, so the entry exists — deliver
|
||||
/// the decoded text (or the error) through the registered sink and drop it,
|
||||
/// like a one-result control request. No scheduler abort on failure: this kind
|
||||
/// never reaches the ring, so there is nothing to stop.
|
||||
fn handle_decode(
|
||||
table: &mut HashMap<Rid, DetokState>,
|
||||
rid: &Rid,
|
||||
token_ids: &[u32],
|
||||
backend: &DetokenizerBackend,
|
||||
) {
|
||||
if let Some(mut st) = table.remove(rid) {
|
||||
let item = match backend.decode_once(token_ids) {
|
||||
Ok(text) => EgressItem::Data(text.into()),
|
||||
Err(e) => EgressItem::Error(e),
|
||||
};
|
||||
let _ = st.sink.try_send(item);
|
||||
st.fsm = RequestState::Completed;
|
||||
}
|
||||
}
|
||||
|
||||
/// Control-request result: deliver the JSON payload to the sink verbatim as a
|
||||
/// single `Done` frame — no detokenization, no streaming.
|
||||
fn handle_result(table: &mut HashMap<Rid, DetokState>, rid: &Rid, payload: bytes::Bytes) {
|
||||
@@ -463,6 +502,58 @@ mod tests {
|
||||
assert_eq!(t, "a STOP");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_once_rejects_skip_mode() {
|
||||
let error = DetokenizerBackend::Skip.decode_once(&[1]).unwrap_err();
|
||||
assert!(matches!(error, Error::Validation(_)));
|
||||
assert!(error.to_string().contains("skip_tokenizer_init=True"));
|
||||
}
|
||||
|
||||
/// A `Decode` job answers through the REGISTERED sink and consumes the
|
||||
/// entry — the `RequestKind::Detokenize` egress contract. Uses the `Skip`
|
||||
/// backend, whose decode error must arrive as an `Error` item (not vanish):
|
||||
/// dropping it leaves the submitter awaiting a reply forever. (Unlike
|
||||
/// `handle_fail` there is deliberately no abort lane in the signature —
|
||||
/// this kind never reached the ring, so there is no scheduler work to stop.)
|
||||
#[test]
|
||||
fn decode_answers_via_registered_sink_and_consumes_the_entry() {
|
||||
let (tx, mut rx) = mpsc::channel::<EgressItem>(4);
|
||||
let mut table = HashMap::new();
|
||||
table.insert(
|
||||
Rid::from("d1"),
|
||||
DetokState {
|
||||
sink: EgressSink::Local(tx),
|
||||
decode_logprob_text: false,
|
||||
no_stop_trim: false,
|
||||
decoder: None,
|
||||
fsm: RequestState::Queued,
|
||||
},
|
||||
);
|
||||
|
||||
handle_decode(
|
||||
&mut table,
|
||||
&Rid::from("d1"),
|
||||
&[1],
|
||||
&DetokenizerBackend::Skip,
|
||||
);
|
||||
|
||||
let Ok(EgressItem::Error(err)) = rx.try_recv() else {
|
||||
panic!("the decode error must reach the sink, not vanish");
|
||||
};
|
||||
assert!(matches!(err, Error::Validation(_)));
|
||||
assert!(!table.contains_key(&Rid::from("d1")), "entry consumed");
|
||||
|
||||
// Unregistered rid (raced with an abort's Deregister): nothing to
|
||||
// answer to — must be a no-op, not a panic.
|
||||
handle_decode(
|
||||
&mut table,
|
||||
&Rid::from("d2"),
|
||||
&[1],
|
||||
&DetokenizerBackend::Skip,
|
||||
);
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
|
||||
/// Two requests on the SAME shard keep separate entries. This is what a
|
||||
/// A shard-hash collision now degrades to: the hash partitions, the rid
|
||||
/// identifies. Keying the table by the hash made colliding rids one entry, so
|
||||
|
||||
@@ -74,6 +74,13 @@ pub enum DetokMsg {
|
||||
/// One decode step's chunks for *this shard*. Batched because `tm-egress` blocks
|
||||
/// per send, so one message per request cost ~1.3 µs × batch (5.1x at 4096).
|
||||
Chunks(Vec<ChunkEvent>),
|
||||
/// Decode a complete token-id sequence — the backend of
|
||||
/// [`RequestKind::Detokenize`], the one request kind the detok stage itself
|
||||
/// answers (it never reaches the scheduler ring). Sent by tm-ingress right
|
||||
/// after the same rid's `Register` on the same channel (FIFO), so the shard
|
||||
/// delivers the text through the registered sink like a control `Result`
|
||||
/// and drops the entry.
|
||||
Decode { rid: Rid, token_ids: Vec<u32> },
|
||||
/// Control result: one already-serialized payload delivered to the sink verbatim.
|
||||
Result { rid: Rid, payload: bytes::Bytes },
|
||||
/// Terminal per-request failure → an `Error` to the sink (a 400, not a crash).
|
||||
|
||||
@@ -53,6 +53,11 @@ pub enum EgressItem {
|
||||
/// A control-request result: one verbatim payload (e.g. `/server_info`),
|
||||
/// delivered as-is with no per-protocol formatting.
|
||||
Control(Bytes),
|
||||
/// Reply to an internal service request (`RequestKind::Detokenize`): raw
|
||||
/// bytes for the SUBMITTER to consume (e.g. the decoded prompt text), not
|
||||
/// client-bound JSON like `Control` and not a generation frame. Generation
|
||||
/// and control drains never see it.
|
||||
Data(Bytes),
|
||||
/// Terminal failure: handler emits an error frame (stream) or status (unary).
|
||||
Error(Error),
|
||||
}
|
||||
|
||||
@@ -68,6 +68,8 @@ pub enum FinishReason {
|
||||
/// This arm is why the outer enum is untagged: a finish reason added Python-side
|
||||
/// must not fail the header decode, which rejects the whole frame — every
|
||||
/// request in the batch, not just the one that carried it.
|
||||
// Keep the native frame compact even when HTTP/rendering dependencies turn
|
||||
// on serde_json's large `preserve_order` map representation.
|
||||
Unknown(Box<serde_json::Map<String, serde_json::Value>>),
|
||||
}
|
||||
|
||||
@@ -78,6 +80,16 @@ impl From<FinishKind> for FinishReason {
|
||||
}
|
||||
|
||||
impl FinishReason {
|
||||
/// Returns the wire type without reserializing the finish reason.
|
||||
pub fn kind_name(&self) -> Option<&str> {
|
||||
match self {
|
||||
FinishReason::Known(FinishKind::Stop { .. }) => Some("stop"),
|
||||
FinishReason::Known(FinishKind::Length { .. }) => Some("length"),
|
||||
FinishReason::Known(FinishKind::Abort(_)) => Some("abort"),
|
||||
FinishReason::Unknown(fields) => fields.get("type").and_then(|value| value.as_str()),
|
||||
}
|
||||
}
|
||||
|
||||
/// The stop this request matched, if it stopped on one. `None` for
|
||||
/// length/abort and for an unknown type.
|
||||
pub fn matched(&self) -> Option<&Matched> {
|
||||
|
||||
@@ -359,6 +359,9 @@ impl GenerateBody {
|
||||
rid,
|
||||
text,
|
||||
input_ids,
|
||||
// Native text prompts keep the post-processor specials; the
|
||||
// chat flow sets this explicitly.
|
||||
skip_special_tokens: false,
|
||||
sampling_params,
|
||||
stream,
|
||||
// Python `GenerateReqInput` defaults.
|
||||
@@ -394,6 +397,13 @@ pub enum RequestKind {
|
||||
/// A control endpoint (e.g. `/server_info`, `/health`): no tokenization, and
|
||||
/// the egress is a single non-streamed JSON result.
|
||||
Control(Box<ControlRequest>),
|
||||
/// Internal service call: decode a complete token-id sequence to text. Walks
|
||||
/// the same FSM as every request (validate → register → Queued), but the
|
||||
/// stage that answers it is the detok shard itself, never the scheduler
|
||||
/// ring; the result arrives on the registered sink as one `Data` payload
|
||||
/// (the raw UTF-8 text). First caller: `/v1/completions` `echo` for
|
||||
/// token-id prompts; a future `/detokenize` parity endpoint maps 1:1.
|
||||
Detokenize { token_ids: TokenIds },
|
||||
}
|
||||
|
||||
/// A single in-flight `/generate` request (per-item from
|
||||
@@ -422,6 +432,12 @@ pub struct GenerateRequest {
|
||||
pub text: Option<String>,
|
||||
/// Client-supplied token ids, or filled by the Tokenizer stage.
|
||||
pub input_ids: Option<TokenIds>,
|
||||
/// Template-rendered prompts (chat) already contain their role/special
|
||||
/// tokens, so the tokenizer pool strips the auto-added BOS/EOS prefix —
|
||||
/// the Rust analogue of Python's `add_special_tokens=False` at the
|
||||
/// chat-template encode site (`serving_chat._encode_messages`). Consumed
|
||||
/// by the pool before the header is built; never reaches the scheduler wire.
|
||||
pub skip_special_tokens: bool,
|
||||
/// Sampling params (defaults when the client sent none, as in Python);
|
||||
/// normalized + verified at ingress, then serialized into the header.
|
||||
pub sampling_params: SamplingParams,
|
||||
|
||||
@@ -19,7 +19,7 @@ mod config;
|
||||
mod runnable;
|
||||
mod threads;
|
||||
|
||||
pub use config::{RuntimeConfig, RustServerServerArgs, ServerArgs};
|
||||
pub use config::{DefaultSamplingParams, RuntimeConfig, RustServerServerArgs, ServerArgs};
|
||||
|
||||
use crate::message::DetokMsg;
|
||||
use crate::ring::{
|
||||
|
||||
@@ -56,7 +56,7 @@ impl Default for RuntimeConfig {
|
||||
}
|
||||
|
||||
/// The scheduler's startup blob (`RustServer._build_server_args`) parsed once into
|
||||
/// typed fields: values are post-`__post_init__`, unknown keys (e.g. `api_key`) are dropped.
|
||||
/// typed fields: values are post-`__post_init__`; unrelated unknown keys are dropped.
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
pub struct ServerArgs {
|
||||
/// HF repo id / local dir of the model, reported by `/get_model_info`.
|
||||
@@ -83,6 +83,21 @@ pub struct ServerArgs {
|
||||
pub log_level: String,
|
||||
#[serde(default)]
|
||||
pub log_level_http: Option<String>,
|
||||
/// Optional built-in chat-template name or path to a Jinja/legacy JSON
|
||||
/// template file. Without an override, uses the tokenizer config template.
|
||||
#[serde(default)]
|
||||
pub chat_template: Option<String>,
|
||||
/// Parser selected by `--tool-call-parser`.
|
||||
#[serde(default)]
|
||||
pub tool_call_parser: Option<String>,
|
||||
/// Reasoning splitter selected by `--reasoning-parser` (e.g. deepseek-r1).
|
||||
/// When set, chat completions strip the model's reasoning markers out of
|
||||
/// `content` into `reasoning_content` — both unary and streaming.
|
||||
#[serde(default)]
|
||||
pub reasoning_parser: Option<String>,
|
||||
/// Python's global default for whether an SSE stream ends with a usage chunk.
|
||||
#[serde(default)]
|
||||
pub stream_response_default_include_usage: bool,
|
||||
/// Pinned tokenizer threads / detok shards (Python asserts both ≥ 1).
|
||||
#[serde(default = "default_worker_num")]
|
||||
pub tokenizer_worker_num: usize,
|
||||
@@ -143,6 +158,37 @@ pub struct ModelConfig {
|
||||
/// boot ([`ServerArgs::validate_mandatory`]).
|
||||
#[serde(default)]
|
||||
pub vocab_size: Option<u64>,
|
||||
/// Resolved default sampling parameters, stamped by
|
||||
/// `RustServer._build_server_args` from Python's
|
||||
/// `ModelConfig.get_default_sampling_params()`. Already gated on
|
||||
/// `--sampling-defaults`: holds the model's generation_config.json values
|
||||
/// in "model" mode, and is empty in "openai" mode. Consumed when a chat
|
||||
/// request omits `temperature`/`top_p` — the conversion must not skip
|
||||
/// straight to the OpenAI terminal defaults.
|
||||
#[serde(default)]
|
||||
pub default_sampling_params: DefaultSamplingParams,
|
||||
}
|
||||
|
||||
/// One `SamplingParams` field per key `get_default_sampling_params()` may emit
|
||||
/// (`repetition_penalty`, `temperature`, `top_k`, `top_p`, `min_p`), filtered
|
||||
/// to values the generation config actually sets — hence all `Option`.
|
||||
///
|
||||
/// `top_k` / `min_p` / `repetition_penalty` are parsed for parity with the
|
||||
/// Python dict but not yet consumed: the Dynamo chat request type only carries
|
||||
/// `temperature` and `top_p`, so the conversion resolves just those two.
|
||||
#[derive(Debug, Default, serde::Deserialize)]
|
||||
#[allow(dead_code)]
|
||||
pub struct DefaultSamplingParams {
|
||||
#[serde(default)]
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub top_p: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub top_k: Option<i64>,
|
||||
#[serde(default)]
|
||||
pub min_p: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub repetition_penalty: Option<f64>,
|
||||
}
|
||||
|
||||
fn join_host_port(host: &str, port: u16) -> String {
|
||||
|
||||
@@ -24,6 +24,13 @@ use crate::tokenizer_manager::TmEvent;
|
||||
/// (read-only) across all pinned workers.
|
||||
pub trait TextTokenizer: Send + Sync {
|
||||
fn encode(&self, text: &str) -> Result<TokenIds, Error>;
|
||||
|
||||
/// The special tokens this tokenizer auto-prepends on every `encode` —
|
||||
/// Python's `encode("")` probe (`serving_chat._tokenizer_auto_adds_specials`).
|
||||
/// Empty when it adds none (tiktoken backends, no BOS/EOS post-processor).
|
||||
fn auto_specials(&self) -> Vec<i32> {
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Load the tokenizer shared (Arc-backed) by the encode pool and detok shards.
|
||||
@@ -42,7 +49,6 @@ pub fn load_tokenizer(
|
||||
let path = tokenizer_path.ok_or_else(|| {
|
||||
"no tokenizer configured: set tokenizer_path or enable skip_tokenizer_init".to_string()
|
||||
})?;
|
||||
|
||||
let file = resolve_model_file(path, revision, "tokenizer.json")
|
||||
.ok_or_else(|| format!("tokenizer.json not found for '{path}'"))?;
|
||||
let tokenizer = dynamo_tokenizers::Tokenizer::from_file_with_options(
|
||||
@@ -124,14 +130,40 @@ impl TextTokenizer for DynamoTokenizer {
|
||||
// Vocab ids are non-negative and fit in i32.
|
||||
Ok(encoding.token_ids().iter().map(|&id| id as i32).collect())
|
||||
}
|
||||
|
||||
/// The post-processor prepends exactly what `encode("")` returns, so the
|
||||
/// probe is the same prefix [`strip_auto_specials`] removes.
|
||||
fn auto_specials(&self) -> Vec<i32> {
|
||||
self.inner
|
||||
.encode("")
|
||||
.map(|encoding| encoding.token_ids().iter().map(|&id| id as i32).collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove one leading run of auto-added specials — exactly what an
|
||||
/// `add_special_tokens=false` encode would have produced, without a second
|
||||
/// tokenizer instance (the post-processor always prepends the same prefix, so
|
||||
/// a template-rendered copy of those tokens is preserved).
|
||||
fn strip_auto_specials(mut ids: Vec<i32>, auto_specials: &[i32]) -> Vec<i32> {
|
||||
if ids.starts_with(auto_specials) {
|
||||
ids.drain(..auto_specials.len());
|
||||
}
|
||||
ids
|
||||
}
|
||||
|
||||
/// One tokenizer worker: pulls a `Request` off the shared inbox, fills
|
||||
/// `input_ids`, returns it to the TokenizerManager. Pinned; backend shared.
|
||||
///
|
||||
/// The `auto_specials` prefix (probed once at construction, Python's
|
||||
/// `encode("")` probe) is stripped from template-rendered prompts —
|
||||
/// [`GenerateRequest`]'s `skip_special_tokens` — so chat prompts gain no
|
||||
/// extra BOS/EOS while native text keeps the post-processor specials.
|
||||
pub struct TokenizerWorker {
|
||||
rx: flume::Receiver<Request>,
|
||||
tm: flume::Sender<TmEvent>,
|
||||
tokenizer: Arc<dyn TextTokenizer>,
|
||||
auto_specials: Vec<i32>,
|
||||
}
|
||||
|
||||
impl TokenizerWorker {
|
||||
@@ -140,7 +172,13 @@ impl TokenizerWorker {
|
||||
tm: flume::Sender<TmEvent>,
|
||||
tokenizer: Arc<dyn TextTokenizer>,
|
||||
) -> Self {
|
||||
Self { rx, tm, tokenizer }
|
||||
let auto_specials = tokenizer.auto_specials();
|
||||
Self {
|
||||
rx,
|
||||
tm,
|
||||
tokenizer,
|
||||
auto_specials,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,7 +208,11 @@ impl Runnable for TokenizerWorker {
|
||||
}
|
||||
match self.tokenizer.encode(g.text.as_deref().unwrap_or("")) {
|
||||
Ok(ids) => {
|
||||
g.input_ids = Some(ids);
|
||||
g.input_ids = Some(if g.skip_special_tokens {
|
||||
strip_auto_specials(ids, &self.auto_specials)
|
||||
} else {
|
||||
ids
|
||||
});
|
||||
Event::TokenizeDone
|
||||
}
|
||||
Err(err) => Event::Error(err),
|
||||
@@ -248,4 +290,62 @@ mod tests {
|
||||
"must be the max TOKEN count (3), not the byte count (8)"
|
||||
);
|
||||
}
|
||||
|
||||
/// The strip reproduces `add_special_tokens=false`: one leading run of
|
||||
/// auto-added specials is removed, a template-rendered copy is kept, and
|
||||
/// tokenizers with no auto specials (empty probe) are untouched.
|
||||
#[test]
|
||||
fn strip_auto_specials_matches_add_special_tokens_false() {
|
||||
assert_eq!(strip_auto_specials(vec![0, 0, 1, 2], &[0]), vec![0, 1, 2]);
|
||||
assert_eq!(strip_auto_specials(vec![1, 2], &[0]), vec![1, 2]);
|
||||
assert_eq!(strip_auto_specials(vec![1, 2], &[]), vec![1, 2]);
|
||||
assert_eq!(strip_auto_specials(vec![0], &[0, 9]), vec![0]);
|
||||
}
|
||||
|
||||
/// Word tokens plus a prepended BOS marker (id 0) — like an HF tokenizer
|
||||
/// whose post-processor adds specials.
|
||||
struct MarkedTokenizer;
|
||||
impl TextTokenizer for MarkedTokenizer {
|
||||
fn encode(&self, text: &str) -> Result<TokenIds, Error> {
|
||||
Ok(vec![0, text.len() as i32])
|
||||
}
|
||||
fn auto_specials(&self) -> Vec<i32> {
|
||||
vec![0]
|
||||
}
|
||||
}
|
||||
|
||||
/// `skip_special_tokens` strips the probed prefix: template-rendered
|
||||
/// prompts (chat) must not gain a BOS the template didn't render — Python's
|
||||
/// `add_special_tokens=False` at the chat-template encode site.
|
||||
#[test]
|
||||
fn skip_special_tokens_strips_the_auto_added_specials() {
|
||||
let run = |skip_special_tokens: bool| {
|
||||
let (req_tx, req_rx) = flume::unbounded::<Request>();
|
||||
let (tm_tx, tm_rx) = flume::unbounded::<TmEvent>();
|
||||
req_tx
|
||||
.send(Request {
|
||||
rid: "1".into(),
|
||||
state: RequestState::Tokenizing,
|
||||
sink: EgressSink::Local(tokio::sync::mpsc::channel(4).0),
|
||||
kind: RequestKind::Generate(Box::new(GenerateRequest {
|
||||
rid: "1".into(),
|
||||
text: Some("hi".into()),
|
||||
skip_special_tokens,
|
||||
..Default::default()
|
||||
})),
|
||||
})
|
||||
.expect("send");
|
||||
drop(req_tx);
|
||||
TokenizerWorker::new(req_rx, tm_tx, Arc::new(MarkedTokenizer)).run();
|
||||
let TmEvent::Tokenized(req) = tm_rx.try_recv().expect("returned") else {
|
||||
panic!("expected Tokenized");
|
||||
};
|
||||
let RequestKind::Generate(g) = &req.kind else {
|
||||
panic!("expected generate");
|
||||
};
|
||||
g.input_ids.clone().expect("tokenized")
|
||||
};
|
||||
assert_eq!(run(false), vec![0, 2], "native prompts keep specials");
|
||||
assert_eq!(run(true), vec![2], "rendered prompts lose the auto BOS");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -204,10 +204,10 @@ impl Ingress {
|
||||
registered = true;
|
||||
// `validate` advanced Received → Validating; keep driving.
|
||||
}
|
||||
// Control skips normalization (no sampling params) straight to the
|
||||
// pre-send checks; generate goes to Normalizing.
|
||||
// Control and detokenize skip normalization (no sampling params)
|
||||
// straight to the pre-send checks; generate goes to Normalizing.
|
||||
RequestState::Validating => match &req.kind {
|
||||
RequestKind::Control(_) => {
|
||||
RequestKind::Control(_) | RequestKind::Detokenize { .. } => {
|
||||
let _ = req
|
||||
.state
|
||||
.apply(Event::Validated(ValidationOutcome::AlreadyTokenized));
|
||||
@@ -221,8 +221,8 @@ impl Ingress {
|
||||
RequestState::Normalizing => {
|
||||
let outcome = {
|
||||
let RequestKind::Generate(g) = &mut req.kind else {
|
||||
// Unreachable (control never reaches here); reject so a
|
||||
// bug can't leak/hang a registered request.
|
||||
// Unreachable (control/detokenize never reach here);
|
||||
// reject so a bug can't leak/hang a registered request.
|
||||
self.fail(
|
||||
&mut req,
|
||||
Error::Internal("non-generate request in Normalizing".into()),
|
||||
@@ -282,14 +282,16 @@ impl Ingress {
|
||||
}
|
||||
let _ = req.state.apply(Event::PreSendValidated); // → Queued
|
||||
}
|
||||
// Push the wire message (control frame or generate payload) to the ring.
|
||||
// Hand the request to the stage that answers it: the scheduler
|
||||
// ring (generate payload or control frame), or — for detokenize
|
||||
// — the detok shard itself.
|
||||
RequestState::Queued => {
|
||||
// `matches!` reads the discriminant without holding a borrow,
|
||||
// so `req` can be moved into the push below.
|
||||
if matches!(req.kind, RequestKind::Generate(_)) {
|
||||
self.push_to_ring(req);
|
||||
} else {
|
||||
self.push_control_to_ring(req);
|
||||
// The patterns bind nothing, so the match reads only the
|
||||
// discriminant and `req` can be moved into each push.
|
||||
match req.kind {
|
||||
RequestKind::Generate(_) => self.push_to_ring(req),
|
||||
RequestKind::Control(_) => self.push_control_to_ring(req),
|
||||
RequestKind::Detokenize { .. } => self.push_detokenize_to_shard(req),
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -323,7 +325,7 @@ impl Ingress {
|
||||
g.return_text_in_logprobs.unwrap_or(false),
|
||||
g.sampling_params.no_stop_trim,
|
||||
),
|
||||
RequestKind::Control(_) => (false, false),
|
||||
RequestKind::Control(_) | RequestKind::Detokenize { .. } => (false, false),
|
||||
};
|
||||
self.senders
|
||||
.detok_for(&req.rid)
|
||||
@@ -336,6 +338,34 @@ impl Ingress {
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
/// Hand a `Detokenize` request to its owning detok shard — the stage that
|
||||
/// answers this kind (it never touches the scheduler ring). The shard
|
||||
/// already holds this rid's sink: `register_detok` queued `Register` on the
|
||||
/// same channel from this same thread, so FIFO gives Register → Decode.
|
||||
fn push_detokenize_to_shard(&self, mut req: Request) {
|
||||
let RequestKind::Detokenize { token_ids } = &req.kind else {
|
||||
self.fail(
|
||||
&mut req,
|
||||
Error::Internal("non-detokenize request reached push_detokenize_to_shard".into()),
|
||||
true,
|
||||
);
|
||||
return;
|
||||
};
|
||||
// Infallible: `validate` rejected out-of-range ids at `Received`.
|
||||
let token_ids: Vec<u32> = token_ids.iter().map(|&id| id as u32).collect();
|
||||
if self
|
||||
.senders
|
||||
.detok_for(&req.rid)
|
||||
.send(DetokMsg::Decode {
|
||||
rid: req.rid.clone(),
|
||||
token_ids,
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
self.fail(&mut req, Error::Internal("detok shard gone".into()), true);
|
||||
}
|
||||
}
|
||||
|
||||
/// Push a bare control request (`[tag, rid, nil]`) onto the ingress ring. The
|
||||
/// scheduler dispatches it (e.g. `GetInternalStateReq`) and replies via the
|
||||
/// egress ring as a single `Result`.
|
||||
@@ -485,6 +515,18 @@ fn validate(req: &mut Request, limits: &Limits) -> Result<(), Error> {
|
||||
}
|
||||
}
|
||||
|
||||
// Detokenize ids must fit the shard's `&[u32]` decode domain. No vocab
|
||||
// bound — parity with the retired direct decode service: an unknown id is
|
||||
// the tokenizer's error to report, and nothing here reaches the scheduler's
|
||||
// embedding lookup.
|
||||
if let RequestKind::Detokenize { token_ids } = &req.kind {
|
||||
for &id in token_ids {
|
||||
if u32::try_from(id).is_err() {
|
||||
return Err(Error::Validation(format!("Token ID {id} is out of range")));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The scheduler only computes hidden states when launched for it, so without
|
||||
// this the request would 200 with `meta_info.hidden_states` silently absent
|
||||
// (Python `TokenizerManager._validate_one_request`).
|
||||
@@ -943,6 +985,68 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// A `Detokenize` request terminates at the detok stage, and the shard must
|
||||
/// see its `Register` BEFORE its `Decode` — the shard delivers the result
|
||||
/// through the sink registered under that rid, so a `Decode` that arrives
|
||||
/// unregistered is silently dropped and the caller waits forever. Both
|
||||
/// messages ride one channel from this one thread, which is the FIFO this
|
||||
/// pins. Nothing may reach the scheduler ring.
|
||||
#[test]
|
||||
fn detokenize_flows_register_then_decode_and_skips_the_ring() {
|
||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
ingress.drive(Request {
|
||||
rid: "41".into(),
|
||||
state: RequestState::Received,
|
||||
sink: EgressSink::Local(tx),
|
||||
kind: RequestKind::Detokenize {
|
||||
token_ids: vec![7, 8, 9],
|
||||
},
|
||||
});
|
||||
assert!(
|
||||
matches!(detok_rx.try_recv(), Ok(DetokMsg::Register { rid, .. }) if rid.as_str() == "41"),
|
||||
"the sink must be registered before the decode job",
|
||||
);
|
||||
assert!(
|
||||
matches!(
|
||||
detok_rx.try_recv(),
|
||||
Ok(DetokMsg::Decode { rid, token_ids })
|
||||
if rid.as_str() == "41" && token_ids == [7, 8, 9]
|
||||
),
|
||||
"the decode job follows, ids intact",
|
||||
);
|
||||
assert!(
|
||||
consumer.drain(16).headers.is_empty(),
|
||||
"must never reach the scheduler"
|
||||
);
|
||||
assert!(rx.try_recv().is_err(), "no egress until the shard answers");
|
||||
}
|
||||
|
||||
/// Negative ids cannot decode (the shard's domain is `&[u32]`): rejected by
|
||||
/// `validate` at `Received` — an `Error` to the sink, and the shard sees
|
||||
/// NOTHING (validation runs before registration, so there is no entry to
|
||||
/// leak and no decode job to drop).
|
||||
#[test]
|
||||
fn detokenize_negative_ids_reject_before_registration() {
|
||||
let (ingress, detok_rx, consumer, _tm_tx) = make_ingress();
|
||||
let (tx, mut rx) = mpsc::channel(8);
|
||||
ingress.drive(Request {
|
||||
rid: "43".into(),
|
||||
state: RequestState::Received,
|
||||
sink: EgressSink::Local(tx),
|
||||
kind: RequestKind::Detokenize {
|
||||
token_ids: vec![1, -1],
|
||||
},
|
||||
});
|
||||
let Ok(EgressItem::Error(err)) = rx.try_recv() else {
|
||||
panic!("sink must receive the validation error");
|
||||
};
|
||||
assert_eq!(err.http_status(), 400);
|
||||
assert!(err.to_string().contains("out of range"), "{err}");
|
||||
assert!(detok_rx.try_recv().is_err(), "shard never hears of it");
|
||||
assert!(consumer.drain(16).headers.is_empty());
|
||||
}
|
||||
|
||||
/// A dropped ring push is survivable, and this pins WHY. The ring is bounded,
|
||||
/// so under load the scheduler never learns to stop and keeps generating; its
|
||||
/// chunks then arrive for a rid the detok table no longer holds and are
|
||||
|
||||
Reference in New Issue
Block a user