v1 host mesh and proc mesh python integration (#1346)

Summary:
Pull Request resolved: #1346

Write the python API for v1 host mesh and v1 proc mesh to (relatively closely) match the python APIs for v0.
ghstack-source-id: 313115893
exported-using-ghexport

Reviewed By: colin2328, zdevito

Differential Revision: D83223400

fbshipit-source-id: 50e5f47c613ed350efa1a9a617a418acc0b42da3

Author: samlurye

monarch_hyperactor/Cargo.toml

@@ -22,6 +22,7 @@ async-trait = "0.1.86"
 bincode = "1.3.3"
 clap = { version = "4.5.42", features = ["derive", "env", "string", "unicode", "wrap_help"] }
 erased-serde = "0.3.27"
+fastrand = "2.1.1"
 fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 futures = { version = "0.3.31", features = ["async-await", "compat"] }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }

monarch_hyperactor/src/shape.rs

@@ -12,6 +12,7 @@ use ndslice::Point;
 use ndslice::Region;
 use ndslice::Shape;
 use ndslice::Slice;
+use ndslice::View;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
@@ -94,6 +95,13 @@ impl PyExtent {
     fn keys<'py>(&self, py: Python<'py>) -> PyResult<PyObject> {
         Ok(self.inner.labels().into_bound_py_any(py)?.into())
     }
+
+    #[getter]
+    fn region(&self) -> PyRegion {
+        PyRegion {
+            inner: self.inner.region(),
+        }
+    }
 }
 
 impl From<Extent> for PyExtent {
@@ -139,13 +147,32 @@ impl PyRegion {
         }
     }
 
+    #[getter]
     fn labels(&self) -> Vec<String> {
         self.inner.labels().to_vec()
     }
 
     fn slice(&self) -> PySlice {
         self.inner.slice().clone().into()
     }
+
+    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
+        let bytes = bincode::serialize(&self.inner)
+            .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 = py
+            .import("monarch._rust_bindings.monarch_hyperactor.shape")?
+            .getattr("Region")?
+            .getattr("from_bytes")?;
+        Ok((from_bytes, py_bytes))
+    }
+
+    #[staticmethod]
+    fn from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<Self> {
+        Ok(bincode::deserialize::<Region>(bytes.as_bytes())
+            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?
+            .into())
+    }
 }
 
 impl From<Region> for PyRegion {

monarch_hyperactor/src/v1/actor_mesh.rs

@@ -9,12 +9,16 @@
 use hyperactor::ActorRef;
 use hyperactor_mesh::v1::actor_mesh::ActorMesh;
 use hyperactor_mesh::v1::actor_mesh::ActorMeshRef;
+use ndslice::Region;
 use ndslice::Selection;
+use ndslice::Slice;
+use ndslice::selection::structurally_equal;
 use ndslice::view::Ranked;
 use ndslice::view::RankedSliceable;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyException;
 use pyo3::exceptions::PyNotImplementedError;
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
@@ -89,9 +93,8 @@ impl ActorMeshProtocol for PythonActorMeshImpl {
     }
 
     fn supervision_event(&self) -> PyResult<Option<PyShared>> {
-        Err(PyErr::new::<PyNotImplementedError, _>(
-            "supervision_event is not implemented yet for v1::PythonActorMeshImpl",
-        ))
+        // FIXME: implement supervision events for v1 actor mesh.
+        Ok(None)
     }
 
     fn new_with_region(&self, region: &PyRegion) -> PyResult<Box<dyn ActorMeshProtocol>> {
@@ -113,13 +116,37 @@ impl ActorMeshProtocol for ActorMeshRef<PythonActor> {
     fn cast(
         &self,
         message: PythonMessage,
-        _selection: Selection,
+        selection: Selection,
         instance: &PyInstance,
     ) -> PyResult<()> {
-        instance_dispatch!(instance, |cx_instance| {
-            self.cast(cx_instance, message.clone())
-                .map_err(|err| PyException::new_err(err.to_string()))?;
-        });
+        if structurally_equal(&selection, &Selection::All(Box::new(Selection::True))) {
+            instance_dispatch!(instance, |cx_instance| {
+                self.cast(cx_instance, message.clone())
+                    .map_err(|err| PyException::new_err(err.to_string()))?;
+            });
+        } else if structurally_equal(&selection, &Selection::Any(Box::new(Selection::True))) {
+            let region = Ranked::region(self);
+            let random_rank = fastrand::usize(0..region.num_ranks());
+            let offset = region
+                .slice()
+                .get(random_rank)
+                .map_err(anyhow::Error::from)?;
+            let singleton_region = Region::new(
+                Vec::new(),
+                Slice::new(offset, Vec::new(), Vec::new()).map_err(anyhow::Error::from)?,
+            );
+            instance_dispatch!(instance, |cx_instance| {
+                self.sliced(singleton_region)
+                    .cast(cx_instance, message.clone())
+                    .map_err(|err| PyException::new_err(err.to_string()))?;
+            });
+        } else {
+            return Err(PyRuntimeError::new_err(format!(
+                "invalid selection: {:?}",
+                selection
+            )));
+        }
+
         Ok(())
     }
 

monarch_hyperactor/src/v1/host_mesh.rs

@@ -177,7 +177,9 @@ impl PyHostMesh {
         let bytes = bincode::serialize(&self.mesh_ref()?)
             .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 = wrap_pyfunction!(py_host_mesh_from_bytes, py)?.into_any();
+        let from_bytes =
+            PyModule::import(py, "monarch._rust_bindings.monarch_hyperactor.v1.host_mesh")?
+                .getattr("py_host_mesh_from_bytes")?;
         Ok((from_bytes, py_bytes))
     }
 }

monarch_hyperactor/src/v1/proc_mesh.rs

@@ -11,6 +11,7 @@ use hyperactor_mesh::v1::proc_mesh::ProcMesh;
 use hyperactor_mesh::v1::proc_mesh::ProcMeshRef;
 use monarch_types::PickledPyObject;
 use ndslice::View;
+use ndslice::view::RankedSliceable;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyException;
 use pyo3::exceptions::PyNotImplementedError;
@@ -181,6 +182,12 @@ impl PyProcMesh {
             "v1::PyProcMesh::stop not implemented",
         ))
     }
