[14/n] tensor engine {handle_cast, handle} -> handle

Pull Request resolved: #491

We unified the message handling logic in rust, so we can keep that unification for the python bindings. We keep the handle_cast variant, rename it handle.

We provide a definition for cast coordinates even when a message wasn't cast: the message is considered to be the  rank of a zero dimension shape.
ghstack-source-id: 295545965
@exported-using-ghexport

Differential Revision: [D78061501](https://our.internmc.facebook.com/intern/diff/D78061501/)

Author: zdevito

hyperactor_mesh/examples/dining_philosophers.rs

@@ -142,7 +142,7 @@ impl Handler<PhilosopherMessage> for PhilosopherActor {
         cx: &Context<Self>,
         message: PhilosopherMessage,
     ) -> Result<(), anyhow::Error> {
-        let (rank, _) = cx.cast_info()?;
+        let (rank, _) = cx.cast_info();
         self.rank = rank;
         match message {
             PhilosopherMessage::Start(waiter) => {

hyperactor_mesh/src/actor_mesh.rs

@@ -441,7 +441,7 @@ pub(crate) mod test_util {
             cx: &Context<Self>,
             GetRank(ok, reply): GetRank,
         ) -> Result<(), anyhow::Error> {
-            let (rank, _) = cx.cast_info()?;
+            let (rank, _) = cx.cast_info();
             reply.send(cx, rank)?;
             anyhow::ensure!(ok, "intentional error!"); // If `!ok` exit with `Err()`.
             Ok(())

hyperactor_mesh/src/comm/multicast.rs

@@ -215,32 +215,20 @@ pub fn set_cast_info_on_headers(headers: &mut Attrs, rank: usize, shape: Shape,
 }
 
 pub trait CastInfo {
-    /// Get the cast rank and cast shape, returning an error
-    /// if the relevant info isn't available.
-    fn cast_info(&self) -> anyhow::Result<(usize, Shape)>;
-
-    /// Get the cast rank and cast shape, returning None
-    /// if the relevant info isn't available.
-    fn maybe_cast_info(&self) -> Option<(usize, Shape)>;
+    /// Get the cast rank and cast shape.
+    /// If something wasn't explicitly sent via a cast, then
+    /// we represent it as the only member of a 0-dimensonal cast shape,
+    /// which is the same as a singleton.
+    fn cast_info(&self) -> (usize, Shape);
 }
 
 impl<A: Actor> CastInfo for Context<'_, A> {
-    fn cast_info(&self) -> anyhow::Result<(usize, Shape)> {
+    fn cast_info(&self) -> (usize, Shape) {
         let headers = self.headers();
-        let rank = headers
-            .get(CAST_RANK)
-            .ok_or_else(|| anyhow::anyhow!("{} not found in headers", CAST_RANK.name()))?;
-        let shape = headers
-            .get(CAST_SHAPE)
-            .ok_or_else(|| anyhow::anyhow!("{} not found in headers", CAST_SHAPE.name()))?
-            .clone();
-        Ok((*rank, shape))
-    }
-
-    fn maybe_cast_info(&self) -> Option<(usize, Shape)> {
-        let headers = self.headers();
-        headers
-            .get(CAST_RANK)
-            .map(|rank| headers.get(CAST_SHAPE).map(|shape| (*rank, shape.clone())))?
+        match (headers.get(CAST_RANK), headers.get(CAST_SHAPE)) {
+            (Some(rank), Some(shape)) => (*rank, shape.clone()),
+            (None, None) => (0, Shape::unity()),
+            _ => panic!("Expected either both rank and shape or neither"),
+        }
     }
 }

monarch_hyperactor/src/actor.rs

@@ -28,6 +28,7 @@ use hyperactor::message::Unbind;
 use hyperactor_mesh::comm::multicast::CastInfo;
 use monarch_types::PickledPyObject;
 use monarch_types::SerializablePyErr;
+use ndslice::Shape;
 use pyo3::conversion::IntoPyObjectExt;
 use pyo3::exceptions::PyBaseException;
 use pyo3::exceptions::PyRuntimeError;
@@ -408,34 +409,21 @@ impl Handler<PythonMessage> for PythonActor {
         let (sender, receiver) = oneshot::channel();
 
         let future = Python::with_gil(|py| -> Result<_, SerializablePyErr> {
-            let awaitable = match cx.maybe_cast_info() {
-                Some((rank, shape)) => self.actor.call_method(
-                    py,
-                    "handle_cast",
-                    (
-                        mailbox,
-                        rank,
-                        PyShape::from(shape),
-                        message,
-                        PanicFlag {
-                            sender: Some(sender),
-                        },
-                    ),
-                    None,
-                )?,
-                None => self.actor.call_method(
-                    py,
-                    "handle",
-                    (
-                        mailbox,
-                        message,
-                        PanicFlag {
-                            sender: Some(sender),
-                        },
-                    ),
-                    None,
-                )?,
-            };
+            let (rank, shape) = cx.cast_info();
+            let awaitable = self.actor.call_method(
+                py,
+                "handle",
+                (
+                    mailbox,
+                    rank,
+                    PyShape::from(shape),
+                    message,
+                    PanicFlag {
+                        sender: Some(sender),
+                    },
+                ),
+                None,
+            )?;
 
             pyo3_async_runtimes::into_future_with_locals(
                 self.get_task_locals(py),

monarch_rdma/examples/parameter_server.rs

@@ -303,7 +303,7 @@ impl Handler<WorkerInit> for WorkerActor {
         cx: &Context<Self>,
         WorkerInit(ps_ref, rdma_managers): WorkerInit,
     ) -> Result<(), anyhow::Error> {
-        let (rank, _) = cx.cast_info()?;
+        let (rank, _) = cx.cast_info();
 
         println!("[worker_actor_{}] initializing", rank);
 
@@ -336,7 +336,7 @@ impl Handler<WorkerStep> for WorkerActor {
         cx: &Context<Self>,
         WorkerStep(reply): WorkerStep,
     ) -> Result<(), anyhow::Error> {
-        let (rank, _) = cx.cast_info()?;
+        let (rank, _) = cx.cast_info();
 
         for (grad_value, weight) in self
             .local_gradients
@@ -385,7 +385,7 @@ impl Handler<WorkerUpdate> for WorkerActor {
         cx: &Context<Self>,
         WorkerUpdate(reply): WorkerUpdate,
     ) -> Result<(), anyhow::Error> {
-        let (rank, _) = cx.cast_info()?;
+        let (rank, _) = cx.cast_info();
 
         println!(
             "[worker_actor_{}] pulling new weights from parameter server (before: {:?})",
@@ -418,7 +418,7 @@ impl Handler<WorkerUpdate> for WorkerActor {
 impl Handler<Log> for WorkerActor {
     /// Logs the worker's weights
     async fn handle(&mut self, cx: &Context<Self>, _: Log) -> Result<(), anyhow::Error> {
-        let (rank, _) = cx.cast_info()?;
+        let (rank, _) = cx.cast_info();
         println!("[worker_actor_{}] weights: {:?}", rank, self.weights_data);
         Ok(())
     }

monarch_tensor_worker/src/lib.rs

@@ -266,7 +266,7 @@ impl Handler<AssignRankMessage> for WorkerActor {
         cx: &hyperactor::Context<Self>,
         _: AssignRankMessage,
     ) -> anyhow::Result<()> {
-        let (rank, shape) = cx.cast_info()?;
+        let (rank, shape) = cx.cast_info();
         self.rank = rank;
         self.respond_with_python_message = true;
         Python::with_gil(|py| {

python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi

@@ -179,9 +179,6 @@ class PanicFlag:
 
 class Actor(Protocol):
     async def handle(
-        self, mailbox: Mailbox, message: PythonMessage, panic_flag: PanicFlag
-    ) -> None: ...
-    async def handle_cast(
         self,
         mailbox: Mailbox,
         rank: int,

python/monarch/_src/actor/actor_mesh.py

@@ -580,11 +580,6 @@ def __init__(self) -> None:
         self.instance: object | None = None
 
     async def handle(
-        self, mailbox: Mailbox, message: PythonMessage, panic_flag: PanicFlag
-    ) -> None:
-        return await self.handle_cast(mailbox, 0, singleton_shape, message, panic_flag)
-
-    async def handle_cast(
         self,
         mailbox: Mailbox,
         rank: int,

python/tests/_monarch/test_hyperactor.py

@@ -27,11 +27,6 @@
 
 class MyActor:
     async def handle(
-        self, mailbox: Mailbox, message: PythonMessage, panic_flag: PanicFlag
-    ) -> None:
-        raise NotImplementedError()
-
-    async def handle_cast(
         self,
         mailbox: Mailbox,
         rank: int,

python/tests/_monarch/test_mailbox.py

@@ -118,11 +118,6 @@ async def recv_message() -> str:
 
 class MyActor:
     async def handle(
-        self, mailbox: Mailbox, message: PythonMessage, panic_flag: PanicFlag
-    ) -> None:
-        return None
-
-    async def handle_cast(
         self,
         mailbox: Mailbox,
         rank: int,