v1: migrate to ValueMesh accumulator (meta-pytorch#1486)

Summary:
Pull Request resolved: meta-pytorch#1486

his change replaces the old `RankedValues<Status>` path with `ValueOverlay<Status>` updates reduced into a `ValueMesh<Status>` snapshot. `GetRankStatus.reply` now carries overlays; the comm reducer merges them (right-wins) and the accumulator applies them into a full `StatusMesh`, seeded to `NotExist` via the accumulator template (`StatusMesh::from_single(region, NotExist)`) and `ValueMesh::default()` + first-update install. agents emit a single-rank overlay (or an empty overlay when the rank is unknown). `wait()` consumes snapshots, declares completion when no `NotExist` remain for the target ranks, and on timeout returns the last snapshot (falling back to the seed if nothing arrived). termination detection now checks `values().any(is_terminating)`. legacy shims convert `StatusMesh` back to `RankedValues` only where error types still require it. removed the `RankedValues` accumulator/reducer and its registration.

Differential Revision: D84318425

Author: shayne-fletcher

hyperactor_mesh/src/proc_mesh/mesh_agent.rs

@@ -548,6 +548,9 @@ impl Handler<resource::GetRankStatus> for ProcMeshAgent {
         cx: &Context<Self>,
         get_rank_status: resource::GetRankStatus,
     ) -> anyhow::Result<()> {
+        use crate::resource::Status;
+        use crate::v1::StatusOverlay;
+
         let (rank, status) = match self.actor_states.get(&get_rank_status.name) {
             Some(ActorInstanceState {
                 spawn: Ok(actor_id),
@@ -560,9 +563,9 @@ impl Handler<resource::GetRankStatus> for ProcMeshAgent {
                 (
                     *create_rank,
                     if supervision_events.is_empty() {
-                        resource::Status::Running
+                        Status::Running
                     } else {
-                        resource::Status::Failed(format!(
+                        Status::Failed(format!(
                             "because of supervision events: {:?}",
                             supervision_events
                         ))
@@ -572,12 +575,21 @@ impl Handler<resource::GetRankStatus> for ProcMeshAgent {
             Some(ActorInstanceState {
                 spawn: Err(e),
                 create_rank,
-            }) => (*create_rank, resource::Status::Failed(e.to_string())),
+            }) => (*create_rank, Status::Failed(e.to_string())),
             // TODO: represent unknown rank
-            None => (usize::MAX, resource::Status::NotExist),
+            None => (usize::MAX, Status::NotExist),
+        };
+
+        // Send a sparse overlay update. If rank is unknown, emit an
+        // empty overlay.
+        let overlay = if rank == usize::MAX {
+            StatusOverlay::new()
+        } else {
+            StatusOverlay::try_from_runs(vec![(rank..(rank + 1), status)])
+                .expect("valid single-run overlay")
         };
+        get_rank_status.reply.send(cx, overlay)?;
 
-        get_rank_status.reply.send(cx, (rank, status).into())?;
         Ok(())
     }
 }

hyperactor_mesh/src/resource.rs

@@ -14,7 +14,6 @@ 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;
@@ -31,19 +30,17 @@ use hyperactor::PortRef;
 use hyperactor::RefClient;
 use hyperactor::RemoteMessage;
 use hyperactor::Unbind;
-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;
+use ndslice::Region;
+use ndslice::ViewExt;
 use serde::Deserialize;
 use serde::Serialize;
 
 use crate::v1::Name;
+use crate::v1::StatusOverlay;
 
 /// The current lifecycle status of a resource.
 #[derive(
@@ -119,8 +116,12 @@ impl Bind for Rank {
     }
 }
 
-/// Get the status of a resource at a rank. This message is designed to be
-/// cast and efficiently accumulated.
+/// Get the status of a resource across the mesh.
+///
+/// This message is cast to all ranks; each rank replies with a sparse
+/// status **overlay**. The comm reducer merges overlays (right-wins)
+/// and the accumulator applies them to produce **full StatusMesh
+/// snapshots** on the receiver side.
 #[derive(
     Clone,
     Debug,
@@ -136,34 +137,50 @@ impl Bind for Rank {
 pub struct GetRankStatus {
     /// The name of the resource.
     pub name: Name,
-    /// The status of the rank.
+    /// Sparse status updates (overlays) from a rank.
     #[binding(include)]
-    pub reply: PortRef<RankedValues<Status>>,
+    pub reply: PortRef<StatusOverlay>,
 }
 
 impl GetRankStatus {
     pub async fn wait(
-        mut rx: PortReceiver<RankedValues<Status>>,
+        mut rx: PortReceiver<crate::v1::StatusMesh>,
         num_ranks: usize,
         max_idle_time: Duration,
-    ) -> Result<RankedValues<Status>, RankedValues<Status>> {
+        region: Region, // used only for fallback
+    ) -> Result<crate::v1::StatusMesh, crate::v1::StatusMesh> {
+        debug_assert_eq!(region.num_ranks(), num_ranks, "region/num_ranks mismatch");
+
         let mut alarm = hyperactor::time::Alarm::new();
         alarm.arm(max_idle_time);
-        let mut statuses = RankedValues::default();
+
+        // Fallback snapshot if we time out before receiving anything.
+        let mut snapshot =
+            crate::v1::StatusMesh::from_single(region, crate::resource::Status::NotExist);
+
         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),
+                _ = sleeper.sleep() => return Err(snapshot),
+                next = rx.recv() => {
+                    match next {
+                        Ok(mesh) => { snapshot = mesh; }   // latest-wins snapshot
+                        Err(_)   => return Err(snapshot),
                     }
                 }
             }
+
             alarm.arm(max_idle_time);
-            if statuses.rank(num_ranks) == num_ranks {
-                break Ok(statuses);
+
+            // Completion: once every rank (among the first
+            // `num_ranks`) has reported at least something (i.e.
+            // moved off NotExist).
+            if snapshot
+                .values()
+                .take(num_ranks)
+                .all(|s| !matches!(s, crate::resource::Status::NotExist))
+            {
+                break Ok(snapshot);
             }
         }
     }
@@ -291,7 +308,7 @@ 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))
+            .flat_map(|(range, value)| std::iter::repeat_n(value, range.end - range.start))
     }
 }
 
@@ -475,44 +492,6 @@ impl<T> FromIterator<(Range<usize>, T)> for RankedValues<T> {
     }
 }
 
