clawft_kernel/workload_runtime/infer/
launch.rs1use 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#[derive(Debug, Clone)]
15pub struct Launch {
16 pub program: PathBuf,
18 pub args: Vec<String>,
20 pub dir: PathBuf,
22 pub env: Vec<(String, String)>,
24}
25
26fn refuse(m: impl Into<String>) -> RuntimeError {
27 RuntimeError::AdmissionRefused(m.into())
28}
29
30pub 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
49fn 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
65pub 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}