: actor_mesh: use commactor in cast_slices (#479)

Summary:

- in D77963903, routing was updated to use the `(ActorMeshId, ActorId)` pair as the stream key for sequence number tracking. this change allows different sub-slices of a mesh to safely share a common stream identity as long as they belong to the same logical `ActorMeshId`, avoiding issues like message reordering or duplication due to slice mismatch. as a result, this diff removes the now-unnecessary logic that intersected the user-provided selection with a reified view of the actor mesh's slice. correctness depends on all casts to a given `ActorMeshId`  being evaluated consistently against that mesh's slice.

- previously, `cast_slices` didn’t perform a true cast; it sent messages point-to-point to each rank in the input slices. now, we replace it with a `Selection` constructed via `Selection::of_ranks` and `dsl::union`, and invoke `cast`. the unused `cast_slices` is removed.

Reviewed By: mariusae

Differential Revision: D77953855

Author: shayne-fletcher

hyperactor_mesh/src/actor_mesh.rs

@@ -37,8 +37,6 @@ use ndslice::Selection;
 use ndslice::Shape;
 use ndslice::ShapeError;
 use ndslice::Slice;
-use ndslice::dsl;
-use ndslice::selection::ReifyView;
 use serde::Deserialize;
 use serde::Serialize;
 use tokio::sync::mpsc;
@@ -62,7 +60,7 @@ pub(crate) fn actor_mesh_cast<M: Castable + Clone, A>(
     caps: &impl cap::CanSend,
     actor_mesh_id: ActorMeshId,
     actor_mesh_shape: &Shape,
-    proc_mesh_shape: &Shape,
+    _proc_mesh_shape: &Shape,
     actor_name: &str,
     sender: &ActorId,
     comm_actor_ref: &ActorRef<CommActor>,
@@ -77,6 +75,7 @@ where
         "message_variant" => message.arm().unwrap_or_default(),
     ));
 
+    let slice = actor_mesh_shape.slice().clone();
     let message = CastMessageEnvelope::new(
         actor_mesh_id,
         sender.clone(),
@@ -86,29 +85,10 @@ where
         None, // TODO: reducer typehash
     )?;
 
-    // Sub-set the selection to the selection that represents the mesh's view
-    // of the root mesh. We need to do this because the comm actor uses the
-    // slice as the stream key; thus different sub-slices will result in potentially
-    // out of order delivery.
-    //
-    // TODO: We should repair this by introducing an explicit stream key, associated
-    // with the root mesh.
-    let selection_of_slice = proc_mesh_shape
-        .slice()
-        .reify_view(actor_mesh_shape.slice())
-        .expect("invalid slice");
-    let selection = dsl::intersection(selection, selection_of_slice);
-
     comm_actor_ref.send(
         caps,
         CastMessage {
-            dest: Uslice {
-                // TODO: currently this slice is being used as the stream key
-                // in comm actor. We should change it to an explicit id, maintained
-                // by the root proc mesh.
-                slice: proc_mesh_shape.slice().clone(),
-                selection,
-            },
+            dest: Uslice { slice, selection },
             message,
         },
     )?;

monarch_extension/src/mesh_controller.rs

@@ -7,6 +7,7 @@
  */
 
 use std::collections::BTreeMap;
+use std::collections::BTreeSet;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::collections::VecDeque;
@@ -32,6 +33,7 @@ use hyperactor::PortRef;
 use hyperactor::cap::CanSend;
 use hyperactor::mailbox::MailboxSenderError;
 use hyperactor_mesh::Mesh;
+use hyperactor_mesh::actor_mesh::ActorMesh;
 use hyperactor_mesh::actor_mesh::RootActorMesh;
 use hyperactor_mesh::shared_cell::SharedCell;
 use hyperactor_mesh::shared_cell::SharedCellRef;
@@ -53,7 +55,9 @@ use monarch_messages::worker::WorkerMessage;
 use monarch_messages::worker::WorkerParams;
 use monarch_tensor_worker::AssignRankMessage;
 use monarch_tensor_worker::WorkerActor;
+use ndslice::Selection;
 use ndslice::Slice;
+use ndslice::selection;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use tokio::sync::Mutex;
@@ -250,6 +254,7 @@ impl Invocation {
         }
     }
 
