xserv/crates/xserv-server/src/main.rs

mod api;
mod engine;
mod pp_engine;
mod tp_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 chat_template: api::ChatTemplate,
    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] [--tp N] [--pp 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 tp: usize = args.iter()
        .position(|a| a == "--tp")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(1)
        .max(1);
    let pp: usize = args.iter()
        .position(|a| a == "--pp")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(1)
        .max(1);
    if tp > 1 && pp > 1 {
        eprintln!("--tp and --pp cannot be combined yet (2D TP×PP is future work)");
        std::process::exit(1);
    }
    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 || {
        if pp > 1 {
            // Pipeline-parallel path: stage-0 coordinator + worker stage threads.
            pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx);
        } else if tp <= 1 {
            let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
            engine.run(rx);
        } else {
            // Tensor-parallel path: rank-0 coordinator + worker rank threads.
            tp_engine::run_tp(&model_dir_clone, tp, max_seq_len, rx);
        }
    });

    let model_type = model_config.model_type.clone().unwrap_or_default();
    let chat_template = api::ChatTemplate::load(&model_dir, &model_type);
    let state = Arc::new(AppState {
        model_name,
        chat_template,
        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();
}