Python actor mesh supervision support (#522)

Summary:
Pull Request resolved: #522

This diff exposes the Rust ActorMesh supervision API to Python ActorMesh. It also wires the supervision events to endpoint calls, including *call()/call_one()/choose()/stream()*. So when a supervision error happens, all inflight calls will get notified, new calls will be failed.

The current diff fails the whole mesh when any actor fails in the mesh. A followup diff will provide more granular management here, such that we may only fail the calls that actually expect reply from the failed actors.

Reviewed By: colin2328

Differential Revision: D77434080

fbshipit-source-id: 654528b7882c1f5f7762a1059d52764a5282859a

Author: moonli

hyperactor_mesh/src/actor_mesh.rs

@@ -179,7 +179,9 @@ pub struct RootActorMesh<'a, A: RemoteActor> {
     proc_mesh: ProcMeshRef<'a>,
     name: String,
     pub(crate) ranks: Vec<ActorRef<A>>, // temporary until we remove `ArcActorMesh`.
-    actor_supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
+    // The receiver of supervision events. It is None if it has been transferred to
+    // an actor event observer.
+    actor_supervision_rx: Option<mpsc::UnboundedReceiver<ActorSupervisionEvent>>,
 }
 
 impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
@@ -193,7 +195,7 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
             proc_mesh: ProcMeshRef::Borrowed(proc_mesh),
             name,
             ranks,
-            actor_supervision_rx,
+            actor_supervision_rx: Some(actor_supervision_rx),
         }
     }
 
@@ -207,7 +209,7 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
             proc_mesh: ProcMeshRef::Shared(Box::new(proc_mesh)),
             name,
             ranks,
-            actor_supervision_rx,
+            actor_supervision_rx: Some(actor_supervision_rx),
         }
     }
 
@@ -216,21 +218,35 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
         self.proc_mesh.client().open_port()
     }
 
-    /// An event stream of proc events. Each ProcMesh can produce only one such
+    /// An event stream of actor events. Each RootActorMesh can produce only one such
     /// stream, returning None after the first call.
+    pub fn events(&mut self) -> Option<ActorSupervisionEvents> {
+        self.actor_supervision_rx
+            .take()
+            .map(|actor_supervision_rx| ActorSupervisionEvents {
+                actor_supervision_rx,
+                mesh_id: self.id(),
+            })
+    }
+}
+
+/// Supervision event stream for actor mesh. It emits actor supervision events.
+pub struct ActorSupervisionEvents {
+    // The receiver of supervision events from proc mesh.
+    actor_supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
+    // The name of the actor mesh.
+    mesh_id: ActorMeshId,
+}
+
+impl ActorSupervisionEvents {
     pub async fn next(&mut self) -> Option<ActorSupervisionEvent> {
         let result = self.actor_supervision_rx.recv().await;
-        match result.as_ref() {
-            Some(event) => {
-                tracing::debug!("Received supervision event: {event:?}");
-            }
-            None => {
-                tracing::info!(
-                    "Supervision stream for actor mesh {} was closed!",
-                    self.name
-                );
-            }
-        };
+        if result.is_none() {
+            tracing::info!(
+                "supervision stream for actor mesh {:?} was closed!",
+                self.mesh_id
+            );
+        }
         result
     }
 }

hyperactor_mesh/src/proc_mesh.rs

@@ -20,6 +20,7 @@ use hyperactor::Mailbox;
 use hyperactor::Named;
 use hyperactor::RemoteMessage;
 use hyperactor::WorldId;
+use hyperactor::actor::ActorStatus;
 use hyperactor::actor::RemoteActor;
 use hyperactor::actor::remote::Remote;
 use hyperactor::cap;
@@ -451,7 +452,7 @@ impl ProcMesh {
 }
 
 /// Proc lifecycle events.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum ProcEvent {
     /// The proc of the given rank was stopped with the provided reason.
     Stopped(usize, ProcStopReason),
@@ -506,6 +507,21 @@ impl ProcEvents {
                         continue;
                     };
 
