make client this_proc() real (#2144)

Summary:

This builds on the previous change to implement a *real* in-process host and local proc mesh. Thus, nothing is fake anymore in the v1 path: this_proc() behaves the same everywhere, and is managed in precisely the same way.

This also means that we can do things like hand out references to the local host (or proc), and have remove actors spawn (procs, actors) on it!

It will also simplify integration: since the client proc is now a host, it has a single front-end address that multiplexes all its procs (including the local one).

We keep around the old codepaths for the benefit of v0; this can be significantly simplified again once we drop v0 support.

Differential Revision: D89196041

Author: pzhan9

hyperactor_mesh/proptest-regressions/actor_mesh.txt

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 7da0353e3138258986cf2598af0836656a1f7c3399a9ffa18ca93cf983b3e64c # shrinks to extent = Extent { inner: ExtentData { labels: ["d/0"], sizes: [1] } }

hyperactor_mesh/src/bootstrap.rs

@@ -266,7 +266,10 @@ async fn halt<R>() -> R {
     unreachable!()
 }
 
-/// Bootstrap a host in this process, returning a handle to the mesh agent:
+/// Bootstrap a host in this process, returning a handle to the mesh agent.
+///
+/// To obtain the local proc, use `GetLocalProc` on the returned host mesh agent,
+/// then use `GetProc` on the returned proc mesh agent.
 ///
 /// - `addr`: the listening address of the host; this is used to bind the frontend address;
 /// - `command`: optional bootstrap command to spawn procs, otherwise [`BootstrapProcManager::current`];

hyperactor_mesh/src/proc_mesh/mesh_agent.rs

@@ -795,6 +795,26 @@ impl Handler<NewClientInstance> for ProcMeshAgent {
     }
 }
 
+/// A handler to get a clone of the proc managed by this agent.
+/// This is used to obtain the local proc from a host mesh.
+#[derive(Debug, hyperactor::Handler, hyperactor::HandleClient)]
+pub struct GetProc {
+    #[reply]
+    pub proc: PortHandle<Proc>,
+}
+
+#[async_trait]
+impl Handler<GetProc> for ProcMeshAgent {
+    async fn handle(
+        &mut self,
+        _cx: &Context<Self>,
+        GetProc { proc }: GetProc,
+    ) -> anyhow::Result<()> {
+        proc.send(self.proc.clone())?;
+        Ok(())
+    }
+}
+
 /// A mailbox sender that initially queues messages, and then relays them to
 /// an underlying sender once configured.
 #[derive(Clone)]

hyperactor_mesh/src/v1.rs

@@ -28,6 +28,7 @@ pub use host_mesh::HostMeshRef;
 use hyperactor::ActorId;
 use hyperactor::ActorRef;
 use hyperactor::Named;
+use hyperactor::ProcId;
 use hyperactor::host::HostError;
 use hyperactor::mailbox::MailboxSenderError;
 use hyperactor::reference;
