Semi-private python API for overriding handle_undeliverable_message inside PythonActor (#797)

Summary:

This diff makes undeliverable message handling overridable for python actors, using the newly introduced `Actor._handle_undeliverable_message` method. Previously, the rust implementation of `PythonActor` simply used the default `Actor::handle_undeliverable_message` implementation. Now, `PythonActor` overrides `handle_undeliverable_message` to call into the corresponding method on the underlying python class.

Differential Revision: D79841379

Author: Sam Lurye

monarch_hyperactor/src/actor.rs

@@ -21,6 +21,8 @@ use hyperactor::Handler;
 use hyperactor::Instance;
 use hyperactor::Named;
 use hyperactor::OncePortHandle;
+use hyperactor::mailbox::MessageEnvelope;
+use hyperactor::mailbox::Undeliverable;
 use hyperactor::message::Bind;
 use hyperactor::message::Bindings;
 use hyperactor::message::Unbind;
@@ -50,6 +52,7 @@ use crate::local_state_broker::BrokerId;
 use crate::local_state_broker::LocalStateBrokerMessage;
 use crate::mailbox::EitherPortRef;
 use crate::mailbox::PyMailbox;
+use crate::mailbox::PythonUndeliverableMessageEnvelope;
 use crate::proc::InstanceWrapper;
 use crate::proc::PyActorId;
 use crate::proc::PyProc;
@@ -498,6 +501,26 @@ impl Actor for PythonActor {
         );
         Ok(())
     }
+
+    async fn handle_undeliverable_message(
+        &mut self,
+        cx: &Instance<Self>,
+        envelope: Undeliverable<MessageEnvelope>,
+    ) -> Result<(), anyhow::Error> {
+        assert_eq!(envelope.0.sender(), cx.self_id());
+
+        Python::with_gil(|py| {
+            self.actor
+                .call_method(
+                    py,
+                    "_handle_undeliverable_message",
+                    (PythonUndeliverableMessageEnvelope { inner: envelope },),
+                    None,
+                )
+                .map_err(|err| anyhow::Error::from(SerializablePyErr::from(py, &err)))
+        })
+        .map(|_| ())
+    }
 }
 
 /// Create a new TaskLocals with its own asyncio event loop in a dedicated thread.

monarch_hyperactor/src/mailbox.rs

@@ -420,10 +420,38 @@ impl PythonPortReceiver {
     module = "monarch._rust_bindings.monarch_hyperactor.mailbox"
 )]
 pub(crate) struct PythonUndeliverableMessageEnvelope {
-    #[allow(dead_code)] // At this time, field `inner` isn't read.
     pub(crate) inner: Undeliverable<MessageEnvelope>,
 }
 
+#[pymethods]
+impl PythonUndeliverableMessageEnvelope {
+    fn __repr__(&self) -> String {
+        format!(
+            "UndeliverableMessageEnvelope(sender={}, dest={}, error={})",
+            self.inner.0.sender(),
+            self.inner.0.dest(),
+            self.error_msg()
+        )
+    }
+
+    fn sender(&self) -> PyActorId {
+        PyActorId {
+            inner: self.inner.0.sender().clone(),
+        }
+    }
+
+    fn dest(&self) -> PyPortId {
+        self.inner.0.dest().clone().into()
+    }
+
+    fn error_msg(&self) -> String {
+        self.inner
+            .0
+            .error()
+            .map_or("None".to_string(), |e| e.to_string())
+    }
+}
+
 #[derive(Debug)]
 #[pyclass(
     name = "UndeliverablePortReceiver",
@@ -713,5 +741,6 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PythonOncePortHandle>()?;
     hyperactor_mod.add_class::<PythonOncePortRef>()?;
     hyperactor_mod.add_class::<PythonOncePortReceiver>()?;
+    hyperactor_mod.add_class::<PythonUndeliverableMessageEnvelope>()?;
     Ok(())
 }

python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi

@@ -210,15 +210,6 @@ class PythonMessage:
     @property
     def kind(self) -> PythonMessageKind: ...
 
-class UndeliverableMessageEnvelope:
-    """
-    An envelope representing a message that could not be delivered.
-
-    This object is opaque; its contents are not accessible from Python.
-    """
-
-    ...
-
 @final
 class PythonActorHandle:
     """

python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi

@@ -8,10 +8,7 @@
 
 from typing import final, Protocol
 
-from monarch._rust_bindings.monarch_hyperactor.actor import (
-    PythonMessage,
-    UndeliverableMessageEnvelope,
-)
+from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage
 
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
@@ -20,9 +17,9 @@ from monarch._rust_bindings.monarch_hyperactor.shape import Shape
 
 @final
 class PortId:
-    def __init__(self, actor_id: ActorId, index: int) -> None:
+    def __init__(self, *, actor_id: ActorId, port: int) -> None:
         """
-        Create a new port id given an actor id and an index.
+        Create a new port id given an actor id and a port index.
         """
         ...
     def __repr__(self) -> str: ...
@@ -68,6 +65,12 @@ class PortRef:
     A reference to a remote port over which PythonMessages can be sent.
     """
 
