Move supervision directly to PortReceiver object (#771)

Summary:
Pull Request resolved: #771

To de-futureify actor spawning, I have to be able to create a port before the rust actormesh exists. So to monitor that port the monitor has to first wait for the actormesh and then wait for the monitor. This refactors how monitors are handled to make it possible to write this.

As an added benefit, it removes a bunch of copy/paste for difference receiver types.
ghstack-source-id: 301241546
exported-using-ghexport

Reviewed By: mariusae

Differential Revision: D79695335

fbshipit-source-id: d14cb519ea00702e1de7f29f63ad3a038380eb8a

Author: zdevito

monarch_hyperactor/src/actor_mesh.rs

@@ -194,29 +194,24 @@ impl PythonActorMesh {
             .map(ActorRef::into_actor_id)
             .map(PyActorId::from))
     }
-
-    fn supervise(&self, py: Python<'_>, receiver: Bound<'_, PyAny>) -> PyResult<PyObject> {
-        if let Ok(r) = receiver.extract::<PyRef<PythonPortReceiver>>() {
-            let rx = SupervisedPythonPortReceiver {
-                inner: r.inner(),
-                monitor: ActorMeshMonitor {
-                    receiver: SharedCell::from(Mutex::new(self.user_monitor_sender.subscribe())),
-                },
-            };
-            rx.into_py_any(py)
-        } else if let Ok(r) = receiver.extract::<PyRef<PythonOncePortReceiver>>() {
-            let rx = SupervisedPythonOncePortReceiver {
-                inner: r.inner(),
-                monitor: ActorMeshMonitor {
-                    receiver: SharedCell::from(Mutex::new(self.user_monitor_sender.subscribe())),
+    fn supervision_event(&self) -> PyResult<PyPythonTask> {
+        let mut receiver = self.user_monitor_sender.subscribe();
+        PyPythonTask::new(async move {
+            let event = receiver.recv().await;
+            let event = match event {
+                Ok(Some(event)) => PyActorSupervisionEvent::from(event.clone()),
+                Ok(None) | Err(_) => PyActorSupervisionEvent {
+                    // Dummy actor as placeholder 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(),
                 },
             };
-            rx.into_py_any(py)
-        } else {
-            Err(PyTypeError::new_err(
-                "Expected a PortReceiver or OncePortReceiver",
-            ))
-        }
+            Ok(PyErr::new::<SupervisionError, _>(format!(
+                "supervision error: {:?}",
+                event
+            )))
+        })
     }
 
     #[pyo3(signature = (**kwargs))]
@@ -406,106 +401,6 @@ impl Drop for PythonActorMesh {
     }
 }
 
-#[derive(Debug, Clone)]
-struct ActorMeshMonitor {
-    receiver: SharedCell<Mutex<tokio::sync::broadcast::Receiver<Option<ActorSupervisionEvent>>>>,
-}
-
-impl ActorMeshMonitor {
-    pub async fn next(&self) -> Result<PyActorSupervisionEvent, PyErr> {
-        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;
-        Ok(match event {
-            Ok(Some(event)) => PyActorSupervisionEvent::from(event.clone()),
-            Ok(None) | Err(_) => PyActorSupervisionEvent {
-                // Dummy actor as placeholder 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(),
-            },
-        })
-    }
-}
-
-// Values of this type can only be created by calling
-// `PythonActorMesh::supervise()`.
-#[pyclass(
-    name = "SupervisedPortReceiver",
-    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
-)]
-struct SupervisedPythonPortReceiver {
-    inner: Arc<tokio::sync::Mutex<PortReceiver<PythonMessage>>>,
-    monitor: ActorMeshMonitor,
-}
-
-#[pymethods]
-impl SupervisedPythonPortReceiver {
-    fn __repr__(&self) -> &'static str {
-        "<SupervisedPortReceiver>"
-    }
-
-    fn recv_task(&mut self) -> PyPythonTask {
-        let receiver = self.inner.clone();
-        let monitor = self.monitor.clone();
-        PythonTask::new(async move {
-            let mut receiver = receiver.lock().await;
-            let result = tokio::select! {
-                result = receiver.recv() => {
-                    result.map_err(|err| PyErr::new::<PyEOFError, _>(format!("port closed: {}", err)))
-                }
-                event = monitor.next() => {
-                    Python::with_gil(|_py| {
-                        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 type can only be created by calling
-// `PythonActorMesh::supervise()`.
-#[pyclass(
-    name = "SupervisedOncePortReceiver",
-    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
-)]
-struct SupervisedPythonOncePortReceiver {
-    inner: Arc<std::sync::Mutex<Option<OncePortReceiver<PythonMessage>>>>,
-    monitor: ActorMeshMonitor,
-}
-
-#[pymethods]
-impl SupervisedPythonOncePortReceiver {
-    fn __repr__(&self) -> &'static str {
-        "<SupervisedOncePortReceiver>"
-    }
-
-    fn recv_task(&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();
-        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() => {
-                    Python::with_gil(|_py| {
-                        Err(PyErr::new::<SupervisionError, _>(format!("supervision error: {:?}", event)))
-                    })
-                }
-            };
-            result.and_then(|message: PythonMessage| Python::with_gil(|py| message.into_py_any(py)))
-        }).into())
-    }
-}
-
 #[pyclass(
     name = "ActorSupervisionEvent",
     module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
