[16/n] tensor engine: enum-ify PythonMessage

PythonMessage was getting a bit stringly typed. I've sectioned into 4 enum case (CallMethod, Result, Exception, Uninit) which capture when you have to pay attention to things like rank/response_port.

I have to further add functionality in a follow-up to get actor ordering correct.

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

ghstack-source-id: 296136316
Pull Request resolved: #529

Author: zdevito

monarch_extension/src/mesh_controller.rs

@@ -37,6 +37,7 @@ use hyperactor_mesh::actor_mesh::RootActorMesh;
 use hyperactor_mesh::shared_cell::SharedCell;
 use hyperactor_mesh::shared_cell::SharedCellRef;
 use monarch_hyperactor::actor::PythonMessage;
+use monarch_hyperactor::actor::PythonMessageKind;
 use monarch_hyperactor::mailbox::PyPortId;
 use monarch_hyperactor::ndslice::PySlice;
 use monarch_hyperactor::proc_mesh::PyProcMesh;
@@ -334,7 +335,7 @@ impl Invocation {
                             Some(PortInfo { port, ranks }) => {
                                 *unreported_exception = None;
                                 for rank in ranks.iter() {
-                                    let msg = exception.as_ref().clone().with_rank(rank);
+                                    let msg = exception.as_ref().clone().into_rank(rank);
                                     port.send(sender, msg)?;
                                 }
                             }
@@ -527,7 +528,7 @@ impl History {
                 .call1((exception.backtrace, traceback, rank))
                 .unwrap();
             let data: Vec<u8> = pickle.call1((exe,)).unwrap().extract().unwrap();
-            PythonMessage::new_from_buf("exception".to_string(), data, None, Some(rank))
+            PythonMessage::new_from_buf(PythonMessageKind::Exception { rank: Some(rank) }, data)
         }));
 
         let mut invocation = invocation.lock().unwrap();
@@ -570,7 +571,10 @@ impl History {
                     Some(exception) => exception.as_ref().clone(),
                     None => {
                         // the byte string is just a Python None
-                        PythonMessage::new("result".to_string(), b"\x80\x04N.", None, None)
+                        PythonMessage::new_from_buf(
+                            PythonMessageKind::Result { rank: None },
+                            b"\x80\x04N.".to_vec(),
+                        )
                     }
                 };
                 port.send(sender, result)?;

monarch_hyperactor/src/actor.rs

@@ -171,41 +171,60 @@ impl PickledMessageClientActor {
     }
 }
 
+#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
+#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq)]
+pub enum PythonMessageKind {
+    CallMethod {
+        name: String,
+        response_port: Option<EitherPortRef>,
+    },
+    Result {
+        rank: Option<usize>,
+    },
+    Exception {
+        rank: Option<usize>,
+    },
+    Uninit {},
+}
+
+impl Default for PythonMessageKind {
+    fn default() -> Self {
+        PythonMessageKind::Uninit {}
+    }
+}
+
 #[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")]
