[15/n] Pass monarch tensors to actor endpoints, part 1

Pull Request resolved: #518

This makes it possible to send a monarch tensor to an actor endpoint that is defiend over the same proc mesh as the tensor.

The send is done locally so the actor can do work with the tensors.

The stream actor in the tensor engine is sent a SendResultOfActorCall message, which it will forward to the actor, binding to the message the real local tensors that were passed as arguments.

The stream actor that owns the tensor waits for called the actor to finish since the actor 'owns' the tensor through the duration of the call.

Known limitation: this message to the actor can go out of order w.r.t to other messages sent from the owner of the tensor engine because the real message is being sent from the stream actor. The next PR will fix this limitation by sending _both_ the tensor engine and the actor a message at the same time. The actor will get a 'wait for SendResultOfActorCall' message, at which point it will stop processing any messages except for the SentResultOfActorCall message it is suppose to be waiting for. This way the correct order is preserved from the perspective of the tensor engine stream and the actor.
ghstack-source-id: 295768514
@exported-using-ghexport

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

Author: zdevito

monarch_extension/src/convert.rs

@@ -176,6 +176,35 @@ impl<'a> MessageParser<'a> {
     fn parseWorkerMessageList(&self, name: &str) -> PyResult<Vec<WorkerMessage>> {
         self.attr(name)?.try_iter()?.map(|x| convert(x?)).collect()
     }
+    #[allow(non_snake_case)]
+    fn parseLocalStateList(&self, name: &str) -> PyResult<Vec<worker::LocalState>> {
+        let local_state_list = self.attr(name)?;
+        let mut result = Vec::new();
+
+        // Get the PyMailbox class for type checking
+        let mailbox_class = self
+            .current
+            .py()
+            .import("monarch._rust_bindings.monarch_hyperactor.mailbox")?
+            .getattr("Mailbox")?;
+
+        for item in local_state_list.try_iter()? {
+            let item = item?;
+            // Check if it's a Ref by trying to extract it
+            if let Ok(ref_val) = create_ref(item.clone()) {
+                result.push(worker::LocalState::Ref(ref_val));
+            } else if item.is_instance(&mailbox_class)? {
+                // It's a PyMailbox instance
+                result.push(worker::LocalState::Mailbox);
+            } else {
+                return Err(PyValueError::new_err(format!(
+                    "Expected Ref or Mailbox in local_state, got: {}",
+                    item.get_type().name()?
+                )));
+            }
+        }
+        Ok(result)
+    }
     fn parse_error_reason(&self, name: &str) -> PyResult<Option<(Option<ActorId>, String)>> {
         let err = self.attr(name)?;
         if err.is_none() {
@@ -432,6 +461,19 @@ fn create_map(py: Python) -> HashMap<u64, FnType> {
             p.parseWorkerMessageList("commands")?,
         ))
     });
+    m.insert(key("SendResultOfActorCall"), |p| {
+        Ok(WorkerMessage::SendResultOfActorCall(
+            worker::ActorCallParams {
+                seq: p.parseSeq("seq")?,
+                actor: p.parse("actor")?,
+                index: p.parse("actor_index")?,
+                python_message: p.parse("python_message")?,
+                local_state: p.parseLocalStateList("local_state")?,
+                mutates: p.parseRefList("mutates")?,
+                stream: p.parseStreamRef("stream")?,
+            },
+        ))
+    });
     m
 }
 

monarch_extension/src/tensor_worker.rs

@@ -1372,6 +1372,7 @@ pub(crate) fn worker_message_to_py(py: Python<'_>, message: &WorkerMessage) -> P
         }
         WorkerMessage::SetRefUnitTestsOnly { .. } => unimplemented!(),
         WorkerMessage::GetRefUnitTestsOnly { .. } => unimplemented!(),
+        WorkerMessage::SendResultOfActorCall { .. } => unimplemented!(),
     }
 }
 

monarch_hyperactor/src/actor.rs

@@ -175,10 +175,10 @@ impl PickledMessageClientActor {
 #[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")]
 #[derive(Default, Clone, Serialize, Deserialize, Named, PartialEq)]
 pub struct PythonMessage {
-    pub(crate) method: String,
-    pub(crate) message: ByteBuf,
-    response_port: Option<EitherPortRef>,
-    rank: Option<usize>,
+    pub method: String,
+    pub message: ByteBuf,
+    pub response_port: Option<EitherPortRef>,
+    pub rank: Option<usize>,
 }
 
 impl PythonMessage {
@@ -289,7 +289,7 @@ impl PythonActorHandle {
         PythonMessage { cast = true },
     ],
 )]
