Plumbing: Update some send APIs on the python side with PythonInstance (#1362)

Summary:
Pull Request resolved: #1362

This diff is part of the effect to adding sequencing logic to sender actor. See D83371710 for details.

This diff specifically updates the send method in `PythonPortRef`, `PythonOncePortRef`, `PythonActorRef`, and `PythonActorHandle` to have the `instance: &PyInstance`, so it can be plumbed to the Rust backend later.

Reviewed By: mariusae, moonli

Differential Revision: D83413791

fbshipit-source-id: 0163c8f3dfc859ea6be8e39da552f9a0281f3418

Author: pzhan9

monarch_hyperactor/src/actor.rs

@@ -13,7 +13,6 @@ use std::sync::Arc;
 use std::sync::OnceLock;
 
 use async_trait::async_trait;
-use bytes::Bytes;
 use hyperactor::Actor;
 use hyperactor::ActorHandle;
 use hyperactor::ActorId;
@@ -50,9 +49,9 @@ use tokio::sync::mpsc::UnboundedSender;
 use tokio::sync::oneshot;
 use tracing::Instrument;
 
-use crate::buffers::Buffer;
 use crate::buffers::FrozenBuffer;
 use crate::config::SHARED_ASYNCIO_RUNTIME;
+use crate::context::PyInstance;
 use crate::local_state_broker::BrokerId;
 use crate::local_state_broker::LocalStateBrokerMessage;
 use crate::mailbox::EitherPortRef;
@@ -275,10 +274,9 @@ impl PythonMessage {
             _ => panic!("PythonMessage is not a response but {:?}", self),
         }
     }
