actor: move more port receiver supervision to rust (#578)

Summary:
Pull Request resolved: #578

this diff moves supervision logic from python into rust, aligning with the goal of eliminating complex supervision wiring in python.

the essential change is that:
```
class ActorEndpoint(...):
    def _port(self, once: bool = False) -> "PortTuple[R]":
        monitor = (
            None
            if self._actor_mesh._actor_mesh is None
            else self._actor_mesh._actor_mesh.monitor()
        )
        return PortTuple.create(self._mailbox, monitor, once)
```
becomes:
```
class ActorEndpoint(...):
    def _port(self, once: bool = False) -> PortTuple[R]:
        p, r = PortTuple.create(self._mailbox, once)
        return PortTuple(p, PortReceiver(self._mailbox, self._supervise(r._receiver)))
```
`_supervise(...)` dispatches to new Rust helpers:
```python
mesh.supervise_port(...) and mesh.supervise_once_port(...)
```
which wrap the receivers with supervision logic (including selection between message arrival and supervision events), completely eliminating the need for python-side constructs like `ActorMeshMonitor`.

most of the python complexity introduced in D77434080 is removed. the only meaningful addition is `_supervise(...)`, a small overrideable hook that defaults to a no-op and cleanly delegates to rust when supervision is desired.

- the creation and wiring of the monitor stream is now fully in rust.
- the responsibility of wrapping receivers with supervision is now fully in rust.
- python no longer constructs or passes supervision monitors; rust now owns the full wiring, and python receives already-wrapped receivers with supervision behavior embedded

this is a strict improvement: lower complexity, cleaner override points and supervision is entirely managed in rust.

Differential Revision: D78528860

Author: shayne-fletcher

monarch_hyperactor/src/actor_mesh.rs

@@ -176,14 +176,32 @@ impl PythonActorMesh {
             .map(PyActorId::from))
     }
 
-    // Start monitoring the actor mesh by subscribing to its supervision events. For each supervision
-    // event, it is consumed by PythonActorMesh first, then gets sent to the monitor for user to consume.
-    fn monitor<'py>(&self, py: Python<'py>) -> PyResult<PyObject> {
-        let receiver = self.user_monitor_sender.subscribe();
-        let monitor_instance = PyActorMeshMonitor {
-            receiver: SharedCell::from(Mutex::new(receiver)),
+    fn supervise_port<'py>(
+        &self,
+        py: Python<'py>,
+        receiver: &PythonPortReceiver,
+    ) -> PyResult<PyObject> {
+        let rx = MonitoredPythonPortReceiver {
+            inner: receiver.inner(),
+            monitor: PyActorMeshMonitor {
+                receiver: SharedCell::from(Mutex::new(self.user_monitor_sender.subscribe())),
+            },
+        };
+        rx.into_py_any(py)
+    }
+
+    fn supervise_once_port<'py>(
+        &self,
+        py: Python<'py>,
+        receiver: &PythonOncePortReceiver,
+    ) -> PyResult<PyObject> {
+        let rx = MonitoredPythonOncePortReceiver {
+            inner: receiver.inner(),
+            monitor: PyActorMeshMonitor {
+                receiver: SharedCell::from(Mutex::new(self.user_monitor_sender.subscribe())),
+            },
         };
-        Ok(monitor_instance.into_py(py))
+        rx.into_py_any(py)
     }
 
     #[pyo3(signature = (**kwargs))]
@@ -335,29 +353,43 @@ impl Drop for PythonActorMesh {
     }
 }
 
+// `PyActorMeshMonitor` is not accessed directly from Python. It is
+// marked with `#[pyclass]` so it can be used as a field inside
+// `MonitoredPythonPortReceiver`.
 #[pyclass(
     name = "ActorMeshMonitor",
     module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
 )]
