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

perf: reduce CompletedTrial cloning overhead

Browse Source 
- 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
This commit is contained in:
Manuel Raimann
2026-02-12 14:39:28 +01:00
parent 74cf0643eb
commit 5da63d8976
5 changed files with 87 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/multi_objective.rs
+2 -12
View File
@@ -321,18 +321,8 @@ impl MultiObjectiveStudy {
}
/// Records a completed trial.
fn complete_trial(&self, mut trial: Trial, values: Vec<f64>) {
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<f64>) {
let mo_trial = trial.into_multi_objective_trial(values, TrialState::Complete);
self.completed_trials.write().push(mo_trial);
}
src/objective.rs
+5
View File
@@ -117,6 +117,11 @@ pub trait Objective<V: PartialOrd = f64> {
/// 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.
src/study.rs
+37 -75
View File
@@ -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<usize> = 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")]
src/trial.rs
+37 -10
View File
@@ -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<V>(self, value: V, state: TrialState) -> CompletedTrial<V> {
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<f64>,
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)]
tests/study/objective.rs
+6 -2
View File
@@ -20,7 +20,9 @@ fn test_callback_early_stopping() {
Ok(x)
}
fn after_trial(&self, study: &Study<f64>, _trial: &CompletedTrial<f64>) -> 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<f64>, _trial: &CompletedTrial<f64>) -> 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(())