Gahow Wang
d8493bd70f
Rewrote engine.rs from scratch: - Scheduler loop: admit → prefill → decode → finish → check new requests - Multiple sequences run concurrently (max_batch_size configurable) - Each sequence has independent GpuKVCache - Non-blocking try_recv() for new requests during decode iterations - Dynamic join: new requests enter batch immediately, don't wait for others Verified with concurrent test (tools/test_concurrent.py): - 3 concurrent requests: wall_time=3.8s, concurrency_ratio=2.82x ✓ - 5 concurrent requests: wall_time=6.1s, concurrency_ratio=4.04x ✓ - All outputs are coherent and correct Design doc (docs/12-continuous-batching.md) fully rewritten with: - Detailed scheduler loop pseudocode - Data structures (Sequence, Scheduler) - Acceptance criteria with specific test cases - Clear separation from Phase 13 (HTTP layer) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
116 lines
3.1 KiB
Rust
116 lines
3.1 KiB
Rust
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).expect("engine channel closed");
|
|
|
|
// 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
|
|
}
|