Add casting support to resource::GetState messages

Summary:
Part of: #1209

Use casting to implement the `supervision_events` API instead of iterating over all
ProcMeshAgents. This will scale better as the size of the ProcMesh increases.

Unfortunately, the Actor name for the mesh agent is not compatible with the v1::Name struct
due to the missing uuid. Make `v1::Name` an enum to allow reserved names to be used for
things like ActorMeshes.

Differential Revision: D82687236

Author: dulinriley

hyperactor_mesh/src/proc_mesh/mesh_agent.rs

@@ -20,6 +20,7 @@ use enum_as_inner::EnumAsInner;
 use hyperactor::Actor;
 use hyperactor::ActorHandle;
 use hyperactor::ActorId;
+use hyperactor::Bind;
 use hyperactor::Context;
 use hyperactor::Data;
 use hyperactor::HandleClient;
@@ -31,6 +32,7 @@ use hyperactor::PortHandle;
 use hyperactor::PortRef;
 use hyperactor::ProcId;
 use hyperactor::RefClient;
+use hyperactor::Unbind;
 use hyperactor::actor::ActorStatus;
 use hyperactor::actor::remote::Remote;
 use hyperactor::channel;
@@ -167,7 +169,7 @@ impl State {
     handlers=[
         MeshAgentMessage,
         resource::CreateOrUpdate<ActorSpec>,
-        resource::GetState<ActorState>
+        resource::GetState<ActorState> { cast = true },
     ]
 )]
 pub struct ProcMeshAgent {
@@ -406,7 +408,7 @@ pub struct ActorSpec {
 }
 
 /// Actor state.
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named, Bind, Unbind)]
 pub struct ActorState {
     /// The actor's ID.
     pub actor_id: ActorId,

hyperactor_mesh/src/resource.rs

@@ -16,6 +16,10 @@ use hyperactor::Handler;
 use hyperactor::Named;
 use hyperactor::PortRef;
 use hyperactor::RefClient;
+use hyperactor::RemoteMessage;
+use hyperactor::message::Bind;
+use hyperactor::message::Bindings;
+use hyperactor::message::Unbind;
 use serde::Deserialize;
 use serde::Serialize;
 
@@ -70,3 +74,24 @@ pub struct GetState<S> {
     #[reply]
     pub reply: PortRef<State<S>>,
 }
+
+// Cannot derive Bind and Unbind for this generic, implement manually.
+impl<T> Unbind for GetState<T>
+where
+    T: RemoteMessage,
+    T: Unbind,
+{
+    fn unbind(&self, bindings: &mut Bindings) -> anyhow::Result<()> {
+        self.reply.unbind(bindings)
+    }
+}
+
+impl<T> Bind for GetState<T>
+where
+    T: RemoteMessage,
+    T: Bind,
+{
+    fn bind(&mut self, bindings: &mut Bindings) -> anyhow::Result<()> {
+        self.reply.bind(bindings)
+    }
+}

hyperactor_mesh/src/v1.rs

@@ -164,27 +164,51 @@ pub type Result<T> = std::result::Result<T, Error>;
     Serialize,
     Deserialize
 )]
-pub struct Name(pub String, pub ShortUuid);
+pub enum Name {
+    /// Normal names for most actors.
+    Suffixed(String, ShortUuid),
+    /// Reserved names for system actors without UUIDs.
+    Reserved(String),
+}
 
 impl Name {
     /// Create a new `Name` from a user-provided base name.
     pub fn new(name: impl Into<String>) -> Self {
+        Self::new_with_uuid(name, Some(ShortUuid::generate()))
+    }
+
+    /// Create a Reserved `Name` with no uuid. Only for use by system actors.
+    pub(crate) fn new_reserved(name: impl Into<String>) -> Self {
+        Self::new_with_uuid(name, None)
+    }
+
+    fn new_with_uuid(name: impl Into<String>, uuid: Option<ShortUuid>) -> Self {
         let mut name = name.into();
         if name.is_empty() {
             name = "unnamed".to_string();
         }
-        let uuid = ShortUuid::generate();
-        Self(name, uuid)
+        if let Some(uuid) = uuid {
+            Self::Suffixed(name, uuid)
+        } else {
+            Self::Reserved(name)
+        }
     }
 
     /// The name portion of this `Name`.
     pub fn name(&self) -> &str {
-        &self.0
+        match self {
+            Self::Suffixed(n, _) => n,
+            Self::Reserved(n) => n,
+        }
     }
 
     /// The UUID portion of this `Name`.
+    /// Only valid for Name::Suffixed, if called on Name::Reserved it'll panic.
     pub fn uuid(&self) -> &ShortUuid {
-        &self.1
+        match self {
+            Self::Suffixed(_, uuid) => uuid,
+            Self::Reserved(_) => panic!("Reserved name has no UUID"),
+        }
     }
 }
 
