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

monarch_hyperactor/src/actor_mesh.rs

@@ -21,6 +21,7 @@ use hyperactor_mesh::Mesh;
 use hyperactor_mesh::RootActorMesh;
 use hyperactor_mesh::actor_mesh::ActorMesh;
 use hyperactor_mesh::actor_mesh::ActorSupervisionEvents;
+use hyperactor_mesh::dashmap::DashMap;
 use hyperactor_mesh::reference::ActorMeshRef;
 use hyperactor_mesh::sel;
 use hyperactor_mesh::shared_cell::SharedCell;
@@ -171,6 +172,8 @@ pub(crate) struct PythonActorMeshImpl {
     unhealthy_event: Arc<std::sync::Mutex<Unhealthy<ActorSupervisionEvent>>>,
     user_monitor_sender: tokio::sync::broadcast::Sender<Option<ActorSupervisionEvent>>,
     monitor: tokio::task::JoinHandle<()>,
+    /// Needed for slices to track health as ActorMeshes only track the latest supervision event
+    crashed_ranks: Arc<DashMap<usize, ActorSupervisionEvent>>,
 }
 
 impl PythonActorMeshImpl {
@@ -185,10 +188,12 @@ impl PythonActorMeshImpl {
         let (user_monitor_sender, _) =
             tokio::sync::broadcast::channel::<Option<ActorSupervisionEvent>>(1);
         let unhealthy_event = Arc::new(std::sync::Mutex::new(Unhealthy::SoFarSoGood));
-        let monitor = tokio::spawn(PythonActorMeshImpl::actor_mesh_monitor(
+        let crashed_ranks = Arc::new(DashMap::new());
+        let monitor = tokio::spawn(Self::actor_mesh_monitor(
             events,
             user_monitor_sender.clone(),
             Arc::clone(&unhealthy_event),
+            crashed_ranks.clone(),
         ));
         PythonActorMeshImpl {
             inner,
@@ -197,6 +202,7 @@ impl PythonActorMeshImpl {
             unhealthy_event,
             user_monitor_sender,
             monitor,
+            crashed_ranks,
         }
     }
     /// Monitor of the actor mesh. It processes supervision errors for the mesh, and keeps mesh
@@ -205,14 +211,19 @@ impl PythonActorMeshImpl {
         mut events: ActorSupervisionEvents,
         user_sender: tokio::sync::broadcast::Sender<Option<ActorSupervisionEvent>>,
         unhealthy_event: Arc<std::sync::Mutex<Unhealthy<ActorSupervisionEvent>>>,
+        crashed_ranks: Arc<DashMap<usize, ActorSupervisionEvent>>,
     ) {
         loop {
             let event = events.next().await;
             tracing::debug!("actor_mesh_monitor received supervision event: {event:?}");
             let mut inner_unhealthy_event = unhealthy_event.lock().unwrap();
             match &event {
                 None => *inner_unhealthy_event = Unhealthy::StreamClosed,
-                Some(event) => *inner_unhealthy_event = Unhealthy::Crashed(event.clone()),
+                Some(event) => {
+                    // Ignore if the crashed actor is not a part of the mesh.
+                    crashed_ranks.insert(event.actor_id.rank(), event.clone());
+                    *inner_unhealthy_event = Unhealthy::Crashed(event.clone())
+                }
             }
 
             // Ignore the sender error when there is no receiver,
@@ -236,7 +247,15 @@ impl PythonActorMeshImpl {
 
     fn bind(&self) -> PyResult<PythonActorMeshRef> {
         let mesh = self.try_inner()?;
-        Ok(PythonActorMeshRef { inner: mesh.bind() })
+        let root_health_state = Some(RootHealthState {
+            user_monitor_sender: self.user_monitor_sender.clone(),
+            unhealthy_event: Arc::clone(&self.unhealthy_event),
+            crashed_ranks: Arc::clone(&self.crashed_ranks),
+        });
+        Ok(PythonActorMeshRef {
+            inner: mesh.bind(),
+            root_health_state,
+        })
     }
 }
 
@@ -351,13 +370,59 @@ impl PythonActorMeshImpl {
     }
 }
 
+#[derive(Clone, Debug)]
+struct RootHealthState {
+    user_monitor_sender: tokio::sync::broadcast::Sender<Option<ActorSupervisionEvent>>,
+    unhealthy_event: Arc<std::sync::Mutex<Unhealthy<ActorSupervisionEvent>>>,
+    crashed_ranks: Arc<DashMap<usize, ActorSupervisionEvent>>,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 struct PythonActorMeshRef {
     inner: ActorMeshRef<PythonActor>,
+    #[serde(skip)]
+    // If the reference has been serialized and sent over the wire
+    // we no longer have access to the underlying mesh's state
+    root_health_state: Option<RootHealthState>,
 }
 
 impl ActorMeshProtocol for PythonActorMeshRef {
     fn cast(&self, message: PythonMessage, selection: Selection, client: Mailbox) -> PyResult<()> {
+        match &self.root_health_state {
+            Some(RootHealthState {
+                crashed_ranks,
+                unhealthy_event,
+                ..
+            }) => {
+                // iterate through all crashed ranks in the root mesh and take first rank
+                // that is in the sliced mesh
+                match self
+                    .inner
+                    .shape()
+                    .slice()
+                    .iter()
+                    .find_map(|rank| crashed_ranks.get(&rank).map(|entry| entry.value().clone()))
+                {
+                    Some(event) => {
+                        return Err(SupervisionError::new_err(format!(
+                            "Actor {:?} is unhealthy with reason: {}",
+                            event.actor_id, event.actor_status
+                        )));
+                    }
+                    None => {
+                        if matches!(
+                            &*unhealthy_event.lock().unwrap_or_else(|e| e.into_inner()),
+                            Unhealthy::StreamClosed
+                        ) {
+                            return Err(SupervisionError::new_err(
+                                "actor mesh is stopped due to proc mesh shutdown".to_string(),
+                            ));
+                        }
+                    }
+                }
+            }
+            None => (),
+        }
         self.inner
             .cast(&client, selection, message.clone())
             .map_err(|err| PyException::new_err(err.to_string()))?;
@@ -369,9 +434,41 @@ impl ActorMeshProtocol for PythonActorMeshRef {
             .inner
             .new_with_shape(shape.get_inner().clone())
             .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?;
-        Ok(Box::new(Self { inner: sliced }))
+        Ok(Box::new(Self {
+            inner: sliced,
+            root_health_state: self.root_health_state.clone(),
+        }))
     }
 
+    fn supervision_event(&self) -> PyResult<Option<PyShared>> {
+        match &self.root_health_state {
+            Some(RootHealthState {
+                user_monitor_sender,
+                ..
+            }) => {
+                let mut receiver = 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(),
+                        },
+                    };
+                    Ok(PyErr::new::<SupervisionError, _>(format!(
+                        "Actor {:?} exited because of the following reason: {}",
+                        event.actor_id, event.actor_status
+                    )))
+                })
+                .map(|mut x| x.spawn().map(Some))?
+            }
+            None => Ok(None),
+        }
+    }
     fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
         let bytes =
             bincode::serialize(self).map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?;
@@ -504,7 +601,7 @@ impl ActorMeshProtocol for AsyncActorMesh {
         let mesh = self.mesh.clone();
         Ok(Box::new(AsyncActorMesh::new(
             self.queue.clone(),
-            false,
+            self.supervised,
             async { Ok(mesh.await?.new_with_shape(shape)?) },
         )))
     }

python/tests/test_actor_error.py

@@ -685,3 +685,61 @@ async def test_supervision_with_sending_error():
         await actor_mesh.check.call()
     with pytest.raises(SupervisionError, match="Actor .* is unhealthy with reason"):
         await actor_mesh.check_with_payload.call(payload="a")
+
+
+async def test_slice_supervision() -> None:
+    pm = await proc_mesh(gpus=4)
+    healthy_mesh = await pm.spawn("healthy", HealthyActor)
+    error_mesh = await pm.spawn("error", ErrorActor)
+    slice_1 = error_mesh.slice(gpus=slice(2, 4))
+    slice_2 = error_mesh.slice(gpus=2)
+    slice_3 = error_mesh.slice(gpus=3)
+
+    # Trigger supervision error on gpus=3
+    with pytest.raises(SupervisionError, match="did not handle supervision event"):
+        await slice_3.fail_with_supervision_error.call()
+
+    # Mesh containing all gpus is unhealthy
+    with pytest.raises(SupervisionError, match="Actor .* is unhealthy with reason:"):
+        await error_mesh.check.call()
+
+    # Slice containing only gpus=3 is unhealthy
+    with pytest.raises(SupervisionError, match="Actor .* is unhealthy with reason:"):
+        await slice_3.check.call()
+
+    # Slice containing gpus=3 is unhealthy
+    with pytest.raises(SupervisionError, match="Actor .* is unhealthy with reason:"):
+        await slice_1.check.call()
+
+    # Slice not containing gpus=3 is healthy
+    check = await slice_2.check.call()
+    for _, item in check.items():
+        assert item == "this is a healthy check"
+
+    # Other actor mesh on the same proc mesh is healthy
+    check = await healthy_mesh.check.call()
+    for _, item in check.items():
+        assert item == "this is a healthy check"
+
+
+async def test_mesh_slices_inherit_parent_errors() -> None:
+    pm = await proc_mesh(gpus=4)
+    error_mesh = await pm.spawn("error", ErrorActor)
+    slice_1 = error_mesh.slice(gpus=slice(2, 4))
+
+    # Trigger supervision error on gpus=2, 3, 4
+    with pytest.raises(SupervisionError):
+        await slice_1.fail_with_supervision_error.call()
+
+    # Newly created slice containing gpu=3 is unhealthy
+    slice_2 = error_mesh.slice(gpus=3)
+    with pytest.raises(SupervisionError):
+        await slice_2.check.call()
+
+    # Newly created slice containing gpu=1 is healthy
+    slice_3 = error_mesh.slice(gpus=1)
+    check = await slice_3.check.call()
+    for _, item in check.items():
+        assert item == "this is a healthy check"
+
+    await pm.stop()