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

feat: add Objective trait and unify optimize API

Browse Source 
- Add `Objective<V>` trait with lifecycle hooks (`before_trial`,
  `after_trial`, `max_retries`) in new `src/objective.rs`
- Replace 14+ optimize variants with 6 methods: `optimize`,
  `optimize_with`, and async/parallel counterparts
- `optimize*` methods accept closures directly (FnMut for sync, Fn for
  async); `optimize_with*` methods accept `impl Objective<V>` for
  struct-based objectives with hooks and retries
- Remove `optimize_until`, `optimize_with_callback`,
  `optimize_with_retries`, `optimize_with_checkpoint`, and all
  deprecated `_with_sampler` methods
This commit is contained in:
Manuel Raimann
2026-02-12 12:25:20 +01:00
parent ee59c9cdd0
commit 964b4d5749
8 changed files with 1023 additions and 1469 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/async_parallel.rs
+13 -19
View File
@@ -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()?;
examples/early_stopping.rs
+33 -18
View File
@@ -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<f64> for EarlyStopObjective {
type Error = Error;
fn evaluate(&self, trial: &mut Trial) -> Result<f64> {
let v = self.x.suggest(trial)?;
Ok((v - 3.0).powi(2))
}
fn after_trial(&self, _study: &Study<f64>, trial: &CompletedTrial<f64>) -> 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<f64> = 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!(
src/lib.rs
+3
View File
@@ -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,
src/objective.rs
+132
View File
@@ -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<f64> = 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<f64> for QuadraticWithEarlyStopping {
//! type Error = Error;
//!
//! fn evaluate(&self, trial: &mut Trial) -> Result<f64> {
//! let v = self.x.suggest(trial)?;
//! Ok((v - 3.0).powi(2))
//! }
//!
//! fn after_trial(&self, _study: &Study<f64>, trial: &CompletedTrial<f64>) -> ControlFlow<()> {
//! if trial.value < self.target {
//! ControlFlow::Break(())
//! } else {
//! ControlFlow::Continue(())
//! }
//! }
//! }
//!
//! let study: Study<f64> = 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<V: PartialOrd = f64> {
/// 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<V, Self::Error>;
/// 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<V>) -> 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<V>, _trial: &CompletedTrial<V>) -> 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
}
}
src/study.rs
+572 -971
View File
File diff suppressed because it is too large Load Diff
tests/async_tests.rs
+29 -89
View File
@@ -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<f64> = 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<f64> = 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::<f64, &str>("always fails")
})
.await;
@@ -150,27 +118,8 @@ async fn test_optimize_parallel_all_failures() {
let study: Study<f64> = 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<f64> = 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::<f64, &str>("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");
tests/integration.rs
+227 -314
View File
@@ -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<f64> = 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<f64> for EarlyStopAfter5 {
type Error = Error;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, Error> {
let x = self.x_param.suggest(trial)?;
Ok(x)
}
fn after_trial(&self, study: &Study<f64>, _trial: &CompletedTrial<f64>) -> ControlFlow<()> {
if study.n_trials() >= 5 {
ControlFlow::Break(())
} else {
ControlFlow::Continue(())
}
}
}
let study: Study<f64> = 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<f64> = Study::new(Direction::Minimize);
let result = study.optimize_with_callback(
5,
|_trial| Err::<f64, &str>("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<f64> = Study::new(Direction::Minimize);
let result = study.optimize_with_sampler(5, |_trial| Err::<f64, &str>("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<f64> = Study::new(Direction::Minimize);
let result = study.optimize_with_callback_sampler(
5,
|_trial| Err::<f64, &str>("always fails"),
|_study, _trial| ControlFlow::Continue(()),
);
let result = study.optimize(5, |_trial| Err::<f64, &str>("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<f64> = Study::with_sampler(Direction::Minimize, sampler);
// create_trial() creates trial with sampler integration for Study<f64>
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<f64> = 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<f64> = 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<f64> for StopImmediately {
type Error = Error;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, Error> {
let x = self.x_param.suggest(trial)?;
Ok(x)
}
fn after_trial(
&self,
_study: &Study<f64>,
_trial: &CompletedTrial<f64>,
) -> ControlFlow<()> {
ControlFlow::Break(())
}
}
let study: Study<f64> = 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<f64> for StopAfter3 {
type Error = Error;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, Error> {
let x = self.x_param.suggest(trial)?;
Ok(x)
}
fn after_trial(&self, study: &Study<f64>, _trial: &CompletedTrial<f64>) -> ControlFlow<()> {
if study.n_trials() >= 3 {
ControlFlow::Break(())
} else {
ControlFlow::Continue(())
}
}
}
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = 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<f64> = 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<f64> = 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<f64> = 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<f64> = 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<f64> = 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<f64> = Study::new(Direction::Minimize);
let result = study.optimize_until(Duration::from_millis(50), |_trial| {
Err::<f64, &str>("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<i32> = 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<f64> = 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<AtomicU32>,
}
impl Objective<f64> for SuccessObj {
type Error = Error;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, Error> {
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<f64> = 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<f64> = 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::<f64, _>("always fails")
});
use optimizer::Objective;
struct AlwaysFailObj {
x_param: FloatParam,
call_count: Arc<AtomicU32>,
}
impl Objective<f64> for AlwaysFailObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
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<f64> = 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<f64> = 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::<f64, _>("transient error")
});
struct AlwaysFailObj {
x_param: FloatParam,
}
impl Objective<f64> for AlwaysFailObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
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<f64> = 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<f64> = 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<Mutex<Vec<f64>>>,
call_count: Arc<AtomicU32>,
}
impl Objective<f64> for RetryObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
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::<f64, _>("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<f64> = 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<f64> = 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<AtomicU32>,
}
impl Objective<f64> for FailFirstObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
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::<f64, _>("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<f64> = 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)
tests/serde_tests.rs
+14 -58
View File
@@ -173,73 +173,27 @@ fn round_trip_preserves_trial_id_counter() {
}
#[test]
fn checkpoint_file_created_at_interval() {
let study: Study<f64> = 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<f64> = 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<f64> = 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<f64> = 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<f64> = 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<f64> = Study::load(&checkpoint).unwrap();
// Save and reload
study.save(&save_path).unwrap();
let loaded: Study<f64> = 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<f64> = 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();
}