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    /// GPU each slot's weights live on (replica mode: slot i → card i).
26    /// Empty = single-device, every slot on the process default.
27    devices: Vec<usize>,
28    sem: Arc<Semaphore>,
29}
30
31/// A checked-out slot: holds both the concurrency permit and the
32/// pipeline lock until dropped.
33pub struct SlotGuard {
34    pub pipe: OwnedMutexGuard<Pipeline>,
35    /// The card this slot's weights are on. The handler thread is
36    /// pinned to it for the whole request — the engine resolves its
37    /// device context (and therefore its weight cache) through that pin.
38    pub device: usize,
39    // Keep the permit after the mutex guard so drop glue unlocks the
40    // pipeline before another waiter can acquire the permit.
41    _permit: OwnedSemaphorePermit,
42}
43
44impl PipelinePool {
45    pub fn new(pipelines: Vec<Pipeline>) -> Self {
46        let n = pipelines.len();
47        Self::with_devices(pipelines, vec![cortiq_engine::gpu::default_device(); n])
48    }
49
50    /// Replica mode: `devices[i]` is the card slot i was loaded on.
51    pub fn with_devices(pipelines: Vec<Pipeline>, devices: Vec<usize>) -> Self {
52        assert!(
53            !pipelines.is_empty(),
54            "pipeline pool needs at least one slot"
55        );
56        assert_eq!(
57            pipelines.len(),
58            devices.len(),
59            "one device per slot: {} pipelines, {} devices",
60            pipelines.len(),
61            devices.len()
62        );
63        let sem = Arc::new(Semaphore::new(pipelines.len()));
64        Self {
65            slots: pipelines
66                .into_iter()
67                .map(|p| Arc::new(Mutex::new(p)))
68                .collect(),
69            devices,
70            sem,
71        }
72    }
73
74    pub fn n_slots(&self) -> usize {
75        self.slots.len()
76    }
77
78    /// Wait for a free slot and check it out. With `permits == slots`,
79    /// holding a permit guarantees the try_lock scan finds a free slot.
80    pub async fn acquire(&self) -> SlotGuard {
81        let permit = self
82            .sem
83            .clone()
84            .acquire_owned()
85            .await
86            .expect("slot semaphore closed");
87        for (i, s) in self.slots.iter().enumerate() {
88            if let Ok(pipe) = s.clone().try_lock_owned() {
89                let device = self.devices[i];
90                // Pin the caller's thread: everything this request does
91                // downstream — including the worker pool, which carries
92                // the pin with each dispatch — addresses this card.
93                cortiq_engine::gpu::set_current_device(device);
94                return SlotGuard {
95                    pipe,
96                    device,
97                    _permit: permit,
98                };
99            }
100        }
101        unreachable!("semaphore permit held but every slot is locked")
102    }
103}
104
105/// Shared application state: runtime (masks, metrics), a tokenizer
106/// handle that never blocks on generation, and the slot pool.
107pub struct AppState {
108    pub runtime: CortiqRuntime,
109    pub tokenizer: Arc<cortiq_engine::tokenizer::Tokenizer>,
110    pub slots: PipelinePool,
111    /// Network pipeline-split worker (serve --peer). One worker holds one
112    /// KV session, so peer mode runs with exactly one slot; the mutex is
113    /// never contended (the slot semaphore already serializes) but keeps
114    /// the type honest.
115    pub remote: Option<Arc<std::sync::Mutex<cortiq_net::RemoteSegment>>>,
116}
117
118/// Liveness probe — returns 200 as soon as the server is accepting
119/// connections. Used by process managers that embed `cortiq serve` (e.g.
120/// a gateway spawning it as a local model server) to know when it is ready.
121/// Also advertises the loaded model's capabilities so managers can route
122/// capability-gated traffic (tool calling) without manual configuration:
123/// tools are "supported" when the model's chat template has a tools branch.
124async fn healthz(State(st): State<Arc<AppState>>) -> Json<serde_json::Value> {
125    let tools = st
126        .tokenizer
127        .chat_template
128        .as_deref()
129        .map(|t| t.contains("tool"))
130        .unwrap_or(false);
131    Json(serde_json::json!({
132        "status": "ok",
133        "capabilities": { "tools": tools }
134    }))
135}
136
137/// Build the full router with all endpoints.
138pub fn build_router(state: Arc<AppState>) -> Router {
139    Router::new()
140        .route("/healthz", get(healthz))
141        .merge(openai::routes())
142        .merge(api::routes())
143        .merge(dashboard::routes())
144        .layer(CorsLayer::permissive())
145        .with_state(state)
146}
cortiq_server/lib.rs

cortiq_server/
lib.rs
