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

feat: add timeout-based optimization methods (optimize_until)

Browse Source 
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)
This commit is contained in:
Manuel Raimann
2026-02-11 17:06:08 +01:00
parent c586650df4
commit bc278d83a3
2 changed files with 479 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/study.rs
+320
View File
@@ -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<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::<_, optimizer::Error>(x * x)
/// })
/// .unwrap();
///
/// assert!(study.n_trials() > 0);
/// ```
pub fn optimize_until<F, E>(&self, duration: Duration, mut objective: F) -> crate::Result<()>
where
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
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<f64> = 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<F, C, E>(
&self,
duration: Duration,
mut objective: F,
mut callback: C,
) -> crate::Result<()>
where
V: Clone + Default,
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
C: FnMut(&Study<V>, &CompletedTrial<V>) -> 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<V, E>)`.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
#[cfg(feature = "async")]
pub async fn optimize_until_async<F, Fut, E>(
&self,
duration: Duration,
objective: F,
) -> crate::Result<()>
where
F: Fn(Trial) -> Fut,
Fut: Future<Output = core::result::Result<(Trial, V), E>>,
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<F, Fut, E>(
&self,
duration: Duration,
concurrency: usize,
objective: F,
) -> crate::Result<()>
where
F: Fn(Trial) -> Fut + Send + Sync + 'static,
Fut: Future<Output = core::result::Result<(Trial, V), E>> + 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(())
}
}
tests/integration.rs
+159
View File
@@ -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<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);
}