Add Dataset.register_batch and wrappers for task ids (#9895)

### Related

* closes rerun-io/dataplatform#689

### What

- Introduces `Dataset.register_batch()` to register multiple RRDs at
once.
- Introduces `Task` and `Tasks`, which are wrapper over one,
respectively a bunch of, task ids. For now, they just have a
`.wait(timeout_secs)` method.
- Add an optional `timeout_secs` argument to `Dataset.register()`

Author: abey79

rerun_py/rerun_bindings/rerun_bindings.pyi

@@ -1038,8 +1038,36 @@ class Dataset(Entry):
     def partition_url(self, partition_id: str) -> str:
         """Return the URL for the given partition."""
 
-    def register(self, recording_uri: str) -> None:
-        """Register a RRD URI to the dataset."""
+    def register(self, recording_uri: str, timeout_secs: int = 60) -> None:
+        """
+        Register a RRD URI to the dataset and wait for completion.
+
+        This method registers a single recording to the dataset and blocks until the registration is
+        complete, or after a timeout (in which case, a `TimeoutError` is raised).
+
+        Parameters
+        ----------
+        recording_uri: str
+            The URI of the RRD to register
+
+        timeout_secs: int
+            The timeout after which this method returns.
+
+        """
+
+    def register_batch(self, recording_uris: list[str]) -> Tasks:
+        """
+        Register a batch of RRD URIs to the dataset and return a handle to the tasks.
+
+        This method initiates the registration of multiple recordings to the dataset, and returns
+        the corresponding task ids in a [`Tasks`] object.
+
+        Parameters
+        ----------
+        recording_uris: list[str]
+            The URIs of the RRDs to register
+
+        """
 
     def download_partition(self, partition_id: str) -> Recording:
         """Download a partition from the dataset."""
@@ -1305,6 +1333,36 @@ class DataFusionTable:
     def name(self) -> str:
         """Name of this table."""
 
+class Task:
+    """A handle on a remote task."""
+
+    @property
+    def id(self) -> str:
+        """The task id."""
+
+    def wait(self, timeout_secs: int) -> None:
+        """
+        Block until the task is completed or the timeout is reached.
+
+        A `TimeoutError` is raised if the timeout is reached.
+        """
+
+class Tasks:
+    """A collection of [`Task`]."""
+
+    def wait(self, timeout_secs: int) -> None:
+        """
+        Block until all tasks are completed or the timeout is reached.
+
+        A `TimeoutError` is raised if the timeout is reached.
+        """
+
+    def __len__(self) -> int:
+        """Return the number of tasks."""
+
+    def __getitem__(self, index: int) -> Task:
+        """Return the task at the given index."""
+
 #####################################################################################################################
 ## SEND_TABLE                                                                                                      ##
 #####################################################################################################################

rerun_py/rerun_sdk/rerun/catalog.py

@@ -8,5 +8,6 @@
     EntryId as EntryId,
     EntryKind as EntryKind,
     Table as Table,
+    Task as Task,
     VectorDistanceMetric as VectorDistanceMetric,
 )

rerun_py/src/catalog/connection_handle.rs

@@ -219,6 +219,7 @@ impl ConnectionHandle {
         })
     }
 