+    def __init__(self, port_id: PortId) -> None:
+        """
+        Create a new port ref given a port id.
+        """
+        ...
+
     def send(self, mailbox: Mailbox, message: PythonMessage) -> None:
         """Send a single message to the port's receiver."""
         ...
@@ -220,3 +223,27 @@ class Reducer(Protocol):
 
     This method's Rust counterpart is `CommReducer::reduce`.
     """
+
+class UndeliverableMessageEnvelope:
+    """
+    An envelope representing a message that could not be delivered.
+    """
+
+    def __repr__(self) -> str: ...
+    def sender(self) -> ActorId:
+        """
+        The actor id of the sender.
+        """
+        ...
+
+    def dest(self) -> PortId:
+        """
+        The port id of the destination.
+        """
+        ...
+
+    def error_msg(self) -> str:
+        """
+        The error message describing why the message could not be delivered.
+        """
+        ...

python/monarch/_src/actor/actor_mesh.py

@@ -58,6 +58,7 @@
     OncePortRef,
     PortReceiver as HyPortReceiver,
     PortRef,
+    UndeliverableMessageEnvelope,
 )
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask, Shared
 
@@ -969,6 +970,17 @@ def _post_mortem_debug(self, exc_tb) -> None:
                 pdb_wrapper.post_mortem(exc_tb)
                 self._maybe_exit_debugger(do_continue=False)
 
+    def _handle_undeliverable_message(
+        self, message: UndeliverableMessageEnvelope
+    ) -> None:
+        handle_undeliverable = getattr(
+            self.instance, "_handle_undeliverable_message", None
+        )
+        if handle_undeliverable is not None:
+            handle_undeliverable(message)
+        else:
+            raise RuntimeError(f"a message was undeliverable and returned: {message}")
+
 
 def _is_mailbox(x: object) -> bool:
     if hasattr(x, "__monarch_ref__"):
@@ -1011,6 +1023,11 @@ def _new_with_shape(self, shape: Shape) -> Self:
             "actor implementations are not meshes, but we can't convince the typechecker of it..."
         )
 
+    def _handle_undeliverable_message(
+        self, message: UndeliverableMessageEnvelope
+    ) -> None:
+        raise RuntimeError(f"a message was undeliverable and returned: {message}")
+
 
 class ActorMesh(MeshTrait, Generic[T]):
     def __init__(

python/tests/test_python_actors.py

@@ -18,14 +18,24 @@
 import time
 import unittest
 from types import ModuleType
-from typing import cast
+from typing import cast, Tuple
 
 import pytest
 
 import torch
+from monarch._rust_bindings.monarch_hyperactor.actor import (
+    PythonMessage,
+    PythonMessageKind,
+)
+from monarch._rust_bindings.monarch_hyperactor.mailbox import (
+    PortId,
+    PortRef,
+    UndeliverableMessageEnvelope,
+)
+from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
 
-from monarch._src.actor.actor_mesh import ActorMesh, Channel, Port
+from monarch._src.actor.actor_mesh import ActorMesh, Channel, MonarchContext, Port
 
 from monarch.actor import (
     Accumulator,
@@ -832,3 +842,60 @@ def s(t):
     b = PythonTask.spawn_blocking(lambda: s(0))
     r = PythonTask.select_one([a.task(), b.task()]).block_on()
     assert r == (0, 1)
+
+
+class UndeliverableMessageReceiver(Actor):
+    def __init__(self):
+        self._messages = asyncio.Queue()
+
+    @endpoint
+    async def receive_undeliverable(
+        self, sender: ActorId, dest: PortId, error_msg: str
+    ) -> None:
+        await self._messages.put((sender, dest, error_msg))
+
+    @endpoint
+    async def get_messages(self) -> Tuple[ActorId, PortId, str]:
+        return await self._messages.get()
+
+
+class UndeliverableMessageSender(Actor):
+    def __init__(self, receiver: UndeliverableMessageReceiver):
+        self._receiver = receiver
+
+    @endpoint
+    def send_undeliverable(self) -> None:
+        mailbox = MonarchContext.get().mailbox
+        port_id = PortId(
+            actor_id=ActorId(
+                world_name=mailbox.actor_id.world_name, rank=0, actor_name="bogus"
+            ),
+            port=1234,
+        )
+        port_ref = PortRef(port_id)
+        port_ref.send(
+            mailbox,
+            PythonMessage(PythonMessageKind.Result(None), b"123"),
+        )
+
+    def _handle_undeliverable_message(
+        self, message: UndeliverableMessageEnvelope
+    ) -> None:
+        self._receiver.receive_undeliverable.call_one(
+            message.sender(), message.dest(), message.error_msg()
+        ).get()
+
+
+@pytest.mark.timeout(60)
+async def test_undeliverable_message() -> None:
+    pm = proc_mesh(gpus=1)
+    receiver = pm.spawn("undeliverable_receiver", UndeliverableMessageReceiver).get()
+    sender = pm.spawn(
+        "undeliverable_sender", UndeliverableMessageSender, receiver
+    ).get()
+    sender.send_undeliverable.call().get()
+    sender, dest, error_msg = receiver.get_messages.call_one().get()
+    assert sender.actor_name == "undeliverable_sender"
+    assert dest.actor_id.actor_name == "bogus"
+    assert error_msg is not None
+    pm.stop().get()