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    /// Serves `daemon` as `name` beside the models already loaded.
227    pub fn insert_loaded(&self, name: &str, daemon: Arc<Daemon>) {
228        self.models.write().expect("registry").insert(name.to_string(), daemon);
229    }
230
231    pub fn default_name(&self) -> String {
232        self.default_model.read().expect("registry").clone()
233    }
234
235    pub fn names(&self) -> Vec<String> {
236        let mut v: Vec<String> = self.models.read().expect("registry").keys().cloned().collect();
237        v.sort();
238        v
239    }
240
241    /// The model `requested` names under [`crate::openai::ModelNaming::Loose`]: exact, else
242    /// by case, else by substring.
243    pub fn loose_name(&self, requested: &str) -> Option<String> {
244        let mut all = self.names();
245        all.extend(self.known.read().expect("registry").keys().cloned());
246        all.sort();
247        all.dedup();
248        let wanted = requested.to_lowercase();
249        all.iter()
250            .find(|n| *n == requested)
251            .or_else(|| all.iter().find(|n| n.to_lowercase() == wanted))
252            .or_else(|| all.iter().find(|n| n.contains(requested)))
253            .cloned()
254    }
255
256    pub fn resolve(&self, model: Option<&str>) -> Option<(String, Arc<Daemon>)> {
257        self.try_resolve(model).ok()
258    }
259
260    pub fn try_resolve(&self, model: Option<&str>) -> Result<(String, Arc<Daemon>), ResolveError> {
261        let name = {
262            let models = self.models.read().expect("registry");
263            match model {
264                Some(m) if models.contains_key(m) => m.to_string(),
265                Some(m) => {
266                    let source = self.known.read().expect("registry").get(m).cloned();
267                    let source = source.ok_or(ResolveError::NotFound)?;
268                    drop(models);
269                    let runtime = self.known_runtimes.read().expect("registry").get(m).cloned();
270                    self.load_as_on(m, &source.to_string_lossy(), runtime).map_err(ResolveError::Load)?;
271                    m.to_string()
272                }
273                None => self.default_model.read().expect("registry").clone(),
274            }
275        };
276        let daemon = self.models.read().expect("registry").get(&name).map(Arc::clone).ok_or(ResolveError::NotFound)?;
277        let daemon = self.fresh(&name, daemon);
278        if self.track_idle.load(Ordering::Relaxed) {
279            self.last_used.lock().expect("registry").insert(name.clone(), Instant::now());
280        }
281        Ok((name, daemon))
282    }
283
284    pub fn register_known(&self, name: &str, path: &Path) {
285        self.known.write().expect("registry").insert(name.to_string(), path.to_path_buf());
286    }
287
288    pub fn is_known(&self, name: &str) -> bool {
289        self.known.read().expect("registry").contains_key(name)
290    }
291
292    pub fn all_names(&self) -> Vec<String> {
293        let mut names = self.names();
294        for k in self.known.read().expect("registry").keys() {
295            if !names.contains(k) {
296                names.push(k.clone());
297            }
298        }
299        names.sort();
300        names
301    }
302
303    pub fn track_idle(&self) {
304        self.track_idle.store(true, Ordering::Relaxed);
305    }
306
307    pub fn sweep_idle(&self, max_idle: Duration) -> Vec<String> {
308        let default = self.default_name();
309        let now = Instant::now();
310        let candidates: Vec<String> = {
311            let last = self.last_used.lock().expect("registry");
312            self.names()
313                .into_iter()
314                .filter(|n| *n != default && self.is_known(n))
315                .filter(|n| {
316                    last.get(n).is_none_or(|t| now.duration_since(*t) > max_idle)
317                })
318                .collect()
319        };
320        let mut unloaded = Vec::new();
321        for name in candidates {
322            if self.unload(&name).unwrap_or(false) {
323                self.last_used.lock().expect("registry").remove(&name);
324                unloaded.push(name);
325            }
326        }
327        unloaded
328    }
329
330    pub fn is_loaded(&self, name: &str) -> bool {
331        self.models.read().expect("registry").contains_key(name)
332    }
333
334    pub(crate) fn peek_loaded(&self, name: &str) -> Option<Arc<Daemon>> {
335        self.models.read().expect("registry").get(name).cloned()
336    }
337
338    pub fn load(&self, name: &str) -> Result<(), DaemonError> {
339        self.load_on(name, None)
340    }
341
342    pub fn load_on(&self, name: &str, runtime: Option<crate::runtime_pick::RuntimeId>) -> Result<(), DaemonError> {
343        let known = self.known.read().expect("registry").get(name).map(|p| p.to_string_lossy().into_owned());
344        match known {
345            Some(source) => {
346                let runtime = runtime.or_else(|| self.known_runtimes.read().expect("registry").get(name).cloned());
347                self.load_as_on(name, &source, runtime)?;
348            }
349            None => {
350                self.load_as_on(name, name, runtime)?;
351            }
352        }
353        *self.default_model.write().expect("registry") = name.to_string();
354        Ok(())
355    }
356
357    pub fn load_as(&self, name: &str, source: &str) -> Result<bool, DaemonError> {
358        self.load_as_on(name, source, None)
359    }
360
361    pub fn load_as_on(
362        &self,
363        name: &str,
364        source: &str,
365        runtime: Option<crate::runtime_pick::RuntimeId>,
366    ) -> Result<bool, DaemonError> {
367        let _guard = self.load_lock.lock().expect("load lock");
368        if let Some(loaded) = self.models.read().expect("registry").get(name) {
369            let have = self
370                .sources
371                .read()
372                .expect("registry")
373                .get(name)
374                .and_then(|(_, r)| *r)
375                .or_else(|| loaded.runtime_id().map(crate::runtime_pick::RuntimeId::new));
376            if let (Some(want), Some(have)) = (runtime, have) {
377                let same = want.base() == have.base() && (want.install().is_none() || want.install() == have.install());
378                if !same {
379                    return Err(DaemonError::Unsupported(crate::Refusal::new(
380                        "unsupported_runtime",
381                        Some("runtime"),
382                        format!("model '{name}' is loaded on runtime {have}; unload it to load it on {}", want.name()),
383                    )));
384                }
385            }
386            return Ok(false);
387        }
388        let daemon = (self.loader)(source, runtime)?;
389        let ran_on = runtime.filter(|r| r.install().is_some()).or_else(|| daemon.runtime_id().map(crate::runtime_pick::RuntimeId::new));
390        self.models.write().expect("registry").insert(name.to_string(), daemon);
391        self.sources.write().expect("registry").insert(name.to_string(), (source.to_string(), ran_on));
392        Ok(true)
393    }
394
395    pub fn shutdown_all(&self) {
396        let daemons: Vec<_> = self.models.read().expect("registry").values().cloned().collect();
397        for d in &daemons {
398            d.signal_shutdown();
399        }
400        for d in &daemons {
401            d.shutdown();
402        }
403    }
404
405    pub fn unload(&self, name: &str) -> Result<bool, DaemonError> {
406        let mut models = self.models.write().expect("registry");
407        if !models.contains_key(name) {
408            return Ok(false);
409        }
410        if models.len() == 1 {
411            return Err(DaemonError::Config("cannot unload the last remaining model"));
412        }
413        let removed = models.remove(name);
414        // It stays known by its name, so the next request for it, or a load
415        // by that name, brings it back as it was.
416        if let Some((source, runtime)) = self.sources.write().expect("registry").remove(name) {
417            let path = PathBuf::from(&source);
418            if path.exists() {
419                self.known.write().expect("registry").entry(name.to_string()).or_insert(path);
420                if let Some(rt) = runtime {
421                    self.known_runtimes.write().expect("registry").entry(name.to_string()).or_insert(rt);
422                }
423            }
424        }
425        let mut def = self.default_model.write().expect("registry");
426        if *def == name {
427            *def = models.keys().next().cloned().unwrap_or_default();
428        }
429        drop(models);
430        drop(def);
431        drop(removed);
432        Ok(true)
433    }
434}
435
436#[cfg(test)]
437mod distinct_name_tests {
438    use super::*;
439
440    #[test]
441    fn same_stem_artifacts_get_distinct_names() {
442        let mut taken = std::collections::HashSet::new();
443        let base = Path::new("/models/foo.base");
444        let gguf = Path::new("/models/foo.gguf");
445        let a = distinct_name(base, &taken);
446        assert_eq!(a, "foo");
447        taken.insert(a);
448        let b = distinct_name(gguf, &taken);
449        assert_eq!(b, "foo.gguf", "the second keeps its extension");
450        taken.insert(b);
451        assert_eq!(distinct_name(Path::new("/other/foo.gguf"), &taken), "other/foo.gguf");
452    }
453
454    #[test]
455    fn two_variants_of_one_hub_model_are_named_by_variant() {
456        let repo = std::env::temp_dir().join(format!("superfluid-variants-{}", std::process::id())).join("org/Model");
457        let mut taken = std::collections::HashSet::new();
458        for variant in ["default-q4", "default-q8"] {
459            std::fs::create_dir_all(repo.join(variant)).unwrap();
460            let sidecar = format!(r#"{{"id":"org/Model","variant":"{variant}"}}"#);
461            std::fs::write(repo.join(variant).join("hub.json"), sidecar).unwrap();
462        }
463        let q4 = distinct_name(&repo.join("default-q4/model.base"), &taken);
464        assert_eq!(q4, "org/Model");
465        taken.insert(q4);
466        assert_eq!(distinct_name(&repo.join("default-q8/model.base"), &taken), "org/Model:default-q8");
467    }
468}
469
470#[cfg(test)]
471mod state_dir_tests {
472    use super::*;
473
474    fn scratch(tag: &str) -> PathBuf {
475        let d = std::env::temp_dir().join(format!("superfluid-state-dir-{tag}-{}", std::process::id()));
476        let _ = std::fs::remove_dir_all(&d);
477        std::fs::create_dir_all(&d).unwrap();
478        d
479    }
480
481    #[test]
482    fn ids_that_sanitize_alike_keep_their_state_apart() {
483        for order in [["m.gguf", "m_gguf"], ["m_gguf", "m.gguf"]] {
484            let models = scratch("alike");
485            let first = state_dir(&models, order[0]).unwrap();
486            let second = state_dir(&models, order[1]).unwrap();
487            assert_eq!(first, models.join("m_gguf"), "the first to ask has the plain name");
488            assert_ne!(first, second, "{order:?} share a state directory");
489            assert!(second.is_dir() && second.parent() == Some(models.as_path()), "{}", second.display());
490            assert_eq!(state_dir(&models, order[1]).unwrap(), second);
491            assert_eq!(state_dir(&models, order[0]).unwrap(), first);
492            let _ = std::fs::remove_dir_all(&models);
493        }
494        let models = scratch("three");
495        let dirs: std::collections::HashSet<PathBuf> = ["a/b", "a.b", "a_b"].iter().map(|id| state_dir(&models, id).unwrap()).collect();
496        assert_eq!(dirs.len(), 3, "{dirs:?}");
497        let _ = std::fs::remove_dir_all(&models);
498    }
499
500    #[test]
501    fn a_state_directory_from_before_stays_where_it_is() {
502        let models = scratch("before");
503        let old = models.join("Qwen3_6-27B");
504        std::fs::create_dir_all(old.join("park")).unwrap();
505        std::fs::write(old.join("wal.log"), b"records").unwrap();
506        let dir = state_dir(&models, "Qwen3.6-27B").unwrap();
507        assert_eq!(dir, old);
508        assert_eq!(std::fs::read(dir.join("wal.log")).unwrap(), b"records");
509        assert_ne!(state_dir(&models, "Qwen3_6-27B").unwrap(), old, "another id does not get it afterwards");
510        assert_eq!(state_dir(&models, "Qwen3.6-27B").unwrap(), old);
511        let _ = std::fs::remove_dir_all(&models);
512    }
513
514    #[test]
515    fn an_id_names_no_path_outside_the_models_directory() {
516        let models = scratch("inside");
517        for id in ["../x", "..", "a/../../b", "/etc/passwd", "."] {
518            let dir = state_dir(&models, id).unwrap();
519            assert_eq!(dir.parent(), Some(models.as_path()), "{id}: {}", dir.display());
520            assert!(dir.is_dir());
521        }
522        let _ = std::fs::remove_dir_all(&models);
523    }
524}
superfluid_daemon/registry.rs

superfluid_daemon/
registry.rs
