: actor: port receiver supervision (#578)

Summary: Pull Request resolved: #578

Differential Revision: D78528860

Author: shayne-fletcher

monarch_hyperactor/src/actor_mesh.rs

@@ -190,14 +190,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())),
+            },
         };
-        Ok(monitor_instance.into_py(py))
+        Ok(rx.into_py(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(rx.into_py(py))
     }
 
     #[pyo3(signature = (**kwargs))]
@@ -349,27 +367,22 @@ 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
     }
@@ -407,25 +420,21 @@ async fn get_next(
     Ok(Python::with_gil(|py| supervision_event.into_py(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>> {
@@ -461,24 +470,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

@@ -50,12 +50,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 +301,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 +397,14 @@ 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
+        return (
+            mesh.supervise_once_port(r)
+            if isinstance(r, OncePortReceiver)
+            else mesh.supervise_port(r)
+        )
+
     def _send(
         self,
         args: Tuple[Any, ...],
@@ -430,12 +436,8 @@ 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)
+        return PortTuple(p, PortReceiver(self._mailbox, self._supervise(r._receiver)))
 
 
 class Accumulator(Generic[P, R, A]):
@@ -589,21 +591,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 +604,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)
             )
 
 
@@ -648,18 +630,10 @@ class PortReceiver(Generic[R]):
     def __init__(
         self,
         mailbox: Mailbox,
-        receiver: MonitoredPortReceiver
-        | MonitoredOncePortReceiver
-        | HyPortReceiver
-        | OncePortReceiver,
+        receiver: HyPortReceiver | OncePortReceiver,
     ) -> None:
         self._mailbox: Mailbox = mailbox
-        self._receiver: (
-            MonitoredPortReceiver
-            | MonitoredOncePortReceiver
-            | HyPortReceiver
-            | OncePortReceiver
-        ) = receiver
+        self._receiver: HyPortReceiver | OncePortReceiver = receiver
 
     async def _recv(self) -> R:
         return self._process(await self._receiver.recv())

python/monarch/common/remote.py

@@ -144,7 +144,7 @@ def _port(self, once: bool = False) -> "PortTuple[R]":
                 "Cannot create raw port objects with an old-style tensor engine controller."
             )
         mailbox: Mailbox = mesh_controller._mailbox
-        return PortTuple.create(mailbox, None, once)
+        return PortTuple.create(mailbox, once)
 
     @property
     def _resolvable(self):

python/monarch/mesh_controller.py

@@ -149,9 +149,7 @@ def fetch(
         defs: Tuple["Tensor", ...],
         uses: Tuple["Tensor", ...],
     ) -> "OldFuture":  # the OldFuture is a lie
-        sender, receiver = PortTuple.create(
-            self._mesh_controller._mailbox, None, once=True
-        )
+        sender, receiver = PortTuple.create(self._mesh_controller._mailbox, once=True)
 
         ident = self.new_node(defs, uses, cast("OldFuture", sender))
         process = mesh._process(shard)