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

feat: add optimize_with_retries() for automatic retry of failed trials

Browse Source
This commit is contained in:
Manuel Raimann
2026-02-11 17:51:36 +01:00
parent df5fcedecf
commit db0314a1d1
2 changed files with 257 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/study.rs
+124
View File
@@ -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<ParamId, ParamValue>) -> 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<f64> = 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::<f64, _>(optimizer::Error::Internal("transient"))
/// } else {
/// Ok(x * x)
/// }
/// })
/// .unwrap();
///
/// assert_eq!(study.n_trials(), 5);
/// ```
pub fn optimize_with_retries<F, E>(
&self,
n_trials: usize,
max_retries: usize,
mut objective: F,
) -> crate::Result<()>
where
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
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<V> Study<V>
tests/integration.rs
+133
View File
@@ -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<f64> = 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<f64> = 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::<f64, _>("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<f64> = 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::<f64, _>("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<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);
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::<f64, _>("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<f64> = 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::<f64, _>("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<f64> = 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);
}