Migrate to PythonActorMesh and PythonActorMeshRef

Summary:
This diff swaps `_ActorMeshRefImpl` with `PythonActorMesh[Ref]`.

The swap itself should be straightforward since `PythonActorMesh[Ref]` should be drop-in replacements for `_ActorMeshRefImpl`. Most of the complexity in this diff is from how I tried to add a toggle between them, just in case there is any bugs with  `PythonActorMesh[Ref]`, so we can quickly switch back to `_ActorMeshRefImpl`. What I did is:

1. Add wrapper classes `EitherPyActorMesh[Ref]`, whose underlying type can be either `PythonActorMesh[Ref]` or `_ActorMeshRefImpl`;
2. a env var `USE_STANDIN_ACTOR_MESH` is used to which one would be used when instantiating `EitherPyActorMesh[Ref]`.

The landing of this diff would mean all Python-side mesh API calls should go through Rust-side's `cast` code path, except several usages of `ActorIdRef`.

Differential Revision: D78355743

Author: pzhan9

python/monarch/_src/actor/actor_mesh.py

@@ -13,6 +13,7 @@
 import inspect
 import itertools
 import logging
+import os
 import pickle
 import random
 import traceback
@@ -56,6 +57,7 @@
     MonitoredOncePortReceiver,
     MonitoredPortReceiver,
     PythonActorMesh,
+    PythonActorMeshRef,
 )
 from monarch._rust_bindings.monarch_hyperactor.mailbox import (
     Mailbox,
@@ -65,9 +67,9 @@
     PortRef,
 )
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
+from monarch._rust_bindings.monarch_hyperactor.selection import Selection as HySelection
 from monarch._rust_bindings.monarch_hyperactor.shape import Point as HyPoint, Shape
 from monarch._rust_bindings.monarch_hyperactor.supervision import SupervisionError
-
 from monarch._rust_bindings.monarch_hyperactor.telemetry import enter_span, exit_span
 from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator
 from monarch._src.actor.future import Future
@@ -146,6 +148,28 @@ def set(debug_context: "DebugContext") -> None:
 Selection = Literal["all", "choose"] | int  # TODO: replace with real selection objects
 
 
+def to_hy_sel(selection: Selection, shape: Shape) -> HySelection:
+    if selection == "choose":
+        dim = len(shape.labels)
+        assert dim > 0
+        query = ",".join(["?"] * dim)
+        return HySelection.from_string(f"{query}")
+    elif selection == "all":
+        return HySelection.from_string("*")
+    else:
+        raise ValueError(f"invalid selection: {selection}")
+
+
+# A temporary gate used by the PythonActorMesh/PythonActorMeshRef migration.
+# We can use this gate to quickly roll back to using _ActorMeshRefImpl, if we
+# encounter any issues with the migration.
+#
+# This should be removed once we confirm PythonActorMesh/PythonActorMeshRef is
+# working correctly in production.
+def _use_standin_mesh() -> bool:
+    return bool(os.getenv("USE_STANDIN_ACTOR_MESH", default=False))
+
+
 # standin class for whatever is the serializable python object we use
 # to name an actor mesh. Hacked up today because ActorMesh
 # isn't plumbed to non-clients
@@ -158,6 +182,10 @@ def __init__(
         shape: Shape,
         actor_ids: List[ActorId],
     ) -> None:
+        if not _use_standin_mesh():
+            raise ValueError(
+                "ActorMeshRefImpl should only be used when USE_STANDIN_ACTOR_MESH is set"
+            )
         self._mailbox = mailbox
         self._actor_mesh = hy_actor_mesh
         # actor meshes do not have a way to look this up at the moment,
