Add return_undeliverable property to PortRef and MessageEnvelope, and fix race in tensor engine shutdown (#1777)

Summary:
Pull Request resolved: #1777

Sometimes when we send a message, we want it to be fully fire-and-forget,
including if the destination is not even reachable. This is typically only used in
scenarios like:
* When shutting down the system, we try to ask a process to nicely shut itself down
before ungracefully killing it. If the message is undeliverable, we can just proceed with
killing the process (it's probably already dead anyways)
* Replying to a message. If the sender is down, there's nothing the current actor can do about it

This should be used sparingly as it could hide real errors, like your messages not getting sent.

This diff adds a `return_undeliverable: bool` property on `MessageEnvelope` and `PortRef`. When the property is set on `PortRef`, any `MessageEnvelope` sent via that `PortRef` will have an equivalent value for `return_undeliverable`. Any envelope with `return_undeliverable == true` will not be returned to its sender on delivery failure.

This is useful for messages like `GetRankStatus` and `GetState`, where the receiver shouldn't fail its reply fails to be delivered. It is also useful during proc termination, when the host mesh agent sends `StopAll` to the proc mesh agent; if the proc mesh agent is already dead, the message won't be delivered, but that shouldn't crash the host mesh agent.

Unrelatedly, this diff also fixes a race condition with host/proc mesh shutdown vs. tensor engine shutdown. Basically, `DeviceMesh.exit` was sending a fire-and-forget `WorkerMessage::Exit` via `Controller.drain_and_stop()`. But if you simultaneously try to shut down the host/proc mesh, then the worker exit message might fail to deliver, crashing the process. With this diff, `Controller.drain_and_stop()` synchronously calls `ActorMesh::stop` on the worker actor mesh so that there can't be a race with host/proc mesh shutdown (at least not from the same thread).

ghstack-source-id: 322545046
exported-using-ghexport

Reviewed By: mariusae, dulinriley, shayne-fletcher

Differential Revision: D86315780

fbshipit-source-id: 06c4aa92331e7f11c64f1ea8b13c52c2e7f0c153

Author: samlurye

hyperactor/src/context.rs

@@ -61,14 +61,15 @@ pub trait Actor: Mailbox {
 pub(crate) trait MailboxExt: Mailbox {
     /// Post a message to the provided destination with the provided headers, and data.
     /// All messages posted from actors should use this implementation.
-    fn post(&self, dest: PortId, headers: Attrs, data: Serialized);
+    fn post(&self, dest: PortId, headers: Attrs, data: Serialized, return_undeliverable: bool);
 
     /// Split a port, using a provided reducer spec, if provided.
     fn split(
         &self,
         port_id: PortId,
         reducer_spec: Option<ReducerSpec>,
         reducer_opts: Option<ReducerOpts>,
+        return_undeliverable: bool,
     ) -> anyhow::Result<PortId>;
 }
 
@@ -80,7 +81,7 @@ static CAN_SEND_WARNED_MAILBOXES: OnceLock<DashSet<ActorId>> = OnceLock::new();
 
 /// Only actors CanSend because they need a return port.
 impl<T: Actor + Send + Sync> MailboxExt for T {
-    fn post(&self, dest: PortId, headers: Attrs, data: Serialized) {
+    fn post(&self, dest: PortId, headers: Attrs, data: Serialized, return_undeliverable: bool) {
         let return_handle = self.mailbox().bound_return_handle().unwrap_or_else(|| {
             let actor_id = self.mailbox().actor_id();
             if CAN_SEND_WARNED_MAILBOXES
@@ -97,7 +98,9 @@ impl<T: Actor + Send + Sync> MailboxExt for T {
             mailbox::monitored_return_handle()
         });
 
-        let envelope = MessageEnvelope::new(self.mailbox().actor_id().clone(), dest, data, headers);
+        let mut envelope =
+            MessageEnvelope::new(self.mailbox().actor_id().clone(), dest, data, headers);
+        envelope.set_return_undeliverable(return_undeliverable);
         MailboxSender::post(self.mailbox(), envelope, return_handle);
     }
 
@@ -106,11 +109,20 @@ impl<T: Actor + Send + Sync> MailboxExt for T {
         port_id: PortId,
         reducer_spec: Option<ReducerSpec>,
         reducer_opts: Option<ReducerOpts>,
+        return_undeliverable: bool,
     ) -> anyhow::Result<PortId> {
-        fn post(mailbox: &mailbox::Mailbox, port_id: PortId, msg: Serialized) {
+        fn post(
+            mailbox: &mailbox::Mailbox,
+            port_id: PortId,
+            msg: Serialized,
+            return_undeliverable: bool,
+        ) {
+            let mut envelope =
+                MessageEnvelope::new(mailbox.actor_id().clone(), port_id, msg, Attrs::new());
+            envelope.set_return_undeliverable(return_undeliverable);
             mailbox::MailboxSender::post(
                 mailbox,
-                MessageEnvelope::new(mailbox.actor_id().clone(), port_id, msg, Attrs::new()),
+                envelope,
                 // TODO(pzhang) figure out how to use upstream's return handle,
                 // instead of getting a new one like this.
                 // This is okay for now because upstream is currently also using
@@ -135,7 +147,7 @@ impl<T: Actor + Send + Sync> MailboxExt for T {
             dyn Fn(Serialized) -> Result<(), (Serialized, anyhow::Error)> + Send + Sync,
         > = match reducer {
             None => Box::new(move |serialized: Serialized| {
-                post(&mailbox, port_id.clone(), serialized);
+                post(&mailbox, port_id.clone(), serialized, return_undeliverable);
                 Ok(())
             }),
             Some(reducer) => {
@@ -154,7 +166,9 @@ impl<T: Actor + Send + Sync> MailboxExt for T {
                             let mut buf = buffer.lock().unwrap();
                             match buf.reduce() {
                                 None => (),
-                                Some(Ok(reduced)) => post(&mailbox, port_id.clone(), reduced),
+                                Some(Ok(reduced)) => {
+                                    post(&mailbox, port_id.clone(), reduced, return_undeliverable)
+                                }
                                 // We simply ignore errors here, and let them be propagated
                                 // later in the enqueueing function.
                                 //
@@ -197,7 +211,7 @@ impl<T: Actor + Send + Sync> MailboxExt for T {
                         None => Ok(()),
                         Some(Ok(reduced)) => {
                             alarm.lock().unwrap().disarm();
-                            post(&mailbox, port_id.clone(), reduced);
+                            post(&mailbox, port_id.clone(), reduced, return_undeliverable);
                             Ok(())
                         }
                         Some(Err(e)) => Err((buf.pop().unwrap(), e)),

hyperactor/src/mailbox.rs

@@ -213,6 +213,10 @@ pub struct MessageEnvelope {
 
     /// Decremented at every `MailboxSender` hop.
     ttl: u8,
+
+    /// If true, undeliverable messages should be returned to sender. Else, they
+    /// are dropped.
+    return_undeliverable: bool,
     // TODO: add typename, source, seq, etc.
 }
 
@@ -226,6 +230,8 @@ impl MessageEnvelope {
             errors: Vec::new(),
             headers,
             ttl: crate::config::global::get(crate::config::MESSAGE_TTL_DEFAULT),
+            // By default, all undeliverable messages should be returned to the sender.
+            return_undeliverable: true,
         }
     }
 
@@ -248,6 +254,8 @@ impl MessageEnvelope {
             dest,
             errors: Vec::new(),
             ttl: crate::config::global::get(crate::config::MESSAGE_TTL_DEFAULT),
+            // By default, all undeliverable messages should be returned to the sender.
+            return_undeliverable: true,
         })
     }
 
@@ -333,6 +341,13 @@ impl MessageEnvelope {
         self.sender = sender;
     }
 
+    /// Set to true if you want this message to be returned to sender if it cannot
+    /// reach dest. This is the default.
+    /// Set to false if you want the message to be dropped instead.
+    pub fn set_return_undeliverable(&mut self, return_undeliverable: bool) {
+        self.return_undeliverable = return_undeliverable;
+    }
+
     /// The message has been determined to be undeliverable with the
     /// provided error. Mark the envelope with the error and return to
     /// sender.
@@ -393,6 +408,7 @@ impl MessageEnvelope {
             errors,
             headers,
             ttl,
+            return_undeliverable,
         } = self;
 
         (
@@ -402,6 +418,7 @@ impl MessageEnvelope {
                 errors,
                 headers,
                 ttl,
+                return_undeliverable,
             },
             data,
         )
@@ -414,6 +431,7 @@ impl MessageEnvelope {
             errors,
             headers,
             ttl,
+            return_undeliverable,
         } = metadata;
 
         Self {
@@ -423,8 +441,13 @@ impl MessageEnvelope {
             errors,
             headers,
             ttl,
+            return_undeliverable,
         }
     }
+
+    fn return_undeliverable(&self) -> bool {
+        self.return_undeliverable
+    }
 }
 
 impl fmt::Display for MessageEnvelope {
@@ -452,6 +475,7 @@ pub struct MessageMetadata {
     errors: Vec<DeliveryError>,
     headers: Attrs,
     ttl: u8,
+    return_undeliverable: bool,
 }
 
 /// Errors that occur during mailbox operations. Each error is associated
@@ -1534,6 +1558,7 @@ impl MailboxSender for Mailbox {
                     dest,
                     errors: metadata_errors,
                     ttl,
+                    return_undeliverable,
                 } = metadata;
 
                 // We use the entry API here so that we can remove the
@@ -1562,6 +1587,7 @@ impl MailboxSender for Mailbox {
                                 dest,
                                 errors: metadata_errors,
                                 ttl,
+                                return_undeliverable,
                             },
                             data,
                         )
@@ -3340,15 +3366,15 @@ mod tests {
 
         // Split it twice on actor1
         let port_id1 = port_id
-            .split(&actor1, reducer_spec.clone(), reducer_opts.clone())
+            .split(&actor1, reducer_spec.clone(), reducer_opts.clone(), true)
             .unwrap();
         let port_id2 = port_id
-            .split(&actor1, reducer_spec.clone(), reducer_opts.clone())
+            .split(&actor1, reducer_spec.clone(), reducer_opts.clone(), true)
             .unwrap();
 
         // A split port id can also be split
         let port_id2_1 = port_id2
-            .split(&actor1, reducer_spec, reducer_opts.clone())
+            .split(&actor1, reducer_spec, reducer_opts.clone(), true)
             .unwrap();
 
         Setup {
@@ -3470,7 +3496,7 @@ mod tests {
         let port_id = port_handle.bind().port_id().clone();
         // Split it
         let reducer_spec = accum::sum::<u64>().reducer_spec();
-        let split_port_id = port_id.split(&actor, reducer_spec, None).unwrap();
+        let split_port_id = port_id.split(&actor, reducer_spec, None, true).unwrap();
 
         // Send 9 messages.
         for msg in [1, 5, 3, 4, 2, 91, 92, 93, 94] {

hyperactor/src/mailbox/undeliverable.rs

@@ -98,9 +98,11 @@ pub(crate) fn return_undeliverable(
     return_handle: PortHandle<Undeliverable<MessageEnvelope>>,
     envelope: MessageEnvelope,
 ) {
-    let envelope_copy = envelope.clone();
-    if (return_handle.send(Undeliverable(envelope))).is_err() {
-        UndeliverableMailboxSender.post(envelope_copy, /*unsued*/ return_handle)
+    if envelope.return_undeliverable() {
+        let envelope_copy = envelope.clone();
+        if (return_handle.send(Undeliverable(envelope))).is_err() {
+            UndeliverableMailboxSender.post(envelope_copy, /*unused*/ return_handle)
+        }
     }
 }
 
@@ -109,7 +111,7 @@ pub(crate) fn return_undeliverable(
 pub enum UndeliverableMessageError {
     /// Delivery of a message to its destination failed.
     #[error(
-        "a message from {} to {} was undeliverable and returned: {:?}: {envelope}", 
+        "a message from {} to {} was undeliverable and returned: {:?}: {envelope}",
         .envelope.sender(),
         .envelope.dest(),
         .envelope.error_msg()

hyperactor/src/proc.rs

@@ -1036,7 +1036,7 @@ impl<A: Actor> Instance<A> {
 
     /// Send a message to the actor running on the proc.
     pub fn post(&self, port_id: PortId, headers: Attrs, message: Serialized) {
-        <Self as context::MailboxExt>::post(self, port_id, headers, message)
+        <Self as context::MailboxExt>::post(self, port_id, headers, message, true)
     }
 
     /// Send a message to the actor itself with a delay usually to trigger some event.

hyperactor/src/reference.rs

@@ -891,7 +891,7 @@ impl PortId {
     pub fn send(&self, cx: &impl context::Actor, serialized: Serialized) {
         let mut headers = Attrs::new();
         crate::mailbox::headers::set_send_timestamp(&mut headers);
-        cx.post(self.clone(), headers, serialized);
+        cx.post(self.clone(), headers, serialized, true);
     }
 
     /// Send a serialized message to this port, provided a sending capability,
@@ -904,7 +904,7 @@ impl PortId {
         mut headers: Attrs,
     ) {
         crate::mailbox::headers::set_send_timestamp(&mut headers);
-        cx.post(self.clone(), headers, serialized);
+        cx.post(self.clone(), headers, serialized, true);
     }
 
     /// Split this port, returning a new port that relays messages to the port
@@ -914,8 +914,14 @@ impl PortId {
         cx: &impl context::Actor,
         reducer_spec: Option<ReducerSpec>,
         reducer_opts: Option<ReducerOpts>,
+        return_undeliverable: bool,
     ) -> anyhow::Result<PortId> {
-        cx.split(self.clone(), reducer_spec, reducer_opts)
+        cx.split(
+            self.clone(),
+            reducer_spec,
+            reducer_opts,
+            return_undeliverable,
+        )
     }
 }
 
@@ -964,6 +970,7 @@ pub struct PortRef<M> {
     )]
     reducer_opts: Option<ReducerOpts>,
     phantom: PhantomData<M>,
+    return_undeliverable: bool,
 }
 
 impl<M: RemoteMessage> PortRef<M> {
@@ -975,6 +982,7 @@ impl<M: RemoteMessage> PortRef<M> {
             reducer_spec: None,
             reducer_opts: None,
             phantom: PhantomData,
+            return_undeliverable: true,
         }
     }
 
@@ -986,6 +994,7 @@ impl<M: RemoteMessage> PortRef<M> {
             reducer_spec,
             reducer_opts: None, // TODO: provide attest_reducible_opts
             phantom: PhantomData,
+            return_undeliverable: true,
         }
     }
 
@@ -1052,13 +1061,24 @@ impl<M: RemoteMessage> PortRef<M> {
     ) {
         crate::mailbox::headers::set_send_timestamp(&mut headers);
         crate::mailbox::headers::set_rust_message_type::<M>(&mut headers);
-        cx.post(self.port_id.clone(), headers, message);
+        cx.post(
+            self.port_id.clone(),
+            headers,
+            message,
+            self.return_undeliverable,
+        );
     }
 
     /// Convert this port into a sink that can be used to send messages using the given capability.
     pub fn into_sink<C: context::Actor>(self, cx: C) -> PortSink<C, M> {
         PortSink::new(cx, self)
     }
+
+    /// Set whether or not messages sent to this port that are undeliverable
+    /// should be returned to the sender.
+    pub fn return_undeliverable(&mut self, return_undeliverable: bool) {
+        self.return_undeliverable = return_undeliverable;
+    }
 }
 
 impl<M: RemoteMessage> Clone for PortRef<M> {
@@ -1068,6 +1088,7 @@ impl<M: RemoteMessage> Clone for PortRef<M> {
             reducer_spec: self.reducer_spec.clone(),
             reducer_opts: self.reducer_opts.clone(),
             phantom: PhantomData,
+            return_undeliverable: self.return_undeliverable,
         }
     }
 }
@@ -1080,7 +1101,12 @@ impl<M: RemoteMessage> fmt::Display for PortRef<M> {
 
 /// The parameters extracted from [`PortRef`] to [`Bindings`].
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Named)]
-pub struct UnboundPort(pub PortId, pub Option<ReducerSpec>, pub Option<ReducerOpts>);
+pub struct UnboundPort(
+    pub PortId,
+    pub Option<ReducerSpec>,
+    pub Option<ReducerOpts>,
+    pub bool, // return_undeliverable
+);
 
 impl UnboundPort {
     /// Update the port id of this binding.
@@ -1095,6 +1121,7 @@ impl<M: RemoteMessage> From<&PortRef<M>> for UnboundPort {
             port_ref.port_id.clone(),
             port_ref.reducer_spec.clone(),
             port_ref.reducer_opts.clone(),
+            port_ref.return_undeliverable,
         )
     }
 }
@@ -1111,6 +1138,7 @@ impl<M: RemoteMessage> Bind for PortRef<M> {
         self.port_id = bound.0;
         self.reducer_spec = bound.1;
         self.reducer_opts = bound.2;
+        self.return_undeliverable = bound.3;
         Ok(())
     }
 }
@@ -1163,7 +1191,7 @@ impl<M: RemoteMessage> OncePortRef<M> {
                 MailboxSenderErrorKind::Serialize(err.into()),
             )
         })?;
-        cx.post(self.port_id.clone(), headers, serialized);
+        cx.post(self.port_id.clone(), headers, serialized, true);
         Ok(())
     }
 }