+
+    fn sliced(&self, region: &PyRegion) -> PyResult<Self> {
+        Ok(Self::new_ref(
+            self.mesh_ref()?.sliced(region.as_inner().clone()),
+        ))
+    }
 }
 
 #[derive(Clone)]

python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi

@@ -257,7 +257,7 @@ class PortProtocol(Generic[R], Protocol):
 class Actor(Protocol):
     async def handle(
         self,
-        context: Any,
+        ctx: Any,
         method: MethodSpecifier,
         message: bytes,
         panic_flag: PanicFlag,

python/monarch/_rust_bindings/monarch_hyperactor/shape.pyi

@@ -173,6 +173,8 @@ class Extent(collections.abc.Mapping):
     def __iter__(self) -> "Iterator[str]": ...
     def __getitem__(self, label: str) -> int: ...
     def __len__(self) -> int: ...
+    @property
+    def region(self) -> "Region": ...
 
 class Point(collections.abc.Mapping):
     """
@@ -204,6 +206,7 @@ class Region:
     """
     def __init__(self, labels: Sequence[str], slice: Slice) -> None: ...
     def as_shape(self) -> "Shape": ...
+    @property
     def labels(self) -> List[str]:
         """
         The labels for each dimension of the region.
@@ -215,3 +218,5 @@ class Region:
         The slice of the region.
         """
         ...
+
+    def __reduce__(self) -> Any: ...

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

@@ -80,3 +80,8 @@ class ProcMesh:
         ...
 
     def __repr__(self) -> str: ...
+    def sliced(self, region: Region) -> "ProcMesh":
+        """
+        Returns a new mesh that is a slice of this mesh with the given region.
+        """
+        ...

python/monarch/_src/actor/actor_mesh.py

@@ -92,6 +92,10 @@
     from monarch._rust_bindings.monarch_hyperactor.actor_mesh import ActorMeshProtocol
     from monarch._rust_bindings.monarch_hyperactor.mailbox import PortReceiverBase
     from monarch._src.actor.proc_mesh import _ControllerController, ProcMesh
+    from monarch._src.actor.v1.proc_mesh import (
+        _ControllerController as _ControllerControllerV1,
+        ProcMesh as ProcMeshV1,
+    )
 from monarch._src.actor.telemetry import get_monarch_tracer
 
 CallMethod = PythonMessageKind.CallMethod
@@ -142,7 +146,7 @@ def actor_id(self) -> ActorId:
         ...
 
     @property
-    def proc(self) -> "ProcMesh":
+    def proc(self) -> "ProcMesh | ProcMeshV1":
         """
         The singleton proc mesh that corresponds to just this actor.
         """
@@ -155,14 +159,14 @@ def proc(self) -> "ProcMesh":
     The actors __init__ message.
     """
     rank: Point
-    proc_mesh: "ProcMesh"
-    _controller_controller: "_ControllerController"
+    proc_mesh: "ProcMesh | ProcMeshV1"
+    _controller_controller: "_ControllerController | _ControllerControllerV1"
 
     # this property is used to hold the handles to actors and processes launched by this actor
     # in order to keep them alive until this actor exits.
-    _children: "Optional[List[ActorMesh | ProcMesh]]"
+    _children: "Optional[List[ActorMesh | ProcMesh | ProcMeshV1]]"
 
-    def _add_child(self, child: "ActorMesh | ProcMesh") -> None:
+    def _add_child(self, child: "ActorMesh | ProcMesh | ProcMeshV1") -> None:
         if self._children is None:
             self._children = [child]
         else:
@@ -212,13 +216,16 @@ def context() -> Context:
     if c is None:
         c = Context._root_client_context()
         _context.set(c)
+
+        # FIXME: Switch to the v1 APIs when it becomes the default.
         from monarch._src.actor.host_mesh import create_local_host_mesh
         from monarch._src.actor.proc_mesh import _get_controller_controller
 
         c.actor_instance.proc_mesh, c.actor_instance._controller_controller = (
             _get_controller_controller()
         )
-        c.actor_instance.proc_mesh._host_mesh = create_local_host_mesh()
+
+        c.actor_instance.proc_mesh._host_mesh = create_local_host_mesh()  # type: ignore
     return c
 
 
@@ -281,7 +288,7 @@ def __init__(
         self,
         actor_mesh: "ActorMeshProtocol",
         shape: Shape,
-        proc_mesh: "Optional[ProcMesh]",
+        proc_mesh: "Optional[ProcMesh] | Optional[ProcMeshV1]",
         name: MethodSpecifier,
         impl: Callable[Concatenate[Any, P], Awaitable[R]],
         propagator: Propagator,
@@ -931,7 +938,7 @@ def __init__(
         Class: Type[T],
         inner: "ActorMeshProtocol",
         shape: Shape,
-        proc_mesh: "Optional[ProcMesh]",
+        proc_mesh: "Optional[ProcMesh] | Optional[ProcMeshV1]",
     ) -> None:
         self.__name__: str = Class.__name__
         self._class: Type[T] = Class
@@ -986,8 +993,9 @@ def _create(
         Class: Type[T],
         actor_mesh: "PythonActorMesh",
         shape: Shape,
-        proc_mesh: "ProcMesh",
-        controller_controller: Optional["_ControllerController"],
+        proc_mesh: "ProcMesh | ProcMeshV1",
+        controller_controller: Optional["_ControllerController"]
+        | Optional["_ControllerControllerV1"],
         # args and kwargs are passed to the __init__ method of the user defined
         # python actor object.
         *args: Any,
@@ -1019,7 +1027,7 @@ def from_actor_id(
         return cls(Class, _SingletonActorAdapator(actor_id), singleton_shape, None)
 
     def __reduce_ex__(self, protocol: ...) -> "Tuple[Type[ActorMesh], Tuple[Any, ...]]":
-        return ActorMesh, (self._class, self._inner, self._shape, None)
+        return ActorMesh, (self._class, self._inner, self._shape, self._proc_mesh)
 
     @property
     def _ndslice(self) -> NDSlice:

python/monarch/_src/actor/host_mesh.py

@@ -25,15 +25,19 @@ def this_host() -> "HostMesh":
 
     This is just shorthand for looking it up via the context
     """
-    return context().actor_instance.proc.host_mesh
+    host_mesh = context().actor_instance.proc.host_mesh
+    assert isinstance(host_mesh, HostMesh), "expected v0 HostMesh, got v1 HostMesh"
+    return host_mesh
 
 
 def this_proc() -> "ProcMesh":
     """
     The current singleton process that this specific actor is
     running on
     """
-    return context().actor_instance.proc
+    proc = context().actor_instance.proc
+    assert isinstance(proc, ProcMesh), "expected v1 ProcMesh, got v0 ProcMesh"
+    return proc
 
 
 def create_local_host_mesh() -> "HostMesh":