: actor: port receiver monitoring is made non-optional

Differential Revision: D78528860

Author: shayne-fletcher

monarch_hyperactor/src/actor_mesh.rs

@@ -190,14 +190,33 @@ impl PythonActorMesh {
             .map(PyActorId::from))
     }
 
+    // Return a monitored port receiver.
+    fn supervise_port<'py>(&self, py: Python<'py> , receiver: &PythonPortReceiver) -> PyResult<PyObject> {
+        let monitor = self._monitor(py)?;
+        let rx = MonitoredPythonPortReceiver {
+            inner: receiver.inner(),
+            monitor,
+        };
+        Ok(rx.into_py(py))
+    }
+
+    fn supervise_once_port<'py>(&self, py: Python<'py> , receiver: &PythonOncePortReceiver) -> PyResult<PyObject> {
+        let monitor = self._monitor(py)?;
+        let rx = MonitoredPythonOncePortReceiver {
+            inner: receiver.inner(),
+            monitor,
+        };
+        Ok(rx.into_py(py))
+    }
+
     // 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> {
+    fn _monitor<'py>(&self, py: Python<'py>) -> PyResult<PyActorMeshMonitor> {
         let receiver = self.user_monitor_sender.subscribe();
         let monitor_instance = PyActorMeshMonitor {
             receiver: SharedCell::from(Mutex::new(receiver)),
         };
-        Ok(monitor_instance.into_py(py))
+        Ok(monitor_instance)
     }
 
     #[pyo3(signature = (**kwargs))]
@@ -353,21 +372,21 @@ impl Drop for PythonActorMesh {
     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
-    }
+// #[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())
-    }
-}
+//     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())
+//     }
+// }
 
 impl PyActorMeshMonitor {
     pub async fn next(&self) -> PyResult<PyObject> {
@@ -407,28 +426,24 @@ 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>> {
+   fn recv<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let receiver = self.inner.clone();
         let monitor = self.monitor.clone();
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
@@ -461,24 +476,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."""