torchmonarch-nightly 2025.8.2__cp310-cp310-manylinux2014_x86_64.whl → 2025.9.4__cp310-cp310-manylinux2014_x86_64.whl

monarch/_src/actor/actor_mesh.py

@@ -6,6 +6,7 @@
 
 # pyre-unsafe
 
+import abc
 import collections
 import contextvars
 import functools
@@ -14,17 +15,18 @@ import itertools
 import logging
 import random
 import traceback
+from abc import abstractmethod, abstractproperty
 
 from dataclasses import dataclass
-from traceback import extract_tb, StackSummary
+from pprint import pformat
+from textwrap import indent
+from traceback import TracebackException
 from typing import (
     Any,
-    AsyncGenerator,
     Awaitable,
     Callable,
     cast,
     Concatenate,
-    Coroutine,
     Dict,
     Generator,
     Generic,
@@ -32,7 +34,6 @@ from typing import (
     Iterator,
     List,
     Literal,
-    NamedTuple,
     Optional,
     overload,
     ParamSpec,
@@ -48,24 +49,25 @@ from monarch._rust_bindings.monarch_hyperactor.actor import (
     PythonMessage,
     PythonMessageKind,
 )
-from monarch._rust_bindings.monarch_hyperactor.actor_mesh import PythonActorMesh
+from monarch._rust_bindings.monarch_hyperactor.actor_mesh import (
+    PythonActorMesh,
+    PythonActorMeshImpl,
+)
 from monarch._rust_bindings.monarch_hyperactor.mailbox import (
     Mailbox,
-    OncePortReceiver,
+    OncePortReceiver as HyOncePortReceiver,  # noqa: F401
     OncePortRef,
-    PortReceiver as HyPortReceiver,
+    PortReceiver as HyPortReceiver,  # noqa: F401
     PortRef,
+    UndeliverableMessageEnvelope,
 )
-
-if TYPE_CHECKING:
-    from monarch._rust_bindings.monarch_hyperactor.actor import PortProtocol
-    from monarch._rust_bindings.monarch_hyperactor.mailbox import PortReceiverBase
-
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
+from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask, Shared
+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.debugger.pdb_wrapper import PdbWrapper
 from monarch._src.actor.endpoint import (
     Endpoint,
     EndpointProperty,
@@ -74,21 +76,28 @@ from monarch._src.actor.endpoint import (
     Propagator,
     Selection,
 )
-from monarch._src.actor.future import Future
-from monarch._src.actor.pdb_wrapper import PdbWrapper
-
+from monarch._src.actor.future import DeprecatedNotAFuture, Future
 from monarch._src.actor.pickle import flatten, unflatten
-
+from monarch._src.actor.python_extension_methods import rust_struct
 from monarch._src.actor.shape import MeshTrait, NDSlice
 from monarch._src.actor.sync_state import fake_sync_state
-
+from monarch._src.actor.telemetry import METER
 from monarch._src.actor.tensor_engine_shim import actor_rref, actor_send
+from typing_extensions import Self
 
 if TYPE_CHECKING:
-    from monarch._src.actor.proc_mesh import ProcMesh
+    from monarch._rust_bindings.monarch_hyperactor.actor import PortProtocol
+    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.telemetry import get_monarch_tracer
+
+CallMethod = PythonMessageKind.CallMethod
 
 logger: logging.Logger = logging.getLogger(__name__)
 
+TRACER = get_monarch_tracer()
+
 Allocator = ProcessAllocator | LocalAllocator
 
 try:
@@ -106,22 +115,103 @@ class Point(HyPoint, collections.abc.Mapping):
     pass
 
 
-@dataclass
-class MonarchContext:
-    mailbox: Mailbox
-    proc_id: str
-    point: Point
+@rust_struct("monarch_hyperactor::mailbox::Instance")
+class Instance(abc.ABC):
+    @abstractproperty
+    def _mailbox(self) -> Mailbox:
+        """
+        This can be removed once we fix all the uses of mailbox to just use context instead.
+        """
+        ...
+
+    @property
+    def proc_id(self) -> str:
+        """
+        The proc_id of the current actor.
+        """
+        return self.actor_id.proc_id
+
+    @abstractproperty
+    def actor_id(self) -> ActorId:
+        """
+        The actor_id of the current actor.
+        """
+        ...
+
+    @property
+    def proc(self) -> "ProcMesh":
+        """
+        The singleton proc mesh that corresponds to just this actor.
+        """
+
+        return self.proc_mesh.slice(**self.rank)
+
+    """
+    Every actor is spawned over some mesh of processes. This identifies the point in that mesh where
+    the current actor was spawned. In other words, it is the `monarch.current_rank()` of
+    The actors __init__ message.
+    """
+    rank: Point
+    proc_mesh: "ProcMesh"
+    _controller_controller: "_ControllerController"
+
+    # 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]]"
+
+    def _add_child(self, child: "ActorMesh | ProcMesh") -> None:
+        if self._children is None:
+            self._children = [child]
+        else:
+            self._children.append(child)
+
+
+@rust_struct("monarch_hyperactor::mailbox::Context")
+class Context:
+    @property
+    def actor_instance(self) -> Instance:
+        """
+        Information about the actor currently running in this context.
+        """
+        ...
+
+    @property
+    def message_rank(self) -> Point:
+        """
+        Every message is sent as some broadcast of messages. This call identifies the
+        point in this space where the current actor is participating.
+
+        This is not the same self.actor_instance.rank: if the message was sent to some slice of
+        actors this identifies where the actor appears in the slice and not the identity of the actor.
+
+        These Point objects always exist. For singletons it will have 0 dimensions.
+        """
+        ...
 
     @staticmethod
-    def get() -> "MonarchContext":
-        return _context.get()
+    def _root_client_context() -> "Context": ...
 
 
-_context: contextvars.ContextVar[MonarchContext] = contextvars.ContextVar(
+_context: contextvars.ContextVar[Context] = contextvars.ContextVar(
     "monarch.actor_mesh._context"
 )
 
 
+def context() -> Context:
+    c = _context.get(None)
+    if c is None:
+        c = Context._root_client_context()
+        _context.set(c)
+        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()
+    return c
+
+
 @dataclass
 class DebugContext:
     pdb_wrapper: Optional[PdbWrapper] = None
@@ -149,6 +239,37 @@ A = TypeVar("A")
 _load_balancing_seed = random.Random(4)
 
 
+class _SingletonActorAdapator:
+    def __init__(self, inner: ActorId, shape: Optional[Shape] = None) -> None:
+        self._inner: ActorId = inner
+        if shape is None:
+            shape = singleton_shape
+        self._shape = shape
+
+    def cast(
+        self,
+        message: PythonMessage,
+        selection: str,
+        mailbox: Mailbox,
+    ) -> None:
+        mailbox.post(self._inner, message)
+
+    def new_with_shape(self, shape: Shape) -> "ActorMeshProtocol":
+        return _SingletonActorAdapator(self._inner, self._shape)
+
+    def supervision_event(self) -> "Optional[Shared[Exception]]":
+        return None
+
+    def stop(self) -> "PythonTask[None]":
+        raise NotImplementedError("stop()")
+
+    def initialized(self) -> "PythonTask[None]":
+        async def empty():
+            pass
+
+        return PythonTask.from_coroutine(empty())
+
+
 # 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
@@ -156,7 +277,7 @@ class _ActorMeshRefImpl:
     def __init__(
         self,
         mailbox: Mailbox,
-        hy_actor_mesh: Optional[PythonActorMesh],
+        hy_actor_mesh: Optional[PythonActorMeshImpl],
         proc_mesh: "Optional[ProcMesh]",
         shape: Shape,
         actor_ids: List[ActorId],
@@ -171,29 +292,19 @@ class _ActorMeshRefImpl:
 
     @staticmethod
     def from_hyperactor_mesh(
-        mailbox: Mailbox, hy_actor_mesh: PythonActorMesh, proc_mesh: "ProcMesh"
+        mailbox: Mailbox,
+        shape: Shape,
+        hy_actor_mesh: PythonActorMeshImpl,
+        proc_mesh: "ProcMesh",
     ) -> "_ActorMeshRefImpl":
-        shape: Shape = hy_actor_mesh.shape
         return _ActorMeshRefImpl(
             mailbox,
             hy_actor_mesh,
             proc_mesh,
-            hy_actor_mesh.shape,
+            shape,
             [cast(ActorId, hy_actor_mesh.get(i)) for i in range(len(shape))],
         )
 
-    @staticmethod
-    def from_actor_id(mailbox: Mailbox, actor_id: ActorId) -> "_ActorMeshRefImpl":
-        return _ActorMeshRefImpl(mailbox, None, None, singleton_shape, [actor_id])
-
-    @staticmethod
-    def from_actor_ref_with_shape(
-        ref: "_ActorMeshRefImpl", shape: Shape
-    ) -> "_ActorMeshRefImpl":
-        return _ActorMeshRefImpl(
-            ref._mailbox, None, None, shape, ref._please_replace_me_actor_ids
-        )
-
     def __getstate__(
         self,
     ) -> Tuple[Shape, List[ActorId], Mailbox]:
@@ -214,22 +325,19 @@ class _ActorMeshRefImpl:
         if self._actor_mesh is not None:
             if self._actor_mesh.stopped:
                 raise SupervisionError(
-                    "actor mesh is not in a healthy state: `ActorMesh` has been stopped"
+                    "actor mesh is unhealthy with reason: actor mesh is stopped due to proc mesh shutdown. "
+                    "`PythonActorMesh` has already been stopped."
                 )
 
             event = self._actor_mesh.get_supervision_event()
             if event is not None:
-                raise SupervisionError(f"actor mesh is not in a healthy state: {event}")
-
-    def send(self, rank: int, message: PythonMessage) -> None:
-        self._check_state()
-        actor = self._please_replace_me_actor_ids[rank]
-        self._mailbox.post(actor, message)
+                raise SupervisionError(f"actor mesh is unhealthy with reason: {event}")
 
     def cast(
         self,
         message: PythonMessage,
-        selection: Selection,
+        selection: str,
+        mailbox: Mailbox,
     ) -> None:
         self._check_state()
 
@@ -279,14 +387,44 @@ class _ActorMeshRefImpl:
         actor_id0 = self._please_replace_me_actor_ids[0]
         return actor_id0.actor_name, actor_id0.pid
 
-    async def stop(self):
-        await self._actor_mesh.stop()
+    @property
+    def shape(self) -> Shape:
+        return self._shape
+
+    @property
+    def proc_mesh(self) -> Optional["ProcMesh"]:
+        return self._proc_mesh
+
+    def new_with_shape(self, shape: Shape) -> "_ActorMeshRefImpl":
+        return _ActorMeshRefImpl(
+            self._mailbox, None, None, shape, self._please_replace_me_actor_ids
+        )
+
+    def supervision_event(self) -> "Optional[Shared[Exception]]":
+        if self._actor_mesh is None:
+            return None
+        return self._actor_mesh.supervision_event()
+
+    def stop(self) -> PythonTask[None]:
+        async def task():
+            if self._actor_mesh is not None:
+                self._actor_mesh.stop()
+
+        return PythonTask.from_coroutine(task())
+
+    def initialized(self) -> PythonTask[None]:
+        async def task():
+            pass
+
+        return PythonTask.from_coroutine(task())
 
 
 class ActorEndpoint(Endpoint[P, R]):
     def __init__(
         self,
-        actor_mesh_ref: _ActorMeshRefImpl,
+        actor_mesh: "ActorMeshProtocol",
+        shape: Shape,
+        proc_mesh: "Optional[ProcMesh]",
         name: MethodSpecifier,
         impl: Callable[Concatenate[Any, P], Awaitable[R]],
         mailbox: Mailbox,
@@ -294,16 +432,14 @@ class ActorEndpoint(Endpoint[P, R]):
         explicit_response_port: bool,
     ) -> None:
         super().__init__(propagator)
-        self._actor_mesh = actor_mesh_ref
+        self._actor_mesh = actor_mesh
         self._name = name
+        self._shape = shape
+        self._proc_mesh = proc_mesh
         self._signature: inspect.Signature = inspect.signature(impl)
         self._mailbox = mailbox
         self._explicit_response_port = explicit_response_port
 
-    def _supervise(self, r: HyPortReceiver | OncePortReceiver) -> Any:
-        mesh = self._actor_mesh._actor_mesh
-        return r if mesh is None else mesh.supervise(r)
-
     def _call_name(self) -> Any:
         return self._name
 
@@ -334,19 +470,17 @@ class ActorEndpoint(Endpoint[P, R]):
                 ),
                 bytes,
             )
-            self._actor_mesh.cast(message, selection)
+            self._actor_mesh.cast(message, selection, self._mailbox)
         else:
             actor_send(self, bytes, objects, port, selection)
-        shape = self._actor_mesh._shape
+        shape = self._shape
         return Extent(shape.labels, shape.ndslice.sizes)
 
-    def _port(self, once: bool = False) -> "PortTuple[R]":
-        p, r = PortTuple.create(self._mailbox, once)
-        if TYPE_CHECKING:
-            assert isinstance(
-                r._receiver, (HyPortReceiver | OncePortReceiver)
-            ), "unexpected receiver type"
-        return PortTuple(p, PortReceiver(self._mailbox, self._supervise(r._receiver)))
+    def _port(self, once: bool = False) -> "Tuple[Port[R], PortReceiver[R]]":
+        p, r = super()._port(once=once)
+        monitor: Optional[Shared[Exception]] = self._actor_mesh.supervision_event()
+        r._set_monitor(monitor)
+        return (p, r)
 
     def _rref(self, args, kwargs):
         self._check_arguments(args, kwargs)
@@ -386,28 +520,48 @@ def as_endpoint(
         if explicit_response_port
         else MethodSpecifier.ReturnsResponse
     )
-    return ActorEndpoint(
-        not_an_endpoint._ref._actor_mesh_ref,
+    return not_an_endpoint._ref._endpoint(
         kind(not_an_endpoint._name),
         getattr(not_an_endpoint._ref, not_an_endpoint._name),
-        not_an_endpoint._ref._mailbox,
         propagate,
         explicit_response_port,
     )
 
 
 class Accumulator(Generic[P, R, A]):
+    """
+    Accumulate the result of a broadcast invocation of an endpoint
+    across a sliced mesh.
+
+    Usage:
+            >>> counter = Accumulator(Actor.increment, 0, lambda x, y: x + y)
+    """
+
     def __init__(
         self, endpoint: Endpoint[P, R], identity: A, combine: Callable[[A, R], A]
     ) -> None:
+        """
+        Args:
+            endpoint: Endpoint to accumulate the result of.
+            identity: Initial value of the accumulated value before the first combine invocation.
+            combine: Lambda invoked for combining the result of the endpoint with the accumulated value.
+        """
         self._endpoint: Endpoint[P, R] = endpoint
         self._identity: A = identity
         self._combine: Callable[[A, R], A] = combine
 
     def accumulate(self, *args: P.args, **kwargs: P.kwargs) -> "Future[A]":
-        gen: Generator[Coroutine[None, None, R], None, None] = self._endpoint._stream(
-            *args, **kwargs
-        )
+        """
+        Accumulate the result of the endpoint invocation.
+
+        Args:
+            args: Arguments to pass to the endpoint.
+            kwargs: Keyword arguments to pass to the endpoint.
+
+        Returns:
+            Future that resolves to the accumulated value.
+        """
+        gen: Generator[Future[R], None, None] = self._endpoint.stream(*args, **kwargs)
 
         async def impl() -> A:
             value = self._identity
@@ -420,7 +574,7 @@ class Accumulator(Generic[P, R, A]):
 
 class ValueMesh(MeshTrait, Generic[R]):
     """
-    Container of return values, indexed by rank.
+    A mesh that holds the result of an endpoint invocation.
     """
 
     def __init__(self, shape: Shape, values: List[R]) -> None:
@@ -431,6 +585,18 @@ class ValueMesh(MeshTrait, Generic[R]):
         return ValueMesh(shape, self._values)
 
     def item(self, **kwargs) -> R:
+        """
+        Get the value at the given coordinates.
+
+        Args:
+            kwargs: Coordinates to get the value at.
+
+        Returns:
+            Value at the given coordinate.
+
+        Raises:
+            KeyError: If invalid coordinates are provided.
+        """
         coordinates = [kwargs.pop(label) for label in self._labels]
         if kwargs:
             raise KeyError(f"item has extra dimensions: {list(kwargs.keys())}")
@@ -438,17 +604,22 @@ class ValueMesh(MeshTrait, Generic[R]):
         return self._values[self._ndslice.nditem(coordinates)]
 
     def items(self) -> Iterable[Tuple[Point, R]]:
-        for rank in self._shape.ranks():
-            yield Point(rank, self._shape), self._values[rank]
+        """
+        Generator that returns values for the provided coordinates.
+
+        Returns:
+            Values at all coordinates.
+        """
+        extent = self._shape.extent
+        for i, rank in enumerate(self._shape.ranks()):
+            yield Point(i, extent), self._values[rank]
 
     def __iter__(self) -> Iterator[Tuple[Point, R]]:
         return iter(self.items())
 
-    def __len__(self) -> int:
-        return len(self._shape)
-
     def __repr__(self) -> str:
-        return f"ValueMesh({self._shape})"
+        body = indent(pformat(tuple(self.items())), "  ")
+        return f"ValueMesh({self._shape.extent}):\n{body}"
 
     @property
     def _ndslice(self) -> NDSlice:
@@ -467,14 +638,27 @@ def send(
     selection: Selection = "all",
 ) -> None:
     """
-    Fire-and-forget broadcast invocation of the endpoint across all actors in the mesh.
+        Fire-and-forget broadcast invocation of the endpoint across a given selection of the mesh.
 
-    This sends the message to all actors but does not wait for any result.
+        This sends the message to all actors but does not wait for any result. Use the port provided to
+        send the response back to the caller.
+
+    Args:
+        endpoint: Endpoint to invoke.
+        args: Arguments to pass to the endpoint.
+        kwargs: Keyword arguments to pass to the endpoint.
+        port: Handle to send the response to.
+        selection: Selection query representing a subset of the mesh.
     """
     endpoint._send(args, kwargs, port, selection)
 
 
 class Port(Generic[R]):
+    """
+    Handle used to send reliable in-order messages through a channel to
+    a PortReceiver.
+    """
+
     def __init__(
         self,
         port_ref: PortRef | OncePortRef,
@@ -486,6 +670,13 @@ class Port(Generic[R]):
         self._rank = rank
 
     def send(self, obj: R) -> None:
+        """
+            Fire-and-forget send R-typed objects in order
+            through a channel to its corresponding PortReceiver.
+
+        Args:
+            obj: R-typed object to send.
+        """
         self._port_ref.send(
             self._mailbox,
             PythonMessage(PythonMessageKind.Result(self._rank), _pickle(obj)),
@@ -522,62 +713,65 @@ R = TypeVar("R")
 
 T = TypeVar("T")
 
-if TYPE_CHECKING:
-    # Python <= 3.10 cannot inherit from Generic[R] and NamedTuple at the same time.
-    # we only need it for type checking though, so copypasta it until 3.11.
-    class PortTuple(NamedTuple, Generic[R]):
-        sender: "Port[R]"
-        receiver: "PortReceiver[R]"
-
-        @staticmethod
-        def create(mailbox: Mailbox, once: bool = False) -> "PortTuple[Any]":
-            handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
-            port_ref = handle.bind()
-            return PortTuple(
-                Port(port_ref, mailbox, rank=None),
-                PortReceiver(mailbox, receiver),
-            )
-else:
-
-    class PortTuple(NamedTuple):
-        sender: "Port[Any]"
-        receiver: "PortReceiver[Any]"
-
-        @staticmethod
-        def create(mailbox: Mailbox, once: bool = False) -> "PortTuple[Any]":
-            handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
-            port_ref = handle.bind()
-            return PortTuple(
-                Port(port_ref, mailbox, rank=None),
-                PortReceiver(mailbox, receiver),
-            )
-
 
 # advance lower-level API for sending messages. This is intentially
 # not part of the Endpoint API because they way it accepts arguments
 # and handles concerns is different.
-def port(endpoint: Endpoint[P, R], once: bool = False) -> "PortTuple[R]":
-    return endpoint._port(once)
+class Channel(Generic[R]):
+    """
+    An advanced low level API for a communication channel used for message passing
+    between actors.
 
+    Provides static methods to create communication channels with port pairs
+    for sending and receiving messages of type R.
+    """
 
-def ranked_port(
-    endpoint: Endpoint[P, R], once: bool = False
-) -> Tuple["Port[R]", "RankedPortReceiver[R]"]:
-    p, receiver = port(endpoint, once)
-    return p, RankedPortReceiver[R](receiver._mailbox, receiver._receiver)
+    @staticmethod
+    def open(once: bool = False) -> Tuple["Port[R]", "PortReceiver[R]"]:
+        """ """
+        mailbox = context().actor_instance._mailbox
+        handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
+        port_ref = handle.bind()
+        return (
+            Port(port_ref, mailbox, rank=None),
+            PortReceiver(mailbox, receiver),
+        )
+
+    @staticmethod
+    def open_ranked(once: bool = False) -> Tuple["Port[R]", "RankedPortReceiver[R]"]:
+        send, recv = Channel[R].open()
+        return (send, recv.ranked())
 
 
 class PortReceiver(Generic[R]):
+    """
+    Receiver for messages sent through a communication channel.
+
+    Handles receiving R-typed objects sent from a corresponding Port.
+    Asynchronously message reception with optional supervision
+    monitoring for error handling.
+    """
+
     def __init__(
         self,
         mailbox: Mailbox,
         receiver: "PortReceiverBase",
+        monitor: "Optional[Shared[Exception]]" = None,
     ) -> None:
         self._mailbox: Mailbox = mailbox
+        self._monitor = monitor
         self._receiver = receiver
 
     async def _recv(self) -> R:
-        return self._process(await self._receiver.recv_task())
+        awaitable = self._receiver.recv_task()
+        if self._monitor is None:
+            result = await awaitable
+        else:
+            # type: ignore
+            result, i = await PythonTask.select_one([self._monitor.task(), awaitable])
+            if i == 0:
+                raise result
+        return self._process(result)
 
     def _process(self, msg: PythonMessage) -> R:
         # TODO: Try to do something more structured than a cast here
@@ -593,6 +787,12 @@ class PortReceiver(Generic[R]):
     def recv(self) -> "Future[R]":
         return Future(coro=self._recv())
 
+    def ranked(self) -> "RankedPortReceiver[R]":
+        return RankedPortReceiver[R](self._mailbox, self._receiver, self._monitor)
+
+    def _set_monitor(self, monitor: "Optional[Shared[Exception]]"):
+        self._monitor = monitor
+
 
 class RankedPortReceiver(PortReceiver[Tuple[int, R]]):
     def _process(self, msg: PythonMessage) -> Tuple[int, R]:
@@ -614,6 +814,8 @@ singleton_shape = Shape([], NDSlice(offset=0, sizes=[], strides=[]))
 #  we need to signal to the consumer of the PythonTask object that the thread really isn't in an async context.
 # We do this by blanking out the running event loop during the call to the synchronous actor function.
 
+MESSAGES_HANDLED = METER.create_counter("py_mesages_handled")
+
 
 class _Actor:
     """
@@ -637,30 +839,28 @@ class _Actor:
 
     async def handle(
         self,
-        mailbox: Mailbox,
-        rank: int,
-        shape: Shape,
-        method_spec: MethodSpecifier,
+        ctx: Context,
+        method: MethodSpecifier,
         message: bytes,
         panic_flag: PanicFlag,
         local_state: Iterable[Any],
-        port: "PortProtocol",
+        response_port: "PortProtocol[Any]",
     ) -> None:
+        MESSAGES_HANDLED.add(1)
         # response_port can be None. If so, then sending to port will drop the response,
         # and raise any exceptions to the caller.
         try:
-            ctx: MonarchContext = MonarchContext(
-                mailbox, mailbox.actor_id.proc_id, Point(rank, shape)
-            )
             _context.set(ctx)
 
             DebugContext.set(DebugContext())
 
             args, kwargs = unflatten(message, local_state)
 
-            match method_spec:
+            match method:
                 case MethodSpecifier.Init():
-                    Class, *args = args
+                    ins = ctx.actor_instance
+                    Class, ins.proc_mesh, ins._controller_controller, *args = args
+                    ins.rank = ctx.message_rank
                     try:
                         self.instance = Class(*args, **kwargs)
                     except Exception as e:
@@ -668,13 +868,13 @@ class _Actor:
                             e, f"Remote actor {Class}.__init__ call failed."
                         )
                         raise e
-                    port.send(None)
+                    response_port.send(None)
                     return None
-                case MethodSpecifier.ReturnsResponse(name=method):
+                case MethodSpecifier.ReturnsResponse(name=method_name):
                     pass
-                case MethodSpecifier.ExplicitPort(name=method):
-                    args = (port, *args)
-                    port = DroppingPort()
+                case MethodSpecifier.ExplicitPort(name=method_name):
+                    args = (response_port, *args)
+                    response_port = DroppingPort()
 
             if self.instance is None:
                 # This could happen because of the following reasons. Both
@@ -687,52 +887,50 @@ class _Actor:
                 #    should never happen. It indicates either a bug in the
                 #    message delivery mechanism, or the framework accidentally
                 #    mixed the usage of cast and direct send.
-                error_message = f"Actor object is missing when executing method {method} on actor {mailbox.actor_id}."
+
+                error_message = f"Actor object is missing when executing method {method_name} on actor {ctx.actor_instance.actor_id}."
                 if self._saved_error is not None:
                     error_message += (
                         f" This is likely due to an earlier error: {self._saved_error}"
                     )
                 raise AssertionError(error_message)
-            the_method = getattr(self.instance, method)
+
+            the_method = getattr(self.instance, method_name)
             if isinstance(the_method, EndpointProperty):
-                module = the_method._method.__module__
                 the_method = functools.partial(the_method._method, self.instance)
-            else:
-                module = the_method.__module__
 
             if inspect.iscoroutinefunction(the_method):
 
                 async def instrumented():
-                    enter_span(
-                        module,
-                        method,
-                        str(ctx.mailbox.actor_id),
-                    )
-                    try:
-                        result = await the_method(*args, **kwargs)
-                        self._maybe_exit_debugger()
-                    except Exception as e:
-                        logging.critical(
-                            "Unhandled exception in actor endpoint",
-                            exc_info=e,
-                        )
-                        raise e
-                    exit_span()
+                    with TRACER.start_as_current_span(
+                        method_name,
+                        attributes={"actor_id": str(ctx.actor_instance.actor_id)},
+                    ):
+                        try:
+                            result = await the_method(*args, **kwargs)
+                            self._maybe_exit_debugger()
+                        except Exception as e:
+                            logging.critical(
+                                "Unhandled exception in actor endpoint",
+                                exc_info=e,
+                            )
+                            raise e
                     return result
 
                 result = await instrumented()
             else:
-                enter_span(module, method, str(ctx.mailbox.actor_id))
-                with fake_sync_state():
-                    result = the_method(*args, **kwargs)
-                self._maybe_exit_debugger()
-                exit_span()
-
-            port.send(result)
+                with TRACER.start_as_current_span(
+                    method_name,
+                    attributes={"actor_id": str(ctx.actor_instance.actor_id)},
+                ):
+                    with fake_sync_state():
+                        result = the_method(*args, **kwargs)
+                    self._maybe_exit_debugger()
+
+            response_port.send(result)
         except Exception as e:
             self._post_mortem_debug(e.__traceback__)
-            traceback.print_exc()
-            port.exception(ActorError(e))
+            response_port.exception(ActorError(e))
         except BaseException as e:
             self._post_mortem_debug(e.__traceback__)
             # A BaseException can be thrown in the case of a Rust panic.
@@ -754,21 +952,33 @@ class _Actor:
         DebugContext.set(DebugContext())
 
     def _post_mortem_debug(self, exc_tb) -> None:
-        from monarch._src.actor.debugger import DebugManager
+        from monarch._src.actor.debugger.debugger import debug_controller
 
         if (pdb_wrapper := DebugContext.get().pdb_wrapper) is not None:
             with fake_sync_state():
-                ctx = MonarchContext.get()
+                ctx = context()
+                msg_rank = ctx.message_rank
                 pdb_wrapper = PdbWrapper(
-                    ctx.point.rank,
-                    ctx.point.shape.coordinates(ctx.point.rank),
-                    ctx.mailbox.actor_id,
-                    DebugManager.ref().get_debug_client.call_one().get(),
+                    msg_rank.rank,
+                    {k: msg_rank[k] for k in msg_rank},
+                    ctx.actor_instance.actor_id,
+                    debug_controller(),
                 )
                 DebugContext.set(DebugContext(pdb_wrapper))
                 pdb_wrapper.post_mortem(exc_tb)
                 self._maybe_exit_debugger(do_continue=False)
 
+    def _handle_undeliverable_message(
+        self, message: UndeliverableMessageEnvelope
+    ) -> bool:
+        handle_undeliverable = getattr(
+            self.instance, "_handle_undeliverable_message", None
+        )
+        if handle_undeliverable is not None:
+            return handle_undeliverable(message)
+        else:
+            return False
+
 
 def _is_mailbox(x: object) -> bool:
     if hasattr(x, "__monarch_ref__"):
@@ -787,7 +997,7 @@ def _pickle(obj: object) -> bytes:
     return msg
 
 
-class Actor(MeshTrait):
+class Actor(MeshTrait, DeprecatedNotAFuture):
     @functools.cached_property
     def logger(cls) -> logging.Logger:
         lgr = logging.getLogger(cls.__class__.__name__)
@@ -806,20 +1016,48 @@ class Actor(MeshTrait):
             "actor implementations are not meshes, but we can't convince the typechecker of it..."
         )
 
-    def _new_with_shape(self, shape: Shape) -> "ActorMeshRef":
+    def _new_with_shape(self, shape: Shape) -> Self:
         raise NotImplementedError(
             "actor implementations are not meshes, but we can't convince the typechecker of it..."
         )
 
+    @property
+    def initialized(self):
+        raise NotImplementedError(
+            "actor implementations are not meshes, but we can't convince the typechecker of it..."
+        )
+
+    def _handle_undeliverable_message(
+        self, message: UndeliverableMessageEnvelope
+    ) -> bool:
+        # Return False to indicate that the undeliverable message was not handled.
+        return False
+
+
+class ActorMesh(MeshTrait, Generic[T], DeprecatedNotAFuture):
+    """
+    A group of actor instances of the same class.
+
+    Represents a collection of T-typed actor instances spawned at most once per process
+    that can be communicated with collectively or individually. Provides
+    methods for spawning actors, managing their lifecycle, and creating
+    endpoints for method invocation across the mesh.
+    """
 
-class ActorMeshRef(MeshTrait):
     def __init__(
-        self, Class: Type[T], actor_mesh_ref: _ActorMeshRefImpl, mailbox: Mailbox
+        self,
+        Class: Type[T],
+        inner: "ActorMeshProtocol",
+        mailbox: Mailbox,
+        shape: Shape,
+        proc_mesh: "Optional[ProcMesh]",
     ) -> None:
         self.__name__: str = Class.__name__
         self._class: Type[T] = Class
-        self._actor_mesh_ref: _ActorMeshRefImpl = actor_mesh_ref
+        self._inner: "ActorMeshProtocol" = inner
         self._mailbox: Mailbox = mailbox
+        self._shape = shape
+        self._proc_mesh = proc_mesh
         for attr_name in dir(self._class):
             attr_value = getattr(self._class, attr_name, None)
             if isinstance(attr_value, EndpointProperty):
@@ -832,11 +1070,9 @@ class ActorMeshRef(MeshTrait):
                 setattr(
                     self,
                     attr_name,
-                    ActorEndpoint(
-                        self._actor_mesh_ref,
+                    self._endpoint(
                         kind(attr_name),
                         attr_value._method,
-                        self._mailbox,
                         attr_value._propagator,
                         attr_value._explicit_response_port,
                     ),
@@ -847,53 +1083,95 @@ class ActorMeshRef(MeshTrait):
             return NotAnEndpoint(self, attr)
         raise AttributeError(attr)
 
-    def _create(
+    def _endpoint(
         self,
-        args: Iterable[Any],
-        kwargs: Dict[str, Any],
-    ) -> None:
+        name: MethodSpecifier,
+        impl: Callable[Concatenate[Any, P], Awaitable[R]],
+        propagator: Any,
+        explicit_response_port: bool,
+    ):
+        return ActorEndpoint(
+            self._inner,
+            self._shape,
+            self._proc_mesh,
+            name,
+            impl,
+            self._mailbox,
+            propagator,
+            explicit_response_port,
+        )
+
+    @classmethod
+    def _create(
+        cls,
+        Class: Type[T],
+        actor_mesh: "PythonActorMesh | PythonActorMeshImpl",
+        mailbox: Mailbox,
+        shape: Shape,
+        proc_mesh: "ProcMesh",
+        controller_controller: Optional["_ControllerController"],
+        # args and kwargs are passed to the __init__ method of the user defined
+        # python actor object.
+        *args: Any,
+        **kwargs: Any,
+    ) -> "ActorMesh[T]":
+        if isinstance(actor_mesh, PythonActorMeshImpl):
+            actor_mesh = _ActorMeshRefImpl.from_hyperactor_mesh(
+                mailbox, shape, actor_mesh, proc_mesh
+            )
+
+        mesh = cls(Class, actor_mesh, mailbox, shape, proc_mesh)
+
         async def null_func(*_args: Iterable[Any], **_kwargs: Dict[str, Any]) -> None:
             return None
 
-        ep = ActorEndpoint(
-            self._actor_mesh_ref,
+        # send __init__ message to the mesh to initialize the user defined
+        # python actor object.
+        ep = mesh._endpoint(
             MethodSpecifier.Init(),
             null_func,
-            self._mailbox,
             None,
             False,
         )
-        send(ep, (self._class, *args), kwargs)
-
-    def __reduce_ex__(
-        self, protocol: ...
-    ) -> "Tuple[Type[ActorMeshRef], Tuple[Any, ...]]":
-        return ActorMeshRef, (
-            self._class,
-            self._actor_mesh_ref,
-            self._mailbox,
+        send(ep, (mesh._class, proc_mesh, controller_controller, *args), kwargs)
+
+        return mesh
+
+    @classmethod
+    def from_actor_id(
+        cls,
+        Class: Type[T],
+        actor_id: ActorId,
+        mailbox: Mailbox,
+    ) -> "ActorMesh[T]":
+        return cls(
+            Class, _SingletonActorAdapator(actor_id), mailbox, singleton_shape, None
         )
 
+    def __reduce_ex__(self, protocol: ...) -> "Tuple[Type[ActorMesh], Tuple[Any, ...]]":
+        return ActorMesh, (self._class, self._inner, self._mailbox, self._shape, None)
+
     @property
     def _ndslice(self) -> NDSlice:
-        return self._actor_mesh_ref._shape.ndslice
+        return self._shape.ndslice
 
     @property
     def _labels(self) -> Iterable[str]:
-        return self._actor_mesh_ref._shape.labels
+        return self._shape.labels
 
-    def _new_with_shape(self, shape: Shape) -> "ActorMeshRef":
-        return ActorMeshRef(
-            self._class,
-            _ActorMeshRefImpl.from_actor_ref_with_shape(self._actor_mesh_ref, shape),
-            self._mailbox,
-        )
+    def _new_with_shape(self, shape: Shape) -> "ActorMesh[T]":
+        sliced = self._inner.new_with_shape(shape)
+        return ActorMesh(self._class, sliced, self._mailbox, shape, self._proc_mesh)
 
     def __repr__(self) -> str:
-        return f"ActorMeshRef(class={self._class}, shape={self._actor_mesh_ref._shape})"
+        return f"ActorMesh(class={self._class}, shape={self._shape}), inner={type(self._inner)})"
 
-    async def stop(self):
-        await self._actor_mesh_ref.stop()
+    def stop(self) -> "Future[None]":
+        return Future(coro=self._inner.stop())
+
+    @property
+    def initialized(self) -> Future[None]:
+        return Future(coro=self._inner.initialized())
 
 
 class ActorError(Exception):
@@ -909,27 +1187,35 @@ class ActorError(Exception):
         message: str = "A remote actor call has failed.",
     ) -> None:
         self.exception = exception
-        self.actor_mesh_ref_frames: StackSummary = extract_tb(exception.__traceback__)
+        # Need to stringify the exception early, because the PyPI package
+        # exceptiongroup may monkeypatch the "TracebackException" class for python
+        # versions < 3.11. If it gets unpickled in a different scope without
+        # using that monkeypatch, it'll have an exception in "format()".
+        # Store the traceback string instead which shouldn't change between machines.
+        actor_mesh_ref_tb = TracebackException.from_exception(exception).format()
+        # Replace any traceback lines to indicate it's a remote call traceback.
+        actor_mesh_ref_tb = (
+            s.replace(
+                "Traceback (most recent call last):",
+                "Traceback of where the remote call failed (most recent call last):",
+            )
+            for s in actor_mesh_ref_tb
+        )
+        self.exception_formatted = "".join(actor_mesh_ref_tb)
         self.message = message
 
     def __str__(self) -> str:
-        exe = str(self.exception)
-        actor_mesh_ref_tb = "".join(traceback.format_list(self.actor_mesh_ref_frames))
-        return (
-            f"{self.message}\n"
-            f"Traceback of where the remote call failed (most recent call last):\n{actor_mesh_ref_tb}{type(self.exception).__name__}: {exe}"
-        )
+        return f"{self.message}\n {self.exception_formatted}"
 
 
 def current_actor_name() -> str:
-    return str(MonarchContext.get().mailbox.actor_id)
+    return str(context().actor_instance.actor_id)
 
 
 def current_rank() -> Point:
-    ctx = MonarchContext.get()
-    return ctx.point
+    return context().message_rank
 
 
 def current_size() -> Dict[str, int]:
-    ctx = MonarchContext.get()
-    return dict(zip(ctx.point.shape.labels, ctx.point.shape.ndslice.sizes))
+    r = context().message_rank.extent
+    return {k: r[k] for k in r}