Skip to content

Repository navigation

Add error_on_overwrite option in attr-related storage apis #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
c-bata merged 1 commit into from
Dec 5, 2025
Merged

Add error_on_overwrite option in attr-related storage apis #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions rustuna_core/src/error.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ pub enum ErrorKind {
DuplicatedStudy,
StudyNotFound,
TrialNotFound,
AttrOverwriteNotAllowed,
InvalidObjectiveValues,
TrialAlreadyFinished,
StorageInternalError,
Expand Down
46 changes: 38 additions & 8 deletions rustuna_core/src/storage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,21 @@ pub trait Storage: Send + Sync {
// are designed to receive multiple attributes for bulk insert operations.
// Furthermore, the `user_attrs` and `system_attrs` are merged into a single HashMap,
// which simplifies the implementation process for third-party storages.
fn set_study_attrs(&mut self, study_id: u32, attrs: Attrs) -> Result<()>;
fn set_trial_attrs(&mut self, study_id: u32, trial_number: u32, attrs: Attrs) -> Result<()>;
// Note: Some backend implementations (e.g., SQLite) may partially apply attributes across
// user/system tables when error_on_overwrite is true.
fn set_study_attrs(
&mut self,
study_id: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()>;
fn set_trial_attrs(
&mut self,
study_id: u32,
trial_number: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()>;
fn get_joint_search_space(&mut self, study_id: u32) -> Result<HashMap<String, Distribution>>;
}

Expand Down Expand Up @@ -191,26 +204,43 @@ impl Storage for InMemoryStorage {
Ok(trial)
}

fn set_study_attrs(&mut self, study_id: u32, attrs: Attrs) -> Result<()> {
fn set_study_attrs(
&mut self,
study_id: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()> {
let study = self
.studies
.iter_mut()
.find(|s| s.id == study_id)
.ok_or(Error::new(ErrorKind::StudyNotFound))?;
attrs.into_iter().for_each(|(key, value)| {
for (key, value) in attrs {
if error_on_overwrite && study.attrs.contains_key(&key) {
return Err(Error::new(ErrorKind::AttrOverwriteNotAllowed));
}
study.attrs.insert(key, value);
});
}
Ok(())
}

fn set_trial_attrs(&mut self, study_id: u32, trial_number: u32, attrs: Attrs) -> Result<()> {
fn set_trial_attrs(
&mut self,
study_id: u32,
trial_number: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()> {
let trial = get_mut_trials_by_study_id(&mut self.trials, study_id)?
.get_mut(trial_number as usize)
.ok_or(Error::new(ErrorKind::TrialNotFound))?;
check_trial_is_updatable(trial)?;
attrs.into_iter().for_each(|(key, value)| {
for (key, value) in attrs {
if error_on_overwrite && trial.attrs.contains_key(&key) {
return Err(Error::new(ErrorKind::AttrOverwriteNotAllowed));
}
trial.attrs.insert(key, value);
});
}
Ok(())
}

Expand Down
2 changes: 1 addition & 1 deletion rustuna_core/src/study.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ impl Study {
for (key, value) in attrs {
a.insert(AttrKey::User(key), value);
}
guard.set_study_attrs(self.id, a)?;
guard.set_study_attrs(self.id, a, false)?;
Ok(())
}
}
Expand Down
6 changes: 3 additions & 3 deletions rustuna_core/src/trial.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ impl Trial {
let mut guard = storage
.write()
.map_err(|_| Error::new(ErrorKind::Unexpected))?;
guard.set_study_attrs(study_id, category_labels)?;
guard.set_study_attrs(study_id, category_labels, false)?;
Ok(c)
}

Expand Down Expand Up @@ -157,7 +157,7 @@ impl Trial {

let key = AttrKey::User(key.to_string());
attrs.insert(key.clone(), value.clone());
guard.set_trial_attrs(self.study_id, self.number, attrs)?;
guard.set_trial_attrs(self.study_id, self.number, attrs, false)?;
self.cached_trial.attrs.insert(key, value);
Ok(())
}
Expand Down Expand Up @@ -226,7 +226,7 @@ mod tests {
let mut guard = storage.write().unwrap();
let mut attrs = Attrs::new();
attrs.insert(AttrKey::System("key".to_string()), "system".to_string());
guard.set_trial_attrs(study.id, 0, attrs).unwrap();
guard.set_trial_attrs(study.id, 0, attrs, false).unwrap();

// Check the attributes
let trial = guard.get_trial(study.id, 0).unwrap();
Expand Down
17 changes: 12 additions & 5 deletions rustuna_pyo3/src/pyobject_storage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ impl PyObjectStorage {
}?;
if sync_attrs {
self.cache
.set_study_attrs(cache_study_id, src_study.attrs.clone())?;
.set_study_attrs(cache_study_id, src_study.attrs.clone(), false)?;
}
Ok(())
}
Expand Down Expand Up @@ -247,7 +247,7 @@ impl PyObjectStorage {
rustuna_core::Error::new(rustuna_core::ErrorKind::StorageError)
})?;
self.cache
.set_study_attrs(cache_study_id, src_study.attrs)?;
.set_study_attrs(cache_study_id, src_study.attrs, false)?;
}
Ok(())
}
Expand Down Expand Up @@ -305,7 +305,7 @@ impl PyObjectStorage {
};
let cache_trial = cache_trial.clone();
self.cache
.set_trial_attrs(cache_study_id, cache_trial.number, src_trial.attrs)?;
.set_trial_attrs(cache_study_id, cache_trial.number, src_trial.attrs, false)?;

for (name, distribution) in src_trial.distributions {
let internal_repr =
Expand Down Expand Up @@ -509,10 +509,12 @@ impl Storage for PyObjectStorage {
&mut self,
study_id: u32,
attrs: rustuna_core::attr::Attrs,
_error_on_overwrite: bool,
) -> rustuna_core::Result<()> {
// TODO(c-bata): Emit warnings if error_on_overwrite is true, since Optuna storage cannot support it.
self.obj_set_study_attrs(study_id, attrs.clone())
.map_err(|_| rustuna_core::Error::new(rustuna_core::ErrorKind::StorageError))?;
match self.cache.set_study_attrs(study_id, attrs) {
match self.cache.set_study_attrs(study_id, attrs, false) {
Ok(_) => Ok(()),
Err(e) => match e.kind {
rustuna_core::ErrorKind::StudyNotFound => {
Expand All @@ -529,10 +531,15 @@ impl Storage for PyObjectStorage {
study_id: u32,
trial_number: u32,
attrs: rustuna_core::attr::Attrs,
_error_on_overwrite: bool,
) -> rustuna_core::Result<()> {
// TODO(c-bata): Emit warnings if error_on_overwrite is true, since Optuna storage cannot support it.
self.obj_set_trial_attrs(study_id, trial_number, attrs.clone())
.map_err(|_| rustuna_core::Error::new(rustuna_core::ErrorKind::StorageError))?;
match self.cache.set_trial_attrs(study_id, trial_number, attrs) {
match self
.cache
.set_trial_attrs(study_id, trial_number, attrs, false)
{
Ok(_) => Ok(()),
Err(e) => match e.kind {
rustuna_core::ErrorKind::StudyNotFound | rustuna_core::ErrorKind::TrialNotFound => {
Expand Down
38 changes: 28 additions & 10 deletions rustuna_pyo3/src/storage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,10 @@ impl PyStorage {
distribution: PyDistribution,
value: f64,
) -> PyResult<()> {
let mut guard = self.storage.write().unwrap();
let mut guard = self
.storage
.write()
.map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?;
let distribution: Distribution = distribution.into();
guard
.set_trial_param(study_id, trial_number, &name, &distribution, value)
Expand All @@ -105,7 +108,6 @@ impl PyStorage {
param_name: String,
choices: Vec<PyObject>,
) -> PyResult<()> {
// TODO(c-bata): Add validation to detect changes of the choices.

@c-bata c-bata Dec 5, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: The main motivation for introducing error_on_overwrite option is resolving this todo item.

let category_labels = Python::with_gil(|py| {
let mut labels: Vec<CategoryLabel> = Vec::with_capacity(choices.len());
for choice in choices {
Expand All @@ -121,10 +123,26 @@ impl PyStorage {
.storage
.write()
.map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?;
guard
.set_study_attrs(study_id, attrs)
.map_err(err_to_exceptions)?;
Ok(())
match guard.set_study_attrs(study_id, attrs, true) {
Ok(_) => Ok(()),
Err(e) => {
if matches!(e.kind, rustuna_core::ErrorKind::AttrOverwriteNotAllowed) {
let study = guard.get_study(study_id).map_err(err_to_exceptions)?;
let existing_labels =
get_category_labels(&study.attrs, &param_name, category_labels.len());
if let Some(existing) = existing_labels {
if existing == category_labels {
return Ok(());
}
}
return Err(PyValueError::new_err(format!(
"Cannot overwrite category labels for parameter '{}'",
param_name
)));
}
Err(err_to_exceptions(e))
}
}
}

fn get_category_labels(
Expand Down Expand Up @@ -234,7 +252,7 @@ impl PyStorage {
})?;
let mut guard = self.storage.write().unwrap();
guard
.set_study_attrs(study_id, system_attrs)
.set_study_attrs(study_id, system_attrs, false)
.map_err(err_to_exceptions)?;
Ok(())
}
Expand All @@ -246,7 +264,7 @@ impl PyStorage {
})?;
let mut guard = self.storage.write().unwrap();
guard
.set_study_attrs(study_id, user_attrs)
.set_study_attrs(study_id, user_attrs, false)
.map_err(err_to_exceptions)?;
Ok(())
}
Expand All @@ -263,7 +281,7 @@ impl PyStorage {
})?;
let mut guard = self.storage.write().unwrap();
guard
.set_trial_attrs(study_id, trial_number, system_attrs)
.set_trial_attrs(study_id, trial_number, system_attrs, false)
.map_err(err_to_exceptions)?;
Ok(())
}
Expand All @@ -280,7 +298,7 @@ impl PyStorage {
})?;
let mut guard = self.storage.write().unwrap();
guard
.set_trial_attrs(study_id, trial_number, user_attrs)
.set_trial_attrs(study_id, trial_number, user_attrs, false)
.map_err(err_to_exceptions)?;
Ok(())
}
Expand Down
2 changes: 1 addition & 1 deletion rustuna_samplers/src/nsgaii.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@ impl Sampler for NSGAIISampler {
AttrKey::System("generation".to_string()),
(child_generation as f64).to_string(),
);
guard.set_trial_attrs(ctx.study_id, ctx.trial_number, attrs)?;
guard.set_trial_attrs(ctx.study_id, ctx.trial_number, attrs, false)?;
drop(guard);

if parent_generation < 0 {
Expand Down
53 changes: 42 additions & 11 deletions rustuna_storages/src/cache.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,19 @@ pub trait CachedStorageBackend: Send + Sync {
fn get_studies(&mut self) -> Result<Vec<PersistedStudy>>;
fn get_study(&mut self, study_id: u32) -> Result<PersistedStudy>;
fn get_trial(&mut self, study_id: u32, trial_number: u32) -> Result<PersistedTrial>;
fn set_study_attrs(&mut self, study_id: u32, attrs: Attrs) -> Result<()>;
fn set_trial_attrs(&mut self, study_id: u32, trial_number: u32, attrs: Attrs) -> Result<()>;
fn set_study_attrs(
&mut self,
study_id: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()>;
fn set_trial_attrs(
&mut self,
study_id: u32,
trial_number: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()>;

// Return trials that need refreshing: unfinished trials in `included_numbers`
// and trials with trial_number greater than `trial_number_greater_than`.
Expand Down Expand Up @@ -280,8 +291,14 @@ impl rustuna_core::storage::Storage for CachedStorage {
Ok(trial)
}

fn set_study_attrs(&mut self, study_id: u32, attrs: Attrs) -> Result<()> {
self.backend.set_study_attrs(study_id, attrs.clone())?;
fn set_study_attrs(
&mut self,
study_id: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()> {
self.backend
.set_study_attrs(study_id, attrs.clone(), error_on_overwrite)?;
self.studies = self.backend.get_studies()?;
let study = self
.studies
Expand All @@ -294,9 +311,15 @@ impl rustuna_core::storage::Storage for CachedStorage {
Ok(())
}

fn set_trial_attrs(&mut self, study_id: u32, trial_number: u32, attrs: Attrs) -> Result<()> {
fn set_trial_attrs(
&mut self,
study_id: u32,
trial_number: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()> {
self.backend
.set_trial_attrs(study_id, trial_number, attrs.clone())?;
.set_trial_attrs(study_id, trial_number, attrs.clone(), error_on_overwrite)?;
self.refresh_trials(study_id)?;
let trials = self
.trials
Expand Down Expand Up @@ -415,17 +438,25 @@ mod tests {
Ok(self.inner.get_trial(study_id, trial_number)?.clone())
}

fn set_study_attrs(&mut self, study_id: u32, attrs: Attrs) -> Result<()> {
self.inner.set_study_attrs(study_id, attrs)
fn set_study_attrs(
&mut self,
study_id: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()> {
self.inner
.set_study_attrs(study_id, attrs, error_on_overwrite)
}

fn set_trial_attrs(
&mut self,
study_id: u32,
trial_number: u32,
attrs: Attrs,
error_on_overwrite: bool,
) -> Result<()> {
self.inner.set_trial_attrs(study_id, trial_number, attrs)
self.inner
.set_trial_attrs(study_id, trial_number, attrs, error_on_overwrite)
}
}

Expand Down Expand Up @@ -608,7 +639,7 @@ mod tests {

let mut s_attrs = Attrs::new();
s_attrs.insert(AttrKey::User("foo".to_string()), "bar".to_string());
storage.set_study_attrs(study_id, s_attrs)?;
storage.set_study_attrs(study_id, s_attrs, false)?;
let study = storage.get_study(study_id)?;
assert_eq!(
study.attrs.get(&AttrKey::User("foo".to_string())).unwrap(),
Expand All @@ -617,7 +648,7 @@ mod tests {

let mut t_attrs = Attrs::new();
t_attrs.insert(AttrKey::System("key".to_string()), "val".to_string());
storage.set_trial_attrs(study_id, 0, t_attrs)?;
storage.set_trial_attrs(study_id, 0, t_attrs, false)?;
let trial = storage.get_trial(study_id, 0)?;
assert_eq!(
trial
Expand Down
Loading