+                    // Need to send this event to actor meshes to notify them of the proc's death.
+                    // TODO(albertli): only send this event to all root actor meshes if any of them use this proc.
+                    for entry in self.actor_event_router.iter() {
+                        // Make a dummy actor supervision event.
+                        let client_proc = ProcId(WorldId(format!("{}_manager", self.event_state.alloc.world_id().name())), 0);
+                        let client = client_proc.actor_id("client", 0);
+                        let event = ActorSupervisionEvent::new(
+                            client.clone(),
+                            ActorStatus::Failed(format!("proc {} is stopped", proc_id))
+                        );
+                        if entry.value().send(event).is_err() {
+                            tracing::warn!("unable to transmit supervision event to actor {}", client);
+                        }
+                    }
+
                     break Some(ProcEvent::Stopped(*rank, reason));
                 }
                 Ok(event) = self.event_state.supervision_events.recv() => {
@@ -521,17 +537,15 @@ impl ProcEvents {
                     };
                     // transmit to the correct root actor mesh.
                     {
-                        let Some(tx) = self.actor_event_router.get(actor_id.name()) else {
+                        if let Some(tx) = self.actor_event_router.get(actor_id.name()) {
+                            if tx.send(event).is_err() {
+                                tracing::warn!("unable to transmit supervision event to actor {}", actor_id);
+                            }
+                        } else {
                             tracing::warn!("received supervision event for unregistered actor {}", actor_id);
-                            continue;
-                        };
-                        let Ok(_) = tx.send(event) else {
-                            tracing::warn!("unable to transmit supervision event to actor {}", actor_id);
-                            continue;
-                        };
+                        }
                     }
-                    // TODO: Actor supervision events need to be wired to the frontend.
-                    // TODO: This event should be handled by the proc mesh if unhandled by actor mesh.
+                    // Send this event to Python proc mesh to keep its health status up to date.
                     break Some(ProcEvent::Crashed(*rank, actor_status.to_string()))
                 }
             }
@@ -738,6 +752,7 @@ mod tests {
         let mut events = mesh.events().unwrap();
 
         let mut actors = mesh.spawn::<TestActor>("failing", &()).await.unwrap();
+        let mut actor_events = actors.events().unwrap();
 
         actors
             .cast(
@@ -751,7 +766,7 @@ mod tests {
             ProcEvent::Crashed(0, reason) if reason.contains("failmonkey")
         );
 
-        let event = actors.next().await.unwrap();
+        let mut event = actor_events.next().await.unwrap();
         assert_matches!(event.actor_status(), ActorStatus::Failed(_));
         assert_eq!(event.actor_id().1, "failing".to_string());
         assert_eq!(event.actor_id().2, 0);
@@ -767,7 +782,9 @@ mod tests {
         );
 
         assert!(events.next().await.is_none());
-        assert!(actors.next().await.is_none());
+        event = actor_events.next().await.unwrap();
+        assert_matches!(event.actor_status(), ActorStatus::Failed(_));
+        assert_eq!(event.actor_id().2, 0);
     }
 
     #[timed_test::async_timed_test(timeout_secs = 5)]

hyperactor_mesh/src/shared_cell.rs

@@ -156,6 +156,16 @@ impl<T> SharedCell<T> {
         SharedCellRef::from(self.inner.clone().try_read_owned()?)
     }
 
+    /// Execute given closure with write access to the underlying data. If the cell is empty, returns an error.
+    pub async fn with_mut<F, R>(&self, f: F) -> Result<R, EmptyCellError>
+    where
+        F: FnOnce(&mut T) -> R,
+    {
+        let mut inner = self.inner.write(true).await;
+        let value = inner.value.as_mut().ok_or(EmptyCellError {})?;
+        Ok(f(value))
+    }
+
     /// Take the item out of the cell, leaving it in an unusable state.
     pub async fn take(&self) -> Result<T, EmptyCellError> {
         let mut inner = self.inner.write(true).await;

hyperactor_telemetry/src/lib.rs

@@ -505,13 +505,13 @@ pub fn initialize_logging(clock: impl TelemetryClock + Send + 'static) {
             } else {
                 None
             })
+            .with(file_layer)
+            .with(stderr_layer)
             .with(if is_layer_enabled(DISABLE_RECORDER_TRACING) {
                 Some(recorder().layer())
             } else {
                 None
             })
-            .with(file_layer)
-            .with(stderr_layer)
             .try_init()
         {
             tracing::debug!("logging already initialized for this process: {}", err);

monarch_extension/src/lib.rs

@@ -79,6 +79,11 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
         "monarch_hyperactor.selection",
     )?)?;
 
+    monarch_hyperactor::supervision::register_python_bindings(&get_or_add_new_module(
+        module,
+        "monarch_hyperactor.supervision",
+    )?)?;
+
     #[cfg(feature = "tensor_engine")]
     {
         client::register_python_bindings(&get_or_add_new_module(