diff --git a/examples/async_parallel.rs b/examples/async_parallel.rs index f3c3949..8a31f30 100644 --- a/examples/async_parallel.rs +++ b/examples/async_parallel.rs @@ -1,7 +1,8 @@ //! Async parallel optimization — evaluate multiple trials concurrently. //! //! Uses `optimize_parallel` with tokio to run several trials at once, -//! reducing wall-clock time when the objective involves I/O or async work. +//! reducing wall-clock time when the objective involves blocking work. +//! Each sync closure is internally wrapped in `spawn_blocking`. //! //! Run with: `cargo run --example async_parallel --features async` @@ -19,25 +20,18 @@ async fn main() -> optimizer::Result<()> { println!("Running {n_trials} trials with {concurrency} concurrent workers..."); + let xc = x.clone(); + let yc = y.clone(); study - .optimize_parallel(n_trials, concurrency, { - let x = x.clone(); - let y = y.clone(); - move |mut trial| { - let x = x.clone(); - let y = y.clone(); - async move { - let xv = x.suggest(&mut trial)?; - let yv = y.suggest(&mut trial)?; - - // Simulate async I/O (e.g. calling an external service) - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - - let value = xv * xv + yv * yv; - Ok::<_, optimizer::Error>((trial, value)) - } - } - }) + .optimize_parallel( + n_trials, + concurrency, + move |trial: &mut optimizer::Trial| { + let xv = xc.suggest(trial)?; + let yv = yc.suggest(trial)?; + Ok::<_, optimizer::Error>(xv * xv + yv * yv) + }, + ) .await?; let best = study.best_trial()?; diff --git a/examples/early_stopping.rs b/examples/early_stopping.rs index 7f57c88..efbf78c 100644 --- a/examples/early_stopping.rs +++ b/examples/early_stopping.rs @@ -1,8 +1,8 @@ //! Early stopping — halt an entire study once a target is reached. //! -//! Use `optimize_with_callback` to inspect each completed trial and return -//! `ControlFlow::Break(())` when the study should stop (e.g. a quality -//! threshold is met or a time budget is exhausted). +//! Implements the [`Objective`] trait on a custom struct and uses the +//! [`after_trial`](Objective::after_trial) hook to return +//! `ControlFlow::Break(())` when the best value drops below a threshold. //! //! Run with: `cargo run --example early_stopping` @@ -10,26 +10,41 @@ use std::ops::ControlFlow; use optimizer::prelude::*; +/// An objective that minimises `(x - 3)^2` and stops early once the +/// value drops below `target`. +struct EarlyStopObjective { + x: FloatParam, + target: f64, +} + +impl Objective for EarlyStopObjective { + type Error = Error; + + fn evaluate(&self, trial: &mut Trial) -> Result { + let v = self.x.suggest(trial)?; + Ok((v - 3.0).powi(2)) + } + + fn after_trial(&self, _study: &Study, trial: &CompletedTrial) -> ControlFlow<()> { + if trial.value < self.target { + println!("Target {} reached at trial #{}", self.target, trial.id); + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + } +} + fn main() -> optimizer::Result<()> { let study: Study = Study::new(Direction::Minimize); let x = FloatParam::new(-10.0, 10.0).name("x"); - let target = 0.01; + let objective = EarlyStopObjective { + x: x.clone(), + target: 0.01, + }; - study.optimize_with_callback( - 100, // upper bound — we expect to stop much earlier - |trial| { - let xv = x.suggest(trial)?; - Ok::<_, Error>((xv - 3.0).powi(2)) - }, - |_study, completed| { - if completed.value < target { - println!("Target {target} reached at trial #{}", completed.id); - return ControlFlow::Break(()); - } - ControlFlow::Continue(()) - }, - )?; + study.optimize_with(100, objective)?; let best = study.best_trial()?; println!( diff --git a/src/lib.rs b/src/lib.rs index 0636f54..a696b6a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -113,6 +113,7 @@ mod fanova; mod importance; mod kde; pub mod multi_objective; +pub mod objective; mod param; pub mod parameter; pub mod pareto; @@ -127,6 +128,7 @@ mod visualization; pub use error::{Error, Result, TrialPruned}; pub use fanova::{FanovaConfig, FanovaResult}; +pub use objective::Objective; #[cfg(feature = "derive")] pub use optimizer_derive::Categorical; #[cfg(feature = "serde")] @@ -150,6 +152,7 @@ pub mod prelude { pub use crate::multi_objective::{ MultiObjectiveSampler, MultiObjectiveStudy, MultiObjectiveTrial, }; + pub use crate::objective::Objective; pub use crate::parameter::{ BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamValue, Parameter, diff --git a/src/objective.rs b/src/objective.rs new file mode 100644 index 0000000..d1ce710 --- /dev/null +++ b/src/objective.rs @@ -0,0 +1,132 @@ +//! The [`Objective`] trait defines what gets optimized. +//! +//! For simple closures, pass them directly to +//! [`Study::optimize`](crate::Study::optimize): +//! +//! ``` +//! use optimizer::prelude::*; +//! +//! let study: Study = Study::new(Direction::Minimize); +//! let x = FloatParam::new(-10.0, 10.0).name("x"); +//! +//! study +//! .optimize(50, |trial| { +//! let v = x.suggest(trial)?; +//! Ok::<_, Error>((v - 3.0).powi(2)) +//! }) +//! .unwrap(); +//! ``` +//! +//! For richer control — early stopping, retries, or per-trial logging — +//! implement [`Objective`] on a struct and pass it to +//! [`Study::optimize_with`](crate::Study::optimize_with): +//! +//! ``` +//! use std::ops::ControlFlow; +//! +//! use optimizer::Objective; +//! use optimizer::prelude::*; +//! +//! struct QuadraticWithEarlyStopping { +//! x: FloatParam, +//! target: f64, +//! } +//! +//! impl Objective for QuadraticWithEarlyStopping { +//! type Error = Error; +//! +//! fn evaluate(&self, trial: &mut Trial) -> Result { +//! let v = self.x.suggest(trial)?; +//! Ok((v - 3.0).powi(2)) +//! } +//! +//! fn after_trial(&self, _study: &Study, trial: &CompletedTrial) -> ControlFlow<()> { +//! if trial.value < self.target { +//! ControlFlow::Break(()) +//! } else { +//! ControlFlow::Continue(()) +//! } +//! } +//! } +//! +//! let study: Study = Study::new(Direction::Minimize); +//! let obj = QuadraticWithEarlyStopping { +//! x: FloatParam::new(-10.0, 10.0).name("x"), +//! target: 1.0, +//! }; +//! study.optimize_with(200, obj).unwrap(); +//! assert!(study.best_value().unwrap() < 1.0); +//! ``` + +use core::ops::ControlFlow; + +use crate::sampler::CompletedTrial; +use crate::study::Study; +use crate::trial::Trial; + +/// Defines an objective function with lifecycle hooks for optimization. +/// +/// The only required method is [`evaluate`](Objective::evaluate), which +/// computes the objective value for a given trial. Optional hooks provide +/// early stopping ([`before_trial`](Objective::before_trial), +/// [`after_trial`](Objective::after_trial)) and automatic retries +/// ([`max_retries`](Objective::max_retries)). +/// +/// # When to use `Objective` vs a closure +/// +/// - **Closure** — pass directly to [`Study::optimize`](crate::Study::optimize) +/// for simple evaluate-only objectives. +/// - **`Objective` struct** — implement this trait when you need hooks +/// (`before_trial`, `after_trial`) or retries. +/// +/// # Thread safety +/// +/// The async optimization methods (`optimize_async`, `optimize_parallel`) +/// additionally require `Send + Sync + 'static` on the objective. The +/// sync `optimize` method has no thread-safety requirements. +pub trait Objective { + /// The error type returned by [`evaluate`](Objective::evaluate). + type Error: ToString + 'static; + + /// Evaluate the objective function for a single trial. + /// + /// Sample parameters from `trial` via + /// [`Parameter::suggest`](crate::parameter::Parameter::suggest) and + /// return the objective value. Return `Err(TrialPruned)` to prune a + /// trial early. + /// + /// # Errors + /// + /// Any error whose type implements `ToString`. Pruning errors + /// (`Error::TrialPruned` or `TrialPruned`) are handled specially — + /// the trial is recorded as pruned rather than failed. + fn evaluate(&self, trial: &mut Trial) -> Result; + + /// Called before each trial is created. + /// + /// Return `ControlFlow::Break(())` to stop the optimization loop + /// before the next trial starts. + /// + /// Default: always continues. + fn before_trial(&self, _study: &Study) -> ControlFlow<()> { + ControlFlow::Continue(()) + } + + /// Called after each **completed** trial (not failed or pruned). + /// + /// Return `ControlFlow::Break(())` to stop the optimization loop. + /// + /// Default: always continues. + fn after_trial(&self, _study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { + ControlFlow::Continue(()) + } + + /// Maximum number of retries for a failed trial. + /// + /// When `evaluate` returns a non-pruning error and retries remain, + /// the same parameter configuration is re-evaluated. Set to `0` + /// (the default) to disable retries. + fn max_retries(&self) -> usize { + 0 + } +} diff --git a/src/study.rs b/src/study.rs index 3827942..fd9acb6 100644 --- a/src/study.rs +++ b/src/study.rs @@ -2,14 +2,10 @@ use core::any::Any; use core::fmt; -#[cfg(feature = "async")] -use core::future::Future; use core::marker::PhantomData; use core::ops::ControlFlow; -use core::time::Duration; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; -use std::time::Instant; use parking_lot::{Mutex, RwLock}; @@ -883,26 +879,16 @@ where completed } - /// Run optimization with the given objective function. + /// Run optimization with a closure. /// - /// This method runs `n_trials` evaluations sequentially. For each trial: - /// 1. A new trial is created - /// 2. The objective function is called with the trial - /// 3. If successful, the trial is recorded as completed - /// 4. If the objective returns an error, the trial is recorded as failed - /// - /// Failed trials do not stop the optimization; the process continues with - /// the next trial. - /// - /// # Arguments - /// - /// * `n_trials` - The number of trials to run. - /// * `objective` - A closure that takes a mutable reference to a `Trial` and - /// returns the objective value or an error. + /// Runs up to `n_trials` evaluations of `objective` sequentially. + /// For lifecycle hooks (early stopping, retries), implement the + /// [`Objective`](crate::Objective) trait and use + /// [`optimize_with`](Self::optimize_with) instead. /// /// # Errors /// - /// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials). + /// Returns `Error::NoCompletedTrials` if no trials completed successfully. /// /// # Examples /// @@ -911,29 +897,25 @@ where /// use optimizer::sampler::random::RandomSampler; /// use optimizer::{Direction, Study}; /// - /// // Minimize x^2 /// let sampler = RandomSampler::with_seed(42); /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); - /// /// let x_param = FloatParam::new(-10.0, 10.0); /// /// study - /// .optimize(10, |trial| { + /// .optimize(10, |trial: &mut optimizer::Trial| { /// let x = x_param.suggest(trial)?; /// Ok::<_, optimizer::Error>(x * x) /// }) /// .unwrap(); /// - /// // At least one trial should have completed /// assert!(study.n_trials() > 0); - /// let best = study.best_value().unwrap(); - /// assert!(best >= 0.0); + /// assert!(study.best_value().unwrap() >= 0.0); /// ``` pub fn optimize(&self, n_trials: usize, mut objective: F) -> crate::Result<()> where - F: FnMut(&mut Trial) -> core::result::Result, + F: FnMut(&mut Trial) -> Result, E: ToString + 'static, - V: Default, + V: Clone + Default, { #[cfg(feature = "tracing")] let _span = @@ -941,478 +923,6 @@ where for _ in 0..n_trials { let mut trial = self.create_trial(); - - match objective(&mut trial) { - Ok(value) => { - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - self.complete_trial(trial, value); - - #[cfg(feature = "tracing")] - { - tracing::info!(trial_id, "trial completed"); - let trials = self.storage.trials_arc().read(); - if trials - .iter() - .filter(|t| t.state == TrialState::Complete) - .count() - == 1 - || trials.last().map(|t| t.id) == self.best_id(&trials) - { - tracing::info!(trial_id, "new best value found"); - } - } - } - Err(e) => { - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - if is_trial_pruned(&e) { - self.prune_trial(trial); - trace_info!(trial_id, "trial pruned"); - } else { - self.fail_trial(trial, e.to_string()); - trace_debug!(trial_id, "trial failed"); - } - } - } - } - - // Return error if no trials completed successfully - let has_complete = self - .storage - .trials_arc() - .read() - .iter() - .any(|t| t.state == TrialState::Complete); - if !has_complete { - return Err(crate::Error::NoCompletedTrials); - } - - Ok(()) - } - - /// Run optimization asynchronously with the given objective function. - /// - /// This method runs `n_trials` evaluations sequentially, but the objective - /// function can be async (e.g., for I/O-bound operations like network requests - /// or file operations). - /// - /// The objective function takes ownership of the `Trial` and must return it - /// along with the result. This allows async operations to use the trial - /// across await points. - /// - /// # Arguments - /// - /// * `n_trials` - The number of trials to run. - /// * `objective` - A function that takes a `Trial` and returns a `Future` - /// that resolves to a tuple of `(Trial, Result)`. - /// - /// # Errors - /// - /// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials). - /// - /// # Examples - /// - /// ``` - /// use optimizer::parameter::{FloatParam, Parameter}; - /// use optimizer::sampler::random::RandomSampler; - /// use optimizer::{Direction, Study}; - /// - /// # #[cfg(feature = "async")] - /// # async fn example() -> optimizer::Result<()> { - /// // Minimize x^2 with async objective - /// let sampler = RandomSampler::with_seed(42); - /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); - /// - /// let x_param = FloatParam::new(-10.0, 10.0); - /// - /// study - /// .optimize_async(10, |mut trial| { - /// let x_param = x_param.clone(); - /// async move { - /// let x = x_param.suggest(&mut trial)?; - /// // Simulate async work (e.g., network request) - /// let value = x * x; - /// Ok::<_, optimizer::Error>((trial, value)) - /// } - /// }) - /// .await?; - /// - /// // At least one trial should have completed - /// assert!(study.n_trials() > 0); - /// # Ok(()) - /// # } - /// ``` - #[cfg(feature = "async")] - pub async fn optimize_async( - &self, - n_trials: usize, - objective: F, - ) -> crate::Result<()> - where - F: Fn(Trial) -> Fut, - Fut: Future>, - E: ToString, - { - #[cfg(feature = "tracing")] - let _span = - tracing::info_span!("optimize_async", n_trials, direction = ?self.direction).entered(); - - for _ in 0..n_trials { - let trial = self.create_trial(); - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - - match objective(trial).await { - Ok((trial, value)) => { - self.complete_trial(trial, value); - trace_info!(trial_id, "trial completed"); - } - Err(e) => { - // For async, we don't have the trial back on error - // We'll just count this as a failed trial without recording it - let _ = e.to_string(); - trace_debug!(trial_id, "trial failed"); - } - } - } - - // Return error if no trials completed successfully - let has_complete = self - .storage - .trials_arc() - .read() - .iter() - .any(|t| t.state == TrialState::Complete); - if !has_complete { - return Err(crate::Error::NoCompletedTrials); - } - - Ok(()) - } - - /// Run optimization with bounded parallelism for concurrent trial evaluation. - /// - /// This method runs up to `concurrency` trials simultaneously, allowing - /// efficient use of async I/O-bound objective functions. A semaphore limits - /// the number of concurrent evaluations. - /// - /// The objective function takes ownership of the `Trial` and must return it - /// along with the result. This allows async operations to use the trial - /// across await points. - /// - /// # Arguments - /// - /// * `n_trials` - The total number of trials to run. - /// * `concurrency` - The maximum number of trials to run simultaneously. - /// * `objective` - A function that takes a `Trial` and returns a `Future` - /// that resolves to a tuple of `(Trial, V)` or an error. - /// - /// # Errors - /// - /// Returns `Error::NoCompletedTrials` if all trials failed (no successful trials). - /// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics. - /// - /// # Examples - /// - /// ``` - /// use optimizer::parameter::{FloatParam, Parameter}; - /// use optimizer::sampler::random::RandomSampler; - /// use optimizer::{Direction, Study}; - /// - /// # #[cfg(feature = "async")] - /// # async fn example() -> optimizer::Result<()> { - /// // Minimize x^2 with parallel async evaluation - /// let sampler = RandomSampler::with_seed(42); - /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); - /// - /// let x_param = FloatParam::new(-10.0, 10.0); - /// - /// study - /// .optimize_parallel(10, 4, move |mut trial| { - /// let x_param = x_param.clone(); - /// async move { - /// let x = x_param.suggest(&mut trial)?; - /// // Async objective function (e.g., network request) - /// let value = x * x; - /// Ok::<_, optimizer::Error>((trial, value)) - /// } - /// }) - /// .await?; - /// - /// // All trials should have completed - /// assert_eq!(study.n_trials(), 10); - /// # Ok(()) - /// # } - /// ``` - #[cfg(feature = "async")] - pub async fn optimize_parallel( - &self, - n_trials: usize, - concurrency: usize, - objective: F, - ) -> crate::Result<()> - where - F: Fn(Trial) -> Fut + Send + Sync + 'static, - Fut: Future> + Send, - E: ToString + Send + 'static, - V: Send + 'static, - { - use tokio::sync::Semaphore; - - #[cfg(feature = "tracing")] - let _span = tracing::info_span!("optimize_parallel", n_trials, concurrency, direction = ?self.direction).entered(); - - let semaphore = Arc::new(Semaphore::new(concurrency)); - let objective = Arc::new(objective); - - let mut handles = Vec::with_capacity(n_trials); - - for _ in 0..n_trials { - let permit = semaphore - .clone() - .acquire_owned() - .await - .map_err(|e| crate::Error::TaskError(e.to_string()))?; - let trial = self.create_trial(); - let objective = Arc::clone(&objective); - - let handle = tokio::spawn(async move { - let result = objective(trial).await; - drop(permit); // Release semaphore permit when done - result - }); - - handles.push(handle); - } - - // Wait for all tasks and record results - for handle in handles { - match handle - .await - .map_err(|e| crate::Error::TaskError(e.to_string()))? - { - Ok((trial, value)) => { - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - self.complete_trial(trial, value); - trace_info!(trial_id, "trial completed"); - } - Err(e) => { - let _ = e.to_string(); - } - } - } - - // Return error if no trials completed successfully - let has_complete = self - .storage - .trials_arc() - .read() - .iter() - .any(|t| t.state == TrialState::Complete); - if !has_complete { - return Err(crate::Error::NoCompletedTrials); - } - - Ok(()) - } - - /// Run optimization with a callback for monitoring progress. - /// - /// This method is similar to `optimize`, but calls a callback function after - /// each completed trial. The callback can inspect the study state and the - /// completed trial, and can optionally stop optimization early by returning - /// `ControlFlow::Break(())`. - /// - /// # Arguments - /// - /// * `n_trials` - The maximum number of trials to run. - /// * `objective` - A closure that takes a mutable reference to a `Trial` and - /// returns the objective value or an error. - /// * `callback` - A closure called after each successful trial. Returns - /// `ControlFlow::Continue(())` to proceed or `ControlFlow::Break(())` to stop. - /// - /// # Errors - /// - /// Returns `Error::NoCompletedTrials` if no trials completed successfully - /// before optimization stopped (either by completing all trials or early stopping). - /// Returns `Error::Internal` if a completed trial is not found after adding (internal invariant violation). - /// - /// # Examples - /// - /// ``` - /// use std::ops::ControlFlow; - /// - /// use optimizer::parameter::{FloatParam, Parameter}; - /// use optimizer::sampler::random::RandomSampler; - /// use optimizer::{Direction, Study}; - /// - /// // Stop early when we find a good enough value - /// let sampler = RandomSampler::with_seed(42); - /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); - /// - /// let x_param = FloatParam::new(-10.0, 10.0); - /// - /// study - /// .optimize_with_callback( - /// 100, - /// |trial| { - /// let x = x_param.suggest(trial)?; - /// Ok::<_, optimizer::Error>(x * x) - /// }, - /// |_study, completed_trial| { - /// // Stop early if we find a value less than 1.0 - /// if completed_trial.value < 1.0 { - /// ControlFlow::Break(()) - /// } else { - /// ControlFlow::Continue(()) - /// } - /// }, - /// ) - /// .unwrap(); - /// - /// // May have stopped early, but should have at least one trial - /// assert!(study.n_trials() > 0); - /// ``` - pub fn optimize_with_callback( - &self, - n_trials: usize, - mut objective: F, - mut callback: C, - ) -> crate::Result<()> - where - V: Clone + Default, - F: FnMut(&mut Trial) -> core::result::Result, - C: FnMut(&Study, &CompletedTrial) -> ControlFlow<()>, - E: ToString + 'static, - { - #[cfg(feature = "tracing")] - let _span = - tracing::info_span!("optimize", n_trials, direction = ?self.direction).entered(); - - for _ in 0..n_trials { - let mut trial = self.create_trial(); - - match objective(&mut trial) { - Ok(value) => { - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - self.complete_trial(trial, value); - - #[cfg(feature = "tracing")] - { - tracing::info!(trial_id, "trial completed"); - let trials = self.storage.trials_arc().read(); - if trials - .iter() - .filter(|t| t.state == TrialState::Complete) - .count() - == 1 - || trials.last().map(|t| t.id) == self.best_id(&trials) - { - tracing::info!(trial_id, "new best value found"); - } - } - - // Get the just-completed trial for the callback - let trials = self.storage.trials_arc().read(); - let Some(completed) = trials.last() else { - return Err(crate::Error::Internal( - "completed trial not found after adding", - )); - }; - - // Call the callback and check if we should stop - // Note: We need to drop the read lock before calling callback - // to avoid potential deadlock if callback accesses the study - let completed_clone = completed.clone(); - drop(trials); - - if let ControlFlow::Break(()) = callback(self, &completed_clone) { - break; - } - } - Err(e) => { - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - if is_trial_pruned(&e) { - self.prune_trial(trial); - trace_info!(trial_id, "trial pruned"); - } else { - self.fail_trial(trial, e.to_string()); - trace_debug!(trial_id, "trial failed"); - } - } - } - } - - // Return error if no trials completed successfully - let has_complete = self - .storage - .trials_arc() - .read() - .iter() - .any(|t| t.state == TrialState::Complete); - if !has_complete { - return Err(crate::Error::NoCompletedTrials); - } - - Ok(()) - } - /// Run optimization until the given duration has elapsed. - /// - /// Trials that are already running when the timeout is reached will - /// complete — we never interrupt mid-trial. The actual elapsed time - /// may therefore slightly exceed the specified duration. - /// - /// # Arguments - /// - /// * `duration` - The maximum wall-clock time to spend on optimization. - /// * `objective` - A closure that takes a mutable reference to a `Trial` and - /// returns the objective value or an error. - /// - /// # Errors - /// - /// Returns `Error::NoCompletedTrials` if no trials completed successfully - /// before the timeout. - /// - /// # Examples - /// - /// ``` - /// use std::time::Duration; - /// - /// use optimizer::parameter::{FloatParam, Parameter}; - /// use optimizer::sampler::random::RandomSampler; - /// use optimizer::{Direction, Study}; - /// - /// let sampler = RandomSampler::with_seed(42); - /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); - /// - /// let x_param = FloatParam::new(-10.0, 10.0); - /// - /// study - /// .optimize_until(Duration::from_millis(100), |trial| { - /// let x = x_param.suggest(trial)?; - /// Ok::<_, optimizer::Error>(x * x) - /// }) - /// .unwrap(); - /// - /// assert!(study.n_trials() > 0); - /// ``` - pub fn optimize_until(&self, duration: Duration, mut objective: F) -> crate::Result<()> - where - F: FnMut(&mut Trial) -> core::result::Result, - E: ToString + 'static, - V: Default, - { - #[cfg(feature = "tracing")] - let _span = tracing::info_span!("optimize", duration_secs = duration.as_secs(), direction = ?self.direction).entered(); - - let deadline = Instant::now() + duration; - while Instant::now() < deadline { - let mut trial = self.create_trial(); - match objective(&mut trial) { Ok(value) => { #[cfg(feature = "tracing")] @@ -1420,209 +930,16 @@ where self.complete_trial(trial, value); trace_info!(trial_id, "trial completed"); } - Err(e) => { + Err(e) if is_trial_pruned(&e) => { #[cfg(feature = "tracing")] let trial_id = trial.id(); - if is_trial_pruned(&e) { - self.prune_trial(trial); - trace_info!(trial_id, "trial pruned"); - } else { - self.fail_trial(trial, e.to_string()); - trace_debug!(trial_id, "trial failed"); - } - } - } - } - - let has_complete = self - .storage - .trials_arc() - .read() - .iter() - .any(|t| t.state == TrialState::Complete); - if !has_complete { - return Err(crate::Error::NoCompletedTrials); - } - - Ok(()) - } - - /// Run optimization until the given duration has elapsed, with a callback. - /// - /// Like [`optimize_until`](Self::optimize_until), but calls a callback after - /// each completed trial. The callback can stop optimization early by returning - /// `ControlFlow::Break(())`. - /// - /// # Arguments - /// - /// * `duration` - The maximum wall-clock time to spend on optimization. - /// * `objective` - A closure that takes a mutable reference to a `Trial` and - /// returns the objective value or an error. - /// * `callback` - A closure called after each successful trial. Returns - /// `ControlFlow::Continue(())` to proceed or `ControlFlow::Break(())` to stop. - /// - /// # Errors - /// - /// Returns `Error::NoCompletedTrials` if no trials completed successfully. - /// Returns `Error::Internal` if a completed trial is not found after adding. - /// - /// # Examples - /// - /// ``` - /// use std::ops::ControlFlow; - /// use std::time::Duration; - /// - /// use optimizer::parameter::{FloatParam, Parameter}; - /// use optimizer::sampler::random::RandomSampler; - /// use optimizer::{Direction, Study}; - /// - /// let sampler = RandomSampler::with_seed(42); - /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); - /// - /// let x_param = FloatParam::new(-10.0, 10.0); - /// - /// study - /// .optimize_until_with_callback( - /// Duration::from_secs(1), - /// |trial| { - /// let x = x_param.suggest(trial)?; - /// Ok::<_, optimizer::Error>(x * x) - /// }, - /// |_study, completed_trial| { - /// if completed_trial.value < 1.0 { - /// ControlFlow::Break(()) - /// } else { - /// ControlFlow::Continue(()) - /// } - /// }, - /// ) - /// .unwrap(); - /// - /// assert!(study.n_trials() > 0); - /// ``` - pub fn optimize_until_with_callback( - &self, - duration: Duration, - mut objective: F, - mut callback: C, - ) -> crate::Result<()> - where - V: Clone + Default, - F: FnMut(&mut Trial) -> core::result::Result, - C: FnMut(&Study, &CompletedTrial) -> ControlFlow<()>, - E: ToString + 'static, - { - #[cfg(feature = "tracing")] - let _span = tracing::info_span!("optimize", duration_secs = duration.as_secs(), direction = ?self.direction).entered(); - - let deadline = Instant::now() + duration; - while Instant::now() < deadline { - let mut trial = self.create_trial(); - - match objective(&mut trial) { - Ok(value) => { - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - self.complete_trial(trial, value); - - #[cfg(feature = "tracing")] - { - tracing::info!(trial_id, "trial completed"); - let trials = self.storage.trials_arc().read(); - if trials - .iter() - .filter(|t| t.state == TrialState::Complete) - .count() - == 1 - || trials.last().map(|t| t.id) == self.best_id(&trials) - { - tracing::info!(trial_id, "new best value found"); - } - } - - let trials = self.storage.trials_arc().read(); - let Some(completed) = trials.last() else { - return Err(crate::Error::Internal( - "completed trial not found after adding", - )); - }; - - let completed_clone = completed.clone(); - drop(trials); - - if let ControlFlow::Break(()) = callback(self, &completed_clone) { - break; - } + self.prune_trial(trial); + trace_info!(trial_id, "trial pruned"); } Err(e) => { #[cfg(feature = "tracing")] let trial_id = trial.id(); - if is_trial_pruned(&e) { - self.prune_trial(trial); - trace_info!(trial_id, "trial pruned"); - } else { - self.fail_trial(trial, e.to_string()); - trace_debug!(trial_id, "trial failed"); - } - } - } - } - - let has_complete = self - .storage - .trials_arc() - .read() - .iter() - .any(|t| t.state == TrialState::Complete); - if !has_complete { - return Err(crate::Error::NoCompletedTrials); - } - - Ok(()) - } - - /// Run optimization asynchronously until the given duration has elapsed. - /// - /// The async variant of [`optimize_until`](Self::optimize_until). Trials are - /// run sequentially, but the objective function can be async (useful for - /// I/O-bound evaluations). - /// - /// # Arguments - /// - /// * `duration` - The maximum wall-clock time to spend on optimization. - /// * `objective` - A function that takes a `Trial` and returns a `Future` - /// that resolves to a tuple of `(Trial, V)` or an error. - /// - /// # Errors - /// - /// Returns `Error::NoCompletedTrials` if no trials completed successfully. - #[cfg(feature = "async")] - pub async fn optimize_until_async( - &self, - duration: Duration, - objective: F, - ) -> crate::Result<()> - where - F: Fn(Trial) -> Fut, - Fut: Future>, - E: ToString, - { - #[cfg(feature = "tracing")] - let _span = tracing::info_span!("optimize_until_async", duration_secs = duration.as_secs(), direction = ?self.direction).entered(); - - let deadline = Instant::now() + duration; - while Instant::now() < deadline { - let trial = self.create_trial(); - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - - match objective(trial).await { - Ok((trial, value)) => { - self.complete_trial(trial, value); - trace_info!(trial_id, "trial completed"); - } - Err(e) => { - let _ = e.to_string(); + self.fail_trial(trial, e.to_string()); trace_debug!(trial_id, "trial failed"); } } @@ -1641,112 +958,13 @@ where Ok(()) } - /// Run optimization with bounded parallelism until the given duration has elapsed. + /// Run optimization with an [`Objective`](crate::Objective) implementation. /// - /// The parallel variant of [`optimize_until`](Self::optimize_until). Runs up to - /// `concurrency` trials simultaneously using async tasks. New trials are spawned - /// as long as the deadline has not been reached; trials already running when the - /// deadline passes will complete. - /// - /// # Arguments - /// - /// * `duration` - The maximum wall-clock time to spend spawning new trials. - /// * `concurrency` - The maximum number of trials to run simultaneously. - /// * `objective` - A function that takes a `Trial` and returns a `Future` - /// that resolves to a tuple of `(Trial, V)` or an error. - /// - /// # Errors - /// - /// Returns `Error::NoCompletedTrials` if no trials completed successfully. - /// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics. - #[cfg(feature = "async")] - pub async fn optimize_until_parallel( - &self, - duration: Duration, - concurrency: usize, - objective: F, - ) -> crate::Result<()> - where - F: Fn(Trial) -> Fut + Send + Sync + 'static, - Fut: Future> + Send, - E: ToString + Send + 'static, - V: Send + 'static, - { - use tokio::sync::Semaphore; - - #[cfg(feature = "tracing")] - let _span = tracing::info_span!("optimize_until_parallel", duration_secs = duration.as_secs(), concurrency, direction = ?self.direction).entered(); - - let deadline = Instant::now() + duration; - let semaphore = Arc::new(Semaphore::new(concurrency)); - let objective = Arc::new(objective); - - let mut handles = Vec::new(); - - while Instant::now() < deadline { - let permit = semaphore - .clone() - .acquire_owned() - .await - .map_err(|e| crate::Error::TaskError(e.to_string()))?; - let trial = self.create_trial(); - let objective = Arc::clone(&objective); - - let handle = tokio::spawn(async move { - let result = objective(trial).await; - drop(permit); - result - }); - - handles.push(handle); - } - - for handle in handles { - match handle - .await - .map_err(|e| crate::Error::TaskError(e.to_string()))? - { - Ok((trial, value)) => { - #[cfg(feature = "tracing")] - let trial_id = trial.id(); - self.complete_trial(trial, value); - trace_info!(trial_id, "trial completed"); - } - Err(e) => { - let _ = e.to_string(); - } - } - } - - let has_complete = self - .storage - .trials_arc() - .read() - .iter() - .any(|t| t.state == TrialState::Complete); - if !has_complete { - return Err(crate::Error::NoCompletedTrials); - } - - Ok(()) - } - - /// Run optimization with automatic retry for failed trials. - /// - /// If the objective function returns an error, the same parameter - /// configuration is retried up to `max_retries` times. Only after all - /// retries are exhausted is the trial recorded as permanently failed. - /// - /// `n_trials` counts unique parameter configurations, not total - /// evaluations. A trial retried 3 times still counts as 1 toward the - /// `n_trials` limit. - /// - /// # Arguments - /// - /// * `n_trials` - The number of unique configurations to evaluate. - /// * `max_retries` - Maximum retry attempts per failed trial. - /// * `objective` - A closure that takes a mutable reference to a `Trial` - /// and returns the objective value or an error. + /// Like [`optimize`](Self::optimize), but accepts a struct implementing + /// [`Objective`](crate::Objective) for lifecycle hooks + /// ([`before_trial`](crate::Objective::before_trial), + /// [`after_trial`](crate::Objective::after_trial)) and automatic retries + /// ([`max_retries`](crate::Objective::max_retries)). /// /// # Errors /// @@ -1755,59 +973,98 @@ where /// # Examples /// /// ``` - /// use optimizer::parameter::{FloatParam, Parameter}; - /// use optimizer::sampler::random::RandomSampler; - /// use optimizer::{Direction, Study}; + /// use std::ops::ControlFlow; /// - /// let sampler = RandomSampler::with_seed(42); - /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); - /// let x_param = FloatParam::new(-10.0, 10.0); + /// use optimizer::prelude::*; /// - /// let call_count = std::cell::Cell::new(0u32); - /// study - /// .optimize_with_retries(5, 2, |trial| { - /// let x = x_param.suggest(trial)?; - /// call_count.set(call_count.get() + 1); - /// // Fail once every other call to exercise retry - /// if call_count.get() % 2 == 0 { - /// Err::(optimizer::Error::Internal("transient")) + /// struct QuadraticObj { + /// x: FloatParam, + /// target: f64, + /// } + /// + /// impl Objective for QuadraticObj { + /// type Error = Error; + /// fn evaluate(&self, trial: &mut Trial) -> Result { + /// let v = self.x.suggest(trial)?; + /// Ok((v - 3.0).powi(2)) + /// } + /// fn after_trial(&self, _: &Study, t: &CompletedTrial) -> ControlFlow<()> { + /// if t.value < self.target { + /// ControlFlow::Break(()) /// } else { - /// Ok(x * x) + /// ControlFlow::Continue(()) /// } - /// }) - /// .unwrap(); + /// } + /// } /// - /// assert_eq!(study.n_trials(), 5); + /// let study: Study = Study::new(Direction::Minimize); + /// let obj = QuadraticObj { + /// x: FloatParam::new(-10.0, 10.0), + /// target: 1.0, + /// }; + /// study.optimize_with(200, obj).unwrap(); + /// assert!(study.best_value().unwrap() < 1.0); /// ``` - pub fn optimize_with_retries( + #[allow(clippy::needless_pass_by_value)] + pub fn optimize_with( &self, n_trials: usize, - max_retries: usize, - mut objective: F, + objective: impl crate::objective::Objective, ) -> crate::Result<()> where - F: FnMut(&mut Trial) -> core::result::Result, - E: ToString + 'static, - V: Default, + V: Clone + Default, { #[cfg(feature = "tracing")] - let _span = tracing::info_span!("optimize_with_retries", n_trials, max_retries, direction = ?self.direction).entered(); + let _span = + tracing::info_span!("optimize_with", n_trials, direction = ?self.direction).entered(); + + let max_retries = objective.max_retries(); for _ in 0..n_trials { + if let ControlFlow::Break(()) = objective.before_trial(self) { + break; + } + let mut trial = self.create_trial(); let mut retries = 0; loop { - match objective(&mut trial) { + match objective.evaluate(&mut trial) { Ok(value) => { #[cfg(feature = "tracing")] let trial_id = trial.id(); self.complete_trial(trial, value); - trace_info!(trial_id, "trial completed"); + + #[cfg(feature = "tracing")] + { + tracing::info!(trial_id, "trial completed"); + let trials = self.storage.trials_arc().read(); + if trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .count() + == 1 + || trials.last().map(|t| t.id) == self.best_id(&trials) + { + tracing::info!(trial_id, "new best value found"); + } + } + + // Fire after_trial hook + let trials = self.storage.trials_arc().read(); + if let Some(completed) = trials.last() { + let completed_clone = completed.clone(); + drop(trials); + if let ControlFlow::Break(()) = + objective.after_trial(self, &completed_clone) + { + // Return early — at least one trial completed. + return Ok(()); + } + } break; } - Err(_) if retries < max_retries => { + Err(e) if !is_trial_pruned(&e) && retries < max_retries => { retries += 1; - // Create a new trial with the same parameters trial = self.create_trial_with_params(trial.params().clone()); } Err(e) => { @@ -1818,7 +1075,7 @@ where trace_info!(trial_id, "trial pruned"); } else { self.fail_trial(trial, e.to_string()); - trace_debug!(trial_id, "trial permanently failed"); + trace_debug!(trial_id, "trial failed"); } break; } @@ -1839,6 +1096,489 @@ where Ok(()) } + + /// Run async optimization with a closure. + /// + /// Each evaluation is wrapped in + /// [`spawn_blocking`](tokio::task::spawn_blocking), keeping the async + /// runtime responsive for CPU-bound objectives. Trials run sequentially. + /// + /// For lifecycle hooks, use [`optimize_with_async`](Self::optimize_with_async). + /// + /// # Errors + /// + /// Returns `Error::NoCompletedTrials` if no trials completed successfully. + /// Returns `Error::TaskError` if a spawned blocking task panics. + /// + /// # Examples + /// + /// ``` + /// use optimizer::parameter::{FloatParam, Parameter}; + /// use optimizer::sampler::random::RandomSampler; + /// use optimizer::{Direction, Study}; + /// + /// # #[cfg(feature = "async")] + /// # async fn example() -> optimizer::Result<()> { + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); + /// let x_param = FloatParam::new(-10.0, 10.0); + /// + /// study + /// .optimize_async(10, move |trial: &mut optimizer::Trial| { + /// let x = x_param.suggest(trial)?; + /// Ok::<_, optimizer::Error>(x * x) + /// }) + /// .await?; + /// + /// assert!(study.n_trials() > 0); + /// # Ok(()) + /// # } + /// ``` + #[cfg(feature = "async")] + pub async fn optimize_async(&self, n_trials: usize, objective: F) -> crate::Result<()> + where + F: Fn(&mut Trial) -> Result + Send + Sync + 'static, + E: ToString + Send + 'static, + V: Clone + Default + Send + 'static, + { + #[cfg(feature = "tracing")] + let _span = + tracing::info_span!("optimize_async", n_trials, direction = ?self.direction).entered(); + + let objective = Arc::new(objective); + + for _ in 0..n_trials { + let obj = Arc::clone(&objective); + let mut trial = self.create_trial(); + let result = tokio::task::spawn_blocking(move || { + let res = obj(&mut trial); + (trial, res) + }) + .await + .map_err(|e| crate::Error::TaskError(e.to_string()))?; + + match result { + (t, Ok(value)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.complete_trial(t, value); + trace_info!(trial_id, "trial completed"); + } + (t, Err(e)) if is_trial_pruned(&e) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.prune_trial(t); + trace_info!(trial_id, "trial pruned"); + } + (t, Err(e)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.fail_trial(t, e.to_string()); + trace_debug!(trial_id, "trial failed"); + } + } + } + + let has_complete = self + .storage + .trials_arc() + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } + + /// Run async optimization with an [`Objective`](crate::Objective) implementation. + /// + /// Like [`optimize_async`](Self::optimize_async), but accepts a struct + /// implementing [`Objective`](crate::Objective) for lifecycle hooks and + /// automatic retries. + /// + /// # Errors + /// + /// Returns `Error::NoCompletedTrials` if no trials completed successfully. + /// Returns `Error::TaskError` if a spawned blocking task panics. + #[cfg(feature = "async")] + pub async fn optimize_with_async(&self, n_trials: usize, objective: O) -> crate::Result<()> + where + O: crate::objective::Objective + Send + Sync + 'static, + O::Error: Send, + V: Clone + Default + Send + 'static, + { + #[cfg(feature = "tracing")] + let _span = + tracing::info_span!("optimize_with_async", n_trials, direction = ?self.direction) + .entered(); + + let objective = Arc::new(objective); + let max_retries = objective.max_retries(); + + for _ in 0..n_trials { + if let ControlFlow::Break(()) = objective.before_trial(self) { + break; + } + + let mut trial = self.create_trial(); + let mut retries = 0; + loop { + let obj = Arc::clone(&objective); + let result = tokio::task::spawn_blocking(move || { + let res = obj.evaluate(&mut trial); + (trial, res) + }) + .await + .map_err(|e| crate::Error::TaskError(e.to_string()))?; + + match result { + (t, Ok(value)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.complete_trial(t, value); + trace_info!(trial_id, "trial completed"); + + // Fire after_trial hook + let trials = self.storage.trials_arc().read(); + if let Some(completed) = trials.last() { + let completed_clone = completed.clone(); + drop(trials); + if let ControlFlow::Break(()) = + objective.after_trial(self, &completed_clone) + { + return Ok(()); + } + } + break; + } + (t, Err(e)) if !is_trial_pruned(&e) && retries < max_retries => { + retries += 1; + trial = self.create_trial_with_params(t.params().clone()); + } + (t, Err(e)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + if is_trial_pruned(&e) { + self.prune_trial(t); + trace_info!(trial_id, "trial pruned"); + } else { + self.fail_trial(t, e.to_string()); + trace_debug!(trial_id, "trial failed"); + } + break; + } + } + } + } + + let has_complete = self + .storage + .trials_arc() + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } + + /// Run parallel optimization with a closure. + /// + /// Spawns up to `concurrency` evaluations concurrently using + /// [`spawn_blocking`](tokio::task::spawn_blocking). Results are + /// collected via a [`JoinSet`](tokio::task::JoinSet). + /// + /// For lifecycle hooks, use + /// [`optimize_with_parallel`](Self::optimize_with_parallel). + /// + /// # Errors + /// + /// Returns `Error::NoCompletedTrials` if no trials completed successfully. + /// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics. + /// + /// # Examples + /// + /// ``` + /// use optimizer::parameter::{FloatParam, Parameter}; + /// use optimizer::sampler::random::RandomSampler; + /// use optimizer::{Direction, Study}; + /// + /// # #[cfg(feature = "async")] + /// # async fn example() -> optimizer::Result<()> { + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); + /// let x_param = FloatParam::new(-10.0, 10.0); + /// + /// study + /// .optimize_parallel(10, 4, move |trial: &mut optimizer::Trial| { + /// let x = x_param.suggest(trial)?; + /// Ok::<_, optimizer::Error>(x * x) + /// }) + /// .await?; + /// + /// assert_eq!(study.n_trials(), 10); + /// # Ok(()) + /// # } + /// ``` + #[cfg(feature = "async")] + #[allow(clippy::missing_panics_doc)] + pub async fn optimize_parallel( + &self, + n_trials: usize, + concurrency: usize, + objective: F, + ) -> crate::Result<()> + where + F: Fn(&mut Trial) -> Result + Send + Sync + 'static, + E: ToString + Send + 'static, + V: Clone + Default + Send + 'static, + { + use tokio::sync::Semaphore; + use tokio::task::JoinSet; + + #[cfg(feature = "tracing")] + let _span = tracing::info_span!("optimize_parallel", n_trials, concurrency, direction = ?self.direction).entered(); + + let objective = Arc::new(objective); + let semaphore = Arc::new(Semaphore::new(concurrency)); + let mut join_set: JoinSet<(Trial, Result)> = JoinSet::new(); + let mut spawned = 0; + + while spawned < n_trials { + // If the join set is full, drain one result to free a slot. + while join_set.len() >= concurrency { + let result = join_set + .join_next() + .await + .expect("join_set should not be empty") + .map_err(|e| crate::Error::TaskError(e.to_string()))?; + match result { + (t, Ok(value)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.complete_trial(t, value); + trace_info!(trial_id, "trial completed"); + } + (t, Err(e)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + if is_trial_pruned(&e) { + self.prune_trial(t); + trace_info!(trial_id, "trial pruned"); + } else { + self.fail_trial(t, e.to_string()); + trace_debug!(trial_id, "trial failed"); + } + } + } + } + + let permit = semaphore + .clone() + .acquire_owned() + .await + .map_err(|e| crate::Error::TaskError(e.to_string()))?; + + let mut trial = self.create_trial(); + let obj = Arc::clone(&objective); + join_set.spawn(async move { + let result = tokio::task::spawn_blocking(move || { + let res = obj(&mut trial); + (trial, res) + }) + .await + .expect("spawn_blocking should not panic"); + drop(permit); + result + }); + spawned += 1; + } + + // Drain remaining in-flight tasks. + while let Some(result) = join_set.join_next().await { + let result = result.map_err(|e| crate::Error::TaskError(e.to_string()))?; + match result { + (t, Ok(value)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.complete_trial(t, value); + trace_info!(trial_id, "trial completed"); + } + (t, Err(e)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + if is_trial_pruned(&e) { + self.prune_trial(t); + trace_info!(trial_id, "trial pruned"); + } else { + self.fail_trial(t, e.to_string()); + trace_debug!(trial_id, "trial failed"); + } + } + } + } + + let has_complete = self + .storage + .trials_arc() + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } + + /// Run parallel optimization with an [`Objective`](crate::Objective) implementation. + /// + /// Like [`optimize_parallel`](Self::optimize_parallel), but accepts a struct + /// implementing [`Objective`](crate::Objective) for lifecycle hooks and + /// automatic retries. The [`after_trial`](crate::Objective::after_trial) + /// hook fires as each result arrives — returning `Break` stops spawning + /// new trials while in-flight tasks drain. + /// + /// # Errors + /// + /// Returns `Error::NoCompletedTrials` if no trials completed successfully. + /// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics. + #[cfg(feature = "async")] + #[allow(clippy::missing_panics_doc, clippy::too_many_lines)] + pub async fn optimize_with_parallel( + &self, + n_trials: usize, + concurrency: usize, + objective: O, + ) -> crate::Result<()> + where + O: crate::objective::Objective + Send + Sync + 'static, + O::Error: Send, + V: Clone + Default + Send + 'static, + { + use tokio::sync::Semaphore; + use tokio::task::JoinSet; + + #[cfg(feature = "tracing")] + let _span = tracing::info_span!("optimize_with_parallel", n_trials, concurrency, direction = ?self.direction).entered(); + + let objective = Arc::new(objective); + let semaphore = Arc::new(Semaphore::new(concurrency)); + let mut join_set: JoinSet<(Trial, Result)> = JoinSet::new(); + let mut spawned = 0; + + 'spawn: while spawned < n_trials { + if let ControlFlow::Break(()) = objective.before_trial(self) { + break; + } + + // If the join set is full, drain one result to free a slot. + while join_set.len() >= concurrency { + let result = join_set + .join_next() + .await + .expect("join_set should not be empty") + .map_err(|e| crate::Error::TaskError(e.to_string()))?; + match result { + (t, Ok(value)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.complete_trial(t, value); + trace_info!(trial_id, "trial completed"); + + let trials = self.storage.trials_arc().read(); + if let Some(completed) = trials.last() { + let completed_clone = completed.clone(); + drop(trials); + if let ControlFlow::Break(()) = + objective.after_trial(self, &completed_clone) + { + break 'spawn; + } + } + } + (t, Err(e)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + if is_trial_pruned(&e) { + self.prune_trial(t); + trace_info!(trial_id, "trial pruned"); + } else { + self.fail_trial(t, e.to_string()); + trace_debug!(trial_id, "trial failed"); + } + } + } + } + + let permit = semaphore + .clone() + .acquire_owned() + .await + .map_err(|e| crate::Error::TaskError(e.to_string()))?; + + let mut trial = self.create_trial(); + let obj = Arc::clone(&objective); + join_set.spawn(async move { + let result = tokio::task::spawn_blocking(move || { + let res = obj.evaluate(&mut trial); + (trial, res) + }) + .await + .expect("spawn_blocking should not panic"); + drop(permit); + result + }); + spawned += 1; + } + + // Drain remaining in-flight tasks. + while let Some(result) = join_set.join_next().await { + let result = result.map_err(|e| crate::Error::TaskError(e.to_string()))?; + match result { + (t, Ok(value)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + self.complete_trial(t, value); + trace_info!(trial_id, "trial completed"); + // Still fire after_trial for bookkeeping, but don't break — we're draining. + let trials = self.storage.trials_arc().read(); + if let Some(completed) = trials.last() { + let completed_clone = completed.clone(); + drop(trials); + let _ = objective.after_trial(self, &completed_clone); + } + } + (t, Err(e)) => { + #[cfg(feature = "tracing")] + let trial_id = t.id(); + if is_trial_pruned(&e) { + self.prune_trial(t); + trace_info!(trial_id, "trial pruned"); + } else { + self.fail_trial(t, e.to_string()); + trace_debug!(trial_id, "trial failed"); + } + } + } + } + + let has_complete = self + .storage + .trials_arc() + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } } impl Study @@ -2324,111 +2064,6 @@ where } } -// Specialized implementation for Study that provides deprecated `_with_sampler` aliases. -// -// For Study, the generic methods from `impl Study` (like `optimize()`, -// `create_trial()`) now automatically use the sampler via the `trial_factory`. -// The `_with_sampler` method names are deprecated in favor of the generic names. -#[allow(clippy::missing_errors_doc)] -impl Study { - /// Deprecated: use `create_trial()` instead. - /// - /// The generic `create_trial()` now automatically integrates with the sampler - /// for `Study`. - #[deprecated( - since = "0.2.0", - note = "use `create_trial()` instead — it now uses the sampler automatically for Study" - )] - #[must_use] - pub fn create_trial_with_sampler(&self) -> Trial { - self.create_trial() - } - - /// Deprecated: use `optimize()` instead. - /// - /// The generic `optimize()` now automatically integrates with the sampler - /// for `Study`. - #[deprecated( - since = "0.2.0", - note = "use `optimize()` instead — it now uses the sampler automatically for Study" - )] - pub fn optimize_with_sampler(&self, n_trials: usize, objective: F) -> crate::Result<()> - where - F: FnMut(&mut Trial) -> core::result::Result, - E: ToString + 'static, - { - self.optimize(n_trials, objective) - } - - /// Deprecated: use `optimize_with_callback()` instead. - /// - /// The generic `optimize_with_callback()` now automatically integrates with the - /// sampler for `Study`. - #[deprecated( - since = "0.2.0", - note = "use `optimize_with_callback()` instead — it now uses the sampler automatically for Study" - )] - pub fn optimize_with_callback_sampler( - &self, - n_trials: usize, - objective: F, - callback: C, - ) -> crate::Result<()> - where - F: FnMut(&mut Trial) -> core::result::Result, - C: FnMut(&Study, &CompletedTrial) -> ControlFlow<()>, - E: ToString + 'static, - { - self.optimize_with_callback(n_trials, objective, callback) - } - - /// Deprecated: use `optimize_async()` instead. - /// - /// The generic `optimize_async()` now automatically integrates with the sampler - /// for `Study`. - #[cfg(feature = "async")] - #[deprecated( - since = "0.2.0", - note = "use `optimize_async()` instead — it now uses the sampler automatically for Study" - )] - pub async fn optimize_async_with_sampler( - &self, - n_trials: usize, - objective: F, - ) -> crate::Result<()> - where - F: Fn(Trial) -> Fut, - Fut: Future>, - E: ToString, - { - self.optimize_async(n_trials, objective).await - } - - /// Deprecated: use `optimize_parallel()` instead. - /// - /// The generic `optimize_parallel()` now automatically integrates with the - /// sampler for `Study`. - #[cfg(feature = "async")] - #[deprecated( - since = "0.2.0", - note = "use `optimize_parallel()` instead — it now uses the sampler automatically for Study" - )] - pub async fn optimize_parallel_with_sampler( - &self, - n_trials: usize, - concurrency: usize, - objective: F, - ) -> crate::Result<()> - where - F: Fn(Trial) -> Fut + Send + Sync + 'static, - Fut: Future> + Send, - E: ToString + Send + 'static, - { - self.optimize_parallel(n_trials, concurrency, objective) - .await - } -} - impl Study { /// Create a study with a custom sampler, pruner, and storage backend. /// @@ -2719,40 +2354,6 @@ impl Study { } } -#[cfg(feature = "serde")] -impl Study { - /// Run optimization with automatic checkpointing every `interval` trials. - /// - /// This is convenience sugar over [`optimize_with_callback`](Self::optimize_with_callback) - /// combined with [`save`](Self::save). The checkpoint is written atomically so - /// a crash mid-write will never leave a corrupt file. - /// - /// # Errors - /// - /// Returns an error if the optimization itself fails (see - /// [`optimize`](Self::optimize) for details). Checkpoint I/O errors are - /// silently ignored (best-effort). - pub fn optimize_with_checkpoint( - &self, - n_trials: usize, - checkpoint_interval: usize, - checkpoint_path: impl AsRef, - objective: F, - ) -> crate::Result<()> - where - F: FnMut(&mut Trial) -> core::result::Result, - E: ToString + 'static, - { - let path = checkpoint_path.as_ref().to_owned(); - self.optimize_with_callback(n_trials, objective, |study, _trial| { - if study.n_trials().is_multiple_of(checkpoint_interval) { - let _ = study.save(&path); - } - ControlFlow::Continue(()) - }) - } -} - #[cfg(feature = "serde")] impl Study { /// Load a study from a JSON file. diff --git a/tests/async_tests.rs b/tests/async_tests.rs index b99b1df..ba9754f 100644 --- a/tests/async_tests.rs +++ b/tests/async_tests.rs @@ -17,12 +17,9 @@ async fn test_optimize_async_basic() { let x_param = FloatParam::new(-10.0, 10.0); study - .optimize_async(10, move |mut trial| { - let x_param = x_param.clone(); - async move { - let x = x_param.suggest(&mut trial)?; - Ok::<_, Error>((trial, x * x)) - } + .optimize_async(10, move |trial: &mut optimizer::Trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x * x) }) .await .expect("async optimization should succeed"); @@ -45,12 +42,9 @@ async fn test_optimize_async_with_tpe() { let x_param = FloatParam::new(-5.0, 5.0); study - .optimize_async(15, move |mut trial| { - let x_param = x_param.clone(); - async move { - let x = x_param.suggest(&mut trial)?; - Ok::<_, Error>((trial, x * x)) - } + .optimize_async(15, move |trial: &mut optimizer::Trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x * x) }) .await .expect("async optimization with sampler should succeed"); @@ -68,12 +62,9 @@ async fn test_optimize_parallel() { let x_param = FloatParam::new(-10.0, 10.0); study - .optimize_parallel(20, 4, move |mut trial| { - let x_param = x_param.clone(); - async move { - let x = x_param.suggest(&mut trial)?; - Ok::<_, Error>((trial, x * x)) - } + .optimize_parallel(20, 4, move |trial: &mut optimizer::Trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x * x) }) .await .expect("parallel optimization should succeed"); @@ -95,14 +86,10 @@ async fn test_optimize_parallel_with_tpe() { let y_param = FloatParam::new(-5.0, 5.0); study - .optimize_parallel(15, 3, move |mut trial| { - let x_param = x_param.clone(); - let y_param = y_param.clone(); - async move { - let x = x_param.suggest(&mut trial)?; - let y = y_param.suggest(&mut trial)?; - Ok::<_, Error>((trial, x * x + y * y)) - } + .optimize_parallel(15, 3, move |trial: &mut optimizer::Trial| { + let x = x_param.suggest(trial)?; + let y = y_param.suggest(trial)?; + Ok::<_, Error>(x * x + y * y) }) .await .expect("parallel optimization with sampler should succeed"); @@ -115,27 +102,8 @@ async fn test_optimize_async_all_failures() { let study: Study = Study::new(Direction::Minimize); let result = study - .optimize_async(5, |trial| async move { - let _ = trial; - Err::<(_, f64), &str>("always fails") - }) - .await; - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "should return NoCompletedTrials when all trials fail" - ); -} - -#[tokio::test] -#[allow(deprecated)] -async fn test_optimize_async_with_sampler_all_failures() { - let study: Study = Study::new(Direction::Minimize); - - let result = study - .optimize_async_with_sampler(5, |trial| async move { - let _ = trial; - Err::<(_, f64), &str>("always fails") + .optimize_async(5, |_trial: &mut optimizer::Trial| { + Err::("always fails") }) .await; @@ -150,27 +118,8 @@ async fn test_optimize_parallel_all_failures() { let study: Study = Study::new(Direction::Minimize); let result = study - .optimize_parallel(5, 2, |trial| async move { - let _ = trial; - Err::<(_, f64), &str>("always fails") - }) - .await; - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "should return NoCompletedTrials when all trials fail" - ); -} - -#[tokio::test] -#[allow(deprecated)] -async fn test_optimize_parallel_with_sampler_all_failures() { - let study: Study = Study::new(Direction::Minimize); - - let result = study - .optimize_parallel_with_sampler(5, 2, |trial| async move { - let _ = trial; - Err::<(_, f64), &str>("always fails") + .optimize_parallel(5, 2, |_trial: &mut optimizer::Trial| { + Err::("always fails") }) .await; @@ -190,16 +139,13 @@ async fn test_optimize_async_partial_failures() { let x_param = FloatParam::new(0.0, 10.0); study - .optimize_async(10, move |mut trial| { + .optimize_async(10, move |trial: &mut optimizer::Trial| { let count = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - let x_param = x_param.clone(); - async move { - if count.is_multiple_of(2) { - let x = x_param.suggest(&mut trial)?; - Ok::<_, Error>((trial, x)) - } else { - Err(Error::NoCompletedTrials) // Use as error type - } + if count.is_multiple_of(2) { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) + } else { + Err(Error::NoCompletedTrials) // Use as error type } }) .await @@ -218,12 +164,9 @@ async fn test_optimize_parallel_high_concurrency() { // Run with concurrency higher than n_trials study - .optimize_parallel(5, 10, move |mut trial| { - let x_param = x_param.clone(); - async move { - let x = x_param.suggest(&mut trial)?; - Ok::<_, Error>((trial, x)) - } + .optimize_parallel(5, 10, move |trial: &mut optimizer::Trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) }) .await .expect("should handle high concurrency"); @@ -240,12 +183,9 @@ async fn test_optimize_parallel_single_concurrency() { // Run with concurrency of 1 (sequential) study - .optimize_parallel(10, 1, move |mut trial| { - let x_param = x_param.clone(); - async move { - let x = x_param.suggest(&mut trial)?; - Ok::<_, Error>((trial, x)) - } + .optimize_parallel(10, 1, move |trial: &mut optimizer::Trial| { + let x = x_param.suggest(trial)?; + Ok::<_, Error>(x) }) .await .expect("should work with single concurrency"); diff --git a/tests/integration.rs b/tests/integration.rs index b895606..a231230 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -570,28 +570,36 @@ fn test_tpe_with_integer_parameters() { #[test] fn test_callback_early_stopping() { - use std::cell::Cell; use std::ops::ControlFlow; - let study: Study = Study::new(Direction::Minimize); - let trials_run = Cell::new(0); - let x_param = FloatParam::new(0.0, 10.0); + use optimizer::Objective; + use optimizer::sampler::CompletedTrial; + struct EarlyStopAfter5 { + x_param: FloatParam, + } + + impl Objective for EarlyStopAfter5 { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + Ok(x) + } + fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { + if study.n_trials() >= 5 { + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + } + } + + let study: Study = Study::new(Direction::Minimize); study - .optimize_with_callback( + .optimize_with( 100, - |trial| { - trials_run.set(trials_run.get() + 1); - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }, - |_study, _trial| { - // Stop after 5 trials - if trials_run.get() >= 5 { - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) - } + EarlyStopAfter5 { + x_param: FloatParam::new(0.0, 10.0), }, ) .expect("optimization should succeed"); @@ -769,48 +777,10 @@ fn test_optimize_all_trials_fail() { } #[test] -fn test_optimize_with_callback_all_trials_fail() { - use std::ops::ControlFlow; - +fn test_optimize_with_all_trials_fail() { let study: Study = Study::new(Direction::Minimize); - let result = study.optimize_with_callback( - 5, - |_trial| Err::("always fails"), - |_study, _trial| ControlFlow::Continue(()), - ); - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "should return NoCompletedTrials when all trials fail" - ); -} - -#[test] -#[allow(deprecated)] -fn test_optimize_with_sampler_all_trials_fail() { - let study: Study = Study::new(Direction::Minimize); - - let result = study.optimize_with_sampler(5, |_trial| Err::("always fails")); - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "should return NoCompletedTrials when all trials fail" - ); -} - -#[test] -#[allow(deprecated)] -fn test_optimize_with_callback_sampler_all_trials_fail() { - use std::ops::ControlFlow; - - let study: Study = Study::new(Direction::Minimize); - - let result = study.optimize_with_callback_sampler( - 5, - |_trial| Err::("always fails"), - |_study, _trial| ControlFlow::Continue(()), - ); + let result = study.optimize(5, |_trial| Err::("always fails")); assert!( matches!(result, Err(Error::NoCompletedTrials)), @@ -964,27 +934,6 @@ fn test_tpe_with_step_distributions() { assert!(best.value < 100.0, "should find reasonable solution"); } -#[test] -#[allow(deprecated)] -fn test_create_trial_vs_create_trial_with_sampler() { - let sampler = RandomSampler::with_seed(42); - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - - // create_trial() creates trial with sampler integration for Study - let trial1 = study.create_trial(); - assert_eq!(trial1.id(), 0); - - // create_trial_with_sampler() is deprecated but still works - let trial2 = study.create_trial_with_sampler(); - assert_eq!(trial2.id(), 1); - - // Both should work for suggesting parameters - let x_param = FloatParam::new(0.0, 1.0); - let mut trial3 = study.create_trial(); - let x = x_param.suggest(&mut trial3).unwrap(); - assert!((0.0..=1.0).contains(&x)); -} - #[test] fn test_manual_trial_completion() { let study: Study = Study::new(Direction::Minimize); @@ -1058,19 +1007,34 @@ fn test_tpe_empty_good_or_bad_values_fallback() { fn test_callback_early_stopping_on_first_trial() { use std::ops::ControlFlow; - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); + use optimizer::Objective; + use optimizer::sampler::CompletedTrial; + struct StopImmediately { + x_param: FloatParam, + } + + impl Objective for StopImmediately { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + Ok(x) + } + fn after_trial( + &self, + _study: &Study, + _trial: &CompletedTrial, + ) -> ControlFlow<()> { + ControlFlow::Break(()) + } + } + + let study: Study = Study::new(Direction::Minimize); study - .optimize_with_callback( + .optimize_with( 100, - |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }, - |_study, _trial| { - // Stop immediately after first trial - ControlFlow::Break(()) + StopImmediately { + x_param: FloatParam::new(0.0, 10.0), }, ) .expect("optimization should succeed"); @@ -1082,23 +1046,35 @@ fn test_callback_early_stopping_on_first_trial() { fn test_callback_sampler_early_stopping() { use std::ops::ControlFlow; + use optimizer::Objective; + use optimizer::sampler::CompletedTrial; + + struct StopAfter3 { + x_param: FloatParam, + } + + impl Objective for StopAfter3 { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + Ok(x) + } + fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { + if study.n_trials() >= 3 { + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + } + } + let sampler = RandomSampler::with_seed(42); let study: Study = Study::with_sampler(Direction::Minimize, sampler); - let x_param = FloatParam::new(0.0, 10.0); - study - .optimize_with_callback( + .optimize_with( 100, - |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }, - |study, _trial| { - if study.n_trials() >= 3 { - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) - } + StopAfter3 { + x_param: FloatParam::new(0.0, 10.0), }, ) .expect("optimization should succeed"); @@ -1364,165 +1340,6 @@ fn test_completed_trial_get() { assert!((1..=10).contains(&n_val)); } -// ============================================================================= -// Tests for timeout-based optimization -// ============================================================================= - -#[test] -fn test_optimize_until_runs_for_approximately_specified_duration() { - use std::time::{Duration, Instant}; - - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(-10.0, 10.0); - - let duration = Duration::from_millis(200); - let start = Instant::now(); - - study - .optimize_until(duration, |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x * x) - }) - .unwrap(); - - let elapsed = start.elapsed(); - assert!( - elapsed >= duration, - "should run for at least the specified duration, elapsed: {elapsed:?}" - ); - // Allow generous upper bound — the last trial may overshoot - assert!( - elapsed < duration + Duration::from_millis(200), - "should not overshoot excessively, elapsed: {elapsed:?}" - ); -} - -#[test] -fn test_optimize_until_completes_at_least_one_trial() { - use std::time::Duration; - - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(-10.0, 10.0); - - study - .optimize_until(Duration::from_millis(100), |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x * x) - }) - .unwrap(); - - assert!( - study.n_trials() >= 1, - "should complete at least one trial, got {}", - study.n_trials() - ); -} - -#[test] -fn test_optimize_until_works_with_minimize() { - use std::time::Duration; - - let sampler = RandomSampler::with_seed(42); - let study: Study = Study::with_sampler(Direction::Minimize, sampler); - let x_param = FloatParam::new(-10.0, 10.0); - - study - .optimize_until(Duration::from_millis(100), |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x * x) - }) - .unwrap(); - - let best = study.best_value().unwrap(); - assert!(best >= 0.0, "x^2 should be non-negative"); -} - -#[test] -fn test_optimize_until_works_with_maximize() { - use std::time::Duration; - - let sampler = RandomSampler::with_seed(42); - let study: Study = Study::with_sampler(Direction::Maximize, sampler); - let x_param = FloatParam::new(0.0, 10.0); - - study - .optimize_until(Duration::from_millis(100), |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }) - .unwrap(); - - let best = study.best_value().unwrap(); - assert!(best >= 0.0); -} - -#[test] -fn test_optimize_until_with_callback_early_stopping() { - use std::ops::ControlFlow; - use std::time::Duration; - - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - - study - .optimize_until_with_callback( - Duration::from_secs(10), // long timeout — callback should stop early - |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x) - }, - |study, _trial| { - if study.n_trials() >= 5 { - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) - } - }, - ) - .unwrap(); - - assert_eq!( - study.n_trials(), - 5, - "callback should have stopped after 5 trials" - ); -} - -#[test] -fn test_optimize_until_all_trials_fail() { - use std::time::Duration; - - let study: Study = Study::new(Direction::Minimize); - - let result = study.optimize_until(Duration::from_millis(50), |_trial| { - Err::("always fails") - }); - - assert!( - matches!(result, Err(Error::NoCompletedTrials)), - "should return NoCompletedTrials when all trials fail" - ); -} - -#[test] -fn test_optimize_until_with_non_f64_value_type() { - use std::time::Duration; - - let study: Study = Study::new(Direction::Minimize); - let x_param = IntParam::new(-10, 10); - - study - .optimize_until(Duration::from_millis(100), |trial| { - let x = x_param.suggest(trial)?; - Ok::<_, Error>(x.abs() as i32) - }) - .unwrap(); - - assert!(study.n_trials() >= 1); - let best = study.best_trial().unwrap(); - assert!(best.value >= 0); -} - // ============================================================================= // Tests for top_trials // ============================================================================= @@ -1966,55 +1783,111 @@ fn test_display_matches_summary() { } // ============================================================================= -// Tests: optimize_with_retries +// Tests: optimize_with retries via Objective trait // ============================================================================= #[test] fn test_retries_successful_trials_not_retried() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - let call_count = std::cell::Cell::new(0u32); + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; - study - .optimize_with_retries(5, 3, |trial| { - let x = x_param.suggest(trial)?; - call_count.set(call_count.get() + 1); - Ok::<_, Error>(x * x) - }) - .unwrap(); + use optimizer::Objective; + + struct SuccessObj { + x_param: FloatParam, + call_count: Arc, + } + + impl Objective for SuccessObj { + type Error = Error; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial)?; + self.call_count.fetch_add(1, Ordering::Relaxed); + Ok(x * x) + } + fn max_retries(&self) -> usize { + 3 + } + } + + let study: Study = Study::new(Direction::Minimize); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = SuccessObj { + x_param: FloatParam::new(0.0, 10.0), + call_count: Arc::clone(&call_count), + }; + + study.optimize_with(5, obj).unwrap(); // All trials succeed on first try — exactly 5 calls - assert_eq!(call_count.get(), 5); + assert_eq!(call_count.load(Ordering::Relaxed), 5); assert_eq!(study.n_trials(), 5); } #[test] fn test_retries_failed_trials_retried_up_to_max() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - let call_count = std::cell::Cell::new(0u32); + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; - let result = study.optimize_with_retries(1, 3, |trial| { - let _ = x_param.suggest(trial).unwrap(); - call_count.set(call_count.get() + 1); - Err::("always fails") - }); + use optimizer::Objective; + + struct AlwaysFailObj { + x_param: FloatParam, + call_count: Arc, + } + + impl Objective for AlwaysFailObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + self.call_count.fetch_add(1, Ordering::Relaxed); + Err("always fails".to_string()) + } + fn max_retries(&self) -> usize { + 3 + } + } + + let study: Study = Study::new(Direction::Minimize); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = AlwaysFailObj { + x_param: FloatParam::new(0.0, 10.0), + call_count: Arc::clone(&call_count), + }; + + let result = study.optimize_with(1, obj); // 1 initial attempt + 3 retries = 4 total calls - assert_eq!(call_count.get(), 4); + assert_eq!(call_count.load(Ordering::Relaxed), 4); // No trials completed assert!(matches!(result, Err(Error::NoCompletedTrials))); } #[test] fn test_retries_permanently_failed_after_exhaustion() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); + use optimizer::Objective; - let result = study.optimize_with_retries(3, 2, |trial| { - let _ = x_param.suggest(trial).unwrap(); - Err::("transient error") - }); + struct AlwaysFailObj { + x_param: FloatParam, + } + + impl Objective for AlwaysFailObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + Err("transient error".to_string()) + } + fn max_retries(&self) -> usize { + 2 + } + } + + let study: Study = Study::new(Direction::Minimize); + let obj = AlwaysFailObj { + x_param: FloatParam::new(0.0, 10.0), + }; + + let result = study.optimize_with(3, obj); assert!( matches!(result, Err(Error::NoCompletedTrials)), @@ -2029,26 +1902,47 @@ fn test_retries_permanently_failed_after_exhaustion() { #[test] fn test_retries_uses_same_parameters() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - let seen_values = std::cell::RefCell::new(Vec::new()); - let call_count = std::cell::Cell::new(0u32); + use std::sync::atomic::{AtomicU32, Ordering}; + use std::sync::{Arc, Mutex}; - study - .optimize_with_retries(1, 2, |trial| { - let x = x_param.suggest(trial).map_err(|e| e.to_string())?; - seen_values.borrow_mut().push(x); - call_count.set(call_count.get() + 1); + use optimizer::Objective; + + struct RetryObj { + x_param: FloatParam, + seen_values: Arc>>, + call_count: Arc, + } + + impl Objective for RetryObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + self.seen_values.lock().unwrap().push(x); + let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1; // Fail first two attempts, succeed on third - if call_count.get() < 3 { - Err::("transient".to_string()) + if count < 3 { + Err("transient".to_string()) } else { Ok(x * x) } - }) - .unwrap(); + } + fn max_retries(&self) -> usize { + 2 + } + } - let values = seen_values.borrow(); + let study: Study = Study::new(Direction::Minimize); + let seen_values = Arc::new(Mutex::new(Vec::new())); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = RetryObj { + x_param: FloatParam::new(0.0, 10.0), + seen_values: Arc::clone(&seen_values), + call_count: Arc::clone(&call_count), + }; + + study.optimize_with(1, obj).unwrap(); + + let values = seen_values.lock().unwrap(); assert_eq!(values.len(), 3, "should be called 3 times (1 + 2 retries)"); // All three calls should have gotten the same parameter value assert_eq!(values[0], values[1]); @@ -2057,25 +1951,44 @@ fn test_retries_uses_same_parameters() { #[test] fn test_retries_n_trials_counts_unique_configs() { - let study: Study = Study::new(Direction::Minimize); - let x_param = FloatParam::new(0.0, 10.0); - let call_count = std::cell::Cell::new(0u32); + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; - study - .optimize_with_retries(3, 2, |trial| { - let x = x_param.suggest(trial).map_err(|e| e.to_string())?; - call_count.set(call_count.get() + 1); + use optimizer::Objective; + + struct FailFirstObj { + x_param: FloatParam, + call_count: Arc, + } + + impl Objective for FailFirstObj { + type Error = String; + fn evaluate(&self, trial: &mut Trial) -> Result { + let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?; + let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1; // Fail first attempt of each config, succeed on retry - if call_count.get() % 2 == 1 { - Err::("transient".to_string()) + if count % 2 == 1 { + Err("transient".to_string()) } else { Ok(x * x) } - }) - .unwrap(); + } + fn max_retries(&self) -> usize { + 2 + } + } + + let study: Study = Study::new(Direction::Minimize); + let call_count = Arc::new(AtomicU32::new(0)); + let obj = FailFirstObj { + x_param: FloatParam::new(0.0, 10.0), + call_count: Arc::clone(&call_count), + }; + + study.optimize_with(3, obj).unwrap(); // 3 unique configs, each needing 2 calls = 6 total calls - assert_eq!(call_count.get(), 6); + assert_eq!(call_count.load(Ordering::Relaxed), 6); // But only 3 completed trials assert_eq!(study.n_trials(), 3); } @@ -2087,7 +2000,7 @@ fn test_retries_with_zero_max_retries_same_as_optimize() { let call_count = std::cell::Cell::new(0u32); study - .optimize_with_retries(5, 0, |trial| { + .optimize(5, |trial| { let x = x_param.suggest(trial)?; call_count.set(call_count.get() + 1); Ok::<_, Error>(x * x) diff --git a/tests/serde_tests.rs b/tests/serde_tests.rs index 202bc19..b734152 100644 --- a/tests/serde_tests.rs +++ b/tests/serde_tests.rs @@ -173,73 +173,27 @@ fn round_trip_preserves_trial_id_counter() { } #[test] -fn checkpoint_file_created_at_interval() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(-10.0, 10.0).name("x"); - - let dir = tempdir(); - let checkpoint = dir.join("checkpoint.json"); - - study - .optimize_with_checkpoint(10, 3, &checkpoint, |trial| { - let v = x.suggest(trial)?; - Ok::<_, optimizer::Error>(v * v) - }) - .unwrap(); - - // Checkpoint should exist (written at trials 3, 6, 9) - assert!(checkpoint.exists(), "checkpoint file was not created"); - - // Load it and verify it's valid - let loaded: Study = Study::load(&checkpoint).unwrap(); - // Last checkpoint was at trial 9, so it should have 9 trials - assert_eq!(loaded.n_trials(), 9); - - std::fs::remove_dir_all(&dir).ok(); -} - -#[test] -fn checkpoint_overwrites_previous() { - let study: Study = Study::new(Direction::Minimize); - let x = FloatParam::new(0.0, 1.0); - - let dir = tempdir(); - let checkpoint = dir.join("checkpoint.json"); - - study - .optimize_with_checkpoint(6, 3, &checkpoint, |trial| { - let v = x.suggest(trial)?; - Ok::<_, optimizer::Error>(v) - }) - .unwrap(); - - // The checkpoint at trial 6 should overwrite the one from trial 3 - let loaded: Study = Study::load(&checkpoint).unwrap(); - assert_eq!(loaded.n_trials(), 6); - - std::fs::remove_dir_all(&dir).ok(); -} - -#[test] -fn resume_from_checkpoint_continues_trial_ids() { +fn save_and_resume_continues_trial_ids() { let study: Study = Study::new(Direction::Minimize); let x = FloatParam::new(-5.0, 5.0).name("x"); let dir = tempdir(); - let checkpoint = dir.join("resume.json"); + let save_path = dir.join("resume.json"); - // Run 10 trials with checkpointing + // Run 10 trials study - .optimize_with_checkpoint(10, 5, &checkpoint, |trial| { + .optimize(10, |trial| { let v = x.suggest(trial)?; Ok::<_, optimizer::Error>(v * v) }) .unwrap(); - // Load and continue - let loaded: Study = Study::load(&checkpoint).unwrap(); + // Save and reload + study.save(&save_path).unwrap(); + let loaded: Study = Study::load(&save_path).unwrap(); assert_eq!(loaded.n_trials(), 10); + // Continue with 5 more trials let remaining = 15 - loaded.n_trials(); loaded .optimize(remaining, |trial| { @@ -261,24 +215,26 @@ fn resume_from_checkpoint_continues_trial_ids() { } #[test] -fn atomic_write_no_temp_file_left_behind() { +fn save_uses_atomic_write() { let study: Study = Study::new(Direction::Minimize); let x = FloatParam::new(0.0, 1.0); let dir = tempdir(); - let checkpoint = dir.join("atomic.json"); + let save_path = dir.join("atomic.json"); study - .optimize_with_checkpoint(3, 3, &checkpoint, |trial| { + .optimize(3, |trial| { let v = x.suggest(trial)?; Ok::<_, optimizer::Error>(v) }) .unwrap(); + study.save(&save_path).unwrap(); + // The temp file should have been renamed, not left behind let tmp_path = dir.join(".atomic.json.tmp"); assert!(!tmp_path.exists(), "temp file was not cleaned up"); - assert!(checkpoint.exists(), "checkpoint file was not created"); + assert!(save_path.exists(), "save file was not created"); std::fs::remove_dir_all(&dir).ok(); }