monarch_hyperactor/

host_mesh.rs

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::HashMap;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::OnceLock;
use std::time::Duration;

use hyperactor::ActorHandle;
use hyperactor::Endpoint as _;
use hyperactor::Gateway;
use hyperactor::Instance;
use hyperactor::Proc;
use hyperactor::ProcAddr;
use hyperactor::ProcId;
use hyperactor::channel::ChannelAddr;
use hyperactor::id::Label;
use hyperactor_mesh::ProcMeshRef;
use hyperactor_mesh::bootstrap::BootstrapCommand;
use hyperactor_mesh::bootstrap::ProcBind;
use hyperactor_mesh::bootstrap::host;
use hyperactor_mesh::host_mesh;
use hyperactor_mesh::host_mesh::HostMesh;
use hyperactor_mesh::host_mesh::HostMeshRef;
use hyperactor_mesh::host_mesh::PerRankBootstrapFn;
use hyperactor_mesh::host_mesh::host_agent::GetLocalProcClient;
use hyperactor_mesh::host_mesh::host_agent::HostAgent;
use hyperactor_mesh::host_mesh::host_agent::ShutdownHost;
use hyperactor_mesh::mesh_admin::MeshAdminMessageClient;
use hyperactor_mesh::mesh_id::ActorMeshId;
use hyperactor_mesh::mesh_id::HostMeshId;
use hyperactor_mesh::mesh_id::ProcMeshId;
use hyperactor_mesh::proc_agent::GetProcClient;
use hyperactor_mesh::proc_mesh::ProcRef;
use hyperactor_mesh::proc_mesh::telemetry_actor_mesh_id;
use hyperactor_mesh::shared_cell::SharedCell;
use hyperactor_mesh::transport::default_bind_spec;
use hyperactor_telemetry::hash_to_u64;
use ndslice::View;
use ndslice::view::RankedSliceable;
use pyo3::IntoPyObjectExt;
use pyo3::exceptions::PyException;
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use pyo3::types::PyMapping;

use crate::actor::PythonActor;
use crate::actor::to_py_error;
use crate::context::PyInstance;
use crate::context::stop_instance_and_wait;
use crate::handle::PyHandle;
use crate::proc_mesh::PyProcMesh;
use crate::pytokio::PyPythonTask;
use crate::runtime::GilSite;
use crate::runtime::monarch_with_gil;
use crate::runtime::monarch_with_gil_blocking;
use crate::shape::PyExtent;
use crate::shape::PyPoint;
use crate::shape::PyRegion;

/// The command used to start a proc.
///
/// Arguments:
/// - `program`: The program to execute.
/// - `arg0`: Optionally, the program's arg0.
/// - `args`: Command line arguments.
/// - `env`: Environment variables as a mapping. A `None` value removes
///   the variable from the inherited environment.
/// - `inherit_env`: If true (the default), the proc inherits the
///   environment of the host agent that spawns it, and `env` is applied
///   on top. If false, the proc starts from an empty environment.
///
/// Monarch always sets its own launch variables (e.g.,
/// `HYPERACTOR_MESH_BOOTSTRAP_MODE`) after applying `env`, and some
/// process launchers add their own (e.g., systemd service metadata).
///
/// The `env` property returns a copy: `cmd.env["X"] = "1"` does not
/// change the command. Use `cmd.with_env({"X": "1"})` to merge, or
/// assign `cmd.env = {...}` to replace the whole mapping.
#[pyclass(
    name = "BootstrapCommand",
    module = "monarch._rust_bindings.monarch_hyperactor.host_mesh"
)]
#[derive(Clone)]
pub struct PyBootstrapCommand {
    #[pyo3(get, set)]
    pub program: String,
    #[pyo3(get, set)]
    pub arg0: Option<String>,
    #[pyo3(get, set)]
    pub args: Vec<String>,
    /// Environment overrides; `None` removes the variable from the
    /// inherited environment.
    pub env: HashMap<String, Option<String>>,
    #[pyo3(get, set)]
    pub inherit_env: bool,
}

/// Extract environment overrides from any Python mapping of `str` to
/// `str | None`.
fn extract_env(env: &Bound<'_, PyAny>) -> PyResult<HashMap<String, Option<String>>> {
    env.cast::<PyMapping>()?
        .items()?
        .iter()
        .map(|item| item.extract())
        .collect()
}

