9c3d741a08
Move 17 library crates into crates/, CLI binary into bin/fxt, consolidate 10 test crates into testing/, split config crate from deployment config files. Root directory reduced from 38+ to ~17 directories. All Cargo.toml paths and build.rs proto refs updated. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
382 lines
12 KiB
Rust
382 lines
12 KiB
Rust
//! Checkpoint archival and backup/restore operation tests
|
|
//!
|
|
//! Tests comprehensive checkpoint management scenarios including:
|
|
//! - Large checkpoint uploads
|
|
//! - Checkpoint retrieval with verification
|
|
//! - Backup and restore workflows
|
|
//! - Concurrent checkpoint operations
|
|
//! - Error handling for checkpoint operations
|
|
|
|
use object_store::memory::InMemory;
|
|
use object_store::ObjectStore;
|
|
use std::sync::Arc;
|
|
use storage::object_store_backend::ObjectStoreBackend;
|
|
use storage::Storage;
|
|
|
|
/// Helper to create test backend with in-memory store
|
|
fn create_test_backend() -> ObjectStoreBackend {
|
|
let in_memory_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
|
|
storage::object_store_backend::test_helpers::new_for_testing(
|
|
in_memory_store,
|
|
"checkpoints-bucket".to_string(),
|
|
)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_upload_and_download() {
|
|
let backend = create_test_backend();
|
|
|
|
// Simulate checkpoint data (10MB)
|
|
let checkpoint_data = vec![0xAB; 10 * 1024 * 1024];
|
|
let checkpoint_path = backend.get_checkpoint_path("mamba2", "epoch_100");
|
|
|
|
// Upload checkpoint
|
|
backend
|
|
.store(&checkpoint_path, &checkpoint_data)
|
|
.await
|
|
.expect("Failed to upload checkpoint");
|
|
|
|
// Download checkpoint
|
|
let downloaded = backend
|
|
.retrieve(&checkpoint_path)
|
|
.await
|
|
.expect("Failed to download checkpoint");
|
|
|
|
assert_eq!(
|
|
downloaded.len(),
|
|
checkpoint_data.len(),
|
|
"Downloaded checkpoint size mismatch"
|
|
);
|
|
assert_eq!(
|
|
downloaded, checkpoint_data,
|
|
"Downloaded checkpoint data mismatch"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_metadata_storage() {
|
|
let backend = create_test_backend();
|
|
|
|
let checkpoint_path = backend.get_checkpoint_path("dqn", "epoch_50");
|
|
let metadata_path = backend.get_metadata_path("dqn", "v1.0");
|
|
|
|
// Store checkpoint
|
|
let checkpoint_data = b"checkpoint weights";
|
|
backend
|
|
.store(&checkpoint_path, checkpoint_data)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Store metadata
|
|
let metadata = serde_json::json!({
|
|
"model": "dqn",
|
|
"epoch": 50,
|
|
"loss": 0.123,
|
|
"accuracy": 0.95
|
|
});
|
|
let metadata_bytes = serde_json::to_vec(&metadata).unwrap();
|
|
backend
|
|
.store(&metadata_path, &metadata_bytes)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Verify both exist
|
|
assert!(backend.exists(&checkpoint_path).await.unwrap());
|
|
assert!(backend.exists(&metadata_path).await.unwrap());
|
|
|
|
// Retrieve and verify metadata
|
|
let retrieved_metadata = backend.retrieve(&metadata_path).await.unwrap();
|
|
let parsed: serde_json::Value = serde_json::from_slice(&retrieved_metadata).unwrap();
|
|
assert_eq!(parsed["model"], "dqn");
|
|
assert_eq!(parsed["epoch"], 50);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_backup_workflow() {
|
|
let backend = create_test_backend();
|
|
|
|
// Primary checkpoint
|
|
let primary_path = "models/ppo/checkpoints/primary/epoch_100.safetensors";
|
|
let checkpoint_data = vec![0xCD; 5 * 1024 * 1024]; // 5MB
|
|
|
|
// Store primary checkpoint
|
|
backend.store(primary_path, &checkpoint_data).await.unwrap();
|
|
|
|
// Create backup
|
|
let backup_path = "models/ppo/backups/epoch_100_backup.safetensors";
|
|
let retrieved = backend.retrieve(primary_path).await.unwrap();
|
|
backend.store(backup_path, &retrieved).await.unwrap();
|
|
|
|
// Verify both exist
|
|
assert!(backend.exists(primary_path).await.unwrap());
|
|
assert!(backend.exists(backup_path).await.unwrap());
|
|
|
|
// Verify backup integrity
|
|
let backup_data = backend.retrieve(backup_path).await.unwrap();
|
|
assert_eq!(backup_data, checkpoint_data);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_restore_from_backup() {
|
|
let backend = create_test_backend();
|
|
|
|
let backup_path = "backups/tft/epoch_200.safetensors";
|
|
let restore_path = "models/tft/checkpoints/epoch_200_restored.safetensors";
|
|
|
|
let backup_data = vec![0xEF; 8 * 1024 * 1024]; // 8MB
|
|
|
|
// Store backup
|
|
backend.store(backup_path, &backup_data).await.unwrap();
|
|
|
|
// Simulate restore operation
|
|
let restored_data = backend.retrieve(backup_path).await.unwrap();
|
|
backend.store(restore_path, &restored_data).await.unwrap();
|
|
|
|
// Verify restored checkpoint
|
|
let final_data = backend.retrieve(restore_path).await.unwrap();
|
|
assert_eq!(final_data.len(), backup_data.len());
|
|
assert_eq!(final_data, backup_data);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_versioning() {
|
|
let backend = create_test_backend();
|
|
|
|
let model_name = "mamba2";
|
|
let versions = ["v1.0", "v1.1", "v2.0"];
|
|
|
|
// Store multiple checkpoint versions
|
|
for version in &versions {
|
|
let path = backend.get_checkpoint_path(model_name, version);
|
|
let data = format!("checkpoint_{}", version).into_bytes();
|
|
backend.store(&path, &data).await.unwrap();
|
|
}
|
|
|
|
// List all checkpoints
|
|
let prefix = format!("models/{}/checkpoints/", model_name);
|
|
let checkpoints = backend.list(&prefix).await.unwrap();
|
|
|
|
assert_eq!(checkpoints.len(), versions.len());
|
|
|
|
// Verify each version exists
|
|
for version in &versions {
|
|
let path = backend.get_checkpoint_path(model_name, version);
|
|
assert!(backend.exists(&path).await.unwrap());
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_deletion() {
|
|
let backend = create_test_backend();
|
|
|
|
let checkpoint_path = backend.get_checkpoint_path("dqn", "old_epoch_10");
|
|
let checkpoint_data = b"old checkpoint data";
|
|
|
|
// Store checkpoint
|
|
backend
|
|
.store(&checkpoint_path, checkpoint_data)
|
|
.await
|
|
.unwrap();
|
|
assert!(backend.exists(&checkpoint_path).await.unwrap());
|
|
|
|
// Delete checkpoint
|
|
let deleted = backend.delete(&checkpoint_path).await.unwrap();
|
|
assert!(deleted);
|
|
|
|
// Verify deletion
|
|
assert!(!backend.exists(&checkpoint_path).await.unwrap());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_cleanup_old_versions() {
|
|
let backend = create_test_backend();
|
|
|
|
let model_name = "ppo";
|
|
let max_checkpoints = 3;
|
|
|
|
// Store 5 checkpoints (should keep only 3 latest)
|
|
for i in 1..=5 {
|
|
let path = backend.get_checkpoint_path(model_name, &format!("epoch_{}", i * 10));
|
|
let data = format!("checkpoint_{}", i).into_bytes();
|
|
backend.store(&path, &data).await.unwrap();
|
|
}
|
|
|
|
let prefix = format!("models/{}/checkpoints/", model_name);
|
|
let all_checkpoints = backend.list(&prefix).await.unwrap();
|
|
assert_eq!(all_checkpoints.len(), 5);
|
|
|
|
// Simulate cleanup: delete oldest checkpoints
|
|
let mut to_delete = all_checkpoints.clone();
|
|
to_delete.sort();
|
|
while to_delete.len() > max_checkpoints {
|
|
let oldest = to_delete.remove(0);
|
|
backend.delete(&oldest).await.unwrap();
|
|
}
|
|
|
|
// Verify only max_checkpoints remain
|
|
let remaining = backend.list(&prefix).await.unwrap();
|
|
assert_eq!(remaining.len(), max_checkpoints);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_concurrent_checkpoint_operations() {
|
|
let backend = Arc::new(create_test_backend());
|
|
let mut handles = vec![];
|
|
|
|
// Concurrent checkpoint uploads
|
|
for i in 0..5 {
|
|
let backend = Arc::clone(&backend);
|
|
let handle = tokio::spawn(async move {
|
|
let path = format!("models/concurrent_test/checkpoints/epoch_{}.bin", i);
|
|
let data = vec![i as u8; 1024 * 1024]; // 1MB each
|
|
backend.store(&path, &data).await.unwrap();
|
|
|
|
// Verify upload
|
|
let retrieved = backend.retrieve(&path).await.unwrap();
|
|
assert_eq!(retrieved.len(), data.len());
|
|
});
|
|
handles.push(handle);
|
|
}
|
|
|
|
for handle in handles {
|
|
handle.await.unwrap();
|
|
}
|
|
|
|
// Verify all checkpoints exist
|
|
let checkpoints = backend
|
|
.list("models/concurrent_test/checkpoints/")
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(checkpoints.len(), 5);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_integrity_verification() {
|
|
let backend = create_test_backend();
|
|
|
|
let checkpoint_path = backend.get_checkpoint_path("tft", "epoch_150");
|
|
let checkpoint_data = vec![0x42; 15 * 1024 * 1024]; // 15MB
|
|
|
|
// Calculate expected checksum (SHA-256)
|
|
use sha2::{Digest, Sha256};
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(&checkpoint_data);
|
|
let expected_checksum = format!("{:x}", hasher.finalize());
|
|
|
|
// Store checkpoint
|
|
backend
|
|
.store(&checkpoint_path, &checkpoint_data)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Retrieve and verify checksum
|
|
let retrieved = backend.retrieve(&checkpoint_path).await.unwrap();
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(&retrieved);
|
|
let actual_checksum = format!("{:x}", hasher.finalize());
|
|
|
|
assert_eq!(actual_checksum, expected_checksum);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_partial_upload_failure() {
|
|
let backend = create_test_backend();
|
|
|
|
let checkpoint_path = backend.get_checkpoint_path("mamba2", "epoch_partial");
|
|
let checkpoint_data = vec![0xAA; 20 * 1024 * 1024]; // 20MB
|
|
|
|
// Upload checkpoint
|
|
let result = backend.store(&checkpoint_path, &checkpoint_data).await;
|
|
assert!(result.is_ok());
|
|
|
|
// Verify metadata reflects correct size
|
|
let metadata = backend.metadata(&checkpoint_path).await.unwrap();
|
|
assert_eq!(metadata.size, checkpoint_data.len() as u64);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_list_with_pagination() {
|
|
let backend = create_test_backend();
|
|
|
|
// Store 20 checkpoints
|
|
for i in 1..=20 {
|
|
let path = format!("models/pagination_test/checkpoints/epoch_{}.bin", i);
|
|
backend.store(&path, &vec![i as u8; 1024]).await.unwrap();
|
|
}
|
|
|
|
// List all checkpoints
|
|
let all_checkpoints = backend
|
|
.list("models/pagination_test/checkpoints/")
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(all_checkpoints.len(), 20);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_empty_content() {
|
|
let backend = create_test_backend();
|
|
|
|
let checkpoint_path = backend.get_checkpoint_path("empty_test", "epoch_0");
|
|
let empty_data = b"";
|
|
|
|
// Store empty checkpoint
|
|
backend.store(&checkpoint_path, empty_data).await.unwrap();
|
|
|
|
// Retrieve and verify
|
|
let retrieved = backend.retrieve(&checkpoint_path).await.unwrap();
|
|
assert_eq!(retrieved.len(), 0);
|
|
|
|
// Verify metadata
|
|
let metadata = backend.metadata(&checkpoint_path).await.unwrap();
|
|
assert_eq!(metadata.size, 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_overwrite_protection() {
|
|
let backend = create_test_backend();
|
|
|
|
let checkpoint_path = backend.get_checkpoint_path("overwrite_test", "epoch_100");
|
|
let original_data = b"original checkpoint v1";
|
|
let new_data = b"updated checkpoint v2";
|
|
|
|
// Store original
|
|
backend
|
|
.store(&checkpoint_path, original_data)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Overwrite (should succeed in S3)
|
|
backend.store(&checkpoint_path, new_data).await.unwrap();
|
|
|
|
// Verify overwrite
|
|
let retrieved = backend.retrieve(&checkpoint_path).await.unwrap();
|
|
assert_eq!(retrieved, new_data);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_checkpoint_metadata_size_validation() {
|
|
let backend = create_test_backend();
|
|
|
|
let sizes = [
|
|
1024, // 1KB
|
|
1024 * 1024, // 1MB
|
|
10 * 1024 * 1024, // 10MB
|
|
100 * 1024 * 1024, // 100MB (large checkpoint)
|
|
];
|
|
|
|
for (idx, &size) in sizes.iter().enumerate() {
|
|
let path = format!("models/size_test/checkpoint_{}.bin", idx);
|
|
let data = vec![0xFF; size];
|
|
|
|
backend.store(&path, &data).await.unwrap();
|
|
|
|
let metadata = backend.metadata(&path).await.unwrap();
|
|
assert_eq!(
|
|
metadata.size, size as u64,
|
|
"Size mismatch for checkpoint {}",
|
|
idx
|
|
);
|
|
}
|
|
}
|