clawft_kernel/workload_runtime/infer/
runtime.rs

1//! [`InferRuntime`]: the `infer.llamacpp`, `infer.mlx-lm` and
2//! `infer.ollama` adapters (ADR-101 section 4) behind one type.
3//!
4//! Two constructions of the same adapter:
5//!
6//! - **Adopted** ([`InferConfig::adopted`], id `infer.<x>.adopted`): the
7//!   adapter registers and health-checks a server somebody else started.
8//!   It never starts, stops or signals it: `start` only confirms the
9//!   server answers, `stop` is refused, `unload` forgets the registration.
10//! - **Managed** ([`InferConfig::managed`], id `infer.<x>`): llama.cpp and
11//!   mlx-lm are launched through the model lab's launcher under
12//!   [`Supervised`]; Ollama is driven through its API. Weights come only
13//!   from the model registry ([`crate::model_manifest::ModelRegistry`]),
14//!   never from a path in the spec.
15//!
16//! Adapters make no governance decisions; wrap one in
17//! [`crate::workload_runtime::WorkloadHost`] to gate and chain it.
18
19use std::collections::HashMap;
20use std::path::PathBuf;
21use std::sync::RwLock;
22use std::time::{Duration, Instant};
23
24use async_trait::async_trait;
25use clawft_types::placement::Capability;
26use tokio::sync::Mutex;
27use tokio::task::JoinHandle;
28
29use super::capabilities;
30use super::config::{InferConfig, InferMode, ManagedConfig};
31use super::exposure::port_in_use;
32use super::launch::{Launch, build_launch, check_format};
33use super::ollama::SharedLoad;
34use super::probe::{ServerClient, ServerReport};
35use super::spec::{InferFlavor, InferenceSpec};
36use crate::model_manifest::{ModelError, ResolvedModel};
37use crate::workload_governance::NetworkPolicy;
38use crate::workload_pkg::manifest::valid_token;
39use crate::workload_runtime::evidence::RunEvidence;
40use crate::workload_runtime::supervise::Supervised;
41use crate::workload_runtime::types::{
42    Admission, ControlMode, InstanceHandle, InstanceStatus, RuntimeError, VerifiedWorkload,
43    WorkloadConfig, WorkloadRuntime,
44};
45
46/// How a managed instance is run.
47pub(super) enum ManagedPlan {
48    /// A launcher script under supervision.
49    Process(Launch),
50    /// Ollama, driven through its API; the tag is what it calls the model.
51    Ollama { tag: String },
52}
53
54/// Managed-mode state of one instance.
55pub(super) struct Managed {
56    pub plan: ManagedPlan,
57    /// Display name of the model (registry name).
58    pub model_name: String,
59    pub proc: Option<Supervised>,
60    /// Set by `start`, cleared by `stop`: whether a dead server is a fault.
61    pub wanted: bool,
62    pub restarts: u32,
63    pub exited_at: Option<Instant>,
64    /// When the current process was spawned.
65    pub started_at: Option<Instant>,
66    pub last: Option<RunEvidence>,
67    /// Set when the server was found listening beyond loopback and was
68    /// stopped for it; the instance will not start again until reloaded.
69    pub exposed: Option<Vec<std::net::IpAddr>>,
70    pub load: Option<(SharedLoad, JoinHandle<()>)>,
71    /// Ollama: this adapter brought the model into memory (it was not
72    /// resident when asked), so only then may it take it out again.
73    pub we_loaded: std::sync::Arc<std::sync::atomic::AtomicBool>,
74}
75
76/// One registered instance.
77pub(super) struct Instance {
78    pub spec: InferenceSpec,
79    pub client: ServerClient,
80    pub managed: Option<Managed>,
81}
82
83/// The result of planning a workload (shared by admit and load).
84struct Planned {
85    spec: InferenceSpec,
86    managed: Option<(ManagedPlan, String)>,
87    notes: Vec<String>,
88}
89
90/// The adapter.
91pub struct InferRuntime {
92    pub(super) cfg: InferConfig,
93    id: String,
94    caps: RwLock<Vec<Capability>>,
95    pub(super) instances: Mutex<HashMap<String, Instance>>,
96}
97
98fn refuse(m: impl Into<String>) -> RuntimeError {
99    RuntimeError::AdmissionRefused(m.into())
100}
101
102fn model_err(what: &str, e: ModelError) -> RuntimeError {
103    refuse(format!("{what}: {e}"))
104}
105
106impl InferRuntime {
107    /// A new adapter. It claims no capabilities until
108    /// [`probe_capabilities`](Self::probe_capabilities) has run.
109    pub fn new(cfg: InferConfig) -> Self {
110        let id = match cfg.mode {
111            InferMode::Adopted => format!("{}.adopted", cfg.flavor.id()),
112            InferMode::Managed(_) => cfg.flavor.id().to_string(),
113        };
114        Self {
115            cfg,
116            id,
117            caps: RwLock::new(Vec::new()),
118            instances: Mutex::new(HashMap::new()),
119        }
120    }
121
122    pub(super) fn managed_cfg(&self) -> Option<&ManagedConfig> {
123        match &self.cfg.mode {
124            InferMode::Managed(m) => Some(m),
125            InferMode::Adopted => None,
126        }
127    }
128
129    fn flavor(&self) -> InferFlavor {
130        self.cfg.flavor
131    }
132
133    /// Check the runtime is there and refresh what `provides()` returns.
134    pub async fn probe_capabilities(&self) -> Vec<Capability> {
135        let port = self
136            .cfg
137            .probe_port
138            .unwrap_or_else(|| self.flavor().default_port());
139        let caps = capabilities::probe(&self.cfg, port).await;
140        if let Ok(mut g) = self.caps.write() {
141            *g = caps.clone();
142        }
143        caps
144    }
145
146    fn client_for(&self, spec: &InferenceSpec) -> Result<ServerClient, RuntimeError> {
147        ServerClient::new(spec.base_url()?, self.cfg.probe_timeout)
148    }
149
150    fn resolve(&self, m: &ManagedConfig, name: &str) -> Result<ResolvedModel, RuntimeError> {
151        m.models
152            .resolve(name)
153            .map_err(|e| model_err(&format!("model {name}"), e))
154    }
155
156    async fn plan(&self, w: &VerifiedWorkload) -> Result<Planned, RuntimeError> {
157        let spec = w.inference_spec(&self.id)?.clone();
158        spec.validate()?;
159        if spec.runtime != self.flavor() {
160            return Err(refuse(format!(
161                "{} cannot run a {} workload",
162                self.id,
163                spec.runtime.id()
164            )));
165        }
166        if let Some((lo, hi)) = self.cfg.allowed_ports {
167            let p = spec.port()?;
168            if !(lo..=hi).contains(&p) {
169                return Err(refuse(format!(
170                    "port {p} is outside this adapter's allowed range {lo}-{hi}"
171                )));
172            }
173        }
174        let client = self.client_for(&spec)?;
175        let mut notes = vec![format!("endpoint {}", client.base())];
176        let Some(m) = self.managed_cfg() else {
177            // Adopted: the server must already answer.
178            let r = client.probe(spec.runtime).await;
179            if let super::probe::Health::Unreachable(why) = &r.health {
180                return Err(refuse(format!(
181                    "no {} server answering at {}: {why}",
182                    spec.runtime.id(),
183                    client.base()
184                )));
185            }
186            notes.push(format!("adopting running server ({:?})", r.health));
187            notes.push(format!("{} model(s) listed", r.models.len()));
188            if let Some(v) = &r.version {
189                notes.push(format!("version {v}"));
190            }
191            return Ok(Planned {
192                spec,
193                managed: None,
194                notes,
195            });
196        };
197        let name = spec
198            .model
199            .clone()
200            .ok_or_else(|| refuse("managed mode needs a model"))?;
201        let model = self.resolve(m, &name)?;
202        check_format(spec.runtime, &model)?;
203        notes.push(format!(
204            "model {} ({}) verified",
205            model.name,
206            model.format.as_str()
207        ));
208        if spec.runtime == InferFlavor::Ollama {
209            let r = client.probe(spec.runtime).await;
210            if matches!(r.health, super::probe::Health::Unreachable(_)) {
211                return Err(refuse(format!(
212                    "ollama is not answering at {}; the adapter drives it but does not start it",
213                    client.base()
214                )));
215            }
216            let tag = model
217                .body
218                .source
219                .ollama_tag
220                .clone()
221                .unwrap_or_else(|| model.name.clone());
222            if !r.lists(&tag) {
223                return Err(refuse(format!(
224                    "ollama has no model {tag}: pull it first (the adapter never downloads)"
225                )));
226            }
227            return Ok(Planned {
228                spec,
229                managed: Some((ManagedPlan::Ollama { tag }, model.name)),
230                notes,
231            });
232        }
233        let draft = match &spec.serve.draft_model {
234            Some(d) => Some(self.resolve(m, d)?),
235            None => None,
236        };
237        let prog: &PathBuf = m
238            .serve_program
239            .as_ref()
240            .ok_or_else(|| refuse("managed mode needs a serve program"))?;
241        if !capabilities_is_exec(prog) {
242            return Err(refuse(format!(
243                "serve program {} is not executable",
244                prog.display()
245            )));
246        }
247        let (ip, port) = (spec.bind_ip()?, spec.port()?);
248        if port_in_use(ip, port).await {
249            return Err(refuse(format!(
250                "something is already listening on {ip}:{port}; adopt it instead of managing over it"
251            )));
252        }
253        let probe_id = "plan";
254        let launch = build_launch(spec.runtime, &spec, &model, draft.as_ref(), m, probe_id)?;
255        notes.push(format!(
256            "launch {} with {} args",
257            launch.program.display(),
258            launch.args.len()
259        ));
260        Ok(Planned {
261            spec,
262            managed: Some((ManagedPlan::Process(launch), model.name)),
263            notes,
264        })
265    }
266}
267
268fn capabilities_is_exec(p: &std::path::Path) -> bool {
269    use std::os::unix::fs::PermissionsExt;
270    std::fs::metadata(p).is_ok_and(|m| m.is_file() && m.permissions().mode() & 0o111 != 0)
271}
272
273/// Instance id: `<role>-<8 hex of identity hash>-<node>`.
274fn instance_id(spec: &InferenceSpec, mode: &str, node_id: &str) -> Result<String, RuntimeError> {
275    if !valid_token(node_id, 64) {
276        return Err(RuntimeError::InvalidConfig(
277            "node id must be a plain token".into(),
278        ));
279    }
280    let mut h = blake3::Hasher::new();
281    h.update(spec.role.as_bytes());
282    h.update(spec.runtime.id().as_bytes());
283    h.update(spec.model.as_deref().unwrap_or("").as_bytes());
284    h.update(spec.base_url()?.as_bytes());
285    h.update(mode.as_bytes());
286    Ok(format!(
287        "{}-{}-{}",
288        spec.role,
289        &h.finalize().to_hex()[..8],
290        node_id
291    ))
292}
293
294#[async_trait]
295impl WorkloadRuntime for InferRuntime {
296    fn id(&self) -> &str {
297        &self.id
298    }
299
300    fn provides(&self) -> Vec<Capability> {
301        self.caps.read().map(|c| c.clone()).unwrap_or_default()
302    }
303
304    async fn admit(&self, w: &VerifiedWorkload) -> Result<Admission, RuntimeError> {
305        let p = self.plan(w).await?;
306        Ok(Admission {
307            runtime: self.id.clone(),
308            arch: std::env::consts::ARCH.into(),
309            emulated: false,
310            notes: p.notes,
311        })
312    }
313
314    async fn load(
315        &self,
316        w: &VerifiedWorkload,
317        cfg: &WorkloadConfig,
318    ) -> Result<InstanceHandle, RuntimeError> {
319        if cfg.mode != crate::workload_runtime::types::RunMode::Listener || !cfg.args.is_empty() {
320            return Err(RuntimeError::InvalidConfig(
321                "inference instances run as listeners and take their arguments from the spec"
322                    .into(),
323            ));
324        }
325        let mode = if self.managed_cfg().is_some() {
326            "managed"
327        } else {
328            "adopted"
329        };
330        let p = self.plan(w).await?;
331        let id = instance_id(&p.spec, mode, &cfg.node_id)?;
332        let client = self.client_for(&p.spec)?;
333        let mut managed = None;
334        if let Some((plan, model_name)) = p.managed {
335            let mut plan = plan;
336            if let (ManagedPlan::Process(l), Some(m)) = (&mut plan, self.managed_cfg()) {
337                l.dir = m.data_root.join(&id);
338                std::fs::create_dir_all(&l.dir).map_err(|e| {
339                    RuntimeError::Backend(format!("create {}: {e}", l.dir.display()))
340                })?;
341            }
342            managed = Some(Managed {
343                plan,
344                model_name,
345                proc: None,
346                wanted: false,
347                restarts: 0,
348                exited_at: None,
349                started_at: None,
350                last: None,
351                exposed: None,
352                load: None,
353                we_loaded: Default::default(),
354            });
355        }
356        let mut g = self.instances.lock().await;
357        if g.contains_key(&id) {
358            return Err(RuntimeError::InvalidState(format!(
359                "{id} is already loaded"
360            )));
361        }
362        let (ip, port) = (p.spec.bind_ip()?, p.spec.port()?);
363        if g.values().any(|i| {
364            i.spec.bind_ip().ok() == Some(ip)
365                && i.spec.port().ok() == Some(port)
366                && (managed.is_some() || i.managed.is_some())
367        }) {
368            return Err(refuse(format!(
369                "{ip}:{port} is already used by another instance"
370            )));
371        }
372        g.insert(
373            id.clone(),
374            Instance {
375                spec: p.spec.clone(),
376                client,
377                managed,
378            },
379        );
380        Ok(InstanceHandle {
381            runtime: self.id.clone(),
382            instance_id: id,
383            workload_id: p.spec.role,
384            store_installed: false,
385        })
386    }
387
388    async fn start(&self, h: &InstanceHandle) -> Result<(), RuntimeError> {
389        self.start_instance(h).await
390    }
391
392    async fn stop(&self, h: &InstanceHandle, grace: Duration) -> Result<RunEvidence, RuntimeError> {
393        self.stop_instance(h, grace).await
394    }
395
396    async fn unload(&self, h: InstanceHandle) -> Result<(), RuntimeError> {
397        self.unload_instance(&h).await
398    }
399
400    async fn status(&self, h: &InstanceHandle) -> InstanceStatus {
401        self.status_of(h).await
402    }
403
404    fn control_mode(&self) -> ControlMode {
405        match self.cfg.mode {
406            InferMode::Adopted => ControlMode::Adopted,
407            InferMode::Managed(_) => ControlMode::Managed,
408        }
409    }
410
411    /// A managed server binds loopback only (a wider bind is refused at
412    /// validation). An adopted server's own binding cannot be observed
413    /// from here (the model lab starts its servers with `--host 0.0.0.0`),
414    /// so it is reported as LAN-reachable rather than assumed private.
415    fn network_exposure(&self) -> NetworkPolicy {
416        match self.cfg.mode {
417            InferMode::Adopted => NetworkPolicy::Lan,
418            InferMode::Managed(_) => NetworkPolicy::None,
419        }
420    }
421
422    async fn adopt(&self, h: &InstanceHandle, w: &VerifiedWorkload) -> Result<(), RuntimeError> {
423        if self.managed_cfg().is_some() {
424            return Err(RuntimeError::Unsupported(format!(
425                "{} cannot re-adopt {}: a managed server's process is not recoverable",
426                self.id, h.instance_id
427            )));
428        }
429        let p = self.plan(w).await?;
430        let client = self.client_for(&p.spec)?;
431        self.instances
432            .lock()
433            .await
434            .entry(h.instance_id.clone())
435            .or_insert(Instance {
436                spec: p.spec,
437                client,
438                managed: None,
439            });
440        Ok(())
441    }
442
443    async fn console(&self, h: &InstanceHandle, _c: &str) -> Result<RunEvidence, RuntimeError> {
444        Err(RuntimeError::Unsupported(format!(
445            "{} has no console run ({})",
446            self.id, h.instance_id
447        )))
448    }
449}
450
451impl InferRuntime {
452    /// Full probe of an instance's server.
453    pub async fn health(&self, h: &InstanceHandle) -> Result<ServerReport, RuntimeError> {
454        let (client, flavor) = {
455            let g = self.instances.lock().await;
456            let i = g
457                .get(&h.instance_id)
458                .ok_or_else(|| RuntimeError::UnknownInstance(h.instance_id.clone()))?;
459            (i.client.clone(), i.spec.runtime)
460        };
461        Ok(client.probe(flavor).await)
462    }
463
464    /// Base URL of an instance (`http://127.0.0.1:<port>`): what the
465    /// stable-address proxy forwards to.
466    pub async fn endpoint(&self, h: &InstanceHandle) -> Option<String> {
467        self.instances
468            .lock()
469            .await
470            .get(&h.instance_id)
471            .map(|i| i.client.base().to_string())
472    }
473
474    /// The model id the server itself knows the instance by (the Ollama tag,
475    /// the registry name for a launched server, the spec's model when
476    /// adopted). The mesh-serving side pins requests to it.
477    pub async fn served_model(&self, h: &InstanceHandle) -> Option<String> {
478        let g = self.instances.lock().await;
479        let i = g.get(&h.instance_id)?;
480        match &i.managed {
481            Some(Managed {
482                plan: ManagedPlan::Ollama { tag },
483                ..
484            }) => Some(tag.clone()),
485            Some(m) => Some(m.model_name.clone()),
486            None => i.spec.model.clone(),
487        }
488    }
489
490    /// The spec of a loaded instance.
491    pub async fn spec_of(&self, h: &InstanceHandle) -> Option<InferenceSpec> {
492        self.instances
493            .lock()
494            .await
495            .get(&h.instance_id)
496            .map(|i| i.spec.clone())
497    }
498
499    /// Health plus one real token (ADR-101: "a tiny completion succeeds").
500    /// Costs compute, so it is separate from [`status`](WorkloadRuntime::status).
501    pub async fn deep_health(&self, h: &InstanceHandle) -> Result<(), RuntimeError> {
502        let (client, model, flavor) = {
503            let g = self.instances.lock().await;
504            let i = g
505                .get(&h.instance_id)
506                .ok_or_else(|| RuntimeError::UnknownInstance(h.instance_id.clone()))?;
507            let model = i
508                .managed
509                .as_ref()
510                .map(|m| m.model_name.clone())
511                .or_else(|| i.spec.model.clone());
512            (i.client.clone(), model, i.spec.runtime)
513        };
514        let client = ServerClient::new(client.base().to_string(), self.cfg.deep_timeout)?;
515        let report = client.probe(flavor).await;
516        let model = model
517            .or_else(|| report.models.first().cloned())
518            .ok_or_else(|| RuntimeError::Backend("no model to ask".into()))?;
519        client
520            .tiny_completion(&model)
521            .await
522            .map_err(|e| RuntimeError::Backend(format!("deep health: {e}")))
523    }
524}
clawft_kernel/workload_runtime/infer/runtime.rs

clawft_kernel/workload_runtime/infer/
runtime.rs