#[pymethods]
impl PyBootstrapCommand {
    #[new]
    #[pyo3(signature = (program, arg0, args, env, inherit_env = true))]
    fn new(
        program: String,
        arg0: Option<String>,
        args: Vec<String>,
        env: &Bound<'_, PyAny>,
        inherit_env: bool,
    ) -> PyResult<Self> {
        Ok(Self {
            program,
            arg0,
            args,
            env: extract_env(env)?,
            inherit_env,
        })
    }

    #[getter]
    fn env(&self) -> HashMap<String, Option<String>> {
        self.env.clone()
    }

    #[setter]
    fn set_env(&mut self, env: &Bound<'_, PyAny>) -> PyResult<()> {
        self.env = extract_env(env)?;
        Ok(())
    }

    fn __repr__(&self) -> String {
        format!(
            "BootstrapCommand(program='{}', args={:?}, env={:?}, inherit_env={})",
            self.program,
            self.args,
            self.env,
            if self.inherit_env { "True" } else { "False" }
        )
    }

    /// Return a copy of this command with `env` merged on top of its
    /// environment. Keys in `env` override any conflicting keys in the
    /// existing environment; a `None` value removes the variable.
    fn with_env(&self, env: &Bound<'_, PyAny>) -> PyResult<Self> {
        let mut new_env = self.env.clone();
        new_env.extend(extract_env(env)?);
        Ok(Self {
            env: new_env,
            ..self.clone()
        })
    }
}

impl PyBootstrapCommand {
    pub fn to_rust(&self) -> BootstrapCommand {
        BootstrapCommand {
            program: PathBuf::from(&self.program),
            arg0: self.arg0.clone(),
            args: self.args.clone(),
            env: self
                .env
                .iter()
                .filter_map(|(k, v)| Some((k.clone(), v.clone()?)))
                .collect(),
            env_remove: self
                .env
                .iter()
                .filter(|(_, v)| v.is_none())
                .map(|(k, _)| k.clone())
                .collect(),
            inherit_env: self.inherit_env,
        }
    }
}

#[pyclass(
    name = "HostMesh",
    module = "monarch._rust_bindings.monarch_hyperactor.host_mesh"
)]
#[expect(
    clippy::large_enum_variant,
    reason = "PyO3 #[pyclass] enum; Box wrapping interacts with PyO3 codegen and Python interop — separate diff"
)]
pub(crate) enum PyHostMesh {
    Owned(PyHostMeshImpl),
    Ref(PyHostMeshRefImpl),
}

impl PyHostMesh {
    pub(crate) fn new_owned(inner: HostMesh) -> Self {
        Self::Owned(PyHostMeshImpl(SharedCell::from(inner)))
    }

    pub(crate) fn new_ref(inner: HostMeshRef) -> Self {
        Self::Ref(PyHostMeshRefImpl(inner))
    }

    pub(crate) fn mesh_ref(&self) -> Result<HostMeshRef, anyhow::Error> {
        match self {
            PyHostMesh::Owned(inner) => Ok(inner.0.borrow()?.clone()),
            PyHostMesh::Ref(inner) => Ok(inner.0.clone()),
        }
    }
}

#[pymethods]
impl PyHostMesh {
    #[pyo3(signature = (instance, name, per_host, proc_bind = None, per_rank_bootstrap = None))]
    fn spawn_nonblocking(
        &self,
        _py: Python<'_>,
        instance: &PyInstance,
        name: String,
        per_host: &PyExtent,
        proc_bind: Option<Vec<HashMap<String, String>>>,
        per_rank_bootstrap: Option<Py<PyAny>>,
    ) -> PyResult<PyPythonTask> {
        let host_mesh = self.mesh_ref()?.clone();
        let per_rank_bootstrap: Option<Box<PerRankBootstrapFn>> = per_rank_bootstrap
            .map(|callable| -> PyResult<Box<PerRankBootstrapFn>> {
                Ok(Box::new(move |point| {
                    monarch_with_gil_blocking(GilSite::Bootstrap, |py| {
                        let result =
                            callable
                                .bind(py)
                                .call1((PyPoint::from(point),))
                                .map_err(|e| {
                                    anyhow::anyhow!("per-rank bootstrap callable raised: {}", e)
                                })?;
                        let cmd: PyBootstrapCommand = result.extract().map_err(|e| {
                            anyhow::anyhow!(
                                "per-rank bootstrap callable did not return BootstrapCommand: {}",
                                e
                            )
                        })?;
                        Ok(cmd.to_rust())
                    })
                }))
            })
            .transpose()?;
        let instance = instance.clone();
        let per_host = per_host.clone().into();
        let proc_bind = proc_bind.map(|v| v.into_iter().map(ProcBind::from).collect());
        let mesh_impl = async move {
            let proc_mesh = host_mesh
                .spawn(
                    instance.deref(),
                    &name,
                    per_host,
                    proc_bind,
                    per_rank_bootstrap,
                )
                .await
                .map_err(to_py_error)?;
            Ok(PyProcMesh::new_owned(proc_mesh))
        };
        PyPythonTask::new(mesh_impl)
    }

