: python: actor: shutdown_context() terminates this_host() (#2178)

Summary:

now `shutdown_context()` terminates the OS process running `this_host()`.

Differential Revision: D89498869

Author: shayne-fletcher

hyperactor/src/init.rs

@@ -19,7 +19,7 @@ static RUNTIME: OnceLock<tokio::runtime::Handle> = OnceLock::new();
 ///
 /// Panics if the runtime has not been initialized *and* the caller is not in an
 /// async context.
-pub(crate) fn get_runtime() -> tokio::runtime::Handle {
+pub fn get_runtime() -> tokio::runtime::Handle {
     match RUNTIME.get() {
         Some(handle) => handle.clone(),
         None => tokio::runtime::Handle::current(),

hyperactor/src/lib.rs

@@ -142,6 +142,8 @@ pub use hyperactor_telemetry::declare_static_timer;
 pub use hyperactor_telemetry::key_value;
 pub use hyperactor_telemetry::kv_pairs;
 #[doc(inline)]
+pub use init::get_runtime;
+#[doc(inline)]
 pub use init::initialize;
 #[doc(inline)]
 pub use init::initialize_with_current_runtime;

monarch_hyperactor/src/v1/host_mesh.rs

@@ -344,13 +344,14 @@ fn py_host_mesh_from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<PyHostMesh> {
 }
 
 #[pyfunction]
-fn shutdown_local_host_mesh() -> PyResult<PyPythonTask> {
+fn shutdown_local_host_mesh() -> PyResult<()> {
     let agent = HOST_MESH_AGENT_FOR_HOST
         .get()
         .ok_or_else(|| PyException::new_err("No local host mesh to shutdown"))?
         .clone();
 
-    PyPythonTask::new(async move {
+    // Block on the async shutdown operation
+    hyperactor::get_runtime().block_on(async move {
         // Create a temporary instance to send the shutdown message
         let temp_proc = hyperactor::Proc::local();
         let (instance, _) = temp_proc
@@ -365,8 +366,11 @@ fn shutdown_local_host_mesh() -> PyResult<PyPythonTask> {
             .await
             .map_err(|e| PyException::new_err(e.to_string()))?;
 
-        Ok(())
-    })
+        Ok::<(), pyo3::PyErr>(())
+    })?;
+
+    // Exit the process
+    std::process::exit(0)
 }
 
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {

python/monarch/_rust_bindings/monarch_hyperactor/v1/host_mesh.pyi

@@ -114,15 +114,19 @@ def bootstrap_host(
     """
     ...
 
-def shutdown_local_host_mesh() -> PythonTask[None]:
+def shutdown_local_host_mesh() -> None:
     """
-    Shutdown the local host mesh created by bootstrap_host().
+    Shutdown the local host mesh created by bootstrap_host() and exit the process.
+
+    This function blocks until shutdown completes. On successful shutdown, it
+    calls exit(0) and never returns. On failure, it raises an exception.
 
     Sends ShutdownHost message to the local host mesh agent with:
     - timeout: 10 seconds grace period before SIGTERM escalation
     - max_in_flight: 16 concurrent child terminations
 
     Raises:
         RuntimeError: If no local host mesh exists (bootstrap_host not called)
+                      or if shutdown fails
     """
     ...

python/monarch/_src/actor/actor_mesh.py

@@ -388,7 +388,9 @@ def shutdown_context() -> "Future[None]":
                 shutdown_local_host_mesh,
             )
 
-            return Future(coro=shutdown_local_host_mesh())
+            # This function blocks and then exits the process
+            shutdown_local_host_mesh()
+            # Never reached
         except RuntimeError:
             # No local host mesh to shutdown
             pass