-pub(super) struct PythonActor {
+pub struct PythonActor {
     /// The Python object that we delegate message handling to. An instance of
     /// `monarch.actor_mesh._Actor`.
     pub(super) actor: PyObject,
@@ -399,16 +399,30 @@ impl PanicFlag {
 }
 
 #[async_trait]
-impl Handler<PythonMessage> for PythonActor {
-    async fn handle(&mut self, cx: &Context<Self>, message: PythonMessage) -> anyhow::Result<()> {
-        let mailbox = PyMailbox {
+impl Handler<LocalPythonMessage> for PythonActor {
+    async fn handle(
+        &mut self,
+        cx: &Context<Self>,
+        message: LocalPythonMessage,
+    ) -> anyhow::Result<()> {
+        let mailbox: PyMailbox = PyMailbox {
             inner: cx.mailbox_for_py().clone(),
         };
         // Create a channel for signaling panics in async endpoints.
         // See [Panics in async endpoints].
         let (sender, receiver) = oneshot::channel();
 
         let future = Python::with_gil(|py| -> Result<_, SerializablePyErr> {
+            let mailbox = mailbox.into_bound_py_any(py).unwrap();
+            let local_state: Option<Vec<Bound<'_, PyAny>>> = message.local_state.map(|state| {
+                state
+                    .into_iter()
+                    .map(|e| match e {
+                        LocalState::Mailbox => mailbox.clone(),
+                        LocalState::PyObject(obj) => obj.into_bound_py_any(py).unwrap(),
+                    })
+                    .collect()
+            });
             let (rank, shape) = cx.cast_info();
             let awaitable = self.actor.call_method(
                 py,
@@ -417,10 +431,11 @@ impl Handler<PythonMessage> for PythonActor {
                     mailbox,
                     rank,
                     PyShape::from(shape),
-                    message,
+                    message.message,
                     PanicFlag {
                         sender: Some(sender),
                     },
+                    local_state,
                 ),
                 None,
             )?;
@@ -439,6 +454,20 @@ impl Handler<PythonMessage> for PythonActor {
     }
 }
 
+#[async_trait]
+impl Handler<PythonMessage> for PythonActor {
+    async fn handle(&mut self, cx: &Context<Self>, message: PythonMessage) -> anyhow::Result<()> {
+        self.handle(
+            cx,
+            LocalPythonMessage {
+                message,
+                local_state: None,
+            },
+        )
+        .await
+    }
+}
+
 /// Helper struct to make a Python future passable in an actor message.
 ///
 /// Also so that we don't have to write this massive type signature everywhere
@@ -538,6 +567,17 @@ impl AsyncEndpointInvocationHandler for AsyncEndpointTask {
         Ok(())
     }
 }
+#[derive(Debug)]
+pub enum LocalState {
+    Mailbox,
+    PyObject(PyObject),
+}
+
+#[derive(Debug)]
+pub struct LocalPythonMessage {
+    pub message: PythonMessage,
+    pub local_state: Option<Vec<LocalState>>,
+}
 
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     hyperactor_mod.add_class::<PickledMessage>()?;

monarch_hyperactor/src/mailbox.rs

@@ -313,7 +313,7 @@ impl PythonUndeliverablePortHandle {
     name = "PortRef",
     module = "monarch._rust_bindings.monarch_hyperactor.mailbox"
 )]
-pub(super) struct PythonPortRef {
+pub struct PythonPortRef {
     pub(crate) inner: PortRef<PythonMessage>,
 }
 
@@ -453,7 +453,7 @@ impl PythonOncePortHandle {
     name = "OncePortRef",
     module = "monarch._rust_bindings.monarch_hyperactor.mailbox"
 )]
-pub(crate) struct PythonOncePortRef {
+pub struct PythonOncePortRef {
     pub(crate) inner: Option<OncePortRef<PythonMessage>>,
 }
 

monarch_messages/Cargo.toml

@@ -12,6 +12,7 @@ anyhow = "1.0.98"
 derive_more = { version = "1.0.0", features = ["full"] }
 enum-as-inner = "0.6.0"
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
+monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" }
 monarch_types = { version = "0.0.0", path = "../monarch_types" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
 pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] }

monarch_messages/src/worker.rs

@@ -27,6 +27,7 @@ use hyperactor::Named;
 use hyperactor::RefClient;
 use hyperactor::Unbind;
 use hyperactor::reference::ActorId;
+use monarch_hyperactor::actor::PythonMessage;
 use monarch_types::SerializablePyErr;
 use ndslice::Slice;
 use pyo3::exceptions::PyValueError;
@@ -402,6 +403,33 @@ pub struct CallFunctionParams {
     pub remote_process_groups: Vec<Ref>,
 }
 
+/// The local state that has to be restored into the python_message during
+/// its unpickling.
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum LocalState {
+    Ref(Ref),
+    Mailbox,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ActorCallParams {
+    pub seq: Seq,
+    /// The actor to call is (proc_id of stream worker, 'actor', 'index')
+    pub actor: String,
+    pub index: usize,
+    /// The PythonMessage object to send to the actor.
+    pub python_message: PythonMessage,
+    /// Referenceable objects to pass to the actor,
+    /// these will be put into the PythonMessage
+    /// during its unpickling. Because unpickling also needs to
+    /// restore mailboxes, we also have to keep track of which
+    /// members should just be the mailbox object.
+    pub local_state: Vec<LocalState>,
+
+    /// Tensors that will be mutated by the call.
+    pub mutates: Vec<Ref>,
+    pub stream: StreamRef,
+}
 /// Type of reduction for [`WorkerMessage::Reduce`].
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum Reduction {
@@ -642,21 +670,30 @@ pub enum WorkerMessage {
 
     /// First use of the borrow on the receiving stream. This is a marker for
     /// synchronization.
-    BorrowFirstUse { borrow: u64 },
+    BorrowFirstUse {
+        borrow: u64,
+    },
 
     /// Last use of the borrow on the receiving stream. This is a marker for
     /// synchronization.
-    BorrowLastUse { borrow: u64 },
+    BorrowLastUse {
+        borrow: u64,
+    },
 
     /// Drop the borrow and free the resources associated with it.
-    BorrowDrop { borrow: u64 },
+    BorrowDrop {
+        borrow: u64,
+    },
 
     /// Delete these refs from the worker state.
     DeleteRefs(Vec<Ref>),
 
     /// A [`ControllerMessage::Status`] will be send to the controller
     /// when all streams have processed all the message sent before this one.
-    RequestStatus { seq: Seq, controller: bool },
+    RequestStatus {
+        seq: Seq,
+        controller: bool,
+    },
 
     /// Perform a reduction operation, using an efficient communication backend.
     /// Only NCCL is supported for now.
@@ -758,6 +795,7 @@ pub enum WorkerMessage {
         stream: StreamRef,
     },
 
+    SendResultOfActorCall(ActorCallParams),
     PipeRecv {
         seq: Seq,
         /// Result refs.

monarch_simulator/src/worker.rs

@@ -312,6 +312,14 @@ impl WorkerMessageHandler for WorkerActor {
         Ok(())
     }
 
+    async fn send_result_of_actor_call(
+        &mut self,
+        cx: &hyperactor::Context<Self>,
+        params: ActorCallParams,
+    ) -> Result<()> {
+        bail!("unimplemented: send_result_of_actor_call");
+    }
+
     async fn command_group(
         &mut self,
         cx: &hyperactor::Context<Self>,

monarch_tensor_worker/src/lib.rs

@@ -72,6 +72,7 @@ use monarch_messages::controller::ControllerActor;
 use monarch_messages::controller::ControllerMessageClient;
 use monarch_messages::controller::Seq;
 use monarch_messages::wire_value::WireValue;
+use monarch_messages::worker::ActorCallParams;
 use monarch_messages::worker::CallFunctionError;
 use monarch_messages::worker::CallFunctionParams;
 use monarch_messages::worker::Factory;
@@ -860,6 +861,17 @@ impl WorkerMessageHandler for WorkerActor {
             .await
     }
 
+    async fn send_result_of_actor_call(
+        &mut self,
+        cx: &hyperactor::Context<Self>,
+        params: ActorCallParams,
+    ) -> Result<()> {
+        let stream = self.try_get_stream(params.stream)?;
+        stream
+            .send_result_of_actor_call(cx, cx.self_id().clone(), params)
+            .await?;
+        Ok(())
+    }
     async fn split_comm(
         &mut self,
         cx: &hyperactor::Context<Self>,