rusty_stack/core/
registry.rs

1//! Installed-component registry — the single source of truth for what Rusty
2//! has installed, where, and whether a *core* component is sealed against
3//! override.
4//!
5//! Persisted at `~/.mlstack/installed.json`. This module is deliberately
6//! self-contained (load / save / mutate) so it can be threaded through the
7//! installer dispatch (see `installer::run_native_installer`) and consulted by:
8//! - dep-sourcing / no-override (refuse to reinstall a sealed core component)
9//! - verification (is this component already known-good?)
10//! - uninstall (enumerate pip packages + locations to remove)
11//!
12//! # Design notes
13//!
14//! - IDs are free-form `String`s matching the existing
15//!   `component_status::is_component_installed_by_id` identifiers (`rocm`,
16//!   `pytorch`, `triton`, `aiter`, …).
17//! - `sealed` marks a **core** component (rocm / pytorch / triton) that later
18//!   installers must reuse, never override. See Stage 3 of the v0.3.0
19//!   remediation (tenet: single-source deps / no-override).
20//! - Timestamps are caller-supplied ISO-8601 strings (chrono `Utc::now`) so the
21//!   registry stays pure with respect to I/O.
22
23use std::collections::HashSet;
24use std::fs;
25use std::path::PathBuf;
26
27use serde::{Deserialize, Serialize};
28
29use crate::platform::environment::mlstack_root;
30
31/// A single installed component record.
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct InstalledComponent {
34    /// Component identifier (e.g. `rocm`, `pytorch`, `triton`, `aiter`).
35    pub id: String,
36    /// Installed version, if known.
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub version: Option<String>,
39    /// Source index / wheel index URL used (for pip components), if any.
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub source_index: Option<String>,
42    /// Install location (venv prefix, `/opt/rocm`, clone dir, …), if recorded.
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub location: Option<String>,
45    /// `true` for sealed **core** components that must never be overridden.
46    #[serde(default)]
47    pub sealed: bool,
48    /// ISO-8601 install timestamp (caller-supplied).
49    #[serde(default, skip_serializing_if = "Option::is_none")]
50    pub installed_at: Option<String>,
51    /// Free-form pip package names this component owns (for uninstall).
52    #[serde(default, skip_serializing_if = "Vec::is_empty")]
53    pub pip_packages: Vec<String>,
54}
55
56/// The persisted registry of installed components.
57#[derive(Debug, Clone, Default, Serialize, Deserialize)]
58pub struct InstalledComponentRegistry {
59    #[serde(default)]
60    components: Vec<InstalledComponent>,
61}
62
63impl InstalledComponentRegistry {
64    /// Path to the persisted registry file: `~/.mlstack/installed.json`.
65    pub fn path() -> PathBuf {
66        mlstack_root().join("installed.json")
67    }
68
69    /// Load the registry from disk, returning an empty one if absent or
70    /// unreadable (never panics — a corrupt registry degrades to empty).
71    pub fn load() -> Self {
72        let path = Self::path();
73        match fs::read_to_string(&path) {
74            Ok(contents) => serde_json::from_str(&contents).unwrap_or_default(),
75            Err(_) => Self::default(),
76        }
77    }
78
79    /// Persist the registry to disk atomically, creating `~/.mlstack/` first.
80    ///
81    /// Writes to `installed.json.tmp` in the same directory, then `fs::rename`s
82    /// it into place. `rename` is atomic on the same filesystem, so an
83    /// interrupted write (crash, disk full, power loss) can never leave a
84    /// corrupt `installed.json` that would make `load()` return empty and risk
85    /// clobbering a sealed core component (Tenet 2).
86    pub fn save(&self) -> anyhow::Result<()> {
87        let path = Self::path();
88        if let Some(parent) = path.parent() {
89            fs::create_dir_all(parent)?;
90        }
91        let json = serde_json::to_string_pretty(self)?;
92        // Temp file in the SAME directory so the rename is atomic (same fs).
93        let temp_path = path.with_extension("json.tmp");
94        fs::write(&temp_path, json)?;
95        fs::rename(&temp_path, &path)?;
96        Ok(())
97    }
98
99    /// Load then save — normalizes the file on disk (used after manual edits).
100    pub fn reload_and_save() -> anyhow::Result<()> {
101        Self::load().save()
102    }
103
104    /// Record (or update) a component. If `sealed` is requested on a component
105    /// that is already sealed, the seal is preserved (never accidentally
106    /// unsealed via re-mark without explicit [`Self::unseal`]).
107    pub fn mark_installed(&mut self, component: InstalledComponent) {
108        if let Some(existing) = self.components.iter_mut().find(|c| c.id == component.id) {
109            let preserve_seal = existing.sealed && !component.sealed;
110            *existing = component;
111            if preserve_seal {
112                existing.sealed = true;
113            }
114        } else {
115            self.components.push(component);
116        }
117    }
118
119    /// Convenience: mark a core component as installed + sealed.
120    pub fn seal_core(
121        &mut self,
122        id: &str,
123        version: Option<String>,
124        source_index: Option<String>,
125        location: Option<String>,
126        pip_packages: Vec<String>,
127        installed_at: Option<String>,
128    ) {
129        self.mark_installed(InstalledComponent {
130            id: id.to_string(),
131            version,
132            source_index,
133            location,
134            sealed: true,
135            installed_at,
136            pip_packages,
137        });
138    }
139
140    /// Lookup a component by id.
141    pub fn get(&self, id: &str) -> Option<&InstalledComponent> {
142        self.components.iter().find(|c| c.id == id)
143    }
144
145    /// Is a component with this id recorded as installed?
146    pub fn is_installed(&self, id: &str) -> bool {
147        self.components.iter().any(|c| c.id == id)
148    }
149
150    /// Is this a sealed core component?
151    pub fn is_sealed(&self, id: &str) -> bool {
152        self.get(id).is_some_and(|c| c.sealed)
153    }
154
155    /// All recorded component ids.
156    pub fn installed_ids(&self) -> HashSet<String> {
157        self.components.iter().map(|c| c.id.clone()).collect()
158    }
159
160    /// Every pip package name owned by any recorded component (for uninstall).
161    pub fn all_pip_packages(&self) -> Vec<String> {
162        let mut seen = HashSet::new();
163        let mut out = Vec::new();
164        for c in &self.components {
165            for p in &c.pip_packages {
166                if seen.insert(p.clone()) {
167                    out.push(p.clone());
168                }
169            }
170        }
171        out
172    }
173
174    /// Explicitly unseal a core component (only via deliberate user action,
175    /// e.g. `--unseal-core`). Returns `true` if it was sealed and is now not.
176    pub fn unseal(&mut self, id: &str) -> bool {
177        if let Some(c) = self.components.iter_mut().find(|c| c.id == id) {
178            let was = c.sealed;
179            c.sealed = false;
180            was
181        } else {
182            false
183        }
184    }
185
186    /// Remove a component record (does not uninstall — see Stage 6).
187    pub fn remove(&mut self, id: &str) -> bool {
188        let before = self.components.len();
189        self.components.retain(|c| c.id != id);
190        self.components.len() != before
191    }
192
193    /// Clear all records.
194    pub fn clear(&mut self) {
195        self.components.clear();
196    }
197
198    /// Number of recorded components.
199    pub fn len(&self) -> usize {
200        self.components.len()
201    }
202
203    /// Is the registry empty?
204    pub fn is_empty(&self) -> bool {
205        self.components.is_empty()
206    }
207
208    /// Iterate over recorded components.
209    pub fn iter(&self) -> impl Iterator<Item = &InstalledComponent> {
210        self.components.iter()
211    }
212}
213
214/// The set of ids treated as **core** (sealed-on-install) components.
215pub const CORE_COMPONENT_IDS: &[&str] = &["rocm", "pytorch", "triton"];
216
217/// Additional **install-once** components the tenet names ("aiter, flash
218/// attention, etc.") — expensive ROCm builds that, once installed, must be
219/// reused and never overridden. Sealed alongside the cores.
220pub const SEALED_INSTALL_ONCE_IDS: &[&str] = &[
221    "aiter",
222    "flash-attn",
223    "flash_attention",
224    "flash_attention_amd",
225    "rccl",
226    "migraphx",
227    "bitsandbytes",
228];
229
230/// Is the given id a core component?
231pub fn is_core_component(id: &str) -> bool {
232    CORE_COMPONENT_IDS.contains(&id)
233}
234
235/// Should this component be **sealed** on install (core OR install-once)?
236///
237/// A sealed component is reused (not reinstalled) on subsequent installs and
238/// protected from force-reinstall unless `MLSTACK_UNSEAL_CORE=1` — the tenet's
239/// "once installed, NEVER overridden" guarantee.
240pub fn should_seal_component(id: &str) -> bool {
241    is_core_component(id) || SEALED_INSTALL_ONCE_IDS.contains(&id)
242}
243
244// ===========================================================================
245// Tests
246// ===========================================================================
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    fn reg(id: &str, sealed: bool) -> InstalledComponent {
253        InstalledComponent {
254            id: id.to_string(),
255            version: Some("1.0".into()),
256            source_index: None,
257            location: None,
258            sealed,
259            installed_at: None,
260            pip_packages: vec![format!("{id}-pkg")],
261        }
262    }
263
264    #[test]
265    fn mark_and_lookup() {
266        let mut r = InstalledComponentRegistry::default();
267        r.mark_installed(reg("pytorch", true));
268        assert!(r.is_installed("pytorch"));
269        assert!(r.is_sealed("pytorch"));
270        assert!(!r.is_installed("rocm"));
271        assert_eq!(r.len(), 1);
272    }
273
274    #[test]
275    fn seal_is_preserved_on_remark() {
276        let mut r = InstalledComponentRegistry::default();
277        r.mark_installed(reg("pytorch", true));
278        // A later non-sealed remark with NEW non-seal fields must update those
279        // fields while preserving the seal (not skip the update entirely).
280        let updated = InstalledComponent {
281            id: "pytorch".to_string(),
282            version: Some("2.9.0".into()),
283            source_index: Some("https://download.pytorch.org/whl/rocm7.6".into()),
284            location: None,
285            sealed: false,
286            installed_at: None,
287            pip_packages: vec!["torch".into(), "torchvision".into()],
288        };
289        r.mark_installed(updated);
290        assert!(
291            r.is_sealed("pytorch"),
292            "seal must survive a non-sealed remark"
293        );
294        let got = r.get("pytorch").expect("component present");
295        assert_eq!(got.version.as_deref(), Some("2.9.0"), "version must update");
296        assert_eq!(
297            got.source_index.as_deref(),
298            Some("https://download.pytorch.org/whl/rocm7.6")
299        );
300        assert_eq!(got.pip_packages, vec!["torch", "torchvision"]);
301    }
302
303    #[test]
304    fn unseal_is_explicit() {
305        let mut r = InstalledComponentRegistry::default();
306        r.mark_installed(reg("pytorch", true));
307        assert!(r.unseal("pytorch"));
308        assert!(!r.is_sealed("pytorch"));
309    }
310
311    #[test]
312    fn all_pip_packages_dedups() {
313        let mut r = InstalledComponentRegistry::default();
314        let mut a = reg("pytorch", true);
315        a.pip_packages = vec!["torch".into(), "torchvision".into()];
316        let mut b = reg("vllm", false);
317        b.pip_packages = vec!["vllm".into(), "torch".into()]; // torch dup
318        r.mark_installed(a);
319        r.mark_installed(b);
320        let pkgs = r.all_pip_packages();
321        let torchs = pkgs.iter().filter(|p| *p == "torch").count();
322        assert_eq!(torchs, 1, "torch must appear once: {pkgs:?}");
323        assert!(pkgs.contains(&"vllm".to_string()));
324    }
325
326    #[test]
327    fn remove_and_clear() {
328        let mut r = InstalledComponentRegistry::default();
329        r.mark_installed(reg("rocm", true));
330        r.mark_installed(reg("pytorch", true));
331        assert!(r.remove("rocm"));
332        assert!(!r.is_installed("rocm"));
333        r.clear();
334        assert!(r.is_empty());
335    }
336
337    #[test]
338    fn load_missing_is_empty() {
339        // Serialize: this test mutates the process-global MLSTACK_USER_HOME.
340        let _guard = crate::test_support::lock_env();
341        // Point HOME at a temp dir so ~/.mlstack/installed.json is absent.
342        let dir = tempfile::tempdir().unwrap();
343        let saved_home = std::env::var("MLSTACK_USER_HOME").ok();
344        std::env::set_var(
345            "MLSTACK_USER_HOME",
346            dir.path().to_string_lossy().to_string(),
347        );
348        let r = InstalledComponentRegistry::load();
349        assert!(r.is_empty());
350        match saved_home {
351            Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
352            None => std::env::remove_var("MLSTACK_USER_HOME"),
353        }
354    }
355
356    #[test]
357    fn save_then_load_roundtrip() {
358        // Serialize: this test mutates the process-global MLSTACK_USER_HOME.
359        let _guard = crate::test_support::lock_env();
360        let dir = tempfile::tempdir().unwrap();
361        let saved_home = std::env::var("MLSTACK_USER_HOME").ok();
362        std::env::set_var(
363            "MLSTACK_USER_HOME",
364            dir.path().to_string_lossy().to_string(),
365        );
366
367        let mut r = InstalledComponentRegistry::default();
368        r.seal_core(
369            "pytorch",
370            Some("2.4.0".into()),
371            Some("https://download.pytorch.org/whl/rocm6.2".into()),
372            Some("/home/u/.mlstack/global".into()),
373            vec!["torch".into()],
374            Some("2026-06-24T00:00:00Z".into()),
375        );
376        r.save().unwrap();
377        assert!(InstalledComponentRegistry::path().exists());
378
379        let loaded = InstalledComponentRegistry::load();
380        assert!(loaded.is_sealed("pytorch"));
381        assert_eq!(
382            loaded.get("pytorch").unwrap().version.as_deref(),
383            Some("2.4.0")
384        );
385
386        match saved_home {
387            Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
388            None => std::env::remove_var("MLSTACK_USER_HOME"),
389        }
390    }
391
392    #[test]
393    fn core_ids_classify_correctly() {
394        assert!(is_core_component("rocm"));
395        assert!(is_core_component("pytorch"));
396        assert!(is_core_component("triton"));
397        assert!(!is_core_component("vllm"));
398        assert!(!is_core_component("deepspeed"));
399    }
400}
rusty_stack/core/registry.rs

rusty_stack/core/
registry.rs