-
-    async fn resolve_indirect_call<T: Actor>(
+    async fn resolve_indirect_call(
         self,
-        cx: &Context<'_, T>,
+        cx: &Context<'_, PythonActor>,
     ) -> anyhow::Result<ResolvedCallMethod> {
         match self.kind {
             PythonMessageKind::CallMethodIndirect {
@@ -331,6 +329,7 @@ impl PythonMessage {
                     .call_method1("repeat", (mailbox.clone(),))
                     .unwrap()
                     .unbind();
+                let instance: PyInstance = cx.into();
                 let response_port = response_port
                     .map_or_else(
                         || {
@@ -343,7 +342,7 @@ impl PythonMessage {
                             let point = cx.cast_point();
                             py.import("monarch._src.actor.actor_mesh")
                                 .unwrap()
-                                .call_method1("Port", (x, mailbox, point.rank()))
+                                .call_method1("Port", (x, instance, point.rank()))
                                 .unwrap()
                         },
                     )
@@ -435,7 +434,8 @@ pub(super) struct PythonActorHandle {
 #[pymethods]
 impl PythonActorHandle {
     // TODO: do the pickling in rust
-    fn send(&self, message: &PythonMessage) -> PyResult<()> {
+    // TODO(pzhang) Use instance after its required by PortHandle.
+    fn send(&self, _instance: &PyInstance, message: &PythonMessage) -> PyResult<()> {
         self.inner
             .send(message.clone())
             .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
@@ -708,7 +708,11 @@ impl Handler<HandlePanic> for PythonActorPanicWatcher {
 
 #[async_trait]
 impl Handler<PythonMessage> for PythonActor {
-    async fn handle(&mut self, cx: &Context<Self>, message: PythonMessage) -> anyhow::Result<()> {
+    async fn handle(
+        &mut self,
+        cx: &Context<PythonActor>,
+        message: PythonMessage,
+    ) -> anyhow::Result<()> {
         let resolved = message.resolve_indirect_call(cx).await?;
 
         // Create a channel for signaling panics in async endpoints.

monarch_hyperactor/src/actor_mesh.rs

@@ -618,7 +618,7 @@ impl ActorMeshProtocol for AsyncActorMesh {
                     let port = py
                         .import("monarch._src.actor.actor_mesh")
                         .unwrap()
-                        .call_method1("Port", (port_ref, instance._mailbox(), 0))
+                        .call_method1("Port", (port_ref, instance, 0))
                         .unwrap();
                     port.call_method1("exception", (pyerr.value(py),)).unwrap();
                 }),

monarch_hyperactor/src/mailbox.rs

@@ -47,6 +47,7 @@ use serde::Serialize;
 
 use crate::actor::PythonMessage;
 use crate::actor::PythonMessageKind;
+use crate::context::PyInstance;
 use crate::proc::PyActorId;
 use crate::pytokio::PyPythonTask;
 use crate::pytokio::PythonTask;
@@ -264,7 +265,8 @@ pub(super) struct PythonPortHandle {
 
 #[pymethods]
 impl PythonPortHandle {
-    fn send(&self, message: PythonMessage) -> PyResult<()> {
+    // TODO(pzhang) Use instance after its required by PortHandle.
+    fn send(&self, _instance: &PyInstance, message: PythonMessage) -> PyResult<()> {
         self.inner
             .send(message)
             .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
@@ -318,9 +320,9 @@ impl PythonPortRef {
         Ok((slf.get_type(), (id,)))
     }
 
-    fn send(&self, mailbox: &PyMailbox, message: PythonMessage) -> PyResult<()> {
+    fn send(&self, instance: &PyInstance, message: PythonMessage) -> PyResult<()> {
         self.inner
-            .send(&mailbox.inner, message)
+            .send(&instance._mailbox().inner, message)
             .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
         Ok(())
     }
@@ -530,13 +532,13 @@ impl PythonOncePortRef {
         Ok((slf.get_type(), (id,)))
     }
 
-    fn send(&mut self, mailbox: &PyMailbox, message: PythonMessage) -> PyResult<()> {
+    fn send(&mut self, instance: &PyInstance, message: PythonMessage) -> PyResult<()> {
         let Some(port_ref) = self.inner.take() else {
             return Err(PyErr::new::<PyValueError, _>("OncePortRef is already used"));
         };
 
         port_ref
-            .send(&mailbox.inner, message)
+            .send(&instance._mailbox().inner, message)
             .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
         Ok(())
     }

python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi

@@ -9,12 +9,10 @@
 from typing import final, Protocol
 
 from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage
-
+from monarch._rust_bindings.monarch_hyperactor.context import Instance
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
 
-from monarch._rust_bindings.monarch_hyperactor.shape import Shape
-
 @final
 class PortId:
     def __init__(self, *, actor_id: ActorId, port: int) -> None:
@@ -52,7 +50,7 @@ class PortHandle:
     A handle to a port over which PythonMessages can be sent.
     """
 
-    def send(self, message: PythonMessage) -> None:
+    def send(self, instance: Instance, message: PythonMessage) -> None:
         """Send a message to the port's receiver."""
 
     def bind(self) -> PortRef:
@@ -71,7 +69,7 @@ class PortRef:
         """
         ...
 
-    def send(self, mailbox: Mailbox, message: PythonMessage) -> None:
+    def send(self, instance: Instance, message: PythonMessage) -> None:
         """Send a single message to the port's receiver."""
         ...
 
@@ -123,7 +121,7 @@ class OncePortRef:
     A reference to a remote once port over which a single PythonMessages can be sent.
     """
 
-    def send(self, mailbox: Mailbox, message: PythonMessage) -> None:
+    def send(self, instance: Instance, message: PythonMessage) -> None:
         """Send a single message to the port's receiver."""
         ...
 

python/monarch/_src/actor/actor_mesh.py

@@ -550,11 +550,11 @@ class Port(Generic[R]):
     def __init__(
         self,
         port_ref: PortRef | OncePortRef,
-        mailbox: Mailbox,
+        instance: HyInstance,
         rank: Optional[int],
     ) -> None:
         self._port_ref = port_ref
-        self._mailbox = mailbox
+        self._instance = instance
         self._rank = rank
 
     def send(self, obj: R) -> None:
@@ -566,18 +566,35 @@ def send(self, obj: R) -> None:
             obj: R-typed object to send.
         """
         self._port_ref.send(
-            self._mailbox,
+            self._instance,
             PythonMessage(PythonMessageKind.Result(self._rank), _pickle(obj)),
         )
 
     def exception(self, obj: Exception) -> None:
         # we deliver each error exactly once, so if there is no port to respond to,
         # the error is sent to the current actor as an exception.
         self._port_ref.send(
-            self._mailbox,
+            self._instance,
             PythonMessage(PythonMessageKind.Exception(self._rank), _pickle(obj)),
         )
 
+    def __reduce__(self):
+        """
+        When Port is sent over the wire, we do not want to send the actor instance
+        from the current context. Instead, we want to reconstruct the Port with
+        the receiver's context, since that is where the message will be sent
+        from through this port.
+        """
+
+        def _reconstruct_port(port_ref, rank):
+            instance = context().actor_instance._as_rust()
+            return Port(port_ref, instance, rank)
+
+        return (
+            _reconstruct_port,
+            (self._port_ref, self._rank),
+        )
+
 
 class DroppingPort:
     """
@@ -617,11 +634,14 @@ class Channel(Generic[R]):
     @staticmethod
     def open(once: bool = False) -> Tuple["Port[R]", "PortReceiver[R]"]:
         """ """
-        mailbox = context().actor_instance._mailbox
+        actor_instance = context().actor_instance
+        mailbox = actor_instance._mailbox
         handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
         port_ref = handle.bind()
+        hy_instance = actor_instance._as_rust()
+        port = Port(port_ref, hy_instance, None)
         return (
-            Port(port_ref, mailbox, rank=None),
+            port,
             PortReceiver(mailbox, receiver),
         )
 

python/tests/_monarch/test_mailbox.py

@@ -115,7 +115,7 @@ def _python_task_test(
 
 @_python_task_test
 async def test_accumulator() -> None:
-    proc_mesh = await allocate()
+    proc_mesh: ProcMesh = await allocate()
     mailbox: Mailbox = Instance._as_py(proc_mesh.client)._mailbox
 
     def my_accumulate(state: str, update: int) -> str:
@@ -128,7 +128,7 @@ def my_accumulate(state: str, update: int) -> str:
 
     def post_message(value: int) -> None:
         port_ref.send(
-            mailbox,
+            proc_mesh.client,
             PythonMessage(
                 PythonMessageKind.CallMethod(
                     MethodSpecifier.ReturnsResponse("test_accumulator"), None

python/tests/test_python_actors.py

@@ -1554,14 +1554,14 @@ async def get_messages(self) -> Tuple[str, str, str]:
 class UndeliverableMessageSender(Actor):
     @endpoint
     def send_undeliverable(self) -> None:
-        mailbox = context().actor_instance._mailbox
+        actor_instance = context().actor_instance
         port_id = PortId(
             actor_id=ActorId(world_name="bogus", rank=0, actor_name="bogus"),
             port=1234,
         )
         port_ref = PortRef(port_id)
         port_ref.send(
-            mailbox,
+            actor_instance._as_rust(),
             PythonMessage(PythonMessageKind.Result(None), b"123"),
         )