Python SDK: Add timeout_sec argument to flush (#11295)

Author: emilk

clippy.toml

@@ -52,6 +52,9 @@ disallowed-methods = [
   { path = "std::panic::catch_unwind", reason = "We compile with `panic = 'abort'`" },
   { path = "std::thread::spawn", reason = "Use `std::thread::Builder` and name the thread" },
 
+  { path = "std::time::Duration::from_secs_f32", reason = "Use try_from_secs_f32 instead to avoid panics" },
+  { path = "std::time::Duration::from_secs_f64", reason = "Use try_from_secs_f64 instead to avoid panics" },
+
   { path = "arrow::compute::concat", reason = "Use `re_arrow_util::arrow_util::concat_arrays` instead, which has better memory management" },
   { path = "arrow::compute::filter", reason = "Use `re_arrow_util::arrow_util::filter_array` instead" },
   { path = "arrow::compute::take", reason = "Use `re_arrow_util::arrow_util::take_array` instead" },

crates/store/re_chunk/src/batcher.rs

@@ -248,7 +248,13 @@ impl ChunkBatcherConfig {
                     err: Box::new(err),
                 })?;
 
-            new.flush_tick = Duration::from_secs_f64(flush_duration_secs);
+            new.flush_tick = Duration::try_from_secs_f64(flush_duration_secs).map_err(|err| {
+                ChunkBatcherError::ParseConfig {
+                    name: Self::ENV_FLUSH_TICK,
+                    value: s.clone(),
+                    err: Box::new(err),
+                }
+            })?;
         }
 
         if let Ok(s) = std::env::var(Self::ENV_FLUSH_NUM_BYTES) {

crates/top/re_sdk/src/recording_stream.rs

@@ -2538,7 +2538,12 @@ impl RecordingStream {
     /// - [`Self::reset_time`]
     #[inline]
     pub fn set_duration_secs(&self, timeline: impl Into<TimelineName>, secs: impl Into<f64>) {
-        self.set_time(timeline, std::time::Duration::from_secs_f64(secs.into()));
+        let secs = secs.into();
+        if let Ok(duration) = std::time::Duration::try_from_secs_f64(secs) {
+            self.set_time(timeline, duration);
+        } else {
+            re_log::error_once!("set_duration_secs: can't set time to {secs}");
+        }
     }
 
     /// Set a timestamp as seconds since Unix epoch (1970-01-01 00:00:00 UTC).

crates/viewer/re_ui/src/notifications.rs

@@ -341,7 +341,7 @@ impl NotificationUi {
 }
 
 fn base_ttl() -> Duration {
-    Duration::from_secs_f64(4.0)
+    Duration::from_secs(4)
 }
 
 struct Toasts {
@@ -363,7 +363,9 @@ impl Toasts {
 
     /// Shows and updates all toasts
     fn show(&self, egui_ctx: &egui::Context, notifications: &mut [Notification]) {
-        let dt = Duration::from_secs_f32(egui_ctx.input(|i| i.unstable_dt));
+        let dt = Duration::try_from_secs_f32(egui_ctx.input(|i| i.unstable_dt))
+            .unwrap_or(std::time::Duration::from_millis(100));
+
         let mut offset = egui::vec2(-8.0, 32.0);
 
         let mut first_nonzero_ttl = None;

examples/rust/objectron/src/main.rs

@@ -310,9 +310,9 @@ struct Args {
     per_frame_sleep: Option<Duration>,
 }
 
-fn parse_duration(arg: &str) -> Result<std::time::Duration, std::num::ParseFloatError> {
+fn parse_duration(arg: &str) -> anyhow::Result<std::time::Duration> {
     let seconds = arg.parse()?;
-    Ok(std::time::Duration::from_secs_f64(seconds))
+    Ok(std::time::Duration::try_from_secs_f64(seconds)?)
 }
 
 fn run(rec: &rerun::RecordingStream, args: &Args) -> anyhow::Result<()> {

rerun_py/rerun_bindings/rerun_bindings.pyi

@@ -1083,7 +1083,8 @@ def disconnect(recording: PyRecordingStream | None = None) -> None:
     Subsequent log messages will be buffered and either sent on the next call to `connect_grpc` or `spawn`.
     """
 
-def flush(blocking: bool, recording: PyRecordingStream | None = None) -> None:
+# TODO(#11294): remove `blocking` argument and put the `*` first
+def flush(blocking: bool = True, recording: PyRecordingStream | None = None, *, timeout_sec: float = 1e38) -> None:
     """Block until outstanding data has been flushed to the sink."""
 
 #

rerun_py/rerun_sdk/rerun/recording_stream.py

@@ -448,7 +448,7 @@ def __exit__(
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        self.flush(blocking=True)
+        self.flush()
 
         current_recording = active_recording_stream.get(None)
 
@@ -471,17 +471,22 @@ def __exit__(
     def to_native(self) -> bindings.PyRecordingStream:
         return self.inner
 
-    def flush(self, blocking: bool = True) -> None:
+    # TODO(#11294): remove `blocking` argument
+    def flush(self, blocking: bool = True, *, timeout_sec: float = 1e38) -> None:
         """
         Initiates a flush the batching pipeline and optionally waits for it to propagate to the underlying file descriptor (if any).
 
         Parameters
         ----------
         blocking:
             If true, the flush will block until the flush is complete.
+        timeout_sec:
+            Wait at most this many seconds.
+            If the timeout is reached, an error is raised.
+            If set to zero, the flush will be started but not waited for.
 
         """
-        bindings.flush(blocking, recording=self.to_native())
+        bindings.flush(blocking=blocking, recording=self.to_native(), timeout_sec=timeout_sec)
 
     def __del__(self) -> None:  # type: ignore[no-untyped-def]
         recording = self.to_native()
@@ -491,7 +496,7 @@ def __del__(self) -> None:  # type: ignore[no-untyped-def]
         #
         # See: https://github.com/rerun-io/rerun/issues/6223 for context on why this is necessary.
         if not recording.is_forked_child():
-            bindings.flush(blocking=False, recording=recording)  # NOLINT
+            bindings.flush(timeout_sec=0.0, recording=recording)  # NOLINT
 
     # any free function taking a `RecordingStream` as the first argument can also be a method
     binary_stream = binary_stream

rerun_py/src/python_bridge.rs

@@ -298,7 +298,13 @@ impl DurationLike {
     fn into_duration(self) -> Duration {
         match self {
             Self::Int(i) => Duration::from_secs(i as u64),
-            Self::Float(f) => Duration::from_secs_f64(f),
+            Self::Float(f) => match duration_from_sec(f) {
+                Ok(duration) => duration,
+                Err(err) => {
+                    re_log::error_once!("{err}");
+                    Duration::ZERO
+                }
+            },
             Self::Duration(d) => d,
         }
     }
@@ -1374,7 +1380,7 @@ impl PyBinarySinkStorage {
         // Release the GIL in case any flushing behavior needs to cleanup a python object.
         py.allow_threads(|| -> PyResult<_> {
             if flush {
-                let timeout = timeout_from_sec(flush_timeout_sec)?;
+                let timeout = duration_from_sec(flush_timeout_sec as _)?;
                 self.inner
                     .flush(timeout)
                     .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
@@ -1405,7 +1411,7 @@ impl PyBinarySinkStorage {
     fn flush(&self, py: Python<'_>, timeout_sec: f32) -> PyResult<()> {
         // Release the GIL in case any flushing behavior needs to cleanup a python object.
         py.allow_threads(|| -> PyResult<_> {
-            let timeout = timeout_from_sec(timeout_sec)?;
+            let timeout = duration_from_sec(timeout_sec as _)?;
             self.inner
                 .flush(timeout)
                 .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
@@ -1415,13 +1421,13 @@ impl PyBinarySinkStorage {
     }
 }
 
-fn timeout_from_sec(seconds: f32) -> PyResult<Duration> {
+fn duration_from_sec(seconds: f64) -> PyResult<Duration> {
     if seconds.is_nan() {
-        Err(PyRuntimeError::new_err("timeout_sec must not be NaN"))
+        Err(PyRuntimeError::new_err("duration must not be NaN"))
     } else if seconds < 0.0 {
-        Err(PyRuntimeError::new_err("timeout_sec must be non-negative"))
+        Err(PyRuntimeError::new_err("duration must be non-negative"))
     } else {
-        Ok(Duration::try_from_secs_f32(seconds).unwrap_or(Duration::MAX))
+        Ok(Duration::try_from_secs_f64(seconds).unwrap_or(Duration::MAX))
     }
 }
 
@@ -1607,24 +1613,33 @@ fn disconnect(py: Python<'_>, recording: Option<&PyRecordingStream>) {
 }
 
 /// Block until outstanding data has been flushed to the sink.
+// TODO(#11294): remove `blocking` argument and put the `*` first
 #[pyfunction]
-#[pyo3(signature = (blocking, recording=None))]
-fn flush(py: Python<'_>, blocking: bool, recording: Option<&PyRecordingStream>) -> PyResult<()> {
+#[pyo3(signature = (blocking = true, recording = None, *, timeout_sec = 1e38))] // Can't use infinity here because of python_check_signatures.py
+fn flush(
+    py: Python<'_>,
+    blocking: bool,
+    recording: Option<&PyRecordingStream>,
+    timeout_sec: f32,
+) -> PyResult<()> {
     let Some(recording) = get_data_recording(recording) else {
         return Ok(());
     };
 
     // Release the GIL in case any flushing behavior needs to cleanup a python object.
-    py.allow_threads(|| -> Result<(), SinkFlushError> {
-        if blocking {
-            recording.flush_blocking()?;
+    py.allow_threads(|| -> PyResult<()> {
+        if !blocking || timeout_sec == 0.0 {
+            recording
+                .flush_async()
+                .map_err(|err: SinkFlushError| PyRuntimeError::new_err(err.to_string()))?;
         } else {
-            recording.flush_async()?;
+            recording
+                .flush_with_timeout(duration_from_sec(timeout_sec as _)?)
+                .map_err(|err: SinkFlushError| PyRuntimeError::new_err(err.to_string()))?;
         }
         flush_garbage_queue();
         Ok(())
     })
-    .map_err(|err: SinkFlushError| PyRuntimeError::new_err(err.to_string()))
 }
 
 // --- Components ---

rerun_py/tests/unit/test_batcher_config.py

@@ -104,7 +104,7 @@ def test_flush_never() -> None:
         assert sz2 == sz1, "Expected the file size to stay the same"
         assert sz3 == sz2, "Expected the file size to stay the same"
 
-        rec.flush(blocking=True)
+        rec.flush()
 
         sz4 = os.stat(rec_path).st_size
 

tests/rust/plot_dashboard_stress/src/main.rs

@@ -12,6 +12,8 @@
 //! pixi run rs-plot-dashboard --num-plots 10 --num-series-per-plot 5 --num-points-per-series 5000 --freq 1000
 //! ```
 
+#![expect(clippy::disallowed_methods)]
+
 use rerun::external::re_log;
 
 #[derive(Debug, clap::ValueEnum, Clone)]