From 0a6f2345a8e5be6b0233dedabd9782888aac933f Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 22:58:38 +0100 Subject: [PATCH] feat: add Storage trait and JSONL journal backend Replace the internal Vec> with a pluggable Storage trait. MemoryStorage is the default (no behavior change for existing users). Behind the `journal` feature flag, JournalStorage persists trials to a JSONL file with fs2 file locking for multi-process safety. - Storage trait with push(), trials_arc(), refresh() methods - MemoryStorage wraps Arc>>> - JournalStorage appends JSON lines with exclusive file locks - Study::with_sampler_and_storage() general constructor - Study::with_journal() convenience constructor (journal feature) - Refresh from storage on create_trial() for multi-process discovery - MSRV bumped to 1.89 --- .github/workflows/ci.yml | 8 +- Cargo.toml | 8 +- optimizer-derive/Cargo.toml | 2 +- src/error.rs | 5 + src/lib.rs | 7 ++ src/storage/journal.rs | 164 +++++++++++++++++++++++++++ src/storage/memory.rs | 47 ++++++++ src/storage/mod.rs | 48 ++++++++ src/study.rs | 219 ++++++++++++++++++++++++++---------- tests/journal_tests.rs | 215 +++++++++++++++++++++++++++++++++++ 10 files changed, 657 insertions(+), 66 deletions(-) create mode 100644 src/storage/journal.rs create mode 100644 src/storage/memory.rs create mode 100644 src/storage/mod.rs create mode 100644 tests/journal_tests.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cd4eb15..2c52ea9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -139,14 +139,14 @@ jobs: fail_ci_if_error: true msrv: - name: MSRV (1.88) + name: MSRV (1.89) runs-on: ubuntu-latest steps: - uses: actions/checkout@v6 - - name: Install Rust 1.88 + - name: Install Rust 1.89 run: | - rustup override set 1.88 - rustup update 1.88 + rustup override set 1.89 + rustup update 1.89 - uses: Swatinem/rust-cache@v2 - name: Check compilation run: cargo check --all-features diff --git a/Cargo.toml b/Cargo.toml index eba6611..5f28c9f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ members = ["optimizer-derive"] name = "optimizer" version = "0.8.1" edition = "2024" -rust-version = "1.88" +rust-version = "1.89" license = "MIT" authors = ["Manuel Raimann = core::result::Result; diff --git a/src/lib.rs b/src/lib.rs index 1187440..b740da7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -233,6 +233,7 @@ pub mod pareto; pub mod pruner; mod rng_util; pub mod sampler; +pub mod storage; mod study; mod trial; mod types; @@ -269,6 +270,9 @@ pub use sampler::random::RandomSampler; #[cfg(feature = "sobol")] pub use sampler::sobol::SobolSampler; pub use sampler::tpe::TpeSampler; +#[cfg(feature = "journal")] +pub use storage::JournalStorage; +pub use storage::{MemoryStorage, Storage}; pub use study::Study; #[cfg(feature = "serde")] pub use study::StudySnapshot; @@ -314,6 +318,9 @@ pub mod prelude { #[cfg(feature = "sobol")] pub use crate::sampler::sobol::SobolSampler; pub use crate::sampler::tpe::TpeSampler; + #[cfg(feature = "journal")] + pub use crate::storage::JournalStorage; + pub use crate::storage::{MemoryStorage, Storage}; pub use crate::study::Study; #[cfg(feature = "serde")] pub use crate::study::StudySnapshot; diff --git a/src/storage/journal.rs b/src/storage/journal.rs new file mode 100644 index 0000000..6b0fb87 --- /dev/null +++ b/src/storage/journal.rs @@ -0,0 +1,164 @@ +//! JSONL-based journal storage backend. + +use core::marker::PhantomData; +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, Write}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use fs2::FileExt; +use parking_lot::{Mutex, RwLock}; +use serde::Serialize; +use serde::de::DeserializeOwned; + +use super::{MemoryStorage, Storage}; +use crate::sampler::CompletedTrial; + +/// A storage backend that appends completed trials as JSON lines to a file. +/// +/// Trials are kept in memory for fast read access and simultaneously +/// persisted to a JSONL file. Multiple processes can safely share +/// the same file: writes use an exclusive file lock, reads use a +/// shared file lock. +/// +/// The type parameter `V` is the objective value type (typically `f64`). +/// It must be serializable so that trials can be written to disk. +/// +/// # Examples +/// +/// ```no_run +/// use optimizer::storage::JournalStorage; +/// +/// let storage: JournalStorage = JournalStorage::new("trials.jsonl"); +/// ``` +pub struct JournalStorage { + memory: MemoryStorage, + path: PathBuf, + /// Serialise in-process writes so we only hold the file lock briefly. + write_lock: Mutex<()>, + _marker: PhantomData, +} + +impl JournalStorage { + /// Creates a new journal storage that writes to the given path. + /// + /// The file does not need to exist yet — it will be created on the + /// first write. Existing trials in the file are **not** loaded + /// until [`refresh`](Storage::refresh) is called (which happens + /// automatically at the start of each trial via the [`Study`](crate::Study)). + /// + /// To pre-load existing trials, use [`JournalStorage::open`]. + #[must_use] + pub fn new(path: impl AsRef) -> Self { + Self { + memory: MemoryStorage::new(), + path: path.as_ref().to_path_buf(), + write_lock: Mutex::new(()), + _marker: PhantomData, + } + } + + /// Opens an existing journal file and loads all stored trials. + /// + /// If the file does not exist, returns an empty storage (no error). + /// + /// # Errors + /// + /// Returns a [`Storage`](crate::Error::Storage) error if the file + /// exists but cannot be read or parsed. + pub fn open(path: impl AsRef) -> crate::Result { + let path = path.as_ref().to_path_buf(); + let trials = load_trials_from_file(&path)?; + Ok(Self { + memory: MemoryStorage::with_trials(trials), + path, + write_lock: Mutex::new(()), + _marker: PhantomData, + }) + } + + /// Append a single trial to the JSONL file (best-effort). + fn write_to_file(&self, trial: &CompletedTrial) -> crate::Result<()> { + let _guard = self.write_lock.lock(); + + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&self.path) + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + file.lock_exclusive() + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + let line = + serde_json::to_string(trial).map_err(|e| crate::Error::Storage(e.to_string()))?; + + writeln!(file, "{line}").map_err(|e| crate::Error::Storage(e.to_string()))?; + file.flush() + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + file.unlock() + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + Ok(()) + } +} + +impl Storage for JournalStorage { + fn push(&self, trial: CompletedTrial) { + // Best-effort persist; the trial stays in memory regardless. + let _ = self.write_to_file(&trial); + self.memory.push(trial); + } + + fn trials_arc(&self) -> &Arc>>> { + self.memory.trials_arc() + } + + fn refresh(&self) -> bool { + let Ok(loaded) = load_trials_from_file::(&self.path) else { + return false; + }; + let mut guard = self.memory.trials_arc().write(); + if loaded.len() > guard.len() { + *guard = loaded; + true + } else { + false + } + } +} + +/// Read all trials from a JSONL file. Returns an empty vec if the +/// file does not exist. +fn load_trials_from_file( + path: &Path, +) -> crate::Result>> { + let file = match File::open(path) { + Ok(f) => f, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) => return Err(crate::Error::Storage(e.to_string())), + }; + + file.lock_shared() + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + let reader = BufReader::new(&file); + let mut trials = Vec::new(); + + for line in reader.lines() { + let line = line.map_err(|e| crate::Error::Storage(e.to_string()))?; + let line = line.trim(); + if line.is_empty() { + continue; + } + let trial: CompletedTrial = + serde_json::from_str(line).map_err(|e| crate::Error::Storage(e.to_string()))?; + trials.push(trial); + } + + file.unlock() + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + Ok(trials) +} diff --git a/src/storage/memory.rs b/src/storage/memory.rs new file mode 100644 index 0000000..d511faf --- /dev/null +++ b/src/storage/memory.rs @@ -0,0 +1,47 @@ +use std::sync::Arc; + +use parking_lot::RwLock; + +use super::Storage; +use crate::sampler::CompletedTrial; + +/// In-memory trial storage (the default). +/// +/// This is a thin wrapper around `Arc>>>`. +pub struct MemoryStorage { + trials: Arc>>>, +} + +impl MemoryStorage { + /// Creates a new, empty in-memory store. + #[must_use] + pub fn new() -> Self { + Self { + trials: Arc::new(RwLock::new(Vec::new())), + } + } + + /// Creates an in-memory store pre-populated with `trials`. + #[must_use] + pub fn with_trials(trials: Vec>) -> Self { + Self { + trials: Arc::new(RwLock::new(trials)), + } + } +} + +impl Default for MemoryStorage { + fn default() -> Self { + Self::new() + } +} + +impl Storage for MemoryStorage { + fn push(&self, trial: CompletedTrial) { + self.trials.write().push(trial); + } + + fn trials_arc(&self) -> &Arc>>> { + &self.trials + } +} diff --git a/src/storage/mod.rs b/src/storage/mod.rs new file mode 100644 index 0000000..faf2d1c --- /dev/null +++ b/src/storage/mod.rs @@ -0,0 +1,48 @@ +//! Trial storage backends. +//! +//! The [`Storage`] trait defines how completed trials are stored and +//! accessed. [`MemoryStorage`] keeps trials in memory (the default). +//! With the `journal` feature enabled, [`JournalStorage`] appends +//! trials to a JSONL file with file-level locking so multiple +//! processes can safely share state. + +#[cfg(feature = "journal")] +mod journal; + +use std::sync::Arc; + +#[cfg(feature = "journal")] +pub use journal::JournalStorage; +use parking_lot::RwLock; + +mod memory; +pub use memory::MemoryStorage; + +use crate::sampler::CompletedTrial; + +/// Trait for storing and retrieving completed trials. +/// +/// Every [`Study`](crate::Study) owns an `Arc>`. The +/// default implementation is [`MemoryStorage`], which keeps trials in +/// a plain `Vec` behind a read-write lock. +/// +/// Implementations must be safe to use from multiple threads. +pub trait Storage: Send + Sync { + /// Append a completed trial to the store. + fn push(&self, trial: CompletedTrial); + + /// Return a reference to the in-memory trial buffer. + /// + /// All implementations must maintain an `Arc>>` that + /// reflects the current set of trials. Callers may acquire a read + /// lock for efficient, allocation-free access. + fn trials_arc(&self) -> &Arc>>>; + + /// Reload from an external source (e.g. a file written by another + /// process). Returns `true` if the in-memory buffer was updated. + /// + /// The default implementation is a no-op that returns `false`. + fn refresh(&self) -> bool { + false + } +} diff --git a/src/study.rs b/src/study.rs index 805bfb8..d6784cc 100644 --- a/src/study.rs +++ b/src/study.rs @@ -49,8 +49,8 @@ where sampler: Arc, /// The pruner used to decide whether to stop trials early. pruner: Arc, - /// Completed trials (wrapped in Arc for sharing with Trial). - completed_trials: Arc>>>, + /// Trial storage backend (default: [`MemoryStorage`](crate::storage::MemoryStorage)). + storage: Arc>, /// Counter for generating unique trial IDs. next_trial_id: AtomicU64, /// Optional factory for creating sampler-aware trials. @@ -84,7 +84,7 @@ where #[must_use] pub fn new(direction: Direction) -> Self where - V: 'static, + V: Send + Sync + 'static, { Self::with_sampler(direction, RandomSampler::new()) } @@ -109,7 +109,7 @@ where #[must_use] pub fn minimize(sampler: impl Sampler + 'static) -> Self where - V: 'static, + V: Send + Sync + 'static, { Self::with_sampler(Direction::Minimize, sampler) } @@ -134,7 +134,7 @@ where #[must_use] pub fn maximize(sampler: impl Sampler + 'static) -> Self where - V: 'static, + V: Send + Sync + 'static, { Self::with_sampler(Direction::Maximize, sampler) } @@ -158,40 +158,28 @@ where /// ``` pub fn with_sampler(direction: Direction, sampler: impl Sampler + 'static) -> Self where - V: 'static, + V: Send + Sync + 'static, { - let sampler: Arc = Arc::new(sampler); - let completed_trials = Arc::new(RwLock::new(Vec::new())); - - let pruner: Arc = Arc::new(NopPruner); - - // For Study, set up a trial factory that provides sampler integration. - // This uses Any downcasting to check at runtime whether V = f64. - let trial_factory = Self::make_trial_factory(&sampler, &completed_trials, &pruner); - - Self { + Self::with_sampler_and_storage( direction, sampler, - pruner, - completed_trials, - next_trial_id: AtomicU64::new(0), - trial_factory, - enqueued_params: Arc::new(Mutex::new(VecDeque::new())), - } + crate::storage::MemoryStorage::::new(), + ) } /// Builds a trial factory for sampler integration when `V = f64`. fn make_trial_factory( sampler: &Arc, - completed_trials: &Arc>>>, + storage: &Arc>, pruner: &Arc, ) -> Option Trial + Send + Sync>> where V: 'static, { - // Try to downcast the completed_trials Arc to the f64 specialization. + // Try to downcast the storage's trial buffer to the f64 specialization. // This succeeds only when V = f64, enabling automatic sampler integration. - let any_ref: &dyn Any = completed_trials; + let trials_arc = storage.trials_arc(); + let any_ref: &dyn Any = trials_arc; let f64_trials: Option<&Arc>>>> = any_ref.downcast_ref(); f64_trials.map(|trials| { @@ -210,6 +198,42 @@ where }) } + /// Creates a study with a custom sampler and storage backend. + /// + /// This is the most general constructor — all other constructors + /// delegate to this one. + pub fn with_sampler_and_storage( + direction: Direction, + sampler: impl Sampler + 'static, + storage: impl crate::storage::Storage + 'static, + ) -> Self + where + V: 'static, + { + let sampler: Arc = Arc::new(sampler); + let pruner: Arc = Arc::new(NopPruner); + let storage: Arc> = Arc::new(storage); + let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner); + + let next_id = storage + .trials_arc() + .read() + .iter() + .map(|t| t.id) + .max() + .map_or(0, |id| id + 1); + + Self { + direction, + sampler, + pruner, + storage, + next_trial_id: AtomicU64::new(next_id), + trial_factory, + enqueued_params: Arc::new(Mutex::new(VecDeque::new())), + } + } + /// Returns the optimization direction. pub fn direction(&self) -> Direction { self.direction @@ -254,18 +278,19 @@ where pruner: impl Pruner + 'static, ) -> Self where - V: 'static, + V: Send + Sync + 'static, { let sampler: Arc = Arc::new(sampler); let pruner: Arc = Arc::new(pruner); - let completed_trials = Arc::new(RwLock::new(Vec::new())); - let trial_factory = Self::make_trial_factory(&sampler, &completed_trials, &pruner); + let storage: Arc> = + Arc::new(crate::storage::MemoryStorage::::new()); + let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner); Self { direction, sampler, pruner, - completed_trials, + storage, next_trial_id: AtomicU64::new(0), trial_factory, enqueued_params: Arc::new(Mutex::new(VecDeque::new())), @@ -277,8 +302,7 @@ where V: 'static, { self.sampler = Arc::new(sampler); - self.trial_factory = - Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner); + self.trial_factory = Self::make_trial_factory(&self.sampler, &self.storage, &self.pruner); } /// Sets a new pruner for the study. @@ -291,8 +315,7 @@ where V: 'static, { self.pruner = Arc::new(pruner); - self.trial_factory = - Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner); + self.trial_factory = Self::make_trial_factory(&self.sampler, &self.storage, &self.pruner); } /// Returns a reference to the study's pruner. @@ -401,6 +424,13 @@ where /// assert_eq!(trial2.id(), 1); /// ``` pub fn create_trial(&self) -> Trial { + if self.storage.refresh() { + let trials = self.storage.trials_arc().read(); + if let Some(max_id) = trials.iter().map(|t| t.id).max() { + self.next_trial_id.fetch_max(max_id + 1, Ordering::SeqCst); + } + } + let id = self.next_trial_id(); let mut trial = if let Some(factory) = &self.trial_factory { factory(id) @@ -455,7 +485,8 @@ where ); completed.state = TrialState::Complete; completed.constraints = trial.constraint_values().to_vec(); - self.completed_trials.write().push(completed); + + self.storage.push(completed); } /// Records a failed trial with an error message. @@ -565,7 +596,8 @@ where ); completed.state = TrialState::Pruned; completed.constraints = trial.constraint_values().to_vec(); - self.completed_trials.write().push(completed); + + self.storage.push(completed); } /// Returns an iterator over all completed trials. @@ -596,7 +628,7 @@ where where V: Clone, { - self.completed_trials.read().clone() + self.storage.trials_arc().read().clone() } /// Returns the number of completed trials. @@ -619,12 +651,13 @@ where /// assert_eq!(study.n_trials(), 1); /// ``` pub fn n_trials(&self) -> usize { - self.completed_trials.read().len() + self.storage.trials_arc().read().len() } /// Returns the number of pruned trials. pub fn n_pruned_trials(&self) -> usize { - self.completed_trials + self.storage + .trials_arc() .read() .iter() .filter(|t| t.state == TrialState::Pruned) @@ -703,7 +736,7 @@ where where V: Clone, { - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); let direction = self.direction; let best = trials @@ -767,7 +800,7 @@ where where V: Clone, { - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); let direction = self.direction; let mut completed: Vec<_> = trials .iter() @@ -848,7 +881,7 @@ where #[cfg(feature = "tracing")] { tracing::info!(trial_id, "trial completed"); - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); if trials .iter() .filter(|t| t.state == TrialState::Complete) @@ -876,7 +909,8 @@ where // Return error if no trials completed successfully let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -975,7 +1009,8 @@ where // Return error if no trials completed successfully let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1101,7 +1136,8 @@ where // Return error if no trials completed successfully let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1197,7 +1233,7 @@ where #[cfg(feature = "tracing")] { tracing::info!(trial_id, "trial completed"); - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); if trials .iter() .filter(|t| t.state == TrialState::Complete) @@ -1210,7 +1246,7 @@ where } // Get the just-completed trial for the callback - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); let Some(completed) = trials.last() else { return Err(crate::Error::Internal( "completed trial not found after adding", @@ -1243,7 +1279,8 @@ where // Return error if no trials completed successfully let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1328,7 +1365,8 @@ where } let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1420,7 +1458,7 @@ where #[cfg(feature = "tracing")] { tracing::info!(trial_id, "trial completed"); - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); if trials .iter() .filter(|t| t.state == TrialState::Complete) @@ -1432,7 +1470,7 @@ where } } - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); let Some(completed) = trials.last() else { return Err(crate::Error::Internal( "completed trial not found after adding", @@ -1461,7 +1499,8 @@ where } let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1519,7 +1558,8 @@ where } let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1608,7 +1648,8 @@ where } let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1716,7 +1757,8 @@ where // Return error if no trials completed successfully let has_complete = self - .completed_trials + .storage + .trials_arc() .read() .iter() .any(|t| t.state == TrialState::Complete); @@ -1746,7 +1788,7 @@ where pub fn to_csv(&self, mut writer: impl std::io::Write) -> std::io::Result<()> { use std::collections::BTreeMap; - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); // Collect all unique parameter labels (sorted for deterministic column order). let mut param_columns: BTreeMap = BTreeMap::new(); @@ -1873,7 +1915,7 @@ where pub fn summary(&self) -> String { use fmt::Write; - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); let n_complete = trials .iter() .filter(|t| t.state == TrialState::Complete) @@ -1967,7 +2009,7 @@ where use crate::param::ParamValue; use crate::types::TrialState; - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); let complete: Vec<_> = trials .iter() .filter(|t| t.state == TrialState::Complete) @@ -2064,7 +2106,7 @@ where use crate::param::ParamValue; use crate::types::TrialState; - let trials = self.completed_trials.read(); + let trials = self.storage.trials_arc().read(); let complete: Vec<_> = trials .iter() .filter(|t| t.state == TrialState::Complete) @@ -2258,6 +2300,63 @@ impl Study { } } +impl Study { + /// Creates a study with a custom sampler, pruner, and storage backend. + pub fn with_sampler_pruner_and_storage( + direction: Direction, + sampler: impl Sampler + 'static, + pruner: impl Pruner + 'static, + storage: impl crate::storage::Storage + 'static, + ) -> Self { + let sampler: Arc = Arc::new(sampler); + let pruner: Arc = Arc::new(pruner); + let storage: Arc> = Arc::new(storage); + let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner); + + let next_id = storage + .trials_arc() + .read() + .iter() + .map(|t| t.id) + .max() + .map_or(0, |id| id + 1); + + Self { + direction, + sampler, + pruner, + storage, + next_trial_id: AtomicU64::new(next_id), + trial_factory, + enqueued_params: Arc::new(Mutex::new(VecDeque::new())), + } + } +} + +#[cfg(feature = "journal")] +impl Study +where + V: PartialOrd + Send + Sync + serde::Serialize + serde::de::DeserializeOwned + 'static, +{ + /// Creates a study backed by a JSONL journal file. + /// + /// Any existing trials in the file are loaded into memory and the + /// trial ID counter is set to one past the highest stored ID. New + /// trials are written through to the file on completion. + /// + /// # Errors + /// + /// Returns a [`Storage`](crate::Error::Storage) error if loading fails. + pub fn with_journal( + direction: Direction, + sampler: impl Sampler + 'static, + path: impl AsRef, + ) -> crate::Result { + let storage = crate::storage::JournalStorage::::open(path)?; + Ok(Self::with_sampler_and_storage(direction, sampler, storage)) + } +} + #[cfg(feature = "visualization")] impl Study { /// Generates an HTML report with interactive Plotly.js charts. @@ -2382,7 +2481,7 @@ impl Study { } #[cfg(feature = "serde")] -impl Study { +impl Study { /// Loads a study from a JSON file. /// /// The loaded study uses a `RandomSampler` by default. Call @@ -2397,7 +2496,7 @@ impl Study { let snapshot: StudySnapshot = serde_json::from_reader(file) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; let study = Study::new(snapshot.direction); - *study.completed_trials.write() = snapshot.trials; + *study.storage.trials_arc().write() = snapshot.trials; study .next_trial_id .store(snapshot.next_trial_id, Ordering::Relaxed); diff --git a/tests/journal_tests.rs b/tests/journal_tests.rs new file mode 100644 index 0000000..cabc453 --- /dev/null +++ b/tests/journal_tests.rs @@ -0,0 +1,215 @@ +//! Integration tests for the journal storage backend. + +use std::collections::HashMap; +use std::sync::Arc; + +use optimizer::parameter::{FloatParam, Parameter}; +use optimizer::sampler::CompletedTrial; +use optimizer::sampler::random::RandomSampler; +use optimizer::storage::{JournalStorage, Storage}; +use optimizer::{Direction, Study}; + +fn temp_path() -> std::path::PathBuf { + let mut path = std::env::temp_dir(); + path.push(format!( + "optimizer_journal_test_{}.jsonl", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + path +} + +fn sample_trial(id: u64, value: f64) -> CompletedTrial { + CompletedTrial::new(id, HashMap::new(), HashMap::new(), HashMap::new(), value) +} + +#[test] +fn roundtrip_single_trial() { + let path = temp_path(); + let storage = JournalStorage::new(&path); + + storage.push(sample_trial(0, 42.0)); + + let loaded = storage.trials_arc().read().clone(); + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].id, 0); + assert_eq!(loaded[0].value, 42.0); + + // Also verify via a fresh open from disk + let storage2 = JournalStorage::::open(&path).unwrap(); + let loaded2 = storage2.trials_arc().read().clone(); + assert_eq!(loaded2.len(), 1); + assert_eq!(loaded2[0].value, 42.0); + + std::fs::remove_file(&path).ok(); +} + +#[test] +fn append_multiple_trials() { + let path = temp_path(); + let storage = JournalStorage::new(&path); + + for i in 0..5 { + storage.push(sample_trial(i, i as f64)); + } + + // Reload from disk + let storage2 = JournalStorage::::open(&path).unwrap(); + let loaded = storage2.trials_arc().read().clone(); + assert_eq!(loaded.len(), 5); + for (i, trial) in loaded.iter().enumerate() { + assert_eq!(trial.id, i as u64); + assert_eq!(trial.value, i as f64); + } + + std::fs::remove_file(&path).ok(); +} + +#[test] +fn missing_file_returns_empty() { + let path = temp_path(); + let storage = JournalStorage::::open(&path).unwrap(); + + let loaded = storage.trials_arc().read().clone(); + assert!(loaded.is_empty()); +} + +#[test] +fn concurrent_writes() { + let path = temp_path(); + let storage = Arc::new(JournalStorage::new(&path)); + + let mut handles = Vec::new(); + for thread_id in 0..4u64 { + let s = Arc::clone(&storage); + handles.push(std::thread::spawn(move || { + for i in 0..25u64 { + let id = thread_id * 25 + i; + s.push(sample_trial(id, id as f64)); + } + })); + } + for h in handles { + h.join().unwrap(); + } + + // Reload from disk to verify persistence + let storage2 = JournalStorage::::open(&path).unwrap(); + let loaded = storage2.trials_arc().read().clone(); + assert_eq!(loaded.len(), 100); + + // Verify all IDs are present (order may vary) + let mut ids: Vec = loaded.iter().map(|t| t.id).collect(); + ids.sort(); + assert_eq!(ids, (0..100).collect::>()); + + std::fs::remove_file(&path).ok(); +} + +#[test] +fn study_with_journal_integration() { + let path = temp_path(); + let x = FloatParam::new(-10.0, 10.0); + + // First "process": run some trials + { + let study = + Study::with_journal(Direction::Minimize, RandomSampler::with_seed(1), &path).unwrap(); + study + .optimize(5, |trial| { + let val = x.suggest(trial)?; + Ok::<_, optimizer::Error>(val * val) + }) + .unwrap(); + assert_eq!(study.n_trials(), 5); + } + + // Second "process": loads the same file, sees existing trials + let study2 = + Study::with_journal(Direction::Minimize, RandomSampler::with_seed(2), &path).unwrap(); + assert_eq!(study2.n_trials(), 5); + + // Continue optimizing + study2 + .optimize(5, |trial| { + let val = x.suggest(trial)?; + Ok::<_, optimizer::Error>(val * val) + }) + .unwrap(); + assert_eq!(study2.n_trials(), 10); + + // Verify all 10 written to disk + let storage3 = JournalStorage::::open(&path).unwrap(); + let loaded = storage3.trials_arc().read().clone(); + assert_eq!(loaded.len(), 10); + + std::fs::remove_file(&path).ok(); +} + +#[test] +fn ids_are_unique_after_reload() { + let path = temp_path(); + + // First batch + { + let study = + Study::with_journal(Direction::Minimize, RandomSampler::with_seed(1), &path).unwrap(); + study + .optimize(3, |trial| { + let _ = FloatParam::new(0.0, 1.0).suggest(trial)?; + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); + } + + // Second batch — IDs should continue from 3 + let study = + Study::with_journal(Direction::Minimize, RandomSampler::with_seed(2), &path).unwrap(); + study + .optimize(3, |trial| { + let _ = FloatParam::new(0.0, 1.0).suggest(trial)?; + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); + + let all = study.trials(); + let mut ids: Vec = all.iter().map(|t| t.id).collect(); + ids.sort(); + // All 6 IDs should be unique + ids.dedup(); + assert_eq!(ids.len(), 6); + + std::fs::remove_file(&path).ok(); +} + +#[test] +fn pruned_trials_are_stored() { + let path = temp_path(); + let study = + Study::with_journal(Direction::Minimize, RandomSampler::with_seed(1), &path).unwrap(); + + // Complete one, prune one + let x = FloatParam::new(0.0, 1.0); + study + .optimize(3, |trial| { + let _ = x.suggest(trial)?; + if trial.id() == 1 { + Err(optimizer::TrialPruned)?; + } + Ok::<_, optimizer::Error>(1.0) + }) + .unwrap(); + + let storage2 = JournalStorage::::open(&path).unwrap(); + let loaded = storage2.trials_arc().read().clone(); + assert_eq!(loaded.len(), 3); + assert!( + loaded + .iter() + .any(|t| t.state == optimizer::TrialState::Pruned) + ); + + std::fs::remove_file(&path).ok(); +}