Use async functions in sync context to avoid code duplication.

Pull Request resolved: #565

We currently have a lot of code duplication for blocking/nonblocking variants.

This PR lets the blocking variants be defined in terms of the async ones.

It first directly exposes the tokio Future object to python, letting it choose to synchronously block on it, or turn it into a asyncio.Future. Then when we are running async code to do pipe receives in a sync context, we can skip running a real event loop and just step the coroutine manually.

See [avoiding async code duplication] in the diff for a description of the approach.


Currently there is jank around the fact that we are running synchronous actor code on a thread that already has a running event loop. To indicate to the synchronous actor that it is ok to `get()` despite this setup, we temporarily blank out the running asyncio loop while a asynchronous actor is running a message. We need to eventually figure out a way to get the synchronous actors running on a thread that truly is not the asyncio event loop.

ghstack-source-id: 297098923

Differential Revision: [D78466520](https://our.internmc.facebook.com/intern/diff/D78466520/)

Author: zdevito

monarch_hyperactor/src/actor.rs

@@ -33,9 +33,11 @@ use monarch_types::SerializablePyErr;
 use pyo3::conversion::IntoPyObjectExt;
 use pyo3::exceptions::PyBaseException;
 use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::PyStopIteration;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use pyo3::types::PyDict;
+use pyo3::types::PyIterator;
 use pyo3::types::PyList;
 use pyo3::types::PyType;
 use serde::Deserialize;
@@ -616,18 +618,19 @@ impl Handler<PythonMessage> for PythonActor {
 /// Helper struct to make a Python future passable in an actor message.
 ///
 /// Also so that we don't have to write this massive type signature everywhere
-struct PythonTask {
+
+pub(crate) struct PythonTask {
     future: Mutex<Pin<Box<dyn Future<Output = PyResult<PyObject>> + Send + 'static>>>,
 }
 
 impl PythonTask {
-    fn new(fut: impl Future<Output = PyResult<PyObject>> + Send + 'static) -> Self {
+    pub(crate) fn new(fut: impl Future<Output = PyResult<PyObject>> + Send + 'static) -> Self {
         Self {
             future: Mutex::new(Box::pin(fut)),
         }
     }
 
-    async fn take(self) -> Pin<Box<dyn Future<Output = PyResult<PyObject>> + Send + 'static>> {
+    fn take(self) -> Pin<Box<dyn Future<Output = PyResult<PyObject>> + Send + 'static>> {
         self.future.into_inner()
     }
 }
@@ -640,6 +643,74 @@ impl fmt::Debug for PythonTask {
     }
 }
 
+#[pyclass(
+    name = "PythonTask",
+    module = "monarch._rust_bindings.monarch_hyperactor.actor"
+)]
+pub(crate) struct PyPythonTask {
+    inner: Option<PythonTask>,
+}
+
+impl From<PythonTask> for PyPythonTask {
+    fn from(task: PythonTask) -> Self {
+        Self { inner: Some(task) }
+    }
+}
+
+#[pyclass(
+    name = "JustStopWithValueIterator",
+    module = "monarch._rust_bindings.monarch_hyperactor.actor"
+)]
+struct JustStopWithValueIterator {
+    value: Option<PyObject>,
+}
+
+#[pymethods]
+impl JustStopWithValueIterator {
+    fn __next__<'py>(&mut self) -> PyResult<PyObject> {
+        Err(PyStopIteration::new_err(self.value.take().unwrap()))
+    }
+}
+
+#[pymethods]
+impl PyPythonTask {
+    fn into_future(&mut self, py: Python<'_>) -> PyResult<PyObject> {
+        let task = self
+            .inner
+            .take()
+            .map(|task| task.take())
+            .expect("PythonTask already consumed");
+        Ok(pyo3_async_runtimes::tokio::future_into_py(py, task)?.unbind())
+    }
+    fn block_on(&mut self, py: Python<'_>) -> PyResult<PyObject> {
+        let task = self
+            .inner
+            .take()
+            .map(|task| task.take())
+            .expect("PythonTask already consumed");
+        signal_safe_block_on(py, task)?
+    }
+
+    /// In an async context this turns the tokio::Future into
+    /// an asyncio Future and awaits it.
+    /// In a synchronous context, this just blocks on the future and
+    /// immediately returns the value without pausing caller coroutine.
+    /// See [avoiding async code duplication] for justitifcation.
+    fn __await__(&mut self, py: Python<'_>) -> PyResult<PyObject> {
+        let lp = py
+            .import("asyncio.events")
+            .unwrap()
+            .call_method0("_get_running_loop")
+            .unwrap();
+        if lp.is_none() {
+            let value = self.block_on(py)?;
+            Ok(JustStopWithValueIterator { value: Some(value) }.into_py_any(py)?)
+        } else {
+            self.into_future(py)?.call_method0(py, "__await__")
+        }
+    }
+}
+
 async fn handle_async_endpoint_panic(
     panic_sender: UnboundedSender<anyhow::Result<(), SerializablePyErr>>,
     task: PythonTask,
@@ -664,7 +735,7 @@ async fn handle_async_endpoint_panic(
             Err(_) => pending().await,
         }
     };
