Fix this_host for v1 to only return a single host instead of the whole host mesh (meta-pytorch#1459)

Summary:
Pull Request resolved: meta-pytorch#1459

As titled.
ghstack-source-id: 314719294
exported-using-ghexport

Reviewed By: zdevito

Differential Revision: D84091757

fbshipit-source-id: d487b2debde10a9b7827743082c91fe875418f69

Author: samlurye

monarch_hyperactor/src/ndslice.rs

@@ -201,6 +201,12 @@ impl PySlice {
     fn new_row_major(sizes: Vec<usize>) -> PySlice {
         ndslice::Slice::new_row_major(sizes).into()
     }
+
+    fn get(&self, index: usize) -> PyResult<usize> {
+        self.inner
+            .get(index)
+            .map_err(|err| PyValueError::new_err(err.to_string()))
+    }
 }
 
 impl From<&PySlice> for ndslice::Slice {

monarch_hyperactor/src/shape.rs

@@ -102,6 +102,14 @@ impl PyExtent {
             inner: self.inner.region(),
         }
     }
+
+    fn __eq__(&self, other: &Bound<'_, PyAny>) -> PyResult<bool> {
+        if let Ok(other) = other.extract::<PyExtent>() {
+            Ok(self.inner == other.inner)
+        } else {
+            Ok(false)
+        }
+    }
 }
 
 impl From<Extent> for PyExtent {
@@ -173,6 +181,21 @@ impl PyRegion {
             .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?
             .into())
     }
+
+    fn point_of_base_rank(&self, rank: usize) -> PyResult<PyPoint> {
+        self.inner
+            .point_of_base_rank(rank)
+            .map_pyerr()
+            .map(PyPoint::from)
+    }
+
+    fn __eq__(&self, other: &Bound<'_, PyAny>) -> PyResult<bool> {
+        if let Ok(other) = other.extract::<PyRegion>() {
+            Ok(self.inner == other.inner)
+        } else {
+            Ok(false)
+        }
+    }
 }
 
 impl From<Region> for PyRegion {

monarch_hyperactor/src/v1/host_mesh.rs

@@ -182,21 +182,25 @@ impl PyHostMesh {
                 .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()?)
+    }
 }
 
 #[derive(Clone)]
 #[pyclass(
     name = "HostMeshImpl",
     module = "monarch._rust_bindings.monarch_hyperactor.v1.host_mesh"
 )]
-struct PyHostMeshImpl(SharedCell<HostMesh>);
+pub(crate) struct PyHostMeshImpl(SharedCell<HostMesh>);
 
 #[derive(Debug, Clone)]
 #[pyclass(
     name = "HostMeshRefImpl",
     module = "monarch._rust_bindings.monarch_hyperactor.v1.host_mesh"
 )]