    fn with_bootstrap(&self, bootstrap_command: &PyBootstrapCommand) -> PyResult<Self> {
        match self {
            PyHostMesh::Owned(inner) => {
                let cmd = bootstrap_command.to_rust();
                inner
                    .0
                    .try_with_mut(|mesh| mesh.set_bootstrap(cmd))
                    .map_err(|e| PyException::new_err(e.to_string()))?;
                Ok(Self::Owned(inner.clone()))
            }
            PyHostMesh::Ref(_) => Ok(Self::new_ref(
                self.mesh_ref()?.with_bootstrap(bootstrap_command.to_rust()),
            )),
        }
    }

    fn sliced(&self, region: &PyRegion) -> PyResult<Self> {
        Ok(Self::new_ref(
            self.mesh_ref()?.sliced(region.as_inner().clone()),
        ))
    }

    #[getter]
    fn region(&self) -> PyResult<PyRegion> {
        Ok(PyRegion::from(self.mesh_ref()?.region()))
    }

    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
        let mesh_ref = self.mesh_ref()?;
        if crate::pickle::push_mesh_reference_if_active(crate::actor::MeshRef::Host(Box::new(
            mesh_ref.clone(),
        ))) {
            let pop_fn = PyModule::import(py, "monarch._rust_bindings.monarch_hyperactor.pickle")?
                .getattr("pop_mesh_reference")?;
            return Ok((pop_fn, pyo3::types::PyTuple::empty(py).into_any()));
        }
        let bytes = bincode::serde::encode_to_vec(&mesh_ref, bincode::config::legacy())
            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?;
        let py_bytes = (PyBytes::new(py, &bytes),).into_bound_py_any(py).unwrap();
        let from_bytes =
            PyModule::import(py, "monarch._rust_bindings.monarch_hyperactor.host_mesh")?
                .getattr("py_host_mesh_from_bytes")?;
        Ok((from_bytes, py_bytes))
    }

    fn __eq__(&self, other: &PyHostMesh) -> PyResult<bool> {
        Ok(self.mesh_ref()? == other.mesh_ref()?)
    }

    fn shutdown(&self, instance: &PyInstance) -> PyResult<PyHandle> {
        match self {
            PyHostMesh::Owned(inner) => {
                let instance = instance.clone();
                let mesh_borrow = inner.0.clone();
                let fut = async move {
                    match mesh_borrow.take().await {
                        Ok(mut mesh) => {
                            mesh.shutdown(instance.deref()).await?;
                            Ok(())
                        }
                        Err(_) => {
                            // Don't return an exception, silently ignore the stop request
                            // because it was already done.
                            tracing::info!("shutdown was already called on host mesh");
                            Ok(())
                        }
                    }
                };
                Ok(PyHandle::spawn(fut))
            }
            PyHostMesh::Ref(_) => Err(PyRuntimeError::new_err(
                "cannot shut down `HostMesh` that is a reference instead of owned",
            )),
        }
    }

    fn stop(&self, instance: &PyInstance) -> PyResult<PyHandle> {
        match self {
            PyHostMesh::Owned(inner) => {
                let instance = instance.clone();
                let mesh_borrow = inner.0.clone();
                let fut = async move {
                    match mesh_borrow.take().await {
                        Ok(mut mesh) => {
                            mesh.stop(instance.deref()).await?;
                            Ok(())
                        }
                        Err(_) => {
                            tracing::info!("stop was already called on host mesh");
                            Ok(())
                        }
                    }
                };
                Ok(PyHandle::spawn(fut))
            }
            PyHostMesh::Ref(_) => Err(PyRuntimeError::new_err(
                "cannot stop `HostMesh` that is a reference instead of owned",
            )),
        }
    }
}

