[hyperactor] mesh: ensure liveness for proc and actor spawning

Make proc and actor spawning subject to a liveness parameter: we expect the spawn to either complete, or else yile updated regularly (configured by a new parameter `PROC_SPAWN_MAX_IDLE` and `ACTOR_SPAWN_MAX_IDLE`. This is a bridge solution to provide good error messaging, and to prevent applications from halting. Once we have owned host meshes w/ their own comm actors, we can use the general purpose liveness mechanism for this.

We also clean up the error messages, refactor the status fence, and print the status of each individual proc on failure.

Differential Revision: [D83882826](https://our.internmc.facebook.com/intern/diff/D83882826/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D83882826/)!

ghstack-source-id: 314019055
Pull Request resolved: #1426

Author: mariusae

hyperactor/src/lib.rs

@@ -92,7 +92,7 @@ pub mod supervision;
 pub mod sync;
 /// Test utilities
 pub mod test_utils;
-mod time;
+pub mod time;
 
 pub use actor::Actor;
 pub use actor::ActorHandle;

hyperactor/src/time.rs

@@ -7,6 +7,7 @@
  */
 
 //! This module contains various utilities for dealing with time.
+//! (This probably belongs in a separate crate.)
 
 use std::sync::Arc;
 use std::sync::Mutex;

hyperactor_mesh/src/resource.rs

@@ -10,12 +10,17 @@
 //! in hyperactor meshes.
 
 use core::slice::GetDisjointMutIndex as _;
+use std::collections::HashMap;
+use std::fmt;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::marker::PhantomData;
 use std::mem::replace;
 use std::mem::take;
+use std::ops::Deref;
+use std::ops::DerefMut;
 use std::ops::Range;
+use std::time::Duration;
 
 use enum_as_inner::EnumAsInner;
 use hyperactor::Bind;
@@ -30,6 +35,8 @@ use hyperactor::accum::Accumulator;
 use hyperactor::accum::CommReducer;
 use hyperactor::accum::ReducerFactory;
 use hyperactor::accum::ReducerSpec;
+use hyperactor::mailbox::MailboxError;
+use hyperactor::mailbox::PortReceiver;
 use hyperactor::message::Bind;
 use hyperactor::message::Bindings;
 use hyperactor::message::Unbind;
@@ -50,7 +57,8 @@ use crate::v1::Name;
     PartialEq,
     Eq,
     Hash,
-    EnumAsInner
+    EnumAsInner,
+    strum::Display
 )]
 pub enum Status {
     /// The resource does not exist.
@@ -65,12 +73,17 @@ pub enum Status {
     Stopped,
     /// The resource has failed, with an error message.
     Failed(String),
+    /// The resource has been declared failed after a timeout.
+    Timeout(Duration),
 }
 
 impl Status {
     /// Returns whether the status is a terminating status.
     pub fn is_terminating(&self) -> bool {
-        matches!(self, Status::Stopping | Status::Stopped | Status::Failed(_))
+        matches!(
+            self,
+            Status::Stopping | Status::Stopped | Status::Failed(_) | Status::Timeout(_)
+        )
     }
 }
 
@@ -128,6 +141,34 @@ pub struct GetRankStatus {
     pub reply: PortRef<RankedValues<Status>>,
 }
 
+impl GetRankStatus {
+    pub async fn wait(
+        mut rx: PortReceiver<RankedValues<Status>>,
+        num_ranks: usize,
+        max_idle_time: Duration,
+    ) -> Result<RankedValues<Status>, RankedValues<Status>> {
+        let mut alarm = hyperactor::time::Alarm::new();
+        alarm.arm(max_idle_time);
+        let mut statuses = RankedValues::default();
+        loop {
+            let mut sleeper = alarm.sleeper();
+            tokio::select! {
+                _ = sleeper.sleep() => return Err(statuses),
+                new_statuses = rx.recv() => {
+                    match new_statuses {
+                        Ok(new_statuses) => statuses = new_statuses,
+                        Err(_) => return Err(statuses),
+                    }
+                }
+            }
+            alarm.arm(max_idle_time);
+            if statuses.rank(num_ranks) == num_ranks {
+                break Ok(statuses);
+            }
+        }
+    }
+}
+
 /// The state of a resource.
 #[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
 pub struct State<S> {
@@ -214,6 +255,14 @@ pub struct RankedValues<T> {
     intervals: Vec<(Range<usize>, T)>,
 }
 
+impl<T: PartialEq> PartialEq for RankedValues<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.intervals == other.intervals
+    }
+}
+
+impl<T: Eq> Eq for RankedValues<T> {}
+
 impl<T> Default for RankedValues<T> {
     fn default() -> Self {
         Self {
@@ -238,6 +287,28 @@ impl<T> RankedValues<T> {
     }
 }
 
+impl<T: Clone> RankedValues<T> {
+    pub fn materialized_iter(&self, until: usize) -> impl Iterator<Item = &T> + '_ {
+        assert_eq!(self.rank(until), until, "insufficient rank");
+        self.iter()
+            .flat_map(|(range, value)| std::iter::repeat(value).take(range.end - range.start))
+    }
+}
+
+impl<T: Hash + Eq + Clone> RankedValues<T> {
+    /// Invert this ranked values into a [`ValuesByRank<T>`].
+    pub fn invert(&self) -> ValuesByRank<T> {
+        let mut inverted: HashMap<T, Vec<Range<usize>>> = HashMap::new();
+        for (range, value) in self.iter() {
+            inverted
+                .entry(value.clone())
+                .or_default()
+                .push(range.clone());
+        }
+        ValuesByRank { values: inverted }
+    }
+}
+
 impl<T: Eq + Clone> RankedValues<T> {
     /// Merge `other` into this set of ranked values. Values in `other` that overlap
     /// with `self` take prededence.
@@ -298,6 +369,11 @@ impl<T: Eq + Clone> RankedValues<T> {
         }
     }
 
+    /// Merge the contents of this RankedValues into another RankedValues.
+    pub fn merge_into(self, other: &mut Self) {
+        other.merge_from(self);
+    }
+
     fn append(&mut self, range: Range<usize>, value: T) {
         if let Some(last) = self.intervals.last_mut()
             && last.0.end == range.start
@@ -310,6 +386,15 @@ impl<T: Eq + Clone> RankedValues<T> {
     }
 }
 
+impl RankedValues<Status> {
+    pub fn first_terminating(&self) -> Option<(usize, Status)> {
+        self.intervals
+            .iter()
+            .find(|(_, status)| status.is_terminating())
+            .map(|(range, status)| (range.start, status.clone()))
+    }
+}
+
 impl<T> From<(usize, T)> for RankedValues<T> {
     fn from((rank, value): (usize, T)) -> Self {
         Self {
@@ -318,6 +403,67 @@ impl<T> From<(usize, T)> for RankedValues<T> {
     }
 }
 
+impl<T> From<(Range<usize>, T)> for RankedValues<T> {
+    fn from((range, value): (Range<usize>, T)) -> Self {
+        Self {
+            intervals: vec![(range, value)],
+        }
+    }
+}
+
+/// An inverted index of RankedValues, providing all ranks for
+/// which each unique T-typed value appears.
+#[derive(Clone, Debug)]
+pub struct ValuesByRank<T> {
+    values: HashMap<T, Vec<Range<usize>>>,
+}
+
+impl<T: Eq + Hash> PartialEq for ValuesByRank<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.values == other.values
+    }
+}
+
+impl<T: Eq + Hash> Eq for ValuesByRank<T> {}
+
+impl<T: fmt::Display> fmt::Display for ValuesByRank<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut first_value = true;
+        for (value, ranges) in self.iter() {
+            if first_value {
+                first_value = false;
+            } else {
+                write!(f, ";")?;
+            }
+            write!(f, "{}=", value)?;
+            let mut first_range = true;
+            for range in ranges.iter() {
+                if first_range {
+                    first_range = false;
+                } else {
+                    write!(f, ",")?;
+                }
+                write!(f, "{}..{}", range.start, range.end)?;
+            }
+        }
+        Ok(())
+    }
+}
+
+impl<T> Deref for ValuesByRank<T> {
+    type Target = HashMap<T, Vec<Range<usize>>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.values
+    }
+}
+
+impl<T> DerefMut for ValuesByRank<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.values
+    }
+}
+
 /// Enabled for test only because we have to guarantee that the input
 /// iterator is well-formed.
 #[cfg(test)]
