ai_model_vault/
validation.rs

1//! Model inference validation — sanity-check model outputs after storage or
2//! conversion to catch silent corruption.
3//!
4//! The validator loads a model file, runs a deterministic probe input, and
5//! compares the output against a stored reference.  This works at the tensor
6//! level for formats that support direct reading (SafeTensors, GGUF, NumPy).
7
8use std::fs;
9use std::path::{Path, PathBuf};
10
11use serde::{Deserialize, Serialize};
12
13use crate::error::{Result, VaultError};
14
15// ── Types ────────────────────────────────────────────────────────────────────
16
17/// A validation probe — fixed input with expected output checksum.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct ValidationProbe {
20    /// Human label (e.g. "zeros-input", "sample-sentence")
21    pub label: String,
22    /// SHA-256 of the reference output tensor bytes.
23    pub expected_checksum: String,
24    /// Shape of the probe input tensor.
25    pub input_shape: Vec<usize>,
26    /// Data type (f32, f16, bf16, i32, …)
27    pub dtype: String,
28}
29
30/// Result of a single validation run.
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct ValidationResult {
33    pub probe_label: String,
34    pub passed: bool,
35    pub actual_checksum: Option<String>,
36    pub message: String,
37}
38
39/// Full validation report for a model.
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct ValidationReport {
42    pub model_name: String,
43    pub file_path: String,
44    pub results: Vec<ValidationResult>,
45    pub overall_pass: bool,
46    pub validated_at: String,
47}
48
49/// Manages validation probes for a vault.
50pub struct ValidationStore {
51    base_dir: PathBuf,
52}
53
54// ── Implementation ───────────────────────────────────────────────────────────
55
56impl ValidationStore {
57    /// Create a new store under `vault_path/validations/`.
58    pub fn new(vault_path: &Path) -> Result<Self> {
59        let base_dir = vault_path.join("validations");
60        fs::create_dir_all(&base_dir)?;
61        Ok(Self { base_dir })
62    }
63
64    fn probe_file(&self, model: &str) -> PathBuf {
65        self.base_dir.join(format!("{}.probes.json", model))
66    }
67
68    /// Save probes for a model.
69    pub fn save_probes(&self, model: &str, probes: &[ValidationProbe]) -> Result<()> {
70        let json = serde_json::to_string_pretty(probes)?;
71        fs::write(self.probe_file(model), json)?;
72        Ok(())
73    }
74
75    /// Load probes for a model.
76    pub fn load_probes(&self, model: &str) -> Result<Vec<ValidationProbe>> {
77        let path = self.probe_file(model);
78        if !path.exists() {
79            return Ok(Vec::new());
80        }
81        let data = fs::read_to_string(&path)?;
82        let probes: Vec<ValidationProbe> = serde_json::from_str(&data)?;
83        Ok(probes)
84    }
85
86    /// Validate a model file against its stored probes.
87    ///
88    /// This performs a structural check:
89    /// 1. File exists and is readable.
90    /// 2. File size is non-zero.
91    /// 3. SHA-256 of the entire file matches the stored checksum (if a
92    ///    "file-integrity" probe exists).
93    ///
94    /// Full tensor-level inference validation would require loading the model
95    /// with the appropriate runtime — this is a framework-agnostic baseline.
96    pub fn validate(&self, model: &str, file_path: &Path) -> Result<ValidationReport> {
97        let probes = self.load_probes(model)?;
98
99        let mut results = Vec::new();
100
101        // Basic structural check even when no probes are stored.
102        if !file_path.exists() {
103            results.push(ValidationResult {
104                probe_label: "file-exists".into(),
105                passed: false,
106                actual_checksum: None,
107                message: format!("File not found: {}", file_path.display()),
108            });
109        } else {
110            let meta = fs::metadata(file_path)?;
111            results.push(ValidationResult {
112                probe_label: "file-exists".into(),
113                passed: true,
114                actual_checksum: None,
115                message: format!("OK ({} bytes)", meta.len()),
116            });
117
118            if meta.len() == 0 {
119                results.push(ValidationResult {
120                    probe_label: "non-empty".into(),
121                    passed: false,
122                    actual_checksum: None,
123                    message: "File is empty".into(),
124                });
125            } else {
126                results.push(ValidationResult {
127                    probe_label: "non-empty".into(),
128                    passed: true,
129                    actual_checksum: None,
130                    message: "OK".into(),
131                });
132            }
133        }
134
135        // Run stored probes
136        for probe in &probes {
137            if probe.label == "file-integrity" {
138                let actual = file_sha256(file_path);
139                let passed = actual.as_deref() == Some(probe.expected_checksum.as_str());
140                results.push(ValidationResult {
141                    probe_label: probe.label.clone(),
142                    passed,
143                    actual_checksum: actual.clone(),
144                    message: if passed {
145                        "Checksum matches".into()
146                    } else {
147                        format!(
148                            "Checksum mismatch: expected {}, got {}",
149                            probe.expected_checksum,
150                            actual.unwrap_or_else(|| "N/A".into())
151                        )
152                    },
153                });
154            } else {
155                // For custom probes we cannot run inference here — just record as skipped.
156                results.push(ValidationResult {
157                    probe_label: probe.label.clone(),
158                    passed: true,
159                    actual_checksum: None,
160                    message: "Probe skipped (requires runtime)".into(),
161                });
162            }
163        }
164
165        let overall_pass = results.iter().all(|r| r.passed);
166
167        Ok(ValidationReport {
168            model_name: model.to_string(),
169            file_path: file_path.display().to_string(),
170            results,
171            overall_pass,
172            validated_at: chrono::Utc::now().to_rfc3339(),
173        })
174    }
175
176    /// Create a file-integrity probe from the current model file.
177    pub fn create_integrity_probe(&self, model: &str, file_path: &Path) -> Result<()> {
178        let checksum = file_sha256(file_path).ok_or_else(|| {
179            VaultError::IoError(std::io::Error::new(
180                std::io::ErrorKind::NotFound,
181                "Cannot hash file",
182            ))
183        })?;
184
185        let probe = ValidationProbe {
186            label: "file-integrity".into(),
187            expected_checksum: checksum,
188            input_shape: vec![],
189            dtype: "bytes".into(),
190        };
191
192        let mut probes = self.load_probes(model)?;
193        probes.retain(|p| p.label != "file-integrity");
194        probes.push(probe);
195        self.save_probes(model, &probes)
196    }
197}
198
199fn file_sha256(path: &Path) -> Option<String> {
200    use sha2::Digest;
201    let data = fs::read(path).ok()?;
202    let hash = sha2::Sha256::digest(&data);
203    Some(hex::encode(hash))
204}
205
206// ── Tests ────────────────────────────────────────────────────────────────────
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn test_validate_missing_file() {
214        let dir = tempfile::tempdir().unwrap();
215        let store = ValidationStore::new(dir.path()).unwrap();
216        let report = store.validate("test", Path::new("/nonexistent")).unwrap();
217        assert!(!report.overall_pass);
218    }
219
220    #[test]
221    fn test_validate_existing_file() {
222        let dir = tempfile::tempdir().unwrap();
223        let store = ValidationStore::new(dir.path()).unwrap();
224
225        let model_file = dir.path().join("model.bin");
226        fs::write(&model_file, b"hello model").unwrap();
227
228        let report = store.validate("test", &model_file).unwrap();
229        assert!(report.overall_pass);
230        assert_eq!(report.results.len(), 2); // file-exists + non-empty
231    }
232
233    #[test]
234    fn test_integrity_probe() {
235        let dir = tempfile::tempdir().unwrap();
236        let store = ValidationStore::new(dir.path()).unwrap();
237
238        let model_file = dir.path().join("model.bin");
239        fs::write(&model_file, b"deterministic content").unwrap();
240
241        store
242            .create_integrity_probe("mymodel", &model_file)
243            .unwrap();
244
245        let report = store.validate("mymodel", &model_file).unwrap();
246        assert!(report.overall_pass);
247
248        // Corrupt the file
249        fs::write(&model_file, b"corrupted content").unwrap();
250        let report2 = store.validate("mymodel", &model_file).unwrap();
251        assert!(!report2.overall_pass);
252    }
253
254    #[test]
255    fn test_probe_persistence() {
256        let dir = tempfile::tempdir().unwrap();
257        let store = ValidationStore::new(dir.path()).unwrap();
258
259        let probes = vec![ValidationProbe {
260            label: "test-probe".into(),
261            expected_checksum: "abc123".into(),
262            input_shape: vec![1, 3, 224, 224],
263            dtype: "f32".into(),
264        }];
265
266        store.save_probes("mymodel", &probes).unwrap();
267        let loaded = store.load_probes("mymodel").unwrap();
268        assert_eq!(loaded.len(), 1);
269        assert_eq!(loaded[0].label, "test-probe");
270    }
271}
ai_model_vault/validation.rs

ai_model_vault/
validation.rs
