Mesh slice supervision (#824)

Summary:
Pull Request resolved: #824

Suppose we have a ProcMesh with 4 gpus. On this mesh we spawn a ActorMesh A, and an ActorMesh B. We create a slice of ActorMesh A SliceA_1 containing only gpu 4 and a slice SliceA_2 containing only gpu 1. If the Actor A on gpu 4 crashes we should have the following health states:
- ActorMesh A is unhealthy (contains A gpu=4)
- SliceA_1 is unhealthy (contains A gpu=4)
- SliceA_2 is healthy (does not contain A gpu=4)
- ActorMesh B is healthy (contains gpu=4 but not Actor A)

Implementation:
1. All supervision event streams are created when a RootActorMesh is spawned. A tx-rx pair are created and the tx is inserted into a map acting as a router.
2. The router now holds onto a Vec of senders instead of a single sender. For each Actor mesh in the router, we can call bind to create another tx-rx pair. The router will manage the tx, sending a message using every tx for a given Actor mesh name every time there is a supervision event. The rx is returned to be used by mesh slices.
3. The spawned RootActorMesh gets a copy of the Arc holding the router so that mesh slices can bind to it
4. PythonActorMeshes contain a monitor which is just a loop that listens from the next supervision event from a stream. If a ::Crashed event comes in, we will update an Arc<> keeping track of the health state. This monitor will now also take in the shape of the mesh it is monitoring, and only update the health state to ::Crashed, if the crashed Actor is within the shape.
5. When a PythonActorMesh is sliced, a PythonActorMeshRef is created. We will add a monitor to PythonActorMeshRefs. It is an Option<> because if it is ever serialized and deserialized, we can no longer monitor it
6. When we cast to an PythonActorMeshRef, we will first check the health state and return a SupervisionError if the mesh is unhealthy

Differential Revision: D79821712

Author: thomasywang

hyperactor_mesh/src/actor_mesh.rs

@@ -54,6 +54,7 @@ use crate::comm::multicast::CastMessageEnvelope;
 use crate::comm::multicast::Uslice;
 use crate::metrics;
 use crate::proc_mesh::ProcMesh;
+use crate::proc_mesh::SupervisionEventBroadcaster;
 use crate::reference::ActorMeshId;
 use crate::reference::ActorMeshRef;
 use crate::reference::ProcMeshId;
@@ -248,6 +249,7 @@ pub struct RootActorMesh<'a, A: RemoteActor> {
     // 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>>,
+    supervision_event_broadcaster: SupervisionEventBroadcaster,
 }
 
 impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
@@ -256,12 +258,14 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
         name: String,
         actor_supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
         ranks: Vec<ActorRef<A>>,
+        supervision_event_broadcaster: SupervisionEventBroadcaster,
     ) -> Self {
         Self {
             proc_mesh: ProcMeshRef::Borrowed(proc_mesh),
             name,
             ranks,
             actor_supervision_rx: Some(actor_supervision_rx),
+            supervision_event_broadcaster,
         }
     }
 
@@ -270,11 +274,13 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
         name: String,
         actor_supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
         ranks: Vec<ActorRef<A>>,
+        supervision_event_broadcaster: SupervisionEventBroadcaster,
     ) -> Self {
         Self {
             proc_mesh: ProcMeshRef::Shared(Box::new(proc_mesh)),
             name,
             ranks,
+            supervision_event_broadcaster,
             actor_supervision_rx: Some(actor_supervision_rx),
         }
     }
@@ -294,6 +300,10 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
                 mesh_id: self.id(),
             })
     }
+
+    pub fn supervision_event_broadcaster(&self) -> SupervisionEventBroadcaster {
+        self.supervision_event_broadcaster.clone()
+    }
 }
 
 /// Supervision event stream for actor mesh. It emits actor supervision events.
@@ -315,6 +325,16 @@ impl ActorSupervisionEvents {
         }
         result
     }
+
+    pub fn new(
+        actor_supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
+        mesh_id: ActorMeshId,
+    ) -> Self {
+        Self {
+            actor_supervision_rx,
+            mesh_id,
+        }
+    }
 }
 
 #[async_trait]

hyperactor_mesh/src/proc_mesh.rs

@@ -101,7 +101,49 @@ pub fn global_root_client() -> &'static Instance<()> {
     instance
 }
 
