Expose the trial queue through sampler context

rustuna_core/src/sampler.rs

@@ -7,19 +7,33 @@ use crate::distribution::Distribution;
 use crate::storage::Storage;
 use crate::study::Direction;
 use crate::trial::TrialStateValues;
+use crate::trial_queue::TrialQueue;
 use crate::Result;
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 /// Lightweight study and trial metadata passed to samplers.
 ///
 /// Rustuna does not pass full `Study` or `FrozenTrial` objects to samplers. Instead, it provides
-/// the identifiers and objective directions required to make sampling decisions, while richer
-/// state can be retrieved through [`Storage`] when necessary.
+/// the identifiers and objective directions required to make sampling decisions, the shared
+/// [`TrialQueue`] for queued trials, and access to richer state through [`Storage`] when necessary.
 pub struct Context {
     pub study_id: u32,
     pub directions: Vec<Direction>,
     pub trial_number: u32,
     pub trial_id: u32,
+    pub trial_queue: Arc<dyn TrialQueue>,
+}
+
+impl std::fmt::Debug for Context {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Context")
+            .field("study_id", &self.study_id)
+            .field("directions", &self.directions)
+            .field("trial_number", &self.trial_number)
+            .field("trial_id", &self.trial_id)
+            .field("trial_queue", &"<TrialQueue>")
+            .finish()
+    }
 }
 
 /// Interface for parameter suggestion algorithms.
