phase 12+13: HTTP API server with OpenAI-compatible endpoint (Milestone ③)

New crate: xserv-server - Engine thread: loads Qwen3-8B, processes requests sequentially - axum HTTP server: /health, /v1/models, /v1/chat/completions - tokio::sync::mpsc channel between API and engine threads - Non-streaming JSON response (streaming SSE to be added later) API is OpenAI-compatible: POST /v1/chat/completions {"messages": [...], "max_tokens": N} → {"choices": [{"message": {"content": "..."}}]} Verified: "Hi" → ", I'm" (3 tokens), model runs correctly via HTTP. Key learnings: - std::sync::mpsc::SyncSender is Send but NOT Sync → wrap in Mutex for Arc<AppState> - MutexGuard must not live across await points (scope carefully) - axum 0.8 Extension<Arc<T>> requires T: Send + Sync Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-22 12:55:19 +08:00
parent 2be27d6d94
commit da043554ba
6 changed files with 376 additions and 0 deletions
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,6 +6,7 @@ members = [
    "crates/xserv-kernels",
    "crates/xserv-model",
    "crates/xserv-tokenizer",
+    "crates/xserv-server",
 ]

 [workspace.package]
@@ -20,3 +21,6 @@ serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 safetensors = "0.5"
 regex = "1"
