Mesh slice supervision (meta-pytorch#824)

Summary:

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

@@ -53,6 +53,7 @@ use crate::comm::multicast::CastMessage;
 use crate::comm::multicast::CastMessageEnvelope;
 use crate::comm::multicast::Uslice;
 use crate::metrics;
+use crate::proc_mesh::ActorEventRouter;
 use crate::proc_mesh::ProcMesh;
 use crate::reference::ActorMeshId;
 use crate::reference::ActorMeshRef;
@@ -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>>,
+    actor_event_router: ActorEventRouter,
 }
 
 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>>,
+        actor_event_router: ActorEventRouter,
     ) -> Self {
         Self {
             proc_mesh: ProcMeshRef::Borrowed(proc_mesh),
             name,
             ranks,
             actor_supervision_rx: Some(actor_supervision_rx),
+            actor_event_router,
         }
     }
 
@@ -270,11 +274,13 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
         name: String,
         actor_supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
         ranks: Vec<ActorRef<A>>,
+        actor_event_router: ActorEventRouter,
     ) -> Self {
         Self {
             proc_mesh: ProcMeshRef::Shared(Box::new(proc_mesh)),
             name,
             ranks,
+            actor_event_router,
             actor_supervision_rx: Some(actor_supervision_rx),
         }
     }
@@ -294,6 +300,10 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
                 mesh_id: self.id(),
             })
     }
+
+    pub fn actor_event_router(&self) -> &ActorEventRouter {
+        &self.actor_event_router
+    }
 }
 
 /// 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

@@ -99,7 +99,33 @@ pub fn global_mailbox() -> Mailbox {
         .clone()
 }
 
-type ActorEventRouter = Arc<DashMap<ActorMeshName, mpsc::UnboundedSender<ActorSupervisionEvent>>>;
+#[derive(Clone, Debug)]
+pub struct ActorEventRouter {
+    inner: Arc<DashMap<ActorMeshName, Vec<mpsc::UnboundedSender<ActorSupervisionEvent>>>>,
+}
+
+impl ActorEventRouter {
+    pub fn bind(&self, name: ActorMeshName) -> mpsc::UnboundedReceiver<ActorSupervisionEvent> {
+        let (tx, rx) = mpsc::unbounded_channel();
+        self.inner.entry(name).or_insert(vec![]).push(tx);
+        rx
+    }
+
+    fn new() -> Self {
+        Self {
+            inner: Arc::new(DashMap::new()),
+        }
+    }
+}
+
+impl Deref for ActorEventRouter {
+    type Target = DashMap<ActorMeshName, Vec<mpsc::UnboundedSender<ActorSupervisionEvent>>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
 /// A ProcMesh maintains a mesh of procs whose lifecycles are managed by
 /// an allocator.
 pub struct ProcMesh {
@@ -335,7 +361,7 @@ impl ProcMesh {
                 alloc: Box::new(alloc),
                 supervision_events,
             }),
-            actor_event_router: Arc::new(DashMap::new()),
+            actor_event_router: ActorEventRouter::new(),
             shape,
             ranks: proc_ids
                 .into_iter()
@@ -434,16 +460,13 @@ impl ProcMesh {
     where
         A::Params: RemoteMessage,
     {
-        let (tx, rx) = mpsc::unbounded_channel::<ActorSupervisionEvent>();
-        {
-            // Instantiate supervision routing BEFORE spawning the actor mesh.
-            self.actor_event_router.insert(actor_name.to_string(), tx);
-        }
+        let rx = self.actor_event_router.bind(actor_name.to_string());
         let root_mesh = RootActorMesh::new(
             self,
             actor_name.to_string(),
             rx,
             Self::spawn_on_procs::<A>(&self.client, self.agents(), actor_name, params).await?,
+            self.actor_event_router.clone(),
         );
         Ok(root_mesh)
     }
@@ -515,10 +538,17 @@ impl ProcMesh {
                 }
             }
         }
+        self.actor_event_router.remove(&mesh_name.to_string());
         Ok(())
     }
 }
 
+impl Drop for ProcMesh {
+    fn drop(&mut self) {
+        self.actor_event_router.clear();
+    }
+}
+
 /// Proc lifecycle events.
 #[derive(Debug, Clone)]
 pub enum ProcEvent {
@@ -596,8 +626,10 @@ 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());
+                        for tx in entry.value().iter() {
+                            if tx.send(event.clone()).is_err() {
+                                tracing::warn!("unable to transmit supervision event to actor {}", entry.key());
+                            }
                         }
                     }
 
@@ -631,9 +663,11 @@ 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);
+                        if let Some(txs) = self.actor_event_router.get(actor_id.name()) {
+                            for tx in txs.iter() {
+                                if tx.send(event.clone()).is_err() {
+                                    tracing::warn!("unable to transmit supervision event to actor {}", actor_id);
+                                }
                             }
                         } else {
                             tracing::warn!("received supervision event for unregistered actor {}", actor_id);
@@ -683,18 +717,16 @@ impl<D: Deref<Target = ProcMesh> + Send + Sync + 'static> SharedSpawnable for D
     where
         A::Params: RemoteMessage,
     {
-        let (tx, rx) = mpsc::unbounded_channel::<ActorSupervisionEvent>();
-        {
-            // Instantiate supervision routing BEFORE spawning the actor mesh.
-            self.actor_event_router.insert(actor_name.to_string(), tx);
-        }
+        let rx = self.actor_event_router.bind(actor_name.to_string());
         let ranks =
             ProcMesh::spawn_on_procs::<A>(&self.client, self.agents(), actor_name, params).await?;
+        let actor_event_router = self.actor_event_router.clone();
         Ok(RootActorMesh::new_shared(
             self,
             actor_name.to_string(),
             rx,
             ranks,
+            actor_event_router,
         ))
     }
 }