-pub struct PyActorMeshMonitor {
+struct PyActorMeshMonitor {
     receiver: SharedCell<Mutex<tokio::sync::broadcast::Receiver<Option<ActorSupervisionEvent>>>>,
 }
 
 #[pymethods]
 impl PyActorMeshMonitor {
-    fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
-        slf
-    }
-
-    pub fn __anext__(&self, py: Python<'_>) -> PyResult<PyObject> {
-        let receiver = self.receiver.clone();
-        Ok(pyo3_async_runtimes::tokio::future_into_py(py, get_next(receiver))?.into())
+    fn __repr__(&self) -> &'static str {
+        "<ActorMeshMonitor>"
     }
 }
 
 impl PyActorMeshMonitor {
     pub async fn next(&self) -> PyResult<PyObject> {
-        get_next(self.receiver.clone()).await
+        let receiver = self.receiver.clone();
+        let receiver = receiver
+            .borrow()
+            .expect("`Actor mesh receiver` is shutdown");
+        let mut receiver = receiver.lock().await;
+        let event = receiver.recv().await.unwrap();
+        let supervision_event = match event {
+            None => PyActorSupervisionEvent {
+                // Dummy actor as place holder to indicate the whole mesh is stopped
+                // TODO(albertli): remove this when pushing all supervision logic to rust.
+                actor_id: id!(default[0].actor[0]).into(),
+                actor_status: "actor mesh is stopped due to proc mesh shutdown".to_string(),
+            },
+            Some(event) => PyActorSupervisionEvent::from(event.clone()),
+        };
+
+        Python::with_gil(|py| supervision_event.into_py_any(py))
     }
 }
 
@@ -369,49 +401,21 @@ impl Clone for PyActorMeshMonitor {
     }
 }
 
-async fn get_next(
-    receiver: SharedCell<Mutex<tokio::sync::broadcast::Receiver<Option<ActorSupervisionEvent>>>>,
-) -> PyResult<PyObject> {
-    let receiver = receiver.clone();
-
-    let receiver = receiver
-        .borrow()
-        .expect("`Actor mesh receiver` is shutdown");
-    let mut receiver = receiver.lock().await;
-    let event = receiver.recv().await.unwrap();
-
-    let supervision_event = match event {
-        None => PyActorSupervisionEvent {
-            // Dummy actor as place holder to indicate the whole mesh is stopped
-            // TODO(albertli): remove this when pushing all supervision logic to rust.
-            actor_id: id!(default[0].actor[0]).into(),
-            actor_status: "actor mesh is stopped due to proc mesh shutdown".to_string(),
-        },
-        Some(event) => PyActorSupervisionEvent::from(event.clone()),
-    };
-
-    Python::with_gil(|py| supervision_event.into_py_any(py))
-}
-
-// TODO(albertli): this is temporary remove this when pushing all supervision logic to rust.
+// Values of this (private) type can only be created by calling
+// `PythonActorMesh::supervise_port()`.
 #[pyclass(
     name = "MonitoredPortReceiver",
     module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
 )]
-pub(super) struct MonitoredPythonPortReceiver {
+struct MonitoredPythonPortReceiver {
     inner: Arc<tokio::sync::Mutex<PortReceiver<PythonMessage>>>,
     monitor: PyActorMeshMonitor,
 }
 
 #[pymethods]
 impl MonitoredPythonPortReceiver {
-    #[new]
-    fn new(receiver: &PythonPortReceiver, monitor: &PyActorMeshMonitor) -> Self {
-        let inner = receiver.inner();
-        MonitoredPythonPortReceiver {
-            inner,
-            monitor: monitor.clone(),
-        }
+    fn __repr__(&self) -> &'static str {
+        "<MonitoredPortReceiver>"
     }
 
     fn recv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
@@ -447,24 +451,21 @@ impl MonitoredPythonPortReceiver {
     }
 }
 
+// Values of this (private) type can only be created by calling
+// `PythonActorMesh::supervise_once_port()`.
 #[pyclass(
     name = "MonitoredOncePortReceiver",
     module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
 )]
