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:
+46
-2
@@ -1688,6 +1688,7 @@ impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
|
|||||||
///
|
///
|
||||||
/// Returns an I/O error if the file cannot be created or written.
|
/// Returns an I/O error if the file cannot be created or written.
|
||||||
pub fn save(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
|
pub fn save(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
|
||||||
|
let path = path.as_ref();
|
||||||
let snapshot = StudySnapshot {
|
let snapshot = StudySnapshot {
|
||||||
version: 1,
|
version: 1,
|
||||||
direction: self.direction,
|
direction: self.direction,
|
||||||
@@ -1695,8 +1696,51 @@ impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
|
|||||||
next_trial_id: self.next_trial_id.load(Ordering::Relaxed),
|
next_trial_id: self.next_trial_id.load(Ordering::Relaxed),
|
||||||
metadata: HashMap::new(),
|
metadata: HashMap::new(),
|
||||||
};
|
};
|
||||||
let file = std::fs::File::create(path)?;
|
|
||||||
serde_json::to_writer_pretty(file, &snapshot).map_err(std::io::Error::other)
|
// Atomic write: write to a temp file in the same directory, then rename.
|
||||||
|
// This prevents corrupt files if the process crashes mid-write.
|
||||||
|
let parent = path.parent().unwrap_or(std::path::Path::new("."));
|
||||||
|
let tmp_path = parent.join(format!(
|
||||||
|
".{}.tmp",
|
||||||
|
path.file_name().unwrap_or_default().to_string_lossy()
|
||||||
|
));
|
||||||
|
let file = std::fs::File::create(&tmp_path)?;
|
||||||
|
serde_json::to_writer_pretty(file, &snapshot).map_err(std::io::Error::other)?;
|
||||||
|
std::fs::rename(&tmp_path, path)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(feature = "serde")]
|
||||||
|
impl<V: PartialOrd + Clone + Default + serde::Serialize> Study<V> {
|
||||||
|
/// Runs optimization with automatic checkpointing every `interval` trials.
|
||||||
|
///
|
||||||
|
/// This is convenience sugar over [`optimize_with_callback`](Self::optimize_with_callback)
|
||||||
|
/// combined with [`save`](Self::save). The checkpoint is written atomically so
|
||||||
|
/// a crash mid-write will never leave a corrupt file.
|
||||||
|
///
|
||||||
|
/// # Errors
|
||||||
|
///
|
||||||
|
/// Returns an error if the optimization itself fails (see
|
||||||
|
/// [`optimize`](Self::optimize) for details). Checkpoint I/O errors are
|
||||||
|
/// silently ignored (best-effort).
|
||||||
|
pub fn optimize_with_checkpoint<F, E>(
|
||||||
|
&self,
|
||||||
|
n_trials: usize,
|
||||||
|
checkpoint_interval: usize,
|
||||||
|
checkpoint_path: impl AsRef<std::path::Path>,
|
||||||
|
objective: F,
|
||||||
|
) -> crate::Result<()>
|
||||||
|
where
|
||||||
|
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
|
||||||
|
E: ToString + 'static,
|
||||||
|
{
|
||||||
|
let path = checkpoint_path.as_ref().to_owned();
|
||||||
|
self.optimize_with_callback(n_trials, objective, |study, _trial| {
|
||||||
|
if study.n_trials().is_multiple_of(checkpoint_interval) {
|
||||||
|
let _ = study.save(&path);
|
||||||
|
}
|
||||||
|
ControlFlow::Continue(())
|
||||||
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -172,6 +172,117 @@ fn round_trip_preserves_trial_id_counter() {
|
|||||||
std::fs::remove_dir_all(&dir).ok();
|
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.
|
/// Helper to create a unique temporary directory.
|
||||||
fn tempdir() -> std::path::PathBuf {
|
fn tempdir() -> std::path::PathBuf {
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
|||||||
Reference in New Issue
Block a user