Added watchdog timeout when stop is issued for stuck actors. (#336)

Summary:
Pull Request resolved: #336

Allocators have a timeout to kill unresponsive procs when stop() is issued.

Differential Revision: D77303504

Author: vidhyav

hyperactor/src/config.rs

@@ -32,6 +32,9 @@ declare_attrs! {
     /// Message delivery timeout
     pub attr MESSAGE_DELIVERY_TIMEOUT: Duration = Duration::from_secs(30);
 
+    /// Process exit timeout
+    pub attr PROCESS_EXIT_TIMEOUT: Duration = Duration::from_secs(10);
+
     /// Message acknowledgment interval
     pub attr MESSAGE_ACK_TIME_INTERVAL: Duration = Duration::from_millis(500);
 

hyperactor/src/sync/monitor.rs

@@ -99,6 +99,7 @@ impl IntoFuture for Handle {
     }
 }
 
+#[derive(Clone)]
 /// A group of tasks that share a common fate. Any tasks that are spawned onto
 /// the group will be aborted if any task fails or if the group is aborted.
 ///

hyperactor_mesh/src/alloc/mod.rs

@@ -228,11 +228,51 @@ pub trait Alloc {
 }
 
 pub mod test_utils {
+    use std::time::Duration;
+
+    use hyperactor::Actor;
+    use hyperactor::Handler;
+    use hyperactor::Instance;
+    use hyperactor::Named;
     use tokio::sync::broadcast::Receiver;
     use tokio::sync::broadcast::Sender;
 
     use super::*;
 
+    // This can't be defined under a `#[cfg(test)]` because there needs to
+    // be an entry in the spawnable actor registry in the executable
+    // 'hyperactor_mesh_test_bootstrap' for the `tests::process` actor
+    // mesh test suite.
+    #[derive(Debug)]
+    #[hyperactor::export(
+        spawn = true,
+        handlers = [
+            Wait
+        ],
+    )]
+    pub struct TestActor;
+
+    #[async_trait]
+    impl Actor for TestActor {
+        type Params = ();
+
+        async fn new(_params: Self::Params) -> Result<Self, anyhow::Error> {
+            Ok(Self)
+        }
+    }
+
+    #[derive(Debug, Serialize, Deserialize, Named, Clone)]
+    pub struct Wait();
+
+    #[async_trait]
+    impl Handler<Wait> for TestActor {
+        async fn handle(&mut self, _: &Instance<Self>, Wait(): Wait) -> Result<(), anyhow::Error> {
+            #[allow(clippy::disallowed_methods)]
+            tokio::time::sleep(Duration::from_secs(60)).await;
+            Ok(())
+        }
+    }
+
     /// Test wrapper around MockAlloc to allow us to block next() calls since
     /// mockall doesn't support returning futures.
     pub struct MockAllocWrapper {
@@ -302,12 +342,29 @@ pub mod test_utils {
 
 #[cfg(test)]
 pub(crate) mod testing {
+    use core::panic;
     use std::collections::HashMap;
     use std::collections::HashSet;
-
+    use std::time::Duration;
+
+    use hyperactor::Mailbox;
+    use hyperactor::actor::remote::Remote;
+    use hyperactor::channel;
+    use hyperactor::mailbox;
+    use hyperactor::mailbox::BoxedMailboxSender;
+    use hyperactor::mailbox::DialMailboxRouter;
+    use hyperactor::mailbox::IntoBoxedMailboxSender;
+    use hyperactor::mailbox::MailboxServer;
+    use hyperactor::mailbox::UndeliverableMailboxSender;
+    use hyperactor::proc::Proc;
+    use hyperactor::reference::Reference;
     use ndslice::shape;
+    use tokio::process::Command;
 
     use super::*;
+    use crate::alloc::test_utils::TestActor;
+    use crate::alloc::test_utils::Wait;
+    use crate::proc_mesh::mesh_agent::MeshAgentMessageClient;
 
     #[macro_export]
     macro_rules! alloc_test_suite {
@@ -367,4 +424,143 @@ pub(crate) mod testing {
         assert!(alloc.next().await.is_none());
         assert_eq!(stopped, running);
     }
+
+    async fn spawn_proc(alloc: &ProcessAlloc) -> (DialMailboxRouter, Mailbox, Proc, ChannelAddr) {
+        let (router_channel_addr, router_rx) = channel::serve(ChannelAddr::any(alloc.transport()))
+            .await
+            .map_err(|err| AllocatorError::Other(err.into()))
+            .unwrap();
+        let router =
+            DialMailboxRouter::new_with_default((UndeliverableMailboxSender {}).into_boxed());
+        router
+            .clone()
+            .serve(router_rx, mailbox::monitored_return_handle());
+
+        let client_proc_id = ProcId(WorldId(format!("test_{}", alloc.world_id().name())), 0);
+        let (client_proc_addr, client_rx) = channel::serve(ChannelAddr::any(alloc.transport()))
+            .await
+            .map_err(|err| AllocatorError::Other(err.into()))
+            .unwrap();
+        let client_proc = Proc::new(
+            client_proc_id.clone(),
+            BoxedMailboxSender::new(router.clone()),
+        );
+        client_proc
+            .clone()
+            .serve(client_rx, mailbox::monitored_return_handle());
+        router.bind(client_proc_id.clone().into(), client_proc_addr);
+        (
+            router,
+            client_proc.attach("test_proc").unwrap(),
+            client_proc,
+            router_channel_addr,
+        )
+    }
+
+    async fn spawn_test_actor(
+        rank: usize,
+        client_proc: &Proc,
+        client: &Mailbox,
+        router_channel_addr: ChannelAddr,
+        mesh_agent: ActorRef<MeshAgent>,
+    ) -> ActorRef<TestActor> {
+        let supervisor = client_proc.attach("supervisor").unwrap();
+        let (supervison_port, _) = supervisor.open_port();
+        let (config_handle, _) = client.open_port();
+        mesh_agent
+            .configure(
+                client,
+                rank,
+                router_channel_addr,
+                supervison_port.bind(),
+                HashMap::new(),
+                config_handle.bind(),
+            )
+            .await
+            .unwrap();
+        let remote = Remote::collect();
+        let actor_type = remote
+            .name_of::<TestActor>()
+            .ok_or(anyhow::anyhow!("actor not registered"))
+            .unwrap()
+            .to_string();
+        let params = &();
+        let (completed_handle, mut completed_receiver) = mailbox::open_port(client);
+        // gspawn actor
+        mesh_agent
+            .gspawn(
+                client,
+                actor_type,
+                "Stuck".to_string(),
+                bincode::serialize(params).unwrap(),
+                completed_handle.bind(),
+            )
+            .await
+            .unwrap();
+        let (_, actor_id) = completed_receiver.recv().await.unwrap();
+        ActorRef::<TestActor>::attest(actor_id)
+    }
+
+    #[tokio::test]
+    async fn test_allocator_stuck_task() {
+        let mut allocator = ProcessAllocator::new(Command::new(
+            buck_resources::get("monarch/hyperactor_mesh/bootstrap").unwrap(),
+        ));
+        // Override config.
+        // Use temporary config for this test
+        let config = hyperactor::config::global::lock();
+        let _guard = config.override_key(
+            hyperactor::config::PROCESS_EXIT_TIMEOUT,
+            Duration::from_secs(1),
+        );
+
+        let mut alloc = allocator
+            .allocate(AllocSpec {
+                shape: shape! { replica = 1 },
+                constraints: Default::default(),
+            })
+            .await
+            .unwrap();
+
+        // Get everything up into running state. We require that we get
+        let mut procs = HashMap::new();
+        let mut running = HashSet::new();
+        let mut actor_ref = None;
+        let (router, client, client_proc, router_addr) = spawn_proc(&alloc).await;
+        while running.is_empty() {
+            match alloc.next().await.unwrap() {
+                ProcState::Created { proc_id, coords } => {
+                    procs.insert(proc_id, coords);
+                }
+                ProcState::Running {
+                    proc_id,
+                    mesh_agent,
+                    addr,
+                } => {
+                    router.bind(Reference::Proc(proc_id.clone()), addr.clone());
+
+                    assert!(procs.contains_key(&proc_id));
+                    assert!(!running.contains(&proc_id));
+
+                    actor_ref = Some(
+                        spawn_test_actor(0, &client_proc, &client, router_addr, mesh_agent).await,
+                    );
+                    running.insert(proc_id);
+                    break;
+                }
+                event => panic!("unexpected event: {:?}", event),
+            }
+        }
+        assert!(actor_ref.unwrap().send(&client, Wait()).is_ok());
+
+        // There is a stuck actor! We should get a watchdog failure.
+        alloc.stop().await.unwrap();
+        let mut stopped = HashSet::new();
+        while let Some(ProcState::Stopped { proc_id, reason }) = alloc.next().await {
+            assert_eq!(reason, ProcStopReason::Watchdog);
+            stopped.insert(proc_id);
+        }
+        assert!(alloc.next().await.is_none());
+        assert_eq!(stopped, running);
+    }
 }

hyperactor_mesh/src/alloc/process.rs

@@ -18,6 +18,7 @@ use std::sync::OnceLock;
 
 use async_trait::async_trait;
 use enum_as_inner::EnumAsInner;
+use futures::executor;
 use hyperactor::ProcId;
 use hyperactor::WorldId;
 use hyperactor::channel;
@@ -28,8 +29,10 @@ use hyperactor::channel::ChannelTx;
 use hyperactor::channel::Rx;
 use hyperactor::channel::Tx;
 use hyperactor::channel::TxStatus;
+use hyperactor::sync::flag;
 use hyperactor::sync::monitor;
 use ndslice::Shape;
+use nix::libc::exit;
 use tokio::io;
 use tokio::process::Command;
 use tokio::sync::Mutex;
@@ -130,8 +133,9 @@ enum ChannelState {
 struct Child {
     channel: ChannelState,
     group: monitor::Group,
-    stdout: LogTailer,
+    exit_flag: Option<flag::Flag>,
     stderr: LogTailer,
+    stdout: LogTailer,
     stop_reason: Arc<OnceLock<ProcStopReason>>,
 }
 
@@ -140,6 +144,7 @@ impl Child {
         mut process: tokio::process::Child,
     ) -> (Self, impl Future<Output = ProcStopReason>) {
         let (group, handle) = monitor::group();
+        let (exit_flag, exit_guard) = flag::guarded();
 
         let stdout = LogTailer::tee(
             MAX_TAIL_LOG_LINES,
@@ -156,6 +161,7 @@ impl Child {
         let child = Self {
             channel: ChannelState::NotConnected,
             group,
+            exit_flag: Some(exit_flag),
             stdout,
             stderr,
             stop_reason: Arc::clone(&stop_reason),
@@ -178,6 +184,7 @@ impl Child {
                 }
                 result = process.wait() => Self::exit_status_to_reason(result),
             };
+            exit_guard.signal();
             stop_reason.get_or_init(|| reason).clone()
         };
 
@@ -240,6 +247,18 @@ impl Child {
         // be killed.
         if let ChannelState::Connected(channel) = &mut self.channel {
             channel.post(message);
+            if let Some(exit_flag) = self.exit_flag.take() {
+                let group = self.group.clone();
+                let stop_reason = self.stop_reason.clone();
+                tokio::spawn(async move {
+                    let exit_timeout =
+                        hyperactor::config::global::get(hyperactor::config::PROCESS_EXIT_TIMEOUT);
+                    if tokio::time::timeout(exit_timeout, exit_flag).await.is_err() {
+                        let _ = stop_reason.set(ProcStopReason::Watchdog);
+                        group.fail();
+                    }
+                });
+            }
         } else {
             self.stop(ProcStopReason::Watchdog);
         }

hyperactor_mesh/src/alloc/remoteprocess.rs

@@ -1053,7 +1053,7 @@ mod test {
         }
     }
 
-    fn set_procstate_execptations(alloc: &mut MockAlloc, shape: Shape) {
+    fn set_procstate_expectations(alloc: &mut MockAlloc, shape: Shape) {
         let alloc_len = shape.slice().len();
         alloc.expect_shape().return_const(shape.clone());
         for i in 0..alloc_len {
@@ -1108,7 +1108,7 @@ mod test {
         alloc.expect_world_id().return_const(world_id.clone());
         alloc.expect_shape().return_const(spec.shape.clone());
 
-        set_procstate_execptations(&mut alloc, spec.shape.clone());
+        set_procstate_expectations(&mut alloc, spec.shape.clone());
 
         // final none
         alloc.expect_next().return_const(None);
@@ -1247,7 +1247,7 @@ mod test {
         alloc.alloc.expect_world_id().return_const(world_id.clone());
         alloc.alloc.expect_shape().return_const(spec.shape.clone());
 
-        set_procstate_execptations(&mut alloc.alloc, spec.shape.clone());
+        set_procstate_expectations(&mut alloc.alloc, spec.shape.clone());
 
         alloc.alloc.expect_next().return_const(None);
         alloc.alloc.expect_stop().times(1).return_once(|| Ok(()));
@@ -1324,7 +1324,7 @@ mod test {
             .return_const(world_id.clone());
         alloc1.alloc.expect_shape().return_const(spec.shape.clone());
 
-        set_procstate_execptations(&mut alloc1.alloc, spec.shape.clone());
+        set_procstate_expectations(&mut alloc1.alloc, spec.shape.clone());
         alloc1.alloc.expect_next().return_const(None);
         alloc1.alloc.expect_stop().times(1).return_once(|| Ok(()));
         // second allocation
@@ -1339,7 +1339,7 @@ mod test {
             .expect_world_id()
             .return_const(world_id.clone());
         alloc2.alloc.expect_shape().return_const(spec.shape.clone());
-        set_procstate_execptations(&mut alloc2.alloc, spec.shape.clone());
+        set_procstate_expectations(&mut alloc2.alloc, spec.shape.clone());
         alloc2.alloc.expect_next().return_const(None);
         alloc2.alloc.expect_stop().times(1).return_once(|| Ok(()));
 
@@ -1444,7 +1444,7 @@ mod test {
         alloc.alloc.expect_world_id().return_const(world_id.clone());
         alloc.alloc.expect_shape().return_const(spec.shape.clone());
 
-        set_procstate_execptations(&mut alloc.alloc, spec.shape.clone());
+        set_procstate_expectations(&mut alloc.alloc, spec.shape.clone());
 
         alloc.alloc.expect_next().return_const(None);
         // we expect a stop due to the failure