@@ -296,8 +324,8 @@ def __init__(
 
     def cast(
         self,
-        message: PythonMessage,
         selection: Selection,
+        message: PythonMessage,
     ) -> None:
         self._mailbox.post(self._actor_id, message)
 
@@ -309,6 +337,110 @@ def monitor(self) -> Optional[ActorMeshMonitor]:
         return None
 
 
+# A temporary wrapper used by the PythonActorMesh/PythonActorMeshRef migration.
+# This wrapper is used to enable switching between PythonActorMeshRef and
+# _ActorMeshRefImpl through the `USE_STANDIN_ACTOR_MESH` env var.
+#
+# This should be removed once we confirm PythonActorMesh/PythonActorMeshRef is
+# working correctly in production.
+class EitherPyActorMeshRef:
+    def __init__(self, inner: PythonActorMeshRef | _ActorMeshRefImpl) -> None:
+        if _use_standin_mesh():
+            assert isinstance(
+                inner, _ActorMeshRefImpl
+            ), "expect _ActorMeshRefImpl because env var USE_STANDIN_ACTOR_MESH is set"
+        else:
+            assert isinstance(
+                inner, PythonActorMeshRef
+            ), "expect PythonActorMeshRef because env var USE_STANDIN_ACTOR_MESH is not set"
+        self._inner: PythonActorMeshRef | _ActorMeshRefImpl = inner
+
+    def cast(
+        self, mailbox: Mailbox, selection: Selection, message: PythonMessage
+    ) -> None:
+        inner = self._inner
+        if isinstance(inner, _ActorMeshRefImpl):
+            inner.cast(message, selection)
+        elif isinstance(inner, PythonActorMeshRef):
+            inner.cast(mailbox, to_hy_sel(selection, self.shape), message)
+        else:
+            raise ValueError(f"unsupported mesh type: {inner.__class__.__name__}")
+
+    def slice(self, **kwargs) -> "EitherPyActorMeshRef":
+        return EitherPyActorMeshRef(self._inner.slice(**kwargs))
+
+    @property
+    def shape(self) -> Shape:
+        return self._inner.shape
+
+    def monitor(self) -> Optional[ActorMeshMonitor]:
+        return None
+
+
+# A temporary wrapper used by the PythonActorMesh/PythonActorMesh migration.
+# This wrapper is used to enable switching between PythonActorMesh and
+# _ActorMeshRefImpl through the `USE_STANDIN_ACTOR_MESH` env var.
+#
+# This should be removed once we confirm PythonActorMesh/PythonActorMeshRef is
+# working correctly in production.
+class EitherPyActorMesh:
+    def __init__(
+        self, actor_mesh: PythonActorMesh, mailbox: Mailbox, proc_mesh: "ProcMesh"
+    ) -> None:
+        if _use_standin_mesh():
+            inner = _ActorMeshRefImpl.from_hyperactor_mesh(
+                mailbox, actor_mesh, proc_mesh
+            )
+        else:
+            inner = actor_mesh
+        self._inner: PythonActorMesh | _ActorMeshRefImpl = inner
+        self._proc_mesh = proc_mesh
+
+    def bind(self) -> "EitherPyActorMeshRef":
+        inner = self._inner
+        if isinstance(inner, PythonActorMesh):
+            return EitherPyActorMeshRef(inner.bind())
+        elif isinstance(inner, _ActorMeshRefImpl):
+            return EitherPyActorMeshRef(inner)
+        else:
+            raise ValueError(f"unsupported mesh type: {inner.__class__.__name__}")
+
+    def cast(self, selection: Selection, message: PythonMessage) -> None:
+        inner = self._inner
+        if isinstance(inner, _ActorMeshRefImpl):
+            inner.cast(message, selection)
+        elif isinstance(inner, PythonActorMesh):
+            inner.cast(to_hy_sel(selection, self.shape), message)
+        else:
+            raise ValueError(f"unsupported mesh type: {inner.__class__.__name__}")
+
+    def slice(self, **kwargs) -> "EitherPyActorMeshRef":
+        return EitherPyActorMeshRef(self._inner.slice(**kwargs))
+
+    @property
+    def shape(self) -> Shape:
+        return self._inner.shape
+
+    def monitor(self) -> Optional[ActorMeshMonitor]:
+        return self._inner.monitor()
+
+    @property
+    def proc_mesh(self) -> "ProcMesh":
+        return self._proc_mesh
+
+    @property
+    def name_pid(self) -> Tuple[str, int]:
+        inner = self._inner
+        if isinstance(inner, _ActorMeshRefImpl):
+            return inner._name_pid
+        elif isinstance(inner, PythonActorMesh):
+            actor_id0 = inner.get(0)
+            assert actor_id0 is not None
+            return actor_id0.actor_name, actor_id0.pid
+        else:
+            raise ValueError(f"unsupported mesh type: {inner.__class__.__name__}")
+
+
 class Extent(NamedTuple):
     labels: Sequence[str]
     sizes: Sequence[int]
@@ -377,26 +509,28 @@ def call(self, *args: P.args, **kwargs: P.kwargs) -> "Future[ValueMesh[R]]":
         extent = self._send(args, kwargs, port=p)
 
         async def process() -> ValueMesh[R]:
-            results: List[R] = [None] * extent.nelements  # pyre-fixme[9]
+            results: Dict[int, R] = dict()
             for _ in range(extent.nelements):
                 rank, value = await r.recv()
                 results[rank] = value
             call_shape = Shape(
                 extent.labels,
                 NDSlice.new_row_major(extent.sizes),
             )
-            return ValueMesh(call_shape, results)
+            sorted_values = [results[rank] for rank in sorted(results)]
+            return ValueMesh(call_shape, sorted_values)
 
         def process_blocking() -> ValueMesh[R]:
-            results: List[R] = [None] * extent.nelements  # pyre-fixme[9]
+            results: Dict[int, R] = dict()
             for _ in range(extent.nelements):
                 rank, value = r.recv().get()
                 results[rank] = value
             call_shape = Shape(
                 extent.labels,
                 NDSlice.new_row_major(extent.sizes),
             )
-            return ValueMesh(call_shape, results)
+            sorted_values = [results[rank] for rank in sorted(results)]
+            return ValueMesh(call_shape, sorted_values)
 
         return Future(process, process_blocking)
 
@@ -428,7 +562,7 @@ def broadcast(self, *args: P.args, **kwargs: P.kwargs) -> None:
 class ActorEndpoint(Endpoint[P, R]):
     def __init__(
         self,
-        actor_mesh_ref: _ActorMeshRefImpl | ActorIdFakeMesh,
+        actor_mesh_ref: EitherPyActorMesh | EitherPyActorMeshRef | ActorIdFakeMesh,
         name: str,
         impl: Callable[Concatenate[Any, P], Awaitable[R]],
         mailbox: Mailbox,
@@ -469,7 +603,15 @@ def _send(
                 ),
                 bytes,
             )
-            self._actor_mesh.cast(message, selection)
+            mesh = self._actor_mesh
+            if isinstance(mesh, EitherPyActorMeshRef):
+                mesh.cast(self._mailbox, selection, message)
+            elif isinstance(mesh, EitherPyActorMesh) or isinstance(
+                mesh, ActorIdFakeMesh
+            ):
+                mesh.cast(selection, message)
+            else:
+                raise ValueError(f"unsupported mesh type: {mesh.__class__.__name__}")
         else:
             importlib.import_module("monarch." + "mesh_controller").actor_send(
                 self, self._name, bytes, refs, port
@@ -931,12 +1073,14 @@ class _ActorMeshTrait(MeshTrait):
     def __init__(
         self,
         Class: Type[T],
-        actor_mesh_ref: _ActorMeshRefImpl | ActorIdFakeMesh,
+        actor_mesh_ref: EitherPyActorMesh | EitherPyActorMeshRef | ActorIdFakeMesh,
         mailbox: Mailbox,
     ) -> None:
         self.__name__: str = Class.__name__
         self._class: Type[T] = Class
-        self._actor_mesh_ref: _ActorMeshRefImpl | ActorIdFakeMesh = actor_mesh_ref
+        self._actor_mesh_ref: (
+            EitherPyActorMesh | EitherPyActorMeshRef | ActorIdFakeMesh
+        ) = actor_mesh_ref
         self._mailbox: Mailbox = mailbox
         for attr_name in dir(self._class):
             attr_value = getattr(self._class, attr_name, None)
@@ -1003,20 +1147,22 @@ class ActorMeshHandle(_ActorMeshTrait, Generic[T]):
     def __init__(
         self,
         Class: Type[T],
-        actor_mesh: _ActorMeshRefImpl,
+        actor_mesh: PythonActorMesh,
         mailbox: Mailbox,
+        proc_mesh: "ProcMesh",
     ) -> None:
-        super().__init__(Class, actor_mesh, mailbox)
+        wrapper = EitherPyActorMesh(actor_mesh, mailbox, proc_mesh)
+        super().__init__(Class, wrapper, mailbox)
 
-    def _inner(self) -> "_ActorMeshRefImpl":
+    def _inner(self) -> "EitherPyActorMesh":
         mesh = self._actor_mesh_ref
         assert isinstance(
-            mesh, _ActorMeshRefImpl
+            mesh, EitherPyActorMesh
         ), f"mesh type is {mesh.__class__.__name__}"
         return mesh
 
     def bind(self) -> "ActorMeshRef[T]":
-        return ActorMeshRef(self._class, self._inner(), self._mailbox)
+        return ActorMeshRef(self._class, self._inner().bind(), self._mailbox)
 
     def _create(
         self,
@@ -1056,22 +1202,22 @@ def proc_mesh(self) -> "Optional[ProcMesh]":
 
     @property
     def name_pid(self) -> Tuple[str, int]:
-        return self._inner()._name_pid
+        return self._inner().name_pid
 
 
 class ActorMeshRef(_ActorMeshTrait, Generic[T]):
     def __init__(
         self,
         Class: Type[T],
-        actor_mesh: _ActorMeshRefImpl,
+        actor_mesh: EitherPyActorMeshRef,
         mailbox: Mailbox,
     ) -> None:
         super().__init__(Class, actor_mesh, mailbox)
 
-    def _inner(self) -> "_ActorMeshRefImpl":
+    def _inner(self) -> "EitherPyActorMeshRef":
         mesh = self._actor_mesh_ref
         assert isinstance(
-            mesh, _ActorMeshRefImpl
+            mesh, EitherPyActorMeshRef
         ), f"mesh type is {mesh.__class__.__name__}"
         return mesh
 

python/monarch/_src/actor/proc_mesh.py

@@ -36,12 +36,7 @@
     ProcMeshMonitor,
 )
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
-from monarch._src.actor.actor_mesh import (
-    _Actor,
-    _ActorMeshRefImpl,
-    Actor,
-    ActorMeshHandle,
-)
+from monarch._src.actor.actor_mesh import _Actor, Actor, ActorMeshHandle
 from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator, SimAllocator
 from monarch._src.actor.code_sync import RsyncMeshClient, WorkspaceLocation
 from monarch._src.actor.code_sync.auto_reload import AutoReloadActor
@@ -180,8 +175,9 @@ def _spawn_blocking(
         actor_mesh = self._proc_mesh.spawn_blocking(name, _Actor)
         service = ActorMeshHandle(
             Class,
-            _ActorMeshRefImpl.from_hyperactor_mesh(self._mailbox, actor_mesh, self),
+            actor_mesh,
             self._mailbox,
+            self._proc_mesh,
         )
         # useful to have this separate, because eventually we can reconstitute ActorMeshRef objects across pickling by
         # doing `ActorMeshRef(Class, actor_handle)` but not calling _create.
@@ -205,8 +201,9 @@ async def _spawn_nonblocking(
         actor_mesh = await self._proc_mesh.spawn_nonblocking(name, _Actor)
         service = ActorMeshHandle(
             Class,
-            _ActorMeshRefImpl.from_hyperactor_mesh(self._mailbox, actor_mesh, self),
+            actor_mesh,
             self._mailbox,
+            self._proc_mesh,
         )
         # useful to have this separate, because eventually we can reconstitute ActorMeshRef objects across pickling by
         # doing `ActorMeshRef(Class, actor_handle)` but not calling _create.