From bc278d83a3dbe1aca6120490cd9c6797ab1f4181 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 17:06:08 +0100 Subject: [PATCH] feat: add timeout-based optimization methods (optimize_until) Add duration-based variants of all optimization methods that run trials until a wall-clock deadline rather than for a fixed trial count: - optimize_until(duration, objective) - optimize_until_with_callback(duration, objective, callback) - optimize_until_async(duration, objective) - optimize_until_parallel(duration, concurrency, objective) --- src/study.rs | 320 +++++++++++++++++++++++++++++++++++++++++++ tests/integration.rs | 159 +++++++++++++++++++++ 2 files changed, 479 insertions(+) diff --git a/src/study.rs b/src/study.rs index 686253e..4ab33d7 100644 --- a/src/study.rs +++ b/src/study.rs @@ -5,7 +5,9 @@ use core::any::Any; use core::future::Future; use core::ops::ControlFlow; use core::sync::atomic::{AtomicU64, Ordering}; +use core::time::Duration; use std::sync::Arc; +use std::time::Instant; use parking_lot::RwLock; @@ -996,6 +998,324 @@ where return Err(crate::Error::NoCompletedTrials); } + Ok(()) + } + /// Runs 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, + { + let deadline = Instant::now() + duration; + while Instant::now() < deadline { + let mut trial = self.create_trial(); + + match objective(&mut trial) { + Ok(value) => { + self.complete_trial(trial, value); + } + Err(e) => { + if is_trial_pruned(&e) { + self.prune_trial(trial); + } else { + self.fail_trial(trial, e.to_string()); + } + } + } + } + + let has_complete = self + .completed_trials + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } + + /// Runs 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, + { + let deadline = Instant::now() + duration; + while Instant::now() < deadline { + let mut trial = self.create_trial(); + + match objective(&mut trial) { + Ok(value) => { + self.complete_trial(trial, value); + + let trials = self.completed_trials.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; + } + } + Err(e) => { + if is_trial_pruned(&e) { + self.prune_trial(trial); + } else { + self.fail_trial(trial, e.to_string()); + } + } + } + } + + let has_complete = self + .completed_trials + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } + + /// Runs 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. + /// + /// # 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, Result)`. + /// + /// # 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, + { + let deadline = Instant::now() + duration; + while Instant::now() < deadline { + let trial = self.create_trial(); + + match objective(trial).await { + Ok((trial, value)) => { + self.complete_trial(trial, value); + } + Err(e) => { + let _ = e.to_string(); + } + } + } + + let has_complete = self + .completed_trials + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } + + /// Runs optimization with bounded parallelism until the given duration has elapsed. + /// + /// 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; + + 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)) => { + self.complete_trial(trial, value); + } + Err(e) => { + let _ = e.to_string(); + } + } + } + + let has_complete = self + .completed_trials + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + Ok(()) } } diff --git a/tests/integration.rs b/tests/integration.rs index deb0c90..015d1ef 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -1363,3 +1363,162 @@ fn test_completed_trial_get() { assert!((-10.0..=10.0).contains(&x_val)); 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); +}