+    /// Wait for the provided tasks to finish.
     pub fn wait_for_tasks(
         &mut self,
         py: Python<'_>,
@@ -251,7 +252,7 @@ impl ConnectionHandle {
                     .decode()
                     .map_err(to_py_err)?
             } else {
-                return Err(PyValueError::new_err("no response from task"));
+                return Err(PyConnectionError::new_err("no response from task"));
             };
 
             // TODO(andrea): this is a bit hideous. Maybe the idea of returning a dataframe rather

rerun_py/src/catalog/dataset.rs

@@ -4,11 +4,11 @@ use arrow::array::{RecordBatch, StringArray};
 use arrow::datatypes::{Field, Schema as ArrowSchema};
 use arrow::pyarrow::PyArrowType;
 use pyo3::{exceptions::PyRuntimeError, pyclass, pymethods, Py, PyAny, PyRef, PyResult, Python};
-use re_grpc_client::redap::get_chunks_response_to_chunk_and_partition_id;
 use tokio_stream::StreamExt as _;
 
 use re_chunk_store::{ChunkStore, ChunkStoreHandle};
 use re_datafusion::{PartitionTableProvider, SearchResultsTableProvider};
+use re_grpc_client::redap::get_chunks_response_to_chunk_and_partition_id;
 use re_log_encoding::codec::wire::encoder::Encode as _;
 use re_log_types::{StoreId, StoreInfo, StoreKind, StoreSource};
 use re_protos::common::v1alpha1::ext::DatasetHandle;
@@ -20,6 +20,7 @@ use re_protos::manifest_registry::v1alpha1::{
 };
 use re_sorbet::{SorbetColumnDescriptors, TimeColumnSelector};
 
+use crate::catalog::task::PyTasks;
 use crate::catalog::{
     dataframe_query::PyDataframeQueryView, to_py_err, PyEntry, VectorDistanceMetricLike, VectorLike,
 };
@@ -94,24 +95,52 @@ impl PyDataset {
         .to_string()
     }
 
-    /// Register a RRD URI to the dataset.
-    fn register(self_: PyRef<'_, Self>, recording_uri: String) -> PyResult<()> {
-        // TODO(#9731): In order to make the `register` method appear synchronous,
-        // we need to hard-code a max timeout for waiting for the corresponding tasks.
-        // 60 seconds is totally arbitrary but should work for interactive uses
-        // for now.
-        //
-        // A more permanent solution is to expose an asynchronous register method, and/or
-        // the timeout directly to the caller.
-        // See also issue #9731
-        const MAX_REGISTER_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
+    /// Register a RRD URI to the dataset and wait for completion.
+    ///
+    /// This method registers a single recording to the dataset and blocks until the registration is
+    /// complete, or after a timeout (in which case, a `TimeoutError` is raised).
+    ///
+    /// Parameters
+    /// ----------
+    /// recording_uri: str
+    ///     The URI of the RRD to register
+    ///
+    /// timeout_secs: int
+    ///     The timeout after which this method returns.
+    #[pyo3(signature = (recording_uri, timeout_secs = 60))]
+    fn register(self_: PyRef<'_, Self>, recording_uri: String, timeout_secs: u64) -> PyResult<()> {
+        let register_timeout = std::time::Duration::from_secs(timeout_secs);
         let super_ = self_.as_super();
         let mut connection = super_.client.borrow(self_.py()).connection().clone();
         let dataset_id = super_.details.id;
 
         let task_ids =
             connection.register_with_dataset(self_.py(), dataset_id, vec![recording_uri])?;
-        connection.wait_for_tasks(self_.py(), &task_ids, MAX_REGISTER_TIMEOUT)
+
+        connection.wait_for_tasks(self_.py(), &task_ids, register_timeout)
+    }
+
+    /// Register a batch of RRD URIs to the dataset and return a handle to the tasks.
+    ///
+    /// This method initiates the registration of multiple recordings to the dataset, and returns
+    /// the corresponding task ids in a [`Tasks`] object.
+    ///
+    /// Parameters
+    /// ----------
+    /// recording_uris: list[str]
+    ///     The URIs of the RRDs to register
+    #[allow(rustdoc::broken_intra_doc_links)]
+    fn register_batch(self_: PyRef<'_, Self>, recording_uris: Vec<String>) -> PyResult<PyTasks> {
+        let super_ = self_.as_super();
+        let mut connection = super_.client.borrow(self_.py()).connection().clone();
+        let dataset_id = super_.details.id;
+
+        let task_ids = connection.register_with_dataset(self_.py(), dataset_id, recording_uris)?;
+
+        Ok(PyTasks {
+            client: super_.client.clone_ref(self_.py()),
+            ids: task_ids,
+        })
     }
 
     /// Download a partition from the dataset.

rerun_py/src/catalog/errors.rs

@@ -13,12 +13,14 @@
 //! - Error type (either built-in such as [`pyo3::exceptions::PyValueError`] or custom) can always
 //!   be used directly using, e.g. `PyValueError::new_err("message")`.
 
-use pyo3::exceptions::{PyConnectionError, PyValueError};
-use pyo3::PyErr;
 use std::error::Error as _;
 
+use pyo3::exceptions::{PyConnectionError, PyTimeoutError, PyValueError};
+use pyo3::PyErr;
+
 use re_grpc_client::redap::ConnectionError;
 use re_protos::manifest_registry::v1alpha1::ext::GetDatasetSchemaResponseError;
+
 // ---
 
 /// Private error type to server as a bridge between various external error type and the
@@ -85,17 +87,22 @@ impl From<ExternalError> for PyErr {
             ExternalError::ConnectionError(err) => PyConnectionError::new_err(err.to_string()),
 
             ExternalError::TonicStatusError(status) => {
-                let mut msg = format!(
-                    "tonic status error: {} (code: {}",
-                    status.message(),
-                    status.code()
-                );
-                if let Some(source) = status.source() {
-                    msg.push_str(&format!(", source: {source})"));
+                if status.code() == tonic::Code::DeadlineExceeded {
+                    PyTimeoutError::new_err("Deadline expired before operation could complete")
                 } else {
-                    msg.push(')');
+                    let mut msg = format!(
+                        "tonic status error: {} (code: {}",
+                        status.message(),
+                        status.code()
+                    );
+                    if let Some(source) = status.source() {
+                        msg.push_str(&format!(", source: {source})"));
+                    } else {
+                        msg.push(')');
+                    }
+
+                    PyConnectionError::new_err(msg)
                 }
-                PyConnectionError::new_err(msg)
             }
 
             ExternalError::UriError(err) => PyValueError::new_err(format!("Invalid URI: {err}")),

rerun_py/src/catalog/mod.rs

@@ -7,6 +7,7 @@ mod dataset;
 mod entry;
 mod errors;
 mod table;
+mod task;
 
 use std::sync::Arc;
 
@@ -18,12 +19,13 @@ use pyo3::{exceptions::PyRuntimeError, prelude::*, Bound, PyResult};
 
 use crate::catalog::dataframe_query::PyDataframeQueryView;
 
-pub use catalog_client::PyCatalogClient;
-pub use connection_handle::ConnectionHandle;
-pub use dataset::PyDataset;
-pub use entry::{PyEntry, PyEntryId, PyEntryKind};
-pub use errors::to_py_err;
-pub use table::PyTable;
+pub(crate) use catalog_client::PyCatalogClient;
+pub(crate) use connection_handle::ConnectionHandle;
+pub(crate) use dataset::PyDataset;
+pub(crate) use entry::{PyEntry, PyEntryId, PyEntryKind};
+pub(crate) use errors::to_py_err;
+pub(crate) use table::PyTable;
+pub(crate) use task::{PyTask, PyTasks};
 
 /// Register the `rerun.catalog` module.
 pub(crate) fn register(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -34,6 +36,8 @@ pub(crate) fn register(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     m.add_class::<PyEntry>()?;
     m.add_class::<PyDataset>()?;
     m.add_class::<PyTable>()?;
+    m.add_class::<PyTask>()?;
+    m.add_class::<PyTasks>()?;
 
     m.add_class::<PyDataframeQueryView>()?;
 

rerun_py/src/catalog/task.rs

@@ -0,0 +1,87 @@
+use pyo3::exceptions::PyIndexError;
+use pyo3::{pyclass, pymethods, Py, PyRef, PyResult, Python};
+
+use re_protos::common::v1alpha1::TaskId;
+
+use crate::catalog::PyCatalogClient;
+
+/// A handle on a remote task.
+#[pyclass(name = "Task")]
+pub struct PyTask {
+    pub client: Py<PyCatalogClient>,
+
+    pub id: TaskId,
+}
+
+/// A handle on a remote task.
+#[pymethods]
+impl PyTask {
+    /// Entry id as a string.
+    pub fn __repr__(&self) -> String {
+        format!("Task({})", self.id.id)
+    }
+
+    /// The task id.
+    #[getter]
+    pub fn id(&self) -> String {
+        self.id.id.clone()
+    }
+
+    /// Block until the task is completed or the timeout is reached.
+    ///
+    /// A `TimeoutError` is raised if the timeout is reached.
+    pub fn wait(&self, py: Python<'_>, timeout_secs: u64) -> PyResult<()> {
+        let mut connection = self.client.borrow(py).connection().clone();
+        let timeout = std::time::Duration::from_secs(timeout_secs);
+        connection.wait_for_tasks(py, &[self.id.clone()], timeout)?;
+
+        Ok(())
+    }
+
+    //TODO(ab): add method to poll about status
+}
+
+/// A collection of [`Task`].
+#[allow(rustdoc::broken_intra_doc_links)]
+#[pyclass(name = "Tasks")]
+pub struct PyTasks {
+    pub client: Py<PyCatalogClient>,
+
+    pub ids: Vec<TaskId>,
+}
+
+#[pymethods]
+impl PyTasks {
+    /// Block until all tasks are completed or the timeout is reached.
+    ///
+    /// A `TimeoutError` is raised if the timeout is reached.
+    pub fn wait(self_: PyRef<'_, Self>, timeout_secs: u64) -> PyResult<()> {
+        let mut connection = self_.client.borrow(self_.py()).connection().clone();
+        let timeout = std::time::Duration::from_secs(timeout_secs);
+        connection.wait_for_tasks(self_.py(), &self_.ids, timeout)?;
+
+        Ok(())
+    }
+
+    //TODO(ab): add method to poll about status (how many are done, etc.)
+
+    //
+    // Sequence methods
+    //
+
+    fn __len__(&self) -> usize {
+        self.ids.len()
+    }
+
+    /// Get the task at the given index.
+    fn __getitem__(&self, py: Python<'_>, index: usize) -> PyResult<PyTask> {
+        if index >= self.ids.len() {
+            return Err(PyIndexError::new_err("Index out of range"));
+        }
+
+        Ok(PyTask {
+            client: self.client.clone_ref(py),
+            id: self.ids[index].clone(),
+        })
+    }
+}

scripts/ci/python_check_signatures.py

@@ -62,10 +62,13 @@ def __eq__(self, other: Any) -> bool:
         if not isinstance(other, APIDef):
             return NotImplemented
 
-        if self.name in ("__init__", "__iter__"):
-            # Ignore the signature of __init__ and __new__ methods
+        if self.name in ("__init__", "__iter__", "__len__"):
+            # pyo3 has a special way to handle these methods that makes it impossible to match everything.
             # TODO(#7779): Remove this special case once we have a better way to handle these methods
             return self.name == other.name and self.signature == other.signature
+        elif self.name in ("__getitem__"):
+            # TODO(#7779): It's somehow even worse for these.
+            return self.name == other.name
         else:
             return self.name == other.name and self.signature == other.signature and self.doc == other.doc