superfluid_daemon/
registry.rs

1//! Multi-model host management.
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{Arc, Mutex, RwLock};
7use std::time::{Duration, Instant};
8
9use crate::{Daemon, DaemonError};
10
11/// A field of the hub sidecar (`hub.json` beside the artifact), trimmed, if
12/// it is there and not empty.
13fn hub_field(path: &Path, field: &str) -> Option<String> {
14    let bytes = std::fs::read(path.parent()?.join("hub.json")).ok()?;
15    let v = serde_json::from_slice::<serde_json::Value>(&bytes).ok()?;
16    let value = v.get(field)?.as_str()?.trim();
17    (!value.is_empty()).then(|| value.to_string())
18}
19
20pub fn model_id_for(path: &Path) -> String {
21    if let Some(id) = hub_field(path, "id") {
22        return id;
23    }
24    let stem = if path.is_dir() { path.file_name() } else { path.file_stem() };
25    stem.map(|s| s.to_string_lossy().into_owned())
26        .filter(|s| !s.is_empty())
27        .unwrap_or_else(|| "superfluid".into())
28}
29
30pub fn distinct_name(path: &Path, taken: &std::collections::HashSet<String>) -> String {
31    let id = model_id_for(path);
32    if !taken.contains(&id) {
33        return id;
34    }
35    // Two variants of one hub model (its q4 and its q8) are told apart by
36    // the variant, as baseRT names them.
37    if let Some(variant) = hub_field(path, "variant") {
38        let qualified = format!("{id}:{variant}");
39        if !taken.contains(&qualified) {
40            return qualified;
41        }
42    }
43    let full = path
44        .file_name()
45        .map(|n| n.to_string_lossy().into_owned())
46        .filter(|n| !n.is_empty() && n != &id)
47        .unwrap_or_else(|| id.clone());
48    if !taken.contains(&full) {
49        return full;
50    }
51    let parent = path
52        .parent()
53        .and_then(|p| p.file_name())
54        .map(|n| n.to_string_lossy().into_owned())
55        .unwrap_or_default();
56    format!("{parent}/{full}")
57}
58
59const STATE_OWNER: &str = "model-id";
60
61pub fn state_dir(models: &Path, id: &str) -> std::io::Result<PathBuf> {
62    use sha2::{Digest, Sha256};
63    let mut name: String = id.chars().map(|c| if c.is_ascii_alphanumeric() || c == '-' || c == '_' { c } else { '_' }).collect();
64    if name.is_empty() {
65        name.push('_');
66    }
67    static CLAIMS: Mutex<()> = Mutex::new(());
68    let _one_at_a_time = CLAIMS.lock().unwrap_or_else(|e| e.into_inner());
69    let plain = models.join(&name);
70    if claim(&plain, id)? {
71        return Ok(plain);
72    }
73    let digest: String = Sha256::digest(id.as_bytes()).iter().take(8).map(|b| format!("{b:02x}")).collect();
74    let own = models.join(format!("{name}-{digest}"));
75    if claim(&own, id)? {
76        return Ok(own);
77    }
78    Err(std::io::Error::new(
79        std::io::ErrorKind::AlreadyExists,
80        format!("no state directory for model '{id}': {} and {} are other models'", plain.display(), own.display()),
81    ))
82}
83
84fn claim(dir: &Path, id: &str) -> std::io::Result<bool> {
85    use std::io::Write;
86    std::fs::create_dir_all(dir)?;
87    let owner = dir.join(STATE_OWNER);
88    match std::fs::read(&owner) {
89        Ok(theirs) => return Ok(theirs == id.as_bytes()),
90        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
91        Err(e) => return Err(e),
92    }
93    let aside = dir.join(format!(".{STATE_OWNER}-new"));
94    let mut file = std::fs::File::create(&aside)?;
95    file.write_all(id.as_bytes())?;
96    file.sync_all()?;
97    drop(file);
98    std::fs::rename(&aside, &owner)?;
99    if let Ok(d) = std::fs::File::open(dir) {
100        let _ = d.sync_all();
101    }
102    Ok(true)
103}
104
105pub type ModelLoader =
106    Box<dyn Fn(&str, Option<crate::runtime_pick::RuntimeId>) -> Result<Arc<Daemon>, DaemonError> + Send + Sync>;
107
108pub type ModelReloader =
109    Box<dyn Fn(&str, Option<crate::runtime_pick::RuntimeId>, &Daemon) -> Result<(), DaemonError> + Send + Sync>;
110
111pub type RuntimeNames = Box<dyn Fn(&str) -> Result<crate::runtime_pick::RuntimeId, String> + Send + Sync>;
112
113#[derive(Debug)]
114pub enum ResolveError {
115    NotFound,
116    Load(DaemonError),
117}
118
119pub struct ModelRegistry {
120    models: RwLock<HashMap<String, Arc<Daemon>>>,
121    loader: ModelLoader,
122    default_model: RwLock<String>,
123    load_lock: Mutex<()>,
124    known: RwLock<HashMap<String, PathBuf>>,
125    known_runtimes: RwLock<HashMap<String, crate::runtime_pick::RuntimeId>>,
126    last_used: Mutex<HashMap<String, Instant>>,
127    track_idle: AtomicBool,
128    runtime_names: RwLock<RuntimeNames>,
129    sources: RwLock<HashMap<String, (String, Option<crate::runtime_pick::RuntimeId>)>>,
130    reload_failed: Mutex<HashMap<String, Instant>>,
131    reloader: RwLock<Option<ModelReloader>>,
132}
133
134const RELOAD_RETRY: Duration = Duration::from_secs(10);
135
136impl ModelRegistry {
137    pub fn with_initial(name: impl Into<String>, daemon: Arc<Daemon>, loader: ModelLoader) -> ModelRegistry {
138        let name = name.into();
139        let mut m = HashMap::new();
140        m.insert(name.clone(), daemon);
141        ModelRegistry {
142            models: RwLock::new(m),
143            loader,
144            default_model: RwLock::new(name),
145            load_lock: Mutex::new(()),
146            known: RwLock::new(HashMap::new()),
147            known_runtimes: RwLock::new(HashMap::new()),
148            last_used: Mutex::new(HashMap::new()),
149            track_idle: AtomicBool::new(false),
150            runtime_names: RwLock::new(Box::new(|name: &str| {
151                crate::runtime_pick::RuntimeId::parse(name).ok_or_else(|| format!("'{name}' is not a runtime id"))
152            })),
153            sources: RwLock::new(HashMap::new()),
154            reload_failed: Mutex::new(HashMap::new()),
155            reloader: RwLock::new(None),
156        }
157    }
158
159    pub fn set_reloader(&self, reloader: ModelReloader) {
160        *self.reloader.write().expect("registry") = Some(reloader);
161    }
162
163    pub fn remember_source(&self, name: &str, source: &str) {
164        self.remember_source_on(name, source, None);
165    }
166
167    pub fn remember_source_on(&self, name: &str, source: &str, runtime: Option<crate::runtime_pick::RuntimeId>) {
168        let runtime = runtime
169            .filter(|r| r.install().is_some())
170            .or_else(|| self.models.read().expect("registry").get(name).and_then(|d| d.runtime_id()).map(crate::runtime_pick::RuntimeId::new));
171        self.sources.write().expect("registry").insert(name.to_string(), (source.to_string(), runtime));
172    }
173
174    fn fresh(&self, name: &str, daemon: Arc<Daemon>) -> Arc<Daemon> {
175        let Some(why) = daemon.needs_reload() else { return daemon };
176        let Some((source, runtime)) = self.sources.read().expect("registry").get(name).cloned() else {
177            return daemon;
178        };
179        if self.reload_failed.lock().expect("registry").get(name).is_some_and(|t| t.elapsed() < RELOAD_RETRY) {
180            return daemon;
181        }
182        let _guard = self.load_lock.lock().expect("load lock");
183        let still_served = self.models.read().expect("registry").get(name).is_some_and(|cur| Arc::ptr_eq(cur, &daemon));
184        if !still_served || daemon.needs_reload().is_none() {
185            return daemon;
186        }
187        let reloader = self.reloader.read().expect("registry");
188        let Some(reload) = reloader.as_ref() else { return daemon };
189        tracing::warn!(model = name, why = %why, "loading the model again");
190        match reload(&source, runtime, &daemon) {
191            Ok(()) => {
192                self.reload_failed.lock().expect("registry").remove(name);
193            }
194            Err(e) => {
195                tracing::error!(model = name, error = %e, "the model did not load again");
196                self.reload_failed.lock().expect("registry").insert(name.to_string(), Instant::now());
197            }
198        }
199        daemon
200    }
201
202    pub fn reload_pending(&self, model: Option<&str>) -> bool {
203        let name = match model {
204            Some(m) => m.to_string(),
205            None => self.default_model.read().expect("registry").clone(),
206        };
207        self.models.read().expect("registry").get(&name).is_some_and(|d| d.needs_reload().is_some())
208    }
209
210    pub fn set_runtime_names(&self, names: RuntimeNames) {
211        *self.runtime_names.write().expect("registry") = names;
212    }
213
214    pub fn runtime_named(&self, name: &str) -> Result<crate::runtime_pick::RuntimeId, String> {
215        (self.runtime_names.read().expect("registry"))(name)
216    }
217
218    pub fn single(name: impl Into<String>, daemon: Arc<Daemon>) -> ModelRegistry {
219        Self::with_initial(
220            name,
221            daemon,
222            Box::new(|_, _| Err(DaemonError::Config("dynamic model loading is not enabled on this server"))),
223        )
224    }
225
226    pub fn default_name(&self) -> String {
227        self.default_model.read().expect("registry").clone()
228    }
229
230    pub fn names(&self) -> Vec<String> {
231        let mut v: Vec<String> = self.models.read().expect("registry").keys().cloned().collect();
232        v.sort();
233        v
234    }
235
236    /// The model `requested` names under [`crate::openai::ModelNaming::Loose`]: exact, else
237    /// by case, else by substring.
238    pub fn loose_name(&self, requested: &str) -> Option<String> {
239        let mut all = self.names();
240        all.extend(self.known.read().expect("registry").keys().cloned());
241        all.sort();
242        all.dedup();
243        let wanted = requested.to_lowercase();
244        all.iter()
245            .find(|n| *n == requested)
246            .or_else(|| all.iter().find(|n| n.to_lowercase() == wanted))
247            .or_else(|| all.iter().find(|n| n.contains(requested)))
248            .cloned()
249    }
250
251    pub fn resolve(&self, model: Option<&str>) -> Option<(String, Arc<Daemon>)> {
252        self.try_resolve(model).ok()
253    }
254
255    pub fn try_resolve(&self, model: Option<&str>) -> Result<(String, Arc<Daemon>), ResolveError> {
256        let name = {
257            let models = self.models.read().expect("registry");
258            match model {
259                Some(m) if models.contains_key(m) => m.to_string(),
260                Some(m) => {
261                    let source = self.known.read().expect("registry").get(m).cloned();
262                    let source = source.ok_or(ResolveError::NotFound)?;
263                    drop(models);
264                    let runtime = self.known_runtimes.read().expect("registry").get(m).cloned();
265                    self.load_as_on(m, &source.to_string_lossy(), runtime).map_err(ResolveError::Load)?;
266                    m.to_string()
267                }
268                None => self.default_model.read().expect("registry").clone(),
269            }
270        };
271        let daemon = self.models.read().expect("registry").get(&name).map(Arc::clone).ok_or(ResolveError::NotFound)?;
272        let daemon = self.fresh(&name, daemon);
273        if self.track_idle.load(Ordering::Relaxed) {
274            self.last_used.lock().expect("registry").insert(name.clone(), Instant::now());
275        }
276        Ok((name, daemon))
277    }
278
279    pub fn register_known(&self, name: &str, path: &Path) {
280        self.known.write().expect("registry").insert(name.to_string(), path.to_path_buf());
281    }
282
283    pub fn is_known(&self, name: &str) -> bool {
284        self.known.read().expect("registry").contains_key(name)
285    }
286
287    pub fn all_names(&self) -> Vec<String> {
288        let mut names = self.names();
289        for k in self.known.read().expect("registry").keys() {
290            if !names.contains(k) {
291                names.push(k.clone());
292            }
293        }
294        names.sort();
295        names
296    }
297
298    pub fn track_idle(&self) {
299        self.track_idle.store(true, Ordering::Relaxed);
300    }
301
302    pub fn sweep_idle(&self, max_idle: Duration) -> Vec<String> {
303        let default = self.default_name();
304        let now = Instant::now();
305        let candidates: Vec<String> = {
306            let last = self.last_used.lock().expect("registry");
307            self.names()
308                .into_iter()
309                .filter(|n| *n != default && self.is_known(n))
310                .filter(|n| {
311                    last.get(n).is_none_or(|t| now.duration_since(*t) > max_idle)
312                })
313                .collect()
314        };
315        let mut unloaded = Vec::new();
316        for name in candidates {
317            if self.unload(&name).unwrap_or(false) {
318                self.last_used.lock().expect("registry").remove(&name);
319                unloaded.push(name);
320            }
321        }
322        unloaded
323    }
324
325    pub fn is_loaded(&self, name: &str) -> bool {
326        self.models.read().expect("registry").contains_key(name)
327    }
328
329    pub(crate) fn peek_loaded(&self, name: &str) -> Option<Arc<Daemon>> {
330        self.models.read().expect("registry").get(name).cloned()
331    }
332
333    pub fn load(&self, name: &str) -> Result<(), DaemonError> {
334        self.load_on(name, None)
335    }
336
337    pub fn load_on(&self, name: &str, runtime: Option<crate::runtime_pick::RuntimeId>) -> Result<(), DaemonError> {
338        let known = self.known.read().expect("registry").get(name).map(|p| p.to_string_lossy().into_owned());
339        match known {
340            Some(source) => {
341                let runtime = runtime.or_else(|| self.known_runtimes.read().expect("registry").get(name).cloned());
342                self.load_as_on(name, &source, runtime)?;
343            }
344            None => {
345                self.load_as_on(name, name, runtime)?;
346            }
347        }
348        *self.default_model.write().expect("registry") = name.to_string();
349        Ok(())
350    }
351
352    pub fn load_as(&self, name: &str, source: &str) -> Result<bool, DaemonError> {
353        self.load_as_on(name, source, None)
354    }
355
356    pub fn load_as_on(
357        &self,
358        name: &str,
359        source: &str,
360        runtime: Option<crate::runtime_pick::RuntimeId>,
361    ) -> Result<bool, DaemonError> {
362        let _guard = self.load_lock.lock().expect("load lock");
363        if let Some(loaded) = self.models.read().expect("registry").get(name) {
364            let have = self
365                .sources
366                .read()
367                .expect("registry")
368                .get(name)
369                .and_then(|(_, r)| *r)
370                .or_else(|| loaded.runtime_id().map(crate::runtime_pick::RuntimeId::new));
371            if let (Some(want), Some(have)) = (runtime, have) {
372                let same = want.base() == have.base() && (want.install().is_none() || want.install() == have.install());
373                if !same {
374                    return Err(DaemonError::Unsupported(crate::Refusal::new(
375                        "unsupported_runtime",
376                        Some("runtime"),
377                        format!("model '{name}' is loaded on runtime {have}; unload it to load it on {}", want.name()),
378                    )));
379                }
380            }
381            return Ok(false);
382        }
383        let daemon = (self.loader)(source, runtime)?;
384        let ran_on = runtime.filter(|r| r.install().is_some()).or_else(|| daemon.runtime_id().map(crate::runtime_pick::RuntimeId::new));
385        self.models.write().expect("registry").insert(name.to_string(), daemon);
386        self.sources.write().expect("registry").insert(name.to_string(), (source.to_string(), ran_on));
387        Ok(true)
388    }
389
390    pub fn shutdown_all(&self) {
391        let daemons: Vec<_> = self.models.read().expect("registry").values().cloned().collect();
392        for d in &daemons {
393            d.signal_shutdown();
394        }
395        for d in &daemons {
396            d.shutdown();
397        }
398    }
399
400    pub fn unload(&self, name: &str) -> Result<bool, DaemonError> {
401        let mut models = self.models.write().expect("registry");
402        if !models.contains_key(name) {
403            return Ok(false);
404        }
405        if models.len() == 1 {
406            return Err(DaemonError::Config("cannot unload the last remaining model"));
407        }
408        let removed = models.remove(name);
409        // It stays known by its name, so the next request for it, or a load
410        // by that name, brings it back as it was.
411        if let Some((source, runtime)) = self.sources.write().expect("registry").remove(name) {
412            let path = PathBuf::from(&source);
413            if path.exists() {
414                self.known.write().expect("registry").entry(name.to_string()).or_insert(path);
415                if let Some(rt) = runtime {
416                    self.known_runtimes.write().expect("registry").entry(name.to_string()).or_insert(rt);
417                }
418            }
419        }
420        let mut def = self.default_model.write().expect("registry");
421        if *def == name {
422            *def = models.keys().next().cloned().unwrap_or_default();
423        }
424        drop(models);
425        drop(def);
426        drop(removed);
427        Ok(true)
428    }
429}
430
431#[cfg(test)]
432mod distinct_name_tests {
433    use super::*;
434
435    #[test]
436    fn same_stem_artifacts_get_distinct_names() {
437        let mut taken = std::collections::HashSet::new();
438        let base = Path::new("/models/foo.base");
439        let gguf = Path::new("/models/foo.gguf");
440        let a = distinct_name(base, &taken);
441        assert_eq!(a, "foo");
442        taken.insert(a);
443        let b = distinct_name(gguf, &taken);
444        assert_eq!(b, "foo.gguf", "the second keeps its extension");
445        taken.insert(b);
446        assert_eq!(distinct_name(Path::new("/other/foo.gguf"), &taken), "other/foo.gguf");
447    }
448
449    #[test]
450    fn two_variants_of_one_hub_model_are_named_by_variant() {
451        let repo = std::env::temp_dir().join(format!("superfluid-variants-{}", std::process::id())).join("org/Model");
452        let mut taken = std::collections::HashSet::new();
453        for variant in ["default-q4", "default-q8"] {
454            std::fs::create_dir_all(repo.join(variant)).unwrap();
455            let sidecar = format!(r#"{{"id":"org/Model","variant":"{variant}"}}"#);
456            std::fs::write(repo.join(variant).join("hub.json"), sidecar).unwrap();
457        }
458        let q4 = distinct_name(&repo.join("default-q4/model.base"), &taken);
459        assert_eq!(q4, "org/Model");
460        taken.insert(q4);
461        assert_eq!(distinct_name(&repo.join("default-q8/model.base"), &taken), "org/Model:default-q8");
462    }
463}
464
465#[cfg(test)]
466mod state_dir_tests {
467    use super::*;
468
469    fn scratch(tag: &str) -> PathBuf {
470        let d = std::env::temp_dir().join(format!("superfluid-state-dir-{tag}-{}", std::process::id()));
471        let _ = std::fs::remove_dir_all(&d);
472        std::fs::create_dir_all(&d).unwrap();
473        d
474    }
475
476    #[test]
477    fn ids_that_sanitize_alike_keep_their_state_apart() {
478        for order in [["m.gguf", "m_gguf"], ["m_gguf", "m.gguf"]] {
479            let models = scratch("alike");
480            let first = state_dir(&models, order[0]).unwrap();
481            let second = state_dir(&models, order[1]).unwrap();
482            assert_eq!(first, models.join("m_gguf"), "the first to ask has the plain name");
483            assert_ne!(first, second, "{order:?} share a state directory");
484            assert!(second.is_dir() && second.parent() == Some(models.as_path()), "{}", second.display());
485            assert_eq!(state_dir(&models, order[1]).unwrap(), second);
486            assert_eq!(state_dir(&models, order[0]).unwrap(), first);
487            let _ = std::fs::remove_dir_all(&models);
488        }
489        let models = scratch("three");
490        let dirs: std::collections::HashSet<PathBuf> = ["a/b", "a.b", "a_b"].iter().map(|id| state_dir(&models, id).unwrap()).collect();
491        assert_eq!(dirs.len(), 3, "{dirs:?}");
492        let _ = std::fs::remove_dir_all(&models);
493    }
494
495    #[test]
496    fn a_state_directory_from_before_stays_where_it_is() {
497        let models = scratch("before");
498        let old = models.join("Qwen3_6-27B");
499        std::fs::create_dir_all(old.join("park")).unwrap();
500        std::fs::write(old.join("wal.log"), b"records").unwrap();
501        let dir = state_dir(&models, "Qwen3.6-27B").unwrap();
502        assert_eq!(dir, old);
503        assert_eq!(std::fs::read(dir.join("wal.log")).unwrap(), b"records");
504        assert_ne!(state_dir(&models, "Qwen3_6-27B").unwrap(), old, "another id does not get it afterwards");
505        assert_eq!(state_dir(&models, "Qwen3.6-27B").unwrap(), old);
506        let _ = std::fs::remove_dir_all(&models);
507    }
508
509    #[test]
510    fn an_id_names_no_path_outside_the_models_directory() {
511        let models = scratch("inside");
512        for id in ["../x", "..", "a/../../b", "/etc/passwd", "."] {
513            let dir = state_dir(&models, id).unwrap();
514            assert_eq!(dir.parent(), Some(models.as_path()), "{id}: {}", dir.display());
515            assert!(dir.is_dir());
516        }
517        let _ = std::fs::remove_dir_all(&models);
518    }
519}
superfluid_daemon/registry.rs

superfluid_daemon/
registry.rs
