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

@@ -177,14 +177,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: ActorMeshMonitor {
+                receiver: SharedCell::from(Mutex::new(self.user_monitor_sender.subscribe())),
+            },
         };
-        Ok(monitor_instance.into_py(py))
+        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: ActorMeshMonitor {
+                receiver: SharedCell::from(Mutex::new(self.user_monitor_sender.subscribe())),
+            },
+        };
+        rx.into_py_any(py)
     }
 
     #[pyo3(signature = (**kwargs))]
@@ -336,83 +354,46 @@ impl Drop for PythonActorMesh {
     }
 }
 
-#[pyclass(
-    name = "ActorMeshMonitor",
-    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
-)]
-pub struct PyActorMeshMonitor {
+#[derive(Debug, Clone)]
+struct ActorMeshMonitor {
     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> {
+impl ActorMeshMonitor {
+    pub async fn next(&self) -> PyActorSupervisionEvent {
         let receiver = self.receiver.clone();
-        Ok(pyo3_async_runtimes::tokio::future_into_py(py, get_next(receiver))?.into())
-    }
-}
-
-impl PyActorMeshMonitor {
-    pub async fn next(&self) -> PyResult<PyObject> {
-        get_next(self.receiver.clone()).await
-    }
-}
-
-impl Clone for PyActorMeshMonitor {
-    fn clone(&self) -> Self {
-        Self {
-            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();
+        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()),
         }
     }
 }
 
-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,
+    monitor: ActorMeshMonitor,
 }
 
 #[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_task<'py>(&mut self) -> PyPythonTask {
@@ -425,32 +406,29 @@ impl MonitoredPythonPortReceiver {
                     result.map_err(|err| PyErr::new::<PyEOFError, _>(format!("port closed: {}", err)))
                 }
                 event = monitor.next() => {
-                    Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event.unwrap())))
+                    Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event)))
                 }
             };
             result.and_then(|message: PythonMessage| Python::with_gil(|py| message.into_py_any(py)))
         }).into()
     }
 }
 
+// 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,
+    monitor: ActorMeshMonitor,
 }
 
 #[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_task<'py>(&mut self) -> PyResult<PyPythonTask> {
@@ -464,7 +442,7 @@ impl MonitoredPythonOncePortReceiver {
                     result.map_err(|err| PyErr::new::<PyEOFError, _>(format!("port closed: {}", err)))
                 }
                 event = monitor.next() => {
-                    Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event.unwrap())))
+                    Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event)))
                 }
             };
             result.and_then(|message: PythonMessage| Python::with_gil(|py| message.into_py_any(py)))
@@ -476,6 +454,7 @@ impl MonitoredPythonOncePortReceiver {
     name = "ActorSupervisionEvent",
     module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
 )]
+#[derive(Debug)]
 pub struct PyActorSupervisionEvent {
     /// Actor ID of the actor where supervision event originates from.
     #[pyo3(get)]
@@ -508,7 +487,6 @@ impl From<ActorSupervisionEvent> for PyActorSupervisionEvent {
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     hyperactor_mod.add_class::<PythonActorMesh>()?;
     hyperactor_mod.add_class::<PythonActorMeshRef>()?;
-    hyperactor_mod.add_class::<PyActorMeshMonitor>()?;
     hyperactor_mod.add_class::<MonitoredPythonPortReceiver>()?;
     hyperactor_mod.add_class::<MonitoredPythonOncePortReceiver>()?;
     hyperactor_mod.add_class::<PyActorSupervisionEvent>()?;

python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi

@@ -89,10 +89,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.
         """
         ...
 
@@ -114,43 +119,17 @@ 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(PortReceiverBase):
+    """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.
-        """
-        ...
 
 @final
 class MonitoredOncePortReceiver(PortReceiverBase):
+    """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.
-        """
-        ...
 
 @final
 class ActorSupervisionEvent:

python/monarch/_src/actor/actor_mesh.py

@@ -40,6 +40,8 @@
     Optional,
     overload,
     ParamSpec,
+    Protocol,
+    runtime_checkable,
     Sequence,
     Tuple,
     Type,
@@ -52,16 +54,12 @@
     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,
     OncePortRef,
+    PortReceiver as HyPortReceiver,
     PortRef,
 )
 
@@ -311,6 +309,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
@@ -393,6 +394,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, ...],
@@ -424,12 +435,12 @@ 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)
+        if TYPE_CHECKING:
+            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]):
@@ -583,21 +594,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:
 
@@ -606,18 +607,9 @@ 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),