server: sampling-param validation, finish_reason normalization, backpressure
Three related hardening changes for the API surface: - validate_request rejects NaN/negative temperature, out-of-range top_p, and absurd top_k before those values reach the CUDA sampling paths. Prevents NaN logits from downstream sampling and matches typical OpenAI-compatible server behavior (400 instead of 500). - normalize_finish_reason maps engine strings to the OpenAI-standard subset. Currently only "error" (from tp/pp engine client-stall) needs normalization — it collapses to null so SDK clients see a clean stream close instead of an unknown finish_reason value. Applied to both streaming (SSE) and non-streaming JSON responses. - Replace the unbounded std::sync::mpsc engine channel with a bounded sync_channel(256) and switch submit_to_engine to try_send. A saturated engine now returns 503 "engine is busy" instead of letting requests pile up in RAM. Also add axum DefaultBodyLimit(4 MiB) so a malicious or misbehaving client cannot exhaust memory with an arbitrary JSON POST.
This commit is contained in:
@@ -5,6 +5,7 @@ mod tp_engine;
|
||||
|
||||
use axum::{
|
||||
Extension, Router,
|
||||
extract::DefaultBodyLimit,
|
||||
routing::{get, post},
|
||||
};
|
||||
use engine::GenerateRequest;
|
||||
@@ -15,7 +16,7 @@ use xserv_model::ModelConfig;
|
||||
pub struct AppState {
|
||||
pub model_name: String,
|
||||
pub chat_template: api::ChatTemplate,
|
||||
pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
|
||||
pub engine_sender: Mutex<mpsc::SyncSender<GenerateRequest>>,
|
||||
pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
|
||||
pub max_seq_len: usize,
|
||||
}
|
||||
@@ -104,8 +105,10 @@ async fn main() {
|
||||
|
||||
let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
||||
|
||||
// Unbounded channel: allows multiple requests to queue up
|
||||
let (tx, rx) = mpsc::channel::<GenerateRequest>();
|
||||
// Bounded channel to backpressure incoming requests when the engine falls
|
||||
// behind, instead of letting them pile up in RAM. try_send in the API
|
||||
// handler surfaces this as 503 to the client.
|
||||
let (tx, rx) = mpsc::sync_channel::<GenerateRequest>(256);
|
||||
|
||||
let model_dir_clone = model_dir.clone();
|
||||
std::thread::spawn(move || {
|
||||
@@ -140,6 +143,7 @@ async fn main() {
|
||||
.route("/health", get(api::health))
|
||||
.route("/v1/models", get(api::list_models))
|
||||
.route("/v1/chat/completions", post(api::chat_completions))
|
||||
.layer(DefaultBodyLimit::max(4 * 1024 * 1024))
|
||||
.layer(Extension(state));
|
||||
|
||||
let addr = format!("0.0.0.0:{port}");
|
||||
|
||||
Reference in New Issue
Block a user