feat: add Storage trait and JSONL journal backend
Replace the internal Vec<CompletedTrial<V>> with a pluggable Storage<V> 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<V> trait with push(), trials_arc(), refresh() methods - MemoryStorage<V> wraps Arc<RwLock<Vec<CompletedTrial<V>>>> - JournalStorage<V> 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
This commit is contained in:
@@ -139,14 +139,14 @@ jobs:
|
|||||||
fail_ci_if_error: true
|
fail_ci_if_error: true
|
||||||
|
|
||||||
msrv:
|
msrv:
|
||||||
name: MSRV (1.88)
|
name: MSRV (1.89)
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v6
|
- uses: actions/checkout@v6
|
||||||
- name: Install Rust 1.88
|
- name: Install Rust 1.89
|
||||||
run: |
|
run: |
|
||||||
rustup override set 1.88
|
rustup override set 1.89
|
||||||
rustup update 1.88
|
rustup update 1.89
|
||||||
- uses: Swatinem/rust-cache@v2
|
- uses: Swatinem/rust-cache@v2
|
||||||
- name: Check compilation
|
- name: Check compilation
|
||||||
run: cargo check --all-features
|
run: cargo check --all-features
|
||||||
|
|||||||
+7
-1
@@ -5,7 +5,7 @@ members = ["optimizer-derive"]
|
|||||||
name = "optimizer"
|
name = "optimizer"
|
||||||
version = "0.8.1"
|
version = "0.8.1"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
rust-version = "1.88"
|
rust-version = "1.89"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
authors = ["Manuel Raimann <raimannma@outlook.de"]
|
authors = ["Manuel Raimann <raimannma@outlook.de"]
|
||||||
description = "A Rust library for optimization algorithms."
|
description = "A Rust library for optimization algorithms."
|
||||||
@@ -26,12 +26,14 @@ serde_json = { version = "1", optional = true }
|
|||||||
tracing = { version = "0.1.29", optional = true }
|
tracing = { version = "0.1.29", optional = true }
|
||||||
sobol_burley = { version = "0.5", optional = true }
|
sobol_burley = { version = "0.5", optional = true }
|
||||||
nalgebra = { version = "0.34", optional = true }
|
nalgebra = { version = "0.34", optional = true }
|
||||||
|
fs2 = { version = "0.4", optional = true }
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = []
|
default = []
|
||||||
async = ["dep:tokio"]
|
async = ["dep:tokio"]
|
||||||
derive = ["dep:optimizer-derive"]
|
derive = ["dep:optimizer-derive"]
|
||||||
serde = ["dep:serde", "dep:serde_json"]
|
serde = ["dep:serde", "dep:serde_json"]
|
||||||
|
journal = ["dep:fs2", "serde"]
|
||||||
tracing = ["dep:tracing"]
|
tracing = ["dep:tracing"]
|
||||||
sobol = ["dep:sobol_burley"]
|
sobol = ["dep:sobol_burley"]
|
||||||
cma-es = ["dep:nalgebra"]
|
cma-es = ["dep:nalgebra"]
|
||||||
@@ -71,3 +73,7 @@ required-features = ["derive"]
|
|||||||
name = "visualization_demo"
|
name = "visualization_demo"
|
||||||
path = "examples/visualization_demo.rs"
|
path = "examples/visualization_demo.rs"
|
||||||
required-features = ["visualization"]
|
required-features = ["visualization"]
|
||||||
|
|
||||||
|
[[test]]
|
||||||
|
name = "journal_tests"
|
||||||
|
required-features = ["journal"]
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
name = "optimizer-derive"
|
name = "optimizer-derive"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
rust-version = "1.88"
|
rust-version = "1.89"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
description = "Derive macros for the optimizer crate"
|
description = "Derive macros for the optimizer crate"
|
||||||
repository = "https://github.com/raimannma/rust-optimizer"
|
repository = "https://github.com/raimannma/rust-optimizer"
|
||||||
|
|||||||
@@ -93,6 +93,11 @@ pub enum Error {
|
|||||||
#[cfg(feature = "async")]
|
#[cfg(feature = "async")]
|
||||||
#[error("async task error: {0}")]
|
#[error("async task error: {0}")]
|
||||||
TaskError(String),
|
TaskError(String),
|
||||||
|
|
||||||
|
/// Returned when a storage operation fails.
|
||||||
|
#[cfg(feature = "journal")]
|
||||||
|
#[error("storage error: {0}")]
|
||||||
|
Storage(String),
|
||||||
}
|
}
|
||||||
|
|
||||||
pub type Result<T> = core::result::Result<T, Error>;
|
pub type Result<T> = core::result::Result<T, Error>;
|
||||||
|
|||||||
@@ -233,6 +233,7 @@ pub mod pareto;
|
|||||||
pub mod pruner;
|
pub mod pruner;
|
||||||
mod rng_util;
|
mod rng_util;
|
||||||
pub mod sampler;
|
pub mod sampler;
|
||||||
|
pub mod storage;
|
||||||
mod study;
|
mod study;
|
||||||
mod trial;
|
mod trial;
|
||||||
mod types;
|
mod types;
|
||||||
@@ -269,6 +270,9 @@ pub use sampler::random::RandomSampler;
|
|||||||
#[cfg(feature = "sobol")]
|
#[cfg(feature = "sobol")]
|
||||||
pub use sampler::sobol::SobolSampler;
|
pub use sampler::sobol::SobolSampler;
|
||||||
pub use sampler::tpe::TpeSampler;
|
pub use sampler::tpe::TpeSampler;
|
||||||
|
#[cfg(feature = "journal")]
|
||||||
|
pub use storage::JournalStorage;
|
||||||
|
pub use storage::{MemoryStorage, Storage};
|
||||||
pub use study::Study;
|
pub use study::Study;
|
||||||
#[cfg(feature = "serde")]
|
#[cfg(feature = "serde")]
|
||||||
pub use study::StudySnapshot;
|
pub use study::StudySnapshot;
|
||||||
@@ -314,6 +318,9 @@ pub mod prelude {
|
|||||||
#[cfg(feature = "sobol")]
|
#[cfg(feature = "sobol")]
|
||||||
pub use crate::sampler::sobol::SobolSampler;
|
pub use crate::sampler::sobol::SobolSampler;
|
||||||
pub use crate::sampler::tpe::TpeSampler;
|
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;
|
pub use crate::study::Study;
|
||||||
#[cfg(feature = "serde")]
|
#[cfg(feature = "serde")]
|
||||||
pub use crate::study::StudySnapshot;
|
pub use crate::study::StudySnapshot;
|
||||||
|
|||||||
@@ -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<f64> = JournalStorage::new("trials.jsonl");
|
||||||
|
/// ```
|
||||||
|
pub struct JournalStorage<V = f64> {
|
||||||
|
memory: MemoryStorage<V>,
|
||||||
|
path: PathBuf,
|
||||||
|
/// Serialise in-process writes so we only hold the file lock briefly.
|
||||||
|
write_lock: Mutex<()>,
|
||||||
|
_marker: PhantomData<V>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<V: Serialize + DeserializeOwned + Send + Sync> JournalStorage<V> {
|
||||||
|
/// 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<Path>) -> 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<Path>) -> crate::Result<Self> {
|
||||||
|
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<V>) -> 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<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for JournalStorage<V> {
|
||||||
|
fn push(&self, trial: CompletedTrial<V>) {
|
||||||
|
// Best-effort persist; the trial stays in memory regardless.
|
||||||
|
let _ = self.write_to_file(&trial);
|
||||||
|
self.memory.push(trial);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>> {
|
||||||
|
self.memory.trials_arc()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn refresh(&self) -> bool {
|
||||||
|
let Ok(loaded) = load_trials_from_file::<V>(&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<V: DeserializeOwned>(
|
||||||
|
path: &Path,
|
||||||
|
) -> crate::Result<Vec<CompletedTrial<V>>> {
|
||||||
|
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<V> =
|
||||||
|
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)
|
||||||
|
}
|
||||||
@@ -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<RwLock<Vec<CompletedTrial<V>>>>`.
|
||||||
|
pub struct MemoryStorage<V> {
|
||||||
|
trials: Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<V> MemoryStorage<V> {
|
||||||
|
/// 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<CompletedTrial<V>>) -> Self {
|
||||||
|
Self {
|
||||||
|
trials: Arc::new(RwLock::new(trials)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<V> Default for MemoryStorage<V> {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<V: Send + Sync> Storage<V> for MemoryStorage<V> {
|
||||||
|
fn push(&self, trial: CompletedTrial<V>) {
|
||||||
|
self.trials.write().push(trial);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>> {
|
||||||
|
&self.trials
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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<dyn Storage<V>>`. 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<V>: Send + Sync {
|
||||||
|
/// Append a completed trial to the store.
|
||||||
|
fn push(&self, trial: CompletedTrial<V>);
|
||||||
|
|
||||||
|
/// Return a reference to the in-memory trial buffer.
|
||||||
|
///
|
||||||
|
/// All implementations must maintain an `Arc<RwLock<Vec<…>>>` that
|
||||||
|
/// reflects the current set of trials. Callers may acquire a read
|
||||||
|
/// lock for efficient, allocation-free access.
|
||||||
|
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>>;
|
||||||
|
|
||||||
|
/// 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
|
||||||
|
}
|
||||||
|
}
|
||||||
+159
-60
@@ -49,8 +49,8 @@ where
|
|||||||
sampler: Arc<dyn Sampler>,
|
sampler: Arc<dyn Sampler>,
|
||||||
/// The pruner used to decide whether to stop trials early.
|
/// The pruner used to decide whether to stop trials early.
|
||||||
pruner: Arc<dyn Pruner>,
|
pruner: Arc<dyn Pruner>,
|
||||||
/// Completed trials (wrapped in Arc for sharing with Trial).
|
/// Trial storage backend (default: [`MemoryStorage`](crate::storage::MemoryStorage)).
|
||||||
completed_trials: Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
storage: Arc<dyn crate::storage::Storage<V>>,
|
||||||
/// Counter for generating unique trial IDs.
|
/// Counter for generating unique trial IDs.
|
||||||
next_trial_id: AtomicU64,
|
next_trial_id: AtomicU64,
|
||||||
/// Optional factory for creating sampler-aware trials.
|
/// Optional factory for creating sampler-aware trials.
|
||||||
@@ -84,7 +84,7 @@ where
|
|||||||
#[must_use]
|
#[must_use]
|
||||||
pub fn new(direction: Direction) -> Self
|
pub fn new(direction: Direction) -> Self
|
||||||
where
|
where
|
||||||
V: 'static,
|
V: Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
Self::with_sampler(direction, RandomSampler::new())
|
Self::with_sampler(direction, RandomSampler::new())
|
||||||
}
|
}
|
||||||
@@ -109,7 +109,7 @@ where
|
|||||||
#[must_use]
|
#[must_use]
|
||||||
pub fn minimize(sampler: impl Sampler + 'static) -> Self
|
pub fn minimize(sampler: impl Sampler + 'static) -> Self
|
||||||
where
|
where
|
||||||
V: 'static,
|
V: Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
Self::with_sampler(Direction::Minimize, sampler)
|
Self::with_sampler(Direction::Minimize, sampler)
|
||||||
}
|
}
|
||||||
@@ -134,7 +134,7 @@ where
|
|||||||
#[must_use]
|
#[must_use]
|
||||||
pub fn maximize(sampler: impl Sampler + 'static) -> Self
|
pub fn maximize(sampler: impl Sampler + 'static) -> Self
|
||||||
where
|
where
|
||||||
V: 'static,
|
V: Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
Self::with_sampler(Direction::Maximize, sampler)
|
Self::with_sampler(Direction::Maximize, sampler)
|
||||||
}
|
}
|
||||||
@@ -158,40 +158,28 @@ where
|
|||||||
/// ```
|
/// ```
|
||||||
pub fn with_sampler(direction: Direction, sampler: impl Sampler + 'static) -> Self
|
pub fn with_sampler(direction: Direction, sampler: impl Sampler + 'static) -> Self
|
||||||
where
|
where
|
||||||
V: 'static,
|
V: Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
Self::with_sampler_and_storage(
|
||||||
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
|
||||||
|
|
||||||
let pruner: Arc<dyn Pruner> = Arc::new(NopPruner);
|
|
||||||
|
|
||||||
// For Study<f64>, 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 {
|
|
||||||
direction,
|
direction,
|
||||||
sampler,
|
sampler,
|
||||||
pruner,
|
crate::storage::MemoryStorage::<V>::new(),
|
||||||
completed_trials,
|
)
|
||||||
next_trial_id: AtomicU64::new(0),
|
|
||||||
trial_factory,
|
|
||||||
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Builds a trial factory for sampler integration when `V = f64`.
|
/// Builds a trial factory for sampler integration when `V = f64`.
|
||||||
fn make_trial_factory(
|
fn make_trial_factory(
|
||||||
sampler: &Arc<dyn Sampler>,
|
sampler: &Arc<dyn Sampler>,
|
||||||
completed_trials: &Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
storage: &Arc<dyn crate::storage::Storage<V>>,
|
||||||
pruner: &Arc<dyn Pruner>,
|
pruner: &Arc<dyn Pruner>,
|
||||||
) -> Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>
|
) -> Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>
|
||||||
where
|
where
|
||||||
V: 'static,
|
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.
|
// 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<RwLock<Vec<CompletedTrial<f64>>>>> = any_ref.downcast_ref();
|
let f64_trials: Option<&Arc<RwLock<Vec<CompletedTrial<f64>>>>> = any_ref.downcast_ref();
|
||||||
|
|
||||||
f64_trials.map(|trials| {
|
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<V> + 'static,
|
||||||
|
) -> Self
|
||||||
|
where
|
||||||
|
V: 'static,
|
||||||
|
{
|
||||||
|
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
||||||
|
let pruner: Arc<dyn Pruner> = Arc::new(NopPruner);
|
||||||
|
let storage: Arc<dyn crate::storage::Storage<V>> = 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.
|
/// Returns the optimization direction.
|
||||||
pub fn direction(&self) -> Direction {
|
pub fn direction(&self) -> Direction {
|
||||||
self.direction
|
self.direction
|
||||||
@@ -254,18 +278,19 @@ where
|
|||||||
pruner: impl Pruner + 'static,
|
pruner: impl Pruner + 'static,
|
||||||
) -> Self
|
) -> Self
|
||||||
where
|
where
|
||||||
V: 'static,
|
V: Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
||||||
let pruner: Arc<dyn Pruner> = Arc::new(pruner);
|
let pruner: Arc<dyn Pruner> = Arc::new(pruner);
|
||||||
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
let storage: Arc<dyn crate::storage::Storage<V>> =
|
||||||
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials, &pruner);
|
Arc::new(crate::storage::MemoryStorage::<V>::new());
|
||||||
|
let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
direction,
|
direction,
|
||||||
sampler,
|
sampler,
|
||||||
pruner,
|
pruner,
|
||||||
completed_trials,
|
storage,
|
||||||
next_trial_id: AtomicU64::new(0),
|
next_trial_id: AtomicU64::new(0),
|
||||||
trial_factory,
|
trial_factory,
|
||||||
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
|
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
|
||||||
@@ -277,8 +302,7 @@ where
|
|||||||
V: 'static,
|
V: 'static,
|
||||||
{
|
{
|
||||||
self.sampler = Arc::new(sampler);
|
self.sampler = Arc::new(sampler);
|
||||||
self.trial_factory =
|
self.trial_factory = Self::make_trial_factory(&self.sampler, &self.storage, &self.pruner);
|
||||||
Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Sets a new pruner for the study.
|
/// Sets a new pruner for the study.
|
||||||
@@ -291,8 +315,7 @@ where
|
|||||||
V: 'static,
|
V: 'static,
|
||||||
{
|
{
|
||||||
self.pruner = Arc::new(pruner);
|
self.pruner = Arc::new(pruner);
|
||||||
self.trial_factory =
|
self.trial_factory = Self::make_trial_factory(&self.sampler, &self.storage, &self.pruner);
|
||||||
Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns a reference to the study's pruner.
|
/// Returns a reference to the study's pruner.
|
||||||
@@ -401,6 +424,13 @@ where
|
|||||||
/// assert_eq!(trial2.id(), 1);
|
/// assert_eq!(trial2.id(), 1);
|
||||||
/// ```
|
/// ```
|
||||||
pub fn create_trial(&self) -> Trial {
|
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 id = self.next_trial_id();
|
||||||
let mut trial = if let Some(factory) = &self.trial_factory {
|
let mut trial = if let Some(factory) = &self.trial_factory {
|
||||||
factory(id)
|
factory(id)
|
||||||
@@ -455,7 +485,8 @@ where
|
|||||||
);
|
);
|
||||||
completed.state = TrialState::Complete;
|
completed.state = TrialState::Complete;
|
||||||
completed.constraints = trial.constraint_values().to_vec();
|
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.
|
/// Records a failed trial with an error message.
|
||||||
@@ -565,7 +596,8 @@ where
|
|||||||
);
|
);
|
||||||
completed.state = TrialState::Pruned;
|
completed.state = TrialState::Pruned;
|
||||||
completed.constraints = trial.constraint_values().to_vec();
|
completed.constraints = trial.constraint_values().to_vec();
|
||||||
self.completed_trials.write().push(completed);
|
|
||||||
|
self.storage.push(completed);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns an iterator over all completed trials.
|
/// Returns an iterator over all completed trials.
|
||||||
@@ -596,7 +628,7 @@ where
|
|||||||
where
|
where
|
||||||
V: Clone,
|
V: Clone,
|
||||||
{
|
{
|
||||||
self.completed_trials.read().clone()
|
self.storage.trials_arc().read().clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns the number of completed trials.
|
/// Returns the number of completed trials.
|
||||||
@@ -619,12 +651,13 @@ where
|
|||||||
/// assert_eq!(study.n_trials(), 1);
|
/// assert_eq!(study.n_trials(), 1);
|
||||||
/// ```
|
/// ```
|
||||||
pub fn n_trials(&self) -> usize {
|
pub fn n_trials(&self) -> usize {
|
||||||
self.completed_trials.read().len()
|
self.storage.trials_arc().read().len()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns the number of pruned trials.
|
/// Returns the number of pruned trials.
|
||||||
pub fn n_pruned_trials(&self) -> usize {
|
pub fn n_pruned_trials(&self) -> usize {
|
||||||
self.completed_trials
|
self.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Pruned)
|
.filter(|t| t.state == TrialState::Pruned)
|
||||||
@@ -703,7 +736,7 @@ where
|
|||||||
where
|
where
|
||||||
V: Clone,
|
V: Clone,
|
||||||
{
|
{
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
let direction = self.direction;
|
let direction = self.direction;
|
||||||
|
|
||||||
let best = trials
|
let best = trials
|
||||||
@@ -767,7 +800,7 @@ where
|
|||||||
where
|
where
|
||||||
V: Clone,
|
V: Clone,
|
||||||
{
|
{
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
let direction = self.direction;
|
let direction = self.direction;
|
||||||
let mut completed: Vec<_> = trials
|
let mut completed: Vec<_> = trials
|
||||||
.iter()
|
.iter()
|
||||||
@@ -848,7 +881,7 @@ where
|
|||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
{
|
{
|
||||||
tracing::info!(trial_id, "trial completed");
|
tracing::info!(trial_id, "trial completed");
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
if trials
|
if trials
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Complete)
|
.filter(|t| t.state == TrialState::Complete)
|
||||||
@@ -876,7 +909,8 @@ where
|
|||||||
|
|
||||||
// Return error if no trials completed successfully
|
// Return error if no trials completed successfully
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -975,7 +1009,8 @@ where
|
|||||||
|
|
||||||
// Return error if no trials completed successfully
|
// Return error if no trials completed successfully
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -1101,7 +1136,8 @@ where
|
|||||||
|
|
||||||
// Return error if no trials completed successfully
|
// Return error if no trials completed successfully
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -1197,7 +1233,7 @@ where
|
|||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
{
|
{
|
||||||
tracing::info!(trial_id, "trial completed");
|
tracing::info!(trial_id, "trial completed");
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
if trials
|
if trials
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Complete)
|
.filter(|t| t.state == TrialState::Complete)
|
||||||
@@ -1210,7 +1246,7 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Get the just-completed trial for the callback
|
// 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 {
|
let Some(completed) = trials.last() else {
|
||||||
return Err(crate::Error::Internal(
|
return Err(crate::Error::Internal(
|
||||||
"completed trial not found after adding",
|
"completed trial not found after adding",
|
||||||
@@ -1243,7 +1279,8 @@ where
|
|||||||
|
|
||||||
// Return error if no trials completed successfully
|
// Return error if no trials completed successfully
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -1328,7 +1365,8 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -1420,7 +1458,7 @@ where
|
|||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
{
|
{
|
||||||
tracing::info!(trial_id, "trial completed");
|
tracing::info!(trial_id, "trial completed");
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
if trials
|
if trials
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Complete)
|
.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 {
|
let Some(completed) = trials.last() else {
|
||||||
return Err(crate::Error::Internal(
|
return Err(crate::Error::Internal(
|
||||||
"completed trial not found after adding",
|
"completed trial not found after adding",
|
||||||
@@ -1461,7 +1499,8 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -1519,7 +1558,8 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -1608,7 +1648,8 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.any(|t| t.state == TrialState::Complete);
|
||||||
@@ -1716,7 +1757,8 @@ where
|
|||||||
|
|
||||||
// Return error if no trials completed successfully
|
// Return error if no trials completed successfully
|
||||||
let has_complete = self
|
let has_complete = self
|
||||||
.completed_trials
|
.storage
|
||||||
|
.trials_arc()
|
||||||
.read()
|
.read()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| t.state == TrialState::Complete);
|
.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<()> {
|
pub fn to_csv(&self, mut writer: impl std::io::Write) -> std::io::Result<()> {
|
||||||
use std::collections::BTreeMap;
|
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).
|
// Collect all unique parameter labels (sorted for deterministic column order).
|
||||||
let mut param_columns: BTreeMap<ParamId, String> = BTreeMap::new();
|
let mut param_columns: BTreeMap<ParamId, String> = BTreeMap::new();
|
||||||
@@ -1873,7 +1915,7 @@ where
|
|||||||
pub fn summary(&self) -> String {
|
pub fn summary(&self) -> String {
|
||||||
use fmt::Write;
|
use fmt::Write;
|
||||||
|
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
let n_complete = trials
|
let n_complete = trials
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Complete)
|
.filter(|t| t.state == TrialState::Complete)
|
||||||
@@ -1967,7 +2009,7 @@ where
|
|||||||
use crate::param::ParamValue;
|
use crate::param::ParamValue;
|
||||||
use crate::types::TrialState;
|
use crate::types::TrialState;
|
||||||
|
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
let complete: Vec<_> = trials
|
let complete: Vec<_> = trials
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Complete)
|
.filter(|t| t.state == TrialState::Complete)
|
||||||
@@ -2064,7 +2106,7 @@ where
|
|||||||
use crate::param::ParamValue;
|
use crate::param::ParamValue;
|
||||||
use crate::types::TrialState;
|
use crate::types::TrialState;
|
||||||
|
|
||||||
let trials = self.completed_trials.read();
|
let trials = self.storage.trials_arc().read();
|
||||||
let complete: Vec<_> = trials
|
let complete: Vec<_> = trials
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Complete)
|
.filter(|t| t.state == TrialState::Complete)
|
||||||
@@ -2258,6 +2300,63 @@ impl Study<f64> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl<V: PartialOrd + Send + Sync + 'static> Study<V> {
|
||||||
|
/// 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<V> + 'static,
|
||||||
|
) -> Self {
|
||||||
|
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
||||||
|
let pruner: Arc<dyn Pruner> = Arc::new(pruner);
|
||||||
|
let storage: Arc<dyn crate::storage::Storage<V>> = 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<V> Study<V>
|
||||||
|
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<std::path::Path>,
|
||||||
|
) -> crate::Result<Self> {
|
||||||
|
let storage = crate::storage::JournalStorage::<V>::open(path)?;
|
||||||
|
Ok(Self::with_sampler_and_storage(direction, sampler, storage))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(feature = "visualization")]
|
#[cfg(feature = "visualization")]
|
||||||
impl Study<f64> {
|
impl Study<f64> {
|
||||||
/// Generates an HTML report with interactive Plotly.js charts.
|
/// Generates an HTML report with interactive Plotly.js charts.
|
||||||
@@ -2382,7 +2481,7 @@ impl<V: PartialOrd + Clone + Default + serde::Serialize> Study<V> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(feature = "serde")]
|
#[cfg(feature = "serde")]
|
||||||
impl<V: PartialOrd + Clone + serde::de::DeserializeOwned + 'static> Study<V> {
|
impl<V: PartialOrd + Send + Sync + Clone + serde::de::DeserializeOwned + 'static> Study<V> {
|
||||||
/// Loads a study from a JSON file.
|
/// Loads a study from a JSON file.
|
||||||
///
|
///
|
||||||
/// The loaded study uses a `RandomSampler` by default. Call
|
/// The loaded study uses a `RandomSampler` by default. Call
|
||||||
@@ -2397,7 +2496,7 @@ impl<V: PartialOrd + Clone + serde::de::DeserializeOwned + 'static> Study<V> {
|
|||||||
let snapshot: StudySnapshot<V> = serde_json::from_reader(file)
|
let snapshot: StudySnapshot<V> = serde_json::from_reader(file)
|
||||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||||
let study = Study::new(snapshot.direction);
|
let study = Study::new(snapshot.direction);
|
||||||
*study.completed_trials.write() = snapshot.trials;
|
*study.storage.trials_arc().write() = snapshot.trials;
|
||||||
study
|
study
|
||||||
.next_trial_id
|
.next_trial_id
|
||||||
.store(snapshot.next_trial_id, Ordering::Relaxed);
|
.store(snapshot.next_trial_id, Ordering::Relaxed);
|
||||||
|
|||||||
@@ -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<f64> {
|
||||||
|
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::<f64>::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::<f64>::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::<f64>::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::<f64>::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<u64> = loaded.iter().map(|t| t.id).collect();
|
||||||
|
ids.sort();
|
||||||
|
assert_eq!(ids, (0..100).collect::<Vec<_>>());
|
||||||
|
|
||||||
|
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::<f64>::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<u64> = 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::<f64>::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();
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user