Add selection parameter to PythonActorMesh.cast (#431)

Summary:
Pull Request resolved: #431

we want to allow python side to pass a selection parameter, rather than always `select all`. For example, the most common use cases is `def call_one`, which we want to pass `select any`.

Reviewed By: mariusae

Differential Revision: D77747932

fbshipit-source-id: a7e07f266e319e564c9517384cb46d32da1c5fdc

Author: pzhan9

monarch_hyperactor/src/actor_mesh.rs

@@ -21,6 +21,7 @@ use crate::actor::PythonMessage;
 use crate::mailbox::PyMailbox;
 use crate::proc::PyActorId;
 use crate::proc_mesh::Keepalive;
+use crate::selection::PySelection;
 use crate::shape::PyShape;
 
 #[pyclass(
@@ -43,10 +44,9 @@ impl PythonActorMesh {
 
 #[pymethods]
 impl PythonActorMesh {
-    fn cast(&self, message: &PythonMessage) -> PyResult<()> {
-        use ndslice::selection::dsl::*;
+    fn cast(&self, selection: &PySelection, message: &PythonMessage) -> PyResult<()> {
         self.try_inner()?
-            .cast(all(true_()), message.clone())
+            .cast(selection.inner().clone(), message.clone())
             .map_err(|err| PyException::new_err(err.to_string()))?;
         Ok(())
     }

monarch_hyperactor/src/selection.rs

@@ -20,6 +20,12 @@ pub struct PySelection {
     inner: Selection,
 }
 
+impl PySelection {
+    pub(crate) fn inner(&self) -> &Selection {
+        &self.inner
+    }
+}
+
 impl From<Selection> for PySelection {
     fn from(inner: Selection) -> Self {
         Self { inner }

python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi

@@ -10,16 +10,15 @@ from typing import final
 
 from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage
 from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox
-
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
-
+from monarch._rust_bindings.monarch_hyperactor.selection import Selection
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape
 
 @final
 class PythonActorMesh:
-    def cast(self, message: PythonMessage) -> None:
+    def cast(self, selection: Selection, message: PythonMessage) -> None:
         """
-        Cast a message to this mesh.
+        Cast a message to the selected actors in the mesh.
         """
 
     def get(self, rank: int) -> ActorId | None:

python/tests/_monarch/test_mailbox.py

@@ -25,9 +25,9 @@
     PortRef,
 )
 from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh
+from monarch._rust_bindings.monarch_hyperactor.selection import Selection
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape
 
-
 S = TypeVar("S")
 U = TypeVar("U")
 
@@ -157,7 +157,10 @@ def my_reduce(state: str, update: str) -> str:
     handle, receiver = proc_mesh.client.open_accum_port(accumulator)
     port_ref = handle.bind()
 
-    actor_mesh.cast(PythonMessage("echo", pickle.dumps("start"), port_ref, None))
+    actor_mesh.cast(
+        Selection.from_string("*"),
+        PythonMessage("echo", pickle.dumps("start"), port_ref, None),
+    )
 
     messge = await asyncio.wait_for(receiver.recv(), timeout=5)
     value = cast(str, pickle.loads(messge.message))