mesh: concurrent spawns, aggregated statuses (#1394)

Summary:
Pull Request resolved: #1394

This adds concurrent spawning for procs, and introduces a principled cast/accumulate status aggregation.

The latter lets us efficiently determine the status of a mesh, and apply whatever policy makes sense in context (e.g., early exit).

We will use GetRankStatus also to apply liveness checks for actors, as we can support this efficiently.

This change also introduces an RLE-encoded data structure, RankedValues, which works well for low-cardinality contexts (e.g., statuses); but to use it more generally for value meshes will require that the merge operation be optimized.
ghstack-source-id: 313885569
exported-using-ghexport

Reviewed By: shayne-fletcher

Differential Revision: D83678884

fbshipit-source-id: 185847b978c2bfc68e2a0f004ee37354ff56ddae

Author: mariusae

hyperactor/src/message.rs

@@ -136,6 +136,15 @@ impl<M> Unbound<M> {
     pub fn new(message: M, bindings: Bindings) -> Self {
         Self { message, bindings }
     }
+
+    /// Use the provided function to update values inside bindings in the same
+    /// order as they were pushed into bindings.
+    pub fn visit_mut<T: Serialize + DeserializeOwned + Named>(
+        &mut self,
+        f: impl FnMut(&mut T) -> anyhow::Result<()>,
+    ) -> anyhow::Result<()> {
+        self.bindings.visit_mut(f)
+    }
 }
 
 impl<M: Bind> Unbound<M> {

hyperactor_mesh/src/bootstrap.rs

@@ -15,6 +15,7 @@ use std::future;
 use std::io;
 use std::io::Write;
 use std::os::unix::process::ExitStatusExt;
+use std::path::PathBuf;
 use std::process::Stdio;
 use std::sync::Arc;
 use std::sync::Mutex;
@@ -220,7 +221,9 @@ pub enum Bootstrap {
     Host {
         /// The address on which to serve the host.
         addr: ChannelAddr,
-
+        /// If specified, use the provided command instead of
+        /// [`BootstrapCommand::current`].
+        command: Option<BootstrapCommand>,
         /// Config snapshot for the child.
         config: Option<Attrs>,
     },
@@ -345,7 +348,11 @@ impl Bootstrap {
                     Err(e) => e.into(),
                 }
             }
-            Bootstrap::Host { addr, config } => {
+            Bootstrap::Host {
+                addr,
+                command,
+                config,
+            } => {
                 if config.is_some() {
                     tracing::debug!(
                         "bootstrap: Host received config snapshot (carried, not applied)"
@@ -354,7 +361,10 @@ impl Bootstrap {
                     tracing::debug!("bootstrap: no config snapshot provided (Host)");
                 }
 
-                let command = ok!(BootstrapCommand::current());
+                let command = match command {
+                    Some(command) => command,
+                    None => ok!(BootstrapCommand::current()),
+                };
                 let manager = BootstrapProcManager::new(command);
                 let (host, _handle) = ok!(Host::serve(manager, addr).await);
                 let addr = host.addr().clone();
@@ -1303,7 +1313,7 @@ impl hyperactor::host::ProcHandle for BootstrapProcHandle {
 /// A specification of the command used to bootstrap procs.
 #[derive(Debug, Named, Serialize, Deserialize, Clone, Default, PartialEq, Eq)]
 pub struct BootstrapCommand {
-    pub program: std::path::PathBuf,
+    pub program: PathBuf,
     pub arg0: Option<String>,
     pub args: Vec<String>,
     pub env: HashMap<String, String>,
@@ -1341,6 +1351,18 @@ impl BootstrapCommand {
     }
 }
 
+impl<T: Into<PathBuf>> From<T> for BootstrapCommand {
+    /// Creates a bootstrap command from the provided path.
+    fn from(s: T) -> Self {
+        Self {
+            program: s.into(),
+            arg0: None,
+            args: vec![],
+            env: HashMap::new(),
+        }
+    }
+}
+
 /// A process manager for launching and supervising **bootstrap
 /// processes** (via the [`bootstrap`] entry point).
 ///
@@ -1956,7 +1978,7 @@ fn debug_sink() -> &'static Mutex<DebugSink> {
 }
 
 /// If true, send `Debug` messages to stderr.
-const DEBUG_TO_STDERR: bool = true;
+const DEBUG_TO_STDERR: bool = false;
 
 /// A bootstrap specific debug writer. If the file /tmp/monarch-bootstrap-debug.log
 /// exists, then the writer's destination is that file; otherwise it discards all writes.
@@ -2059,6 +2081,7 @@ mod tests {
     fn test_bootstrap_mode_env_string_none_config_host() {
         let value = Bootstrap::Host {
             addr: ChannelAddr::any(ChannelTransport::Unix),
+            command: None,
             config: None,
         };
 
@@ -2113,6 +2136,7 @@ mod tests {
         {
             let original = Bootstrap::Host {
                 addr: ChannelAddr::any(ChannelTransport::Unix),
+                command: None,
                 config: Some(attrs.clone()),
             };
             let env_str = original.to_env_safe_string().expect("encode bootstrap");

hyperactor_mesh/src/comm.rs

@@ -8,6 +8,7 @@
 
 use crate::comm::multicast::CAST_ORIGINATING_SENDER;
 use crate::reference::ActorMeshId;
+use crate::resource;
 pub mod multicast;
 
 use std::cmp::Ordering;
@@ -279,6 +280,13 @@ impl CommActor {
         if deliver_here {
             let rank_on_root_mesh = mode.self_rank(cx.self_id())?;
             let cast_rank = message.relative_rank(rank_on_root_mesh)?;
+            // Replace ranks with self ranks.
+            message
+                .data_mut()
+                .visit_mut::<resource::Rank>(|resource::Rank(rank)| {
+                    *rank = Some(cast_rank);
+                    Ok(())
+                })?;
             let cast_shape = message.shape();
             let point = cast_shape
                 .extent()

hyperactor_mesh/src/lib.rs

@@ -11,6 +11,8 @@
 #![feature(assert_matches)]
 #![feature(exit_status_error)]
 #![feature(impl_trait_in_bindings)]
+#![feature(get_disjoint_mut_helpers)]
+#![feature(exact_size_is_empty)]
 
 pub mod actor_mesh;
 pub mod alloc;

hyperactor_mesh/src/proc_mesh/mesh_agent.rs

@@ -164,21 +164,29 @@ impl State {
     }
 }
 
+/// Actor state used for v1 API.
+#[derive(Debug)]
+struct ActorInstanceState {
+    create_rank: usize,
+    spawn: Result<ActorId, anyhow::Error>,
+}
+
 /// A mesh agent is responsible for managing procs in a [`ProcMesh`].
 #[derive(Debug)]
 #[hyperactor::export(
     handlers=[
         MeshAgentMessage,
-        resource::CreateOrUpdate<ActorSpec>,
+        resource::CreateOrUpdate<ActorSpec> { cast = true },
         resource::GetState<ActorState> { cast = true },
+        resource::GetRankStatus { cast = true },
     ]
 )]
 pub struct ProcMeshAgent {
     proc: Proc,
     remote: Remote,
     state: State,
     /// Actors created and tracked through the resource behavior.
-    created: HashMap<Name, Result<ActorId, anyhow::Error>>,
+    actor_states: HashMap<Name, ActorInstanceState>,
     /// If true, and supervisor is None, record supervision events to be reported
     /// to owning actors later.
     record_supervision_events: bool,
@@ -203,7 +211,7 @@ impl ProcMeshAgent {
             proc: proc.clone(),
             remote: Remote::collect(),
             state: State::UnconfiguredV0 { sender },
-            created: HashMap::new(),
+            actor_states: HashMap::new(),
             record_supervision_events: false,
             supervision_events: HashMap::new(),
         };
@@ -216,7 +224,7 @@ impl ProcMeshAgent {
             proc: proc.clone(),
             remote: Remote::collect(),
             state: State::V1,
-            created: HashMap::new(),
+            actor_states: HashMap::new(),
             record_supervision_events: true,
             supervision_events: HashMap::new(),
         };
@@ -442,10 +450,10 @@ pub struct ActorState {
 impl Handler<resource::CreateOrUpdate<ActorSpec>> for ProcMeshAgent {
     async fn handle(
         &mut self,
-        cx: &Context<Self>,
+        _cx: &Context<Self>,
         create_or_update: resource::CreateOrUpdate<ActorSpec>,
     ) -> anyhow::Result<()> {
-        if self.created.contains_key(&create_or_update.name) {
+        if self.actor_states.contains_key(&create_or_update.name) {
             // There is no update.
             return Ok(());
         }
@@ -454,19 +462,63 @@ impl Handler<resource::CreateOrUpdate<ActorSpec>> for ProcMeshAgent {
             actor_type,
             params_data,
         } = create_or_update.spec;
-        self.created.insert(
+        self.actor_states.insert(
             create_or_update.name.clone(),
-            self.remote
-                .gspawn(
-                    &self.proc,
-                    &actor_type,
-                    &create_or_update.name.to_string(),
-                    params_data,
-                )
-                .await,
+            ActorInstanceState {
+                create_rank: create_or_update.rank.unwrap(),
+                spawn: self
+                    .remote
+                    .gspawn(
+                        &self.proc,
+                        &actor_type,
+                        &create_or_update.name.to_string(),
+                        params_data,
+                    )
+                    .await,
+            },
         );
 
-        create_or_update.reply.send(cx, true)?;
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl Handler<resource::GetRankStatus> for ProcMeshAgent {
+    async fn handle(
+        &mut self,
+        cx: &Context<Self>,
+        get_rank_status: resource::GetRankStatus,
+    ) -> anyhow::Result<()> {
+        let (rank, status) = match self.actor_states.get(&get_rank_status.name) {
+            Some(ActorInstanceState {
+                spawn: Ok(actor_id),
+                create_rank,
+            }) => {
+                let supervision_events = self
+                    .supervision_events
+                    .get(actor_id)
+                    .map_or_else(Vec::new, |a| a.clone());
+                (
+                    *create_rank,
+                    if supervision_events.is_empty() {
+                        resource::Status::Running
+                    } else {
+                        resource::Status::Failed(format!(
+                            "because of supervision events: {:?}",
+                            supervision_events
+                        ))
+                    },
+                )
+            }
+            Some(ActorInstanceState {
+                spawn: Err(e),
+                create_rank,
+            }) => (*create_rank, resource::Status::Failed(e.to_string())),
+            // TODO: represent unknown rank
+            None => (usize::MAX, resource::Status::NotExist),
+        };
+
+        get_rank_status.reply.send(cx, (rank, status).into())?;
         Ok(())
     }
 }
@@ -478,12 +530,11 @@ impl Handler<resource::GetState<ActorState>> for ProcMeshAgent {
         cx: &Context<Self>,
         get_state: resource::GetState<ActorState>,
     ) -> anyhow::Result<()> {
-        let rank = self
-            .state
-            .rank()
-            .ok_or_else(|| anyhow::anyhow!("tried to get status of unconfigured proc"))?;
-        let state = match self.created.get(&get_state.name) {
-            Some(Ok(actor_id)) => {
+        let state = match self.actor_states.get(&get_state.name) {
+            Some(ActorInstanceState {
+                create_rank,
+                spawn: Ok(actor_id),
+            }) => {
                 let supervision_events = self
                     .supervision_events
                     .get(actor_id)
@@ -501,12 +552,12 @@ impl Handler<resource::GetState<ActorState>> for ProcMeshAgent {
                     status,
                     state: Some(ActorState {
                         actor_id: actor_id.clone(),
-                        create_rank: rank,
+                        create_rank: *create_rank,
                         supervision_events,
                     }),
                 }
             }
-            Some(Err(e)) => resource::State {
+            Some(ActorInstanceState { spawn: Err(e), .. }) => resource::State {
                 name: get_state.name.clone(),
                 status: resource::Status::Failed(e.to_string()),
                 state: None,