@@ -149,6 +150,9 @@ pub enum Error {
     #[error("error spawning controller actor for mesh {0}: {1}")]
     ControllerActorSpawnError(Name, anyhow::Error),
 
+    #[error("proc {0} must be direct-addressable")]
+    RankedProc(ProcId),
+
     #[error("error: {0} does not exist")]
     NotExist(Name),
 
@@ -253,7 +257,7 @@ impl Name {
     }
 
     /// Create a Reserved `Name` with no uuid. Only for use by system actors.
-    pub(crate) fn new_reserved(name: impl Into<String>) -> Result<Self> {
+    pub fn new_reserved(name: impl Into<String>) -> Result<Self> {
         Ok(Self::new_with_uuid(name, None)?)
     }
 

hyperactor_mesh/src/v1/host_mesh.rs

@@ -147,6 +147,18 @@ impl HostRef {
     }
 }
 
+impl TryFrom<ActorRef<HostMeshAgent>> for HostRef {
+    type Error = v1::Error;
+
+    fn try_from(value: ActorRef<HostMeshAgent>) -> Result<Self, v1::Error> {
+        let proc_id = value.actor_id().proc_id();
+        match proc_id.as_direct() {
+            Some((addr, _)) => Ok(HostRef(addr.clone())),
+            None => Err(v1::Error::RankedProc(proc_id.clone())),
+        }
+    }
+}
+
 impl std::fmt::Display for HostRef {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         self.0.fmt(f)
@@ -734,6 +746,20 @@ impl HostMeshRef {
         }
     }
 
+    /// Create a new HostMeshRef from an arbitrary set of host mesh agents.
+    pub fn from_host_agents(name: Name, agents: Vec<ActorRef<HostMeshAgent>>) -> v1::Result<Self> {
+        Ok(Self {
+            name,
+            region: extent!(hosts = agents.len()).into(),
+            ranks: Arc::new(
+                agents
+                    .into_iter()
+                    .map(HostRef::try_from)
+                    .collect::<v1::Result<_>>()?,
+            ),
+        })
+    }
+
     /// Spawn a ProcMesh onto this host mesh. The per_host extent specifies the shape
     /// of the procs to spawn on each host.
     ///

hyperactor_mesh/src/v1/proc_mesh.rs

@@ -106,7 +106,8 @@ pub struct ProcRef {
 }
 
 impl ProcRef {
-    pub(crate) fn new(proc_id: ProcId, create_rank: usize, agent: ActorRef<ProcMeshAgent>) -> Self {
+    /// Create a new proc ref from the provided id, create rank and agent.
+    pub fn new(proc_id: ProcId, create_rank: usize, agent: ActorRef<ProcMeshAgent>) -> Self {
         Self {
             proc_id,
             create_rank,
@@ -713,6 +714,20 @@ impl ProcMeshRef {
         })
     }
 
+    /// Create a singleton ProcMeshRef, given the provided ProcRef and name.
+    /// This is used to support creating local singleton proc meshes to support `this_proc()`
+    /// in python client actors.
+    pub fn new_singleton(name: Name, proc_ref: ProcRef) -> Self {
+        Self {
+            name,
+            region: Extent::unity().into(),
+            ranks: Arc::new(vec![proc_ref]),
+            host_mesh: None,
+            root_region: None,
+            root_comm_actor: None,
+        }
+    }
+
     pub(crate) fn root_comm_actor(&self) -> Option<&ActorRef<CommActor>> {
         self.root_comm_actor.as_ref()
     }

monarch_hyperactor/src/actor.rs

@@ -520,6 +520,8 @@ impl PythonActor {
         })
     }
 
+    /// Bootstrap the root client actor, creating a new proc for it.
+    /// This is the legacy entry point that creates its own proc.
     pub(crate) fn bootstrap_client(py: Python<'_>) -> (&'static Instance<Self>, ActorHandle<Self>) {
         static ROOT_CLIENT_INSTANCE: OnceLock<Instance<PythonActor>> = OnceLock::new();
 
@@ -530,8 +532,21 @@ impl PythonActor {
         )
         .unwrap();
 
+        Self::bootstrap_client_inner(py, client_proc, &ROOT_CLIENT_INSTANCE)
+    }
+
+    /// Bootstrap the client proc, storing the root client instance in given static.
+    /// This is passed in because we require storage, as the instance is shared.
+    /// This can be simplified when we remove v0.
+    pub(crate) fn bootstrap_client_inner(
+        py: Python<'_>,
+        client_proc: Proc,
+        root_client_instance: &'static OnceLock<Instance<PythonActor>>,
+    ) -> (&'static Instance<Self>, ActorHandle<Self>) {
         // Make this proc reachable through the global router, so that we can use the
         // same client in both direct-addressed and ranked-addressed modes.
+        //
+        // DEPRECATE after v0 removal
         router::global().bind(client_proc.proc_id().clone().into(), client_proc.clone());
 
         let actor_mesh_mod = py
@@ -557,7 +572,7 @@ impl PythonActor {
             )
             .expect("root instance create");
 
-        ROOT_CLIENT_INSTANCE
+        root_client_instance
             .set(client)
             .map_err(|_| "already initialized root client instance")
             .unwrap();
@@ -577,7 +592,7 @@ impl PythonActor {
             )
             .expect("initialize root client");
 
-        let instance = ROOT_CLIENT_INSTANCE.get().unwrap();
+        let instance = root_client_instance.get().unwrap();
 
         get_tokio_runtime().spawn(async move {
             let mut signal_rx = signal_rx;
@@ -626,7 +641,7 @@ impl PythonActor {
             instance.proc().handle_supervision_event(event);
         });
 
-        (ROOT_CLIENT_INSTANCE.get().unwrap(), handle)
+        (root_client_instance.get().unwrap(), handle)
     }
 }
 

monarch_hyperactor/src/context.rs

@@ -126,6 +126,17 @@ impl PyContext {
             rank: Extent::unity().point_of_rank(0).unwrap(),
         })
     }
