Pass comm_actor_ref in ActorMeshRef (#426)

Summary:
Pull Request resolved: #426

tl;dr `comm actor`'s world id can be different from `ProcMesh`'s `world_id`. We should not attest it.

`Comm actor`, which is spawned through `ProcessAllocator`, gets its world ID from a uuid:

https://www.internalfb.com/code/fbsource/[a29e91672b216c626792ec9406e77a922b9d88dc]/fbcode/monarch/hyperactor_mesh/src/alloc/process.rs?lines=108-109%2C420

`ProcMesh` gets its `world_id` from the `Alloc` it uses.

https://www.internalfb.com/code/fbsource/[a29e91672b216c626792ec9406e77a922b9d88dc]/fbcode/monarch/hyperactor_mesh/src/proc_mesh.rs?lines=95%2C99-101%2C300%2C302%2C317

`MastAllocator` uses `RemoteProcessAlloc`, and explicitly uses the `task_group_name` as its `world_id`.

https://www.internalfb.com/code/fbsource/[a29e91672b216c626792ec9406e77a922b9d88dc]/fbcode/monarch/hyperactor_meta/src/alloc.rs?lines=274-275%2C284%2C394%2C397%2C399

For example, this is a log I get from running a test. `ProcMesh` has `test_task_group`, comm actor has `_1C8Rf4TR6jZe`.

> I0703 09:09:51.966694 416137 fbcode/monarch/hyperactor_mesh/src/actor_mesh.rs:154] binding actor mesh ProcMesh { world_id: WorldId("test_task_group"), shape: Shape { labels: ["hosts", "gpus"], slice: Slice { offset: 0, sizes: [1, 2], strides: [2, 1] } }, ranks: [(ProcId(WorldId("_1C8Rf4TR6jZe"), 0), (Unix(Bound("vujWCqQQ55kCU2NQQlCC2q3N" (abstract))), ActorRef { actor_id: ActorId(ProcId(WorldId("_1C8Rf4TR6jZe"), 0), "mesh", 0), phantom: PhantomData<hyperactor_mesh::proc_mesh::mesh_agent::MeshAgent> })), (ProcId(WorldId("_1C8Rf4TR6jZe"), 1), (Unix(Bound("vujWCqQQ55kCU2NQQlCC2q3N" (abstract))), ActorRef { actor_id: ActorId(ProcId(WorldId("_1C8Rf4TR6jZe"), 1), "mesh", 0), phantom: PhantomData<hyperactor_mesh::proc_mesh::mesh_agent::MeshAgent> }))], client: Mailbox { inner: State { actor_id: ActorId(ProcId(WorldId("test_task_group_manager"), 0), "client", 0), open_ports: [17082645012990790806], next_port: 1036 } }, comm_actors: [ActorRef { actor_id: ActorId(ProcId(WorldId("_1C8Rf4TR6jZe"), 0), "comm", 0), phantom: PhantomData<hyperactor_mesh::comm::CommActor> }, ActorRef { actor_id: ActorId(ProcId(WorldId("_1C8Rf4TR6jZe"), 1), "comm", 0), phantom: PhantomData<hyperactor_mesh::comm::CommActor> }] }

Reviewed By: mariusae

Differential Revision: D77737951

fbshipit-source-id: 4f70aad14e984b1bc89bea79bf3a3ed8218048e0

Author: pzhan9

hyperactor_mesh/src/actor_mesh.rs

@@ -155,6 +155,7 @@ pub trait ActorMesh: Mesh {
             ),
             self.shape().clone(),
             self.proc_mesh().shape().clone(),
+            self.proc_mesh().comm_actor().clone(),
         )
     }
 }

hyperactor_mesh/src/reference.rs

@@ -11,12 +11,9 @@ use std::cmp::PartialOrd;
 use std::hash::Hash;
 use std::marker::PhantomData;
 
-use hyperactor::ActorId;
 use hyperactor::ActorRef;
 use hyperactor::Named;
-use hyperactor::ProcId;
 use hyperactor::RemoteHandles;
-use hyperactor::WorldId;
 use hyperactor::actor::RemoteActor;
 use hyperactor::cap;
 use hyperactor::message::Castable;
@@ -79,6 +76,8 @@ pub struct ActorMeshRef<A: RemoteActor> {
     shape: Shape,
     /// The shape of the underlying Proc Mesh.
     proc_mesh_shape: Shape,
+    /// The reference to the comm actor of the underlying Proc Mesh.
+    comm_actor_ref: ActorRef<CommActor>,
     phantom: PhantomData<A>,
 }
 
@@ -87,11 +86,17 @@ impl<A: RemoteActor> ActorMeshRef<A> {
     /// typed reference.  This is usually invoked to provide a guarantee
     /// that an externally-provided mesh ID (e.g., through a command
     /// line argument) is a valid reference.
-    pub(crate) fn attest(mesh_id: ActorMeshId, shape: Shape, proc_mesh_shape: Shape) -> Self {
+    pub(crate) fn attest(
+        mesh_id: ActorMeshId,
+        shape: Shape,
+        proc_mesh_shape: Shape,
+        comm_actor_ref: ActorRef<CommActor>,
+    ) -> Self {
         Self {
             mesh_id,
             shape,
             proc_mesh_shape,
+            comm_actor_ref,
             phantom: PhantomData,
         }
     }
@@ -127,21 +132,16 @@ impl<A: RemoteActor> ActorMeshRef<A> {
     where
         A: RemoteHandles<M> + RemoteHandles<IndexedErasedUnbound<M>>,
     {
-        let world_id = WorldId(self.mesh_id.0.0.clone());
-        let comm_actor_id = ActorId(ProcId(world_id, 0), "comm".to_string(), 0);
-
         actor_mesh_cast::<M, A>(
             caps,
             self.shape(),
             self.proc_mesh_shape(),
             self.name(),
             caps.mailbox().actor_id(),
-            &ActorRef::<CommActor>::attest(comm_actor_id),
+            &self.comm_actor_ref,
             selection,
             message,
-        )?;
-
-        Ok(())
+        )
     }
 }
 
@@ -151,6 +151,7 @@ impl<A: RemoteActor> Clone for ActorMeshRef<A> {
             mesh_id: self.mesh_id.clone(),
             shape: self.shape.clone(),
             proc_mesh_shape: self.proc_mesh_shape.clone(),
+            comm_actor_ref: self.comm_actor_ref.clone(),
             phantom: PhantomData,
         }
     }
@@ -211,6 +212,7 @@ mod tests {
         mesh_id: ActorMeshId,
         shape: Shape,
         proc_mesh_shape: Shape,
+        comm_actor_ref: ActorRef<CommActor>,
     }
 
     #[async_trait]
@@ -223,6 +225,7 @@ mod tests {
                     params.mesh_id,
                     params.shape,
                     params.proc_mesh_shape,
+                    params.comm_actor_ref,
                 ),
             })
         }
@@ -284,6 +287,7 @@ mod tests {
                     ),
                     shape: ping_proc_mesh.shape().clone(),
                     proc_mesh_shape: ping_proc_mesh.shape().clone(),
+                    comm_actor_ref: ping_proc_mesh.comm_actor().clone(),
                 },
             )
             .await
@@ -301,6 +305,7 @@ mod tests {
                     ),
                     shape: pong_proc_mesh.shape().clone(),
                     proc_mesh_shape: pong_proc_mesh.shape().clone(),
+                    comm_actor_ref: pong_proc_mesh.comm_actor().clone(),
                 },
             )
             .await