clawft_kernel/workload_runtime/infer/
probe.rs

1//! Loopback HTTP probes of model servers: health, listed models, and an
2//! optional tiny completion. Never follows redirects, ignores proxy
3//! environment, caps response size, and only ever talks to the address the
4//! caller built from a validated loopback spec.
5
6use std::time::{Duration, Instant};
7
8use serde_json::{Value, json};
9
10use super::spec::InferFlavor;
11use crate::workload_runtime::types::RuntimeError;
12
13/// Largest response body read from a server.
14const MAX_BODY: usize = 1024 * 1024;
15
16/// What a server reported.
17#[derive(Debug, Clone, PartialEq)]
18pub enum Health {
19    /// Reachable and serving.
20    Up,
21    /// Reachable, still loading the model (HTTP 503).
22    Loading,
23    /// Reachable but unhealthy.
24    Unhealthy(String),
25    /// Nothing answered.
26    Unreachable(String),
27}
28
29/// One probe of a server.
30#[derive(Debug, Clone, PartialEq)]
31pub struct ServerReport {
32    /// Health.
33    pub health: Health,
34    /// Model ids the server lists (`/v1/models`, or Ollama `/api/tags`).
35    pub models: Vec<String>,
36    /// Models resident in memory (Ollama `/api/ps`); empty elsewhere.
37    pub loaded: Vec<String>,
38    /// Server version when it says (Ollama).
39    pub version: Option<String>,
40    /// Round trip of the health request.
41    pub latency: Duration,
42}
43
44impl ServerReport {
45    fn down(why: String, latency: Duration) -> Self {
46        Self {
47            health: Health::Unreachable(why),
48            models: Vec::new(),
49            loaded: Vec::new(),
50            version: None,
51            latency,
52        }
53    }
54
55    /// Whether `model` is listed (exact id, or the id's `:latest`-less form).
56    pub fn lists(&self, model: &str) -> bool {
57        self.models.iter().any(|m| same_model(m, model))
58    }
59
60    /// Whether `model` is resident in memory.
61    pub fn resident(&self, model: &str) -> bool {
62        self.loaded.iter().any(|m| same_model(m, model))
63    }
64}
65
66/// Servers name models differently: Ollama adds `:latest` to an untagged
67/// pull, llama-server lists the GGUF path, mlx-lm the HF repo id. A listed
68/// id matches when it is the wanted name, or its last path segment (minus
69/// `.gguf`) contains it.
70fn same_model(listed: &str, wanted: &str) -> bool {
71    if listed == wanted
72        || listed.strip_suffix(":latest") == Some(wanted)
73        || wanted.strip_suffix(":latest") == Some(listed)
74    {
75        return true;
76    }
77    let base = listed.rsplit('/').next().unwrap_or(listed);
78    let base = base.strip_suffix(".gguf").unwrap_or(base);
79    base.to_ascii_lowercase()
80        .contains(&wanted.to_ascii_lowercase())
81}
82
83/// A client bound to one server.
84#[derive(Clone)]
85pub struct ServerClient {
86    http: reqwest::Client,
87    base: String,
88}
89
90impl ServerClient {
91    /// Client for `base` (`http://127.0.0.1:PORT`) with a per-request timeout.
92    pub fn new(base: String, timeout: Duration) -> Result<Self, RuntimeError> {
93        let http = reqwest::Client::builder()
94            .redirect(reqwest::redirect::Policy::none())
95            .no_proxy()
96            // A probe must see the server as it is now, not a pooled socket.
97            .pool_max_idle_per_host(0)
98            .connect_timeout(timeout)
99            .timeout(timeout)
100            .build()
101            .map_err(|e| RuntimeError::Backend(format!("http client: {e}")))?;
102        Ok(Self { http, base })
103    }
104
105    /// Base URL.
106    pub fn base(&self) -> &str {
107        &self.base
108    }
109
110    async fn read(resp: reqwest::Response) -> Result<(u16, String), String> {
111        let status = resp.status().as_u16();
112        let mut resp = resp;
113        let mut buf = Vec::new();
114        while let Some(chunk) = resp.chunk().await.map_err(|e| e.to_string())? {
115            buf.extend_from_slice(&chunk);
116            if buf.len() > MAX_BODY {
117                return Err("response too large".into());
118            }
119        }
120        Ok((status, String::from_utf8_lossy(&buf).into_owned()))
121    }
122
123    /// GET `path`.
124    pub async fn get(&self, path: &str) -> Result<(u16, String), String> {
125        let r = self
126            .http
127            .get(format!("{}{path}", self.base))
128            .send()
129            .await
130            .map_err(|e| e.to_string())?;
131        Self::read(r).await
132    }
133
134    /// POST JSON to `path`.
135    pub async fn post(&self, path: &str, body: &Value) -> Result<(u16, String), String> {
136        let r = self
137            .http
138            .post(format!("{}{path}", self.base))
139            .json(body)
140            .send()
141            .await
142            .map_err(|e| e.to_string())?;
143        Self::read(r).await
144    }
145
146    /// Probe the server the way `flavor` is probed: `/health` then
147    /// `/v1/models` for llama.cpp and mlx-lm (a 404 on `/health` falls back
148    /// to `/v1/models` alone), `/api/version`, `/api/tags` and `/api/ps`
149    /// for Ollama.
150    pub async fn probe(&self, flavor: InferFlavor) -> ServerReport {
151        match flavor {
152            InferFlavor::Ollama => self.probe_ollama().await,
153            _ => self.probe_openai().await,
154        }
155    }
156
157    async fn probe_openai(&self) -> ServerReport {
158        let t = Instant::now();
159        let health = match self.get("/health").await {
160            Ok((200, _)) => Health::Up,
161            Ok((503, _)) => Health::Loading,
162            Ok((404, _)) => match self.get("/v1/models").await {
163                Ok((200, _)) => Health::Up,
164                Ok((503, _)) => Health::Loading,
165                Ok((c, _)) => Health::Unhealthy(format!("/v1/models returned HTTP {c}")),
166                Err(e) => return ServerReport::down(e, t.elapsed()),
167            },
168            Ok((c, _)) => Health::Unhealthy(format!("/health returned HTTP {c}")),
169            Err(e) => return ServerReport::down(e, t.elapsed()),
170        };
171        let latency = t.elapsed();
172        let models = if health == Health::Up {
173            match self.get("/v1/models").await {
174                Ok((200, body)) => parse_ids(&body, "data", "id"),
175                _ => Vec::new(),
176            }
177        } else {
178            Vec::new()
179        };
180        ServerReport {
181            health,
182            models,
183            loaded: Vec::new(),
184            version: None,
185            latency,
186        }
187    }
188
189    async fn probe_ollama(&self) -> ServerReport {
190        let t = Instant::now();
191        let (health, version) = match self.get("/api/version").await {
192            Ok((200, body)) => (
193                Health::Up,
194                serde_json::from_str::<Value>(&body)
195                    .ok()
196                    .and_then(|v| v["version"].as_str().map(str::to_string)),
197            ),
198            Ok((503, _)) => (Health::Loading, None),
199            Ok((c, _)) => (
200                Health::Unhealthy(format!("/api/version returned HTTP {c}")),
201                None,
202            ),
203            Err(e) => return ServerReport::down(e, t.elapsed()),
204        };
205        let latency = t.elapsed();
206        let (mut models, mut loaded) = (Vec::new(), Vec::new());
207        if health == Health::Up {
208            if let Ok((200, body)) = self.get("/api/tags").await {
209                models = parse_ids(&body, "models", "name");
210            }
211            if let Ok((200, body)) = self.get("/api/ps").await {
212                loaded = parse_ids(&body, "models", "name");
213            }
214        }
215        ServerReport {
216            health,
217            models,
218            loaded,
219            version,
220            latency,
221        }
222    }
223
224    /// Ask for one token. Costs real compute, so callers opt in.
225    pub async fn tiny_completion(&self, model: &str) -> Result<(), String> {
226        let body = json!({
227            "model": model,
228            "messages": [{"role": "user", "content": "hi"}],
229            "max_tokens": 1,
230            "stream": false,
231        });
232        match self.post("/v1/chat/completions", &body).await? {
233            (200, _) => Ok(()),
234            (c, _) => Err(format!("completion returned HTTP {c}")),
235        }
236    }
237}
238
239fn parse_ids(body: &str, list: &str, key: &str) -> Vec<String> {
240    serde_json::from_str::<Value>(body)
241        .ok()
242        .and_then(|v| {
243            v[list].as_array().map(|a| {
244                a.iter()
245                    .filter_map(|m| m[key].as_str().map(str::to_string))
246                    .collect()
247            })
248        })
249        .unwrap_or_default()
250}
clawft_kernel/workload_runtime/infer/probe.rs

clawft_kernel/workload_runtime/infer/
probe.rs
