diff --git a/Cargo.toml b/Cargo.toml index a62c23c..6b9ae89 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ tracing = { version = "0.1.29", optional = true } sobol_burley = { version = "0.5", optional = true } nalgebra = { version = "0.34", optional = true } fs2 = { version = "0.4", optional = true } +rusqlite = { version = "0.38", features = ["bundled"], optional = true } [features] default = [] @@ -34,6 +35,7 @@ async = ["dep:tokio"] derive = ["dep:optimizer-derive"] serde = ["dep:serde", "dep:serde_json"] journal = ["dep:fs2", "serde"] +sqlite = ["dep:rusqlite", "serde"] tracing = ["dep:tracing"] sobol = ["dep:sobol_burley"] cma-es = ["dep:nalgebra"] @@ -74,3 +76,7 @@ path = "examples/visualization_demo.rs" [[test]] name = "journal_tests" required-features = ["journal"] + +[[test]] +name = "sqlite_tests" +required-features = ["sqlite"] \ No newline at end of file diff --git a/src/error.rs b/src/error.rs index 130534e..1a917bc 100644 --- a/src/error.rs +++ b/src/error.rs @@ -95,7 +95,7 @@ pub enum Error { TaskError(String), /// Returned when a storage operation fails. - #[cfg(feature = "journal")] + #[cfg(any(feature = "journal", feature = "sqlite"))] #[error("storage error: {0}")] Storage(String), } diff --git a/src/lib.rs b/src/lib.rs index b3c22f0..93eab99 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -269,8 +269,10 @@ pub use sampler::sobol::SobolSampler; pub use sampler::tpe::TpeSampler; #[cfg(feature = "journal")] pub use storage::JournalStorage; +#[cfg(feature = "sqlite")] +pub use storage::SqliteStorage; pub use storage::{MemoryStorage, Storage}; -pub use study::Study; +pub use study::{Study, StudyBuilder}; #[cfg(feature = "serde")] pub use study::StudySnapshot; pub use trial::{AttrValue, Trial}; @@ -315,8 +317,10 @@ pub mod prelude { pub use crate::sampler::tpe::TpeSampler; #[cfg(feature = "journal")] pub use crate::storage::JournalStorage; + #[cfg(feature = "sqlite")] + pub use crate::storage::SqliteStorage; pub use crate::storage::{MemoryStorage, Storage}; - pub use crate::study::Study; + pub use crate::study::{Study, StudyBuilder}; #[cfg(feature = "serde")] pub use crate::study::StudySnapshot; pub use crate::trial::{AttrValue, Trial}; diff --git a/src/storage/mod.rs b/src/storage/mod.rs index faf2d1c..65b5648 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -8,12 +8,16 @@ #[cfg(feature = "journal")] mod journal; +#[cfg(feature = "sqlite")] +mod sqlite; use std::sync::Arc; #[cfg(feature = "journal")] pub use journal::JournalStorage; use parking_lot::RwLock; +#[cfg(feature = "sqlite")] +pub use sqlite::SqliteStorage; mod memory; pub use memory::MemoryStorage; diff --git a/src/storage/sqlite.rs b/src/storage/sqlite.rs new file mode 100644 index 0000000..b1378fe --- /dev/null +++ b/src/storage/sqlite.rs @@ -0,0 +1,135 @@ +//! `SQLite`-backed storage backend for multi-process optimization. + +use core::marker::PhantomData; +use std::path::Path; +use std::sync::Arc; + +use parking_lot::{Mutex, RwLock}; +use rusqlite::Connection; +use serde::Serialize; +use serde::de::DeserializeOwned; + +use super::{MemoryStorage, Storage}; +use crate::sampler::CompletedTrial; + +/// A storage backend that persists completed trials in a `SQLite` database. +/// +/// Uses WAL mode for concurrent readers and a single writer, making it +/// suitable for single-machine multi-process optimization. Compared to +/// [`JournalStorage`](super::JournalStorage), `SQLite` provides proper +/// ACID transactions and better concurrent access. +/// +/// 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::SqliteStorage; +/// +/// let storage: SqliteStorage = SqliteStorage::new("trials.db").unwrap(); +/// ``` +pub struct SqliteStorage { + memory: MemoryStorage, + conn: Mutex, + _marker: PhantomData, +} + +impl SqliteStorage { + /// Creates a new `SQLite` storage at the given path. + /// + /// The database file is created if it does not exist. Any trials + /// already stored in the database are loaded into memory. + /// + /// WAL mode is enabled automatically for better concurrency. + /// + /// # Errors + /// + /// Returns a [`Storage`](crate::Error::Storage) error if the + /// database cannot be opened or the schema cannot be created. + pub fn new(path: impl AsRef) -> crate::Result { + let conn = Connection::open(path).map_err(|e| crate::Error::Storage(e.to_string()))?; + + // WAL mode: concurrent readers, single writer. + conn.pragma_update(None, "journal_mode", "WAL") + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS trials ( + trial_id INTEGER PRIMARY KEY, + data TEXT NOT NULL + );", + ) + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + let trials = load_all(&conn)?; + + Ok(Self { + memory: MemoryStorage::with_trials(trials), + conn: Mutex::new(conn), + _marker: PhantomData, + }) + } + + /// Persist a single trial into the database. + fn write_trial(&self, trial: &CompletedTrial) -> crate::Result<()> { + let data = + serde_json::to_string(trial).map_err(|e| crate::Error::Storage(e.to_string()))?; + + let conn = self.conn.lock(); + conn.execute( + "INSERT OR REPLACE INTO trials (trial_id, data) VALUES (?1, ?2)", + rusqlite::params![i64::try_from(trial.id).unwrap_or(i64::MAX), data], + ) + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + Ok(()) + } +} + +impl Storage for SqliteStorage { + fn push(&self, trial: CompletedTrial) { + // Best-effort persist; the trial stays in memory regardless. + let _ = self.write_trial(&trial); + self.memory.push(trial); + } + + fn trials_arc(&self) -> &Arc>>> { + self.memory.trials_arc() + } + + fn refresh(&self) -> bool { + let conn = self.conn.lock(); + let Ok(loaded) = load_all::(&conn) else { + return false; + }; + let mut guard = self.memory.trials_arc().write(); + if loaded.len() > guard.len() { + *guard = loaded; + true + } else { + false + } + } +} + +/// Load every trial from the database, ordered by id. +fn load_all(conn: &Connection) -> crate::Result>> { + let mut stmt = conn + .prepare("SELECT data FROM trials ORDER BY trial_id") + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + let rows = stmt + .query_map([], |row| row.get::<_, String>(0)) + .map_err(|e| crate::Error::Storage(e.to_string()))?; + + let mut trials = Vec::new(); + for row in rows { + let data = row.map_err(|e| crate::Error::Storage(e.to_string()))?; + let trial: CompletedTrial = + serde_json::from_str(&data).map_err(|e| crate::Error::Storage(e.to_string()))?; + trials.push(trial); + } + + Ok(trials) +} diff --git a/src/study.rs b/src/study.rs index 63c8e68..77b6eaf 100644 --- a/src/study.rs +++ b/src/study.rs @@ -2355,6 +2355,34 @@ where } } +#[cfg(feature = "sqlite")] +impl Study +where + V: PartialOrd + Send + Sync + serde::Serialize + serde::de::DeserializeOwned + 'static, +{ + /// Creates a study backed by a `SQLite` database. + /// + /// Any existing trials in the database 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 database on completion. + /// + /// Uses WAL mode for concurrent readers, making it suitable for + /// single-machine multi-process optimization. + /// + /// # Errors + /// + /// Returns a [`Storage`](crate::Error::Storage) error if the + /// database cannot be opened. + pub fn with_sqlite( + direction: Direction, + sampler: impl Sampler + 'static, + path: impl AsRef, + ) -> crate::Result { + let storage = crate::storage::SqliteStorage::::new(path)?; + Ok(Self::with_sampler_and_storage(direction, sampler, storage)) + } +} + impl Study { /// Generates an HTML report with interactive Plotly.js charts. /// diff --git a/tests/sqlite_tests.rs b/tests/sqlite_tests.rs new file mode 100644 index 0000000..792538b --- /dev/null +++ b/tests/sqlite_tests.rs @@ -0,0 +1,243 @@ +//! Integration tests for the SQLite 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::{SqliteStorage, Storage}; +use optimizer::{Direction, Study}; + +fn temp_path() -> std::path::PathBuf { + let mut path = std::env::temp_dir(); + path.push(format!( + "optimizer_sqlite_test_{}.db", + 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 = SqliteStorage::new(&path).unwrap(); + + 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); + + // Verify via a fresh open from disk + let storage2 = SqliteStorage::::new(&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 = SqliteStorage::new(&path).unwrap(); + + for i in 0..5 { + storage.push(sample_trial(i, i as f64)); + } + + // Reload from disk + let storage2 = SqliteStorage::::new(&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 concurrent_writes() { + let path = temp_path(); + let storage = Arc::new(SqliteStorage::new(&path).unwrap()); + + 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 = SqliteStorage::::new(&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 refresh_picks_up_external_writes() { + let path = temp_path(); + let storage1 = SqliteStorage::new(&path).unwrap(); + let storage2 = SqliteStorage::::new(&path).unwrap(); + + // Write via storage1 + storage1.push(sample_trial(0, 1.0)); + storage1.push(sample_trial(1, 2.0)); + + // storage2 doesn't see them yet in memory + assert_eq!(storage2.trials_arc().read().len(), 0); + + // After refresh, storage2 picks them up + assert!(storage2.refresh()); + assert_eq!(storage2.trials_arc().read().len(), 2); + + // No-op refresh returns false + assert!(!storage2.refresh()); + + std::fs::remove_file(&path).ok(); +} + +#[test] +fn study_with_sqlite_integration() { + let path = temp_path(); + let x = FloatParam::new(-10.0, 10.0); + + // First "process": run some trials + { + let study = + Study::with_sqlite(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_sqlite(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 = SqliteStorage::::new(&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_sqlite(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_sqlite(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(); + 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_sqlite(Direction::Minimize, RandomSampler::with_seed(1), &path).unwrap(); + + 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 = SqliteStorage::::new(&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(); +} + +#[test] +fn handles_many_trials() { + let path = temp_path(); + let storage = SqliteStorage::new(&path).unwrap(); + + for i in 0..1000 { + storage.push(sample_trial(i, i as f64)); + } + + let storage2 = SqliteStorage::::new(&path).unwrap(); + let loaded = storage2.trials_arc().read().clone(); + assert_eq!(loaded.len(), 1000); + + std::fs::remove_file(&path).ok(); +}