diff --git a/src/study.rs b/src/study.rs index ec50734..94e4447 100644 --- a/src/study.rs +++ b/src/study.rs @@ -1688,6 +1688,7 @@ impl Study { /// /// Returns an I/O error if the file cannot be created or written. pub fn save(&self, path: impl AsRef) -> std::io::Result<()> { + let path = path.as_ref(); let snapshot = StudySnapshot { version: 1, direction: self.direction, @@ -1695,8 +1696,51 @@ impl Study { next_trial_id: self.next_trial_id.load(Ordering::Relaxed), 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 Study { + /// 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( + &self, + n_trials: usize, + checkpoint_interval: usize, + checkpoint_path: impl AsRef, + objective: F, + ) -> crate::Result<()> + where + F: FnMut(&mut Trial) -> core::result::Result, + 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(()) + }) } } diff --git a/tests/serde_tests.rs b/tests/serde_tests.rs index cb0e812..66aa745 100644 --- a/tests/serde_tests.rs +++ b/tests/serde_tests.rs @@ -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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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};