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

Summary:
Pull Request resolved: #964

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.
ghstack-source-id: 305073646

Reviewed By: mariusae

Differential Revision: D79841379

fbshipit-source-id: 2d285b52df401ed01813ef57e39e158f0104bea1

Author: samlurye

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,41 @@ 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());
+
+        let (envelope, handled) = Python::with_gil(|py| {
+            let py_envelope = PythonUndeliverableMessageEnvelope {
+                inner: Some(envelope),
+            }
+            .into_bound_py_any(py)?;
+            let handled = self
+                .actor
+                .call_method(py, "_handle_undeliverable_message", (&py_envelope,), None)
+                .map_err(|err| anyhow::Error::from(SerializablePyErr::from(py, &err)))?
+                .extract::<bool>(py)?;
+            Ok::<_, anyhow::Error>((
+                py_envelope
+                    .downcast::<PythonUndeliverableMessageEnvelope>()
+                    .map_err(PyErr::from)?
+                    .try_borrow_mut()
+                    .map_err(PyErr::from)?
+                    .take()?,
+                handled,
+            ))
+        })?;
+
+        if !handled {
+            <Self as Actor>::handle_undeliverable_message(self, cx, envelope).await
+        } else {
+            Ok(())
+        }
+    }
 }
 
 /// Create a new TaskLocals with its own asyncio event loop in a dedicated thread.

monarch_hyperactor/src/mailbox.rs

@@ -421,8 +421,53 @@ 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>,
+    pub(crate) inner: Option<Undeliverable<MessageEnvelope>>,
+}
+
+impl PythonUndeliverableMessageEnvelope {
+    fn inner(&self) -> PyResult<&Undeliverable<MessageEnvelope>> {
+        self.inner.as_ref().ok_or_else(|| {
+            PyErr::new::<PyRuntimeError, _>(
+                "PythonUndeliverableMessageEnvelope was already consumed",
+            )
+        })
+    }
+
+    pub(crate) fn take(&mut self) -> anyhow::Result<Undeliverable<MessageEnvelope>> {
+        self.inner.take().ok_or_else(|| {
+            anyhow::anyhow!("PythonUndeliverableMessageEnvelope was already consumed")
+        })
+    }
+}
+
+#[pymethods]
+impl PythonUndeliverableMessageEnvelope {
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!(
+            "UndeliverableMessageEnvelope(sender={}, dest={}, error={})",
+            self.inner()?.0.sender(),
+            self.inner()?.0.dest(),
+            self.error_msg()?
+        ))
+    }
+
+    fn sender(&self) -> PyResult<PyActorId> {
+        Ok(PyActorId {
+            inner: self.inner()?.0.sender().clone(),
+        })
+    }
+
+    fn dest(&self) -> PyResult<PyPortId> {
+        Ok(self.inner()?.0.dest().clone().into())
+    }
+
+    fn error_msg(&self) -> PyResult<String> {
+        Ok(self
+            .inner()?
+            .0
+            .error()
+            .map_or("None".to_string(), |e| e.to_string()))
+    }
 }
 
 #[derive(Debug)]
@@ -445,7 +490,9 @@ impl PythonUndeliverablePortReceiver {
                 .recv()
                 .await
                 .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
-            Ok(PythonUndeliverableMessageEnvelope { inner: message })
+            Ok(PythonUndeliverableMessageEnvelope {
+                inner: Some(message),
+            })
         })
     }
 
@@ -457,7 +504,9 @@ impl PythonUndeliverablePortReceiver {
         let message = signal_safe_block_on(py, async move { receiver.lock().await.recv().await })?
             .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
 
-        Ok(PythonUndeliverableMessageEnvelope { inner: message })
+        Ok(PythonUndeliverableMessageEnvelope {
+            inner: Some(message),
+        })
     }
 }
 
@@ -796,5 +845,6 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PythonOncePortReceiver>()?;
     hyperactor_mod.add_class::<Instance>()?;
     hyperactor_mod.add_class::<Context>()?;
+    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."""
         ...
@@ -218,3 +221,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

@@ -55,6 +55,7 @@
     OncePortRef,
     PortReceiver as HyPortReceiver,
     PortRef,
+    UndeliverableMessageEnvelope,
 )
 
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
@@ -895,6 +896,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
+    ) -> bool:
+        handle_undeliverable = getattr(
+            self.instance, "_handle_undeliverable_message", None
+        )
+        if handle_undeliverable is not None:
+            return handle_undeliverable(message)
+        else:
+            return False
+
 
 def _is_mailbox(x: object) -> bool:
     if hasattr(x, "__monarch_ref__"):
@@ -943,6 +955,12 @@ def initialized(self):
             "actor implementations are not meshes, but we can't convince the typechecker of it..."
         )
 
+    def _handle_undeliverable_message(
+        self, message: UndeliverableMessageEnvelope
+    ) -> bool:
+        # Return False to indicate that the undeliverable message was not handled.
+        return False
+
 
 class ActorMesh(MeshTrait, Generic[T], DeprecatedNotAFuture):
     def __init__(

python/tests/test_python_actors.py

@@ -19,14 +19,24 @@
 import unittest
 import unittest.mock
 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, context, Port
 from monarch._src.actor.future import Future
 
 from monarch.actor import (
@@ -1127,3 +1137,59 @@ def test_mesh_len():
     proc_mesh = local_proc_mesh(gpus=12).get()
     s = proc_mesh.spawn("sync_actor", SyncActor).get()
     assert 12 == len(s)
+
+
+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 = context().actor_instance._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
+    ) -> bool:
+        self._receiver.receive_undeliverable.call_one(
+            message.sender(), message.dest(), message.error_msg()
+        ).get()
+        return True
+
+
+@pytest.mark.timeout(60)
+async def test_undeliverable_message() -> None:
+    pm = proc_mesh(gpus=1)
+    receiver = pm.spawn("undeliverable_receiver", UndeliverableMessageReceiver)
+    sender = pm.spawn("undeliverable_sender", UndeliverableMessageSender, receiver)
+    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()