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:
@@ -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
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user