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.

Reviewed By: suo

Differential Revision: D77303504

Author: vidhyav

hyperactor/src/config.rs

@@ -32,6 +32,12 @@ 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);
+
+    /// Process exit timeout
+    pub attr ACTOR_EXIT_TIMEOUT: Duration = Duration::from_millis(10);
+
     /// Message acknowledgment interval
     pub attr MESSAGE_ACK_TIME_INTERVAL: Duration = Duration::from_millis(500);
 
@@ -86,6 +92,13 @@ pub fn from_env() -> Attrs {
 
     config[IS_MANAGED_SUBPROCESS] = env::var("HYPERACTOR_MANAGED_SUBPROCESS").is_ok();
 
+    // Load message ack time interval
+    if let Ok(val) = env::var("ACTOR_EXIT_TIMEOUT_MS") {
+        if let Ok(parsed) = val.parse::<u64>() {
+            config[ACTOR_EXIT_TIMEOUT] = Duration::from_millis(parsed);
+        }
+    }
+
     config
 }
 
@@ -112,6 +125,9 @@ pub fn merge(config: &mut Attrs, other: &Attrs) {
     if other.contains_key(MESSAGE_DELIVERY_TIMEOUT) {
         config[MESSAGE_DELIVERY_TIMEOUT] = other[MESSAGE_DELIVERY_TIMEOUT];
     }
+    if other.contains_key(ACTOR_EXIT_TIMEOUT) {
+        config[ACTOR_EXIT_TIMEOUT] = other[ACTOR_EXIT_TIMEOUT];
+    }
     if other.contains_key(MESSAGE_ACK_TIME_INTERVAL) {
         config[MESSAGE_ACK_TIME_INTERVAL] = other[MESSAGE_ACK_TIME_INTERVAL];
     }

hyperactor/src/sync/monitor.rs

@@ -103,6 +103,7 @@ impl IntoFuture for Handle {
 /// the group will be aborted if any task fails or if the group is aborted.
 ///
 /// The group is also aborted if the group itself is dropped.
+#[derive(Clone)]
 pub struct Group(Arc<Mutex<State>>);
 
 /// The status of a group. Groups start out in [`Status::Running`]

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> {
+            loop {
+                std::thread::sleep(Duration::from_secs(60));
+            }
+        }
+    }
+
     /// 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,148 @@ 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)
+    }
+
+    /// In order to simulate stuckness, we have to do two things:
+    /// An actor that is blocked forever AND
+    /// a proc that does not time out when it is asked to wait for
+    /// a stuck actor.
+    #[tokio::test]
+    async fn test_allocator_stuck_task() {
+        // 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 command =
+            Command::new(buck_resources::get("monarch/hyperactor_mesh/bootstrap").unwrap());
+        command.env("ACTOR_EXIT_TIMEOUT_MS", "300000");
+        let mut allocator = ProcessAllocator::new(command);
+        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

@@ -28,6 +28,7 @@ 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 tokio::io;
@@ -130,8 +131,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 +142,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 +159,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 +182,8 @@ impl Child {
                 }
                 result = process.wait() => Self::exit_status_to_reason(result),
             };
+            exit_guard.signal();
+
             stop_reason.get_or_init(|| reason).clone()
         };
 
@@ -234,10 +240,23 @@ impl Child {
         true
     }
 
+    fn spawn_watchdog(&mut self) {
+        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();
+                }
+            });
+        }
+    }
+
     #[hyperactor::instrument_infallible]
     fn post(&mut self, message: Allocator2Process) {
-        // We're here simply assuming that if we're not connected, we're about to
-        // be killed.
         if let ChannelState::Connected(channel) = &mut self.channel {
             channel.post(message);
         } else {
@@ -439,6 +458,7 @@ impl Alloc for ProcessAlloc {
         // for liveness.
         for (_index, child) in self.active.iter_mut() {
             child.post(Allocator2Process::StopAndExit(0));
+            child.spawn_watchdog();
         }
 
         self.running = false;