superfluid_daemon/
linked.rs1use std::path::{Path, PathBuf};
4use std::sync::Arc;
5
6use superfluid_adapter_kit::{Report, Runtime, WorkerArgs};
7
8use crate::runtime_pick::RuntimeId;
9
10#[allow(unused_mut, clippy::vec_init_then_push)]
11pub fn runtimes() -> Vec<&'static dyn Runtime> {
12 let mut all: Vec<&'static dyn Runtime> = Vec::new();
13 #[cfg(feature = "basert")]
14 all.push(&superfluid_adapter_basert::Basert);
15 #[cfg(feature = "llamacpp")]
16 all.push(&superfluid_adapter_llamacpp::Llamacpp);
17 #[cfg(all(feature = "mlx", target_os = "macos", target_arch = "aarch64"))]
18 all.push(&superfluid_adapter_mlx::Mlx);
19 all
20}
21
22const FEATURES: &[(&str, &str)] = &[("basert", "basert"), ("native", "basert"), ("llamacpp", "llamacpp"), ("mlx", "mlx")];
23
24pub fn ids() -> Vec<RuntimeId> {
25 runtimes().into_iter().map(|r| RuntimeId::new(r.info().id)).collect()
26}
27
28pub fn get(id: RuntimeId) -> Option<&'static dyn Runtime> {
29 runtimes().into_iter().find(|r| r.info().id == id.name())
30}
31
32pub fn engine(name: &str) -> Option<&'static dyn Runtime> {
33 let mock: &'static dyn Runtime = &superfluid_adapter_mock::Mock;
34 runtimes().into_iter().chain([mock]).find(|r| {
35 let info = r.info();
36 info.id == name || info.aliases.contains(&name)
37 })
38}
39
40pub fn feature_for(name: &str) -> Option<&'static str> {
41 FEATURES.iter().find(|(r, _)| *r == name).map(|(_, f)| *f)
42}
43
44pub fn sizes_context(id: RuntimeId) -> bool {
45 #[cfg(feature = "basert")]
46 if id.base().name() == superfluid_adapter_basert::Basert.info().id {
47 return true;
48 }
49 let _ = id;
50 false
51}
52
53pub fn suggest_max_context(
54 id: RuntimeId,
55 models: &[PathBuf],
56 speculators: &[(PathBuf, bool, Option<usize>)],
57 lanes: i32,
58 kv_bits: i32,
59) -> Result<Option<i32>, String> {
60 #[cfg(feature = "basert")]
61 if id.base().name() == superfluid_adapter_basert::Basert.info().id {
62 return Ok(superfluid_adapter_basert::suggest_max_context(models, speculators, lanes, kv_bits));
63 }
64 let _ = (models, speculators, lanes, kv_bits);
65 Err(format!(
66 "--max-context auto: the {} runtime sizes no context window here (the native engine does, from a .base bundle's header): pass --max-context <tokens>",
67 id.base().name()
68 ))
69}
70
71pub fn check(id: RuntimeId, model: Option<&Path>) -> Option<Report> {
74 check_with(id, model, &crate::runtime_options::for_runtime(id))
75}
76
77fn check_with(id: RuntimeId, model: Option<&Path>, options: &[(String, String)]) -> Option<Report> {
78 let runtime = get(id)?;
79 let argv: Vec<String> = options.iter().flat_map(|(f, v)| [f.clone(), v.clone()]).collect();
80 let mut args = WorkerArgs::parse(&argv).unwrap_or_default();
82 args.model = model.map(Path::to_path_buf);
83 Some(superfluid_adapter_kit::check(runtime, &args))
84}
85
86pub fn report(id: RuntimeId, model: Option<&Path>) -> Option<serde_json::Value> {
87 check(id, model).map(|r| r.to_json())
88}
89
90pub fn tokenizer(id: RuntimeId, model: &Path) -> Option<Result<Arc<dyn superfluid_engine::Tokenizer>, String>> {
91 get(id).and_then(|r| r.tokenizer(model))
92}
93
94pub fn spawn(id: RuntimeId, model: PathBuf, ctx: i32, max_batch: u32, kv_bits: i32) -> Option<Result<crate::EngineHost, String>> {
95 let runtime = get(id)?;
96 let args = match crate::runtime_options::worker_args(id, &model, ctx, max_batch, kv_bits) {
97 Ok(args) => args,
98 Err(e) => return Some(Err(e)),
99 };
100 Some(crate::EngineHost::try_spawn_loaded(move || runtime.open(&args)).map_err(|e| e.to_string()))
101}
102
103#[cfg(test)]
104mod tests {
105 use super::*;
106
107 #[test]
108 fn every_linked_runtime_has_its_feature_and_answers_for_itself() {
109 for runtime in runtimes() {
110 let info = runtime.info();
111 let id = RuntimeId::new(info.id);
112 assert!(feature_for(info.id).is_some(), "{}", info.id);
113 assert!(get(id).is_some_and(|r| r.info().id == info.id));
114 assert!(engine(info.id).is_some_and(|r| r.info().id == info.id));
115 for alias in info.aliases {
116 assert!(engine(alias).is_some_and(|r| r.info().id == info.id), "{alias}");
117 }
118 assert_eq!(check(id, None).map(|r| r.info.id), Some(info.id));
119 }
120 assert!(get(RuntimeId::new("vllm")).is_none());
121 assert!(check(RuntimeId::new("vllm"), None).is_none());
122 assert!(engine("vllm").is_none());
123 }
124
125 #[test]
126 fn a_linked_check_reads_the_runtime_options() {
127 let named = |id: &str, options: &[(&str, &str)]| {
128 let options: Vec<(String, String)> = options.iter().map(|(f, v)| (f.to_string(), v.to_string())).collect();
129 check_with(RuntimeId::new(id), None, &options).and_then(|r| r.named)
130 };
131 if get(RuntimeId::new("llamacpp")).is_some() {
132 assert_eq!(named("llamacpp", &[("--lib", "/custom/lib")]).as_deref(), Some("--llamacpp-lib"), "the option names the library");
133 }
134 if get(RuntimeId::new("basert")).is_some() {
135 assert_eq!(named("basert", &[("--basert-lib", "/custom/lib")]).as_deref(), Some("--basert-lib"));
136 }
137 assert!(check_with(RuntimeId::new("vllm"), None, &[]).is_none());
138 }
139
140 #[test]
141 fn the_mock_engine_is_always_there_and_never_a_runtime_here() {
142 assert!(engine("mock").is_some_and(|r| r.info().id == "mock"));
143 assert!(get(RuntimeId::new("mock")).is_none());
144 assert!(!ids().contains(&RuntimeId::new("mock")));
145 }
146}