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    "rccl",
225    "migraphx",
226    "bitsandbytes",
227];
228// NOTE: the split Flash Attention backends (flash-attn-triton, flash-attn-ck) are
229// deliberately NOT sealed — installing one over the other is the SUPPORTED backend
230// switch (the ~/.mlstack/flash-attention/.backend marker makes the last-installed
231// active, enforcing mutual exclusivity). The canonical "flash-attn" id stays
232// sealed. "flash_attention_amd" (legacy) was removed from sealing and from
233// component_pip_packages; it remains in uninstall/purge paths only.
234
235/// Is the given id a core component?
236pub fn is_core_component(id: &str) -> bool {
237    CORE_COMPONENT_IDS.contains(&id)
238}
239
240/// Should this component be **sealed** on install (core OR install-once)?
241///
242/// A sealed component is reused (not reinstalled) on subsequent installs and
243/// protected from force-reinstall unless `MLSTACK_UNSEAL_CORE=1` — the tenet's
244/// "once installed, NEVER overridden" guarantee.
245pub fn should_seal_component(id: &str) -> bool {
246    is_core_component(id) || SEALED_INSTALL_ONCE_IDS.contains(&id)
247}
248
249// ===========================================================================
250// Tests
251// ===========================================================================
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    fn reg(id: &str, sealed: bool) -> InstalledComponent {
258        InstalledComponent {
259            id: id.to_string(),
260            version: Some("1.0".into()),
261            source_index: None,
262            location: None,
263            sealed,
264            installed_at: None,
265            pip_packages: vec![format!("{id}-pkg")],
266        }
267    }
268
269    #[test]
270    fn mark_and_lookup() {
271        let mut r = InstalledComponentRegistry::default();
272        r.mark_installed(reg("pytorch", true));
273        assert!(r.is_installed("pytorch"));
274        assert!(r.is_sealed("pytorch"));
275        assert!(!r.is_installed("rocm"));
276        assert_eq!(r.len(), 1);
277    }
278
279    #[test]
280    fn seal_is_preserved_on_remark() {
281        let mut r = InstalledComponentRegistry::default();
282        r.mark_installed(reg("pytorch", true));
283        // A later non-sealed remark with NEW non-seal fields must update those
284        // fields while preserving the seal (not skip the update entirely).
285        let updated = InstalledComponent {
286            id: "pytorch".to_string(),
287            version: Some("2.9.0".into()),
288            source_index: Some("https://download.pytorch.org/whl/rocm7.6".into()),
289            location: None,
290            sealed: false,
291            installed_at: None,
292            pip_packages: vec!["torch".into(), "torchvision".into()],
293        };
294        r.mark_installed(updated);
295        assert!(
296            r.is_sealed("pytorch"),
297            "seal must survive a non-sealed remark"
298        );
299        let got = r.get("pytorch").expect("component present");
300        assert_eq!(got.version.as_deref(), Some("2.9.0"), "version must update");
301        assert_eq!(
302            got.source_index.as_deref(),
303            Some("https://download.pytorch.org/whl/rocm7.6")
304        );
305        assert_eq!(got.pip_packages, vec!["torch", "torchvision"]);
306    }
307
308    #[test]
309    fn unseal_is_explicit() {
310        let mut r = InstalledComponentRegistry::default();
311        r.mark_installed(reg("pytorch", true));
312        assert!(r.unseal("pytorch"));
313        assert!(!r.is_sealed("pytorch"));
314    }
315
316    #[test]
317    fn all_pip_packages_dedups() {
318        let mut r = InstalledComponentRegistry::default();
319        let mut a = reg("pytorch", true);
320        a.pip_packages = vec!["torch".into(), "torchvision".into()];
321        let mut b = reg("vllm", false);
322        b.pip_packages = vec!["vllm".into(), "torch".into()]; // torch dup
323        r.mark_installed(a);
324        r.mark_installed(b);
325        let pkgs = r.all_pip_packages();
326        let torchs = pkgs.iter().filter(|p| *p == "torch").count();
327        assert_eq!(torchs, 1, "torch must appear once: {pkgs:?}");
328        assert!(pkgs.contains(&"vllm".to_string()));
329    }
330
331    #[test]
332    fn remove_and_clear() {
333        let mut r = InstalledComponentRegistry::default();
334        r.mark_installed(reg("rocm", true));
335        r.mark_installed(reg("pytorch", true));
336        assert!(r.remove("rocm"));
337        assert!(!r.is_installed("rocm"));
338        r.clear();
339        assert!(r.is_empty());
340    }
341
342    #[test]
343    fn load_missing_is_empty() {
344        // Serialize: this test mutates the process-global MLSTACK_USER_HOME.
345        let _guard = crate::test_support::lock_env();
346        // Point HOME at a temp dir so ~/.mlstack/installed.json is absent.
347        let dir = tempfile::tempdir().unwrap();
348        let saved_home = std::env::var("MLSTACK_USER_HOME").ok();
349        std::env::set_var(
350            "MLSTACK_USER_HOME",
351            dir.path().to_string_lossy().to_string(),
352        );
353        let r = InstalledComponentRegistry::load();
354        assert!(r.is_empty());
355        match saved_home {
356            Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
357            None => std::env::remove_var("MLSTACK_USER_HOME"),
358        }
359    }
360
361    #[test]
362    fn save_then_load_roundtrip() {
363        // Serialize: this test mutates the process-global MLSTACK_USER_HOME.
364        let _guard = crate::test_support::lock_env();
365        let dir = tempfile::tempdir().unwrap();
366        let saved_home = std::env::var("MLSTACK_USER_HOME").ok();
367        std::env::set_var(
368            "MLSTACK_USER_HOME",
369            dir.path().to_string_lossy().to_string(),
370        );
371
372        let mut r = InstalledComponentRegistry::default();
373        r.seal_core(
374            "pytorch",
375            Some("2.4.0".into()),
376            Some("https://download.pytorch.org/whl/rocm6.2".into()),
377            Some("/home/u/.mlstack/global".into()),
378            vec!["torch".into()],
379            Some("2026-06-24T00:00:00Z".into()),
380        );
381        r.save().unwrap();
382        assert!(InstalledComponentRegistry::path().exists());
383
384        let loaded = InstalledComponentRegistry::load();
385        assert!(loaded.is_sealed("pytorch"));
386        assert_eq!(
387            loaded.get("pytorch").unwrap().version.as_deref(),
388            Some("2.4.0")
389        );
390
391        match saved_home {
392            Some(v) => std::env::set_var("MLSTACK_USER_HOME", v),
393            None => std::env::remove_var("MLSTACK_USER_HOME"),
394        }
395    }
396
397    #[test]
398    fn core_ids_classify_correctly() {
399        assert!(is_core_component("rocm"));
400        assert!(is_core_component("pytorch"));
401        assert!(is_core_component("triton"));
402        assert!(!is_core_component("vllm"));
403        assert!(!is_core_component("deepspeed"));
404    }
405}
rusty_stack/core/registry.rs

rusty_stack/core/
registry.rs
