Add supervision_event monitor to v1 ActorMeshRef (#1454)

Summary:
Pull Request resolved: #1454

Part of #1209

Implement `supervision_event` for the v1 PythonActorMesh and References, for parity
with the v0 API.
This does not yet implement the goal of calling `__supervise__` on the owning PythonActor,
but that'll be the next step.

For parity with v0, had to add a few things:
* For the "monitor", poll the `actor_states` API every few seconds. If any Actor has
collected supervision events, surface those as a SupervisionError to python.
This should be optimized into a push-based API instead, or one based on accumulation
since we only need the first event.
* Before doing a `cast`, have to check if any previous cast had encountered a
supervision error, to avoid sending an undeliverable message
* Disallow spawning new actors on a ProcMeshAgent which has had a supervision event
in its past
* If there's an error that kills the whole proc (process exit, signal, etc.), the agent itself will
be down. Augment the `actor_states` API to handle this by reporting stopped for all remaining
ranks.

Reviewed By: colin2328, pablorfb-meta

Differential Revision: D82494380

fbshipit-source-id: 789586bd89e1d1164438509912b0a565f5174c36

Author: dulinriley

hyperactor_mesh/src/proc_mesh.rs

@@ -59,7 +59,6 @@ use tracing::span;
 
 use crate::CommActor;
 use crate::Mesh;
-use crate::actor_mesh::CAST_ACTOR_MESH_ID;
 use crate::actor_mesh::RootActorMesh;
 use crate::alloc::Alloc;
 use crate::alloc::AllocExt;
@@ -73,7 +72,7 @@ use crate::proc_mesh::mesh_agent::GspawnResult;
 use crate::proc_mesh::mesh_agent::MeshAgentMessageClient;
 use crate::proc_mesh::mesh_agent::ProcMeshAgent;
 use crate::proc_mesh::mesh_agent::StopActorResult;
-use crate::reference::ActorMeshId;
+use crate::proc_mesh::mesh_agent::update_event_actor_id;
 use crate::reference::ProcMeshId;
 use crate::router;
 use crate::shortuuid::ShortUuid;
@@ -770,7 +769,7 @@ impl ProcEvents {
                 // rewriting `actor_id` to a synthetic mesh-level id
                 // so that routing reaches the correct `ActorMesh`
                 // subscribers.
-                Ok(mut event) = self.event_state.supervision_events.recv() => {
+                Ok(event) = self.event_state.supervision_events.recv() => {
                     let had_headers = event.message_headers.is_some();
                     tracing::info!(
                         name = SupervisionEventState::SupervisionEventReceived.as_ref(),
@@ -782,33 +781,7 @@ impl ProcEvents {
                     tracing::debug!(?event, "proc supervision: full event");
 
                     // Normalize events that came via the comm tree.
-                    if let Some(headers) = &event.message_headers {
-                        if let Some(actor_mesh_id) = headers.get(CAST_ACTOR_MESH_ID) {
-                            match actor_mesh_id {
-                                ActorMeshId::V0(proc_mesh_id, actor_name) => {
-                                    let old_actor = event.actor_id.clone();
-                                    event.actor_id = ActorId(
-                                        ProcId::Ranked(WorldId(proc_mesh_id.0.clone()), 0),
-                                        actor_name.clone(),
-                                        0,
-                                    );
-                                    tracing::debug!(
-                                        actor_id = %old_actor,
-                                        "proc supervision: remapped comm-actor id to mesh id from CAST_ACTOR_MESH_ID {}", event.actor_id
-                                    );
-                                }
-                                ActorMeshId::V1(_) => {
-                                    tracing::debug!("proc supervision: headers present but V1 ActorMeshId; leaving actor_id unchanged");
-                                }
-                            }
-                        } else {
-                            tracing::debug!(
-                                "proc supervision: headers present but no CAST_ACTOR_MESH_ID; leaving actor_id unchanged"
-                            );
-                        }
-                    } else {
-                        tracing::debug!("proc supervision: no headers attached; leaving actor_id unchanged");
-                    }
+                    let event = update_event_actor_id(event);
 
                     // Forward the supervision event to the ActorMesh (keyed by its mesh name)
                     // that registered for events in this ProcMesh. The routing table

hyperactor_mesh/src/proc_mesh/mesh_agent.rs

@@ -34,6 +34,7 @@ use hyperactor::PortRef;
 use hyperactor::ProcId;
 use hyperactor::RefClient;
 use hyperactor::Unbind;
+use hyperactor::WorldId;
 use hyperactor::actor::ActorStatus;
 use hyperactor::actor::remote::Remote;
 use hyperactor::channel;
@@ -52,7 +53,9 @@ use hyperactor::supervision::ActorSupervisionEvent;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::actor_mesh::CAST_ACTOR_MESH_ID;
 use crate::proc_mesh::SupervisionEventState;
+use crate::reference::ActorMeshId;
 use crate::resource;
 use crate::v1::Name;
 
@@ -171,6 +174,41 @@ struct ActorInstanceState {
     spawn: Result<ActorId, anyhow::Error>,
 }
 
+/// Normalize events that came via the comm tree. Updates their actor id based on
+/// the message headers for the event.
+pub(crate) fn update_event_actor_id(mut event: ActorSupervisionEvent) -> ActorSupervisionEvent {
+    if let Some(headers) = &event.message_headers {
+        if let Some(actor_mesh_id) = headers.get(CAST_ACTOR_MESH_ID) {
+            match actor_mesh_id {
+                ActorMeshId::V0(proc_mesh_id, actor_name) => {
+                    let old_actor = event.actor_id.clone();
+                    event.actor_id = ActorId(
+                        ProcId::Ranked(WorldId(proc_mesh_id.0.clone()), 0),
+                        actor_name.clone(),
+                        0,
+                    );
+                    tracing::debug!(
+                        actor_id = %old_actor,
+                        "proc supervision: remapped comm-actor id to mesh id from CAST_ACTOR_MESH_ID {}", event.actor_id
+                    );
+                }
+                ActorMeshId::V1(_) => {
+                    tracing::debug!(
+                        "proc supervision: headers present but V1 ActorMeshId; leaving actor_id unchanged"
+                    );
+                }
+            }
+        } else {
+            tracing::debug!(
+                "proc supervision: headers present but no CAST_ACTOR_MESH_ID; leaving actor_id unchanged"
+            );
+        }
+    } else {
+        tracing::debug!("proc supervision: no headers attached; leaving actor_id unchanged");
+    }
+    event
+}
+
 /// A mesh agent is responsible for managing procs in a [`ProcMesh`].
 #[derive(Debug)]
 #[hyperactor::export(
@@ -397,8 +435,13 @@ impl Handler<ActorSupervisionEvent> for ProcMeshAgent {
         cx: &Context<Self>,
         event: ActorSupervisionEvent,
     ) -> anyhow::Result<()> {
+        let event = update_event_actor_id(event);
         if self.record_supervision_events {
-            tracing::info!("Received supervision event: {:?}, recording", event);
+            tracing::info!(
+                "Received supervision event on proc {}: {:?}, recording",
+                self.proc.proc_id(),
+                event
+            );
             self.supervision_events
                 .entry(event.actor_id.clone())
                 .or_default()
@@ -457,6 +500,22 @@ impl Handler<resource::CreateOrUpdate<ActorSpec>> for ProcMeshAgent {
             // There is no update.
             return Ok(());
         }
+        let create_rank = create_or_update.rank.unwrap();
+        // If there have been supervision events for any actors on this proc,
+        // we disallow spawning new actors on it, as this proc may be in an
+        // invalid state.
+        if !self.supervision_events.is_empty() {
+            self.actor_states.insert(
+                create_or_update.name.clone(),
+                ActorInstanceState {
+                    spawn: Err(anyhow::anyhow!(
+                        "Cannot spawn new actors on mesh with supervision events"
+                    )),
+                    create_rank,
+                },
+            );
+            return Ok(());
+        }
 
         let ActorSpec {
             actor_type,
@@ -465,7 +524,7 @@ impl Handler<resource::CreateOrUpdate<ActorSpec>> for ProcMeshAgent {
         self.actor_states.insert(
             create_or_update.name.clone(),
             ActorInstanceState {
-                create_rank: create_or_update.rank.unwrap(),
+                create_rank,
                 spawn: self
                     .remote
                     .gspawn(

hyperactor_mesh/src/v1/actor_mesh.rs

@@ -70,6 +70,10 @@ impl<A: Referable> ActorMesh<A> {
             current_ref,
         }
     }
+
+    pub fn name(&self) -> &Name {
+        &self.name
+    }
 }
 
 impl<A: Referable> Deref for ActorMesh<A> {
@@ -491,7 +495,7 @@ mod tests {
     }
 
     #[async_timed_test(timeout_secs = 30)]
-    async fn test_actor_states() {
+    async fn test_actor_states_with_panic() {
         hyperactor_telemetry::initialize_logging_for_test();
 
         let instance = testing::instance().await;
@@ -526,22 +530,20 @@ mod tests {
         // status such that when a process switches to unhealthy it sets a
         // supervision event.
         let supervision_task = tokio::spawn(async move {
-            match actor_mesh.actor_states(&instance).await {
-                Ok(events) => {
-                    for state in events.values() {
-                        supervisor.send(instance, state.clone()).unwrap();
-                    }
-                }
-                Err(e) => {
-                    println!("error: {:?}", e);
-                }
-            };
+            let events = actor_mesh.actor_states(&instance).await.unwrap();
+            for state in events.values() {
+                supervisor.send(instance, state.clone()).unwrap();
+            }
         });
         // Make sure the task completes first without a panic.
         supervision_task.await.unwrap();
 
         for _ in 0..num_replicas {
-            let state = supervision_receiver.recv().await.unwrap();
+            let state = RealClock
+                .timeout(Duration::from_secs(10), supervision_receiver.recv())
+                .await
+                .expect("timeout")
+                .unwrap();
             if let resource::Status::Failed(s) = state.status {
                 assert!(s.contains("supervision events"));
             } else {
@@ -558,6 +560,143 @@ mod tests {
         }
     }
 
+    #[async_timed_test(timeout_secs = 30)]
+    async fn test_actor_states_with_process_exit() {
+        hyperactor_telemetry::initialize_logging_for_test();
+
+        let instance = testing::instance().await;
+        // Listen for supervision events sent to the parent instance.
+        let (supervision_port, mut supervision_receiver) =
+            instance.open_port::<resource::State<ActorState>>();
+        let supervisor = supervision_port.bind();
+        let num_replicas = 4;
+        let meshes = testing::proc_meshes(instance, extent!(replicas = num_replicas)).await;
+        let proc_mesh = &meshes[1];
+        let child_name = Name::new("child");
+
+        let actor_mesh = proc_mesh
+            .spawn_with_name::<testactor::TestActor>(instance, child_name.clone(), &())
+            .await
+            .unwrap();
+
+        actor_mesh
+            .cast(
+                instance,
+                testactor::CauseSupervisionEvent(testactor::SupervisionEventType::ProcessExit(1)),
+            )
+            .unwrap();
+
+        // Wait for the casted message to cause a process exit on all actors.
+        // We can't use a reply port because the handler for the message will
+        // by definition not complete and send a reply.
+        #[allow(clippy::disallowed_methods)]
+        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+
+        // Now that all ranks have completed, set up a continuous poll of the
+        // status such that when a process switches to unhealthy it sets a
+        // supervision event.
+        let supervision_task = tokio::spawn(async move {
+            let events = actor_mesh.actor_states(&instance).await.unwrap();
+            for state in events.values() {
+                supervisor.send(instance, state.clone()).unwrap();
+            }
+        });
+        // Make sure the task completes first without a panic.
+        RealClock
+            .timeout(Duration::from_secs(10), supervision_task)
+            .await
+            .expect("timeout")
+            .unwrap();
+
+        for _ in 0..num_replicas {
+            let state = RealClock
+                .timeout(Duration::from_secs(10), supervision_receiver.recv())
+                .await
+                .expect("timeout")
+                .unwrap();
+            assert_matches!(state.status, resource::Status::Stopped);
+            let events = state
+                .state
+                .expect("state should be present")
+                .supervision_events;
+            assert_eq!(events.len(), 1);
+            assert_eq!(events[0].actor_status, ActorStatus::Stopped);
+        }
+    }
+
+    #[async_timed_test(timeout_secs = 30)]
+    async fn test_actor_states_on_sliced_mesh() {
+        hyperactor_telemetry::initialize_logging_for_test();
+
+        let instance = testing::instance().await;
+        // Listen for supervision events sent to the parent instance.
+        let (supervision_port, mut supervision_receiver) =
+            instance.open_port::<resource::State<ActorState>>();
+        let supervisor = supervision_port.bind();
+        let num_replicas = 4;
+        let meshes = testing::proc_meshes(instance, extent!(replicas = num_replicas)).await;
+        let proc_mesh = &meshes[1];
+        let child_name = Name::new("child");
+
+        let actor_mesh = proc_mesh
+            .spawn_with_name::<testactor::TestActor>(instance, child_name.clone(), &())
+            .await
+            .unwrap();
+        let sliced = actor_mesh
+            .range("replicas", 1..3)
+            .expect("slice should be valid");
+        let sliced_replicas = sliced.len();
+
+        sliced
+            .cast(
+                instance,
+                testactor::CauseSupervisionEvent(testactor::SupervisionEventType::Panic),
+            )
+            .unwrap();
+
+        // Wait for the casted message to cause a process exit on all actors.
+        // We can't use a reply port because the handler for the message will
+        // by definition not complete and send a reply.
+        #[allow(clippy::disallowed_methods)]
+        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+
+        // Now that all ranks have completed, set up a continuous poll of the
+        // status such that when a process switches to unhealthy it sets a
+        // supervision event.
+        let supervision_task = tokio::spawn(async move {
+            let events = sliced.actor_states(&instance).await.unwrap();
+            for state in events.values() {
+                supervisor.send(instance, state.clone()).unwrap();
+            }
+        });
+        // Make sure the task completes first without a panic.
+        RealClock
+            .timeout(Duration::from_secs(10), supervision_task)
+            .await
+            .expect("timeout")
+            .unwrap();
+
+        for _ in 0..sliced_replicas {
+            let state = RealClock
+                .timeout(Duration::from_secs(10), supervision_receiver.recv())
+                .await
+                .expect("timeout")
+                .unwrap();
+            if let resource::Status::Failed(s) = state.status {
+                assert!(s.contains("supervision events"));
+            } else {
+                panic!("Not failed: {:?}", state.status);
+            }
+            if let Some(ref inner) = state.state {
+                assert!(!inner.supervision_events.is_empty());
+                for event in &inner.supervision_events {
+                    assert_eq!(event.actor_id.name(), format!("{}", child_name.clone()));
+                    assert_matches!(event.actor_status, ActorStatus::Failed(_));
+                }
+            }
+        }
+    }
+
     #[async_timed_test(timeout_secs = 30)]
     async fn test_cast() {
         let config = hyperactor::config::global::lock();