Send SIGTERM to remote process alloc children (#399)

Summary:
Pull Request resolved: #399

- Kill child processes with SIGTERM
- Setup signal listener for SIGTERM and exit gracefully on SIGTERM

Reviewed By: dulinriley

Differential Revision: D77614649

fbshipit-source-id: 3657c131187b09c2017ba43d9c796d0079d36f41

Author: thomasywang

hyperactor_mesh/Cargo.toml

@@ -47,6 +47,8 @@ rand = { version = "0.8", features = ["small_rng"] }
 serde = { version = "1.0.185", features = ["derive", "rc"] }
 serde_bytes = "0.11"
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "unbounded_depth"] }
+signal-hook = "0.3"
+signal-hook-tokio = { version = "0.3", features = ["futures-v0_3"] }
 tempfile = "3.15"
 thiserror = "2.0.12"
 tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }

hyperactor_mesh/src/alloc/process.rs

@@ -34,6 +34,8 @@ use hyperactor::sync::flag;
 use hyperactor::sync::monitor;
 use hyperactor_state::state_actor::StateActor;
 use ndslice::Shape;
+use nix::sys::signal;
+use nix::unistd::Pid;
 use tokio::io;
 use tokio::process::Command;
 use tokio::sync::Mutex;
@@ -204,17 +206,15 @@ impl Child {
         let monitor = async move {
             let reason = tokio::select! {
                 _ = handle => {
-                    match process.kill().await {
-                        Err(e) => {
-                            tracing::error!("error killing process: {}", e);
-                            // In this cased, we're left with little choice but to
-                            // orphan the process.
-                            ProcStopReason::Unknown
-                        },
-                        Ok(_) => {
-                            Self::exit_status_to_reason(process.wait().await)
-                        }
-                    }
+                    let Some(id) = process.id() else {
+                        tracing::error!("could not get child process id");
+                        return ProcStopReason::Unknown;
+                    };
+                    if let Err(e) = signal::kill(Pid::from_raw(id as i32), signal::SIGTERM) {
+                        tracing::error!("failed to kill child process: {}", e);
+                        return ProcStopReason::Unknown;
+                    };
+                    Self::exit_status_to_reason(process.wait().await)
                 }
                 result = process.wait() => Self::exit_status_to_reason(result),
             };
@@ -301,6 +301,11 @@ impl Child {
             self.stop(ProcStopReason::Watchdog);
         }
     }
+
+    #[cfg(test)]
+    fn fail_group(&self) {
+        self.group.fail();
+    }
 }
 
 impl ProcessAlloc {
@@ -516,4 +521,47 @@ mod tests {
     crate::alloc_test_suite!(ProcessAllocator::new(Command::new(
         buck_resources::get("monarch/hyperactor_mesh/bootstrap").unwrap()
     )));
+
+    #[tokio::test]
+    async fn test_sigterm_on_group_fail() {
+        let bootstrap_binary = buck_resources::get("monarch/hyperactor_mesh/bootstrap").unwrap();
+        let mut allocator = ProcessAllocator::new(Command::new(bootstrap_binary));
+
+        let mut alloc = allocator
+            .allocate(AllocSpec {
+                shape: ndslice::shape! { replica = 1 },
+                constraints: Default::default(),
+            })
+            .await
+            .unwrap();
+
+        let proc_id = {
+            loop {
+                match alloc.next().await {
+                    Some(ProcState::Running { proc_id, .. }) => {
+                        break proc_id;
+                    }
+                    Some(ProcState::Failed { description, .. }) => {
+                        panic!("Process allocation failed: {}", description);
+                    }
+                    Some(_other) => {}
+                    None => {
+                        panic!("Allocation ended unexpectedly");
+                    }
+                }
+            }
+        };
+
+        if let Some(child) = alloc.active.get(&proc_id.rank()) {
+            child.fail_group();
+        }
+
+        assert!(matches!(
+            alloc.next().await,
+            Some(ProcState::Stopped {
+                reason: ProcStopReason::Killed(15, false),
+                ..
+            })
+        ));
+    }
 }

hyperactor_mesh/src/bootstrap.rs

@@ -8,6 +8,7 @@
 
 use std::time::Duration;
 
+use futures::StreamExt;
 use hyperactor::ActorRef;
 use hyperactor::Named;
 use hyperactor::ProcId;
@@ -21,6 +22,7 @@ use hyperactor::clock::RealClock;
 use hyperactor::mailbox::MailboxServer;
 use serde::Deserialize;
 use serde::Serialize;
+use signal_hook::consts::signal::SIGTERM;
 
 use crate::proc_mesh::mesh_agent::MeshAgent;
 
@@ -119,6 +121,8 @@ async fn exit_if_missed_heartbeat(bootstrap_index: usize, bootstrap_addr: Channe
 /// Use [`bootstrap_or_die`] to implement this behavior directly.
 pub async fn bootstrap() -> anyhow::Error {
     pub async fn go() -> Result<(), anyhow::Error> {
+        let mut signals = signal_hook_tokio::Signals::new([SIGTERM])?;
+
         let bootstrap_addr: ChannelAddr = std::env::var(BOOTSTRAP_ADDR_ENV)
             .map_err(|err| anyhow::anyhow!("read `{}`: {}", BOOTSTRAP_ADDR_ENV, err))?
             .parse()?;
@@ -141,45 +145,76 @@ pub async fn bootstrap() -> anyhow::Error {
 
         loop {
             let _ = hyperactor::tracing::info_span!("wait_for_next_message_from_mesh_agent");
-            match rx.recv().await? {
-                Allocator2Process::StartProc(proc_id, listen_transport) => {
-                    let (proc, mesh_agent) = MeshAgent::bootstrap(proc_id.clone()).await?;
-                    let (proc_addr, proc_rx) =
-                        channel::serve(ChannelAddr::any(listen_transport)).await?;
-                    // Undeliverable messages get forwarded to the mesh agent.
-                    let handle = proc.clone().serve(proc_rx, mesh_agent.port());
-                    drop(handle); // linter appeasement; it is safe to drop this future
-                    tx.send(Process2Allocator(
-                        bootstrap_index,
-                        Process2AllocatorMessage::StartedProc(
-                            proc_id.clone(),
-                            mesh_agent.bind(),
-                            proc_addr,
-                        ),
-                    ))
-                    .await?;
-                    procs.push(proc);
-                }
-                Allocator2Process::StopAndExit(code) => {
-                    tracing::info!("stopping procs with code {code}");
-                    for mut proc_to_stop in procs {
-                        if let Err(err) = proc_to_stop
-                            .destroy_and_wait(Duration::from_millis(10), None)
-                            .await
-                        {
-                            tracing::error!(
-                                "error while stopping proc {}: {}",
-                                proc_to_stop.proc_id(),
-                                err
-                            );
+            tokio::select! {
+                msg = rx.recv() => {
+                    match msg? {
+                        Allocator2Process::StartProc(proc_id, listen_transport) => {
+                            let (proc, mesh_agent) = MeshAgent::bootstrap(proc_id.clone()).await?;
+                            let (proc_addr, proc_rx) =
+                                channel::serve(ChannelAddr::any(listen_transport)).await?;
+                            // Undeliverable messages get forwarded to the mesh agent.
+                            let handle = proc.clone().serve(proc_rx, mesh_agent.port());
+                            drop(handle); // linter appeasement; it is safe to drop this future
+                            tx.send(Process2Allocator(
+                                bootstrap_index,
+                                Process2AllocatorMessage::StartedProc(
+                                    proc_id.clone(),
+                                    mesh_agent.bind(),
+                                    proc_addr,
+                                ),
+                            ))
+                            .await?;
+                            procs.push(proc);
+                        }
+                        Allocator2Process::StopAndExit(code) => {
+                            tracing::info!("stopping procs with code {code}");
+                            for mut proc_to_stop in procs {
+                                if let Err(err) = proc_to_stop
+                                    .destroy_and_wait(Duration::from_millis(10), None)
+                                    .await
+                                {
+                                    tracing::error!(
+                                        "error while stopping proc {}: {}",
+                                        proc_to_stop.proc_id(),
+                                        err
+                                    );
+                                }
+                            }
+                            tracing::info!("exiting with {code}");
+                            std::process::exit(code);
+                        }
+                        Allocator2Process::Exit(code) => {
+                            tracing::info!("exiting with {code}");
+                            std::process::exit(code);
                         }
                     }
-                    tracing::info!("exiting with {code}");
-                    std::process::exit(code);
                 }
-                Allocator2Process::Exit(code) => {
-                    tracing::info!("exiting with {code}");
-                    std::process::exit(code);
+                signal = signals.next() => {
+                    if signal.is_some_and(|sig| sig == SIGTERM) {
+                        tracing::info!("received SIGTERM, stopping procs");
+                        for mut proc_to_stop in procs {
+                            if let Err(err) = proc_to_stop
+                                .destroy_and_wait(Duration::from_millis(10), None)
+                                .await
+                            {
+                                tracing::error!(
+                                    "error while stopping proc {}: {}",
+                                    proc_to_stop.proc_id(),
+                                    err
+                                );
+                            }
+                        }
+                        // SAFETY: We're setting the handle to SigDfl (defautl system behaviour)
+                        if let Err(err) = unsafe {
+                            nix::sys::signal::signal(nix::sys::signal::SIGTERM, nix::sys::signal::SigHandler::SigDfl)
+                        } {
+                            tracing::error!("failed to signal SIGTERM: {}", err);
+                        }
+                        if let Err(err) = nix::sys::signal::raise(nix::sys::signal::SIGTERM) {
+                            tracing::error!("failed to raise SIGTERM: {}", err);
+                        }
+                        std::process::exit(128 + SIGTERM);
+                    }
                 }
             }
         }