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/{namespace}/{component}/{endpoint}/{slug}/{suffix}/{filename}`
8//! URL on its own `system_status_server`. The route handler reads paths
9//! back out by the same key and streams the bytes to the frontend, which
10//! blake3-verifies them against the MDC.
11//!
12//! The endpoint triple in the key disambiguates multiple `LocalModel`
13//! instances coexisting on one DRT (e.g. different roles attaching the
14//! same base model). `suffix` is the LoRA slug (or `"_base"` for
15//! non-LoRA), scoping each registration so detaching a LoRA doesn't
16//! unregister the base model's files.
17//!
18//! Each entry also stores its `Owner = (instance_id, lora_slug)` so
19//! `unregister_for_owner` can clean up on detach without the caller
20//! threading the model slug. `register` returns
21//! `Err(CollisionError)` on conflict — the caller propagates and the
22//! worker fails to start.
23
24use std::collections::HashMap;
25use std::path::PathBuf;
26use std::sync::Arc;
27
28use parking_lot::RwLock;
29use thiserror::Error;
30
31/// Sentinel `suffix` for non-LoRA registrations. LoRA suffixes are
32/// `Slug::slugify` outputs (`[a-z0-9_-]+`); a name that slugifies to
33/// `_base` would collide with this sentinel and is not supported.
34pub const BASE_SUFFIX: &str = "_base";
35
36/// `(instance_id, lora_slug)`. `None` lora_slug = base model.
37pub type Owner = (u64, Option<String>);
38
39/// `(namespace, component, endpoint, slug, suffix, filename)`.
40type Key = (String, String, String, String, String, String);
41
42#[allow(clippy::too_many_arguments)]
43fn make_key(
44    namespace: &str,
45    component: &str,
46    endpoint: &str,
47    slug: &str,
48    suffix: &str,
49    filename: &str,
50) -> Key {
51    (
52        namespace.to_string(),
53        component.to_string(),
54        endpoint.to_string(),
55        slug.to_string(),
56        suffix.to_string(),
57        filename.to_string(),
58    )
59}
60
61/// Registration collided with a different owner — programmer error.
62#[derive(Debug, Error)]
63#[error("metadata-registry collision on key {key:?}: prior_owner={prior:?}, new_owner={new:?}")]
64pub struct CollisionError {
65    pub key: Key,
66    pub prior: Owner,
67    pub new: Owner,
68}
69
70/// Cloning shares the underlying map.
71#[derive(Clone, Debug, Default)]
72pub struct MetadataArtifactRegistry {
73    entries: Arc<RwLock<HashMap<Key, (PathBuf, Owner)>>>,
74}
75
76impl MetadataArtifactRegistry {
77    pub fn new() -> Self {
78        Self::default()
79    }
80
81    /// Returns `Err(CollisionError)` if a different owner already
82    /// registered this key. Same-owner re-register updates the path.
83    #[allow(clippy::too_many_arguments)]
84    pub fn register(
85        &self,
86        owner: &Owner,
87        namespace: &str,
88        component: &str,
89        endpoint: &str,
90        slug: &str,
91        suffix: &str,
92        filename: &str,
93        path: PathBuf,
94    ) -> Result<(), Box<CollisionError>> {
95        let key = make_key(namespace, component, endpoint, slug, suffix, filename);
96        let mut entries = self.entries.write();
97        if let Some((_, prior)) = entries.get(&key)
98            && prior != owner
99        {
100            return Err(Box::new(CollisionError {
101                key,
102                prior: prior.clone(),
103                new: owner.clone(),
104            }));
105        }
106        entries.insert(key, (path, owner.clone()));
107        tracing::debug!(
108            namespace,
109            component,
110            endpoint,
111            slug,
112            suffix,
113            filename,
114            "registered metadata artifact",
115        );
116        Ok(())
117    }
118
119    #[allow(clippy::too_many_arguments)]
120    pub fn get(
121        &self,
122        namespace: &str,
123        component: &str,
124        endpoint: &str,
125        slug: &str,
126        suffix: &str,
127        filename: &str,
128    ) -> Option<PathBuf> {
129        let key = make_key(namespace, component, endpoint, slug, suffix, filename);
130        self.entries.read().get(&key).map(|(p, _)| p.clone())
131    }
132
133    /// Drop every entry registered by `owner`. No-op if `owner` never
134    /// registered (e.g. self-host was disabled or skipped).
135    pub fn unregister_for_owner(&self, owner: &Owner) {
136        self.entries.write().retain(|_, (_, o)| o != owner);
137    }
138
139    pub fn len(&self) -> usize {
140        self.entries.read().len()
141    }
142
143    pub fn is_empty(&self) -> bool {
144        self.entries.read().is_empty()
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    fn base() -> Owner {
153        (1, None)
154    }
155
156    fn lora(slug: &str) -> Owner {
157        (1, Some(slug.to_string()))
158    }
159
160    #[test]
161    fn register_get_roundtrip() {
162        let reg = MetadataArtifactRegistry::new();
163        let p = PathBuf::from("/tmp/tokenizer.json");
164        reg.register(
165            &base(),
166            "ns",
167            "comp",
168            "ep",
169            "llama-3-8b",
170            "_base",
171            "tokenizer.json",
172            p.clone(),
173        )
174        .unwrap();
175
176        assert_eq!(
177            reg.get("ns", "comp", "ep", "llama-3-8b", "_base", "tokenizer.json"),
178            Some(p)
179        );
180        assert!(
181            reg.get("ns", "comp", "ep", "llama-3-8b", "_base", "missing.json")
182                .is_none()
183        );
184        assert!(
185            reg.get(
186                "ns",
187                "comp",
188                "ep",
189                "llama-3-8b",
190                "lora-v1",
191                "tokenizer.json"
192            )
193            .is_none()
194        );
195    }
196
197    #[test]
198    fn unregister_for_owner_clears_only_that_owner() {
199        let reg = MetadataArtifactRegistry::new();
200        let lora_owner = lora("lora-v1");
201        reg.register(
202            &base(),
203            "ns",
204            "comp",
205            "ep",
206            "m",
207            "_base",
208            "config.json",
209            PathBuf::from("/m/c"),
210        )
211        .unwrap();
212        reg.register(
213            &base(),
214            "ns",
215            "comp",
216            "ep",
217            "m",
218            "_base",
219            "tokenizer.json",
220            PathBuf::from("/m/t"),
221        )
222        .unwrap();
223        reg.register(
224            &lora_owner,
225            "ns",
226            "comp",
227            "ep",
228            "m",
229            "lora-v1",
230            "adapter.json",
231            PathBuf::from("/m/a"),
232        )
233        .unwrap();
234
235        reg.unregister_for_owner(&lora_owner);
236
237        assert!(
238            reg.get("ns", "comp", "ep", "m", "lora-v1", "adapter.json")
239                .is_none()
240        );
241        assert_eq!(
242            reg.get("ns", "comp", "ep", "m", "_base", "config.json"),
243            Some(PathBuf::from("/m/c"))
244        );
245        // Idempotent — second call is a no-op.
246        reg.unregister_for_owner(&lora_owner);
247        assert_eq!(reg.len(), 2);
248    }
249
250    #[test]
251    fn register_returns_err_on_owner_collision() {
252        let reg = MetadataArtifactRegistry::new();
253        let owner_a = (1, None);
254        let owner_b = (2, None);
255        reg.register(
256            &owner_a,
257            "ns",
258            "comp",
259            "ep",
260            "m",
261            "_base",
262            "config.json",
263            PathBuf::from("/a"),
264        )
265        .unwrap();
266        let err = reg
267            .register(
268                &owner_b,
269                "ns",
270                "comp",
271                "ep",
272                "m",
273                "_base",
274                "config.json",
275                PathBuf::from("/b"),
276            )
277            .unwrap_err();
278        assert_eq!(err.prior, owner_a);
279        assert_eq!(err.new, owner_b);
280    }
281
282    #[test]
283    fn register_same_owner_updates_path() {
284        let reg = MetadataArtifactRegistry::new();
285        reg.register(
286            &base(),
287            "ns",
288            "comp",
289            "ep",
290            "m",
291            "_base",
292            "config.json",
293            PathBuf::from("/a"),
294        )
295        .unwrap();
296        reg.register(
297            &base(),
298            "ns",
299            "comp",
300            "ep",
301            "m",
302            "_base",
303            "config.json",
304            PathBuf::from("/b"),
305        )
306        .unwrap();
307        assert_eq!(
308            reg.get("ns", "comp", "ep", "m", "_base", "config.json"),
309            Some(PathBuf::from("/b"))
310        );
311    }
312
313    #[test]
314    fn different_endpoints_coexist() {
315        let reg = MetadataArtifactRegistry::new();
316        let owner_a = (1, None);
317        let owner_b = (2, None);
318        reg.register(
319            &owner_a,
320            "ns",
321            "comp",
322            "ep-a",
323            "m",
324            "_base",
325            "config.json",
326            PathBuf::from("/a"),
327        )
328        .unwrap();
329        // Same (slug, suffix, filename) but different endpoint → no collision.
330        reg.register(
331            &owner_b,
332            "ns",
333            "comp",
334            "ep-b",
335            "m",
336            "_base",
337            "config.json",
338            PathBuf::from("/b"),
339        )
340        .unwrap();
341        assert_eq!(
342            reg.get("ns", "comp", "ep-a", "m", "_base", "config.json"),
343            Some(PathBuf::from("/a"))
344        );
345        assert_eq!(
346            reg.get("ns", "comp", "ep-b", "m", "_base", "config.json"),
347            Some(PathBuf::from("/b"))
348        );
349    }
350}
dynamo_runtime/metadata_registry.rs

dynamo_runtime/
metadata_registry.rs