#[derive(Clone)]
#[pyclass(
    name = "HostMeshImpl",
    module = "monarch._rust_bindings.monarch_hyperactor.host_mesh"
)]
pub(crate) struct PyHostMeshImpl(SharedCell<HostMesh>);

#[derive(Debug, Clone)]
#[pyclass(
    name = "HostMeshRefImpl",
    module = "monarch._rust_bindings.monarch_hyperactor.host_mesh"
)]
pub(crate) struct PyHostMeshRefImpl(HostMeshRef);

impl PyHostMeshRefImpl {
    fn __repr__(&self) -> PyResult<String> {
        Ok(format!("<HostMeshRefImpl {:?}>", self.0))
    }
}

/// Static storage for the root client instance when using host-based bootstrap.
static ROOT_CLIENT_INSTANCE_FOR_HOST: OnceLock<Instance<PythonActor>> = OnceLock::new();

/// Static storage for the host mesh agent created by bootstrap_host().
static HOST_MESH_AGENT_FOR_HOST: OnceLock<ActorHandle<HostAgent>> = OnceLock::new();

/// Static storage for the host shutdown handle created by bootstrap_host().
/// Used during shutdown_context to join the mailbox server and flush
/// receive-side acks.
static HOST_SHUTDOWN_HANDLE: OnceLock<
    tokio::sync::Mutex<Option<hyperactor_mesh::bootstrap::HostShutdownHandle>>,
> = OnceLock::new();

