Use v0's cast implementation for v1 ActorMesh (#1187)

Summary:

As title.

Reviewed By: mariusae

Differential Revision: D82251703

Author: pzhan9

hyperactor_mesh/src/v1.rs

@@ -77,6 +77,9 @@ pub enum Error {
 
     #[error("error while sending message to actor {0}: {1}")]
     SendingError(ActorId, Box<MailboxSenderError>),
+
+    #[error("error while casting message to {0}: {1}")]
+    CastingError(Name, anyhow::Error),
 }
 
 impl From<crate::alloc::AllocatorError> for Error {

hyperactor_mesh/src/v1/actor_mesh.rs

@@ -15,13 +15,20 @@ use hyperactor::RemoteMessage;
 use hyperactor::actor::RemoteActor;
 use hyperactor::context;
 use hyperactor::message::Castable;
+use hyperactor::message::IndexedErasedUnbound;
+use hyperactor_mesh_macros::sel;
+use ndslice::Selection;
+use ndslice::Shape;
+use ndslice::ViewExt;
 use ndslice::view;
 use ndslice::view::Region;
 use ndslice::view::View;
-use ndslice::view::ViewExt;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::CommActor;
+use crate::actor_mesh as v0_actor_mesh;
+use crate::reference::ActorMeshId;
 use crate::v1;
 use crate::v1::Error;
 use crate::v1::Name;
@@ -82,16 +89,41 @@ impl<A: Actor + RemoteActor> ActorMeshRef<A> {
     /// Cast a message to all actors in this mesh.
     pub fn cast<M>(&self, cx: &impl context::Actor, message: M) -> v1::Result<()>
     where
-        M: Castable + RemoteMessage + Clone,
-        A: RemoteHandles<M>,
+        A: RemoteHandles<M> + RemoteHandles<IndexedErasedUnbound<M>>,
+        M: Castable + RemoteMessage,
     {
-        // todo: headers, binding/unbinding/accumulation
-        for actor_ref in self.values() {
-            actor_ref
-                .send(cx, message.clone())
-                .map_err(|e| Error::SendingError(actor_ref.actor_id().clone(), Box::new(e)))?;
+        self.region();
+        let cast_mesh_shape = to_shape(view::Ranked::region(self));
+        let comm_actor_ref = self
+            .proc_mesh
+            .root_mesh_rank_0
+            .attest::<CommActor>(self.proc_mesh.comm_actor_name());
+        let actor_mesh_id = ActorMeshId::V1(self.name.clone());
+        match &self.proc_mesh.root_region {
+            Some(root_region) => {
+                let root_mesh_shape = to_shape(root_region);
+                v0_actor_mesh::cast_to_sliced_mesh::<A, M>(
+                    cx,
+                    actor_mesh_id,
+                    &comm_actor_ref,
+                    &sel!(*),
+                    message,
+                    &cast_mesh_shape,
+                    &root_mesh_shape,
+                )
+                .map_err(|e| Error::CastingError(self.name.clone(), e.into()))
+            }
+            None => v0_actor_mesh::actor_mesh_cast::<A, M>(
+                cx,
+                actor_mesh_id,
+                &comm_actor_ref,
+                sel!(*),
+                &cast_mesh_shape,
+                &cast_mesh_shape,
+                message,
+            )
+            .map_err(|e| Error::CastingError(self.name.clone(), e.into())),
         }
-        Ok(())
     }
 }
 
@@ -120,3 +152,8 @@ impl<A: RemoteActor> view::RankedSliceable for ActorMeshRef<A> {
         ActorMeshRef::new(self.name.clone(), proc_mesh, actor_refs)
     }
 }
+
+fn to_shape(region: &Region) -> Shape {
+    Shape::new(region.labels().to_vec(), region.slice().clone())
+        .expect("Shape::new should not fail because a Region by definition is a valid Shape")
+}

hyperactor_mesh/src/v1/proc_mesh.rs

@@ -32,6 +32,7 @@ use ndslice::view::Region;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::CommActor;
 use crate::alloc::Alloc;
 use crate::alloc::AllocExt;
 use crate::alloc::AllocatedProc;
@@ -90,6 +91,7 @@ impl ProcRef {
 pub struct ProcMesh {
     name: Name,
     allocation: ProcMeshAllocation,
+    comm_actor_name: Name,
 }
 
 impl ProcMesh {
@@ -99,12 +101,21 @@ impl ProcMesh {
         let region = self.allocation.extent().clone().into();
         match &self.allocation {
             ProcMeshAllocation::Allocated { ranks, .. } => {
-                ProcMeshRef::new(self.name.clone(), region, Arc::clone(ranks)).unwrap()
+                let root_mesh_rank_0 = ranks.first().expect("root mesh cannot be empty").clone();
+                ProcMeshRef::new(
+                    self.name.clone(),
+                    region,
+                    Arc::clone(ranks),
+                    self.comm_actor_name.clone(),
+                    None, // this is the root mesh
+                    root_mesh_rank_0,
+                )
+                .unwrap()
             }
         }
     }
 
-    /// Allocate a new ProcMeshRef from the provided alloc.
+    /// Allocate a new ProcMesh from the provided alloc.
     /// Allocate does not require an owning actor because references are not owned.
     /// Allocate a new ProcMesh from the provided alloc.
     pub async fn allocate(
@@ -173,7 +184,7 @@ impl ProcMesh {
             }
         }
 
-        let ranks = running
+        let ranks: Vec<_> = running
             .into_iter()
             .enumerate()
             .map(|(create_rank, allocated)| ProcRef {
@@ -183,13 +194,23 @@ impl ProcMesh {
             })
             .collect();
 
-        Ok(Self {
+        let proc_mesh = Self {
             name: Name::new(name),
             allocation: ProcMeshAllocation::Allocated {
                 alloc: Box::new(alloc),
                 ranks: Arc::new(ranks),
             },
-        })
+            comm_actor_name: Name::new("comm"),
+        };
+        // Spawn a comm actor on each proc, so that they can be used to perform
+        // tree distribution and accumulation.
+        let comm_actor_name = proc_mesh.comm_actor_name.clone();
+        proc_mesh
+            .freeze()
+            .spawn_with_name::<CommActor>(cx, comm_actor_name, &Default::default())
+            .await?;
+
+        Ok(proc_mesh)
     }
 }
 
@@ -237,11 +258,28 @@ pub struct ProcMeshRef {
     name: Name,
     region: Region,
     ranks: Arc<Vec<ProcRef>>,
+    comm_actor_name: Name,
+    // Temporary: used to fit v1 ActorMesh with v0's casting implementation. This
+    // should be removed after we remove the v0 code.
+    // The root region of this mesh. None means this mesh itself is the root.
+    pub(crate) root_region: Option<Region>,
+    // Temporary: used to fit v1 ActorMesh with v0's casting implementation. This
+    // should be removed after we remove the v0 code.
+    // v0 casting requires root mesh rank 0 as the 1st hop, so we need to provide
+    // it here. For v1, this can be removed since v1 can use any rank.
+    pub(crate) root_mesh_rank_0: ProcRef,
 }
 
 impl ProcMeshRef {
-    /// Create a new ProcMeshRef from the given name, region, and ranks.
-    fn new(name: Name, region: Region, ranks: Arc<Vec<ProcRef>>) -> v1::Result<Self> {
+    /// Create a new ProcMeshRef from the given name, region, ranks, and so on.
+    fn new(
+        name: Name,
+        region: Region,
+        ranks: Arc<Vec<ProcRef>>,
+        comm_actor_name: Name,
+        root_region: Option<Region>,
+        root_mesh_rank_0: ProcRef,
+    ) -> v1::Result<Self> {
         if region.num_ranks() != ranks.len() {
             return Err(v1::Error::InvalidRankCardinality {
                 expected: region.num_ranks(),
@@ -252,9 +290,16 @@ impl ProcMeshRef {
             name,
             region,
             ranks,
+            comm_actor_name,
+            root_region,
+            root_mesh_rank_0,
         })
     }
 
+    pub(crate) fn comm_actor_name(&self) -> &Name {
+        &self.comm_actor_name
+    }
+
     /// The current statuses of procs in this mesh.
     #[allow(dead_code)]
     async fn status(&self, cx: &impl context::Actor) -> v1::Result<ValueMesh<bool>> {
@@ -273,6 +318,19 @@ impl ProcMeshRef {
         name: &str,
         params: &A::Params,
     ) -> v1::Result<ActorMesh<A>>
+    where
+        A::Params: RemoteMessage,
+    {
+        self.spawn_with_name(cx, Name::new(name), params).await
+    }
+
+    #[allow(dead_code)]
+    async fn spawn_with_name<A: Actor + RemoteActor>(
+        &self,
+        cx: &impl context::Actor,
+        name: Name,
+        params: &A::Params,
+    ) -> v1::Result<ActorMesh<A>>
     where
         A::Params: RemoteMessage,
     {
@@ -282,7 +340,6 @@ impl ProcMeshRef {
             .ok_or(Error::ActorTypeNotRegistered(type_name::<A>().to_string()))?
             .to_string();
 
-        let name = Name::new(name);
         let serialized_params = bincode::serialize(params)?;
 
         let (completed_handle, mut completed_receiver) = cx.mailbox().open_port();
@@ -292,7 +349,7 @@ impl ProcMeshRef {
                 .gspawn(
                     cx,
                     actor_type.clone(),
-                    name.clone().to_string(),
+                    name.to_string(),
                     serialized_params.clone(),
                     completed_handle.bind(),
                 )
@@ -352,7 +409,15 @@ impl view::RankedSliceable for ProcMeshRef {
             .unwrap()
             .map(|index| self.get(index).unwrap().clone())
             .collect();
-        Self::new(self.name.clone(), region, Arc::new(ranks)).unwrap()
+        Self::new(
+            self.name.clone(),
+            region,
+            Arc::new(ranks),
+            self.comm_actor_name.clone(),
+            Some(self.root_region.as_ref().unwrap_or(&self.region).clone()),
+            self.root_mesh_rank_0.clone(),
+        )
+        .unwrap()
     }
 }
 

hyperactor_mesh/src/v1/testactor.rs

@@ -1,3 +1,11 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 //! This module defines a test actor. It is defined in a separate module
 //! (outside of [`crate::v1::testing`]) to ensure that it is compiled into
 //! the bootstrap binary, which is not built in test mode (and anyway, test mode
@@ -20,14 +28,14 @@ use serde::Serialize;
 #[hyperactor::export(
     spawn = true,
     handlers = [
-        GetActorId,
+        GetActorId { cast = true },
     ]
 )]
 pub struct TestActor;
 
 /// A message that returns the recipient actor's id.
 #[derive(Debug, Clone, Named, Bind, Unbind, Serialize, Deserialize)]
-pub struct GetActorId(pub PortRef<ActorId>);
+pub struct GetActorId(#[binding(include)] pub PortRef<ActorId>);
 
 #[async_trait]
 impl Handler<GetActorId> for TestActor {