dynamo_runtime/
metadata_registry.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Worker-side index from a model metadata file's identity to its
5//! on-disk path. When a worker self-hosts metadata, it registers each
6//! file here and rewrites the MDC's `CheckedFile.path` to a
7//! `/v1/metadata/{slug}/{suffix}/{filename}` URL on its own
8//! `system_status_server`. The route handler reads paths back out by
9//! the same key and streams the bytes to the frontend, which
10//! blake3-verifies them against the MDC.
11//!
12//! `suffix` is the LoRA slug (or `"_base"` for non-LoRA). It scopes
13//! each registration so detaching a LoRA doesn't unregister the base
14//! model's files (or vice versa).
15
16use std::collections::HashMap;
17use std::path::PathBuf;
18use std::sync::Arc;
19
20use parking_lot::RwLock;
21
22/// Sentinel `suffix` for non-LoRA registrations. LoRA suffixes are
23/// `Slug::slugify` outputs (`[a-z0-9_-]+`); a name that slugifies to
24/// `_base` would collide with this sentinel and is not supported.
25pub const BASE_SUFFIX: &str = "_base";
26
27/// `(slug, suffix, filename)`.
28type Key = (String, String, String);
29
30/// Cloning shares the underlying map.
31#[derive(Clone, Debug, Default)]
32pub struct MetadataArtifactRegistry {
33    entries: Arc<RwLock<HashMap<Key, PathBuf>>>,
34}
35
36impl MetadataArtifactRegistry {
37    pub fn new() -> Self {
38        Self {
39            entries: Arc::new(RwLock::new(HashMap::new())),
40        }
41    }
42
43    pub fn register(&self, slug: &str, suffix: &str, filename: &str, path: PathBuf) {
44        let mut entries = self.entries.write();
45        entries.insert(
46            (slug.to_string(), suffix.to_string(), filename.to_string()),
47            path,
48        );
49        tracing::debug!(slug, suffix, filename, "registered metadata artifact");
50    }
51
52    pub fn get(&self, slug: &str, suffix: &str, filename: &str) -> Option<PathBuf> {
53        let entries = self.entries.read();
54        entries
55            .get(&(slug.to_string(), suffix.to_string(), filename.to_string()))
56            .cloned()
57    }
58
59    /// Drop entries for a single registration scoped by `(slug, suffix)`.
60    pub fn unregister(&self, slug: &str, suffix: &str) {
61        let mut entries = self.entries.write();
62        entries.retain(|(s, sx, _), _| !(s == slug && sx == suffix));
63    }
64
65    pub fn len(&self) -> usize {
66        self.entries.read().len()
67    }
68
69    pub fn is_empty(&self) -> bool {
70        self.entries.read().is_empty()
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn register_get_roundtrip() {
80        let reg = MetadataArtifactRegistry::new();
81        let p = PathBuf::from("/tmp/tokenizer.json");
82        reg.register("llama-3-8b", "_base", "tokenizer.json", p.clone());
83
84        assert_eq!(reg.get("llama-3-8b", "_base", "tokenizer.json"), Some(p));
85        assert!(reg.get("llama-3-8b", "_base", "missing.json").is_none());
86        assert!(reg.get("llama-3-8b", "lora-v1", "tokenizer.json").is_none());
87    }
88
89    #[test]
90    fn unregister_only_removes_matching_suffix() {
91        let reg = MetadataArtifactRegistry::new();
92        reg.register("m", "_base", "config.json", PathBuf::from("/m/c"));
93        reg.register("m", "_base", "tokenizer.json", PathBuf::from("/m/t"));
94        reg.register("m", "lora-v1", "config.json", PathBuf::from("/m/c"));
95
96        reg.unregister("m", "_base");
97
98        assert!(reg.get("m", "_base", "config.json").is_none());
99        assert!(reg.get("m", "_base", "tokenizer.json").is_none());
100        // LoRA entry on the same slug survives detach of the base.
101        assert_eq!(
102            reg.get("m", "lora-v1", "config.json"),
103            Some(PathBuf::from("/m/c"))
104        );
105        assert_eq!(reg.len(), 1);
106    }
107}
dynamo_runtime/metadata_registry.rs

dynamo_runtime/
metadata_registry.rs
