Skip to main content

cortiq_server/
lib.rs

1//! Cortiq Server — OpenAI-compatible API + web management dashboard.
2
3pub mod api;
4pub mod dashboard;
5pub mod openai;
6pub mod streaming;
7
8use axum::{Json, Router, routing::get};
9use axum::extract::State;
10use cortiq_engine::{CortiqRuntime, Pipeline};
11use std::sync::Arc;
12use tokio::sync::{Mutex, OwnedMutexGuard, OwnedSemaphorePermit, Semaphore};
13use tower_http::cors::CorsLayer;
14
15/// Fixed pool of pipeline slots over ONE shared mmap'd model (roadmap
16/// §3 «serving полностью сериализован», этап 5.1): the weights are
17/// zero-copy shared through `Arc<CmfModel>`, each slot owns its
18/// KV-cache / recurrent state / sampler / workspace. A request checks a
19/// slot out for the duration of one generation, so up to `slots`
20/// requests decode CONCURRENTLY; excess requests queue fairly on the
21/// semaphore. This is bounded-concurrency serving, not yet continuous
22/// batching (этап 5.2+).
23pub struct PipelinePool {
24    slots: Vec<Arc<Mutex<Pipeline>>>,
25    sem: Arc<Semaphore>,
26}
27
28/// A checked-out slot: holds both the concurrency permit and the
29/// pipeline lock until dropped.
30pub struct SlotGuard {
31    pub pipe: OwnedMutexGuard<Pipeline>,
32    // Keep the permit after the mutex guard so drop glue unlocks the
33    // pipeline before another waiter can acquire the permit.
34    _permit: OwnedSemaphorePermit,
35}
36
37impl PipelinePool {
38    pub fn new(pipelines: Vec<Pipeline>) -> Self {
39        assert!(
40            !pipelines.is_empty(),
41            "pipeline pool needs at least one slot"
42        );
43        let sem = Arc::new(Semaphore::new(pipelines.len()));
44        Self {
45            slots: pipelines
46                .into_iter()
47                .map(|p| Arc::new(Mutex::new(p)))
48                .collect(),
49            sem,
50        }
51    }
52
53    pub fn n_slots(&self) -> usize {
54        self.slots.len()
55    }
56
57    /// Wait for a free slot and check it out. With `permits == slots`,
58    /// holding a permit guarantees the try_lock scan finds a free slot.
59    pub async fn acquire(&self) -> SlotGuard {
60        let permit = self
61            .sem
62            .clone()
63            .acquire_owned()
64            .await
65            .expect("slot semaphore closed");
66        for s in &self.slots {
67            if let Ok(pipe) = s.clone().try_lock_owned() {
68                return SlotGuard {
69                    pipe,
70                    _permit: permit,
71                };
72            }
73        }
74        unreachable!("semaphore permit held but every slot is locked")
75    }
76}
77
78/// Shared application state: runtime (masks, metrics), a tokenizer
79/// handle that never blocks on generation, and the slot pool.
80pub struct AppState {
81    pub runtime: CortiqRuntime,
82    pub tokenizer: Arc<cortiq_engine::tokenizer::Tokenizer>,
83    pub slots: PipelinePool,
84}
85
86/// Liveness probe — returns 200 as soon as the server is accepting
87/// connections. Used by process managers that embed `cortiq serve` (e.g.
88/// a gateway spawning it as a local model server) to know when it is ready.
89/// Also advertises the loaded model's capabilities so managers can route
90/// capability-gated traffic (tool calling) without manual configuration:
91/// tools are "supported" when the model's chat template has a tools branch.
92async fn healthz(State(st): State<Arc<AppState>>) -> Json<serde_json::Value> {
93    let tools = st
94        .tokenizer
95        .chat_template
96        .as_deref()
97        .map(|t| t.contains("tool"))
98        .unwrap_or(false);
99    Json(serde_json::json!({
100        "status": "ok",
101        "capabilities": { "tools": tools }
102    }))
103}
104
105/// Build the full router with all endpoints.
106pub fn build_router(state: Arc<AppState>) -> Router {
107    Router::new()
108        .route("/healthz", get(healthz))
109        .merge(openai::routes())
110        .merge(api::routes())
111        .merge(dashboard::routes())
112        .layer(CorsLayer::permissive())
113        .with_state(state)
114}