From 5da63d89763517cce56ea83ee0393d85ade4c852 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Thu, 12 Feb 2026 14:39:28 +0100 Subject: [PATCH] perf: reduce CompletedTrial cloning overhead - Add Trial::into_completed() and into_multi_objective_trial() to move fields instead of cloning 5 HashMaps/Vecs per trial completion - Fire after_trial callback before pushing to storage, eliminating the clone-from-storage pattern at all 4 call sites - Optimize top_trials(n) to sort indices and clone only N trials instead of cloning all completed trials - Remove unused set_complete/set_pruned methods --- src/multi_objective.rs | 14 +---- src/objective.rs | 5 ++ src/study.rs | 112 +++++++++++++-------------------------- src/trial.rs | 47 ++++++++++++---- tests/study/objective.rs | 8 ++- 5 files changed, 87 insertions(+), 99 deletions(-) diff --git a/src/multi_objective.rs b/src/multi_objective.rs index 9b924f3..d2e702a 100644 --- a/src/multi_objective.rs +++ b/src/multi_objective.rs @@ -321,18 +321,8 @@ impl MultiObjectiveStudy { } /// Records a completed trial. - fn complete_trial(&self, mut trial: Trial, values: Vec) { - trial.set_complete(); - let mo_trial = MultiObjectiveTrial { - id: trial.id(), - params: trial.params().clone(), - distributions: trial.distributions().clone(), - param_labels: trial.param_labels().clone(), - values, - state: TrialState::Complete, - user_attrs: trial.user_attrs().clone(), - constraints: trial.constraint_values().to_vec(), - }; + fn complete_trial(&self, trial: Trial, values: Vec) { + let mo_trial = trial.into_multi_objective_trial(values, TrialState::Complete); self.completed_trials.write().push(mo_trial); } diff --git a/src/objective.rs b/src/objective.rs index c22babf..4cad201 100644 --- a/src/objective.rs +++ b/src/objective.rs @@ -117,6 +117,11 @@ pub trait Objective { /// Called after each **completed** trial (not failed or pruned). /// + /// The trial is passed directly as the argument *before* it is pushed + /// to storage, so `study.n_trials()` and `study.trials()` do not yet + /// include this trial. The trial is always pushed to storage after this + /// callback returns, regardless of the return value. + /// /// Return `ControlFlow::Break(())` to stop the optimization loop. /// /// Default: always continues. diff --git a/src/study.rs b/src/study.rs index dc4c648..2fcba0e 100644 --- a/src/study.rs +++ b/src/study.rs @@ -491,20 +491,8 @@ where /// /// assert_eq!(study.n_trials(), 1); /// ``` - pub fn complete_trial(&self, mut trial: Trial, value: V) { - trial.set_complete(); - let mut completed = CompletedTrial::with_intermediate_values( - trial.id(), - trial.params().clone(), - trial.distributions().clone(), - trial.param_labels().clone(), - value, - trial.intermediate_values().to_vec(), - trial.user_attrs().clone(), - ); - completed.state = TrialState::Complete; - completed.constraints = trial.constraint_values().to_vec(); - + pub fn complete_trial(&self, trial: Trial, value: V) { + let completed = trial.into_completed(value, TrialState::Complete); self.storage.push(completed); } @@ -604,23 +592,11 @@ where /// # Arguments /// /// * `trial` - The trial that was pruned. - pub fn prune_trial(&self, mut trial: Trial) + pub fn prune_trial(&self, trial: Trial) where V: Default, { - trial.set_pruned(); - let mut completed = CompletedTrial::with_intermediate_values( - trial.id(), - trial.params().clone(), - trial.distributions().clone(), - trial.param_labels().clone(), - V::default(), - trial.intermediate_values().to_vec(), - trial.user_attrs().clone(), - ); - completed.state = TrialState::Pruned; - completed.constraints = trial.constraint_values().to_vec(); - + let completed = trial.into_completed(V::default(), TrialState::Pruned); self.storage.push(completed); } @@ -852,15 +828,17 @@ where { let trials = self.storage.trials_arc().read(); let direction = self.direction; - let mut completed: Vec<_> = trials + // Sort indices instead of cloning all trials, then clone only the top N. + let mut indices: Vec = trials .iter() - .filter(|t| t.state == TrialState::Complete) - .cloned() + .enumerate() + .filter(|(_, t)| t.state == TrialState::Complete) + .map(|(i, _)| i) .collect(); // Sort best-first: reverse the compare_trials ordering (which is designed for max_by) - completed.sort_by(|a, b| Self::compare_trials(b, a, direction)); - completed.truncate(n); - completed + indices.sort_by(|&a, &b| Self::compare_trials(&trials[b], &trials[a], direction)); + indices.truncate(n); + indices.iter().map(|&i| trials[i].clone()).collect() } /// Run optimization with an objective. @@ -921,7 +899,12 @@ where Ok(value) => { #[cfg(feature = "tracing")] let trial_id = trial.id(); - self.complete_trial(trial, value); + + let completed = trial.into_completed(value, TrialState::Complete); + + // Fire after_trial hook before pushing to storage + let flow = objective.after_trial(self, &completed); + self.storage.push(completed); #[cfg(feature = "tracing")] { @@ -938,17 +921,8 @@ where } } - // Fire after_trial hook - let trials = self.storage.trials_arc().read(); - if let Some(completed) = trials.last() { - let completed_clone = completed.clone(); - drop(trials); - if let ControlFlow::Break(()) = - objective.after_trial(self, &completed_clone) - { - // Return early — at least one trial completed. - return Ok(()); - } + if let ControlFlow::Break(()) = flow { + return Ok(()); } } Err(e) if is_trial_pruned(&e) => { @@ -1049,19 +1023,14 @@ where (t, Ok(value)) => { #[cfg(feature = "tracing")] let trial_id = t.id(); - self.complete_trial(t, value); + + let completed = t.into_completed(value, TrialState::Complete); + let flow = objective.after_trial(self, &completed); + self.storage.push(completed); trace_info!(trial_id, "trial completed"); - // Fire after_trial hook - let trials = self.storage.trials_arc().read(); - if let Some(completed) = trials.last() { - let completed_clone = completed.clone(); - drop(trials); - if let ControlFlow::Break(()) = - objective.after_trial(self, &completed_clone) - { - return Ok(()); - } + if let ControlFlow::Break(()) = flow { + return Ok(()); } } (t, Err(e)) if is_trial_pruned(&e) => { @@ -1172,18 +1141,14 @@ where (t, Ok(value)) => { #[cfg(feature = "tracing")] let trial_id = t.id(); - self.complete_trial(t, value); + + let completed = t.into_completed(value, TrialState::Complete); + let flow = objective.after_trial(self, &completed); + self.storage.push(completed); trace_info!(trial_id, "trial completed"); - let trials = self.storage.trials_arc().read(); - if let Some(completed) = trials.last() { - let completed_clone = completed.clone(); - drop(trials); - if let ControlFlow::Break(()) = - objective.after_trial(self, &completed_clone) - { - break 'spawn; - } + if let ControlFlow::Break(()) = flow { + break 'spawn; } } (t, Err(e)) => { @@ -1228,15 +1193,12 @@ where (t, Ok(value)) => { #[cfg(feature = "tracing")] let trial_id = t.id(); - self.complete_trial(t, value); - trace_info!(trial_id, "trial completed"); + + let completed = t.into_completed(value, TrialState::Complete); // Still fire after_trial for bookkeeping, but don't break — we're draining. - let trials = self.storage.trials_arc().read(); - if let Some(completed) = trials.last() { - let completed_clone = completed.clone(); - drop(trials); - let _ = objective.after_trial(self, &completed_clone); - } + let _ = objective.after_trial(self, &completed); + self.storage.push(completed); + trace_info!(trial_id, "trial completed"); } (t, Err(e)) => { #[cfg(feature = "tracing")] diff --git a/src/trial.rs b/src/trial.rs index 1ef939e..9d069e0 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -26,6 +26,7 @@ use parking_lot::RwLock; use crate::distribution::Distribution; use crate::error::{Error, Result}; +use crate::multi_objective::MultiObjectiveTrial; use crate::param::ParamValue; use crate::parameter::{ParamId, Parameter}; use crate::pruner::Pruner; @@ -389,21 +390,11 @@ impl Trial { &self.constraint_values } - /// Set the trial state to `Complete`. - pub(crate) fn set_complete(&mut self) { - self.state = TrialState::Complete; - } - /// Set the trial state to `Failed`. pub(crate) fn set_failed(&mut self) { self.state = TrialState::Failed; } - /// Set the trial state to `Pruned`. - pub(crate) fn set_pruned(&mut self) { - self.state = TrialState::Pruned; - } - /// Suggest a parameter value using a [`Parameter`] definition. /// /// This is the primary entry point for sampling parameters. It handles @@ -482,6 +473,42 @@ impl Trial { Ok(result) } + + /// Consume this trial and move its fields into a [`CompletedTrial`]. + /// + /// This avoids cloning the trial's `HashMap`s and `Vec`s by moving + /// ownership directly into the completed trial. + pub(crate) fn into_completed(self, value: V, state: TrialState) -> CompletedTrial { + CompletedTrial { + id: self.id, + params: self.params, + distributions: self.distributions, + param_labels: self.param_labels, + value, + intermediate_values: self.intermediate_values, + state, + user_attrs: self.user_attrs, + constraints: self.constraint_values, + } + } + + /// Consume this trial and move its fields into a [`MultiObjectiveTrial`]. + pub(crate) fn into_multi_objective_trial( + self, + values: Vec, + state: TrialState, + ) -> MultiObjectiveTrial { + MultiObjectiveTrial { + id: self.id, + params: self.params, + distributions: self.distributions, + param_labels: self.param_labels, + values, + state, + user_attrs: self.user_attrs, + constraints: self.constraint_values, + } + } } #[cfg(test)] diff --git a/tests/study/objective.rs b/tests/study/objective.rs index 8d28bf1..e379ec9 100644 --- a/tests/study/objective.rs +++ b/tests/study/objective.rs @@ -20,7 +20,9 @@ fn test_callback_early_stopping() { Ok(x) } fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { - if study.n_trials() >= 5 { + // The current trial has not yet been pushed to storage when + // after_trial fires, so n_trials() == 4 means this is the 5th. + if study.n_trials() >= 4 { ControlFlow::Break(()) } else { ControlFlow::Continue(()) @@ -98,7 +100,9 @@ fn test_callback_sampler_early_stopping() { Ok(x) } fn after_trial(&self, study: &Study, _trial: &CompletedTrial) -> ControlFlow<()> { - if study.n_trials() >= 3 { + // The current trial has not yet been pushed to storage when + // after_trial fires, so n_trials() == 2 means this is the 3rd. + if study.n_trials() >= 2 { ControlFlow::Break(()) } else { ControlFlow::Continue(())