Pass cast rank to python actor (meta-pytorch#747)

Summary:



Rollback Plan:

Differential Revision: D79530146

Author: pzhan9

hyperactor_mesh/src/actor_mesh.rs

@@ -67,6 +67,8 @@ pub(crate) fn actor_mesh_cast<A, M>(
     selection_of_root: Selection,
     root_mesh_shape: &Shape,
     message: M,
+    // None means cast to root mesh.
+    cast_mesh_shape: Option<&Shape>,
 ) -> Result<(), CastError>
 where
     A: RemoteActor + RemoteHandles<IndexedErasedUnbound<M>>,
@@ -80,13 +82,16 @@ where
     let message = CastMessageEnvelope::new::<A, M>(
         actor_mesh_id,
         sender.clone(),
-        root_mesh_shape.clone(),
+        cast_mesh_shape.unwrap_or(root_mesh_shape).clone(),
         message,
     )?;
 
     comm_actor_ref.send(
         caps,
         CastMessage {
+            // Note: `dest` is on the root mesh' shape, which could be different
+            // from the cast mesh's shape if the cast is on a view, e.g. a sliced
+            // mesh.
             dest: Uslice {
                 slice: root_mesh_shape.slice().clone(),
                 selection: selection_of_root,
@@ -136,6 +141,7 @@ where
         sel_of_root,
         root_mesh_shape,
         message,
+        Some(sliced_shape),
     )
 }
 
@@ -161,6 +167,7 @@ pub trait ActorMesh: Mesh<Id = ActorMeshId> {
             selection,                            // the selected actors
             self.shape(),                         // root mesh shape
             message,                              // the message
+            None,                                 // cast mesh shape
         )
     }
 
@@ -407,7 +414,7 @@ impl<A: RemoteActor> ActorMesh for SlicedActorMesh<'_, A> {
             /*sel_of_sliced=*/ &sel,
             /*message=*/ message,
             /*sliced_shape=*/ self.shape(),
-            /*base_shape=*/ self.0.shape(),
+            /*roo_mesh_shape=*/ self.0.shape(),
         )
     }
 }

hyperactor_mesh/src/comm.rs

