Skip to main content

superfluid_daemon/
linked.rs

1//! The runtimes this build links, for source builds (`--features basert`, `llamacpp`, `mlx`).
2
3use 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
71/// The report a worker process of `id` would print, its runtime's options counting from the first probe as they do
72/// in a worker's prefix (`Catalog::locate`): the library they name, the cache types a model is sized with.
73pub 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    // Flag and value pairs always parse; what an option's value means is the adapter's to judge.
81    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}