feat(ml): SHA-256 checksum validation for model checkpoint integrity

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-23 10:23:51 +01:00
parent 15c0fb5397
commit 10032ca13f
2 changed files with 79 additions and 1 deletions

View File

@@ -64,7 +64,7 @@ pub use signer::{CheckpointSigner, SignatureInfo};
#[cfg(feature = "s3-storage")]
pub use storage::S3CheckpointStorage;
pub use storage::{CheckpointStorage, FileSystemStorage, MemoryStorage, StorageStats};
pub use validation::ValidationManager;
pub use validation::{verify_checksum, write_checksum, ValidationManager};
pub use versioning::VersionManager;
/// Checkpoint format options

View File

@@ -3,6 +3,7 @@
//! Provides checksum validation and corruption detection for checkpoints.
use std::collections::HashMap;
use std::path::Path;
use sha2::{Digest, Sha256};
use tracing::{debug, error, warn};
@@ -10,6 +11,43 @@ use tracing::{debug, error, warn};
use super::{CheckpointMetadata, ModelType};
use crate::MLError;
/// Write SHA-256 checksum sidecar file alongside a safetensors checkpoint.
/// Creates `{path}.sha256` containing the hex digest.
pub fn write_checksum(safetensors_path: &Path) -> Result<(), MLError> {
let bytes = std::fs::read(safetensors_path).map_err(|e| {
MLError::CheckpointError(format!("Failed to read file for checksum: {}", e))
})?;
let hash = Sha256::digest(&bytes);
let hex = format!("{:x}", hash);
let checksum_path = safetensors_path.with_extension("sha256");
std::fs::write(&checksum_path, hex.as_bytes()).map_err(|e| {
MLError::CheckpointError(format!("Failed to write checksum: {}", e))
})?;
Ok(())
}
/// Verify SHA-256 checksum of a safetensors file against its `.sha256` sidecar.
/// Returns `Ok(true)` if valid, `Ok(false)` if mismatch.
/// Returns `Ok(true)` if no sidecar exists (backwards compatible).
pub fn verify_checksum(safetensors_path: &Path) -> Result<bool, MLError> {
let checksum_path = safetensors_path.with_extension("sha256");
if !checksum_path.exists() {
warn!(
"No checksum file for {}, skipping integrity check",
safetensors_path.display()
);
return Ok(true);
}
let expected = std::fs::read_to_string(&checksum_path).map_err(|e| {
MLError::CheckpointError(format!("Failed to read checksum: {}", e))
})?;
let bytes = std::fs::read(safetensors_path).map_err(|e| {
MLError::CheckpointError(format!("Failed to read file for verification: {}", e))
})?;
let actual = format!("{:x}", Sha256::digest(&bytes));
Ok(actual.trim() == expected.trim())
}
/// Validation manager for checkpoint integrity
#[derive(Debug)]
pub struct ValidationManager {
@@ -522,4 +560,44 @@ mod tests {
assert!(summary.contains("1 errors"));
assert!(summary.contains("1 warnings"));
}
#[test]
fn test_sidecar_checksum_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let dir = tempfile::TempDir::new()?;
let path = dir.path().join("test_model.safetensors");
std::fs::write(&path, b"fake model data")?;
write_checksum(&path)?;
let checksum_path = path.with_extension("sha256");
assert!(checksum_path.exists());
assert!(verify_checksum(&path)?);
Ok(())
}
#[test]
fn test_sidecar_checksum_detects_corruption() -> Result<(), Box<dyn std::error::Error>> {
let dir = tempfile::TempDir::new()?;
let path = dir.path().join("test_model.safetensors");
std::fs::write(&path, b"original data")?;
write_checksum(&path)?;
// Corrupt the file
std::fs::write(&path, b"corrupted data")?;
assert!(!verify_checksum(&path)?);
Ok(())
}
#[test]
fn test_verify_without_sidecar_returns_true() -> Result<(), Box<dyn std::error::Error>> {
let dir = tempfile::TempDir::new()?;
let path = dir.path().join("test_model.safetensors");
std::fs::write(&path, b"no checksum file")?;
// No .sha256 file exists -- should return true (backwards compatible)
assert!(verify_checksum(&path)?);
Ok(())
}
}