cortiq_server/
lib.rs

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

cortiq_server/
lib.rs