-impl<T: Eq + Clone + Named> Accumulator for RankedValues<T> {
-    type State = Self;
-    type Update = Self;
-
-    fn accumulate(&self, state: &mut Self::State, update: Self::Update) -> anyhow::Result<()> {
-        state.merge_from(update);
-        Ok(())
-    }
-
-    fn reducer_spec(&self) -> Option<ReducerSpec> {
-        Some(ReducerSpec {
-            typehash: <RankedValuesReducer<T> as Named>::typehash(),
-            builder_params: None,
-        })
-    }
-}
-
-#[derive(Named)]
-struct RankedValuesReducer<T>(std::marker::PhantomData<T>);
-
-impl<T: Hash + Eq + Ord + Clone> CommReducer for RankedValuesReducer<T> {
-    type Update = RankedValues<T>;
-
-    fn reduce(&self, mut left: Self::Update, right: Self::Update) -> anyhow::Result<Self::Update> {
-        left.merge_from(right);
-        Ok(left)
-    }
-}
-
-// register for concrete types:
-
-hyperactor::submit! {
-    ReducerFactory {
-        typehash_f: <RankedValuesReducer<Status> as Named>::typehash,
-        builder_f: |_| Ok(Box::new(RankedValuesReducer::<Status>(PhantomData))),
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

hyperactor_mesh/src/v1/host_mesh.rs

@@ -14,6 +14,8 @@ use hyperactor::config;
 use hyperactor::config::CONFIG;
 use hyperactor::config::ConfigAttr;
 use hyperactor::declare_attrs;
+
+use crate::resource::RankedValues;
 pub mod mesh_agent;
 
 use std::collections::HashSet;
@@ -43,7 +45,6 @@ use crate::resource;
 use crate::resource::CreateOrUpdateClient;
 use crate::resource::GetRankStatus;
 use crate::resource::GetRankStatusClient;
-use crate::resource::RankedValues;
 use crate::resource::Status;
 use crate::v1;
 use crate::v1::Name;
@@ -537,21 +538,45 @@ impl HostMeshRef {
             .concat(&per_host)
             .map_err(|err| v1::Error::ConfigurationError(err.into()))?;
 
+        let region: Region = extent.clone().into();
         let mesh_name = Name::new(name);
+
+        // Helper: legacy shim for error types that still require
+        // RankedValues<Status>. TODO(shayne-fletcher): Delete this
+        // shim once Error::ActorSpawnError carries a StatusMesh
+        // (ValueMesh<Status>) directly. At that point, use the mesh
+        // as-is and remove `mesh_to_rankedvalues_*` calls below.
+        fn mesh_to_rankedvalues_with_default(
+            mesh: &crate::v1::StatusMesh,
+            default_fill: Status,
+            len: usize,
+        ) -> RankedValues<Status> {
+            let mut out = RankedValues::from((0..len, default_fill));
+            for (i, s) in mesh.values().enumerate() {
+                if !matches!(s, Status::NotExist) {
+                    out.merge_from(RankedValues::from((i..i + 1, s.clone())));
+                }
+            }
+            out
+        }
+
         let mut procs = Vec::new();
-        let num_ranks = self.region().num_ranks() * per_host.num_ranks();
+        let num_ranks = region.num_ranks();
+        // Accumulator outputs full StatusMesh snapshots; seed with
+        // NotExist.
         let (port, rx) = cx.mailbox().open_accum_port_opts(
-            RankedValues::default(),
+            crate::v1::StatusMesh::from_single(region.clone(), Status::NotExist),
             Some(ReducerOpts {
                 max_update_interval: Some(Duration::from_millis(50)),
             }),
         );
 
-        // We CreateOrUpdate each proc, and then fence on getting statuses back.
-        // This is currently necessary because otherwise there is a race between
-        // the procs being created, and subsequent messages becoming unroutable
-        // (the agent actor manages the local muxer). We can solve this by allowing
-        // buffering in the host-level muxer.
+        // Create or update each proc, then fence on receiving status
+        // overlays. This prevents a race where procs become
+        // addressable before their local muxers are ready, which
+        // could make early messages unroutable. A future improvement
+        // would allow buffering in the host-level muxer to eliminate
+        // the need for this synchronization step.
         for (host_rank, host) in self.ranks.iter().enumerate() {
             for per_host_rank in 0..per_host.num_ranks() {
                 let create_rank = per_host.num_ranks() * host_rank + per_host_rank;
@@ -585,25 +610,42 @@ impl HostMeshRef {
 
         let start_time = RealClock.now();
 
-        match GetRankStatus::wait(rx, num_ranks, config::global::get(PROC_SPAWN_MAX_IDLE)).await {
+        // Wait on accumulated StatusMesh snapshots until complete or
+        // timeout.
+        match GetRankStatus::wait(
+            rx,
+            num_ranks,
+            config::global::get(PROC_SPAWN_MAX_IDLE),
+            region.clone(), // fallback mesh if nothing arrives
+        )
+        .await
+        {
             Ok(statuses) => {
-                if let Some((rank, status)) = statuses.first_terminating() {
+                // If any rank is terminating, surface a
+                // ProcCreationError pointing at that rank.
+                if let Some((rank, status)) = statuses
+                    .values()
+                    .enumerate()
+                    .find(|(_, s)| s.is_terminating())
+                {
                     let proc_name = Name::new(format!("{}-{}", name, rank % per_host.num_ranks()));
                     return Err(v1::Error::ProcCreationError {
                         proc_name,
-                        mesh_agent: self.ranks[rank].mesh_agent(),
+                        mesh_agent: self.ranks[rank / per_host.num_ranks()].mesh_agent(),
                         host_rank: rank / per_host.num_ranks(),
-                        status: status.clone(),
+                        status,
                     });
                 }
             }
             Err(complete) => {
-                // Fill the remaining statuses with a timeout error.
-                let mut statuses =
-                    RankedValues::from((0..num_ranks, Status::Timeout(start_time.elapsed())));
-                statuses.merge_from(complete);
-
-                return Err(v1::Error::ProcSpawnError { statuses });
+                // Fill remaining ranks with a timeout status via the
+                // legacy shim.
+                let legacy = mesh_to_rankedvalues_with_default(
+                    &complete,
+                    Status::Timeout(start_time.elapsed()),
+                    num_ranks,
+                );
+                return Err(v1::Error::ProcSpawnError { statuses: legacy });
             }
         }
 

hyperactor_mesh/src/v1/host_mesh/mesh_agent.rs

@@ -147,20 +147,22 @@ impl Handler<resource::GetRankStatus> for HostMeshAgent {
         cx: &Context<Self>,
         get_rank_status: resource::GetRankStatus,
     ) -> anyhow::Result<()> {
-        let Some(created) = self.created.get(&get_rank_status.name) else {
-            // TODO: how can we get the host's rank here? we should model its absence explicitly.
-            get_rank_status
-                .reply
-                .send(cx, (usize::MAX, resource::Status::NotExist).into())?;
-            return Ok(());
-        };
+        use crate::resource::Status;
+        use crate::v1::StatusOverlay;
 
-        let rank_status = match created {
-            (rank, Ok(_)) => (*rank, resource::Status::Running),
-            (rank, Err(e)) => (*rank, resource::Status::Failed(e.to_string())),
+        let (rank, status) = match self.created.get(&get_rank_status.name) {
+            Some((rank, Ok(_))) => (*rank, Status::Running),
+            Some((rank, Err(e))) => (*rank, Status::Failed(e.to_string())),
+            None => (usize::MAX, Status::NotExist),
         };
-        get_rank_status.reply.send(cx, rank_status.into())?;
 
+        let overlay = if rank == usize::MAX {
+            StatusOverlay::new()
+        } else {
+            StatusOverlay::try_from_runs(vec![(rank..(rank + 1), status)])
+                .expect("valid single-run overlay")
+        };
+        get_rank_status.reply.send(cx, overlay)?;
         Ok(())
     }
 }