monarch_hyperactor/
bootstrap.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 futures::future::try_join_all;
10use hyperactor::Gateway;
11use hyperactor::channel::ChannelAddr;
12use hyperactor::id::Label;
13use hyperactor_mesh::bootstrap::BootstrapCommand;
14use hyperactor_mesh::bootstrap::bootstrap;
15use hyperactor_mesh::bootstrap::halt;
16use hyperactor_mesh::bootstrap::host;
17use hyperactor_mesh::host_mesh::HostMesh;
18use hyperactor_mesh::mesh_id::HostMeshId;
19use monarch_types::MapPyErr;
20use pyo3::Bound;
21use pyo3::PyAny;
22use pyo3::PyResult;
23use pyo3::Python;
24use pyo3::exceptions::PyRuntimeError;
25use pyo3::pyfunction;
26use pyo3::types::PyAnyMethods;
27use pyo3::types::PyModule;
28use pyo3::types::PyModuleMethods;
29use pyo3::wrap_pyfunction;
30
31use crate::host_mesh::PyHostMesh;
32use crate::pytokio::PyPythonTask;
33use crate::runtime::GilSite;
34use crate::runtime::monarch_with_gil;
35
36#[pyfunction]
37#[pyo3(signature = ())]
38pub fn bootstrap_main(py: Python) -> PyResult<Bound<PyAny>> {
39    // SAFETY: this is a correct use of this function.
40    unsafe {
41        fbinit::perform_init();
42    };
43
44    hyperactor::internal_macro_support::tracing::debug!("entering async bootstrap");
45    crate::runtime::future_into_py::<_, i32>(py, async move {
46        // SAFETY:
47        // - Only one of these is ever created.
48        // - This is the entry point of this program, so this will be dropped when
49        // no more FB C++ code is running.
50        let _destroy_guard = unsafe { fbinit::DestroyGuard::new() };
51        bootstrap()
52            .await
53            .map_err(|e| PyRuntimeError::new_err(format!("{:?}", e)))
54    })
55}
56
57#[pyfunction]
58pub fn run_worker_loop_forever(_py: Python<'_>, address: &str) -> PyResult<PyPythonTask> {
59    let (addr, listener) = ChannelAddr::from_zmq_url_with_listener(address)?;
60
61    // Check if we're running in a PAR/XAR build by looking for FB_XAR_INVOKED_NAME environment variable
62    let invoked_name = std::env::var("FB_XAR_INVOKED_NAME");
63
64    let mut env: std::collections::HashMap<String, String> = std::env::vars().collect();
65
66    let command = Some(if let Ok(invoked_name) = invoked_name {
67        // For PAR/XAR builds: use argv[0] from Python's sys.argv as the current executable
68        let current_exe = std::path::PathBuf::from(&invoked_name);
69
70        // For PAR/XAR builds: set PAR_MAIN_OVERRIDE and no additional args
71        env.insert(
72            "PAR_MAIN_OVERRIDE".to_string(),
73            "monarch._src.actor.bootstrap_main".to_string(),
74        );
75        BootstrapCommand {
76            program: current_exe,
77            arg0: Some(invoked_name),
78            args: vec![],
79            env,
80        }
81    } else {
82        // For regular Python builds: use argv[0] to preserve the original
83        // invocation path.  current_exe() resolves symlinks, which breaks
84        // virtual environments — the resolved path doesn't find pyvenv.cfg
85        // so site-packages aren't activated in subprocesses.
86        let current_exe = std::env::args()
87            .next()
88            .map(std::path::PathBuf::from)
89            .or_else(|| std::env::current_exe().ok())
90            .ok_or_else(|| {
91                pyo3::PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
92                    "Failed to determine current executable",
93                )
94            })?;
95        let current_exe_str = current_exe.to_string_lossy().to_string();
96        BootstrapCommand {
97            program: current_exe,
98            arg0: Some(current_exe_str),
99            args: vec![
100                "-m".to_string(),
101                "monarch._src.actor.bootstrap_main".to_string(),
102            ],
103            env,
104        }
105    });
106
107    PyPythonTask::new(async move {
108        let (_agent_handle, shutdown) =
109            host(addr, command, None, true, listener, Gateway::new(), None)
110                .await
111                .map_pyerr()?;
112        shutdown.join().await;
113        halt::<()>().await;
114        Ok(())
115    })
116}
117
118#[pyfunction]
119pub fn attach_to_workers(
120    instance: &crate::context::PyInstance,
121    workers: Vec<Bound<'_, PyPythonTask>>,
122    name: Option<&str>,
123) -> PyResult<PyPythonTask> {
124    let tasks = workers
125        .into_iter()
126        .map(|x| x.borrow_mut().take_task())
127        .collect::<PyResult<Vec<_>>>()?;
128
129    // `Label::strip` (vs. `Label::new`) sanitizes user-supplied names — lowercases,
130    // drops illegal characters, falls back to "nil" if empty. Callers pass names
131    // derived from experiment / job names that may contain uppercase or punctuation;
132    // rejecting them surfaces as an opaque PyException far from the input site.
133    let name = HostMeshId::instance(Label::strip(name.unwrap_or("hosts")));
134    let instance = instance.clone();
135    PyPythonTask::new(async move {
136        let results = try_join_all(tasks).await?;
137
138        let addresses: Result<Vec<ChannelAddr>, anyhow::Error> =
139            monarch_with_gil(GilSite::Bootstrap, |py| {
140                results
141                    .into_iter()
142                    .map(|result| {
143                        let url_str: String = result.bind(py).extract()?;
144                        Ok(ChannelAddr::from_zmq_url(&url_str)?.into_dial_addr())
145                    })
146                    .collect()
147            })
148            .await;
149        let addresses = addresses?;
150
151        let host_mesh = HostMesh::attach(&*instance, name, addresses)
152            .await
153            .map_err(|e| anyhow::anyhow!("attach failed: {}", e))?;
154        Ok(PyHostMesh::new_owned(host_mesh))
155    })
156}
157
158pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
159    let f = wrap_pyfunction!(bootstrap_main, hyperactor_mod)?;
160    f.setattr(
161        "__module__",
162        "monarch._rust_bindings.monarch_hyperactor.bootstrap",
163    )?;
164    hyperactor_mod.add_function(f)?;
165
166    let f = wrap_pyfunction!(run_worker_loop_forever, hyperactor_mod)?;
167    f.setattr(
168        "__module__",
169        "monarch._rust_bindings.monarch_hyperactor.bootstrap",
170    )?;
171    hyperactor_mod.add_function(f)?;
172
173    let f = wrap_pyfunction!(attach_to_workers, hyperactor_mod)?;
174    f.setattr(
175        "__module__",
176        "monarch._rust_bindings.monarch_hyperactor.bootstrap",
177    )?;
178    hyperactor_mod.add_function(f)?;
179
180    Ok(())
181}
monarch_hyperactor/bootstrap.rs

monarch_hyperactor/
bootstrap.rs