-type ActorEventRouter = Arc<DashMap<ActorMeshName, mpsc::UnboundedSender<ActorSupervisionEvent>>>;
+type ActorEventRouter = Arc<DashMap<ActorMeshName, SupervisionEventBroadcaster>>;
+
+/// Wrapper class around a Vec<mpsc::UnboundedSender<ActorSupervisionEvent>>>>
+/// allowing for new listeners to be added and events to be broadcasted
+#[derive(Clone, Debug)]
+pub struct SupervisionEventBroadcaster {
+    inner: Arc<std::sync::Mutex<Vec<mpsc::UnboundedSender<ActorSupervisionEvent>>>>,
+}
+
+impl Deref for SupervisionEventBroadcaster {
+    type Target = Arc<std::sync::Mutex<Vec<mpsc::UnboundedSender<ActorSupervisionEvent>>>>;
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl SupervisionEventBroadcaster {
+    fn new() -> Self {
+        Self {
+            inner: Arc::new(std::sync::Mutex::new(Vec::new())),
+        }
+    }
+    pub fn subscribe(&self) -> mpsc::UnboundedReceiver<ActorSupervisionEvent> {
+        let (tx, rx) = mpsc::unbounded_channel();
+        {
+            self.inner
+                .lock()
+                .unwrap_or_else(|err| err.into_inner())
+                .push(tx);
+        }
+        rx
+    }
+
+    fn broadcast(&self, event: ActorSupervisionEvent) {
+        let mut guard = self.inner.lock().unwrap_or_else(|err| err.into_inner());
+        for tx in guard.iter_mut() {
+            tx.send(event.clone()).unwrap_or_else(|err| {
+                tracing::warn!("failed to send supervision event: {err:?}");
+            });
+        }
+    }
+}
+
 /// A ProcMesh maintains a mesh of procs whose lifecycles are managed by
 /// an allocator.
 pub struct ProcMesh {
@@ -436,16 +478,20 @@ impl ProcMesh {
     where
         A::Params: RemoteMessage,
     {
-        let (tx, rx) = mpsc::unbounded_channel::<ActorSupervisionEvent>();
+        let supervision_event_broadcaster = SupervisionEventBroadcaster::new();
         {
             // Instantiate supervision routing BEFORE spawning the actor mesh.
-            self.actor_event_router.insert(actor_name.to_string(), tx);
+            self.actor_event_router.insert(
+                actor_name.to_string(),
+                supervision_event_broadcaster.clone(),
+            );
         }
         let root_mesh = RootActorMesh::new(
             self,
             actor_name.to_string(),
-            rx,
+            supervision_event_broadcaster.subscribe(),
             Self::spawn_on_procs::<A>(&self.client, self.agents(), actor_name, params).await?,
+            supervision_event_broadcaster,
         );
         Ok(root_mesh)
     }
@@ -517,6 +563,7 @@ impl ProcMesh {
                 }
             }
         }
+        self.actor_event_router.remove(&mesh_name.to_string());
         Ok(())
     }
 }
@@ -598,9 +645,7 @@ impl ProcEvents {
                             message_headers: None,
                             caused_by: None,
                         };
-                        if entry.value().send(event).is_err() {
-                            tracing::warn!("unable to transmit supervision event to actor {}", entry.key());
-                        }
+                        entry.value().broadcast(event);
                     }
 
                     break Some(ProcEvent::Stopped(*rank, reason));
@@ -634,9 +679,7 @@ impl ProcEvents {
                     // transmit to the correct root actor mesh.
                     {
                         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);
-                            }
+                            tx.broadcast(event.clone());
                         } else {
                             tracing::warn!("received supervision event for unregistered actor {}", actor_id);
                         }
@@ -685,18 +728,22 @@ impl<D: Deref<Target = ProcMesh> + Send + Sync + 'static> SharedSpawnable for D
     where
         A::Params: RemoteMessage,
     {
-        let (tx, rx) = mpsc::unbounded_channel::<ActorSupervisionEvent>();
+        let supervision_event_broadcaster = SupervisionEventBroadcaster::new();
         {
             // Instantiate supervision routing BEFORE spawning the actor mesh.
-            self.actor_event_router.insert(actor_name.to_string(), tx);
+            self.actor_event_router.insert(
+                actor_name.to_string(),
+                supervision_event_broadcaster.clone(),
+            );
         }
         let ranks =
             ProcMesh::spawn_on_procs::<A>(&self.client, self.agents(), actor_name, params).await?;
         Ok(RootActorMesh::new_shared(
             self,
             actor_name.to_string(),
-            rx,
+            supervision_event_broadcaster.subscribe(),
             ranks,
+            supervision_event_broadcaster,
         ))
     }
 }