-struct PyHostMeshRefImpl(HostMeshRef);
+pub(crate) struct PyHostMeshRefImpl(HostMeshRef);
 
 impl PyHostMeshRefImpl {
     fn __repr__(&self) -> PyResult<String> {

python/monarch/_rust_bindings/monarch_hyperactor/shape.pyi

@@ -97,6 +97,19 @@ class Slice:
         ...
 
     def __repr__(self) -> str: ...
+    def get(self, index: int) -> int:
+        """
+        Given a logical index in row-major order, compute the physical
+        memory offset according to the slice layout. Inverse of `index`.
+        """
+        ...
+
+    def index(self, value: int) -> int:
+        """
+        Given a physical memory offset, compute the logical index in
+        row-major order. Inverse of `get`.
+        """
+        ...
 
 @final
 class Shape:
@@ -175,6 +188,7 @@ class Extent(collections.abc.Mapping):
     def __len__(self) -> int: ...
     @property
     def region(self) -> "Region": ...
+    def __eq__(self, other: "Extent") -> bool: ...
 
 class Point(collections.abc.Mapping):
     """
@@ -220,3 +234,10 @@ class Region:
         ...
 
     def __reduce__(self) -> Any: ...
+    def point_of_base_rank(self, rank: int) -> "Point":
+        """
+        Get the point in this region that corresponds to the given base rank
+        in the super-region that this region is a subset of.
+        """
+        ...
+    def __eq__(self, other: "Region") -> bool: ...

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

@@ -69,6 +69,7 @@ class HostMesh:
         ...
 
     def __reduce__(self) -> Any: ...
+    def __eq__(self, other: "HostMesh") -> bool: ...
 
 @final
 class BootstrapCommand:

python/monarch/_src/actor/actor_mesh.py

@@ -494,6 +494,16 @@ def items(self) -> Iterable[Tuple[Point, R]]:
         for i, _global_rank in enumerate(self._shape.ranks()):
             yield Point(i, extent), self._hy.get(i)
 
+    def values(self) -> Iterable[R]:
+        """
+        Generator that iterates over just the values in the mesh.
+
+        Returns:
+            Values at all coordinates.
+        """
+        for _, value in self.items():
+            yield value
+
     def __iter__(self) -> Iterator[Tuple[Point, R]]:
         return iter(self.items())
 

python/monarch/_src/actor/v1/host_mesh.py

@@ -47,12 +47,7 @@ def this_host() -> "HostMesh":
 
     This is just shorthand for looking it up via the context
     """
-    proc = this_proc()
-    if proc.host_mesh.is_fake_in_process:
-        return create_local_host_mesh("root_host")
-    host_mesh = proc.host_mesh
-    assert isinstance(host_mesh, HostMesh), "expected v1 HostMesh, got v0 HostMesh"
-    return host_mesh
+    return this_proc().host_mesh
 
 
 def this_proc() -> "ProcMesh":
@@ -241,6 +236,14 @@ def __reduce_ex__(self, protocol: ...) -> Tuple[Any, Tuple[Any, ...]]:
     def is_fake_in_process(self) -> bool:
         return self._is_fake_in_process
 
+    def __eq__(self, other: "HostMesh") -> bool:
+        return (
+            self._hy_host_mesh.block_on() == other._hy_host_mesh.block_on()
+            and self._region == other._region
+            and self.stream_logs == other.stream_logs
+            and self.is_fake_in_process == other.is_fake_in_process
+        )
+
 
 def fake_in_process_host(name: str) -> "HostMesh":
     """

python/monarch/_src/actor/v1/proc_mesh.py

@@ -80,12 +80,14 @@ def __init__(
         hy_proc_mesh: "Shared[HyProcMesh]",
         host_mesh: "HostMesh",
         region: Region,
+        root_region: Region,
         _device_mesh: Optional["DeviceMesh"] = None,
     ) -> None:
         _proc_mesh_registry.add(self)
         self._proc_mesh = hy_proc_mesh
         self._host_mesh = host_mesh
         self._region = region
+        self._root_region = root_region
         self._maybe_device_mesh = _device_mesh
         self._logging_manager = LoggingManager()
         self._controller_controller: Optional["_ControllerController"] = None
@@ -107,7 +109,16 @@ async def task() -> Literal[True]:
 
     @property
     def host_mesh(self) -> "HostMesh":
-        return self._host_mesh
+        if self.extent.nelements != 1:
+            raise NotImplementedError(
+                "`ProcMesh.host_mesh` is not yet supported for non-singleton proc meshes."
+            )
+        elif self._host_mesh.is_fake_in_process:
+            from monarch._src.actor.v1.host_mesh import create_local_host_mesh
+
+            return create_local_host_mesh("root_host")
+        else:
+            return self._host(0)
 
     @property
     def _ndslice(self) -> Slice:
@@ -134,6 +145,7 @@ async def task() -> HyProcMesh:
             PythonTask.from_coroutine(task()).spawn(),
             self._host_mesh,
             shape.region,
+            self._root_region,
             _device_mesh=device_mesh,
         )
 
@@ -176,7 +188,7 @@ def from_host_mesh(
         setup: Callable[[], None] | None = None,
         _attach_controller_controller: bool = True,
     ) -> "ProcMesh":
-        pm = ProcMesh(hy_proc_mesh, host_mesh, region)
+        pm = ProcMesh(hy_proc_mesh, host_mesh, region, region)
 
         if _attach_controller_controller:
             instance = context().actor_instance
@@ -341,7 +353,11 @@ async def __aexit__(
 
     @classmethod
     def _from_initialized_hy_proc_mesh(
-        cls, hy_proc_mesh: HyProcMesh, host_mesh: "HostMesh", region: Region
+        cls,
+        hy_proc_mesh: HyProcMesh,
+        host_mesh: "HostMesh",
+        region: Region,
+        root_region: Region,
     ) -> "ProcMesh":
         async def task() -> HyProcMesh:
             return hy_proc_mesh
@@ -350,13 +366,25 @@ async def task() -> HyProcMesh:
             PythonTask.from_coroutine(task()).spawn(),
             host_mesh,
             region,
+            root_region,
         )
 
     def __reduce_ex__(self, protocol: ...) -> Tuple[Any, Tuple[Any, ...]]:
         return ProcMesh._from_initialized_hy_proc_mesh, (
             self._proc_mesh.block_on(),
-            self.host_mesh,
+            self._host_mesh,
             self._region,
+            self._root_region,
+        )
+
+    def _host(self, proc_rank: int) -> "HostMesh":
+        base_proc_rank = self._region.slice().get(proc_rank)
+        n_procs = len(self._root_region.slice())
+        procs_per_host = n_procs // len(self._host_mesh.region.slice())
+        host_rank = base_proc_rank // procs_per_host
+        base_host_rank = self._host_mesh.region.slice().get(host_rank)
+        return self._host_mesh.slice(
+            **self._host_mesh.region.point_of_base_rank(base_host_rank)
         )
 
 

python/tests/test_host_mesh.py

@@ -74,7 +74,7 @@ def test_spawn_proc_mesh() -> None:
         Extent(["replicas", "hosts"], [2, 4]),
     )
     proc_mesh = host.spawn_procs(name="proc")
-    assert proc_mesh.host_mesh is host
+    assert proc_mesh._host_mesh is host
     assert proc_mesh._ndslice == host._ndslice
     assert tuple(proc_mesh._labels) == host._labels
     hy_proc_mesh = proc_mesh._proc_mesh.block_on()
@@ -89,7 +89,7 @@ def test_spawn_proc_mesh() -> None:
         name="proc_sliced", per_host={"gpus": 3, "just_for_fun": 4}
     )
     hy_sliced_proc = sliced_proc._proc_mesh.block_on()
-    assert sliced_proc.host_mesh is sliced_host
+    assert sliced_proc._host_mesh is sliced_host
     assert sliced_proc._ndslice == Slice(offset=0, sizes=[2, 3, 4], strides=[12, 4, 1])
     assert sliced_proc._labels == ["hosts", "gpus", "just_for_fun"]
     assert hy_sliced_proc.region.labels == sliced_proc._labels

python/tests/test_proc_mesh.py

@@ -128,7 +128,7 @@ def test_nested_meshes() -> None:
     for i, nested in enumerate([nested_0, nested_1]):
         region = cast(
             ProcMesh, cast(ActorMesh[TestActor], nested)._proc_mesh
-        ).host_mesh.region
+        )._host_mesh.region
         assert region.labels == ["hosts"]
         assert region.slice() == Slice(offset=i, sizes=[1], strides=[1])
     res_0 = nested_0.slice(gpus=0).call_on_other_mesh.call_one(nested_1).get()