Additional dims per host when spawning proc mesh on host mesh (#1345)

Summary:
Pull Request resolved: #1345

Add an API to specify an `Extent` for adding additional dims on each host when spawning a proc mesh on a v1 host mesh. The extent of the proc mesh will be the extent of the host mesh + the extent per host.
ghstack-source-id: 312970198
exported-using-ghexport

Reviewed By: mariusae

Differential Revision: D83274881

fbshipit-source-id: 4309309a151c1169d0d4f2f434062512ee0b228a

Author: samlurye

hyperactor_mesh/src/bootstrap.rs

@@ -1709,6 +1709,7 @@ mod tests {
     use hyperactor::context::Mailbox as _;
     use hyperactor::host::ProcHandle;
     use hyperactor::id;
+    use ndslice::Extent;
     use ndslice::ViewExt;
     use ndslice::extent;
     use tokio::process::Command;
@@ -2863,7 +2864,10 @@ mod tests {
         //
         // (4) We collect the per-host procs into a `ProcMesh` and
         // return it.
-        let proc_mesh = host_mesh.spawn(&instance, "p0").await.unwrap();
+        let proc_mesh = host_mesh
+            .spawn(&instance, "p0", Extent::unity())
+            .await
+            .unwrap();
 
         // Note: There is no support for status() in v1.
         // assert!(proc_mesh.status(&instance).await.is_err());

hyperactor_mesh/src/v1/actor_mesh.rs

@@ -341,6 +341,7 @@ mod tests {
     use hyperactor::clock::RealClock;
     use hyperactor::context::Mailbox as _;
     use hyperactor::mailbox;
+    use ndslice::Extent;
     use ndslice::ViewExt;
     use ndslice::extent;
     use ndslice::view::Ranked;
@@ -530,7 +531,10 @@ mod tests {
 
         let instance = testing::instance().await;
         let host_mesh = testing::host_mesh(extent!(host = 4)).await;
-        let proc_mesh = host_mesh.spawn(instance, "test").await.unwrap();
+        let proc_mesh = host_mesh
+            .spawn(instance, "test", Extent::unity())
+            .await
+            .unwrap();
         let actor_mesh = proc_mesh
             .spawn::<testactor::TestActor>(instance, "test", &())
             .await

hyperactor_mesh/src/v1/host_mesh.rs

@@ -9,6 +9,7 @@
 use hyperactor::channel::ChannelTransport;
 pub mod mesh_agent;
 
+use std::collections::HashSet;
 use std::ops::Deref;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -460,31 +461,69 @@ impl HostMeshRef {
         }
     }
 
-    /// Spawn a ProcMesh onto this host mesh.
-    // TODO: add an "additional dims" API
-    pub async fn spawn(&self, cx: &impl context::Actor, name: &str) -> v1::Result<ProcMesh> {
-        let name = Name::new(name);
+    /// Spawn a ProcMesh onto this host mesh. The per_host extent specifies the shape
+    /// of the procs to spawn on each host.
+    pub async fn spawn(
+        &self,
+        cx: &impl context::Actor,
+        name: &str,
+        per_host: Extent,
+    ) -> v1::Result<ProcMesh> {
+        let per_host_labels = per_host.labels().iter().collect::<HashSet<_>>();
+        let host_labels = self.region.labels().iter().collect::<HashSet<_>>();
+        if !per_host_labels
+            .intersection(&host_labels)
+            .collect::<Vec<_>>()
+            .is_empty()
+        {
+            return Err(v1::Error::ConfigurationError(anyhow::anyhow!(
+                "per_host dims overlap with existing dims when spawning proc mesh"
+            )));
+        }
+
+        let labels = self
+            .region
+            .labels()
+            .to_vec()
+            .into_iter()
+            .chain(per_host.labels().to_vec().into_iter())
+            .collect();
+        let sizes = self
+            .region
+            .extent()
+            .sizes()
+            .to_vec()
+            .into_iter()
+            .chain(per_host.sizes().to_vec().into_iter())
+            .collect();
+        let extent =
+            Extent::new(labels, sizes).map_err(|err| v1::Error::ConfigurationError(err.into()))?;
+
+        let mesh_name = Name::new(name);
         let mut procs = Vec::new();
-        for (rank, host) in self.ranks.iter().enumerate() {
-            let _ok = host
-                .mesh_agent()
-                .create_or_update(cx, name.clone(), ())
-                .await
-                .map_err(|e| {
-                    v1::Error::HostMeshAgentConfigurationError(
-                        host.mesh_agent().actor_id().clone(),
-                        format!("failed while creating proc: {}", e),
-                    )
-                })?;
-            procs.push(ProcRef::new(
-                host.named_proc(&name),
-                rank,
-                // TODO: specify or retrieve from state instead, to avoid attestation.
-                ActorRef::attest(host.named_proc(&name).actor_id("agent", 0)),
-            ));
+        for (host_rank, host) in self.ranks.iter().enumerate() {
+            for per_host_rank in 0..per_host.num_ranks() {
+                let proc_name = Name::new(format!("{}-{}", name, per_host_rank));
+                let _ok = host
+                    .mesh_agent()
+                    .create_or_update(cx, proc_name.clone(), ())
+                    .await
+                    .map_err(|e| {
+                        v1::Error::HostMeshAgentConfigurationError(
+                            host.mesh_agent().actor_id().clone(),
+                            format!("failed while creating proc: {}", e),
+                        )
+                    })?;
+                procs.push(ProcRef::new(
+                    host.named_proc(&proc_name),
+                    per_host.num_ranks() * host_rank + per_host_rank,
+                    // TODO: specify or retrieve from state instead, to avoid attestation.
+                    ActorRef::attest(host.named_proc(&proc_name).actor_id("agent", 0)),
+                ));
+            }
         }
 
-        ProcMesh::create_owned_unchecked(cx, name, self.clone(), procs).await
+        ProcMesh::create_owned_unchecked(cx, mesh_name, extent, self.clone(), procs).await
     }
 }
 
@@ -621,12 +660,28 @@ mod tests {
                 .await
                 .unwrap();
 
-            let proc_mesh1 = host_mesh.spawn(instance, "test_1").await.unwrap();
+            let proc_mesh1 = host_mesh
+                .spawn(instance, "test_1", Extent::unity())
+                .await
+                .unwrap();
             let actor_mesh1: ActorMesh<testactor::TestActor> =
                 proc_mesh1.spawn(instance, "test", &()).await.unwrap();
-            let proc_mesh2 = host_mesh.spawn(instance, "test_2").await.unwrap();
+            let proc_mesh2 = host_mesh
+                .spawn(instance, "test_2", extent!(gpus = 3, extra = 2))
+                .await
+                .unwrap();
+            assert_eq!(
+                proc_mesh2.extent(),
+                extent!(replicas = 4, gpus = 3, extra = 2)
+            );
+            assert_eq!(proc_mesh2.values().count(), 24);
             let actor_mesh2: ActorMesh<testactor::TestActor> =
                 proc_mesh2.spawn(instance, "test", &()).await.unwrap();
+            assert_eq!(
+                actor_mesh2.extent(),
+                extent!(replicas = 4, gpus = 3, extra = 2)
+            );
+            assert_eq!(actor_mesh2.values().count(), 24);
 
             // Host meshes can be dereferenced to produce a concrete ref.
             let host_mesh_ref: HostMeshRef = host_mesh.clone();
@@ -637,23 +692,24 @@ mod tests {
             );
 
             // Validate we can cast:
-
-            let (port, mut rx) = instance.mailbox().open_port();
-            actor_mesh1
-                .cast(instance, testactor::GetActorId(port.bind()))
-                .unwrap();
-
-            let mut expected_actor_ids: HashSet<_> = actor_mesh1
-                .values()
-                .map(|actor_ref| actor_ref.actor_id().clone())
-                .collect();
-
-            while !expected_actor_ids.is_empty() {
-                let actor_id = rx.recv().await.unwrap();
-                assert!(
-                    expected_actor_ids.remove(&actor_id),
-                    "got {actor_id}, expect {expected_actor_ids:?}"
-                );
+            for actor_mesh in [&actor_mesh1, &actor_mesh2] {
+                let (port, mut rx) = instance.mailbox().open_port();
+                actor_mesh
+                    .cast(instance, testactor::GetActorId(port.bind()))
+                    .unwrap();
+
+                let mut expected_actor_ids: HashSet<_> = actor_mesh
+                    .values()
+                    .map(|actor_ref| actor_ref.actor_id().clone())
+                    .collect();
+
+                while !expected_actor_ids.is_empty() {
+                    let actor_id = rx.recv().await.unwrap();
+                    assert!(
+                        expected_actor_ids.remove(&actor_id),
+                        "got {actor_id}, expect {expected_actor_ids:?}"
+                    );
+                }
             }
 
             // Now forward a message through all directed edges across the two meshes.
@@ -719,7 +775,10 @@ mod tests {
 
         let instance = testing::instance().await;
         let host_mesh = HostMeshRef::from_hosts(hosts);
-        let proc_mesh = host_mesh.spawn(&instance, "test").await.unwrap();
+        let proc_mesh = host_mesh
+            .spawn(&testing::instance().await, "test", Extent::unity())
+            .await
+            .unwrap();
         let actor_mesh: ActorMesh<testactor::TestActor> = proc_mesh
             .spawn(&testing::instance().await, "test", &())
             .await

hyperactor_mesh/src/v1/proc_mesh.rs

@@ -217,10 +217,10 @@ impl ProcMesh {
     pub(crate) async fn create_owned_unchecked(
         cx: &impl context::Actor,
         name: Name,
+        extent: Extent,
         hosts: HostMeshRef,
         ranks: Vec<ProcRef>,
     ) -> v1::Result<Self> {
-        let extent = hosts.extent();
         Self::create(
             cx,
             name,

monarch_hyperactor/src/shape.rs

@@ -102,6 +102,12 @@ impl From<Extent> for PyExtent {
     }
 }
 
+impl From<PyExtent> for Extent {
+    fn from(py_extent: PyExtent) -> Self {
+        py_extent.inner
+    }
+}
+
 #[derive(Serialize, Deserialize, Clone)]
 #[pyclass(
     name = "Region",

monarch_hyperactor/src/v1/host_mesh.rs

@@ -27,6 +27,7 @@ use crate::alloc::PyAlloc;
 use crate::context::PyInstance;
 use crate::instance_dispatch;
 use crate::pytokio::PyPythonTask;
+use crate::shape::PyExtent;
 use crate::shape::PyRegion;
 use crate::v1::proc_mesh::PyProcMesh;
 
@@ -142,12 +143,18 @@ impl PyHostMesh {
         })
     }
 
-    fn spawn_nonblocking(&self, instance: &PyInstance, name: String) -> PyResult<PyPythonTask> {
+    fn spawn_nonblocking(
+        &self,
+        instance: &PyInstance,
+        name: String,
+        per_host: &PyExtent,
+    ) -> PyResult<PyPythonTask> {
         let host_mesh = self.mesh_ref()?.clone();
         let instance = instance.clone();
+        let per_host = per_host.clone().into();
         let mesh_impl = async move {
             let proc_mesh = instance_dispatch!(instance, async move |cx_instance| {
-                host_mesh.spawn(cx_instance, &name).await
+                host_mesh.spawn(cx_instance, &name, per_host).await
             })
             .map_err(to_py_error)?;
             Ok(PyProcMesh::new_owned(proc_mesh))

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

@@ -12,7 +12,7 @@ from monarch._rust_bindings.monarch_hyperactor.alloc import Alloc
 from monarch._rust_bindings.monarch_hyperactor.context import Instance
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
 
-from monarch._rust_bindings.monarch_hyperactor.shape import Region
+from monarch._rust_bindings.monarch_hyperactor.shape import Extent, Region
 from monarch._rust_bindings.monarch_hyperactor.v1.proc_mesh import ProcMesh
 
 @final
@@ -40,13 +40,15 @@ class HostMesh:
         self,
         instance: Instance,
         name: str,
+        per_host: Extent,
     ) -> PythonTask[ProcMesh]:
         """
         Spawn a new actor on this mesh.
 
         Arguments:
         - `instance`: The instance to use to spawn the mesh.
         - `name`: Name of the proc mesh
+        - `per_host`: Extent describing the shape of the proc mesh on each host.
         """
         ...
 

python/tests/_monarch/test_actor_mesh.py

@@ -25,7 +25,7 @@
     AllocConstraints,
     AllocSpec,
 )
-from monarch._rust_bindings.monarch_hyperactor.shape import Region, Slice
+from monarch._rust_bindings.monarch_hyperactor.shape import Extent, Region, Slice
 from monarch._src.actor.proc_mesh import _get_bootstrap_args, ProcessAllocator
 
 if TYPE_CHECKING:
@@ -256,7 +256,7 @@ async def test_host_mesh() -> None:
     async def run() -> None:
         cmd, args, bootstrap_env = _get_bootstrap_args()
         allocator = ProcessAllocator(cmd, args, bootstrap_env)
-        spec: AllocSpec = AllocSpec(AllocConstraints(), replicas=2, hosts=2, gpus=4)
+        spec: AllocSpec = AllocSpec(AllocConstraints(), hosts=2)
         alloc = allocator.allocate(spec)
 
         host_mesh = await HostMesh.allocate_nonblocking(
@@ -271,34 +271,35 @@ async def run() -> None:
             ),
         ).spawn()
 
-        assert host_mesh.region.labels() == ["replicas", "hosts", "gpus"]
-        assert host_mesh.region.slice() == Slice(
-            offset=0, sizes=[2, 2, 4], strides=[8, 4, 1]
-        )
+        assert host_mesh.region.labels() == ["hosts"]
+        assert host_mesh.region.slice() == Slice(offset=0, sizes=[2], strides=[1])
 
         proc_mesh = await host_mesh.spawn_nonblocking(
-            context().actor_instance._as_rust(), "proc_mesh"
+            context().actor_instance._as_rust(),
+            "proc_mesh",
+            Extent(["gpus", "replicas"], [2, 4]),
         ).spawn()
         actor_mesh = await spawn_actor_mesh(proc_mesh)
 
         await verify_cast_to_call(actor_mesh, context().actor_instance, list(range(16)))
 
-        # Ranks 4, 6, 12, 14 (gpus 0 and 2 on host 1 on both replicas)
         sliced_hm = host_mesh.sliced(
             Region(
-                labels=["replicas", "gpus"],
-                slice=Slice(offset=4, sizes=[2, 2], strides=[8, 2]),
+                labels=["hosts"],
+                slice=Slice(offset=1, sizes=[1], strides=[1]),
             )
         )
 
-        assert sliced_hm.region.labels() == ["replicas", "gpus"]
-        assert sliced_hm.region.slice() == Slice(offset=4, sizes=[2, 2], strides=[8, 2])
+        assert sliced_hm.region.labels() == ["hosts"]
+        assert sliced_hm.region.slice() == Slice(offset=1, sizes=[1], strides=[1])
 
         sliced_pm = await sliced_hm.spawn_nonblocking(
-            context().actor_instance._as_rust(), "sliced_pm"
+            context().actor_instance._as_rust(),
+            "sliced_pm",
+            Extent(["gpus", "replicas"], [2, 3]),
         )
         sliced_am = await spawn_actor_mesh(sliced_pm)
 
-        await verify_cast_to_call(sliced_am, context().actor_instance, list(range(4)))
+        await verify_cast_to_call(sliced_am, context().actor_instance, list(range(6)))
 
     run()