@@ -246,6 +260,7 @@ mod tests {
     use crate::storage::InMemoryStorage;
     use crate::study::create_study;
     use crate::trial::Trial;
+    use crate::trial_queue::InMemoryTrialQueue;
 
     pub struct DummyJointSampler {
         joint_params: HashMap<String, f64>,
@@ -294,12 +309,14 @@ mod tests {
     fn random_sampler_can_be_shared_between_threads() -> Result<()> {
         let sampler: Arc<dyn Sampler> = Arc::new(RandomSampler::seed_from_u64(0));
         let storage: Arc<RwLock<dyn Storage>> = Arc::new(RwLock::new(InMemoryStorage::new()));
+        let trial_queue: Arc<dyn TrialQueue> = Arc::new(InMemoryTrialQueue::new());
         let distribution = Distribution::new_float(-1.0, 1.0, None, false);
         let mut handles = Vec::new();
 
         for thread_id in 0..4 {
             let sampler = Arc::clone(&sampler);
             let storage = Arc::clone(&storage);
+            let trial_queue = Arc::clone(&trial_queue);
             let distribution = distribution.clone();
             handles.push(thread::spawn(move || -> Result<()> {
                 for trial_number in 0..100 {
@@ -308,6 +325,7 @@ mod tests {
                         directions: vec![Direction::Minimize],
                         trial_number,
                         trial_id: thread_id * 100 + trial_number,
+                        trial_queue: Arc::clone(&trial_queue),
                     };
                     let value = sampler.sample_independent(
                         &ctx,

rustuna_core/src/study.rs

@@ -182,6 +182,7 @@ impl Study {
             trial_number: persisted_trial.number,
             trial_id: persisted_trial.id,
             directions: self.directions.clone(),
+            trial_queue: Arc::clone(&self.queue),
         };
         self.sampler.before_trial(&ctx, self.storage.clone())?;
 
@@ -219,6 +220,7 @@ impl Study {
             self.directions.clone(),
             Arc::clone(&self.storage),
             Arc::clone(&self.sampler),
+            Arc::clone(&self.queue),
             joint_params,
         );
         Ok(trial)
@@ -253,6 +255,7 @@ impl Study {
             directions: self.directions.clone(),
             trial_number,
             trial_id,
+            trial_queue: Arc::clone(&self.queue),
         };
         let after_trial_result =
             self.sampler

rustuna_core/src/trial.rs

@@ -9,6 +9,7 @@ use crate::distribution::Distribution;
 use crate::sampler::{Context as SamplerContext, Sampler};
 use crate::storage::Storage;
 use crate::study::Direction;
+use crate::trial_queue::TrialQueue;
 use crate::{Error, ErrorKind, Result};
 
 const CONSTRAINTS_KEY: &str = "constraints";
@@ -27,6 +28,7 @@ pub struct Trial {
     directions: Vec<Direction>,
     storage: Arc<RwLock<dyn Storage>>,
     sampler: Arc<dyn Sampler>,
+    trial_queue: Arc<dyn TrialQueue>,
     joint_params: HashMap<String, (Distribution, f64)>,
     fixed_params: HashMap<String, CategoryLabel>,
     cached_user_attrs: HashMap<String, String>,
@@ -38,6 +40,7 @@ impl Trial {
         directions: Vec<Direction>,
         storage: Arc<RwLock<dyn Storage>>,
         sampler: Arc<dyn Sampler>,
+        trial_queue: Arc<dyn TrialQueue>,
         joint_params: HashMap<String, (Distribution, f64)>,
     ) -> Self {
         let fixed_params = extract_fixed_params(&trial.attrs);
@@ -58,6 +61,7 @@ impl Trial {
             directions,
             storage,
             sampler,
+            trial_queue,
             joint_params,
             fixed_params,
             cached_user_attrs,
@@ -135,6 +139,7 @@ impl Trial {
             trial_number: self.number,
             trial_id: self.id,
             directions: self.directions.clone(),
+            trial_queue: Arc::clone(&self.trial_queue),
         };
         let param_value =
             self.sampler

rustuna_pyo3/rustuna/_rustuna.pyi

@@ -1165,13 +1165,16 @@ class SamplerContext:
     Attributes:
         study_id: Study ID.
         trial_number: Trial number.
+        trial_id: Trial ID.
         directions: Optimization directions.
+        trial_queue: The queue shared with the study.
     """
 
     study_id: int
     trial_number: int
     trial_id: int
     directions: list[StudyDirection]
+    trial_queue: TrialQueueProtocol
 
     def __init__(
         self,

rustuna_pyo3/src/sampler/context.rs

@@ -1,8 +1,12 @@
+use std::sync::Arc;
+
 use pyo3::prelude::*;
 
 use rustuna_core::sampler::Context as SamplerContext;
+use rustuna_core::trial_queue::InMemoryTrialQueue;
 
 use crate::study::PyDirection;
+use crate::trial_queue::to_python::ToPythonTrialQueue;
 
 #[derive(Clone)]
 #[pyclass(name = "SamplerContext", skip_from_py_object)]
@@ -32,6 +36,7 @@ impl PySamplerContext {
                 trial_number,
                 trial_id,
                 directions: directions.into_iter().map(|d| d.into()).collect(),
+                trial_queue: Arc::new(InMemoryTrialQueue::new()),
             },
         })
     }
@@ -55,4 +60,10 @@ impl PySamplerContext {
             .map(|d| d.clone().into())
             .collect()
     }
+
+    #[getter]
+    fn trial_queue(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
+        let queue = ToPythonTrialQueue::new(Arc::clone(&self.context.trial_queue));
+        Ok(Py::new(py, queue)?.into_any())
+    }
 }

rustuna_pyo3/src/storage/mod.rs

@@ -3,5 +3,5 @@ pub mod cached;
 pub mod in_memory;
 pub mod journal;
 pub mod sqlite3;
-pub mod to_python;
+pub(crate) mod to_python;
 pub mod to_rust;

rustuna_pyo3/src/storage/to_python.rs

@@ -12,18 +12,18 @@ use crate::trial::{PyPersistedTrial, PyTrialState};
 #[derive(Clone)]
 #[pyclass(name = "ToPythonStorage", skip_from_py_object)]
 #[pyo3(module = "rustuna")]
-pub struct ToPythonStorage {
+pub(crate) struct ToPythonStorage {
     pub(crate) binding: StorageBinding,
 }
 
 impl ToPythonStorage {
-    pub fn new(storage: Arc<RwLock<dyn Storage>>) -> Self {
+    pub(crate) fn new(storage: Arc<RwLock<dyn Storage>>) -> Self {
         Self {
             binding: StorageBinding { storage },
         }
     }
 
-    pub fn storage(&self) -> Arc<RwLock<dyn Storage>> {
+    pub(crate) fn storage(&self) -> Arc<RwLock<dyn Storage>> {
         self.binding.storage.clone()
     }
 }

rustuna_pyo3/src/trial_queue/mod.rs

@@ -1,4 +1,5 @@
 pub mod directory;
 pub mod inmemory;
 pub mod sqlite3;
+pub(crate) mod to_python;
 pub mod to_rust;

rustuna_pyo3/src/trial_queue/to_python.rs

@@ -0,0 +1,35 @@
+use std::sync::Arc;
+
+use pyo3::prelude::*;
+use rustuna_core::trial_queue::TrialQueue;
+use rustuna_core::ErrorKind;
+
+use crate::exception::err_to_exceptions;
+
+/// Python-facing adapter for the trial queue carried by a sampler context.
+#[derive(Clone)]
+#[pyclass(name = "ToPythonTrialQueue", module = "rustuna", skip_from_py_object)]
+pub(crate) struct ToPythonTrialQueue {
+    queue: Arc<dyn TrialQueue>,
+}
+
+impl ToPythonTrialQueue {
+    pub(crate) fn new(queue: Arc<dyn TrialQueue>) -> Self {
+        Self { queue }
+    }
+}
+
+#[pymethods]
+impl ToPythonTrialQueue {
+    fn enqueue(&self, py: Python<'_>, trial_id: u32) -> PyResult<()> {
+        py.detach(|| self.queue.enqueue(trial_id).map_err(err_to_exceptions))
+    }
+
+    fn dequeue(&self, py: Python<'_>) -> PyResult<Option<u32>> {
+        py.detach(|| match self.queue.dequeue() {
+            Ok(trial_id) => Ok(Some(trial_id)),
+            Err(error) if matches!(error.kind, ErrorKind::TrialQueueEmpty) => Ok(None),
+            Err(error) => Err(err_to_exceptions(error)),
+        })
+    }
+}

rustuna_sampler/src/tpe/sampler.rs

@@ -877,6 +877,10 @@ mod tests {
     use rustuna_core::study::{create_study, Direction};
     use rustuna_core::study::{get_best_trial, get_pareto_front};
 
+    fn test_trial_queue() -> Arc<dyn rustuna_core::trial_queue::TrialQueue> {
+        Arc::new(rustuna_core::trial_queue::InMemoryTrialQueue::new())
+    }
+
     #[test]
     fn test_optimize() {
         let storage = InMemoryStorage::new();
@@ -1376,6 +1380,7 @@ mod tests {
             directions: vec![Direction::Minimize],
             trial_number: 15,
             trial_id: 999,
+            trial_queue: test_trial_queue(),
         };
         let distribution = Distribution::new_float(0.0, 10.0, None, false);
         let cached = TpeSampler::seed_from_u64(7);
@@ -1424,6 +1429,7 @@ mod tests {
             directions: vec![Direction::Minimize],
             trial_number: trial_id,
             trial_id,
+            trial_queue: test_trial_queue(),
         };
         let key = |trial_id| {
             (
@@ -1683,6 +1689,7 @@ mod tests {
             trial_id: 100,
             trial_number: 1,
             directions: vec![Direction::Minimize],
+            trial_queue: test_trial_queue(),
         };
         sampler.before_trial(&ctx, first.clone()).unwrap();
         sampler.before_trial(&ctx, second.clone()).unwrap();
@@ -1727,6 +1734,7 @@ mod tests {
                 trial_id: u32::MAX,
                 trial_number: 2,
                 directions: vec![Direction::Minimize],
+                trial_queue: test_trial_queue(),
             };
             if prime_cache {
                 sampler.snapshot_from_storage(&ctx, &storage).unwrap();
@@ -1762,6 +1770,7 @@ mod tests {
                 trial_id: study_id,
                 trial_number: 0,
                 directions: vec![Direction::Minimize],
+                trial_queue: test_trial_queue(),
             };
             sampler.before_trial(&ctx, storage.clone()).unwrap();
         }