Cleanup internal data-structures when process has been forked

examples/python/multiprocessing/main.py

@@ -10,6 +10,10 @@
 import rerun as rr  # pip install rerun-sdk
 
 
+# Python does not guarantee that the normal atexit-handlers will be called at the
+# termination of a multiprocessing.Process. Explicitly add the `shutdown_at_exit`
+# decorator to ensure data is flushed when the task completes.
+@rr.shutdown_at_exit
 def task(child_index: int) -> None:
     # All processes spawned with `multiprocessing` will automatically
     # be assigned the same default recording_id.
@@ -37,11 +41,6 @@ def main() -> None:
 
     task(0)
 
-    # Using multiprocessing with "fork" results in a hang on shutdown so
-    # always use "spawn"
-    # TODO(https://github.com/rerun-io/rerun/issues/1921)
-    multiprocessing.set_start_method("spawn")
-
     for i in [1, 2, 3]:
         p = multiprocessing.Process(target=task, args=(i,))
         p.start()

rerun_py/rerun_sdk/rerun/__init__.py

@@ -1,6 +1,9 @@
 """The Rerun Python SDK, which is a wrapper around the re_sdk crate."""
 from __future__ import annotations
 
+import functools
+from typing import Any, Callable, TypeVar, cast
+
 # NOTE: The imports determine what is public API. Avoid importing globally anything that is not public API. Use
 # (private) function and local import if needed.
 import rerun_bindings as bindings  # type: ignore[attr-defined]
@@ -331,6 +334,31 @@ def _register_on_fork() -> None:
 
 _register_on_fork()
 
+
+_TFunc = TypeVar("_TFunc", bound=Callable[..., Any])
+
+
+def shutdown_at_exit(func: _TFunc) -> _TFunc:
+    """
+    Decorator to shutdown Rerun cleanly when this function exits.
+
+    Normally, Rerun installs an atexit-handler that attempts to shutdown cleanly and
+    flush all outgoing data before terminating. However, some cases, such as forked
+    processes will always skip this at-exit handler. In these cases, you can use this
+    decorator on the entry-point to your subprocess to ensure cleanup happens as
+    expected without losing data.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            return func(*args, **kwargs)
+        finally:
+            rerun_shutdown()
+
+    return cast(_TFunc, wrapper)
+
+
 # ---