-    let future = task.take().await;
+    let future = task.take();
     let result: anyhow::Result<(), SerializablePyErr> = tokio::select! {
         result = future => {
             match result {
@@ -689,6 +760,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PythonMessageKind>()?;
     hyperactor_mod.add_class::<UnflattenArg>()?;
     hyperactor_mod.add_class::<PanicFlag>()?;
+    hyperactor_mod.add_class::<PyPythonTask>()?;
     Ok(())
 }
 

monarch_hyperactor/src/actor_mesh.rs

@@ -20,6 +20,7 @@ use hyperactor_mesh::actor_mesh::ActorSupervisionEvents;
 use hyperactor_mesh::reference::ActorMeshRef;
 use hyperactor_mesh::shared_cell::SharedCell;
 use hyperactor_mesh::shared_cell::SharedCellRef;
+use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyEOFError;
 use pyo3::exceptions::PyException;
 use pyo3::exceptions::PyNotImplementedError;
@@ -33,8 +34,10 @@ use serde::Deserialize;
 use serde::Serialize;
 use tokio::sync::Mutex;
 
+use crate::actor::PyPythonTask;
 use crate::actor::PythonActor;
 use crate::actor::PythonMessage;
+use crate::actor::PythonTask;
 use crate::mailbox::PyMailbox;
 use crate::mailbox::PythonOncePortReceiver;
 use crate::mailbox::PythonPortReceiver;
@@ -413,36 +416,21 @@ impl MonitoredPythonPortReceiver {
         }
     }
 
-    fn recv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+    fn recv_task<'py>(&mut self) -> PyPythonTask {
         let receiver = self.inner.clone();
         let monitor = self.monitor.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        PythonTask::new(async move {
             let mut receiver = receiver.lock().await;
-            tokio::select! {
+            let result = tokio::select! {
                 result = receiver.recv() => {
                     result.map_err(|err| PyErr::new::<PyEOFError, _>(format!("port closed: {}", err)))
                 }
                 event = monitor.next() => {
                     Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event.unwrap())))
                 }
-            }
-        })
-    }
-
-    fn blocking_recv<'py>(&mut self, py: Python<'py>) -> PyResult<PythonMessage> {
-        let receiver = self.inner.clone();
-        let monitor = self.monitor.clone();
-        signal_safe_block_on(py, async move {
-            let mut receiver = receiver.lock().await;
-            tokio::select! {
-                result = receiver.recv() => {
-                   result.map_err(|err| PyErr::new::<PyEOFError, _>(format!("port closed: {}", err)))
-                }
-                event = monitor.next() => {
-                    Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event.unwrap())))
-                }
-            }
-        })?
+            };
+            result.and_then(|message: PythonMessage| Python::with_gil(|py| message.into_py_any(py)))
+        }).into()
     }
 }
 
@@ -466,38 +454,22 @@ impl MonitoredPythonOncePortReceiver {
         }
     }
 
-    fn recv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+    fn recv_task<'py>(&mut self) -> PyResult<PyPythonTask> {
         let Some(receiver) = self.inner.lock().unwrap().take() else {
             return Err(PyErr::new::<PyValueError, _>("OncePort is already used"));
         };
         let monitor = self.monitor.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            tokio::select! {
+        Ok(PythonTask::new(async move {
+            let result = tokio::select! {
                 result = receiver.recv() => {
                     result.map_err(|err| PyErr::new::<PyEOFError, _>(format!("port closed: {}", err)))
                 }
                 event = monitor.next() => {
                     Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event.unwrap())))
                 }
-            }
-        })
-    }
-
-    fn blocking_recv<'py>(&mut self, py: Python<'py>) -> PyResult<PythonMessage> {
-        let Some(receiver) = self.inner.lock().unwrap().take() else {
-            return Err(PyErr::new::<PyValueError, _>("OncePort is already used"));
-        };
-        let monitor = self.monitor.clone();
-        signal_safe_block_on(py, async move {
-            tokio::select! {
-                result = receiver.recv() => {
-                   result.map_err(|err| PyErr::new::<PyEOFError, _>(format!("port closed: {}", err)))
-                }
-                event = monitor.next() => {
-                    Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event.unwrap())))
-                }
-            }
-        })?
+            };
+            result.and_then(|message: PythonMessage| Python::with_gil(|py| message.into_py_any(py)))
+        }).into())
     }
 }
 

