Adding `recv_async` dora method to retrieve data in python async

apis/python/node/Cargo.toml

@@ -10,9 +10,10 @@ repository.workspace = true
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-default = ["tracing", "metrics", "telemetry"]
+default = ["tracing", "metrics", "telemetry", "async"]
 tracing = ["dora-node-api/tracing"]
 metrics = ["dora-node-api/metrics"]
+async = ["pyo3/experimental-async"]
 telemetry = ["dora-runtime/telemetry"]
 
 [dependencies]

apis/python/node/pyproject.toml

@@ -12,13 +12,13 @@ readme = "README.md"
 dependencies = ['pyarrow']
 
 [dependency-groups]
-dev = ["pytest >=8.1.1", "ruff >=0.9.1"]
+dev = ["pytest >=7.1.1", "ruff >=0.9.1"]
 
 [tool.maturin]
 features = ["pyo3/extension-module"]
 
 [tool.ruff.lint]
 extend-select = [
-  "D",   # pydocstyle
-  "UP"
+    "D",  # pydocstyle
+    "UP",
 ]

apis/python/node/src/lib.rs

@@ -57,14 +57,14 @@ impl Node {
         let dataflow_id = *node.dataflow_id();
         let node_id = node.id().clone();
         let node = DelayedCleanup::new(node);
-        let events = DelayedCleanup::new(events);
+        let events = events;
         let cleanup_handle = NodeCleanupHandle {
-            _handles: Arc::new((node.handle(), events.handle())),
+            _handles: Arc::new(node.handle()),
         };
         Ok(Node {
             events: Events {
                 inner: EventsInner::Dora(events),
-                cleanup_handle,
+                _cleanup_handle: cleanup_handle,
             },
             dataflow_id,
             node_id,
@@ -107,6 +107,43 @@ impl Node {
         }
     }
 
+    /// `.recv_async()` gives you the next input that the node has received asynchronously.
+    /// It does not blocks until the next event becomes available.
+    /// You can use timeout in seconds to return if no input is available.
+    /// It will return an Error if the timeout is reached.
+    /// It will return `None` when all senders has been dropped.
+    ///
+    /// warning::
+    ///     This feature is experimental as pyo3 async (rust-python FFI) is still in development.
+    ///
+    /// ```python
+    /// event = await node.recv_async()
+    /// ```
+    ///
+    /// You can also iterate over the event stream with a loop
+    ///
+    /// :type timeout: float, optional
+    /// :rtype: dict
+    #[pyo3(signature = (timeout=None))]
+    #[allow(clippy::should_implement_trait)]
+    pub async fn recv_async(&mut self, timeout: Option<f32>) -> PyResult<Option<Py<PyDict>>> {
+        let event = self
+            .events
+            .recv_async_timeout(timeout.map(Duration::from_secs_f32))
+            .await;
+        if let Some(event) = event {
+            // Get python
+            Python::with_gil(|py| {
+                let dict = event
+                    .to_py_dict(py)
+                    .context("Could not convert event into a dict")?;
+                Ok(Some(dict))
+            })
+        } else {
+            Ok(None)
+        }
+    }
+
     /// You can iterate over the event stream with a loop
     ///
     /// ```python
@@ -254,30 +291,38 @@ fn err_to_pyany(err: eyre::Report, gil: Python<'_>) -> Py<PyAny> {
 
 struct Events {
     inner: EventsInner,
-    cleanup_handle: NodeCleanupHandle,
+    _cleanup_handle: NodeCleanupHandle,
 }
 
 impl Events {
     fn recv(&mut self, timeout: Option<Duration>) -> Option<PyEvent> {
+        let event = match &mut self.inner {
+            EventsInner::Dora(events) => match timeout {
+                Some(timeout) => events.recv_timeout(timeout).map(MergedEvent::Dora),
+                None => events.recv().map(MergedEvent::Dora),
+            },
+            EventsInner::Merged(events) => futures::executor::block_on(events.next()),
+        };
+        event.map(|event| PyEvent { event })
+    }
+
+    async fn recv_async_timeout(&mut self, timeout: Option<Duration>) -> Option<PyEvent> {
         let event = match &mut self.inner {
             EventsInner::Dora(events) => match timeout {
                 Some(timeout) => events
-                    .get_mut()
-                    .recv_timeout(timeout)
+                    .recv_async_timeout(timeout)
+                    .await
                     .map(MergedEvent::Dora),
-                None => events.get_mut().recv().map(MergedEvent::Dora),
+                None => events.recv_async().await.map(MergedEvent::Dora),
             },
-            EventsInner::Merged(events) => futures::executor::block_on(events.next()),
+            EventsInner::Merged(events) => events.next().await,
         };
-        event.map(|event| PyEvent {
-            event,
-            _cleanup: Some(self.cleanup_handle.clone()),
-        })
+        event.map(|event| PyEvent { event })
     }
 }
 
 enum EventsInner {
-    Dora(DelayedCleanup<EventStream>),
+    Dora(EventStream),
     Merged(Box<dyn Stream<Item = MergedEvent<PyObject>> + Unpin + Send + Sync>),
 }
 

apis/python/operator/src/lib.rs

@@ -19,14 +19,13 @@ use pyo3::{
 /// Dora Event
 pub struct PyEvent {
     pub event: MergedEvent<PyObject>,
-    pub _cleanup: Option<NodeCleanupHandle>,
 }
 
 /// Keeps the dora node alive until all event objects have been dropped.
 #[derive(Clone)]
 #[pyclass]
 pub struct NodeCleanupHandle {
-    pub _handles: Arc<(CleanupHandle<DoraNode>, CleanupHandle<EventStream>)>,
+    pub _handles: Arc<CleanupHandle<DoraNode>>,
 }
 
 /// Owned type with delayed cleanup (using `handle` method).
@@ -139,17 +138,6 @@ impl PyEvent {
             }
         }
 
-        if let Some(cleanup) = self._cleanup.clone() {
-            pydict.insert(
-                "_cleanup",
-                cleanup
-                    .into_pyobject(py)
-                    .context("failed to convert cleanup handle to pyobject")?
-                    .into_any()
-                    .unbind(),
-            );
-        }
-
         Ok(pydict
             .into_py_dict(py)
             .context("Failed to create py_dict")?

binaries/runtime/src/operator/python.rs

@@ -217,7 +217,6 @@ pub fn run(
 
                 let py_event = PyEvent {
                     event: MergedEvent::Dora(event),
-                    _cleanup: None,
                 }
                 .to_py_dict(py)
                 .context("Could not convert event to pydict bound")?;

examples/python-async/README.md

@@ -0,0 +1,10 @@
+# Python Async Example
+
+To get started:
+
+```bash
+uv venv --seed -p 3.11
+uv pip install -e ../../apis/python/node
+dora build dataflow.yaml --uv
+dora run dataflow.yaml --uv
+```

examples/python-async/dataflow.yaml

@@ -0,0 +1,14 @@
+nodes:
+  - id: send_data
+    build: pip install asyncio
+    path: ./send_data.py
+    inputs:
+      tick: dora/timer/millis/10
+    outputs:
+      - data
+
+  - id: receive_data_with_sleep
+    build: pip install numpy pyarrow
+    path: ./receive_data.py
+    inputs:
+      tick: send_data/data

examples/python-async/receive_data.py

@@ -0,0 +1,17 @@
+import asyncio
+
+from dora import Node
+
+
+async def main():
+    node = Node()
+    for _ in range(100):
+        event = await node.recv_async()
+        print(event)
+        # del event
+    print('done!')
+
+
+if __name__ ==  '__main__':
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())

examples/python-async/send_data.py

@@ -0,0 +1,18 @@
+"""TODO: Add docstring."""
+
+import time
+
+import numpy as np
+import pyarrow as pa
+from dora import Node
+
+node = Node()
+
+i = 0
+for event in node:
+    if i == 100:
+        break
+    else:
+        i += 1
+    now = time.perf_counter_ns()
+    node.send_output("data", pa.array([np.uint64(now)]))