monarch_tensor_worker/
test_util.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::io::IsTerminal;
10
11use anyhow::Result;
12use monarch_gil::GilSite;
13use monarch_gil::monarch_with_gil_blocking;
14use pyo3::Python;
15use pyo3::ffi::c_str;
16use tracing_subscriber::fmt::format::FmtSpan;
17
18pub fn test_setup() -> Result<()> {
19    let _ = tracing_subscriber::fmt()
20        .with_thread_ids(true)
21        .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
22        .with_max_level(tracing::Level::DEBUG)
23        .with_ansi(std::io::stderr().is_terminal())
24        .with_writer(std::io::stderr)
25        .try_init();
26
27    // Redirect NCCL_DEBUG log output to a file so it doesn't clash on stdout.
28    // TestX requires stdout to have JSON output on individual lines, and
29    // the NCCL output is not JSON. Because it runs in a different thread, it'll
30    // race on writing to stdout.
31    // Do this regardless of whether NCCL_DEBUG is set or not, because it can
32    // be set after this point in the test. If it doesn't get set, NCCL_DEBUG_FILE
33    // will be ignored.
34    // %h becomes hostname, %p becomes pid.
35    let nccl_debug_file = std::env::temp_dir().join("nccl_debug.%h.%p");
36    tracing::debug!("Set NCCL_DEBUG_FILE to {:?}", nccl_debug_file);
37    // Safety: Can be unsound if there are multiple threads
38    // reading and writing the environment.
39    unsafe {
40        std::env::set_var("NCCL_DEBUG_FILE", nccl_debug_file);
41    }
42    // NOTE(agallagher): Calling `Python::initialize()` appears to
43    // clear `PYTHONPATH` in the env, which we need for test subprocesses
44    // to work.  So, manually preserve it.
45    let py_path = std::env::var("PYTHONPATH");
46    Python::initialize();
47    if let Ok(py_path) = py_path {
48        // SAFETY: Re-setting env var cleared by `Python::initialize()`.
49        unsafe { std::env::set_var("PYTHONPATH", py_path) }
50    }
51
52    // We need to load torch to initialize some internal structures used by
53    // the FFI funcs we use to convert ivalues to/from py objects.
54    monarch_with_gil_blocking(GilSite::Test, |py| {
55        py.run(c_str!("import torch"), None, None)
56    })?;
57
58    Ok(())
59}
monarch_tensor_worker/test_util.rs

monarch_tensor_worker/
test_util.rs
