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}