/// Bootstrap the client host and root client actor.
///
/// This creates a proper Host with BootstrapProcManager, spawns the root client
/// actor on the Host's local_proc.
///
/// The root client actor is the program's singleton root orchestration actor;
/// it is distinct from the root ProcAgent on the same proc. Before creating the
/// actor, bootstrap binds the restricted client-root service API on that
/// ProcAgent and places the resulting `ClientRootRef` in the actor's persistent
/// environment. Descendants inherit the capability, giving Python and non-Python
/// actors the same route to services owned by the program root. Worker
/// ProcAgents do not bind this API.
///
/// Returns a tuple of (HostMesh, ProcMesh, PyInstance) where:
/// - PyHostMesh: the bootstrapped (local) host mesh; and
/// - PyProcMesh: the local ProcMesh on this HostMesh; and
/// - PyInstance: the root client actor instance, on the ProcMesh.
///
/// The HostMesh is served on the default transport unless ``bind_addr`` is set.
///
/// If ``via`` is set to a ZMQ-style address of a remote host's duplex
/// server, the local host's gateway is attached to that remote
/// gateway: outbound traffic to unknown destinations is forwarded over
/// the duplex, and inbound traffic from the duplex is delivered to
/// local procs by the gateway's routing. The local host still owns
/// its own frontend address; ``via`` only adds a forwarding path. The
/// returned ``PyHostMesh`` / ``PyProcMesh`` / ``PyInstance`` all live
/// on the local host's procs.
///
/// If ``bind_addr`` is set, the local host frontend binds that explicit
/// address instead of the process-wide default transport. This is independent
/// of ``via`` and does not alter the runtime configuration inherited by other
/// hosts.
///
/// This should be called only once, at process initialization.
#[pyfunction]
#[pyo3(signature = (bootstrap_cmd, via=None, bind_addr=None))]
fn bootstrap_host(
    bootstrap_cmd: Option<PyBootstrapCommand>,
    via: Option<&str>,
    bind_addr: Option<&str>,
) -> PyResult<PyPythonTask> {
    let bootstrap_cmd = match bootstrap_cmd {
        Some(cmd) => cmd.to_rust(),
        None => BootstrapCommand::current().map_err(|e| PyException::new_err(e.to_string()))?,
    };
    let via_addr = via
        .map(|s| {
            ChannelAddr::from_zmq_url(s)
                .map_err(|e| PyValueError::new_err(format!("via address: {}", e)))
        })
        .transpose()?;
    let bind_addr = bind_addr
        .map(|s| {
            ChannelAddr::from_zmq_url(s)
                .map_err(|e| PyValueError::new_err(format!("bind address: {}", e)))
        })
        .transpose()?
        .unwrap_or_else(|| default_bind_spec().binding_addr());

    PyPythonTask::new(async move {
        // Pass the via address into `host` so the gateway's `serve_via`
        // attach runs *after* the host's own frontend serve but *before*
        // any proc or actor ref is minted. That ordering makes the via
        // session the newest active serve, so refs advertise the
        // routable `Via` location and out-of-cluster return traffic
        // flows back over the duplex. The host owns the resulting serve
        // handle and tears it down on drop.
        let gateway = Gateway::global().clone();
        let service_proc_addr = ProcAddr::new(
            ProcId::singleton(Label::strip(hyperactor::proc::LEGACY_SERVICE_PROC_NAME)),
            bind_addr.into(),
        );

        let (host_mesh_agent, shutdown_handle) = host(
            service_proc_addr,
            Some(bootstrap_cmd),
            None,
            false,
            None,
            gateway,
            via_addr,
        )
        .await
        .map_err(|e| PyException::new_err(e.to_string()))?;

        // Store the agent and shutdown handle for later shutdown
        HOST_MESH_AGENT_FOR_HOST.set(host_mesh_agent.clone()).ok();
        HOST_SHUTDOWN_HANDLE.get_or_init(|| tokio::sync::Mutex::new(Some(shutdown_handle)));

        let host_mesh_id = HostMeshId::singleton(Label::new("local").unwrap());
        let host_mesh = HostMeshRef::from_host_agent(host_mesh_id, host_mesh_agent.bind())
            .map_err(|e| PyException::new_err(e.to_string()))?;

        // Register C so MeshAdminAgent can discover it ("A/C
        // invariant" - hyperactor_mesh/src/mesh_admin.rs).
        hyperactor_mesh::global_context::register_client_host(host_mesh.clone());

        // We require a temporary instance to make a call to the host/proc agent.
        let temp_proc = Proc::isolated();
        let temp_instance = temp_proc.client("temp");

        let local_proc_agent: hyperactor::ActorHandle<hyperactor_mesh::proc_agent::ProcAgent> =
            host_mesh_agent
                .get_local_proc(&temp_instance)
                .await
                .map_err(|e| PyException::new_err(e.to_string()))?;

        let proc_mesh = ProcMeshRef::new_singleton(
            ProcMeshId::singleton(Label::new("local").unwrap()),
            ProcRef::new(
                local_proc_agent.actor_addr().proc_addr(),
                0,
                local_proc_agent.bind(),
            ),
        )
        .map_err(|e| PyException::new_err(e.to_string()))?;

        let local_proc = local_proc_agent
            .get_proc(&temp_instance)
            .await
            .map_err(|e| PyException::new_err(e.to_string()))?;

        // Bind the client-root API on this program's one root ProcAgent and seed
        // the Python root client's environment with the resulting capability, so
        // descendants inherit it exactly as the Rust root's descendants do.
        // Worker ProcAgents never bind the API.
        let client_root = hyperactor_mesh::client_root::ClientRootRef::bind(&local_proc_agent);
        let mut root_env = hyperactor::ActorEnvironment::default();
        root_env
            .set(hyperactor_mesh::client_root::CLIENT_ROOT, client_root)
            .map_err(|e| PyException::new_err(e.to_string()))?;

        let (instance, _handle) = monarch_with_gil(GilSite::Bootstrap, |py| {
            PythonActor::bootstrap_client_inner(
                py,
                local_proc,
                root_env,
                &ROOT_CLIENT_INSTANCE_FOR_HOST,
            )
        })
        .await;

        // Register the client proc mesh for external-runtime discovery
        // (`python_client_root`). `bootstrap_client_inner` has just set
        // `ROOT_CLIENT_INSTANCE_FOR_HOST`, so registering here keeps the pair
        // complete: the root client instance and the client proc are set
        // together, before `proc_mesh` is moved into `PyProcMesh` below.
        hyperactor_mesh::global_context::register_client_proc(proc_mesh.clone());

        // Notify telemetry of the bootstrap host mesh, proc mesh, and client actor.
        {
            let now = std::time::SystemTime::now();

            let host_name_str = host_mesh.id().to_string();
            let host_mesh_id = hash_to_u64(host_mesh.id());
            hyperactor_telemetry::notify_mesh_created(hyperactor_telemetry::MeshEvent {
                id: host_mesh_id,
                timestamp: now,
                class: "Host".to_string(),
                given_name: host_mesh
                    .id()
                    .display_label()
                    .map(|l| l.as_str())
                    .unwrap_or("unnamed")
                    .to_string(),
                full_name: host_name_str,
                shape_json: serde_json::to_string(&host_mesh.region().extent()).unwrap_or_default(),
                parent_mesh_id: None,
                parent_view_json: None,
            });

            let host_agent_addr = host_mesh_agent.actor_addr();
            hyperactor_telemetry::notify_actor_created(hyperactor_telemetry::ActorEvent {
                id: hyperactor_telemetry::hash_to_u64(host_agent_addr.id()),
                timestamp: now,
                mesh_id: host_mesh_id,
                rank: 0,
                full_name: host_agent_addr.to_string(),
                display_name: None,
            });

            let proc_id_str = proc_mesh.id().to_string();
            let proc_mesh_id = hash_to_u64(proc_mesh.id());
            hyperactor_telemetry::notify_mesh_created(hyperactor_telemetry::MeshEvent {
                id: proc_mesh_id,
                timestamp: now,
                class: "Proc".to_string(),
                given_name: proc_mesh
                    .id()
                    .display_label()
                    .map(|l| l.as_str())
                    .unwrap_or("unnamed")
                    .to_string(),
                full_name: proc_id_str,
                shape_json: serde_json::to_string(&proc_mesh.region().extent()).unwrap_or_default(),
                parent_mesh_id: Some(host_mesh_id),
                parent_view_json: None,
            });

            let proc_agent_addr = local_proc_agent.actor_addr();
            hyperactor_telemetry::notify_actor_created(hyperactor_telemetry::ActorEvent {
                id: hyperactor_telemetry::hash_to_u64(proc_agent_addr.id()),
                timestamp: now,
                mesh_id: proc_mesh_id,
                rank: 0,
                full_name: proc_agent_addr.to_string(),
                display_name: None,
            });

            let client_mesh_actor_id = ActorMeshId::singleton(Label::new("client").unwrap());
            let client_mesh_name = format!("{}/client", proc_mesh.id());
            let client_mesh_id = telemetry_actor_mesh_id(proc_mesh.id(), &client_mesh_actor_id);
            hyperactor_telemetry::notify_mesh_created(hyperactor_telemetry::MeshEvent {
                id: client_mesh_id,
                timestamp: now,
                class: <PythonActor as typeuri::Named>::typename().to_string(),
                given_name: "client".to_string(),
                full_name: client_mesh_name,
                shape_json: serde_json::to_string(&proc_mesh.region().extent()).unwrap_or_default(),
                parent_mesh_id: Some(proc_mesh_id),
                parent_view_json: None,
            });

            hyperactor_telemetry::notify_actor_created(hyperactor_telemetry::ActorEvent {
                id: hyperactor_telemetry::hash_to_u64(instance.self_addr().id()),
                timestamp: now,
                mesh_id: client_mesh_id,
                rank: 0,
                full_name: instance.self_addr().to_string(),
                display_name: Some("<root>".to_string()),
            });
        }

        Ok((
            PyHostMesh::new_ref(host_mesh),
            PyProcMesh::new_ref(proc_mesh),
            PyInstance::from(instance),
        ))
    })
}

