feat: add optimize_with_checkpoint() and atomic save writes
Adds a convenience method that combines optimize_with_callback + save to automatically checkpoint every N trials. Also makes save() use atomic writes (write to temp file, then rename) to prevent corrupt files on crash.
This commit is contained in:
@@ -172,6 +172,117 @@ fn round_trip_preserves_trial_id_counter() {
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn checkpoint_file_created_at_interval() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(-10.0, 10.0).name("x");
|
||||
|
||||
let dir = tempdir();
|
||||
let checkpoint = dir.join("checkpoint.json");
|
||||
|
||||
study
|
||||
.optimize_with_checkpoint(10, 3, &checkpoint, |trial| {
|
||||
let v = x.suggest(trial)?;
|
||||
Ok::<_, optimizer::Error>(v * v)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Checkpoint should exist (written at trials 3, 6, 9)
|
||||
assert!(checkpoint.exists(), "checkpoint file was not created");
|
||||
|
||||
// Load it and verify it's valid
|
||||
let loaded: Study<f64> = Study::load(&checkpoint).unwrap();
|
||||
// Last checkpoint was at trial 9, so it should have 9 trials
|
||||
assert_eq!(loaded.n_trials(), 9);
|
||||
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn checkpoint_overwrites_previous() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
let dir = tempdir();
|
||||
let checkpoint = dir.join("checkpoint.json");
|
||||
|
||||
study
|
||||
.optimize_with_checkpoint(6, 3, &checkpoint, |trial| {
|
||||
let v = x.suggest(trial)?;
|
||||
Ok::<_, optimizer::Error>(v)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// The checkpoint at trial 6 should overwrite the one from trial 3
|
||||
let loaded: Study<f64> = Study::load(&checkpoint).unwrap();
|
||||
assert_eq!(loaded.n_trials(), 6);
|
||||
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resume_from_checkpoint_continues_trial_ids() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(-5.0, 5.0).name("x");
|
||||
|
||||
let dir = tempdir();
|
||||
let checkpoint = dir.join("resume.json");
|
||||
|
||||
// Run 10 trials with checkpointing
|
||||
study
|
||||
.optimize_with_checkpoint(10, 5, &checkpoint, |trial| {
|
||||
let v = x.suggest(trial)?;
|
||||
Ok::<_, optimizer::Error>(v * v)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Load and continue
|
||||
let loaded: Study<f64> = Study::load(&checkpoint).unwrap();
|
||||
assert_eq!(loaded.n_trials(), 10);
|
||||
|
||||
let remaining = 15 - loaded.n_trials();
|
||||
loaded
|
||||
.optimize(remaining, |trial| {
|
||||
let v = x.suggest(trial)?;
|
||||
Ok::<_, optimizer::Error>(v * v)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(loaded.n_trials(), 15);
|
||||
|
||||
// Verify no duplicate trial IDs
|
||||
let trials = loaded.trials();
|
||||
let mut ids: Vec<u64> = trials.iter().map(|t| t.id).collect();
|
||||
ids.sort_unstable();
|
||||
ids.dedup();
|
||||
assert_eq!(ids.len(), 15, "duplicate trial IDs found");
|
||||
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn atomic_write_no_temp_file_left_behind() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
let dir = tempdir();
|
||||
let checkpoint = dir.join("atomic.json");
|
||||
|
||||
study
|
||||
.optimize_with_checkpoint(3, 3, &checkpoint, |trial| {
|
||||
let v = x.suggest(trial)?;
|
||||
Ok::<_, optimizer::Error>(v)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// The temp file should have been renamed, not left behind
|
||||
let tmp_path = dir.join(".atomic.json.tmp");
|
||||
assert!(!tmp_path.exists(), "temp file was not cleaned up");
|
||||
assert!(checkpoint.exists(), "checkpoint file was not created");
|
||||
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
}
|
||||
|
||||
/// Helper to create a unique temporary directory.
|
||||
fn tempdir() -> std::path::PathBuf {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
Reference in New Issue
Block a user