From db0314a1d144586c5419f5ec76114da33d937e00 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 17:51:36 +0100 Subject: [PATCH] feat: add optimize_with_retries() for automatic retry of failed trials --- src/study.rs | 124 ++++++++++++++++++++++++++++++++++++++++ tests/integration.rs | 133 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 257 insertions(+) diff --git a/src/study.rs b/src/study.rs index 97e2d87..9dc32f5 100644 --- a/src/study.rs +++ b/src/study.rs @@ -358,6 +358,22 @@ where .map(|t| t.id) } + /// Creates a new trial with pre-set parameter values. + /// + /// The trial gets a new unique ID but reuses the given parameters. When + /// `suggest_param` is called on the resulting trial, fixed values are + /// returned instead of sampling. + fn create_trial_with_params(&self, params: HashMap) -> Trial { + let id = self.next_trial_id(); + let mut trial = if let Some(factory) = &self.trial_factory { + factory(id) + } else { + Trial::new(id) + }; + trial.set_fixed_params(params); + trial + } + /// Returns the number of enqueued parameter configurations. #[must_use] pub fn n_enqueued(&self) -> usize { @@ -1593,6 +1609,114 @@ where Ok(()) } + + /// Runs 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. + /// + /// # Errors + /// + /// Returns `Error::NoCompletedTrials` if no trials completed successfully. + /// + /// # Examples + /// + /// ``` + /// 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); + /// + /// 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")) + /// } else { + /// Ok(x * x) + /// } + /// }) + /// .unwrap(); + /// + /// assert_eq!(study.n_trials(), 5); + /// ``` + pub fn optimize_with_retries( + &self, + n_trials: usize, + max_retries: usize, + 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_with_retries", n_trials, max_retries, direction = ?self.direction).entered(); + + for _ in 0..n_trials { + let mut trial = self.create_trial(); + let mut retries = 0; + loop { + match objective(&mut trial) { + Ok(value) => { + #[cfg(feature = "tracing")] + let trial_id = trial.id(); + self.complete_trial(trial, value); + trace_info!(trial_id, "trial completed"); + break; + } + Err(_) if retries < max_retries => { + retries += 1; + // Create a new trial with the same parameters + trial = self.create_trial_with_params(trial.params().clone()); + } + 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 permanently failed"); + } + break; + } + } + } + } + + // Return error if no trials completed successfully + let has_complete = self + .completed_trials + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } } impl Study diff --git a/tests/integration.rs b/tests/integration.rs index f10a999..e029531 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -1964,3 +1964,136 @@ fn test_display_matches_summary() { assert_eq!(format!("{study}"), study.summary()); } + +// ============================================================================= +// Tests: optimize_with_retries +// ============================================================================= + +#[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); + + 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(); + + // All trials succeed on first try — exactly 5 calls + assert_eq!(call_count.get(), 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); + + 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") + }); + + // 1 initial attempt + 3 retries = 4 total calls + assert_eq!(call_count.get(), 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); + + let result = study.optimize_with_retries(3, 2, |trial| { + let _ = x_param.suggest(trial).unwrap(); + Err::("transient error") + }); + + assert!( + matches!(result, Err(Error::NoCompletedTrials)), + "all trials should permanently fail" + ); + assert_eq!( + study.n_trials(), + 0, + "no completed trials should be recorded" + ); +} + +#[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); + + 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); + // Fail first two attempts, succeed on third + if call_count.get() < 3 { + Err::("transient".to_string()) + } else { + Ok(x * x) + } + }) + .unwrap(); + + let values = seen_values.borrow(); + 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]); + assert_eq!(values[1], values[2]); +} + +#[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); + + 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); + // Fail first attempt of each config, succeed on retry + if call_count.get() % 2 == 1 { + Err::("transient".to_string()) + } else { + Ok(x * x) + } + }) + .unwrap(); + + // 3 unique configs, each needing 2 calls = 6 total calls + assert_eq!(call_count.get(), 6); + // But only 3 completed trials + assert_eq!(study.n_trials(), 3); +} + +#[test] +fn test_retries_with_zero_max_retries_same_as_optimize() { + let study: Study = Study::new(Direction::Minimize); + let x_param = FloatParam::new(0.0, 10.0); + let call_count = std::cell::Cell::new(0u32); + + study + .optimize_with_retries(5, 0, |trial| { + let x = x_param.suggest(trial)?; + call_count.set(call_count.get() + 1); + Ok::<_, Error>(x * x) + }) + .unwrap(); + + assert_eq!(call_count.get(), 5); + assert_eq!(study.n_trials(), 5); +}