hyperactor/init.rs
1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::sync::OnceLock;
10
11use hyperactor_telemetry::DefaultTelemetryClock;
12
13use crate::panic_handler;
14
15/// A global runtime handle used for spawning tasks. Do not use for executing long running or
16/// compute intensive tasks.
17static RUNTIME: OnceLock<tokio::runtime::Handle> = OnceLock::new();
18
19/// Get a handle to the global runtime.
20///
21/// Panics if the runtime has not been initialized *and* the caller is not in an
22/// async context.
23pub(crate) fn get_runtime() -> tokio::runtime::Handle {
24 match RUNTIME.get() {
25 Some(handle) => handle.clone(),
26 None => tokio::runtime::Handle::current(),
27 }
28}
29
30/// Initialize the Hyperactor runtime. Specifically:
31/// - Set up panic handling, so that we get consistent panic stack traces in Actors.
32/// - Initialize logging defaults.
33/// - Store the provided tokio runtime handle for use by the hyperactor system.
34pub fn initialize(handle: tokio::runtime::Handle) {
35 initialize_with_log_prefix(handle, Option::None);
36}
37
38/// Initialize the Hyperactor runtime. Specifically:
39/// - Set up panic handling, so that we get consistent panic stack traces in Actors.
40/// - Initialize logging defaults.
41/// - Store the provided tokio runtime handle for use by the hyperactor system.
42/// - Set the env var whose value should be used to prefix log messages.
43pub fn initialize_with_log_prefix(
44 handle: tokio::runtime::Handle,
45 env_var_log_prefix: Option<String>,
46) {
47 RUNTIME
48 .set(handle)
49 .expect("hyperactor::initialize must only be called once");
50
51 panic_handler::set_panic_hook();
52 hyperactor_telemetry::initialize_logging_with_log_prefix(
53 DefaultTelemetryClock {},
54 env_var_log_prefix,
55 );
56}
57
58/// Initialize the Hyperactor runtime using the current tokio runtime handle.
59pub fn initialize_with_current_runtime() {
60 let handle = tokio::runtime::Handle::current();
61 initialize(handle);
62}