get rid of hyperactor runtime, attempt 2 (#406)

Summary:
Pull Request resolved: #406

Based on the discussion in D77422927, a new attempt:

I still want to avoid creating an extra runtime in hyperactor, which I think is a Pretty Bad thing to do.

So here, I have `hyperactor::initialize` take a runtime handle provided by the caller. This ensures the user has full control over the runtime, but that we still have a reference to a runtime handy in hyperactor when we need one from a sync context.
ghstack-source-id: 293872756
exported-using-ghexport

Reviewed By: mariusae

Differential Revision: D77632354

fbshipit-source-id: 472333d1544d11e1ee7b02b40d17d0bf3b5c785c

Author: suo

controller/src/main.rs

@@ -69,7 +69,7 @@ async fn serve(inp: impl AsyncBufRead + Unpin, mut outp: impl AsyncWrite + Unpin
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    hyperactor::initialize();
+    hyperactor::initialize_with_current_runtime();
 
     match BootstrapCommand::try_parse()? {
         BootstrapCommand::Run(cmd) => controller::bootstrap::run(cmd)?.await??,

hyper/src/main.rs

@@ -63,7 +63,7 @@ async fn main() -> Result<(), anyhow::Error> {
 
 async fn run() -> Result<(), anyhow::Error> {
     let args = Cli::parse();
-    hyperactor::initialize();
+    hyperactor::initialize_with_current_runtime();
 
     match args.command {
         Command::Serve(command) => Ok(command.run().await?),

hyperactor/src/channel/net.rs

@@ -190,7 +190,7 @@ impl<M: RemoteMessage> NetTx<M> {
             dest,
             status,
         };
-        crate::init::RUNTIME.spawn(Self::run(link, receiver, notify));
+        crate::init::get_runtime().spawn(Self::run(link, receiver, notify));
         tx
     }
 

hyperactor/src/init.rs

@@ -6,30 +6,45 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-//! Utilities for launching hyperactor processes.
-
-use std::sync::LazyLock;
 use std::sync::OnceLock;
 
 use crate::clock::ClockKind;
 use crate::panic_handler;
 
-/// A global runtime used in binding async and sync code. Do not use for executing long running or
+/// A global runtime handle used for spawning tasks. Do not use for executing long running or
 /// compute intensive tasks.
-pub(crate) static RUNTIME: LazyLock<tokio::runtime::Runtime> =
-    LazyLock::new(|| tokio::runtime::Runtime::new().expect("failed to create global runtime"));
+static RUNTIME: OnceLock<tokio::runtime::Handle> = OnceLock::new();
+
+/// Get a handle to the global runtime.
+///
+/// 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 {
+    match RUNTIME.get() {
+        Some(handle) => handle.clone(),
+        None => tokio::runtime::Handle::current(),
+    }
+}
 
 /// Initialize the Hyperactor runtime. Specifically:
 /// - Set up panic handling, so that we get consistent panic stack traces in Actors.
 /// - Initialize logging defaults.
-pub fn initialize() {
-    static INITIALIZED: OnceLock<()> = OnceLock::new();
-    INITIALIZED.get_or_init(|| {
-        panic_handler::set_panic_hook();
-        hyperactor_telemetry::initialize_logging(ClockKind::default());
-        #[cfg(target_os = "linux")]
-        linux::initialize();
-    });
+/// - Store the provided tokio runtime handle for use by the hyperactor system.
+pub fn initialize(handle: tokio::runtime::Handle) {
+    RUNTIME
+        .set(handle)
+        .expect("hyperactor::initialize must only be called once");
+
+    panic_handler::set_panic_hook();
+    hyperactor_telemetry::initialize_logging(ClockKind::default());
+    #[cfg(target_os = "linux")]
+    linux::initialize();
+}
+
+/// Initialize the Hyperactor runtime using the current tokio runtime handle.
+pub fn initialize_with_current_runtime() {
+    let handle = tokio::runtime::Handle::current();
+    initialize(handle);
 }
 
 #[cfg(target_os = "linux")]

hyperactor/src/lib.rs

@@ -135,6 +135,8 @@ pub use hyperactor_telemetry::key_value;
 pub use hyperactor_telemetry::kv_pairs;
 #[doc(inline)]
 pub use init::initialize;
+#[doc(inline)]
+pub use init::initialize_with_current_runtime;
 // Re-exported to make this available to callers of the `register!` macro.
 #[doc(hidden)]
 pub use inventory::submit;

hyperactor/src/mailbox.rs

@@ -684,7 +684,7 @@ impl<T: Message> Buffer<T> {
     {
         let (queue, mut next) = mpsc::unbounded_channel();
         let (last_processed, processed) = watch::channel(0);
-        crate::init::RUNTIME.spawn(async move {
+        crate::init::get_runtime().spawn(async move {
             let mut seq = 0;
             while let Some((msg, return_handle)) = next.recv().await {
                 process(msg, return_handle).await;
@@ -931,7 +931,7 @@ impl MailboxClient {
         cancel_token: CancellationToken,
         addr: ChannelAddr,
     ) {
-        crate::init::RUNTIME.spawn(async move {
+        crate::init::get_runtime().spawn(async move {
             loop {
                 tokio::select! {
                     changed = rx.changed() => {

hyperactor/src/mailbox/undeliverable.rs

@@ -63,7 +63,7 @@ pub fn monitored_return_handle() -> PortHandle<Undeliverable<MessageEnvelope>> {
         // Don't reuse `return_handle` for `h`: else it will never get
         // dropped and the task will never return.
         let (h, _) = new_undeliverable_port();
-        crate::init::RUNTIME.spawn(async move {
+        crate::init::get_runtime().spawn(async move {
             while let Ok(Undeliverable(mut envelope)) = rx.recv().await {
                 envelope.try_set_error(DeliveryError::BrokenLink(
                     "message returned to undeliverable port".to_string(),

hyperactor_mesh/src/bootstrap.rs

@@ -78,10 +78,6 @@ pub(crate) enum Allocator2Process {
 /// Use [`bootstrap_or_die`] to implement this behavior directly.
 pub async fn bootstrap() -> anyhow::Error {
     pub async fn go() -> Result<(), anyhow::Error> {
-        // This ensures, among other things, that we correctly exit child processes
-        // when the parent dies.
-        hyperactor::initialize();
-
         let bootstrap_addr: ChannelAddr = std::env::var(BOOTSTRAP_ADDR_ENV)
             .map_err(|err| anyhow::anyhow!("read `{}`: {}", BOOTSTRAP_ADDR_ENV, err))?
             .parse()?;

monarch_extension/src/lib.rs

@@ -64,8 +64,9 @@ fn get_or_add_new_module<'py>(
 #[pymodule]
 #[pyo3(name = "_rust_bindings")]
 pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    ::hyperactor::initialize();
     monarch_hyperactor::runtime::initialize(module.py())?;
+    let runtime = monarch_hyperactor::runtime::get_tokio_runtime();
+    ::hyperactor::initialize(runtime.handle().clone());
 
     monarch_hyperactor::shape::register_python_bindings(&get_or_add_new_module(
         module,

monarch_extension/src/tensor_worker.rs

@@ -26,6 +26,7 @@ use hyperactor::data::Serialized;
 use hyperactor::reference::ActorId;
 use monarch_hyperactor::ndslice::PySlice;
 use monarch_hyperactor::proc::PyActorId;
+use monarch_hyperactor::runtime::get_tokio_runtime;
 use monarch_messages::wire_value::WireValue;
 use monarch_messages::wire_value::func_call_args_to_wire_values;
 use monarch_messages::worker::*;
@@ -1386,22 +1387,17 @@ fn worker_main(py: Python<'_>) -> PyResult<()> {
             BinaryArgs::Pipe => bootstrap_pipe(),
             BinaryArgs::WorkerServer { rd, wr } => {
                 worker_server(
+                    get_tokio_runtime(),
                     // SAFETY: Raw FD passed in from parent.
                     BufReader::new(File::from(unsafe { OwnedFd::from_raw_fd(rd) })),
                     // SAFETY: Raw FD passed in from parent.
                     File::from(unsafe { OwnedFd::from_raw_fd(wr) }),
                 )
             }
-            BinaryArgs::Worker(args) => {
-                let rt = tokio::runtime::Builder::new_multi_thread()
-                    .enable_all()
-                    .build()?;
-                rt.block_on(async move {
-                    hyperactor::initialize();
-                    let _ = bootstrap_worker_proc(args).await?.await;
-                    Ok(())
-                })
-            }
+            BinaryArgs::Worker(args) => get_tokio_runtime().block_on(async move {
+                let _ = bootstrap_worker_proc(args).await?.await;
+                Ok(())
+            }),
         }
         .map_err(|err: anyhow::Error| PyRuntimeError::new_err(err.to_string()))
     })