: actor_mesh: use commactor in cast_slices

Summary: refactor `cast_slices`. the old approach iterated over the input slices and for each rank in each slice, sent a message. now, call `reify_views` on the input slices to produce a single `Selection` (union of ranges). then, cast the message to the selected ranks via the comm actor.

Differential Revision: D77935610

Author: shayne-fletcher

hyperactor_mesh/src/actor_mesh.rs

@@ -25,7 +25,6 @@ use hyperactor::RemoteMessage;
 use hyperactor::Unbind;
 use hyperactor::WorldId;
 use hyperactor::actor::RemoteActor;
-use hyperactor::attrs::Attrs;
 use hyperactor::cap;
 use hyperactor::mailbox::MailboxSenderError;
 use hyperactor::mailbox::PortReceiver;
@@ -49,7 +48,6 @@ use crate::comm::multicast::CastMessage;
 use crate::comm::multicast::CastMessageEnvelope;
 use crate::comm::multicast::DestinationPort;
 use crate::comm::multicast::Uslice;
-use crate::comm::multicast::set_cast_info_on_headers;
 use crate::metrics;
 use crate::proc_mesh::ProcMesh;
 use crate::reference::ActorMeshId;
@@ -237,32 +235,35 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
     /// Until the selection logic is more powerful, we need a way to
     /// replicate the send patterns that the worker actor mesh actually does.
     #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `CastError`.
-    pub fn cast_slices<M: RemoteMessage + Clone>(
-        &self,
-        sel: Vec<Slice>,
-        message: M,
-    ) -> Result<(), CastError>
+    pub fn cast_slices<M>(&self, sel: Vec<Slice>, message: M) -> Result<(), CastError>
     where
+        M: Castable + Clone,
         A: RemoteHandles<M> + RemoteHandles<IndexedErasedUnbound<M>>,
     {
         let _ = metrics::ACTOR_MESH_CAST_DURATION.start(hyperactor::kv_pairs!(
             "message_type" => M::typename(),
             "message_variant" => message.arm().unwrap_or_default(),
         ));
-        for ref slice in sel {
-            for rank in slice.iter() {
-                let mut headers = Attrs::new();
-                set_cast_info_on_headers(
-                    &mut headers,
-                    rank,
-                    self.shape().clone(),
-                    self.proc_mesh.client().actor_id().clone(),
-                );
-                self.ranks[rank]
-                    .send_with_headers(self.proc_mesh.client(), headers, message.clone())
-                    .map_err(|err| CastError::MailboxSenderError(rank, err))?;
-            }
-        }
+
+        let client = self.proc_mesh.client();
+        let snd = client.actor_id().clone();
+        let dst = DestinationPort::new::<A, M>(self.name().to_string());
+        let shape = self.proc_mesh.shape().clone();
+        let base_slice = self.proc_mesh.shape().slice().clone();
+        let slices: &[&Slice] = &sel.iter().collect::<Vec<_>>();
+        let message = CastMessageEnvelope::new(snd, dst, shape, message, None)?;
+
+        self.proc_mesh.comm_actor().send(
+            client,
+            CastMessage {
+                dest: Uslice {
+                    slice: base_slice.clone(),
+                    selection: base_slice.reify_views(slices).expect("invalid slices"),
+                },
+                message,
+            },
+        )?;
+
         Ok(())
     }
 
@@ -528,6 +529,7 @@ mod tests {
             use $crate::sel_from_shape;
             use $crate::sel;
             use $crate::proc_mesh::SharedSpawnable;
+            use $crate::comm::multicast::set_cast_info_on_headers;
             use std::collections::VecDeque;
             use hyperactor::data::Serialized;