Use async functions in sync context to avoid code duplication.

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.

TODO: expand writeup.


Currently there is jank around the fact that we are running synchronous actor code on a thread that already has a running event loop. I put in a workaround with a comment, but we should resolve this by not running the sync code on the asyncio loop.

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

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D78466520/)!

ghstack-source-id: 296756465
Pull Request resolved: #565

Author: zdevito

monarch_hyperactor/src/actor.rs

@@ -480,18 +480,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()
     }
 }
@@ -504,6 +505,40 @@ 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) }
+    }
+}
+
+#[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)?
+    }
+}
+
 /// An ['Actor'] used to monitor the result of an async endpoint. We use an
 /// actor so that:
 /// - Actually waiting on the async endpoint can happen concurrently with other endpoints.
@@ -557,7 +592,7 @@ impl AsyncEndpointInvocationHandler for AsyncEndpointTask {
                 Err(_) => pending().await,
             }
         };
-        let future = task.take().await;
+        let future = task.take();
         let result: Result<(), SerializablePyErr> = tokio::select! {
             result = future => {
                 match result {
@@ -595,6 +630,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PythonMessage>()?;
     hyperactor_mod.add_class::<PythonMessageKind>()?;
     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;
@@ -31,8 +32,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;
@@ -332,36 +335,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()
     }
 }
 
@@ -385,38 +373,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;
@@ -374,23 +378,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()
     }
 }
 
@@ -545,24 +549,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())
     }
 }
 
@@ -704,6 +702,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

@@ -14,6 +14,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
@@ -106,7 +107,7 @@ class ActorMeshMonitor:
         ...
 
 @final
-class MonitoredPortReceiver:
+class MonitoredPortReceiver(PortReceiverBase):
     """
     A monitored receiver to which PythonMessages are sent.
     """
@@ -117,15 +118,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.
     """
@@ -136,13 +130,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