-#[derive(Default, Clone, Serialize, Deserialize, Named, PartialEq)]
+#[derive(Clone, Serialize, Deserialize, Named, PartialEq, Default)]
 pub struct PythonMessage {
-    pub method: String,
-    pub message: ByteBuf,
-    pub response_port: Option<EitherPortRef>,
-    pub rank: Option<usize>,
+    pub kind: PythonMessageKind,
+    pub message: Vec<u8>,
 }
 
 impl PythonMessage {
-    pub fn with_rank(self, rank: usize) -> PythonMessage {
-        PythonMessage {
-            rank: Some(rank),
-            ..self
-        }
+    pub fn new_from_buf(kind: PythonMessageKind, message: Vec<u8>) -> Self {
+        Self { kind, message }
     }
-    pub fn new_from_buf(
-        method: String,
-        message: Vec<u8>,
-        response_port: Option<EitherPortRef>,
-        rank: Option<usize>,
-    ) -> Self {
-        Self {
-            method,
-            message: message.into(),
-            response_port,
-            rank,
+
+    pub fn into_rank(self, rank: usize) -> Self {
+        let rank = Some(rank);
+        match self.kind {
+            PythonMessageKind::Result { .. } => PythonMessage {
+                kind: PythonMessageKind::Result { rank },
+                message: self.message,
+            },
+            PythonMessageKind::Exception { .. } => PythonMessage {
+                kind: PythonMessageKind::Exception { rank },
+                message: self.message,
+            },
+            _ => panic!("PythonMessage is not a response but {:?}", self),
         }
     }
 }
 
 impl std::fmt::Debug for PythonMessage {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("PythonMessage")
-            .field("method", &self.method)
+            .field("kind", &self.kind)
             .field(
                 "message",
                 &hyperactor::data::HexFmt(self.message.as_slice()).to_string(),
@@ -216,48 +235,39 @@ impl std::fmt::Debug for PythonMessage {
 
 impl Unbind for PythonMessage {
     fn unbind(&self, bindings: &mut Bindings) -> anyhow::Result<()> {
-        self.response_port.unbind(bindings)
+        match &self.kind {
+            PythonMessageKind::CallMethod { response_port, .. } => response_port.unbind(bindings),
+            _ => Ok(()),
+        }
     }
 }
 
 impl Bind for PythonMessage {
     fn bind(&mut self, bindings: &mut Bindings) -> anyhow::Result<()> {
-        self.response_port.bind(bindings)
+        match &mut self.kind {
+            PythonMessageKind::CallMethod { response_port, .. } => response_port.bind(bindings),
+            _ => Ok(()),
+        }
     }
 }
 
 #[pymethods]
 impl PythonMessage {
     #[new]
-    #[pyo3(signature = (method, message, response_port, rank))]
-    pub fn new(
-        method: String,
-        message: &[u8],
-        response_port: Option<EitherPortRef>,
-        rank: Option<usize>,
-    ) -> Self {
-        Self::new_from_buf(method, message.into(), response_port, rank)
+    #[pyo3(signature = (kind, message))]
+    pub fn new(kind: PythonMessageKind, message: &[u8]) -> Self {
+        PythonMessage::new_from_buf(kind, message.to_vec())
     }
 
     #[getter]
-    fn method(&self) -> &String {
-        &self.method
+    fn kind(&self) -> PythonMessageKind {
+        self.kind.clone()
     }
 
     #[getter]
     fn message<'a>(&self, py: Python<'a>) -> Bound<'a, PyBytes> {
         PyBytes::new(py, self.message.as_ref())
     }
-
-    #[getter]
-    fn response_port(&self) -> Option<EitherPortRef> {
-        self.response_port.clone()
-    }
-
-    #[getter]
-    fn rank(&self) -> Option<usize> {
-        self.rank
-    }
 }
 
 #[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
@@ -583,6 +593,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PickledMessageClientActor>()?;
     hyperactor_mod.add_class::<PythonActorHandle>()?;
     hyperactor_mod.add_class::<PythonMessage>()?;
+    hyperactor_mod.add_class::<PythonMessageKind>()?;
     hyperactor_mod.add_class::<PanicFlag>()?;
     Ok(())
 }

monarch_hyperactor/src/mailbox.rs

@@ -44,6 +44,7 @@ use serde::Deserialize;
 use serde::Serialize;
 
 use crate::actor::PythonMessage;
+use crate::actor::PythonMessageKind;
 use crate::proc::PyActorId;
 use crate::runtime::signal_safe_block_on;
 use crate::shape::PyShape;
@@ -528,7 +529,8 @@ impl PythonOncePortReceiver {
     Named,
     PartialEq,
     FromPyObject,
-    IntoPyObject
+    IntoPyObject,
+    Debug
 )]
 pub enum EitherPortRef {
     Unbounded(PythonPortRef),
@@ -610,7 +612,7 @@ impl Accumulator for PythonAccumulator {
     fn accumulate(&self, state: &mut Self::State, update: Self::Update) -> anyhow::Result<()> {
         Python::with_gil(|py: Python<'_>| {
             // Initialize state if it is empty.
-            if state.message.is_empty() && state.method.is_empty() {
+            if matches!(state.kind, PythonMessageKind::Uninit {}) {
                 *state = self
                     .accumulator
                     .getattr(py, "initial_state")?

monarch_tensor_worker/src/stream.rs

@@ -39,6 +39,7 @@ use hyperactor::proc::Proc;
 use monarch_hyperactor::actor::LocalPythonMessage;
 use monarch_hyperactor::actor::PythonActor;
 use monarch_hyperactor::actor::PythonMessage;
+use monarch_hyperactor::actor::PythonMessageKind;
 use monarch_hyperactor::mailbox::EitherPortRef;
 use monarch_messages::controller::ControllerMessageClient;
 use monarch_messages::controller::Seq;
@@ -1009,10 +1010,10 @@ impl StreamActor {
                         .extract()
                         .unwrap();
                     Ok(PythonMessage::new_from_buf(
-                        "result".to_string(),
+                        PythonMessageKind::Result {
+                            rank: Some(worker_actor_id.rank()),
+                        },
                         data,
-                        None,
-                        Some(worker_actor_id.rank()),
                     ))
                 })
         });
@@ -1686,12 +1687,13 @@ impl StreamMessageHandler for StreamActor {
         let (send, recv) = cx.open_once_port();
         let send = send.bind();
         let send = EitherPortRef::Once(send.into());
-        let message = PythonMessage {
-            method: params.method,
-            message: params.args_kwargs_tuple.into(),
-            response_port: Some(send),
-            rank: None,
-        };
+        let message = PythonMessage::new_from_buf(
+            PythonMessageKind::CallMethod {
+                name: params.method,
+                response_port: Some(send),
+            },
+            params.args_kwargs_tuple.into(),
+        );
         let local_state: Result<Vec<monarch_hyperactor::actor::LocalState>> =
             Python::with_gil(|py| {
                 params
@@ -1713,31 +1715,44 @@ impl StreamMessageHandler for StreamActor {
             local_state: Some(local_state?),
         };
         actor_handle.send(message)?;
-        let result = recv.recv().await?.with_rank(worker_actor_id.rank());
-        if result.method == "exception" {
-            // If result has "exception" as its kind, then
-            // we need to unpickle and turn it into a WorkerError
-            // and call remote_function_failed otherwise the
-            // controller assumes the object is correct and doesn't handle
-            // dependency tracking correctly.
-            let err = Python::with_gil(|py| -> Result<WorkerError, SerializablePyErr> {
-                let err = py
-                    .import("pickle")
-                    .unwrap()
-                    .call_method1("loads", (result.message.into_vec(),))?;
-                Ok(WorkerError {
-                    worker_actor_id,
-                    backtrace: err.to_string(),
-                })
-            })?;
-            self.controller_actor
-                .remote_function_failed(cx, params.seq, err)
-                .await?;
-        } else {
-            let result = Serialized::serialize(&result).unwrap();
-            self.controller_actor
-                .fetch_result(cx, params.seq, Ok(result))
-                .await?;
+        let result = recv.recv().await?;
+        match result.kind {
+            PythonMessageKind::Exception { .. } => {
+                // If result has "exception" as its kind, then
+                // we need to unpickle and turn it into a WorkerError
+                // and call remote_function_failed otherwise the
+                // controller assumes the object is correct and doesn't handle
+                // dependency tracking correctly.
+                let err = Python::with_gil(|py| -> Result<WorkerError, SerializablePyErr> {
+                    let err = py
+                        .import("pickle")
+                        .unwrap()
+                        .call_method1("loads", (result.message,))?;
+                    Ok(WorkerError {
+                        worker_actor_id,
+                        backtrace: err.to_string(),
+                    })
+                })?;
+                self.controller_actor
+                    .remote_function_failed(cx, params.seq, err)
+                    .await?;
+            }
+            PythonMessageKind::Result { .. } => {
+                let result = PythonMessage::new_from_buf(
+                    PythonMessageKind::Result {
+                        rank: Some(worker_actor_id.rank()),
+                    },
+                    result.message,
+                );
+                let result = Serialized::serialize(&result).unwrap();
+                self.controller_actor
+                    .fetch_result(cx, params.seq, Ok(result))
+                    .await?;
+            }
+            _ => panic!(
+                "Unexpected response kind from PythonActor: {:?}",
+                result.kind
+            ),
         }
         Ok(())
     }