Fix assert with undeliverable message from comm actor (#1606)

Summary:
Pull Request resolved: #1606

Fix a panic in PythonActor::handle_undeliverable_message when the "sender" is the comm
actor. We need to update the sender back to the original "self" actor by using the headers
set in the envelope.

After this fix, instead of a panic we get a supervision error:
```
SupervisionError(
Actor ...wrapper_1xYczTZiTdb1[0] exited because of the following reason:
<PyActorSupervisionEvent: ...wrapper_1xYczTZiTdb1[0]: failed: serving
...wrapper_1xYczTZiTdb1[0]: processing error: a message from ...wrapper_1xYczTZiTdb1[0] to
...fail_1am38hE5fnus[0] was undeliverable and returned: Some("send error: channel closed;
multicast error: comm actor comm_1JgvjFbdpnUf[0] failed to deliver the cast message to the
dest actor; return to its original sender's port ...wrapper_1xYczTZiTdb1[0]
)
```

Not very terse, but better than a panic! This also allows any custom override of
handle_undeliverable_message to work.

Reviewed By: pablorfb-meta

Differential Revision: D84952942

fbshipit-source-id: b2cb36600ca89a03e7a85cfccd46ce6bcf2487cf

Author: dulinriley

hyperactor/src/mailbox.rs

@@ -324,6 +324,13 @@ impl MessageEnvelope {
         self.errors.push(error)
     }
 
+    /// Change the sender on the envelope in case it was set incorrectly. This
+    /// should only be used by CommActor since it is forwarding from another
+    /// sender.
+    pub fn update_sender(&mut self, sender: ActorId) {
+        self.sender = sender;
+    }
+
     /// The message has been determined to be undeliverable with the
     /// provided error. Mark the envelope with the error and return to
     /// sender.

monarch_hyperactor/src/actor.rs

@@ -20,12 +20,17 @@ use hyperactor::Handler;
 use hyperactor::Instance;
 use hyperactor::Named;
 use hyperactor::OncePortHandle;
+use hyperactor::ProcId;
 use hyperactor::mailbox::MessageEnvelope;
 use hyperactor::mailbox::Undeliverable;
 use hyperactor::message::Bind;
 use hyperactor::message::Bindings;
 use hyperactor::message::Unbind;
+use hyperactor::reference::WorldId;
+use hyperactor_mesh::actor_mesh::CAST_ACTOR_MESH_ID;
+use hyperactor_mesh::comm::multicast::CAST_ORIGINATING_SENDER;
 use hyperactor_mesh::comm::multicast::CastInfo;
+use hyperactor_mesh::reference::ActorMeshId;
 use monarch_types::PickledPyObject;
 use monarch_types::SerializablePyErr;
 use pyo3::IntoPyObjectExt;
@@ -493,6 +498,47 @@ impl PythonActor {
     }
 }
 
+/// An undeliverable might have its sender address set as the comm actor instead
+/// of the original sender. Update it based on the headers present in the message
+/// so it matches the sender.
+fn update_undeliverable_envelope_for_casting(
+    mut envelope: Undeliverable<MessageEnvelope>,
+) -> Undeliverable<MessageEnvelope> {
+    let old_actor = envelope.0.sender().clone();
+    // v1 casting
+    if let Some(actor_id) = envelope.0.headers().get(CAST_ORIGINATING_SENDER).cloned() {
+        tracing::debug!(
+            actor_id = %old_actor,
+            "PythonActor::handle_undeliverable_message: remapped comm-actor id to id from CAST_ORIGINATING_SENDER {}", actor_id
+        );
+        envelope.0.update_sender(actor_id);
+    // v0 casting
+    } else if let Some(actor_mesh_id) = envelope.0.headers().get(CAST_ACTOR_MESH_ID) {
+        match actor_mesh_id {
+            ActorMeshId::V0(proc_mesh_id, actor_name) => {
+                let actor_id = ActorId(
+                    ProcId::Ranked(WorldId(proc_mesh_id.0.clone()), 0),
+                    actor_name.clone(),
+                    0,
+                );
+                tracing::debug!(
+                    actor_id = %old_actor,
+                    "PythonActor::handle_undeliverable_message: remapped comm-actor id to mesh id from CAST_ACTOR_MESH_ID {}", actor_id
+                );
+                envelope.0.update_sender(actor_id);
+            }
+            ActorMeshId::V1(_) => {
+                tracing::debug!(
+                    "PythonActor::handle_undeliverable_message: headers present but V1 ActorMeshId; leaving actor_id unchanged"
+                );
+            }
+        }
+    } else {
+        // Do nothing, it wasn't from a comm actor.
+    }
+    envelope
+}
+
 #[async_trait]
 impl Actor for PythonActor {
     type Params = PickledPyObject;
@@ -537,9 +583,21 @@ impl Actor for PythonActor {
     async fn handle_undeliverable_message(
         &mut self,
         ins: &Instance<Self>,
-        envelope: Undeliverable<MessageEnvelope>,
+        mut envelope: Undeliverable<MessageEnvelope>,
     ) -> Result<(), anyhow::Error> {
-        assert_eq!(envelope.0.sender(), ins.self_id());
+        if envelope.0.sender() != ins.self_id() {
+            // This can happen if the sender is comm. Update the envelope.
+            envelope = update_undeliverable_envelope_for_casting(envelope);
+        }
+        assert_eq!(
+            envelope.0.sender(),
+            ins.self_id(),
+            "undeliverable message was returned to the wrong actor. \
+            Return address = {}, src actor = {}, dest actor port = {}",
+            envelope.0.sender(),
+            ins.self_id(),
+            envelope.0.dest()
+        );
 
         let cx = Context::new(ins, envelope.0.headers().clone());