/// The Python-bootstrapped client root: the root client instance plus the
/// client proc mesh registered by `bootstrap_host`. Discovered by
/// [`python_client_root`] for callers that need the stable Python root.
pub struct PythonClientRoot {
    cx: &'static Instance<PythonActor>,
    proc_mesh: ProcMeshRef,
}

impl PythonClientRoot {
    /// Spawn a controllerless owner actor on the client root's proc mesh,
    /// using the root's own stable client instance as the caller context, and
    /// return its typed [`ActorRef<A>`](hyperactor::ActorRef).
    ///
    /// The Python client-root proc mesh is a singleton (one proc), so the
    /// spawned mesh has exactly one rank. This validates that invariant and
    /// returns rank 0's actor reference, encoding the single-owner contract in
    /// the return type and letting callers use the generated `RefClient`
    /// methods. The intermediate `ActorMesh` is dropped once the ref is
    /// extracted; its `Drop` is inert, so this does not tear the owner down.
    ///
    /// Method-only by design: the owner must be spawned with the long-lived
    /// root `cx`, never an arbitrary (possibly short-lived) caller context.
    pub async fn spawn_controllerless_owner<A: hyperactor::RemoteSpawn>(
        &self,
        name: &str,
        params: &A::Params,
    ) -> hyperactor_mesh::Result<hyperactor::ActorRef<A>>
    where
        A::Params: hyperactor::RemoteMessage,
    {
        let mesh = self
            .proc_mesh
            .spawn_controllerless_service(self.cx, name, params)
            .await?;
        // `Ranked` is not imported at module scope: it would make the existing
        // `.region()` calls in this file ambiguous with `View::region`.
        let actual = ndslice::view::Ranked::region(&*mesh).num_ranks();
        if actual != 1 {
            return Err(hyperactor_mesh::Error::InvalidRankCardinality {
                expected: 1,
                actual,
            });
        }
        Ok(ndslice::view::Ranked::get(&*mesh, 0)
            .expect("singleton owner mesh must contain rank 0")
            .clone())
    }
}