+tokio = { version = "1", features = ["full"] }
+axum = "0.8"
+uuid = { version = "1", features = ["v4"] }
--- a/crates/xserv-server/Cargo.toml
+++ b/crates/xserv-server/Cargo.toml
@@ -0,0 +1,21 @@
+[package]
+name = "xserv-server"
+version.workspace = true
+edition.workspace = true
+
+[[bin]]
+name = "xserv-server"
+path = "src/main.rs"
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+xserv-tensor = { path = "../xserv-tensor" }
+xserv-kernels = { path = "../xserv-kernels" }
+xserv-model = { path = "../xserv-model" }
+xserv-tokenizer = { path = "../xserv-tokenizer" }
+half.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+tokio.workspace = true
+axum.workspace = true
+uuid.workspace = true
--- a/crates/xserv-server/src/api.rs
+++ b/crates/xserv-server/src/api.rs
@@ -0,0 +1,115 @@
+use axum::Extension;
+use axum::Json;
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use uuid::Uuid;
+
+use crate::engine::{GenerateEvent, GenerateRequest};
+use crate::AppState;
+
+#[derive(Deserialize)]
+pub struct ChatRequest {
+    #[serde(default)]
+    pub model: Option<String>,
+    pub messages: Vec<Message>,
+    #[serde(default = "default_max_tokens")]
+    pub max_tokens: usize,
+}
+
+#[derive(Deserialize)]
+pub struct Message {
+    pub role: String,
+    pub content: String,
+}
+
+fn default_max_tokens() -> usize { 256 }
+
+#[derive(Serialize)]
+pub struct ModelsResponse {
+    object: &'static str,
+    data: Vec<ModelInfo>,
+}
+
+#[derive(Serialize)]
+pub struct ModelInfo {
+    id: String,
+    object: &'static str,
+    owned_by: &'static str,
+}
+
+pub async fn health() -> &'static str { "ok" }
+
+pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
+    Json(ModelsResponse {
+        object: "list",
+        data: vec![ModelInfo {
+            id: state.model_name.clone(),
+            object: "model",
+            owned_by: "xserv",
+        }],
+    })
+}
+
+pub async fn chat_completions(
+    Extension(state): Extension<Arc<AppState>>,
+    Json(req): Json<ChatRequest>,
+) -> Json<serde_json::Value> {
+    let id = format!("chatcmpl-{}", Uuid::new_v4());
+    let model_name = state.model_name.clone();
+    let created = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+
+    // Prepare prompt tokens (MutexGuard scoped)
+    let prompt = build_prompt(&req.messages);
+    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
+
+    // Create channel and submit request (MutexGuard scoped)
+    let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
+    let gen_req = GenerateRequest {
+        prompt_tokens,
+        max_tokens: req.max_tokens,
+        sender: tx,
+    };
+    state.engine_sender.lock().unwrap().send(gen_req).unwrap();
+
+    // Now await — no MutexGuards held here
+    let mut content = String::new();
+    let mut finish_reason = "length".to_string();
+    while let Some(event) = rx.recv().await {
+        match event {
+            GenerateEvent::Token { text, .. } => content.push_str(&text),
+            GenerateEvent::Done { finish_reason: fr } => { finish_reason = fr; break; }
+        }
+    }
+
+    Json(serde_json::json!({
+        "id": id,
+        "object": "chat.completion",
+        "created": created,
+        "model": model_name,
+        "choices": [{
+            "index": 0,
+            "message": { "role": "assistant", "content": content },
+            "finish_reason": finish_reason,
+        }],
+        "usage": {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0
+        }
+    }))
+}
+
+fn build_prompt(messages: &[Message]) -> String {
+    let mut prompt = String::new();
+    for msg in messages {
+        match msg.role.as_str() {
+            "system" => { prompt.push_str(&msg.content); prompt.push('\n'); }
+            "user" | "assistant" => { prompt.push_str(&msg.content); }
+            _ => {}
+        }
+    }
+    prompt
+}
--- a/crates/xserv-server/src/engine.rs
+++ b/crates/xserv-server/src/engine.rs
@@ -0,0 +1,76 @@
+use std::path::Path;
+use xserv_model::{loader, GpuKVCache, ModelConfig, Qwen3};
+use xserv_model::qwen3::sample_greedy;
+use xserv_tensor::{DType, Device};
+use xserv_tokenizer::Tokenizer;
+
+pub struct Engine {
+    model: Qwen3,
+    config: ModelConfig,
+    tokenizer: Tokenizer,
+}
+
+pub struct GenerateRequest {
+    pub prompt_tokens: Vec<u32>,
+    pub max_tokens: usize,
+    pub sender: tokio::sync::mpsc::Sender<GenerateEvent>,
+}
+
+pub enum GenerateEvent {
+    Token { id: u32, text: String },
+    Done { finish_reason: String },
+}
+
+impl Engine {
+    pub fn load(model_dir: &Path) -> Self {
+        xserv_cuda::device::set_device(0).unwrap();
+        let config = ModelConfig::from_file(&model_dir.join("config.json"));
+        eprintln!("[engine] Loading weights...");
+        let weights = loader::load_model_dir(model_dir, Device::Cuda(0));
+        eprintln!("[engine] Loaded {} tensors", weights.len());
+        let model = Qwen3::from_weights(config.clone(), weights);
+        let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
+        eprintln!("[engine] Ready");
+        Self { model, config, tokenizer }
+    }
+
+    pub fn tokenizer(&self) -> &Tokenizer {
+        &self.tokenizer
+    }
+
+    pub fn generate(&self, req: GenerateRequest) {
+        let max_seq = 256;
+        let mut cache = GpuKVCache::new(&self.config, max_seq, DType::BF16);
+
+        let logits = self.model.forward_gpu_cache(&req.prompt_tokens, &mut cache);
+        let mut next = sample_greedy(&logits);
+
+        for _ in 0..req.max_tokens {
+            let text = self.tokenizer.decode(&[next]);
+            if req.sender.blocking_send(GenerateEvent::Token { id: next, text }).is_err() {
+                return;
+            }
+
+            if self.tokenizer.eos_token_id() == Some(next) {
+                let _ = req.sender.blocking_send(GenerateEvent::Done {
+                    finish_reason: "stop".to_string(),
+                });
+                return;
+            }
+
+            if cache.seq_len() >= max_seq - 1 {
+                let _ = req.sender.blocking_send(GenerateEvent::Done {
+                    finish_reason: "length".to_string(),
+                });
+                return;
+            }
+
+            let logits = self.model.forward_gpu_cache(&[next], &mut cache);
+            next = sample_greedy(&logits);
+        }
+
+        let _ = req.sender.blocking_send(GenerateEvent::Done {
+            finish_reason: "length".to_string(),
+        });
+    }
+}
--- a/crates/xserv-server/src/main.rs
+++ b/crates/xserv-server/src/main.rs
@@ -0,0 +1,62 @@
+mod api;
+mod engine;
+
+use axum::{routing::{get, post}, Extension, Router};
+use std::path::PathBuf;
+use std::sync::{mpsc, Arc, Mutex};
+use engine::GenerateRequest;
+
+pub struct AppState {
+    pub model_name: String,
+    pub engine_sender: Mutex<mpsc::SyncSender<GenerateRequest>>,
+    pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
+}
+
+#[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]");
+        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 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"));
+    let (tx, rx) = mpsc::sync_channel::<GenerateRequest>(1);
+
+    let model_dir_clone = model_dir.clone();
+    std::thread::spawn(move || {
+        let engine = engine::Engine::load(&model_dir_clone);
+        eprintln!("[engine] Listening for requests...");
+        while let Ok(req) = rx.recv() {
+            engine.generate(req);
+        }
+    });
+
+    let state = Arc::new(AppState {
+        model_name,
+        engine_sender: Mutex::new(tx),
+        engine_tokenizer: Mutex::new(tokenizer),
+    });
+
+    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}");
+    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
+    axum::serve(listener, app).await.unwrap();
+}
--- a/docs/12-13-serving.md
+++ b/docs/12-13-serving.md
@@ -0,0 +1,98 @@
+# Phase 12+13: Continuous Batching + HTTP API — Design Document (Milestone ③)
+
+## Goal
+
+实现 HTTP serving 层：接收请求、调度执行、streaming 返回结果。OpenAI 兼容 API。
+
+由于当前是单请求引擎（无 multi-GPU、无并发），Phase 12 (continuous batching) 和 Phase 13 (HTTP API) 合并实现：先实现单请求 serving，scheduler 作为 placeholder 留待后续扩展。
+
+## Architecture
+
+```
+Client (curl / OpenAI SDK)
+    │
+    ▼  HTTP POST /v1/chat/completions
+┌─────────────────────────────────────┐
+│  xserv-api (axum)                   │
+│  - Parse request                    │
+│  - Apply chat template              │
+│  - Submit to engine via channel     │
+│  - Stream SSE chunks from channel   │
+└────────────┬────────────────────────┘
+             │  InferenceRequest → mpsc channel
+             ▼
+┌─────────────────────────────────────┐
+│  xserv-engine (dedicated thread)    │
+│  - Receive requests                 │
+│  - Run model forward (prefill+decode)│
+│  - Send tokens back via channel     │
+└─────────────────────────────────────┘
+```
+
+## Crates
+
+- `xserv-engine`: inference orchestration (model + cache + generate loop)
+- `xserv-api`: HTTP server with axum
+
+Both in one binary: `xserv-server`
+
+## API Endpoints
+
+```
+POST /v1/chat/completions   # main endpoint
+GET  /v1/models             # list models
+GET  /health                # health check
+```
+
+## Request/Response (OpenAI compatible)
+
+Request:
+```json
+{
+  "model": "qwen3-8b",
+  "messages": [{"role": "user", "content": "Hello"}],
+  "stream": true,
+  "max_tokens": 256,
+  "temperature": 1.0
+}
+```
+
+SSE Response:
+```
+data: {"id":"...","choices":[{"delta":{"content":"Hi"},"index":0}]}
+
+data: {"id":"...","choices":[{"delta":{},"finish_reason":"stop"}]}
+
+data: [DONE]
+```
+
+## Engine Design
+
+```rust
+pub struct Engine {
+    model: Qwen3,
+    config: ModelConfig,
+    tokenizer: Tokenizer,
+}
+
+impl Engine {
+    pub fn generate(&self, prompt_tokens: &[u32], params: &SamplingParams,
+                    sender: mpsc::Sender<Token>) { ... }
+}
+```
+
+Engine runs on a dedicated OS thread (avoids async/GPU conflicts).
+API handlers communicate via `tokio::sync::mpsc` channels.
+
+## Sampling
+
+For this phase: greedy only (temperature=0 or 1 with argmax).
+Top-k/top-p sampling added later.
+
+## Test Plan
+
+- [ ] curl streaming request → get SSE chunks
+- [ ] Python OpenAI SDK client works
+- [ ] /v1/models returns model info
+- [ ] /health returns 200
+- [ ] Multiple sequential requests work