-pub(super) struct MonitoredPythonOncePortReceiver {
+struct MonitoredPythonOncePortReceiver {
     inner: Arc<std::sync::Mutex<Option<OncePortReceiver<PythonMessage>>>>,
     monitor: PyActorMeshMonitor,
 }
 
 #[pymethods]
 impl MonitoredPythonOncePortReceiver {
-    #[new]
-    fn new(receiver: &PythonOncePortReceiver, monitor: &PyActorMeshMonitor) -> Self {
-        let inner = receiver.inner();
-        MonitoredPythonOncePortReceiver {
-            inner,
-            monitor: monitor.clone(),
-        }
+    fn __repr__(&self) -> &'static str {
+        "<MonitoredOncePortReceiver>"
     }
 
     fn recv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {

python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi

@@ -88,10 +88,15 @@ class PythonActorMesh:
         """
         ...
 
-    # TODO(albertli): remove this when pushing all supervision logic to Rust
-    def monitor(self) -> ActorMeshMonitor:
+    def supervise_port(self, r: PortReceiver) -> MonitoredPortReceiver:
         """
-        Returns a supervision monitor for this mesh.
+        Return a monitored port receiver.
+        """
+        ...
+
+    def supervise_once_port(self, r: OncePortReceiver) -> MonitoredOncePortReceiver:
+        """
+        Return a monitored once port receiver.
         """
         ...
 
@@ -113,31 +118,11 @@ class PythonActorMesh:
         """
         ...
 
-@final
-class ActorMeshMonitor:
-    def __aiter__(self) -> AsyncIterator["ActorSupervisionEvent"]:
-        """
-        Returns an async iterator for this monitor.
-        """
-        ...
-
-    async def __anext__(self) -> "ActorSupervisionEvent":
-        """
-        Returns the next proc event in the proc mesh.
-        """
-        ...
-
 @final
 class MonitoredPortReceiver:
+    """A monitored receiver to which PythonMessages are sent. Values
+    of this type cannot be constructed directly in Python.
     """
-    A monitored receiver to which PythonMessages are sent.
-    """
-
-    def __init__(self, receiver: PortReceiver, monitor: ActorMeshMonitor) -> None:
-        """
-        Create a new monitored receiver from a PortReceiver.
-        """
-        ...
 
     async def recv(self) -> PythonMessage:
         """Receive a PythonMessage from the port's sender."""
@@ -148,15 +133,9 @@ class MonitoredPortReceiver:
 
 @final
 class MonitoredOncePortReceiver:
+    """A monitored once receiver to which PythonMessages are sent.
+    Values of this type cannot be constructed directly in Python.
     """
-    A variant of monitored PortReceiver that can only receive a single message.
-    """
-
-    def __init__(self, receiver: OncePortReceiver, monitor: ActorMeshMonitor) -> None:
-        """
-        Create a new monitored receiver from a PortReceiver.
-        """
-        ...
 
     async def recv(self) -> PythonMessage:
         """Receive a single PythonMessage from the port's sender."""

python/monarch/_src/actor/actor_mesh.py

@@ -38,6 +38,8 @@
     Optional,
     overload,
     ParamSpec,
+    Protocol,
+    runtime_checkable,
     Sequence,
     Tuple,
     Type,
@@ -50,12 +52,7 @@
     PythonMessage,
     PythonMessageKind,
 )
-from monarch._rust_bindings.monarch_hyperactor.actor_mesh import (
-    ActorMeshMonitor,
-    MonitoredOncePortReceiver,
-    MonitoredPortReceiver,
-    PythonActorMesh,
-)
+from monarch._rust_bindings.monarch_hyperactor.actor_mesh import PythonActorMesh
 from monarch._rust_bindings.monarch_hyperactor.mailbox import (
     Mailbox,
     OncePortReceiver,
@@ -306,6 +303,9 @@ def _send(
     def _port(self, once: bool = False) -> "PortTuple[R]":
         pass
 
+    def _supervise(self, r: HyPortReceiver | OncePortReceiver) -> Any:
+        return r
+
     # the following are all 'adverbs' or different ways to handle the
     # return values of this endpoint. Adverbs should only ever take *args, **kwargs
     # of the original call. If we want to add syntax sugar for something that needs additional
@@ -399,6 +399,16 @@ def __init__(
         self._signature: inspect.Signature = inspect.signature(impl)
         self._mailbox = mailbox
 
+    def _supervise(self, r: HyPortReceiver | OncePortReceiver) -> Any:
+        mesh = self._actor_mesh._actor_mesh
+        if mesh is None:
+            return r
+        return (
+            mesh.supervise_once_port(r)
+            if isinstance(r, OncePortReceiver)
+            else mesh.supervise_port(r)
+        )
+
     def _send(
         self,
         args: Tuple[Any, ...],
@@ -430,12 +440,11 @@ def _send(
         return Extent(shape.labels, shape.ndslice.sizes)
 
     def _port(self, once: bool = False) -> "PortTuple[R]":
-        monitor = (
-            None
-            if self._actor_mesh._actor_mesh is None
-            else self._actor_mesh._actor_mesh.monitor()
-        )
-        return PortTuple.create(self._mailbox, monitor, once)
+        p, r = PortTuple.create(self._mailbox, once)
+        assert isinstance(
+            r._receiver, (HyPortReceiver, OncePortReceiver)
+        ), "unexpected receiver type"
+        return PortTuple(p, PortReceiver(self._mailbox, self._supervise(r._receiver)))
 
 
 class Accumulator(Generic[P, R, A]):
@@ -589,21 +598,11 @@ class PortTuple(NamedTuple, Generic[R]):
         receiver: "PortReceiver[R]"
 
         @staticmethod
-        def create(
-            mailbox: Mailbox, monitor: Optional[ActorMeshMonitor], once: bool = False
-        ) -> "PortTuple[Any]":
+        def create(mailbox: Mailbox, once: bool = False) -> "PortTuple[Any]":
             handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
             port_ref = handle.bind()
-            if monitor is not None:
-                receiver = (
-                    MonitoredOncePortReceiver(receiver, monitor)
-                    if isinstance(receiver, OncePortReceiver)
-                    else MonitoredPortReceiver(receiver, monitor)
-                )
-
             return PortTuple(
-                Port(port_ref, mailbox, rank=None),
-                PortReceiver(mailbox, receiver),
+                Port(port_ref, mailbox, rank=None), PortReceiver(mailbox, receiver)
             )
 else:
 
@@ -612,21 +611,11 @@ class PortTuple(NamedTuple):
         receiver: "PortReceiver[Any]"
 
         @staticmethod
-        def create(
-            mailbox: Mailbox, monitor: Optional[ActorMeshMonitor], once: bool = False
-        ) -> "PortTuple[Any]":
+        def create(mailbox: Mailbox, once: bool = False) -> "PortTuple[Any]":
             handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
             port_ref = handle.bind()
-            if monitor is not None:
-                receiver = (
-                    MonitoredOncePortReceiver(receiver, monitor)
-                    if isinstance(receiver, OncePortReceiver)
-                    else MonitoredPortReceiver(receiver, monitor)
-                )
-
             return PortTuple(
-                Port(port_ref, mailbox, rank=None),
-                PortReceiver(mailbox, receiver),
+                Port(port_ref, mailbox, rank=None), PortReceiver(mailbox, receiver)
             )
 
 
@@ -644,22 +633,19 @@ def ranked_port(
     return p, RankedPortReceiver[R](receiver._mailbox, receiver._receiver)
 
 
+R = TypeVar("R")
+
+
+@runtime_checkable
+class ReceiverLike(Protocol[R]):
+    def blocking_recv(self) -> R: ...
+    async def recv(self) -> R: ...
+
+
 class PortReceiver(Generic[R]):
-    def __init__(
-        self,
-        mailbox: Mailbox,
-        receiver: MonitoredPortReceiver
-        | MonitoredOncePortReceiver
-        | HyPortReceiver
-        | OncePortReceiver,
-    ) -> None:
-        self._mailbox: Mailbox = mailbox
-        self._receiver: (
-            MonitoredPortReceiver
-            | MonitoredOncePortReceiver
-            | HyPortReceiver
-            | OncePortReceiver
-        ) = receiver
+    def __init__(self, mailbox: Mailbox, receiver: ReceiverLike) -> None:
+        self._mailbox = mailbox
+        self._receiver: ReceiverLike = receiver
 
     async def _recv(self) -> R:
         return self._process(await self._receiver.recv())