monarch_hyperactor/src/mailbox.rs

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::future::Future;
 use std::hash::DefaultHasher;
 use std::hash::Hash;
 use std::hash::Hasher;
@@ -45,9 +46,12 @@ use pyo3::types::PyTuple;
 use pyo3::types::PyType;
 use serde::Deserialize;
 use serde::Serialize;
+use tokio::sync::mpsc::UnboundedReceiver;
 
+use crate::actor::PyPythonTask;
 use crate::actor::PythonMessage;
 use crate::actor::PythonMessageKind;
+use crate::actor::PythonTask;
 use crate::proc::PyActorId;
 use crate::runtime::signal_safe_block_on;
 use crate::shape::PyShape;
@@ -380,23 +384,23 @@ pub(super) struct PythonPortReceiver {
     inner: Arc<tokio::sync::Mutex<PortReceiver<PythonMessage>>>,
 }
 
+async fn recv_async(
+    receiver: Arc<tokio::sync::Mutex<PortReceiver<PythonMessage>>>,
+) -> PyResult<PyObject> {
+    receiver
+        .lock()
+        .await
+        .recv()
+        .await
+        .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))
+        .and_then(|message| Python::with_gil(|py| message.into_py_any(py)))
+}
+
 #[pymethods]
 impl PythonPortReceiver {
-    fn recv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+    fn recv_task<'py>(&mut self) -> PyPythonTask {
         let receiver = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            receiver
-                .lock()
-                .await
-                .recv()
-                .await
-                .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))
-        })
-    }
-    fn blocking_recv<'py>(&mut self, py: Python<'py>) -> PyResult<PythonMessage> {
-        let receiver = self.inner.clone();
-        signal_safe_block_on(py, async move { receiver.lock().await.recv().await })?
-            .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))
+        PythonTask::new(recv_async(receiver)).into()
     }
 }
 
@@ -551,24 +555,18 @@ pub(super) struct PythonOncePortReceiver {
 
 #[pymethods]
 impl PythonOncePortReceiver {
-    fn recv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+    fn recv_task<'py>(&mut self) -> PyResult<PyPythonTask> {
         let Some(receiver) = self.inner.lock().unwrap().take() else {
             return Err(PyErr::new::<PyValueError, _>("OncePort is already used"));
         };
-
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let fut = async move {
             receiver
                 .recv()
                 .await
                 .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))
-        })
-    }
-    fn blocking_recv<'py>(&mut self, py: Python<'py>) -> PyResult<PythonMessage> {
-        let Some(receiver) = self.inner.lock().unwrap().take() else {
-            return Err(PyErr::new::<PyValueError, _>("OncePort is already used"));
+                .and_then(|message| Python::with_gil(|py| message.into_py_any(py)))
         };
-        signal_safe_block_on(py, async move { receiver.recv().await })?
-            .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))
+        Ok(PythonTask::new(fut).into())
     }
 }
 
@@ -710,6 +708,5 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PythonOncePortHandle>()?;
     hyperactor_mod.add_class::<PythonOncePortRef>()?;
     hyperactor_mod.add_class::<PythonOncePortReceiver>()?;
-
     Ok(())
 }

python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi

@@ -13,6 +13,7 @@ from monarch._rust_bindings.monarch_hyperactor.mailbox import (
     Mailbox,
     OncePortReceiver,
     PortReceiver,
+    PortReceiverBase,
 )
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
 from monarch._rust_bindings.monarch_hyperactor.selection import Selection
@@ -128,7 +129,7 @@ class ActorMeshMonitor:
         ...
 
 @final
-class MonitoredPortReceiver:
+class MonitoredPortReceiver(PortReceiverBase):
     """
     A monitored receiver to which PythonMessages are sent.
     """
@@ -139,15 +140,8 @@ class MonitoredPortReceiver:
         """
         ...
 
-    async def recv(self) -> PythonMessage:
-        """Receive a PythonMessage from the port's sender."""
-        ...
-    def blocking_recv(self) -> PythonMessage:
-        """Receive a single PythonMessage from the port's sender."""
-        ...
-
 @final
-class MonitoredOncePortReceiver:
+class MonitoredOncePortReceiver(PortReceiverBase):
     """
     A variant of monitored PortReceiver that can only receive a single message.
     """
@@ -158,13 +152,6 @@ class MonitoredOncePortReceiver:
         """
         ...
 
-    async def recv(self) -> PythonMessage:
-        """Receive a single PythonMessage from the port's sender."""
-        ...
-    def blocking_recv(self) -> PythonMessage:
-        """Receive a single PythonMessage from the port's sender."""
-        ...
-
 @final
 class ActorSupervisionEvent:
     @property