clawft_kernel/workload_runtime/infer/
capabilities.rs

1//! Capability probes: what a node may advertise after checking the runtime
2//! is really there (ADR-101 section 4).
3
4use std::time::Duration;
5
6use clawft_types::placement::{AttrValue, Capability, CapabilityId, Provenance};
7
8use super::config::{InferConfig, InferMode};
9use super::probe::{Health, ServerClient};
10use super::spec::InferFlavor;
11
12fn cap(id: &str) -> Option<Capability> {
13    CapabilityId::new(id)
14        .ok()
15        .map(|id| Capability::new(id, Provenance::Probed))
16}
17
18fn is_executable(p: &std::path::Path) -> bool {
19    use std::os::unix::fs::PermissionsExt;
20    std::fs::metadata(p).is_ok_and(|m| m.is_file() && m.permissions().mode() & 0o111 != 0)
21}
22
23async fn run_version(prog: &std::path::Path, args: &[String]) -> Option<String> {
24    let mut cmd = tokio::process::Command::new(prog);
25    cmd.args(args)
26        .env_clear()
27        .env("PATH", "/usr/local/bin:/usr/bin:/bin")
28        .stdin(std::process::Stdio::null())
29        .kill_on_drop(true);
30    let out = tokio::time::timeout(Duration::from_secs(5), cmd.output())
31        .await
32        .ok()?
33        .ok()?;
34    if !out.status.success() {
35        return None;
36    }
37    let text = String::from_utf8_lossy(&out.stdout);
38    Some(
39        text.lines()
40            .next()
41            .unwrap_or("")
42            .chars()
43            .take(120)
44            .collect(),
45    )
46}
47
48/// Probe and build the capability list. Empty when the runtime is not
49/// there: a node claims nothing it did not just check.
50pub async fn probe(cfg: &InferConfig, probe_port: u16) -> Vec<Capability> {
51    let mut version: Option<String> = None;
52    let present = match (&cfg.mode, cfg.flavor) {
53        (InferMode::Managed(m), f) if f != InferFlavor::Ollama => {
54            if !m.serve_program.as_deref().is_some_and(is_executable) {
55                return Vec::new();
56            }
57            match &m.version_probe {
58                Some((p, a)) => match run_version(p, a).await {
59                    Some(v) => {
60                        version = Some(v);
61                        true
62                    }
63                    None => false,
64                },
65                None => true,
66            }
67        }
68        _ => {
69            let Ok(client) =
70                ServerClient::new(format!("http://127.0.0.1:{probe_port}"), cfg.probe_timeout)
71            else {
72                return Vec::new();
73            };
74            let r = client.probe(cfg.flavor).await;
75            version = r.version.clone();
76            // Loading is a server that is there; a broken or absent one
77            // is not a capability.
78            matches!(r.health, Health::Up | Health::Loading)
79        }
80    };
81    if !present {
82        return Vec::new();
83    }
84    let mode = match cfg.mode {
85        InferMode::Adopted => "adopted",
86        InferMode::Managed(_) => "managed",
87    };
88    let formats: Vec<AttrValue> = cfg
89        .flavor
90        .formats()
91        .iter()
92        .map(|f| AttrValue::from(*f))
93        .collect();
94    let mut runtime = match cap(cfg.flavor.capability()) {
95        Some(c) => c
96            .with_attr("mode", mode)
97            .with_attr("formats", AttrValue::List(formats)),
98        None => return Vec::new(),
99    };
100    if let Some(v) = &version {
101        runtime = runtime.with_attr("version", v.as_str());
102    }
103    let mut out = vec![runtime];
104    for f in cfg.flavor.formats() {
105        out.extend(cap(&format!("format.{f}")));
106    }
107    if cfg!(all(target_os = "macos", target_arch = "aarch64")) {
108        out.extend(cap("accel.gpu.metal"));
109    }
110    out
111}
clawft_kernel/workload_runtime/infer/capabilities.rs

clawft_kernel/workload_runtime/infer/
capabilities.rs