+
+    /// Create a context from an existing instance.
+    /// This is used when the root client was bootstrapped via bootstrap_host()
+    /// instead of the default bootstrap_client().
+    #[staticmethod]
+    fn _from_instance(py: Python<'_>, instance: PyInstance) -> PyResult<PyContext> {
+        Ok(PyContext {
+            instance: instance.into_pyobject(py)?.into(),
+            rank: Extent::unity().point_of_rank(0).unwrap(),
+        })
+    }
 }
 
 impl PyContext {

monarch_hyperactor/src/v1/host_mesh.rs

@@ -9,11 +9,20 @@
 use std::collections::HashMap;
 use std::ops::Deref;
 use std::path::PathBuf;
+use std::sync::OnceLock;
 
+use hyperactor::Instance;
+use hyperactor::Proc;
 use hyperactor_mesh::bootstrap::BootstrapCommand;
+use hyperactor_mesh::bootstrap::host;
+use hyperactor_mesh::proc_mesh::default_transport;
+use hyperactor_mesh::proc_mesh::mesh_agent::GetProcClient;
 use hyperactor_mesh::shared_cell::SharedCell;
+use hyperactor_mesh::v1::ProcMeshRef;
 use hyperactor_mesh::v1::host_mesh::HostMesh;
 use hyperactor_mesh::v1::host_mesh::HostMeshRef;
+use hyperactor_mesh::v1::host_mesh::mesh_agent::GetLocalProcClient;
+use hyperactor_mesh::v1::proc_mesh::ProcRef;
 use ndslice::View;
 use ndslice::view::RankedSliceable;
 use pyo3::IntoPyObjectExt;
@@ -24,6 +33,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use pyo3::types::PyType;
 
+use crate::actor::PythonActor;
 use crate::actor::to_py_error;
 use crate::alloc::PyAlloc;
 use crate::context::PyInstance;
@@ -240,6 +250,76 @@ impl PyHostMeshRefImpl {
     }
 }
 
+/// Static storage for the root client instance when using host-based bootstrap.
+static ROOT_CLIENT_INSTANCE_FOR_HOST: OnceLock<Instance<PythonActor>> = 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.
+///
+/// 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.
+///
+/// This should be called only once, at process initialization
+#[pyfunction]
+fn bootstrap_host(bootstrap_cmd: Option<PyBootstrapCommand>) -> 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()))?,
+    };
+
+    PyPythonTask::new(async move {
+        let host_mesh_agent = host(default_transport().any(), Some(bootstrap_cmd), None)
+            .await
+            .map_err(|e| PyException::new_err(e.to_string()))?;
+
+        let host_mesh_name = hyperactor_mesh::v1::Name::new_reserved("local").unwrap();
+        let host_mesh = HostMeshRef::from_host_agents(host_mesh_name, vec![host_mesh_agent.bind()])
+            .map_err(|e| PyException::new_err(e.to_string()))?;
+
+        // We require a temporary instance to make a call to the host/proc agent.
+        let temp_proc = Proc::local();
+        let (temp_instance, _) = temp_proc
+            .instance("temp")
+            .map_err(|e| PyException::new_err(e.to_string()))?;
+
+        let local_proc_agent = host_mesh_agent
+            .get_local_proc(&temp_instance)
+            .await
+            .map_err(|e| PyException::new_err(e.to_string()))?;
+
+        let proc_mesh_name = hyperactor_mesh::v1::Name::new_reserved("local").unwrap();
+        let proc_mesh = ProcMeshRef::new_singleton(
+            proc_mesh_name,
+            ProcRef::new(
+                local_proc_agent.actor_id().proc_id().clone(),
+                0,
+                local_proc_agent.bind(),
+            ),
+        );
+
+        let local_proc = local_proc_agent
+            .get_proc(&temp_instance)
+            .await
+            .map_err(|e| PyException::new_err(e.to_string()))?;
+
+        let (instance, _handle) = Python::with_gil(|py| {
+            PythonActor::bootstrap_client_inner(py, local_proc, &ROOT_CLIENT_INSTANCE_FOR_HOST)
+        });
+
+        Ok((
+            PyHostMesh::new_ref(host_mesh),
+            PyProcMesh::new_ref(proc_mesh),
+            PyInstance::from(instance),
+        ))
+    })
+}
+
 #[pyfunction]
 fn py_host_mesh_from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<PyHostMesh> {
     let r: PyResult<HostMeshRef> = bincode::deserialize(bytes.as_bytes())
@@ -254,6 +334,14 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
         "monarch._rust_bindings.monarch_hyperactor.v1.host_mesh",
     )?;
     hyperactor_mod.add_function(f)?;
+
+    let f2 = wrap_pyfunction!(bootstrap_host, hyperactor_mod)?;
+    f2.setattr(
+        "__module__",
+        "monarch._rust_bindings.monarch_hyperactor.v1.host_mesh",
+    )?;
+    hyperactor_mod.add_function(f2)?;
+
     hyperactor_mod.add_class::<PyHostMesh>()?;
     hyperactor_mod.add_class::<PyBootstrapCommand>()?;
     Ok(())

python/monarch/_rust_bindings/monarch_hyperactor/v1/host_mesh.pyi

@@ -101,3 +101,15 @@ class BootstrapCommand:
         ...
 
     def __repr__(self) -> str: ...
+
+def bootstrap_host(
+    bootstrap_cmd: BootstrapCommand | None,
+) -> PythonTask[(HostMesh, ProcMesh, Instance)]:
+    """
+    Bootstrap a host mesh in this process, returning the host mesh,
+    proc mesh, and client instance.
+
+    Arguments:
+    - `bootstrap_cmd`: The bootstrap command to use to bootstrap the host.
+    """
+    ...