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
This commit is contained in:
+2
-12
@@ -321,18 +321,8 @@ impl MultiObjectiveStudy {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Records a completed trial.
|
/// Records a completed trial.
|
||||||
fn complete_trial(&self, mut trial: Trial, values: Vec<f64>) {
|
fn complete_trial(&self, trial: Trial, values: Vec<f64>) {
|
||||||
trial.set_complete();
|
let mo_trial = trial.into_multi_objective_trial(values, TrialState::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(),
|
|
||||||
};
|
|
||||||
self.completed_trials.write().push(mo_trial);
|
self.completed_trials.write().push(mo_trial);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -117,6 +117,11 @@ pub trait Objective<V: PartialOrd = f64> {
|
|||||||
|
|
||||||
/// Called after each **completed** trial (not failed or pruned).
|
/// 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.
|
/// Return `ControlFlow::Break(())` to stop the optimization loop.
|
||||||
///
|
///
|
||||||
/// Default: always continues.
|
/// Default: always continues.
|
||||||
|
|||||||
+34
-72
@@ -491,20 +491,8 @@ where
|
|||||||
///
|
///
|
||||||
/// assert_eq!(study.n_trials(), 1);
|
/// assert_eq!(study.n_trials(), 1);
|
||||||
/// ```
|
/// ```
|
||||||
pub fn complete_trial(&self, mut trial: Trial, value: V) {
|
pub fn complete_trial(&self, trial: Trial, value: V) {
|
||||||
trial.set_complete();
|
let completed = trial.into_completed(value, TrialState::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();
|
|
||||||
|
|
||||||
self.storage.push(completed);
|
self.storage.push(completed);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -604,23 +592,11 @@ where
|
|||||||
/// # Arguments
|
/// # Arguments
|
||||||
///
|
///
|
||||||
/// * `trial` - The trial that was pruned.
|
/// * `trial` - The trial that was pruned.
|
||||||
pub fn prune_trial(&self, mut trial: Trial)
|
pub fn prune_trial(&self, trial: Trial)
|
||||||
where
|
where
|
||||||
V: Default,
|
V: Default,
|
||||||
{
|
{
|
||||||
trial.set_pruned();
|
let completed = trial.into_completed(V::default(), TrialState::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();
|
|
||||||
|
|
||||||
self.storage.push(completed);
|
self.storage.push(completed);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -852,15 +828,17 @@ where
|
|||||||
{
|
{
|
||||||
let trials = self.storage.trials_arc().read();
|
let trials = self.storage.trials_arc().read();
|
||||||
let direction = self.direction;
|
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()
|
.iter()
|
||||||
.filter(|t| t.state == TrialState::Complete)
|
.enumerate()
|
||||||
.cloned()
|
.filter(|(_, t)| t.state == TrialState::Complete)
|
||||||
|
.map(|(i, _)| i)
|
||||||
.collect();
|
.collect();
|
||||||
// Sort best-first: reverse the compare_trials ordering (which is designed for max_by)
|
// 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));
|
indices.sort_by(|&a, &b| Self::compare_trials(&trials[b], &trials[a], direction));
|
||||||
completed.truncate(n);
|
indices.truncate(n);
|
||||||
completed
|
indices.iter().map(|&i| trials[i].clone()).collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Run optimization with an objective.
|
/// Run optimization with an objective.
|
||||||
@@ -921,7 +899,12 @@ where
|
|||||||
Ok(value) => {
|
Ok(value) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
let trial_id = trial.id();
|
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")]
|
#[cfg(feature = "tracing")]
|
||||||
{
|
{
|
||||||
@@ -938,19 +921,10 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fire after_trial hook
|
if let ControlFlow::Break(()) = flow {
|
||||||
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
Err(e) if is_trial_pruned(&e) => {
|
Err(e) if is_trial_pruned(&e) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
let trial_id = trial.id();
|
let trial_id = trial.id();
|
||||||
@@ -1049,21 +1023,16 @@ where
|
|||||||
(t, Ok(value)) => {
|
(t, Ok(value)) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
let trial_id = t.id();
|
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");
|
trace_info!(trial_id, "trial completed");
|
||||||
|
|
||||||
// Fire after_trial hook
|
if let ControlFlow::Break(()) = flow {
|
||||||
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
(t, Err(e)) if is_trial_pruned(&e) => {
|
(t, Err(e)) if is_trial_pruned(&e) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
let trial_id = t.id();
|
let trial_id = t.id();
|
||||||
@@ -1172,20 +1141,16 @@ where
|
|||||||
(t, Ok(value)) => {
|
(t, Ok(value)) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
let trial_id = t.id();
|
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");
|
trace_info!(trial_id, "trial completed");
|
||||||
|
|
||||||
let trials = self.storage.trials_arc().read();
|
if let ControlFlow::Break(()) = flow {
|
||||||
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;
|
break 'spawn;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
(t, Err(e)) => {
|
(t, Err(e)) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
let trial_id = t.id();
|
let trial_id = t.id();
|
||||||
@@ -1228,15 +1193,12 @@ where
|
|||||||
(t, Ok(value)) => {
|
(t, Ok(value)) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
let trial_id = t.id();
|
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.
|
// Still fire after_trial for bookkeeping, but don't break — we're draining.
|
||||||
let trials = self.storage.trials_arc().read();
|
let _ = objective.after_trial(self, &completed);
|
||||||
if let Some(completed) = trials.last() {
|
self.storage.push(completed);
|
||||||
let completed_clone = completed.clone();
|
trace_info!(trial_id, "trial completed");
|
||||||
drop(trials);
|
|
||||||
let _ = objective.after_trial(self, &completed_clone);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
(t, Err(e)) => {
|
(t, Err(e)) => {
|
||||||
#[cfg(feature = "tracing")]
|
#[cfg(feature = "tracing")]
|
||||||
|
|||||||
+37
-10
@@ -26,6 +26,7 @@ use parking_lot::RwLock;
|
|||||||
|
|
||||||
use crate::distribution::Distribution;
|
use crate::distribution::Distribution;
|
||||||
use crate::error::{Error, Result};
|
use crate::error::{Error, Result};
|
||||||
|
use crate::multi_objective::MultiObjectiveTrial;
|
||||||
use crate::param::ParamValue;
|
use crate::param::ParamValue;
|
||||||
use crate::parameter::{ParamId, Parameter};
|
use crate::parameter::{ParamId, Parameter};
|
||||||
use crate::pruner::Pruner;
|
use crate::pruner::Pruner;
|
||||||
@@ -389,21 +390,11 @@ impl Trial {
|
|||||||
&self.constraint_values
|
&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`.
|
/// Set the trial state to `Failed`.
|
||||||
pub(crate) fn set_failed(&mut self) {
|
pub(crate) fn set_failed(&mut self) {
|
||||||
self.state = TrialState::Failed;
|
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.
|
/// Suggest a parameter value using a [`Parameter`] definition.
|
||||||
///
|
///
|
||||||
/// This is the primary entry point for sampling parameters. It handles
|
/// This is the primary entry point for sampling parameters. It handles
|
||||||
@@ -482,6 +473,42 @@ impl Trial {
|
|||||||
|
|
||||||
Ok(result)
|
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)]
|
#[cfg(test)]
|
||||||
|
|||||||
@@ -20,7 +20,9 @@ fn test_callback_early_stopping() {
|
|||||||
Ok(x)
|
Ok(x)
|
||||||
}
|
}
|
||||||
fn after_trial(&self, study: &Study<f64>, _trial: &CompletedTrial<f64>) -> ControlFlow<()> {
|
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(())
|
ControlFlow::Break(())
|
||||||
} else {
|
} else {
|
||||||
ControlFlow::Continue(())
|
ControlFlow::Continue(())
|
||||||
@@ -98,7 +100,9 @@ fn test_callback_sampler_early_stopping() {
|
|||||||
Ok(x)
|
Ok(x)
|
||||||
}
|
}
|
||||||
fn after_trial(&self, study: &Study<f64>, _trial: &CompletedTrial<f64>) -> ControlFlow<()> {
|
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(())
|
ControlFlow::Break(())
|
||||||
} else {
|
} else {
|
||||||
ControlFlow::Continue(())
|
ControlFlow::Continue(())
|
||||||
|
|||||||
Reference in New Issue
Block a user