Gahow Wang
fc1900a745
Fixes the paged-KV OOM at large --max-seq-len and adds elastic memory: - Size the GPU block pool to available VRAM (cudaMemGetInfo) instead of the worst-case blocks_per_seq * max_batch * 2 reservation, which OOM'd at 8192. - Scheduler tracks waiting/running/swapped sets: block-aware admission, swap-in of resumable sequences when blocks free, and preemption of the newest running sequence to host when the pool can't cover a decode step. - --swap-space-gb (default 8) sizes the pinned host swap pool; XSERV_MAX_KV_BLOCKS forces a small pool to exercise swapping. - api: poison-tolerant lock + clean 503 when the engine thread is gone, instead of cascading mutex-poison panics. Verified on RTX 5090: serves at --max-seq-len 8192 (previously OOM), and a forced 40-block pool drives 48 lossless swap-out/swap-in cycles under concurrency with coherent output. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
94 lines
3.2 KiB
Rust
94 lines
3.2 KiB
Rust
mod api;
|
|
mod engine;
|
|
|
|
use axum::{routing::{get, post}, Extension, Router};
|
|
use std::path::PathBuf;
|
|
use std::sync::{mpsc, Arc, Mutex};
|
|
use engine::GenerateRequest;
|
|
use xserv_model::ModelConfig;
|
|
|
|
pub struct AppState {
|
|
pub model_name: String,
|
|
pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
|
|
pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
|
|
pub max_seq_len: usize,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let args: Vec<String> = std::env::args().collect();
|
|
if args.len() < 2 {
|
|
eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N]");
|
|
std::process::exit(1);
|
|
}
|
|
|
|
let model_dir = PathBuf::from(&args[1]);
|
|
let port: u16 = args.iter()
|
|
.position(|a| a == "--port")
|
|
.and_then(|i| args.get(i + 1))
|
|
.and_then(|s| s.parse().ok())
|
|
.unwrap_or(8080);
|
|
let max_batch: usize = args.iter()
|
|
.position(|a| a == "--max-batch")
|
|
.and_then(|i| args.get(i + 1))
|
|
.and_then(|s| s.parse().ok())
|
|
.unwrap_or(4)
|
|
.max(1);
|
|
let requested_max_seq_len: usize = args.iter()
|
|
.position(|a| a == "--max-seq-len")
|
|
.and_then(|i| args.get(i + 1))
|
|
.and_then(|s| s.parse().ok())
|
|
.unwrap_or(2048)
|
|
.max(1);
|
|
let swap_space_gb: usize = args.iter()
|
|
.position(|a| a == "--swap-space-gb")
|
|
.and_then(|i| args.get(i + 1))
|
|
.and_then(|s| s.parse().ok())
|
|
.unwrap_or(8);
|
|
let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
|
|
let model_max_seq_len = model_config.max_seq_len();
|
|
if model_max_seq_len == 0 {
|
|
eprintln!("model config has invalid max_seq_len=0");
|
|
std::process::exit(1);
|
|
}
|
|
let max_seq_len = requested_max_seq_len.min(model_max_seq_len);
|
|
if max_seq_len != requested_max_seq_len {
|
|
eprintln!(
|
|
"[server] --max-seq-len {requested_max_seq_len} exceeds model limit {model_max_seq_len}; using {max_seq_len}"
|
|
);
|
|
}
|
|
|
|
let model_name = model_dir.file_name()
|
|
.map(|n| n.to_string_lossy().to_string())
|
|
.unwrap_or_else(|| "unknown".to_string());
|
|
|
|
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>();
|
|
|
|
let model_dir_clone = model_dir.clone();
|
|
std::thread::spawn(move || {
|
|
let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
|
|
engine.run(rx);
|
|
});
|
|
|
|
let state = Arc::new(AppState {
|
|
model_name,
|
|
engine_sender: Mutex::new(tx),
|
|
engine_tokenizer: Mutex::new(tokenizer),
|
|
max_seq_len,
|
|
});
|
|
|
|
let app = Router::new()
|
|
.route("/health", get(api::health))
|
|
.route("/v1/models", get(api::list_models))
|
|
.route("/v1/chat/completions", post(api::chat_completions))
|
|
.layer(Extension(state));
|
|
|
|
let addr = format!("0.0.0.0:{port}");
|
|
eprintln!("[server] Listening on {addr} (max_batch={max_batch}, max_seq_len={max_seq_len})");
|
|
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
|
|
axum::serve(listener, app).await.unwrap();
|
|
}
|