From e97d1e3c42c83bb98bc72e80a83ee8520a1fe646 Mon Sep 17 00:00:00 2001 From: c-bata Date: Tue, 3 Feb 2026 10:49:39 +0900 Subject: [PATCH 1/2] Revert "Add no_sync option to avoid syncing with backend" This reverts commit fd71624443c6f3b15e7ece20ed0dd31a3a2d989e. --- rustuna_core/src/storage.rs | 12 ++--- rustuna_core/src/study.rs | 6 +-- rustuna_core/src/trial.rs | 2 +- rustuna_importance/src/fanova.rs | 2 +- rustuna_js/src/study.rs | 2 +- rustuna_pyo3/rustuna/_rustuna.pyi | 10 ++-- rustuna_pyo3/rustuna/converter/_storage.py | 8 +--- rustuna_pyo3/src/pyobject_storage.rs | 22 ++++----- rustuna_pyo3/src/storage.rs | 18 ++------ rustuna_pyo3/src/study.rs | 6 +-- rustuna_pyo3/src/trial.rs | 4 +- rustuna_samplers/src/nsgaii.rs | 4 +- rustuna_samplers/src/tpe/sampler.rs | 6 +-- rustuna_storages/src/cache.rs | 54 ++++++++-------------- rustuna_storages/src/journal/storage.rs | 38 +++++++-------- rustuna_storages/tests/journal.rs | 8 ++-- rustuna_storages/tests/sqlite3.rs | 4 +- 17 files changed, 82 insertions(+), 124 deletions(-) diff --git a/rustuna_core/src/storage.rs b/rustuna_core/src/storage.rs index 6f74b1f4..cf1b3f12 100644 --- a/rustuna_core/src/storage.rs +++ b/rustuna_core/src/storage.rs @@ -33,12 +33,8 @@ pub trait Storage: Send + Sync { // references without relying on unsafe patterns. fn get_studies(&mut self) -> Result<&Vec>; fn get_study(&mut self, study_id: u32) -> Result<&PersistedStudy>; - // Design Note: - // `no_sync` avoids refreshing or syncing from the backend and reads from the - // current in-memory cache. Implementations may return TrialNotFound or StudyNotFound - // if the requested data is not cached yet. - fn get_trials(&mut self, study_id: u32, no_sync: bool) -> Result<&Vec>; - fn get_trial(&mut self, trial_id: u32, no_sync: bool) -> Result<&PersistedTrial>; + fn get_trials(&mut self, study_id: u32) -> Result<&Vec>; + fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial>; // Design Note: // Category labels are stored in study system attrs internally, but exposed via dedicated // APIs for caching efficiency. Since category labels cannot be overwritten once set for @@ -304,11 +300,11 @@ impl Storage for InMemoryStorage { Ok(study) } - fn get_trials(&mut self, study_id: u32, _no_sync: bool) -> Result<&Vec> { + fn get_trials(&mut self, study_id: u32) -> Result<&Vec> { get_trials_by_study_id(&self.trials, study_id) } - fn get_trial(&mut self, trial_id: u32, _no_sync: bool) -> Result<&PersistedTrial> { + fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial> { let (study_id, trial_number) = get_study_id_trial_number_by_trial_id(&self.trial_id_to_study_number, trial_id)?; let trial = get_trials_by_study_id(&self.trials, study_id)? diff --git a/rustuna_core/src/study.rs b/rustuna_core/src/study.rs index e82dffec..16d7c8b0 100644 --- a/rustuna_core/src/study.rs +++ b/rustuna_core/src/study.rs @@ -202,7 +202,7 @@ impl Study { .storage .write() .map_err(|_| Error::new(ErrorKind::StorageError))?; - let trials = guard.get_trials(self.id, false)?; + let trials = guard.get_trials(self.id)?; Ok(trials.clone()) } @@ -278,7 +278,7 @@ pub fn get_best_trial(study: &Study) -> Result { .storage .write() .map_err(|_| Error::new(ErrorKind::StorageError))?; - let trials = guard.get_trials(study.id, false)?; + let trials = guard.get_trials(study.id)?; let best_trial = trials .iter() @@ -313,7 +313,7 @@ pub fn get_pareto_front(study: &Study) -> Result> { .write() .map_err(|_| Error::new(ErrorKind::StorageError))?; let trials = guard - .get_trials(study.id, false)? + .get_trials(study.id)? .iter() .filter(|t| matches!(t.state_values, TrialStateValues::Complete(ref _v))) .collect::>(); diff --git a/rustuna_core/src/trial.rs b/rustuna_core/src/trial.rs index 55ddc187..72c08c48 100644 --- a/rustuna_core/src/trial.rs +++ b/rustuna_core/src/trial.rs @@ -250,7 +250,7 @@ mod tests { guard.set_trial_attrs(trial.id, attrs, false)?; // Check the attributes - let trial = guard.get_trial(trial.id, false)?; + let trial = guard.get_trial(trial.id)?; assert_eq!( trial.attrs.get(&AttrKey::User("key".into())), Some(&"user".to_string()) diff --git a/rustuna_importance/src/fanova.rs b/rustuna_importance/src/fanova.rs index 94a83482..25994cb2 100644 --- a/rustuna_importance/src/fanova.rs +++ b/rustuna_importance/src/fanova.rs @@ -11,7 +11,7 @@ pub fn get_param_importance(study: &Study) -> Result>> { .map_err(|_e| Error::new(ErrorKind::Unexpected))?; // TODO(c-bata): Avoid to clone trials. let completed_trials: Vec = guard - .get_trials(study.id, false)? + .get_trials(study.id)? .clone() .into_iter() .filter(|t| matches!(t.state_values, TrialStateValues::Complete(_))) diff --git a/rustuna_js/src/study.rs b/rustuna_js/src/study.rs index 3c0190d9..11843fd9 100644 --- a/rustuna_js/src/study.rs +++ b/rustuna_js/src/study.rs @@ -45,7 +45,7 @@ impl JsStudy { JsError::new(&format!("Failed to acquire the storage guard: {e:?}")) })?; let trials = guard - .get_trials(self.0.id, false) + .get_trials(self.0.id) .map_err(|e| JsError::new(&format!("Failed to get trials: {:?}", e.kind)))? .clone(); let study_attrs = guard diff --git a/rustuna_pyo3/rustuna/_rustuna.pyi b/rustuna_pyo3/rustuna/_rustuna.pyi index a47b25d7..9c3ce4f6 100644 --- a/rustuna_pyo3/rustuna/_rustuna.pyi +++ b/rustuna_pyo3/rustuna/_rustuna.pyi @@ -409,10 +409,8 @@ class StorageProtocol(Protocol): ) -> None: ... def get_studies(self) -> list[PersistedStudy]: ... def get_study(self, study_id: int) -> PersistedStudy: ... - def get_trials( - self, study_id: int, *, no_sync: bool = ... - ) -> list[PersistedTrial]: ... - def get_trial(self, trial_id: int, *, no_sync: bool = ...) -> PersistedTrial: ... + def get_trials(self, study_id: int) -> list[PersistedTrial]: ... + def get_trial(self, trial_id: int) -> PersistedTrial: ... def get_trial_id_from_study_id_trial_number( self, study_id: int, trial_number: int ) -> int: ... @@ -544,7 +542,7 @@ class Storage: Returns: The study. """ - def get_trials(self, study_id: int, *, no_sync: bool = ...) -> list[PersistedTrial]: + def get_trials(self, study_id: int) -> list[PersistedTrial]: """Get all trials in a study. Args: @@ -562,7 +560,7 @@ class Storage: Returns: Number of trials. """ - def get_trial(self, trial_id: int, *, no_sync: bool = ...) -> PersistedTrial: + def get_trial(self, trial_id: int) -> PersistedTrial: """Get a trial by ID. Args: diff --git a/rustuna_pyo3/rustuna/converter/_storage.py b/rustuna_pyo3/rustuna/converter/_storage.py index 30ee46a7..3b87c046 100644 --- a/rustuna_pyo3/rustuna/converter/_storage.py +++ b/rustuna_pyo3/rustuna/converter/_storage.py @@ -101,9 +101,7 @@ def get_study(self, study_id: int) -> rustuna.PersistedStudy: return to_persisted_study(s) raise KeyError(f"Study {study_id} not found") - def get_trials( - self, study_id: int, *, no_sync: bool = False - ) -> list[rustuna.PersistedTrial]: + def get_trials(self, study_id: int) -> list[rustuna.PersistedTrial]: frozen_trials = self._storage.get_all_trials(study_id) persisted_trials: list[rustuna.PersistedTrial] = [] @@ -113,9 +111,7 @@ def get_trials( self._trial_id_to_study_id[t._trial_id] = study_id return persisted_trials - def get_trial( - self, trial_id: int, *, no_sync: bool = False - ) -> rustuna.PersistedTrial: + def get_trial(self, trial_id: int) -> rustuna.PersistedTrial: with self._lock: study_id = self._trial_id_to_study_id.get(trial_id, -1) if study_id == -1: diff --git a/rustuna_pyo3/src/pyobject_storage.rs b/rustuna_pyo3/src/pyobject_storage.rs index 3df7f340..f9244549 100644 --- a/rustuna_pyo3/src/pyobject_storage.rs +++ b/rustuna_pyo3/src/pyobject_storage.rs @@ -249,7 +249,7 @@ impl PyObjectStorage { ) -> rustuna_core::Result<()> { let cache_n_trials = match cache_n_trials { Some(n) => n, - None => self.cache.get_trials(cache_study_id, false)?.len() as u32, + None => self.cache.get_trials(cache_study_id)?.len() as u32, }; match src_trial.number.cmp(&cache_n_trials) { std::cmp::Ordering::Equal => { @@ -273,7 +273,7 @@ impl PyObjectStorage { } } - let cache_trial = self.cache.get_trial(src_trial.id, false)?.clone(); + let cache_trial = self.cache.get_trial(src_trial.id)?.clone(); if cache_trial.is_finished() { return Ok(()); } @@ -315,10 +315,10 @@ impl PyObjectStorage { .map_err(|_| rustuna_core::Error::new(rustuna_core::ErrorKind::StorageError))?; src_trials.sort_by_key(|trial| trial.number); - let mut cache_n_trials = self.cache.get_trials(cache_study_id, false)?.len() as u32; + let mut cache_n_trials = self.cache.get_trials(cache_study_id)?.len() as u32; for src_trial in src_trials { self.sync_trial(cache_study_id, src_trial, Some(cache_n_trials))?; - cache_n_trials = self.cache.get_trials(cache_study_id, false)?.len() as u32; + cache_n_trials = self.cache.get_trials(cache_study_id)?.len() as u32; } Ok(()) } @@ -385,15 +385,15 @@ impl Storage for PyObjectStorage { let src_trial_id = src_trial.id; if self.is_distributed { self.sync_trials(study_id)?; - return self.cache.get_trial(src_trial_id, false); + return self.cache.get_trial(src_trial_id); } - let cached_n_trials = self.cache.get_trials(study_id, false)?.len() as u32; + let cached_n_trials = self.cache.get_trials(study_id)?.len() as u32; if src_trial.number != cached_n_trials { self.sync_trials(study_id)?; - return self.cache.get_trial(src_trial_id, false); + return self.cache.get_trial(src_trial_id); } self.sync_trial(study_id, src_trial, Some(cached_n_trials))?; - self.cache.get_trial(src_trial_id, false) + self.cache.get_trial(src_trial_id) } fn set_trial_param( @@ -459,17 +459,15 @@ impl Storage for PyObjectStorage { fn get_trials( &mut self, study_id: u32, - _no_sync: bool, ) -> rustuna_core::Result<&Vec> { - self.cache.get_trials(study_id, false) + self.cache.get_trials(study_id) } fn get_trial( &mut self, trial_id: u32, - _no_sync: bool, ) -> rustuna_core::Result<&rustuna_core::trial::PersistedTrial> { - self.cache.get_trial(trial_id, false) + self.cache.get_trial(trial_id) } fn get_category_labels( diff --git a/rustuna_pyo3/src/storage.rs b/rustuna_pyo3/src/storage.rs index 61e1c2e6..9899a44c 100644 --- a/rustuna_pyo3/src/storage.rs +++ b/rustuna_pyo3/src/storage.rs @@ -133,7 +133,7 @@ impl PyStorage { .write() .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; guard - .get_trial(trial_id, false) + .get_trial(trial_id) .map_err(err_to_exceptions)? .study_id }; @@ -249,15 +249,12 @@ impl PyStorage { Ok(study.clone().into()) } - #[pyo3(signature = (study_id, *, no_sync = false))] - fn get_trials(&mut self, study_id: u32, no_sync: bool) -> PyResult> { + fn get_trials(&mut self, study_id: u32) -> PyResult> { let mut guard = self .storage .write() .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; - let trials = guard - .get_trials(study_id, no_sync) - .map_err(err_to_exceptions)?; + let trials = guard.get_trials(study_id).map_err(err_to_exceptions)?; let py_trials: Vec = trials .iter() .map(|t| PyPersistedTrial::from_storage(self.storage.clone(), t)) @@ -265,18 +262,13 @@ impl PyStorage { Ok(py_trials) } - #[pyo3(signature = (trial_id, *, no_sync = false))] - fn get_trial(&mut self, trial_id: u32, no_sync: bool) -> PyResult { + fn get_trial(&mut self, trial_id: u32) -> PyResult { let mut guard = self .storage .write() .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; - if no_sync { - let trial = guard.get_trial(trial_id, true).map_err(err_to_exceptions)?; - return Ok(PyPersistedTrial::from_storage(self.storage.clone(), trial)); - } let trial = guard - .get_trial(trial_id, false) + .get_trial(trial_id) .map_err(err_to_exceptions)? .clone(); let study_attrs = guard diff --git a/rustuna_pyo3/src/study.rs b/rustuna_pyo3/src/study.rs index 06059198..0e3b967d 100644 --- a/rustuna_pyo3/src/study.rs +++ b/rustuna_pyo3/src/study.rs @@ -309,7 +309,7 @@ impl PyStudy { PyRuntimeError::new_err(format!("Failed to get trial id: {:?}", e.kind)) })?; let trial = guard - .get_trial(trial_id, false) + .get_trial(trial_id) .map_err(|e| PyRuntimeError::new_err(format!("Failed to get trial: {:?}", e.kind)))?; Ok(PyPersistedTrial::from_storage( self.study.storage.clone(), @@ -325,7 +325,7 @@ impl PyStudy { .write() .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; let trials_vec = guard - .get_trials(self.study.id, false) + .get_trials(self.study.id) .map_err(|e| PyRuntimeError::new_err(format!("Failed to get trials: {:?}", e.kind)))?; let trials: Vec = trials_vec .iter() @@ -361,7 +361,7 @@ impl PyStudy { .write() .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; let trials_vec = guard - .get_trials(self.study.id, false) + .get_trials(self.study.id) .map_err(|e| PyRuntimeError::new_err(format!("Failed to get trials: {:?}", e.kind)))?; let best_trials = pareto_front_numbers .iter() diff --git a/rustuna_pyo3/src/trial.rs b/rustuna_pyo3/src/trial.rs index adc08310..97ac3a1e 100644 --- a/rustuna_pyo3/src/trial.rs +++ b/rustuna_pyo3/src/trial.rs @@ -230,9 +230,7 @@ impl PyPersistedTrial { let mut guard = storage .write() .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; - let trial = guard - .get_trial(*trial_id, true) - .map_err(err_to_exceptions)?; + let trial = guard.get_trial(*trial_id).map_err(err_to_exceptions)?; f(trial) } } diff --git a/rustuna_samplers/src/nsgaii.rs b/rustuna_samplers/src/nsgaii.rs index 5a7b0497..d5b4ea41 100644 --- a/rustuna_samplers/src/nsgaii.rs +++ b/rustuna_samplers/src/nsgaii.rs @@ -131,7 +131,7 @@ impl NSGAIISampler { let mut guard = storage .write() .map_err(|_e| Error::new(ErrorKind::Unexpected))?; - let trials = guard.get_trials(ctx.study_id, false)?; + let trials = guard.get_trials(ctx.study_id)?; let population_params = trials .iter() .filter(|trial| population_numbers.contains(&trial.number)) @@ -220,7 +220,7 @@ impl Sampler for NSGAIISampler { let mut guard = storage .write() .map_err(|_e| Error::new(ErrorKind::Unexpected))?; - let trials = guard.get_trials(ctx.study_id, false)?; + let trials = guard.get_trials(ctx.study_id)?; let (parent_generation, parent_population_numbers) = self.get_parent_population_numbers(ctx, trials)?; let child_generation = u32::try_from(parent_generation + 1).unwrap(); diff --git a/rustuna_samplers/src/tpe/sampler.rs b/rustuna_samplers/src/tpe/sampler.rs index 7cafcc04..00e42ec8 100644 --- a/rustuna_samplers/src/tpe/sampler.rs +++ b/rustuna_samplers/src/tpe/sampler.rs @@ -84,7 +84,7 @@ impl TpeSampler { let mut guard = storage .write() .map_err(|_e| Error::new(ErrorKind::Unexpected))?; - let trials = guard.get_trials(ctx.study_id, false)?.clone(); + let trials = guard.get_trials(ctx.study_id)?.clone(); drop(guard); let complete_trials: Vec<_> = trials.into_iter().filter(|t| t.is_finished()).collect(); @@ -396,7 +396,7 @@ impl Sampler for TpeSampler { let mut guard = storage .write() .map_err(|_e| Error::new(ErrorKind::Unexpected))?; - let trials = guard.get_trials(ctx.study_id, false)?.clone(); + let trials = guard.get_trials(ctx.study_id)?.clone(); drop(guard); let complete_trials: Vec<_> = trials.into_iter().filter(|t| t.is_finished()).collect(); @@ -425,7 +425,7 @@ impl Sampler for TpeSampler { let mut guard = storage .write() .map_err(|_e| Error::new(ErrorKind::Unexpected))?; - let trials = guard.get_trials(ctx.study_id, false)?.clone(); + let trials = guard.get_trials(ctx.study_id)?.clone(); drop(guard); let complete_trials: Vec<_> = trials.into_iter().filter(|t| t.is_finished()).collect(); diff --git a/rustuna_storages/src/cache.rs b/rustuna_storages/src/cache.rs index 9928d563..7cfdca8b 100644 --- a/rustuna_storages/src/cache.rs +++ b/rustuna_storages/src/cache.rs @@ -319,10 +319,8 @@ impl rustuna_core::storage::Storage for CachedStorage { Ok(study) } - fn get_trials(&mut self, study_id: u32, no_sync: bool) -> Result<&Vec> { - if !no_sync { - self.refresh_trials(study_id)?; - } + fn get_trials(&mut self, study_id: u32) -> Result<&Vec> { + self.refresh_trials(study_id)?; let trials_map = self .trials .get(&study_id) @@ -338,21 +336,7 @@ impl rustuna_core::storage::Storage for CachedStorage { Ok(&self.trials_sorted_buffer) } - fn get_trial(&mut self, trial_id: u32, no_sync: bool) -> Result<&PersistedTrial> { - if no_sync { - let (study_id, trial_number) = self - .trial_id_to_study_number - .get(&trial_id) - .copied() - .ok_or_else(|| Error::new(ErrorKind::TrialNotFound))?; - let trials = self - .trials - .get(&study_id) - .ok_or_else(|| Error::new(ErrorKind::StudyNotFound))?; - return trials - .get(&trial_number) - .ok_or_else(|| Error::new(ErrorKind::TrialNotFound)); - } + fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial> { let (study_id, trial_number) = self.resolve_trial_location(trial_id)?; self.refresh_trials(study_id)?; let trials = self @@ -470,7 +454,7 @@ impl rustuna_core::storage::Storage for CachedStorage { fn get_joint_search_space(&mut self, study_id: u32) -> Result> { let trials_vec = { - let trials = self.get_trials(study_id, false)?; + let trials = self.get_trials(study_id)?; let mut v = trials.clone(); v.sort_by_key(|t| t.number); v @@ -568,7 +552,7 @@ mod tests { included_numbers: &[u32], trial_number_greater_than: i32, ) -> Result> { - let all = self.inner.get_trials(study_id, false)?.clone(); + let all = self.inner.get_trials(study_id)?.clone(); let mut trials = Vec::new(); for t in all { if included_numbers.contains(&t.number) @@ -581,7 +565,7 @@ mod tests { } fn get_trial(&mut self, trial_id: u32) -> Result { - Ok(self.inner.get_trial(trial_id, false)?.clone()) + Ok(self.inner.get_trial(trial_id)?.clone()) } fn set_study_attrs( @@ -665,7 +649,7 @@ mod tests { } #[test] - fn get_study_and_get_studies_no_sync() -> Result<()> { + fn get_study_and_get_studies_use_cache() -> Result<()> { let mut storage = CachedStorage::new(Box::new(DummyBackend::new())); storage.create_new_study("s1", vec![Direction::Minimize])?; storage.create_new_study("s2", vec![Direction::Maximize])?; @@ -703,11 +687,11 @@ mod tests { let t0_id = storage.create_new_trial(study_id)?.id; let t1_id = storage.create_new_trial(study_id)?.id; - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 2); - let t0 = storage.get_trial(t0_id, false)?; + let t0 = storage.get_trial(t0_id)?; assert_eq!(t0.number, 0); - let t1 = storage.get_trial(t1_id, false)?; + let t1 = storage.get_trial(t1_id)?; assert_eq!(t1.number, 1); Ok(()) } @@ -719,7 +703,7 @@ mod tests { storage.create_new_trial(study_id)?; storage.trials.clear(); - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 1); Ok(()) } @@ -750,11 +734,11 @@ mod tests { backend.create_new_trial(study_id)?; let mut storage = CachedStorage::new(Box::new(backend)); - let trials1 = storage.get_trials(study_id, false)?; + let trials1 = storage.get_trials(study_id)?; assert_eq!(trials1.len(), 1); storage.backend.create_new_trial(study_id)?; - let trials2 = storage.get_trials(study_id, false)?; + let trials2 = storage.get_trials(study_id)?; assert_eq!(trials2.len(), 2); Ok(()) } @@ -767,7 +751,7 @@ mod tests { let mut storage = CachedStorage::new(Box::new(backend)); storage.set_trial_state_values(trial_id, TrialStateValues::Complete(vec![1.0]))?; - let trial = storage.get_trial(trial_id, false)?; + let trial = storage.get_trial(trial_id)?; assert!(matches!(trial.state_values, TrialStateValues::Complete(_))); Ok(()) } @@ -813,7 +797,7 @@ mod tests { let mut t_attrs = Attrs::new(); t_attrs.insert(AttrKey::System("key".into()), "val".to_string()); storage.set_trial_attrs(trial_id, t_attrs, false)?; - let trial = storage.get_trial(trial_id, false)?; + let trial = storage.get_trial(trial_id)?; assert_eq!( trial .attrs @@ -837,10 +821,10 @@ mod tests { step: None, log: false, }; - let trial_id = storage.get_trials(study_id, false)?[0].id; + let trial_id = storage.get_trials(study_id)?[0].id; storage.set_trial_param(trial_id, "x", &dist, 0.5)?; - let trial = storage.get_trial(trial_id, false)?; + let trial = storage.get_trial(trial_id)?; assert_eq!(trial.internal_params.get("x"), Some(&0.5)); assert_eq!( trial.distributions.get("x"), @@ -887,14 +871,14 @@ mod tests { // Set new params. storage.set_trial_param(trial_1_id, "x", &distribution_x, 0.5)?; storage.set_trial_param(trial_1_id, "y", &distribution_y_1, 2.0)?; - let trial = storage.get_trial(trial_1_id, false)?; + let trial = storage.get_trial(trial_1_id)?; assert_eq!(trial.internal_params["x"], 0.5); assert_eq!(trial.internal_params["y"], 2.0); // Set params to another trial storage.set_trial_param(trial_2_id, "x", &distribution_x, 0.3)?; storage.set_trial_param(trial_2_id, "z", &distribution_z, 0.1)?; - let trial = storage.get_trial(trial_2_id, false)?; + let trial = storage.get_trial(trial_2_id)?; assert_eq!(trial.internal_params["x"], 0.3); assert_eq!(trial.internal_params["z"], 0.1); diff --git a/rustuna_storages/src/journal/storage.rs b/rustuna_storages/src/journal/storage.rs index 3733a8a5..33887122 100644 --- a/rustuna_storages/src/journal/storage.rs +++ b/rustuna_storages/src/journal/storage.rs @@ -280,10 +280,8 @@ impl Storage for JournalStorage { }) } - fn get_trials(&mut self, study_id: u32, no_sync: bool) -> Result<&Vec> { - if !no_sync { - self.sync_with_backend()?; - } + fn get_trials(&mut self, study_id: u32) -> Result<&Vec> { + self.sync_with_backend()?; let trials = self.replay.trials_by_study.get(&study_id).ok_or_else(|| { Error::with_reason( ErrorKind::StudyNotFound, @@ -293,10 +291,8 @@ impl Storage for JournalStorage { Ok(trials) } - fn get_trial(&mut self, trial_id: u32, no_sync: bool) -> Result<&PersistedTrial> { - if !no_sync { - self.sync_with_backend()?; - } + fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial> { + self.sync_with_backend()?; let (study_id, trial_number) = self .replay .trial_id_to_study_number @@ -1968,7 +1964,7 @@ mod tests { } #[test] - fn get_study_and_get_studies_no_sync() -> Result<()> { + fn get_study_and_get_studies_use_cache() -> Result<()> { let (mut storage, _logs) = new_storage()?; storage.create_new_study("s1", vec![Direction::Minimize])?; storage.create_new_study("s2", vec![Direction::Maximize])?; @@ -1991,7 +1987,7 @@ mod tests { let t1_num = storage.create_new_trial(study)?.number; assert_eq!(t0_num, 0); assert_eq!(t1_num, 1); - let trials = storage.get_trials(study, false)?; + let trials = storage.get_trials(study)?; assert_eq!(trials.len(), 2); Ok(()) } @@ -2003,11 +1999,11 @@ mod tests { let t0_id = storage.create_new_trial(study_id)?.id; let t1_id = storage.create_new_trial(study_id)?.id; - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 2); - let t0 = storage.get_trial(t0_id, false)?; + let t0 = storage.get_trial(t0_id)?; assert_eq!(t0.number, 0); - let t1 = storage.get_trial(t1_id, false)?; + let t1 = storage.get_trial(t1_id)?; assert_eq!(t1.number, 1); Ok(()) } @@ -2020,7 +2016,7 @@ mod tests { let backend = InMemoryJournalBackend { logs: logs.clone() }; let mut storage2 = JournalStorage::new(Box::new(backend))?; - let trials = storage2.get_trials(study_id, false)?; + let trials = storage2.get_trials(study_id)?; assert_eq!(trials.len(), 1); Ok(()) } @@ -2046,7 +2042,7 @@ mod tests { storage.create_new_trial(study_id)?; append_external_log(&logs, create_trial_log(study_id)); - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 2); Ok(()) } @@ -2058,7 +2054,7 @@ mod tests { let trial_id = storage.create_new_trial(study_id)?.id; storage.set_trial_state_values(trial_id, TrialStateValues::Complete(vec![1.0]))?; - let trial = storage.get_trial(trial_id, false)?; + let trial = storage.get_trial(trial_id)?; assert!(matches!(trial.state_values, TrialStateValues::Complete(_))); Ok(()) } @@ -2104,7 +2100,7 @@ mod tests { let mut t_attrs = Attrs::new(); t_attrs.insert(AttrKey::System("key".into()), "val".to_string()); storage.set_trial_attrs(trial_id, t_attrs, false)?; - let trial = storage.get_trial(trial_id, false)?; + let trial = storage.get_trial(trial_id)?; assert_eq!( trial .attrs @@ -2127,10 +2123,10 @@ mod tests { step: None, log: false, }; - let trial_id = storage.get_trials(study_id, false)?[0].id; + let trial_id = storage.get_trials(study_id)?[0].id; storage.set_trial_param(trial_id, "x", &dist, 0.5)?; - let trial = storage.get_trial(trial_id, false)?; + let trial = storage.get_trial(trial_id)?; assert_eq!(trial.internal_params.get("x"), Some(&0.5)); assert_eq!( trial.distributions.get("x"), @@ -2170,13 +2166,13 @@ mod tests { storage.set_trial_param(trial_1_id, "x", &distribution_x, 0.5)?; storage.set_trial_param(trial_1_id, "y", &distribution_y_1, 2.0)?; - let trial = storage.get_trial(trial_1_id, false)?; + let trial = storage.get_trial(trial_1_id)?; assert_eq!(trial.internal_params["x"], 0.5); assert_eq!(trial.internal_params["y"], 2.0); storage.set_trial_param(trial_2_id, "x", &distribution_x, 0.3)?; storage.set_trial_param(trial_2_id, "z", &distribution_z, 0.1)?; - let trial = storage.get_trial(trial_2_id, false)?; + let trial = storage.get_trial(trial_2_id)?; assert_eq!(trial.internal_params["x"], 0.3); assert_eq!(trial.internal_params["z"], 0.1); diff --git a/rustuna_storages/tests/journal.rs b/rustuna_storages/tests/journal.rs index b38f781d..5cc51e38 100644 --- a/rustuna_storages/tests/journal.rs +++ b/rustuna_storages/tests/journal.rs @@ -120,8 +120,8 @@ study.optimize(objective, n_trials=10) studies[0].id }; - let trial_id = storage.get_trials(study_id, false)?[0].id; - let trial0 = storage.get_trial(trial_id, false)?; + let trial_id = storage.get_trials(study_id)?[0].id; + let trial0 = storage.get_trial(trial_id)?; assert_eq!(trial0.number, 0); // Distributions @@ -192,7 +192,7 @@ study.optimize(objective, n_trials=10) assert_eq!(studies.len(), 1); studies[0].id }; - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 10); // Evaluate more 10 trials with load_if_exists=True @@ -220,7 +220,7 @@ study.optimize(objective, n_trials=10) script_continue, )?; let mut storage = JournalStorage::new(Box::new(JournalFileBackend::new(&journal_path, None)?))?; - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 20); assert_eq!(trials[0].distributions.len(), 3); assert_eq!( diff --git a/rustuna_storages/tests/sqlite3.rs b/rustuna_storages/tests/sqlite3.rs index 5a693db8..2eff954e 100644 --- a/rustuna_storages/tests/sqlite3.rs +++ b/rustuna_storages/tests/sqlite3.rs @@ -157,12 +157,12 @@ study.optimize(objective, n_trials=10) assert_eq!(studies.len(), 1); studies[0].id }; - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 10); // Evaluate more 10 trials run_optuna_script(&python, &db_path, script)?; - let trials = storage.get_trials(study_id, false)?; + let trials = storage.get_trials(study_id)?; assert_eq!(trials.len(), 20); assert_eq!(trials[0].distributions.len(), 3); assert_eq!( From c0692a8d6532ac9d3b5fb59d0b63b19e8eeb33bd Mon Sep 17 00:00:00 2001 From: c-bata Date: Tue, 3 Feb 2026 11:40:25 +0900 Subject: [PATCH 2/2] Add get_cached_trial --- rustuna_core/src/storage.rs | 9 +++++++++ rustuna_pyo3/rustuna/_rustuna.pyi | 10 ++++++++++ rustuna_pyo3/rustuna/converter/_storage.py | 3 +++ rustuna_pyo3/src/pyobject_storage.rs | 7 +++++++ rustuna_pyo3/src/storage.rs | 11 +++++++++++ rustuna_pyo3/src/trial.rs | 8 +++++--- rustuna_storages/src/cache.rs | 11 ++++++++++- rustuna_storages/src/journal/storage.rs | 4 ++++ 8 files changed, 59 insertions(+), 4 deletions(-) diff --git a/rustuna_core/src/storage.rs b/rustuna_core/src/storage.rs index cf1b3f12..ce06268c 100644 --- a/rustuna_core/src/storage.rs +++ b/rustuna_core/src/storage.rs @@ -36,6 +36,11 @@ pub trait Storage: Send + Sync { fn get_trials(&mut self, study_id: u32) -> Result<&Vec>; fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial>; // Design Note: + // get_cached_* methods always return references from the in-memory cache without + // synchronizing with backends. These methods are separate from get_* to allow + // callers to use read locks for cache-only reads. + fn get_cached_trial(&self, trial_id: u32) -> Result<&PersistedTrial>; + // Design Note: // Category labels are stored in study system attrs internally, but exposed via dedicated // APIs for caching efficiency. Since category labels cannot be overwritten once set for // a given (study_id, param_name), implementations can safely cache them without @@ -305,6 +310,10 @@ impl Storage for InMemoryStorage { } fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial> { + self.get_cached_trial(trial_id) + } + + fn get_cached_trial(&self, trial_id: u32) -> Result<&PersistedTrial> { let (study_id, trial_number) = get_study_id_trial_number_by_trial_id(&self.trial_id_to_study_number, trial_id)?; let trial = get_trials_by_study_id(&self.trials, study_id)? diff --git a/rustuna_pyo3/rustuna/_rustuna.pyi b/rustuna_pyo3/rustuna/_rustuna.pyi index 9c3ce4f6..142d802e 100644 --- a/rustuna_pyo3/rustuna/_rustuna.pyi +++ b/rustuna_pyo3/rustuna/_rustuna.pyi @@ -411,6 +411,7 @@ class StorageProtocol(Protocol): def get_study(self, study_id: int) -> PersistedStudy: ... def get_trials(self, study_id: int) -> list[PersistedTrial]: ... def get_trial(self, trial_id: int) -> PersistedTrial: ... + def get_cached_trial(self, trial_id: int) -> PersistedTrial: ... def get_trial_id_from_study_id_trial_number( self, study_id: int, trial_number: int ) -> int: ... @@ -563,6 +564,15 @@ class Storage: def get_trial(self, trial_id: int) -> PersistedTrial: """Get a trial by ID. + Args: + trial_id: ID of the trial. + + Returns: + The trial. + """ + def get_cached_trial(self, trial_id: int) -> PersistedTrial: + """Get a cached trial by ID without synchronizing with backends. + Args: trial_id: ID of the trial. diff --git a/rustuna_pyo3/rustuna/converter/_storage.py b/rustuna_pyo3/rustuna/converter/_storage.py index 3b87c046..6d744d9c 100644 --- a/rustuna_pyo3/rustuna/converter/_storage.py +++ b/rustuna_pyo3/rustuna/converter/_storage.py @@ -122,6 +122,9 @@ def get_trial(self, trial_id: int) -> rustuna.PersistedTrial: frozen_trial = self._storage.get_trial(trial_id) return to_persisted_trial(frozen_trial, study_id=study_id) + def get_cached_trial(self, trial_id: int) -> rustuna.PersistedTrial: + return self.get_trial(trial_id) + def delete_study(self, study_id: int) -> None: self._storage.delete_study(study_id) diff --git a/rustuna_pyo3/src/pyobject_storage.rs b/rustuna_pyo3/src/pyobject_storage.rs index f9244549..5656b478 100644 --- a/rustuna_pyo3/src/pyobject_storage.rs +++ b/rustuna_pyo3/src/pyobject_storage.rs @@ -470,6 +470,13 @@ impl Storage for PyObjectStorage { self.cache.get_trial(trial_id) } + fn get_cached_trial( + &self, + trial_id: u32, + ) -> rustuna_core::Result<&rustuna_core::trial::PersistedTrial> { + self.cache.get_cached_trial(trial_id) + } + fn get_category_labels( &mut self, study_id: u32, diff --git a/rustuna_pyo3/src/storage.rs b/rustuna_pyo3/src/storage.rs index 9899a44c..473c0a76 100644 --- a/rustuna_pyo3/src/storage.rs +++ b/rustuna_pyo3/src/storage.rs @@ -279,6 +279,17 @@ impl PyStorage { Ok(PyPersistedTrial::new(trial, study_attrs)) } + fn get_cached_trial(&self, trial_id: u32) -> PyResult { + let guard = self + .storage + .read() + .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; + let trial = guard + .get_cached_trial(trial_id) + .map_err(err_to_exceptions)?; + Ok(PyPersistedTrial::from_storage(self.storage.clone(), trial)) + } + fn get_trial_id_from_study_id_trial_number( &mut self, study_id: u32, diff --git a/rustuna_pyo3/src/trial.rs b/rustuna_pyo3/src/trial.rs index 97ac3a1e..9093ebaf 100644 --- a/rustuna_pyo3/src/trial.rs +++ b/rustuna_pyo3/src/trial.rs @@ -227,10 +227,12 @@ impl PyPersistedTrial { PyPersistedTrialSource::StorageBacked { storage, trial_id, .. } => { - let mut guard = storage - .write() + let guard = storage + .read() .map_err(|_| PyRuntimeError::new_err("Failed to acquire the storage guard"))?; - let trial = guard.get_trial(*trial_id).map_err(err_to_exceptions)?; + let trial = guard + .get_cached_trial(*trial_id) + .map_err(err_to_exceptions)?; f(trial) } } diff --git a/rustuna_storages/src/cache.rs b/rustuna_storages/src/cache.rs index 7cfdca8b..b10442c9 100644 --- a/rustuna_storages/src/cache.rs +++ b/rustuna_storages/src/cache.rs @@ -337,8 +337,17 @@ impl rustuna_core::storage::Storage for CachedStorage { } fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial> { - let (study_id, trial_number) = self.resolve_trial_location(trial_id)?; + let (study_id, _number) = self.resolve_trial_location(trial_id)?; self.refresh_trials(study_id)?; + self.get_cached_trial(trial_id) + } + + fn get_cached_trial(&self, trial_id: u32) -> Result<&PersistedTrial> { + let (study_id, trial_number) = self + .trial_id_to_study_number + .get(&trial_id) + .copied() + .ok_or_else(|| Error::new(ErrorKind::TrialNotFound))?; let trials = self .trials .get(&study_id) diff --git a/rustuna_storages/src/journal/storage.rs b/rustuna_storages/src/journal/storage.rs index 33887122..d6aa78fc 100644 --- a/rustuna_storages/src/journal/storage.rs +++ b/rustuna_storages/src/journal/storage.rs @@ -293,6 +293,10 @@ impl Storage for JournalStorage { fn get_trial(&mut self, trial_id: u32) -> Result<&PersistedTrial> { self.sync_with_backend()?; + self.get_cached_trial(trial_id) + } + + fn get_cached_trial(&self, trial_id: u32) -> Result<&PersistedTrial> { let (study_id, trial_number) = self .replay .trial_id_to_study_number