@@ -208,24 +232,33 @@ impl FromStr for Name {
     type Err = NameParseError;
 
     fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
-        let (name, uuid) = s.split_once('-').ok_or(NameParseError::MissingSeparator)?;
-        if name.is_empty() {
-            return Err(NameParseError::MissingName);
+        if let Some((name, uuid)) = s.split_once('-') {
+            if name.is_empty() {
+                return Err(NameParseError::MissingName);
+            }
+            if uuid.is_empty() {
+                return Err(NameParseError::MissingName);
+            }
+
+            Ok(Name::new_with_uuid(name.to_string(), Some(uuid.parse()?)))
+        } else {
+            if s.is_empty() {
+                return Err(NameParseError::MissingName);
+            }
+            Ok(Name::new_reserved(s))
         }
-        if uuid.is_empty() {
-            return Err(NameParseError::MissingName);
-        }
-
-        let name = name.to_string();
-        let uuid = uuid.parse()?;
-        Ok(Name(name, uuid))
     }
 }
 
 impl std::fmt::Display for Name {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}-", self.name())?;
-        self.uuid().format(f, true /*raw*/)
+        match self {
+            Self::Suffixed(n, uuid) => {
+                write!(f, "{}-", n)?;
+                uuid.format(f, true /*raw*/)
+            }
+            Self::Reserved(n) => write!(f, "{}", n),
+        }
     }
 }
 

hyperactor_mesh/src/v1/actor_mesh.rs

@@ -159,9 +159,10 @@ impl<A: Actor + RemoteActor> ActorMeshRef<A> {
     pub async fn supervision_events(
         &self,
         cx: &impl context::Actor,
-        name: Name,
     ) -> v1::Result<ValueMesh<Vec<ActorSupervisionEvent>>> {
-        self.proc_mesh.supervision_events(cx, name).await
+        self.proc_mesh
+            .supervision_events(cx, self.name.clone())
+            .await
     }
 }
 
@@ -452,12 +453,8 @@ mod tests {
         // status such that when a process switches to unhealthy it sets a
         // supervision event.
         let actor_mesh_ref = actor_mesh.freeze();
-        let child_name_clone = child_name.clone();
         let supervision_task = tokio::spawn(async move {
-            match actor_mesh_ref
-                .supervision_events(&instance, child_name_clone)
-                .await
-            {
+            match actor_mesh_ref.supervision_events(&instance).await {
                 Ok(events) => {
                     for event_list in events.values() {
                         assert!(!event_list.is_empty());

hyperactor_mesh/src/v1/proc_mesh.rs

@@ -357,18 +357,46 @@ impl ProcMeshRef {
         vm.join().await.transpose()
     }
 
+    fn agent_mesh(&self) -> ActorMeshRef<ProcMeshAgent> {
+        let actor_refs = self.values().map(|p| p.agent).collect::<Vec<_>>();
+        let agent_name = actor_refs.first().unwrap().actor_id().name();
+        // This name must match the ProcMeshAgent name, which can change depending on the allocator.
+        ActorMeshRef::new(Name::new_reserved(agent_name), self.clone(), actor_refs)
+    }
+
     /// The supervision events of procs in this mesh.
     pub async fn supervision_events(
         &self,
         cx: &impl context::Actor,
         name: Name,
     ) -> v1::Result<ValueMesh<Vec<ActorSupervisionEvent>>> {
-        let vm: ValueMesh<_> = self.map_into(|proc_ref| {
-            let proc_ref = proc_ref.clone();
-            let name = name.clone();
-            async move { proc_ref.supervision_events(cx, name).await }
-        });
-        vm.join().await.transpose()
+        let agent_mesh = self.agent_mesh();
+        let (port, mut rx) = cx.mailbox().open_port::<resource::State<ActorState>>();
+        agent_mesh.cast(
+            cx,
+            resource::GetState::<ActorState> {
+                name: name.clone(),
+                reply: port.bind(),
+            },
+        )?;
+        let expected = self.ranks.len();
+        let mut states = Vec::with_capacity(expected);
+        for _ in 0..expected {
+            let state = rx.recv().await?;
+            states.push(state);
+        }
+        let vm = states
+            .into_iter()
+            .map(|state| {
+                if let Some(state) = state.state {
+                    state.supervision_events
+                } else {
+                    // Empty vec for ranks with no supervision events.
+                    Vec::new()
+                }
+            })
+            .collect_mesh::<ValueMesh<Vec<_>>>(self.region.clone())?;
+        Ok(vm)
     }
 
     /// Spawn an actor on all of the procs in this mesh, returning a new ActorMesh.