@@ -543,8 +438,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::<SupervisedPythonPortReceiver>()?;
-    hyperactor_mod.add_class::<SupervisedPythonOncePortReceiver>()?;
     hyperactor_mod.add_class::<PyActorSupervisionEvent>()?;
     Ok(())
 }

monarch_hyperactor/src/pytokio.rs

@@ -244,6 +244,23 @@ impl PyPythonTask {
         });
         Ok(PyShared { rx })
     }
+
+    #[staticmethod]
+    fn select_one(mut tasks: Vec<PyRefMut<'_, PyPythonTask>>) -> PyResult<PyPythonTask> {
+        if tasks.is_empty() {
+            return Err(PyValueError::new_err("Cannot select from empty task list"));
+        }
+
+        let mut futures = Vec::new();
+        for task_ref in tasks.iter_mut() {
+            futures.push(task_ref.take_task()?);
+        }
+
+        PyPythonTask::new(async move {
+            let (result, index, _remaining) = futures::future::select_all(futures).await;
+            result.map(|r| (r, index))
+        })
+    }
 }
 
 #[pyclass(

python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi

@@ -6,16 +6,16 @@
 
 # pyre-strict
 
-from typing import AsyncIterator, final
+from typing import AsyncIterator, final, NoReturn
 
 from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage
 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.pytokio import PythonTask
 from monarch._rust_bindings.monarch_hyperactor.selection import Selection
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape
 from typing_extensions import Self
@@ -97,23 +97,15 @@ class PythonActorMesh:
         """
         ...
 
-    def get(self, rank: int) -> ActorId | None:
+    def supervision_event(self) -> PythonTask[Exception]:
         """
-        Get the actor id for the actor at the given rank.
+        Completes with an exception when there is a supervision error.
         """
         ...
 
-    def supervise(
-        self, r: PortReceiver | OncePortReceiver
-    ) -> SupervisedPortReceiver | SupervisedOncePortReceiver:
-        """Return a monitored port receiver.
-
-        A monitored port receiver behaves like a regular port receiver
-        but also observes the health of the actor mesh associated with
-        the sender. If the actor mesh becomes unhealthy, the receiver
-        will yield a supervision error instead of waiting indefinitely
-        for a message.
-
+    def get(self, rank: int) -> ActorId | None:
+        """
+        Get the actor id for the actor at the given rank.
         """
         ...
 
@@ -149,18 +141,6 @@ class PythonActorMesh:
         """
         ...
 
-@final
-class SupervisedPortReceiver(PortReceiverBase):
-    """A monitored receiver to which PythonMessages are sent. Values
-    of this type cannot be constructed directly in Python.
-    """
-
-@final
-class SupervisedOncePortReceiver(PortReceiverBase):
-    """A monitored once receiver to which PythonMessages are sent.
-    Values of this type cannot be constructed directly in Python.
-    """
-
 @final
 class ActorSupervisionEvent:
     @property

python/monarch/_rust_bindings/monarch_hyperactor/pytokio.pyi

@@ -7,7 +7,18 @@
 # pyre-strict
 
 import asyncio
-from typing import Any, Awaitable, Callable, Coroutine, Generator, Generic, TypeVar
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Generator,
+    Generic,
+    NoReturn,
+    Sequence,
+    Tuple,
+    TypeVar,
+)
 
 T = TypeVar("T")
 
@@ -62,6 +73,15 @@ class PythonTask(Generic[T], Awaitable[T]):
         """
         ...
 
+    @staticmethod
+    def select_one(
+        tasks: "Sequence[PythonTask[T]]",
+    ) -> "PythonTask[Tuple[T, int]]":
+        """
+        Run the tasks concurrently and return the first one to finish along with the index of which task it was.
+        """
+        ...
+
 class Shared(Generic[T]):
     """
     The result of a spawned PythonTask, which can be awaited on multiple times like Python Futures.