+    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `MailboxSenderError`.
     fn add_user(
         &mut self,
         sender: &impl CanSend,
@@ -287,6 +292,7 @@ impl Invocation {
         }
     }
 
+    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `MailboxSenderError`.
     fn complete(&mut self, sender: &impl CanSend) -> Result<(), MailboxSenderError> {
         let old_status = std::mem::replace(&mut self.status, Status::Complete {});
         match old_status {
@@ -310,6 +316,7 @@ impl Invocation {
     /// Incomplete, it may have users that will also become errored. This function
     /// will return those users so the error can be propagated. It does not autmoatically
     /// propagate the error to avoid deep recursive invocations.
+    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `MailboxSenderError`.
     fn set_exception(
         &mut self,
         sender: &impl CanSend,
@@ -447,6 +454,7 @@ impl History {
     }
 
     /// Add an invocation to the history.
+    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `MailboxSenderError`.
     pub fn add_invocation(
         &mut self,
         sender: &impl CanSend,
@@ -486,6 +494,7 @@ impl History {
 
     /// Propagate worker error to the invocation with the given Seq. This will also propagate
     /// to all seqs that depend on this seq directly or indirectly.
+    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `MailboxSenderError`.
     pub fn propagate_exception(
         &mut self,
         sender: &impl CanSend,
@@ -539,6 +548,7 @@ impl History {
 
     /// Mark the given rank as completed up to but excluding the given Seq. This will also purge history for
     /// any Seqs that are no longer relevant (completed on all ranks).
+    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `MailboxSenderError`.
     pub fn rank_completed(
         &mut self,
         sender: &impl CanSend,
@@ -623,6 +633,7 @@ impl MeshControllerActor {
     fn workers(&self) -> SharedCellRef<RootActorMesh<'static, WorkerActor>> {
         self.workers.as_ref().unwrap().borrow().unwrap()
     }
+
     fn handle_debug(
         &mut self,
         this: &Context<Self>,
@@ -741,7 +752,8 @@ impl Actor for MeshControllerActor {
         workers
             .borrow()
             .unwrap()
-            .cast_slices(vec![slice.clone()], AssignRankMessage::AssignRank())?;
+            .cast(selection::dsl::true_(), AssignRankMessage::AssignRank())?;
+
         self.workers = Some(workers);
         Ok(())
     }
@@ -796,7 +808,19 @@ impl Handler<ClientToControllerMessage> for MeshControllerActor {
     ) -> anyhow::Result<()> {
         match message {
             ClientToControllerMessage::Send { slices, message } => {
-                self.workers().cast_slices(slices, message)?;
+                let selection = slices
+                    .iter()
+                    .map(|slice| {
+                        Selection::of_ranks(
+                            self.workers().shape().slice(),
+                            &slice.iter().collect::<BTreeSet<usize>>(),
+                        )
+                    })
+                    .collect::<Result<Vec<_>, _>>()?
+                    .into_iter()
+                    .reduce(selection::dsl::union)
+                    .unwrap_or_else(selection::dsl::false_);
+                self.workers().cast(selection, message.clone())?;
             }
             ClientToControllerMessage::Node {
                 seq,
@@ -812,9 +836,8 @@ impl Handler<ClientToControllerMessage> for MeshControllerActor {
                 self.history.drop_refs(refs);
             }
             ClientToControllerMessage::SyncAtExit { port } => {
-                let all_ranks = vec![self.workers().shape().slice().clone()];
-                self.workers().cast_slices(
-                    all_ranks,
+                self.workers().cast(
+                    selection::dsl::true_(),
                     WorkerMessage::RequestStatus {
                         seq: self.history.seq_lower_bound,
                         controller: false,

monarch_tensor_worker/src/lib.rs

@@ -290,7 +290,17 @@ pub enum AssignRankMessage {
 }
 
 #[async_trait]
-#[forward(WorkerMessage)]
+impl Handler<WorkerMessage> for WorkerActor {
+    async fn handle(
+        &mut self,
+        cx: &hyperactor::Context<Self>,
+        message: WorkerMessage,
+    ) -> anyhow::Result<()> {
+        <Self as WorkerMessageHandler>::handle(self, cx, message).await
+    }
+}
+
+#[async_trait]
 impl WorkerMessageHandler for WorkerActor {
     async fn backend_network_init(
         &mut self,