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
71pub fn check(id: RuntimeId, model: Option<&Path>) -> Option<Report> {
72    let args = WorkerArgs { model: model.map(Path::to_path_buf), ..Default::default() };
73    get(id).map(|r| superfluid_adapter_kit::check(r, &args))
74}
75
76pub fn report(id: RuntimeId, model: Option<&Path>) -> Option<serde_json::Value> {
77    check(id, model).map(|r| r.to_json())
78}
79
80pub fn tokenizer(id: RuntimeId, model: &Path) -> Option<Result<Arc<dyn superfluid_engine::Tokenizer>, String>> {
81    get(id).and_then(|r| r.tokenizer(model))
82}
83
84pub fn spawn(id: RuntimeId, model: PathBuf, ctx: i32, max_batch: u32, kv_bits: i32) -> Option<Result<crate::EngineHost, String>> {
85    let runtime = get(id)?;
86    let args = WorkerArgs { model: Some(model), max_context: ctx, max_batch, kv_bits, ..Default::default() };
87    Some(crate::EngineHost::try_spawn_loaded(move || runtime.open(&args)).map_err(|e| e.to_string()))
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn every_linked_runtime_has_its_feature_and_answers_for_itself() {
96        for runtime in runtimes() {
97            let info = runtime.info();
98            let id = RuntimeId::new(info.id);
99            assert!(feature_for(info.id).is_some(), "{}", info.id);
100            assert!(get(id).is_some_and(|r| r.info().id == info.id));
101            assert!(engine(info.id).is_some_and(|r| r.info().id == info.id));
102            for alias in info.aliases {
103                assert!(engine(alias).is_some_and(|r| r.info().id == info.id), "{alias}");
104            }
105            assert_eq!(check(id, None).map(|r| r.info.id), Some(info.id));
106        }
107        assert!(get(RuntimeId::new("vllm")).is_none());
108        assert!(check(RuntimeId::new("vllm"), None).is_none());
109        assert!(engine("vllm").is_none());
110    }
111
112    #[test]
113    fn the_mock_engine_is_always_there_and_never_a_runtime_here() {
114        assert!(engine("mock").is_some_and(|r| r.info().id == "mock"));
115        assert!(get(RuntimeId::new("mock")).is_none());
116        assert!(!ids().contains(&RuntimeId::new("mock")));
117    }
118}