/// Discover the Python-bootstrapped client root, if present.
///
/// Python-root discovery only: reads `ROOT_CLIENT_INSTANCE_FOR_HOST` and
/// [`try_registered_client_proc`](hyperactor_mesh::global_context::try_registered_client_proc)
/// — never `context()` / `this_proc()` / `try_this_host()`, and it does not
/// consult the Rust `GLOBAL_CONTEXT`.
///
/// Returns `Ok(None)` when neither slot is registered, `Ok(Some(_))` when both
/// are, and `Err` when exactly one is (partial registration — the pair is set
/// together by `bootstrap_host`).
pub fn python_client_root() -> Result<Option<PythonClientRoot>, anyhow::Error> {
    match (
        ROOT_CLIENT_INSTANCE_FOR_HOST.get(),
        hyperactor_mesh::global_context::try_registered_client_proc(),
    ) {
        (Some(cx), Some(proc_mesh)) => Ok(Some(PythonClientRoot {
            cx,
            proc_mesh: proc_mesh.clone(),
        })),
        (None, None) => Ok(None),
        (Some(_), None) => Err(anyhow::anyhow!(
            "partial Python client root registration: root client instance is registered but client proc mesh is missing"
        )),
        (None, Some(_)) => Err(anyhow::anyhow!(
            "partial Python client root registration: client proc mesh is registered but root client instance is missing"
        )),
    }
}

#[pyfunction]
fn py_host_mesh_from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<PyHostMesh> {
    let r: PyResult<HostMeshRef> =
        bincode::serde::decode_from_slice(bytes.as_bytes(), bincode::config::legacy())
            .map(|(v, _)| v)
            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()));
    r.map(PyHostMesh::new_ref)
}

#[pyfunction]
#[pyo3(signature = (timeout_secs=None))]
fn shutdown_local_host_mesh(timeout_secs: Option<f64>) -> PyResult<PyHandle> {
    let timeout = timeout_secs
        .map(Duration::try_from_secs_f64)
        .transpose()
        .map_err(|e| PyValueError::new_err(format!("timeout_secs: {}", e)))?
        .unwrap_or(Duration::from_secs(10));
    let agent = HOST_MESH_AGENT_FOR_HOST
        .get()
        .ok_or_else(|| PyRuntimeError::new_err("No local host mesh to shutdown"))?
        .clone();

    Ok(PyHandle::spawn(async move {
        // Create a temporary instance to send the shutdown message
        let temp_proc = hyperactor::Proc::isolated();
        let instance = temp_proc.client("shutdown_requester");

        tracing::info!(
            "sending shutdown_host request to agent {}",
            agent.actor_addr()
        );
        // Use same defaults as HostMesh::shutdown():
        // - MESH_TERMINATE_TIMEOUT = 10 seconds
        // - MESH_TERMINATE_CONCURRENCY = 16

        let (port, _) = instance.open_port::<usize>();
        let mut port = port.bind();
        // We don't need the ack, and this temporary proc doesn't have a mailbox
        // receiver set up anyways. Just ignore the message.
        port.return_undeliverable(false);
        agent.post(
            &instance,
            ShutdownHost {
                timeout,
                max_in_flight: 16,
                rank: hyperactor_mesh::resource::Rank::new(0),
                ack: port,
            },
        );

        // Like the host's own flush, the root client's final flush must not
        // outlive the shutdown timeout when the peer is already gone.
        let flush_timeout =
            hyperactor_config::global::get(hyperactor::config::FORWARDER_FLUSH_TIMEOUT)
                .min(timeout);
        let shutdown = if let Some(lock) = HOST_SHUTDOWN_HANDLE.get() {
            lock.lock().await.take()
        } else {
            None
        };
        if let Some(shutdown) = shutdown {
            if let Some(root) = ROOT_CLIENT_INSTANCE_FOR_HOST.get() {
                shutdown
                    .stop_and_join_after_drain(stop_instance_and_wait(
                        root,
                        "shutdown".to_string(),
                        flush_timeout,
                    ))
                    .await;
            } else {
                tracing::warn!("shutting down partially bootstrapped host without root client");
                shutdown.stop_and_join().await;
            }
        } else if let Some(root) = ROOT_CLIENT_INSTANCE_FOR_HOST.get() {
            stop_instance_and_wait(root, "shutdown".to_string(), flush_timeout).await;
        }

        // Python `None`, not unit: PyO3 converts `()` to an empty tuple.
        Ok(None::<()>)
    }))
}

