gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: move next_trial_id counter from Study into Storage trait

Browse Source
This commit is contained in:
Manuel Raimann
2026-02-11 23:32:15 +01:00
parent d14f4e0792
commit d2e36ec304
5 changed files with 56 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/storage/journal.rs
+7
View File
@@ -115,12 +115,19 @@ impl<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for JournalStorag
self.memory.trials_arc()
}
fn next_trial_id(&self) -> u64 {
self.memory.next_trial_id()
}
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() {
if let Some(max_id) = loaded.iter().map(|t| t.id).max() {
self.memory.bump_next_id(max_id + 1);
}
*guard = loaded;
true
} else {
src/storage/memory.rs
+14
View File
@@ -1,3 +1,4 @@
use core::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use parking_lot::RwLock;
@@ -10,6 +11,7 @@ use crate::sampler::CompletedTrial;
/// This is a thin wrapper around `Arc<RwLock<Vec<CompletedTrial<V>>>>`.
pub struct MemoryStorage<V> {
trials: Arc<RwLock<Vec<CompletedTrial<V>>>>,
next_id: AtomicU64,
}
impl<V> MemoryStorage<V> {
@@ -18,16 +20,24 @@ impl<V> MemoryStorage<V> {
pub fn new() -> Self {
Self {
trials: Arc::new(RwLock::new(Vec::new())),
next_id: AtomicU64::new(0),
}
}
/// Creates an in-memory store pre-populated with `trials`.
#[must_use]
pub fn with_trials(trials: Vec<CompletedTrial<V>>) -> Self {
let next_id = trials.iter().map(|t| t.id).max().map_or(0, |id| id + 1);
Self {
trials: Arc::new(RwLock::new(trials)),
next_id: AtomicU64::new(next_id),
}
}
/// Ensures the ID counter is at least `min_value`.
pub(crate) fn bump_next_id(&self, min_value: u64) {
self.next_id.fetch_max(min_value, Ordering::SeqCst);
}
}
impl<V> Default for MemoryStorage<V> {
@@ -44,4 +54,8 @@ impl<V: Send + Sync> Storage<V> for MemoryStorage<V> {
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>> {
&self.trials
}
fn next_trial_id(&self) -> u64 {
self.next_id.fetch_add(1, Ordering::SeqCst)
}
}
src/storage/mod.rs
+6
View File
@@ -42,6 +42,12 @@ pub trait Storage<V>: Send + Sync {
/// lock for efficient, allocation-free access.
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>>;
/// Atomically returns the next unique trial ID.
///
/// Each call increments an internal counter so that consecutive
/// calls always produce distinct IDs.
fn next_trial_id(&self) -> u64;
/// Reload from an external source (e.g. a file written by another
/// process). Returns `true` if the in-memory buffer was updated.
///
src/storage/sqlite.rs
+7
View File
@@ -98,6 +98,10 @@ impl<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for SqliteStorage
self.memory.trials_arc()
}
fn next_trial_id(&self) -> u64 {
self.memory.next_trial_id()
}
fn refresh(&self) -> bool {
let conn = self.conn.lock();
let Ok(loaded) = load_all::<V>(&conn) else {
@@ -105,6 +109,9 @@ impl<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for SqliteStorage
};
let mut guard = self.memory.trials_arc().write();
if loaded.len() > guard.len() {
if let Some(max_id) = loaded.iter().map(|t| t.id).max() {
self.memory.bump_next_id(max_id + 1);
}
*guard = loaded;
true
} else {
src/study.rs
+22 -46
View File
@@ -6,7 +6,6 @@ use core::fmt;
use core::future::Future;
use core::marker::PhantomData;
use core::ops::ControlFlow;
use core::sync::atomic::{AtomicU64, Ordering};
use core::time::Duration;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
@@ -52,8 +51,6 @@ where
pruner: Arc<dyn Pruner>,
/// Trial storage backend (default: [`MemoryStorage`](crate::storage::MemoryStorage)).
storage: Arc<dyn crate::storage::Storage<V>>,
/// Counter for generating unique trial IDs.
next_trial_id: AtomicU64,
/// Optional factory for creating sampler-aware trials.
/// Set automatically for `Study<f64>` so that `create_trial()` and all
/// optimization methods use the sampler without requiring `_with_sampler` suffixes.
@@ -240,26 +237,18 @@ where
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.
#[must_use]
pub fn direction(&self) -> Direction {
self.direction
}
@@ -316,7 +305,6 @@ where
sampler,
pruner,
storage,
next_trial_id: AtomicU64::new(0),
trial_factory,
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
}
@@ -344,6 +332,7 @@ where
}
/// Returns a reference to the study's pruner.
#[must_use]
pub fn pruner(&self) -> &dyn Pruner {
&*self.pruner
}
@@ -423,7 +412,7 @@ where
/// Generates the next unique trial ID.
pub(crate) fn next_trial_id(&self) -> u64 {
self.next_trial_id.fetch_add(1, Ordering::SeqCst)
self.storage.next_trial_id()
}
/// Creates a new trial with a unique ID.
@@ -448,13 +437,9 @@ where
/// let trial2 = study.create_trial();
/// assert_eq!(trial2.id(), 1);
/// ```
#[must_use]
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);
}
}
self.storage.refresh();
let id = self.next_trial_id();
let mut trial = if let Some(factory) = &self.trial_factory {
@@ -565,6 +550,7 @@ where
/// let value = x_val * x_val;
/// study.tell(trial, Ok::<_, &str>(value));
/// ```
#[must_use]
pub fn ask(&self) -> Trial {
self.create_trial()
}
@@ -649,6 +635,7 @@ where
/// println!("Trial {} has value {:?}", completed.id, completed.value);
/// }
/// ```
#[must_use]
pub fn trials(&self) -> Vec<CompletedTrial<V>>
where
V: Clone,
@@ -675,11 +662,13 @@ where
/// study.complete_trial(trial, 0.5);
/// assert_eq!(study.n_trials(), 1);
/// ```
#[must_use]
pub fn n_trials(&self) -> usize {
self.storage.trials_arc().read().len()
}
/// Returns the number of pruned trials.
#[must_use]
pub fn n_pruned_trials(&self) -> usize {
self.storage
.trials_arc()
@@ -821,6 +810,7 @@ where
/// Only includes completed trials (not failed or pruned).
///
/// If fewer than `n` completed trials exist, returns all of them.
#[must_use]
pub fn top_trials(&self, n: usize) -> Vec<CompletedTrial<V>>
where
V: Clone,
@@ -1987,6 +1977,7 @@ where
///
/// This clones the internal trial list, so it is suitable for
/// analysis and iteration but not for hot paths.
#[must_use]
pub fn iter(&self) -> std::vec::IntoIter<CompletedTrial<V>> {
self.trials().into_iter()
}
@@ -2234,6 +2225,7 @@ impl Study<f64> {
since = "0.2.0",
note = "use `create_trial()` instead — it now uses the sampler automatically for Study<f64>"
)]
#[must_use]
pub fn create_trial_with_sampler(&self) -> Trial {
self.create_trial()
}
@@ -2336,20 +2328,11 @@ impl<V: PartialOrd + Send + Sync + 'static> Study<V> {
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())),
}
@@ -2451,20 +2434,11 @@ impl<V: PartialOrd> StudyBuilder<V> {
let storage: Arc<dyn crate::storage::Storage<V>> = Arc::from(storage);
let trial_factory = Study::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);
Study {
direction: self.direction,
sampler,
pruner,
storage,
next_trial_id: AtomicU64::new(next_id),
trial_factory,
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
}
@@ -2590,11 +2564,13 @@ impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
/// Returns an I/O error if the file cannot be created or written.
pub fn save(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
let path = path.as_ref();
let trials = self.trials();
let next_trial_id = trials.iter().map(|t| t.id).max().map_or(0, |id| id + 1);
let snapshot = StudySnapshot {
version: 1,
direction: self.direction,
trials: self.trials(),
next_trial_id: self.next_trial_id.load(Ordering::Relaxed),
trials,
next_trial_id,
metadata: HashMap::new(),
};
@@ -2660,12 +2636,12 @@ impl<V: PartialOrd + Send + Sync + Clone + serde::de::DeserializeOwned + 'static
let file = std::fs::File::open(path)?;
let snapshot: StudySnapshot<V> = serde_json::from_reader(file)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
let study = Study::new(snapshot.direction);
*study.storage.trials_arc().write() = snapshot.trials;
study
.next_trial_id
.store(snapshot.next_trial_id, Ordering::Relaxed);
Ok(study)
let storage = crate::storage::MemoryStorage::with_trials(snapshot.trials);
Ok(Self::with_sampler_and_storage(
snapshot.direction,
RandomSampler::new(),
storage,
))
}
}