Skip to main content

clawft_kernel/workload_runtime/infer/
launch.rs

1//! Managed-mode launch plan: argv from the spec and the resolved model,
2//! environment, working directory, and the "is something already on this
3//! port" guard that keeps the adapter from ever starting over a server it
4//! does not own.
5
6use std::path::PathBuf;
7
8use super::config::ManagedConfig;
9use super::spec::{InferFlavor, InferenceSpec};
10use crate::model_manifest::{ModelFormat, ResolvedModel};
11use crate::workload_runtime::types::RuntimeError;
12
13/// A launch ready to spawn.
14#[derive(Debug, Clone)]
15pub struct Launch {
16    /// Launcher script.
17    pub program: PathBuf,
18    /// Arguments (argv, never a shell string).
19    pub args: Vec<String>,
20    /// Working directory (created at load).
21    pub dir: PathBuf,
22    /// Full environment.
23    pub env: Vec<(String, String)>,
24}
25
26fn refuse(m: impl Into<String>) -> RuntimeError {
27    RuntimeError::AdmissionRefused(m.into())
28}
29
30/// The model must be in a format the runtime loads.
31pub fn check_format(flavor: InferFlavor, m: &ResolvedModel) -> Result<(), RuntimeError> {
32    let ok = match flavor {
33        InferFlavor::LlamaCpp => m.format == ModelFormat::Gguf,
34        InferFlavor::MlxLm => matches!(m.format, ModelFormat::Mlx | ModelFormat::Safetensors),
35        InferFlavor::Ollama => m.format == ModelFormat::Ollama,
36    };
37    if ok {
38        Ok(())
39    } else {
40        Err(refuse(format!(
41            "{} cannot load a {} model ({})",
42            flavor.id(),
43            m.format.as_str(),
44            m.name
45        )))
46    }
47}
48
49/// Where the server reads weights from: the first GGUF shard (llama.cpp
50/// finds the rest) or the model directory (MLX).
51fn weights_arg(flavor: InferFlavor, m: &ResolvedModel) -> Result<String, RuntimeError> {
52    let p = match flavor {
53        InferFlavor::LlamaCpp => m
54            .shards
55            .first()
56            .cloned()
57            .ok_or_else(|| refuse("model has no shards"))?,
58        _ => m.root.clone(),
59    };
60    p.to_str()
61        .map(str::to_string)
62        .ok_or_else(|| refuse("model path is not UTF-8"))
63}
64
65/// Build the launch for a llama.cpp or mlx-lm instance. The argument shape
66/// follows the model lab's launchers (`serve-llamacpp <model.gguf>
67/// [--draft m] [--kv q] [--ctx n] [--port p] [--host h]`, `serve <model>
68/// --port p --host h`); it lives in this one function.
69pub fn build_launch(
70    flavor: InferFlavor,
71    spec: &InferenceSpec,
72    model: &ResolvedModel,
73    draft: Option<&ResolvedModel>,
74    cfg: &ManagedConfig,
75    instance_id: &str,
76) -> Result<Launch, RuntimeError> {
77    check_format(flavor, model)?;
78    let program = cfg.serve_program.clone().ok_or_else(|| {
79        refuse(format!(
80            "{} managed mode needs a serve program",
81            flavor.id()
82        ))
83    })?;
84    let mut args = vec![weights_arg(flavor, model)?];
85    args.extend(["--port".into(), spec.port()?.to_string()]);
86    args.extend(["--host".into(), spec.bind_ip()?.to_string()]);
87    if flavor == InferFlavor::LlamaCpp {
88        if let Some(c) = spec.serve.ctx {
89            args.extend(["--ctx".into(), c.to_string()]);
90        }
91        if let Some(q) = &spec.serve.kv_quant {
92            args.extend(["--kv".into(), q.clone()]);
93        }
94        if let Some(d) = draft {
95            check_format(flavor, d)?;
96            args.extend(["--draft".into(), weights_arg(flavor, d)?]);
97        }
98    }
99    args.extend(spec.serve.extra_args.iter().cloned());
100    let mut env = cfg.env.clone();
101    for name in &cfg.env_passthrough {
102        if let Ok(v) = std::env::var(name) {
103            env.push((name.clone(), v));
104        }
105    }
106    Ok(Launch {
107        program,
108        args,
109        dir: cfg.data_root.join(instance_id),
110        env,
111    })
112}