@@ -249,11 +249,14 @@ impl CommActor {
 
         // Deliver message here, if necessary.
         if deliver_here {
+            let rank_on_root_mesh = mode.self_rank(cx.self_id());
+            let cast_rank = message.relative_rank(rank_on_root_mesh)?;
+            let cast_shape = message.shape();
             let mut headers = cx.headers().clone();
             set_cast_info_on_headers(
                 &mut headers,
-                mode.self_rank(cx.self_id()),
-                message.shape().clone(),
+                cast_rank,
+                cast_shape.clone(),
                 message.sender().clone(),
             );
             cx.post(

hyperactor_mesh/src/comm/multicast.rs

@@ -21,6 +21,7 @@ use hyperactor::message::Castable;
 use hyperactor::message::ErasedUnbound;
 use hyperactor::message::IndexedErasedUnbound;
 use hyperactor::reference::ActorId;
+use ndslice::Extent;
 use ndslice::Shape;
 use ndslice::Slice;
 use ndslice::selection::Selection;
@@ -120,6 +121,15 @@ impl CastMessageEnvelope {
         &self.shape
     }
 
+    /// Given a rank in the root shape, return the corresponding point in the
+    /// provided shape, which is a view of the root shape.
+    pub(crate) fn relative_rank(&self, rank_on_root_mesh: usize) -> anyhow::Result<usize> {
+        let shape = self.shape();
+        let coords = shape.slice().coordinates(rank_on_root_mesh)?;
+        let extent = Extent::new(shape.labels().to_vec(), shape.slice().sizes().to_vec())?;
+        Ok(extent.point(coords)?.rank())
+    }
+
     /// The unique key used to indicate the stream to which to deliver this message.
     /// Concretely, the comm actors along the path should use this key to manage
     /// sequence numbers and reorder buffers.
@@ -203,9 +213,14 @@ declare_attrs! {
     pub attr CAST_ORIGINATING_SENDER: ActorId;
 }
 
-pub fn set_cast_info_on_headers(headers: &mut Attrs, rank: usize, shape: Shape, sender: ActorId) {
-    headers.set(CAST_RANK, rank);
-    headers.set(CAST_SHAPE, shape);
+pub fn set_cast_info_on_headers(
+    headers: &mut Attrs,
+    cast_rank: usize,
+    cast_shape: Shape,
+    sender: ActorId,
+) {
+    headers.set(CAST_RANK, cast_rank);
+    headers.set(CAST_SHAPE, cast_shape);
     headers.set(CAST_ORIGINATING_SENDER, sender);
 }
 

hyperactor_mesh/src/reference.rs

@@ -148,6 +148,7 @@ impl<A: RemoteActor> ActorMeshRef<A> {
                 selection,
                 &self.root,
                 message,
+                None,
             ),
         }
     }

python/tests/_monarch/test_actor_mesh.py

@@ -7,7 +7,7 @@
 # pyre-unsafe
 
 import pickle
-from typing import Any, Callable, Coroutine, Iterable, List, TYPE_CHECKING
+from typing import Any, Callable, cast, Coroutine, Iterable, List, TYPE_CHECKING
 
 import monarch
 import pytest
@@ -57,6 +57,9 @@ async def allocate() -> ProcMesh:
 
 
 class MyActor:
+    def __init__(self) -> None:
+        self._root_rank: int = -1
+
     async def handle(
         self,
         mailbox: Mailbox,
@@ -68,8 +71,19 @@ async def handle(
         local_state: Iterable[Any],
         response_port: "PortProtocol[Any]",
     ) -> None:
-        assert rank is not None
-        response_port.send(f"rank: {rank}")
+        match method:
+            case MethodSpecifier.Init():
+                self._root_rank = rank
+                response_port.send(None)
+                return None
+            case MethodSpecifier.ReturnsResponse(name=_):
+                response_port.send(self._root_rank)
+                return None
+            case MethodSpecifier.ExplicitPort(name=_):
+                response_port.exception(
+                    NotImplementedError("ExplicitPort is not supported yet")
+                )
+                return None
 
 
 # TODO - re-enable after resolving T232206970
@@ -95,35 +109,70 @@ async def run() -> None:
     run()
 
 
-async def verify_cast(
+async def spawn_actor_mesh(proc_mesh: ProcMesh) -> PythonActorMesh:
+    actor_mesh = await proc_mesh.spawn_nonblocking("test", MyActor)
+    # init actors to record their root ranks
+    receiver: PortReceiver
+    handle, receiver = proc_mesh.client.open_port()
+    port_ref = handle.bind()
+
+    message = PythonMessage(
+        PythonMessageKind.CallMethod(MethodSpecifier.Init(), port_ref),
+        pickle.dumps(None),
+    )
+    actor_mesh.cast(Selection.all(), message)
+    # wait for init to complete
+    for _ in range(len(actor_mesh.shape.ndslice)):
+        await receiver.recv_task()
+
+    return actor_mesh
+
+
+async def cast_to_call(
+    actor_mesh: PythonActorMesh | PythonActorMeshRef,
+    mailbox: Mailbox,
+    message: PythonMessage,
+) -> None:
+    sel = Selection.all()
+    if isinstance(actor_mesh, PythonActorMesh):
+        actor_mesh.cast(sel, message)
+    elif isinstance(actor_mesh, PythonActorMeshRef):
+        actor_mesh.cast(mailbox, sel, message)
+
+
+async def verify_cast_to_call(
     actor_mesh: PythonActorMesh | PythonActorMeshRef,
     mailbox: Mailbox,
-    cast_ranks: List[int],
+    root_ranks: List[int],
 ) -> None:
     receiver: PortReceiver
     handle, receiver = mailbox.open_port()
     port_ref = handle.bind()
 
+    # Now send the real message
     message = PythonMessage(
         PythonMessageKind.CallMethod(MethodSpecifier.ReturnsResponse("echo"), port_ref),
         pickle.dumps("ping"),
     )
-    sel = Selection.from_string("*")
-    if isinstance(actor_mesh, PythonActorMesh):
-        actor_mesh.cast(sel, message)
-    elif isinstance(actor_mesh, PythonActorMeshRef):
-        actor_mesh.cast(mailbox, sel, message)
+    await cast_to_call(actor_mesh, mailbox, message)
 
     rcv_ranks = []
-    for _ in range(len(cast_ranks)):
+    for _ in range(len(root_ranks)):
         message = await receiver.recv_task()
         result_kind = message.kind
         assert isinstance(result_kind, PythonMessageKind.Result)
-        rank = result_kind.rank
-        assert rank is not None
-        rcv_ranks.append(rank)
-    rcv_ranks.sort()
-    assert rcv_ranks == cast_ranks
+        cast_rank = result_kind.rank
+        assert cast_rank is not None
+        root_rank = cast(int, pickle.loads(message.message))
+        rcv_ranks.append((cast_rank, root_rank))
+    rcv_ranks.sort(key=lambda pair: pair[0])
+    recv_cast_ranks, recv_root_ranks = zip(*rcv_ranks)
+    assert recv_root_ranks == tuple(
+        root_ranks
+    ), f"recv_root_ranks={recv_root_ranks}, root_ranks={tuple(root_ranks)}"
+    assert recv_cast_ranks == tuple(
+        range(len(root_ranks))
+    ), f"recv_cast_ranks={recv_cast_ranks}, root_ranks={tuple(root_ranks)}"
     # verify no more messages are received
     with pytest.raises(TimeoutError):
         await receiver.recv_task().with_timeout(1)
@@ -136,8 +185,8 @@ async def test_cast_handle() -> None:
     @run_on_tokio
     async def run() -> None:
         proc_mesh = await allocate()
-        actor_mesh = await proc_mesh.spawn_nonblocking("test", MyActor)
-        await verify_cast(actor_mesh, proc_mesh.client, list(range(3 * 8 * 8)))
+        actor_mesh = await spawn_actor_mesh(proc_mesh)
+        await verify_cast_to_call(actor_mesh, proc_mesh.client, list(range(3 * 8 * 8)))
 
         await proc_mesh.stop_nonblocking()
 
@@ -151,9 +200,11 @@ async def test_cast_ref() -> None:
     @run_on_tokio
     async def run() -> None:
         proc_mesh = await allocate()
-        actor_mesh = await proc_mesh.spawn_nonblocking("test", MyActor)
+        actor_mesh = await spawn_actor_mesh(proc_mesh)
         actor_mesh_ref = actor_mesh.bind()
-        await verify_cast(actor_mesh_ref, proc_mesh.client, list(range(3 * 8 * 8)))
+        await verify_cast_to_call(
+            actor_mesh_ref, proc_mesh.client, list(range(3 * 8 * 8))
+        )
 
         await proc_mesh.stop_nonblocking()
 
@@ -184,7 +235,7 @@ async def verify_slice(
     assert (
         sliced_shape.ranks() == replica_0_ranks + replica_1_ranks
     ), f"left is {sliced_shape.ranks()}"
-    await verify_cast(sliced_mesh, mailbox, sliced_shape.ranks())
+    await verify_cast_to_call(sliced_mesh, mailbox, sliced_shape.ranks())
 
     assert sliced_shape.labels == ["replicas", "hosts", "gpus"]
     assert sliced_shape.ndslice.sizes == [2, 4, 3]
@@ -224,7 +275,8 @@ async def test_slice_actor_mesh_handle() -> None:
     @run_on_tokio
     async def run() -> None:
         proc_mesh = await allocate()
-        actor_mesh = await proc_mesh.spawn_nonblocking("test", MyActor)
+        actor_mesh = await spawn_actor_mesh(proc_mesh)
+
         await verify_slice(actor_mesh, proc_mesh.client)
 
         await proc_mesh.stop_nonblocking()
@@ -239,7 +291,8 @@ async def test_slice_actor_mesh_ref() -> None:
     @run_on_tokio
     async def run() -> None:
         proc_mesh = await allocate()
-        actor_mesh = await proc_mesh.spawn_nonblocking("test", MyActor)
+        actor_mesh = await spawn_actor_mesh(proc_mesh)
+
         actor_mesh_ref = actor_mesh.bind()
         await verify_slice(actor_mesh_ref, proc_mesh.client)