@@ -425,4 +571,16 @@ mod tests {
         assert_eq!(left.rank(70), 62);
         assert_eq!(left.rank(100), 62);
     }
+
+    #[test]
+    fn test_equality() {
+        assert_eq!(
+            RankedValues::from((0..10, 123)),
+            RankedValues::from((0..10, 123))
+        );
+        assert_eq!(
+            RankedValues::from((0..10, Status::Failed("foo".to_string()))),
+            RankedValues::from((0..10, Status::Failed("foo".to_string()))),
+        );
+    }
 }

hyperactor_mesh/src/v1.rs

@@ -17,10 +17,14 @@ pub mod testactor;
 pub mod testing;
 pub mod value_mesh;
 
+use std::collections::HashMap;
+use std::ops::Range;
 use std::str::FromStr;
+use std::time::Duration;
 
 pub use actor_mesh::ActorMesh;
 pub use actor_mesh::ActorMeshRef;
+use enum_as_inner::EnumAsInner;
 pub use host_mesh::HostMeshRef;
 use hyperactor::ActorId;
 use hyperactor::ActorRef;
@@ -33,12 +37,15 @@ use serde::Serialize;
 pub use value_mesh::ValueMesh;
 
 use crate::resource;
+use crate::resource::RankedValues;
+use crate::resource::Status;
+use crate::resource::ValuesByRank;
 use crate::shortuuid::ShortUuid;
 use crate::v1::host_mesh::HostMeshAgent;
 use crate::v1::host_mesh::HostMeshRefParseError;
 
 /// Errors that occur during mesh operations.
-#[derive(Debug, thiserror::Error)]
+#[derive(Debug, EnumAsInner, thiserror::Error)]
 pub enum Error {
     #[error("invalid mesh ref: expected {expected} ranks, but contains {actual} ranks")]
     InvalidRankCardinality { expected: usize, actual: usize },
@@ -96,6 +103,18 @@ pub enum Error {
         status: resource::Status,
     },
 
+    #[error(
+        "error spawning proc mesh: statuses: {}",
+        RankedValues::invert(&*.statuses)
+    )]
+    ProcSpawnError { statuses: RankedValues<Status> },
+
+    #[error(
+        "error spawning actor mesh: statuses: {}",
+        RankedValues::invert(&*.statuses)
+    )]
+    ActorSpawnError { statuses: RankedValues<Status> },
+
     #[error("error: {0} does not exist")]
     NotExist(Name),
 }