clawft_kernel/workload_runtime/infer/
spec.rs

1//! The inference workload spec (ADR-101 section 1) and its boundary
2//! validation.
3
4use std::net::IpAddr;
5
6use serde::{Deserialize, Serialize};
7
8use crate::workload_pkg::manifest::valid_token;
9use crate::workload_runtime::types::{
10    InferencePayload, RuntimeError, VerifiedWorkload, WorkloadSource,
11};
12
13/// Workload kind string for inference endpoints.
14pub const KIND_INFERENCE: &str = "inference";
15
16/// Largest `extra_args` list and longest single argument.
17const MAX_EXTRA_ARGS: usize = 32;
18const MAX_ARG_LEN: usize = 256;
19
20/// Which server software an adapter drives.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
22pub enum InferFlavor {
23    /// `llama-server` (OpenAI-compatible, GGUF), usually through `serve-llamacpp`.
24    #[serde(rename = "infer.llamacpp")]
25    LlamaCpp,
26    /// `mlx_lm.server` (MLX quantisation directories), usually through `serve`.
27    #[serde(rename = "infer.mlx-lm")]
28    MlxLm,
29    /// Ollama, which owns its own process and model load/unload.
30    #[serde(rename = "infer.ollama")]
31    Ollama,
32}
33
34impl InferFlavor {
35    /// Runtime id (`infer.llamacpp`).
36    pub fn id(self) -> &'static str {
37        match self {
38            Self::LlamaCpp => "infer.llamacpp",
39            Self::MlxLm => "infer.mlx-lm",
40            Self::Ollama => "infer.ollama",
41        }
42    }
43
44    /// The capability a node advertises when this runtime works.
45    pub fn capability(self) -> &'static str {
46        match self {
47            Self::LlamaCpp => "runtime.infer.llamacpp",
48            Self::MlxLm => "runtime.infer.mlx-lm",
49            Self::Ollama => "runtime.infer.ollama",
50        }
51    }
52
53    /// Conventional port of a hand-started server. A default only: every
54    /// test and every managed instance sets its own.
55    pub fn default_port(self) -> u16 {
56        match self {
57            Self::LlamaCpp => 8090,
58            Self::MlxLm => 8081,
59            Self::Ollama => 11434,
60        }
61    }
62
63    /// `format.*` capability suffixes this runtime can load.
64    pub fn formats(self) -> &'static [&'static str] {
65        match self {
66            Self::LlamaCpp => &["gguf"],
67            Self::MlxLm => &["mlx", "safetensors"],
68            Self::Ollama => &["ollama", "gguf"],
69        }
70    }
71}
72
73/// `interactive` roles carry a latency bound; `batch` roles do not.
74#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
75#[serde(rename_all = "snake_case")]
76pub enum LatencyClass {
77    /// Voice and agent loops.
78    Interactive,
79    /// Throughput work.
80    #[default]
81    Batch,
82}
83
84/// Memory the instance is expected to hold (placement input).
85#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
86pub struct MemoryBudget {
87    /// Weights.
88    pub weights_bytes: u64,
89    /// KV cache.
90    pub kv_budget_bytes: u64,
91}
92
93/// Serve arguments. The adapter owns host and port; everything else is
94/// mapped to the runtime's flags or passed through `extra_args`.
95#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
96#[serde(deny_unknown_fields)]
97pub struct ServeArgs {
98    /// Bind / probe address. Loopback only in v1 (a wider bind is a
99    /// governed decision, ADR-101 section 7). Default `127.0.0.1`.
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub host: Option<String>,
102    /// Port; the runtime's conventional port when absent.
103    #[serde(default, skip_serializing_if = "Option::is_none")]
104    pub port: Option<u16>,
105    /// Context length.
106    #[serde(default, skip_serializing_if = "Option::is_none")]
107    pub ctx: Option<u32>,
108    /// KV quantisation (`q8_0`).
109    #[serde(default, skip_serializing_if = "Option::is_none")]
110    pub kv_quant: Option<String>,
111    /// Draft model for speculative decoding (a model registry name).
112    #[serde(default, skip_serializing_if = "Option::is_none")]
113    pub draft_model: Option<String>,
114    /// Further runtime arguments, passed as plain argv (never a shell).
115    #[serde(default, skip_serializing_if = "Vec::is_empty")]
116    pub extra_args: Vec<String>,
117}
118
119/// An inference endpoint to adopt or run.
120#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
121#[serde(deny_unknown_fields)]
122pub struct InferenceSpec {
123    /// Stable role name consumers use (`coder-daily`).
124    pub role: String,
125    /// Model registry name or package id. Required to manage an instance;
126    /// optional when adopting (then only the server is observed).
127    #[serde(default, skip_serializing_if = "Option::is_none")]
128    pub model: Option<String>,
129    /// Server software.
130    pub runtime: InferFlavor,
131    /// Serve arguments.
132    #[serde(default)]
133    pub serve: ServeArgs,
134    /// Expected memory.
135    #[serde(default)]
136    pub memory: MemoryBudget,
137    /// Co-residency exclusions (`role:planner`).
138    #[serde(default, skip_serializing_if = "Vec::is_empty")]
139    pub excludes: Vec<String>,
140    /// Latency class.
141    #[serde(default)]
142    pub latency_class: LatencyClass,
143    /// No automatic migration (warm KV).
144    #[serde(default)]
145    pub sticky: bool,
146    /// API dialect.
147    #[serde(default = "default_api")]
148    pub api: String,
149}
150
151fn default_api() -> String {
152    "openai-v1".into()
153}
154
155fn bad(m: impl Into<String>) -> RuntimeError {
156    RuntimeError::InvalidConfig(m.into())
157}
158
159/// Flags a spec may pass through, with how many value tokens each takes.
160/// An allowlist, exact spellings only: a denylist loses to argparse prefix
161/// abbreviations (`--hos`), short aliases (`-md`, `-hf`, `-mu`), synonyms
162/// (`--ctx-size`) and every flag that reads or writes files or opens a
163/// listener (`--lora*`, `--mmproj`, `--api-key`, `--ssl-*`, `--path`,
164/// `--log-file`, `--slot-save-path`). Host, port, model, draft, ctx and kv
165/// are owned by the adapter and are not on any list.
166fn allowed_flags(flavor: InferFlavor) -> &'static [(&'static str, usize)] {
167    match flavor {
168        InferFlavor::LlamaCpp => &[
169            ("--threads", 1),
170            ("-t", 1),
171            ("--threads-batch", 1),
172            ("--batch-size", 1),
173            ("-b", 1),
174            ("--ubatch-size", 1),
175            ("-ub", 1),
176            ("--parallel", 1),
177            ("-np", 1),
178            ("--n-gpu-layers", 1),
179            ("-ngl", 1),
180            ("--flash-attn", 0),
181            ("-fa", 0),
182            ("--cont-batching", 0),
183            ("--mlock", 0),
184            ("--no-mmap", 0),
185            ("--jinja", 0),
186            ("--metrics", 0),
187            ("--no-webui", 0),
188            ("--temp", 1),
189            ("--top-p", 1),
190            ("--top-k", 1),
191            ("--repeat-penalty", 1),
192            ("--seed", 1),
193        ],
194        InferFlavor::MlxLm => &[
195            ("--max-tokens", 1),
196            ("--temp", 1),
197            ("--top-p", 1),
198            ("--top-k", 1),
199            ("--log-level", 1),
200        ],
201        // Ollama is configured through its own environment, not argv.
202        InferFlavor::Ollama => &[],
203    }
204}
205
206/// A value token: a plain word, or a number (which may be negative). Never
207/// something that could be read as another flag, and no whitespace.
208fn valid_flag_value(v: &str) -> bool {
209    !v.is_empty()
210        && v.len() <= MAX_ARG_LEN
211        && (!v.starts_with('-') || v.parse::<f64>().is_ok())
212        && v.bytes()
213            .all(|b| b.is_ascii_alphanumeric() || b"._:+-/".contains(&b))
214}
215
216fn validate_extra_args(flavor: InferFlavor, args: &[String]) -> Result<(), RuntimeError> {
217    if args.len() > MAX_EXTRA_ARGS {
218        return Err(bad("too many extra_args"));
219    }
220    let allowed = allowed_flags(flavor);
221    let mut i = 0;
222    while i < args.len() {
223        let a = &args[i];
224        if a.bytes().any(|b| b.is_ascii_whitespace() || b == 0) {
225            return Err(bad(format!(
226                "extra_args entries are single tokens without whitespace: {a:?}"
227            )));
228        }
229        let Some(&(_, arity)) = allowed.iter().find(|(f, _)| f == a) else {
230            return Err(bad(format!(
231                "extra_args {a:?} is not an allowed {} flag (exact spelling, value as a separate entry)",
232                flavor.id()
233            )));
234        };
235        for v in args.iter().skip(i + 1).take(arity) {
236            if !valid_flag_value(v) {
237                return Err(bad(format!(
238                    "extra_args value {v:?} for {a} is not a plain word or number"
239                )));
240            }
241        }
242        if i + arity >= args.len() && arity > 0 {
243            return Err(bad(format!("extra_args {a} needs {arity} value(s)")));
244        }
245        i += 1 + arity;
246    }
247    Ok(())
248}
249
250/// Largest memory figure accepted (4 TiB): keeps cost arithmetic honest.
251const MAX_MEMORY_BYTES: u64 = 1 << 42;
252
253impl InferenceSpec {
254    /// A spec with defaults for everything but its identity.
255    pub fn new(role: &str, runtime: InferFlavor) -> Self {
256        Self {
257            role: role.into(),
258            model: None,
259            runtime,
260            serve: ServeArgs::default(),
261            memory: MemoryBudget::default(),
262            excludes: Vec::new(),
263            latency_class: LatencyClass::Batch,
264            sticky: false,
265            api: default_api(),
266        }
267    }
268
269    /// The loopback address the server is reached at.
270    pub fn bind_ip(&self) -> Result<IpAddr, RuntimeError> {
271        let raw = self.serve.host.as_deref().unwrap_or("127.0.0.1");
272        let ip: IpAddr = if raw == "localhost" {
273            IpAddr::from([127, 0, 0, 1])
274        } else {
275            raw.parse()
276                .map_err(|_| bad(format!("serve.host {raw:?} is not an IP address")))?
277        };
278        if !ip.is_loopback() {
279            return Err(bad(format!(
280                "serve.host {raw} is not loopback: binding wider than loopback needs a governed \
281                 Permit and is not supported by the inference adapters yet"
282            )));
283        }
284        Ok(ip)
285    }
286
287    /// Port in force.
288    pub fn port(&self) -> Result<u16, RuntimeError> {
289        match self.serve.port {
290            Some(0) => Err(bad("serve.port must not be 0")),
291            Some(p) => Ok(p),
292            None => Ok(self.runtime.default_port()),
293        }
294    }
295
296    /// `http://<ip>:<port>` of the server.
297    pub fn base_url(&self) -> Result<String, RuntimeError> {
298        Ok(format!(
299            "http://{}",
300            std::net::SocketAddr::new(self.bind_ip()?, self.port()?)
301        ))
302    }
303
304    /// Boundary validation.
305    pub fn validate(&self) -> Result<(), RuntimeError> {
306        if !valid_token(&self.role, 64) {
307            return Err(bad("role must be a plain token (1-64 chars)"));
308        }
309        for m in self.model.iter().chain(self.serve.draft_model.iter()) {
310            if !valid_token(m, 128) {
311                return Err(bad("model names must be plain tokens (1-128 chars)"));
312            }
313        }
314        self.bind_ip()?;
315        self.port()?;
316        if self.serve.ctx == Some(0) {
317            return Err(bad("serve.ctx must be positive"));
318        }
319        if let Some(q) = &self.serve.kv_quant
320            && !(valid_token(q, 16) && !q.contains(':') && !q.contains('+'))
321        {
322            return Err(bad("serve.kv_quant must be a plain token"));
323        }
324        if self.serve.draft_model.is_some() && self.runtime != InferFlavor::LlamaCpp {
325            return Err(bad("draft models are supported by infer.llamacpp only"));
326        }
327        validate_extra_args(self.runtime, &self.serve.extra_args)?;
328        if self.memory.weights_bytes > MAX_MEMORY_BYTES
329            || self.memory.kv_budget_bytes > MAX_MEMORY_BYTES
330        {
331            return Err(bad("memory figures must be at most 4 TiB"));
332        }
333        for x in &self.excludes {
334            if !valid_token(x, 64) {
335                return Err(bad("excludes entries must be plain tokens"));
336            }
337        }
338        if self.api != "openai-v1" {
339            return Err(bad("only api \"openai-v1\" is supported"));
340        }
341        Ok(())
342    }
343}
344
345impl VerifiedWorkload {
346    /// An inference workload. Trust comes from the operator-attested model
347    /// manifest (when managed) and the operator's own request; nothing
348    /// signed is fetched.
349    pub fn inference(spec: InferenceSpec) -> Result<Self, RuntimeError> {
350        spec.validate()?;
351        Ok(Self {
352            kind: KIND_INFERENCE.to_string(),
353            id: spec.role.clone(),
354            version: spec.model.clone().unwrap_or_else(|| "adopted".into()),
355            source: WorkloadSource::Inference(Box::new(InferencePayload { spec })),
356        })
357    }
358
359    /// The inference spec, or an admission refusal naming the adapter.
360    pub fn inference_spec(&self, runtime: &str) -> Result<&InferenceSpec, RuntimeError> {
361        match &self.source {
362            WorkloadSource::Inference(p) => Ok(&p.spec),
363            _ => Err(RuntimeError::AdmissionRefused(format!(
364                "{runtime} runs only inference workloads"
365            ))),
366        }
367    }
368}
clawft_kernel/workload_runtime/infer/spec.rs

clawft_kernel/workload_runtime/infer/
spec.rs