/// Opaque capability token for `ActorRef<MeshAdminAgent>` across the
/// Python boundary. No methods, no getters — Python never inspects
/// this. It exists solely to transport the typed ref from
/// `_spawn_admin` to `_start_periodic_snapshots_http`.
#[pyclass(
    name = "PyMeshAdminRef",
    module = "monarch._rust_bindings.monarch_hyperactor.host_mesh"
)]
#[derive(Clone)]
pub struct PyMeshAdminRef(hyperactor::ActorRef<hyperactor_mesh::mesh_admin::MeshAdminAgent>);

impl PyMeshAdminRef {
    pub fn actor_ref(&self) -> hyperactor::ActorRef<hyperactor_mesh::mesh_admin::MeshAdminAgent> {
        self.0.clone()
    }
}

/// `MESH_ADMIN_ADDR` config.
///
/// Python-facing wrapper around
/// [`hyperactor_mesh::host_mesh::spawn_admin`].
#[pyfunction]
fn _spawn_admin(
    host_meshes: Vec<PyRef<'_, PyHostMesh>>,
    instance: &PyInstance,
    admin_addr: Option<String>,
    telemetry_url: Option<String>,
) -> PyResult<PyPythonTask> {
    if host_meshes.is_empty() {
        return Err(PyException::new_err("at least one mesh is required"));
    }

    let admin_addr = admin_addr
        .map(|s| {
            s.parse::<std::net::SocketAddr>()
                .map_err(|e| PyException::new_err(format!("invalid admin_addr '{}': {}", s, e)))
        })
        .transpose()?;

    let mesh_refs = host_meshes
        .iter()
        .map(|m| -> PyResult<HostMeshRef> { Ok(m.mesh_ref()?.clone()) })
        .collect::<PyResult<Vec<HostMeshRef>>>()?;

    let instance = instance.clone();
    PyPythonTask::new(async move {
        let admin_ref =
            host_mesh::spawn_admin(&mesh_refs, instance.deref(), admin_addr, telemetry_url)
                .await
                .map_err(|e| PyException::new_err(e.to_string()))?;
        let admin_url = admin_ref
            .get_admin_addr(instance.deref())
            .await
            .map_err(|e| PyException::new_err(e.to_string()))?
            .addr
            .ok_or_else(|| PyException::new_err("mesh admin agent did not report an address"))?;
        Ok((admin_url, PyMeshAdminRef(admin_ref)))
    })
}

pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
    let f = wrap_pyfunction!(py_host_mesh_from_bytes, hyperactor_mod)?;
    f.setattr(
        "__module__",
        "monarch._rust_bindings.monarch_hyperactor.host_mesh",
    )?;
    hyperactor_mod.add_function(f)?;

    let f2 = wrap_pyfunction!(bootstrap_host, hyperactor_mod)?;
    f2.setattr(
        "__module__",
        "monarch._rust_bindings.monarch_hyperactor.host_mesh",
    )?;
    hyperactor_mod.add_function(f2)?;

    let f3 = wrap_pyfunction!(shutdown_local_host_mesh, hyperactor_mod)?;
    f3.setattr(
        "__module__",
        "monarch._rust_bindings.monarch_hyperactor.host_mesh",
    )?;
    hyperactor_mod.add_function(f3)?;

    let f4 = wrap_pyfunction!(_spawn_admin, hyperactor_mod)?;
    f4.setattr(
        "__module__",
        "monarch._rust_bindings.monarch_hyperactor.host_mesh",
    )?;
    hyperactor_mod.add_function(f4)?;

    hyperactor_mod.add_class::<PyHostMesh>()?;
    hyperactor_mod.add_class::<PyBootstrapCommand>()?;
    hyperactor_mod.add_class::<PyMeshAdminRef>()?;
    Ok(())
}