# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# pyre-strict
import abc
import asyncio
import collections
import contextvars
import functools
import importlib
import inspect
import logging
import threading
import warnings
from abc import abstractmethod, abstractproperty
from dataclasses import dataclass
from functools import cache
from pprint import pformat
from textwrap import indent
from traceback import TracebackException
from typing import (
Any,
Awaitable,
Callable,
cast,
Concatenate,
Dict,
Generic,
Iterable,
Iterator,
List,
Literal,
Optional,
overload,
ParamSpec,
Tuple,
Type,
TYPE_CHECKING,
TypeVar,
)
from monarch._rust_bindings.monarch_hyperactor.actor import (
DroppingPort,
MethodSpecifier,
PanicFlag,
PythonMessage,
PythonMessageKind,
)
from monarch._rust_bindings.monarch_hyperactor.buffers import FrozenBuffer
from monarch._rust_bindings.monarch_hyperactor.channel import BindSpec, ChannelTransport
from monarch._rust_bindings.monarch_hyperactor.config import configure
from monarch._rust_bindings.monarch_hyperactor.context import Instance as HyInstance
from monarch._rust_bindings.monarch_hyperactor.endpoint import ActorEndpoint
from monarch._rust_bindings.monarch_hyperactor.logging import log_endpoint_exception
from monarch._rust_bindings.monarch_hyperactor.mailbox import (
Mailbox,
OncePortRef,
PortRef,
UndeliverableMessageEnvelope,
)
from monarch._rust_bindings.monarch_hyperactor.pickle import (
PendingMessage,
pickle,
PicklingState,
)
from monarch._rust_bindings.monarch_hyperactor.proc import ActorAddr
from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
from monarch._rust_bindings.monarch_hyperactor.shape import Point as HyPoint, Shape
from monarch._rust_bindings.monarch_hyperactor.supervision import MeshFailure
from monarch._src.actor import config
from monarch._src.actor.debugger.pdb_wrapper import PdbWrapper
from monarch._src.actor.endpoint import (
Endpoint,
EndpointProperty,
NotAnEndpoint,
Propagator,
Selection,
)
from monarch._src.actor.future import Future
from monarch._src.actor.mpsc import Receiver # noqa: F401 - used in annotations
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, span
from monarch._src.actor.tensor_engine_shim import actor_rref, create_actor_message_kind
from opentelemetry.metrics import Counter
from typing_extensions import Self
if TYPE_CHECKING:
from monarch._rust_bindings.monarch_hyperactor.actor import (
PortProtocol,
QueuedMessage,
)
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, DeviceMesh, ProcMesh
def _assert_implements_endpoint(x: Endpoint[..., Any]) -> None: ...
def _check_actor_endpoint_satisfies_protocol(ep: ActorEndpoint[..., Any]) -> None:
# pyrefly: ignore [bad-argument-type]
_assert_implements_endpoint(ep)
logger: logging.Logger = logging.getLogger(__name__)
# `@rust_struct` (see `python_extension_methods.rust_struct`) attaches Python
# methods to a Rust pyclass as a side effect of importing the module that
# declares them. The queue-dispatch loop below calls `Receiver.recv` (from
# `mpsc`), which awaits `Event.wait`/`Event.clear` (from `waker`); all three are
# such patched methods. Under lazy imports (native-python builds) a plain
# `from ... import` is deferred until the imported *symbol* is used, but the
# loop only calls these methods on Rust-created *instances* and never touches
# the symbols, so the patches never run and the calls fail with `AttributeError:
# 'Receiver' object has no attribute 'recv'`. Force the modules to load here so
# their patches are applied before any actor dispatches a message.
#
# TODO: ideally `@rust_struct` would make a patched method resolve (and run) its
# defining module on first access, removing the need for these explicit imports.
# That likely requires a change to the lazy-import machinery outside monarch.
importlib.import_module("monarch._src.actor.mpsc")
importlib.import_module("monarch._src.actor.waker")
try:
from __manifest__ import fbmake # noqa
IN_PAR = bool(fbmake.get("par_style"))
except ImportError:
IN_PAR = False
T1 = TypeVar("T1")
T2 = TypeVar("T2")
[docs]class Point(HyPoint, collections.abc.Mapping):
"""A coordinate within a mesh.
A ``Point`` maps each mesh dimension label (for example ``"hosts"`` or
``"gpus"``) to its integer index along that dimension. It behaves as a
read-only mapping, so ``point["gpus"]`` is the index along the ``"gpus"``
dimension. ``current_rank`` and ``Context.message_rank`` return the
``Point`` that locates the current actor within the mesh that received the
message.
"""
pass
@rust_struct("monarch_hyperactor::context::Instance")
class Instance(abc.ABC):
# Optional tensor engine factory for mocking. When set, this is used
# instead of the real tensor engine when spawning device meshes.
_mock_tensor_engine_factory: Optional[Callable[["ProcMesh"], "DeviceMesh"]] = None
def spawn_tensor_engine(self, proc_mesh: "ProcMesh") -> "DeviceMesh":
"""
Spawn a tensor engine for this actor.
If a mock tensor engine factory is set, use it. Otherwise, use the
real tensor engine from mesh_controller.
"""
if self._mock_tensor_engine_factory is not None:
return self._mock_tensor_engine_factory(proc_mesh)
# pyre-ignore[21]: mesh_controller may not be visible to pyre in this target
from monarch.mesh_controller import spawn_tensor_engine as real_spawn
return real_spawn(proc_mesh)
@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) -> ActorAddr:
"""
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"
name: str # the name this actor was given on spawn
class_name: str # the fully qualified class name of the actor.
creator: Optional[
"CreatorInstance"
] # information about the actor who spawned this actor
# None if this actor is the spawning actor.
# 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[Any] | ProcMesh]]"
def _add_child(self, child: "ActorMesh[Any] | ProcMesh") -> None:
if self._children is None:
self._children = [child]
else:
self._children.append(child)
def _as_rust(self) -> HyInstance:
return cast(HyInstance, self)
@staticmethod
def _as_py(ins: HyInstance) -> "Instance":
return cast(Instance, ins)
def _as_creator(self) -> "CreatorInstance":
return CreatorInstance(
self.rank,
self.proc_mesh,
self.proc,
self.name,
self.class_name,
self.creator,
)
def __repr__(self) -> str:
return _qualified_name(self)
@abstractmethod
def abort(self, reason: Optional[str] = None) -> None:
"""
Abort the current actor. This will cause the actor to terminate
with a failure, and a supervision error will propagate to its creator.
"""
...
@abstractmethod
def _execution_start(self, name: str) -> int:
"""
Producer write-side for the mesh `execution` field: record the start
of a handler invocation, returning its in-flight token. Returns 0
(a no-op sentinel) when this instance has no execution tracker.
"""
...
@abstractmethod
def _execution_finish(self, token: int) -> None:
"""
End the handler invocation started by `_execution_start`. A no-op
for the 0 sentinel. Called from `_Actor.handle`'s `finally`.
"""
...
@abstractmethod
def kill(self, reason: Optional[str] = None) -> None:
"""
Terminate the current actor with a failure. A supervision error
propagates to its creator.
"""
...
@abstractmethod
def stop(self, reason: Optional[str] = None) -> None:
"""
Stop this actor instance and its children gracefully. The
actor drains pending messages and transitions to stopped status.
An actor stopping itself when its work is done is a normal
lifecycle pattern (map-reduce shards, batch processing,
worker pools, etc.; see ``sleep_actors.py`` for an example).
This is observable in the mesh: on exit the actor emits
``Signal::ChildStopped`` to its parent (always), so
ProcAgent *does* see the stop, and terminated snapshots
preserve post‑mortem state for introspection.
Use ``ActorMesh.stop()`` when you need coordinated, mesh-wide
shutdown.
"""
...
@abstractmethod
def set_system(self) -> None:
"""Mark this actor as system/infrastructure."""
...
@dataclass
class CreatorInstance:
"""
An instance that can be serialized so it can be passed around
to describe the creation hierarchy of an actor instance.
"""
rank: Point
proc_mesh: "ProcMesh"
proc: "ProcMesh"
name: str
class_name: Optional[str]
creator: Optional["CreatorInstance"]
def __repr__(self) -> str:
return _qualified_name(self)
def _qualified_name(ins: "CreatorInstance | Instance | None") -> str:
names = []
while ins:
class_prefix = "" if ins.class_name is None else f"{ins.class_name} "
rank_postfix = str(ins.rank) if len(ins.rank) > 0 else ""
names.append(f"<{class_prefix}{ins.name}{rank_postfix}>")
ins = ins.creator
return ".".join(reversed(names))
[docs]@rust_struct("monarch_hyperactor::context::Context")
class Context:
"""Runtime information about the currently executing actor.
A ``Context`` is returned by ``context()`` from within an endpoint. It
exposes the running actor through ``actor_instance`` and the position of the
current message within its broadcast through ``message_rank``.
"""
@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 _root_client_context() -> "Context": ...
@staticmethod
def _from_instance(instance: Instance) -> "Context": ...
_context: contextvars.ContextVar[Optional[Context]] = contextvars.ContextVar(
"monarch.actor_mesh._context", default=None
)
class _ActorFilter(logging.Filter):
"""
Logging filter that adds actor context to log messages.
This filter is automatically added to all logging handlers when code runs
inside a Monarch actor. It prefixes log messages with the actor's identity,
e.g. "[actor=<root>.MyActor] my log message".
We skip empty messages because torch.compile uses them for structured trace
logging. If we modify those, torch's formatter fails with "expected empty
string for trace".
"""
def __init__(self) -> None:
super().__init__()
def filter(self, record: logging.LogRecord) -> bool:
try:
if not config.prefix_python_logs_with_actor:
return True
ctx = _context.get()
if ctx is not None:
actor_prefix = f"[actor={ctx.actor_instance}] "
record.actor_prefix = actor_prefix # type: ignore[attr-defined]
# Skip empty messages (used for structured logging, e.g. torch.compile)
if record.msg:
record.msg = f"{actor_prefix}{record.msg}"
except Exception as e:
warnings.warn(
f"failed to add monarch actor information to python logs: {e}",
stacklevel=2,
)
return True
@cache
def _init_context_log_handler() -> None:
af: _ActorFilter = _ActorFilter()
logger = logging.getLogger()
for handler in logger.handlers:
handler.addFilter(af)
_original_addHandler: Any = logging.Logger.addHandler
def _patched_addHandler(self: logging.Logger, hdlr: logging.Handler) -> None:
_original_addHandler(self, hdlr)
if af not in hdlr.filters:
hdlr.addFilter(af)
logging.Logger.addHandler = _patched_addHandler
def _set_context(c: Context) -> contextvars.Token[Optional[Context]]:
_init_context_log_handler()
return _context.set(c)
def _reset_context(c: contextvars.Token[Optional[Context]]) -> None:
_context.reset(c)
T = TypeVar("T")
class _Lazy(Generic[T]):
def __init__(self, init: Callable[[], T]) -> None:
self._lock = threading.Lock()
self._val: Optional[T] = None
self._init = init
def get(self) -> T:
with self._lock:
if not self._val:
self._val = self._init()
return self._val
def try_get(self) -> Optional[T]:
return self._val
def _init_client_context(via: Optional[str] = None) -> Context:
"""
Create a client context that bootstraps an actor instance running on a real
local proc mesh on a real local host mesh.
When ``via`` is a non-empty ZMQ-style address, the local client's
gateway is connected to the gateway serving that address: outbound
traffic forwards over the duplex, and the remote gateway routes
return traffic back over the same duplex. ``this_host()`` still
names the current machine; attach only controls how the host's
procs are reached, not the host's identity. Use ``attach``
to supply ``via`` before the client context is first used.
"""
import atexit
from monarch._rust_bindings.monarch_hyperactor.host_mesh import bootstrap_host
from monarch._src.actor.host_mesh import default_bootstrap_cmd, HostMesh
from monarch._src.actor.proc_mesh import ProcMesh
hy_host_mesh, hy_proc_mesh, hy_instance = bootstrap_host(
default_bootstrap_cmd(), via=via
).block_on()
ctx = Context._from_instance(cast(Instance, hy_instance)) # type: ignore
# Set the context here to avoid recursive context creation:
token = _set_context(ctx)
try:
py_host_mesh = HostMesh._from_rust(hy_host_mesh)
py_proc_mesh = ProcMesh._from_rust(hy_proc_mesh, py_host_mesh)
finally:
_reset_context(token)
ctx.actor_instance.proc_mesh = py_proc_mesh
# Register shutdown_context as an atexit handler. Python atexit handlers
# run in LIFO order. shutdown_tokio_runtime was registered earlier (during
# module init), so this handler runs first — ensuring the actor system is
# cleanly shut down (connections flushed, acks delivered) before the tokio
# runtime is torn down.
#
# The timeout must be short enough that the process exits before
# the test executor's SIGTERM grace period (~2s). Combined with
# the 1s shutdown_tokio_runtime timeout, total atexit budget is
# ~2s, so we allow 1s here.
atexit.register(lambda: shutdown_context().get(timeout=1.0))
return ctx
_client_context: _Lazy[Context] = _Lazy(_init_client_context)
def attach(addr: str) -> None:
"""Bootstrap this process's client by attaching its gateway to the
remote duplex server at ``addr`` (a ZMQ-style address, e.g.
``"ipc:///tmp/sock"`` or ``"tcp://host:port"``).
Outbound client traffic is forwarded over the duplex and the remote
gateway routes return traffic back, so the client need not be
directly reachable by the procs in the mesh — e.g. a client running
outside a Kubernetes cluster whose procs run inside it.
Must be called before the client context is bootstrapped — i.e.
before the first ``context()`` / ``this_host()`` / ``this_proc()``
on the client — because actor and port refs snapshot their
location when they are created. We are working to remove this
sequencing dependency so attach can be configured independently of
first client-context use. Raises ``RuntimeError`` if the client
context has already been bootstrapped. ``this_host()`` still names
the current machine — attach only changes how this host's procs
are reached.
"""
with _client_context._lock:
if _client_context._val is not None:
raise RuntimeError(
"client already bootstrapped; call attach(addr) before "
"the first context()/this_host()/this_proc()"
)
_client_context._val = _init_client_context(via=addr)
_shutdown_done = False
[docs]def shutdown_context() -> "Future[None]":
"""Shutdown global actor context resources.
Idempotent: subsequent calls return an immediately-resolved future.
This is safe to call both explicitly and from atexit.
Returns:
Future[None]: A future that completes when shutdown is
finished. Call with .get() to wait for
completion.
"""
from monarch._src.actor.future import Future
if _shutdown_done:
async def _noop() -> None:
pass
return Future(coro=_noop())
c: Context | None = _context.get()
async def _shutdown_sequence() -> None:
global _shutdown_done
if _shutdown_done:
return
_shutdown_done = True
try:
from monarch._rust_bindings.monarch_hyperactor.host_mesh import (
shutdown_local_host_mesh,
)
# Shutdown the host mesh first, while the client actor is still
# alive to route messages. This drains children and joins the
# mailbox server, flushing receive-side acks.
await shutdown_local_host_mesh()
except RuntimeError:
# No local host mesh to shutdown
pass
# Stop the client actor and wait for it to reach terminal status.
# This ensures pending messages are drained and send-side acks
# are flushed before the tokio runtime is torn down.
if c is not None:
instance = c.actor_instance._as_rust()
await instance.stop_and_wait("shutdown")
_context.set(None)
return Future(coro=_shutdown_sequence())
[docs]def context() -> Context:
"""Return the ``Context`` for the currently executing actor.
Call this from within an endpoint to inspect the running actor and the
current message's position in the mesh. Outside an actor (on the client) it
returns the root client context.
"""
c = _context.get()
if c is None:
from monarch._src.actor.proc_mesh import _get_controller_controller
c = _client_context.get()
_set_context(c)
_, c.actor_instance._controller_controller = _get_controller_controller()
return c
_transport: Optional[BindSpec] = None
_transport_lock = threading.Lock()
[docs]def enable_transport(transport: "ChannelTransport | str") -> None:
"""
Allow monarch to communicate with transport type 'transport'
This must be called before any other calls in the monarch API.
If it isn't called, we will implicitly call
`monarch.enable_transport(ChannelTransport.Unix)` on the first monarch call.
Currently only one transport type may be enabled at one time.
In the future we may allow multiple to be enabled.
Supported transport values:
- ChannelTransport enum: ChannelTransport.Unix, ChannelTransport.TcpWithHostname, etc.
- string short cuts for the ChannelTransport enum:
- "tcp": ChannelTransport.TcpWithHostname
- "ipc": ChannelTransport.Unix
- "metatls": ChannelTransport.MetaTlsWithIpV6
- "metatls-hostname": ChannelTransport.MetaTlsWithHostname
- "tls": ChannelTransport.Tls (uses configurable TLS certs)
- ZMQ-style URL format string for explicit address, e.g.:
- "tcp://127.0.0.1:8080"
For Meta usage, use metatls-hostname
"""
if isinstance(transport, str):
# Handle string shortcuts for the ChannelTransport enum,
resolved = {
"tcp": ChannelTransport.TcpWithHostname,
"ipc": ChannelTransport.Unix,
"metatls": ChannelTransport.MetaTlsWithIpV6,
"metatls-hostname": ChannelTransport.MetaTlsWithHostname,
"tls": ChannelTransport.Tls,
}.get(transport)
if resolved is not None:
transport_config = BindSpec(resolved)
else:
transport_config = BindSpec(transport)
else:
# ChannelTransport enum
transport_config = BindSpec(transport)
global _transport
if _context.get() is not None:
# Context already initialized. It is only an error to enable a *different* transport;
# re-enabling the same transport is a no-op.
with _transport_lock:
if _transport == transport_config:
return
raise RuntimeError(
"`enable_transport()` must be called before any other calls in the monarch API. "
"If it isn't called, we will implicitly call `monarch.enable_transport(ChannelTransport.Unix)` "
"on the first monarch call."
)
with _transport_lock:
if _transport is not None and _transport != transport_config:
raise RuntimeError(
f"Only one transport type may be enabled at one time. "
f"Currently enabled transport type is `{_transport}`. "
f"Attempted to enable transport type `{transport_config}`."
)
_transport = transport_config
# pyre-ignore[6]: BindSpec is accepted by configure. We just do not expose
# it in the method's signature since BindSpec is not a public type.
configure(default_transport=transport_config)
@dataclass
class DebugContext:
pdb_wrapper: Optional[PdbWrapper] = None
@staticmethod
def get() -> "DebugContext":
return _debug_context.get()
@staticmethod
def set(debug_context: "DebugContext") -> None:
_debug_context.set(debug_context)
_debug_context: contextvars.ContextVar[DebugContext] = contextvars.ContextVar(
"monarch.actor_mesh._debug_context"
)
T = TypeVar("T")
P = ParamSpec("P")
R = TypeVar("R")
A = TypeVar("A")
def _check_endpoint_arguments(
method_name: MethodSpecifier,
signature: inspect.Signature,
args: Tuple[Any, ...],
kwargs: Dict[str, Any],
) -> None:
"""
Check that the arguments match the expected signature for the endpoint.
For Init methods, the message args contain an ActorInitArgs wrapper, so we
unpack it and validate the actual constructor arguments.
For ExplicitPort methods, the signature expects (self, port, *args, **kwargs),
so we bind with two None placeholders for self and port.
For other methods, the signature expects (self, *args, **kwargs),
so we bind with one None placeholder for self.
"""
match method_name:
# pyrefly: ignore [invalid-pattern]
case MethodSpecifier.Init():
# For Init, args[0] is ActorInitArgs which wraps the real constructor args
if len(args) != 1 or not isinstance(args[0], ActorInitArgs):
raise TypeError("Init message must contain exactly one ActorInitArgs")
init_args = args[0]
# Validate the actual constructor arguments against the signature
signature.bind(None, *init_args.args, **kwargs)
# pyrefly: ignore [invalid-pattern]
case MethodSpecifier.ExplicitPort():
signature.bind(None, None, *args, **kwargs)
case _:
signature.bind(None, *args, **kwargs)
def _create_endpoint_message(
method_name: MethodSpecifier,
signature: inspect.Signature,
args: Tuple[Any, ...],
kwargs: Dict[str, Any],
port_ref: "Optional[PortRef | OncePortRef]",
proc_mesh: "Optional[ProcMesh]",
) -> PendingMessage:
"""
Create a PythonMessage for sending to an actor endpoint.
Checks arguments, flattens them, and creates the appropriate message based on
whether the arguments contain monarch references.
Returns:
PythonMessage ready to be sent to the actor mesh
"""
_check_endpoint_arguments(method_name, signature, args, kwargs)
pickling_state = pickle(
(args, kwargs),
allow_tensor_engine_references=True,
allow_mesh_references=True,
)
objects = pickling_state.tensor_engine_references()
if not objects:
# pyrefly: ignore [bad-argument-count]
message_kind = PythonMessageKind.CallMethod(method_name, port_ref)
else:
message_kind = create_actor_message_kind(
method_name, proc_mesh, objects, port_ref
)
# pyrefly: ignore [bad-argument-type]
return PendingMessage(message_kind, pickling_state)
def _dispatch_actor_rref(
endpoint: Any,
args: Tuple[Any, ...],
kwargs: Dict[str, Any],
) -> Any:
_check_endpoint_arguments(endpoint._name, endpoint._signature, args, kwargs)
state = pickle((args, kwargs))
return actor_rref(endpoint, state)
@overload
def as_endpoint(
not_an_endpoint: Callable[P, R],
*,
propagate: Propagator = None,
explicit_response_port: Literal[False] = False,
) -> Endpoint[P, R]: ...
@overload
def as_endpoint(
not_an_endpoint: Callable[Concatenate["PortProtocol[R]", P], None],
*,
propagate: Propagator = None,
explicit_response_port: Literal[True],
) -> Endpoint[P, R]: ...
[docs]def as_endpoint(
not_an_endpoint: Any,
*,
propagate: Propagator = None,
explicit_response_port: bool = False,
) -> Any:
"""Treat an actor method that is not an ``@endpoint`` as one.
Use this to call a plain method of a spawned actor through the messaging
adverbs when the method was not decorated with ``@endpoint``, for example
``as_endpoint(actor.method).call(...)``. The options match those of
``endpoint``.
Args:
not_an_endpoint: An unannotated method of a spawned actor.
propagate: Tensor-shape propagation for ``rref``; see ``endpoint``.
explicit_response_port: When ``True``, the method receives a ``Port``
as its first argument and sends its result through it; see
``endpoint``.
Returns:
An ``Endpoint`` exposing the messaging adverbs.
Raises:
ValueError: If ``not_an_endpoint`` is not a method of a spawned actor.
"""
if not isinstance(not_an_endpoint, NotAnEndpoint):
raise ValueError("expected an method of a spawned actor")
kind = (
MethodSpecifier.ExplicitPort
if explicit_response_port
else MethodSpecifier.ReturnsResponse
)
return not_an_endpoint._ref._endpoint(
# pyrefly: ignore [bad-argument-count, bad-argument-type]
kind(not_an_endpoint._name),
getattr(not_an_endpoint._ref, not_an_endpoint._name),
propagate,
)
[docs]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)
"""
[docs] 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
[docs] def accumulate(self, *args: P.args, **kwargs: P.kwargs) -> "Future[A]":
"""
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: Iterator[Future[R]] = self._endpoint.stream(*args, **kwargs)
async def impl() -> A:
value = self._identity
for x in gen:
value = self._combine(value, await x)
return value
return Future(coro=impl())
[docs]@rust_struct("monarch_hyperactor::value_mesh::ValueMesh")
class ValueMesh(MeshTrait, Generic[R]):
"""
A mesh that holds the result of an endpoint invocation.
ValueMesh is returned when calling `.get()` on a Future from an endpoint
invocation on an ActorMesh or ProcMesh, or by awaiting the Future directly.
It contains the return values from all actors in the mesh, organized by
their coordinates.
Iteration:
The most efficient way to iterate over a ValueMesh is using `.items()`,
which yields (point, value) tuples:
>>> for point, result in value_mesh.items():
... rank = point["hosts"] * gpus_per_host + point["gpus"]
... print(f"Rank {rank}: {result}")
You can also iterate over just values:
>>> for result in value_mesh.values():
... process(result)
Accessing specific values:
Use `.item()` to extract a single value from a singleton mesh:
>>> single_value = value_mesh.slice(hosts=0, gpus=0).item()
Or with keyword arguments for multi-dimensional access:
>>> value = value_mesh.item(hosts=0, gpus=0)
Mesh operations:
ValueMesh supports the same operations as other MeshTrait types:
- `.flatten(dimension)`: Flatten to a single dimension
- `.slice(**coords)`: Select a subset of the mesh
Examples:
>>> # Sync API - Get results from all actors
>>> results = actor_mesh.endpoint.call(arg).get()
>>> for point, result in results.items():
... print(f"Actor at {point}: {result}")
>>>
>>> # Async API - Await the future directly
>>> results = await actor_mesh.endpoint.call(arg)
>>> for point, result in results.items():
... print(f"Actor at {point}: {result}")
>>>
>>> # Access a specific actor's result
>>> result_0 = results.item(hosts=0, gpus=0)
>>>
>>> # Flatten and iterate
>>> for point, result in results.flatten("rank").items():
... print(f"Rank {point.rank}: {result}")
"""
def __init__(self, shape: Shape, values: List[R]) -> None: ...
def __len__(self) -> int: ...
@property
def _shape(self) -> Shape: ...
def get(self, rank: int) -> R: ...
def _new_with_shape(self, shape: Shape) -> "ValueMesh[R]":
# Build a map from current global ranks -> local indices.
cur_ranks = list(self._shape.ranks())
pos = {g: i for i, g in enumerate(cur_ranks)}
# For each global rank of the target shape, pull from our
# current local index.
remapped = [self.get(pos[g]) for g in shape.ranks()]
return ValueMesh(shape, remapped)
[docs] def item(self, **kwargs: int) -> 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())}")
global_rank = self._ndslice.nditem(coordinates) # May include offset.
# Map global -> local (position in this shape's rank order).
ranks = list(self._shape.ranks())
try:
local_idx = ranks.index(global_rank)
except ValueError:
# Shouldn't happen if Shape is consistent, but keep a clear
# error.
raise IndexError(f"rank {global_rank} not in current shape")
try:
# pyre-ignore[21]: monarch.common only exists with tensor engine
from monarch.common.device_mesh import no_mesh
except ImportError:
return self.get(local_idx)
with no_mesh.activate():
return self.get(local_idx)
[docs] def items(self) -> Iterable[Tuple[Point, R]]:
"""
Generator that returns values for the provided coordinates.
Returns:
Values at all coordinates.
"""
extent = self._shape.extent
for i, _global_rank in enumerate(self._shape.ranks()):
# pyrefly: ignore [invalid-yield]
yield Point(i, extent), self.get(i)
def values(self) -> Iterable[R]:
"""
Generator that iterates over just the values in the mesh.
Returns:
Values at all coordinates.
"""
...
def __iter__(self) -> Iterator[Tuple[Point, R]]:
return iter(self.items())
def __repr__(self) -> str:
body = indent(pformat(tuple(self.items())), " ")
return f"ValueMesh({self._shape.extent}):\n{body}"
@property
def _ndslice(self) -> NDSlice:
return self._shape.ndslice
@property
# pyrefly: ignore [bad-override]
def _labels(self) -> Iterable[str]:
return self._shape.labels
def __reduce__(self) -> Tuple[Type["ValueMesh[R]"], Tuple[Shape, List[R]]]:
return (ValueMesh, (self._shape, list(self.values())))
[docs]def send(
endpoint: Endpoint[P, R],
args: Tuple[Any, ...],
kwargs: Dict[str, Any],
port: "Optional[Port[R]]" = None,
selection: Selection = "all",
) -> None:
"""
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. 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._port_ref if port is not None else None, selection
)
[docs]@rust_struct("monarch_hyperactor::actor::Port")
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,
instance: HyInstance,
rank: Optional[int],
) -> None: ...
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.
"""
...
def send_message(self, message: PythonMessage) -> None: ...
[docs] async def resolve_and_send(self, result: object) -> None:
# This is Port.send with mesh-reference resolution inserted before the
# already-serialized Result message is posted.
state = pickle(
result,
allow_tensor_engine_references=False,
allow_mesh_references=True,
)
kind = cast(PythonMessageKind, cast(Any, PythonMessageKind.Result)(self._rank))
message = PendingMessage(kind, state)
resolved = await Future(coro=cast(Any, message).resolve())
self.send_message(resolved)
def exception(self, obj: Exception) -> None: ...
def __reduce__(self) -> Tuple[Any, Tuple[Any, ...]]:
"""
When Port is sent over the wire, we do not want to send the actor instance
from the current context. Instead, we want to reconstruct the Port with
the receiver's context, since that is where the message will be sent
from through this port.
"""
def _reconstruct_port(
port_ref: PortRef | OncePortRef, rank: Optional[int]
) -> "Port[R]":
instance = context().actor_instance._as_rust()
return Port(port_ref, instance, rank)
return (
_reconstruct_port,
(self._port_ref, self._rank),
)
@property
def return_undeliverable(self) -> bool:
"""Whether to return undeliverable messages to the sender. By default, this is True.
Setting to False means that if the user of this port cannot deliver the message, it will
not be returned to the sender and will be dropped. Use this sparingly as
it could lead to unexpected hangs instead of errors that propagate back
to the owner of the actor"""
...
@return_undeliverable.setter
def return_undeliverable(self, value: bool) -> None: ...
@property
def _port_ref(self) -> PortRef | OncePortRef: ...
@property
def _rank(self) -> Optional[int]: ...
R = TypeVar("R")
T = TypeVar("T")
# 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.
[docs]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.
"""
[docs] @staticmethod
def open(once: bool = False) -> Tuple["Port[R]", "PortReceiver[R]"]:
""" """
actor_instance = context().actor_instance
mailbox = actor_instance._mailbox
handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
port_ref = handle.bind()
hy_instance = actor_instance._as_rust()
port = Port(port_ref, hy_instance, None)
return (
port,
PortReceiver(mailbox, receiver),
)
[docs] @staticmethod
def open_ranked(once: bool = False) -> Tuple["Port[R]", "RankedPortReceiver[R]"]:
send, recv = Channel[R].open()
return (send, recv.ranked())
[docs]class PortReceiver(Generic[R]):
"""
Receiver for messages sent through a communication channel.
Handles receiving R-typed objects sent from a corresponding Port.
"""
def __init__(
self,
mailbox: Mailbox,
receiver: "PortReceiverBase",
) -> None:
self._mailbox: Mailbox = mailbox
self._receiver = receiver
async def _recv(self) -> R:
return self._process(await self._receiver.recv_task())
def _process(self, msg: PythonMessage) -> R:
payload = cast(R, msg.decode())
match msg.kind:
# pyrefly: ignore [invalid-pattern]
case PythonMessageKind.Result():
return payload
# pyrefly: ignore [invalid-pattern]
case PythonMessageKind.Exception():
e = cast(Exception, payload)
raise e
case _:
raise ValueError(f"Unexpected message kind: {msg.kind}")
[docs] def recv(self) -> "Future[R]":
return Future(coro=self._recv())
[docs] def ranked(self) -> "RankedPortReceiver[R]":
return RankedPortReceiver[R](self._mailbox, self._receiver)
class RankedPortReceiver(PortReceiver[Tuple[int, R]]):
def _process(self, msg: PythonMessage) -> Tuple[int, R]:
rank = getattr(msg.kind, "rank", None)
if rank is None:
raise ValueError(
f"RankedPort receiver got a message without a rank {msg}",
)
# pyrefly: ignore [bad-return]
return rank, super()._process(msg)
singleton_shape = Shape([], NDSlice(offset=0, sizes=[], strides=[]))
# Currently the synchronous function of actors are run on a python thread that has an active event loop.
# Technically it is unsafe for them to block at all because they will block the loop of other
# calls, so all calls to .get() should be failing.
# But in the meantime, to implement get() by reusing async functions,
# 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: Counter = METER.create_counter("py_mesages_handled")
@dataclass
class ActorInitArgs:
Class: Type["Actor"]
proc_mesh: Optional["ProcMesh"]
controller_controller: Optional["_ControllerController"]
name: str
creator: Optional[CreatorInstance]
args: Tuple[Any, ...]
class _QueuePanicFlag:
"""Panic flag for queue dispatch mode.
Unlike the DummyPanicFlag, this one stores the exception so it can
be re-raised after handle() returns, ensuring proper cleanup.
"""
def __init__(self) -> None:
self.panic_exception: BaseException | None = None
def signal_panic(self, ex: BaseException) -> None:
self.panic_exception = ex
async def _dispatch_loop(
actor: Any,
receiver: "Receiver[QueuedMessage]",
self_instance: "Instance",
) -> None:
"""
Message loop for queue-dispatch mode. Called from Rust Actor::init.
Args:
actor: The Python actor object that implements ``handle``.
receiver: Channel receiver for queued messages.
self_instance: The actor's own Instance, used to kill self on
an unhandled exception.
"""
while True:
try:
msg = await receiver.recv()
await _handle_queued_message(actor, msg)
except asyncio.CancelledError:
return
except BaseException as e:
reason = "".join(TracebackException.from_exception(e).format())
self_instance.kill(reason)
raise
async def _handle_queued_message(actor: Any, msg: "QueuedMessage") -> None:
"""Handle a single queued message."""
panic_flag = _QueuePanicFlag()
await actor.handle(
msg.context,
msg.method,
msg.bytes,
panic_flag, # pyre-ignore[6]: _QueuePanicFlag implements PanicFlag protocol
msg.local_state,
msg.refs,
msg.response_port,
)
# If a panic was signaled, re-raise it after handle() has cleaned up.
if panic_flag.panic_exception is not None:
raise panic_flag.panic_exception
class _Actor:
"""
This is the message handling implementation of a Python actor.
The layering goes:
Rust `PythonActor` -> `_Actor` -> user-provided `Actor` instance
Messages are received from the Rust backend, and forwarded to the `handle`
methods on this class.
This class wraps the actual `Actor` instance provided by the user, and
routes messages to it, managing argument serialization/deserialization and
error handling.
"""
def __init__(self) -> None:
self.instance: object | None = None
# TODO: (@pzhang) remove this with T229200522
self._saved_error: ActorError | None = None
self._method_cache: Dict[str, Tuple[Callable[..., Any], bool, bool]] = {}
async def handle(
self,
ctx: Context,
method: MethodSpecifier,
message: FrozenBuffer,
panic_flag: PanicFlag,
local_state: List[Any],
mesh_references: List[Any],
response_port: "PortProtocol[Any]",
) -> None:
MESSAGES_HANDLED.add(1)
# Initialize method_name before try block so it's always defined
method_name = method.name
# response_port can be None. If so, then sending to port will drop the response,
# and raise any exceptions to the caller.
try:
_set_context(ctx)
DebugContext.set(DebugContext())
args, kwargs = PicklingState(
message, local_state, mesh_references
).unpickle()
match method:
# pyrefly: ignore [invalid-pattern]
case MethodSpecifier.Init():
ins = ctx.actor_instance
(args,) = args
init_args = cast(ActorInitArgs, args)
Class = init_args.Class
ins.proc_mesh = cast("ProcMesh", init_args.proc_mesh)
ins._controller_controller = cast(
"_ControllerController", init_args.controller_controller
)
ins.name = init_args.name
ins.creator = init_args.creator
args = init_args.args
ins.rank = ctx.message_rank
ins.class_name = f"{Class.__module__}.{Class.__qualname__}"
try:
self.instance = Class(*args, **kwargs)
# Check if there's a tensor engine mock registered for this actor class.
# If so, set _mock_tensor_engine_factory on the Instance for use by
# Instance.spawn_tensor_engine().
from monarch._src.actor.mock import get_tensor_engine_factory
mock_factory = get_tensor_engine_factory(Class)
if mock_factory is not None:
ins._mock_tensor_engine_factory = (
lambda proc_mesh: mock_factory(proc_mesh)
)
# PY-SYS-2: If Class._is_system_actor is true,
# ins.set_system() must run during
# MethodSpecifier::Init before first
# introspection publish.
if getattr(Class, "_is_system_actor", False):
ins.set_system()
self._maybe_exit_debugger()
except Exception as e:
self._saved_error = ActorError(
e, f"Actor call {ins.name}.{method_name} failed."
)
raise
response_port.send(None)
return
# pyrefly: ignore [invalid-pattern]
case MethodSpecifier.ReturnsResponse():
pass
# pyrefly: ignore [invalid-pattern]
case MethodSpecifier.ExplicitPort():
args = (response_port, *args)
response_port = DroppingPort()
if self.instance is None:
assert self._saved_error is not None
error_message = (
f'Actor object is missing when executing method "{method_name}" on actor {ctx.actor_instance.actor_id}. '
f"This is due to an earlier error: {self._saved_error}."
)
raise AssertionError(error_message)
if method_name not in self._method_cache:
the_method = getattr(self.instance, method_name)
should_instrument = False
if isinstance(the_method, EndpointProperty):
should_instrument = the_method._instrument
the_method = functools.partial(the_method._method, self.instance)
self._method_cache[method_name] = (
the_method,
should_instrument,
inspect.iscoroutinefunction(the_method),
)
the_method, should_instrument, is_coro = self._method_cache[method_name]
# Bracket the real user-method invocation so the actor reports it
# as in-flight (PE-2: only the invocation, never Init/plumbing).
# The finally drains on normal return, exception, panic, and
# cancellation; `_execution_start` returns 0 (a no-op token) when
# the instance has no tracker.
token = ctx.actor_instance._execution_start(method_name)
try:
if is_coro:
if should_instrument:
with span(method_name):
result = await the_method(*args, **kwargs)
else:
result = await the_method(*args, **kwargs)
self._maybe_exit_debugger()
else:
with fake_sync_state():
if should_instrument:
with span(method_name):
result = the_method(*args, **kwargs)
else:
result = the_method(*args, **kwargs)
self._maybe_exit_debugger()
finally:
ctx.actor_instance._execution_finish(token)
response = response_port.resolve_and_send(result)
if isinstance(response, PythonTask):
await Future(coro=response)
else:
await response
except Exception as e:
log_endpoint_exception(e, method_name, ctx.actor_instance.actor_id)
self._post_mortem_debug(e.__traceback__)
response_port.exception(
ActorError(
e,
f"Actor call {ctx.actor_instance.name}.{method_name} failed.",
)
)
return
except BaseException as e:
self._post_mortem_debug(e.__traceback__)
# A BaseException can be thrown in the case of a Rust panic.
# In this case, we need a way to signal the panic to the Rust side.
# See [Panics in async endpoints]
try:
panic_flag.signal_panic(
ActorError(
e,
f"Actor call {ctx.actor_instance.name}.{method_name} failed with BaseException.",
)
)
except Exception:
# The channel might be closed if the Rust side has already detected the error
pass
raise
def _maybe_exit_debugger(self, do_continue: bool = True) -> None:
if (pdb_wrapper := DebugContext.get().pdb_wrapper) is not None:
if do_continue:
pdb_wrapper.clear_all_breaks()
pdb_wrapper.do_continue("")
pdb_wrapper.end_debug_session()
DebugContext.set(DebugContext())
def _post_mortem_debug(self, exc_tb: Any) -> None:
from monarch._src.actor.debugger.debug_controller import debug_controller
if (pdb_wrapper := DebugContext.get().pdb_wrapper) is not None:
with fake_sync_state():
ctx = context()
msg_rank = ctx.message_rank
pdb_wrapper = PdbWrapper(
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, cx: Context, message: UndeliverableMessageEnvelope
) -> bool:
_set_context(cx)
handle_undeliverable = getattr(
self.instance, "_handle_undeliverable_message", None
)
if handle_undeliverable is not None:
return handle_undeliverable(message)
else:
return False
async def __supervise__(self, cx: Context, *args: Any, **kwargs: Any) -> object:
"""Dispatch the user's ``__supervise__``.
Mirrors ``__cleanup__``: both sync and async user methods are
supported. An ``async def`` user method is awaited on the actor's
asyncio event loop; a sync one runs under :func:`fake_sync_state` so
it cannot observe a running loop.
"""
_set_context(cx)
instance = self.instance
if instance is None:
# This could happen because of the following reasons. Both
# indicates a possible bug in the framework:
# 1. the execution of the previous message for "__init__" failed,
# but that error is not surfaced to the caller.
# - TODO(T229200522): there is a known bug. fix it.
# 2. this message is delivered to this actor before the previous
# message of "__init__" is delivered. Out-of-order delivery
# 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 __supervise__ on actor {cx.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)
supervise = getattr(instance, "__supervise__", None)
if not _is_user_override(supervise):
# If there is no __supervise__ override, the default is to return
# None. The supervision error is not handled here and will propagate
# to the next owner.
return None
if inspect.iscoroutinefunction(supervise):
return await supervise(*args, **kwargs)
else:
with fake_sync_state():
# pyrefly: ignore [not-callable]
return supervise(*args, **kwargs)
async def __cleanup__(self, cx: Context, exc: str | Exception | None) -> None:
"""Cleans up any resources owned by this Actor before stopping. Automatically
called even if there is an error"""
_context.set(cx)
instance = self.instance
if instance is None:
# If there is no instance, there's nothing to clean up, the actor
# was never constructed
return None
# Forward a call to cleanup on this actor to the user-provided instance.
cleanup = getattr(instance, "__cleanup__", None)
if not _is_user_override(cleanup):
return None
if isinstance(exc, str):
# Wrap the string in an exception object so the main API of __cleanup__
# is to take an optional exception object.
# The raw string is used for wider compatibility with other error
# types for now.
exc = Exception(exc)
if inspect.iscoroutinefunction(cleanup):
return await cleanup(exc)
else:
with fake_sync_state():
# pyrefly: ignore [not-callable]
return cleanup(exc)
def __repr__(self) -> str:
return f"_Actor(instance={self.instance!r})"
F = TypeVar("F", bound=Callable[..., Any])
def _doc_stub(fn: F) -> F:
"""Mark a method on ``Actor`` as a documentation-only stub.
The base class declares ``__cleanup__`` and ``__supervise__`` so that
``help()`` and IDEs can surface the docstring, but subclasses that do
not override them must be treated as if they provided no implementation.
Runtime code that detects user overrides (the sync/async endpoint check
in :class:`ActorMesh` and the dispatch in :class:`_Actor`) consults this
marker and ignores stub-marked methods.
"""
# pyre-ignore[16]: function attributes are dynamic
fn._monarch_doc_stub = True
return fn
def _is_user_override(method: Any) -> bool:
"""Return True if ``method`` is a real user override, not a doc stub."""
return method is not None and not getattr(method, "_monarch_doc_stub", False)
[docs]class Actor(MeshTrait):
"""Base class for actors.
Subclass ``Actor`` to define an actor, and decorate the methods that form
its public API with ``@endpoint``. Actors are spawned onto a ``ProcMesh``
(for example with ``this_proc().spawn(...)``), after which their endpoints
are invoked remotely through the messaging adverbs. Each actor processes its
messages sequentially and participates in the supervision tree.
"""
@functools.cached_property
def logger(cls) -> logging.Logger:
lgr = logging.getLogger(cls.__class__.__name__)
lgr.setLevel(logging.DEBUG)
return lgr
@property
def _ndslice(self) -> NDSlice:
raise NotImplementedError(
"actor implementations are not meshes, but we can't convince the typechecker of it..."
)
@property
def _labels(self) -> Tuple[str, ...]:
raise NotImplementedError(
"actor implementations are not meshes, but we can't convince the typechecker of it..."
)
# pyrefly: ignore [not-a-type]
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) -> Any:
raise NotImplementedError(
"actor implementations are not meshes, but we can't convince the typechecker of it..."
)
# Methods to be (optionally) overridden by user code
def _handle_undeliverable_message(
self, message: UndeliverableMessageEnvelope
) -> bool:
"""If a message sent by this actor cannot be delivered to its destination, this
method is called. The default implementation returns False, indicating that the
undeliverable message was not handled. Returning True indicates that the message
was handled in some way and does not need to be escalated as an error."""
# Return False to indicate that the undeliverable message was not handled.
return False
@_doc_stub
def __supervise__(self, failure: MeshFailure) -> bool:
"""Called when the actor observes a failure in a resource that it
owns. A resource is a host, proc, actor, or meshes of these.
If a truthy value is returned, the failure is considered handled and will not
propagate any further. If a falsey value is returned, the failure will be
further sent to the owner of this Actor.
Note that this is *not* called for errors within this Actor.
Rank failures in an owned ``ActorMesh`` are reported independently. A
mesh with ``N`` ranks may call this method multiple times for that mesh,
but it will call at most once per failed rank, for at most ``N`` total
rank-scoped failure callbacks. Once ``N`` callbacks have arrived for a
mesh, no more rank-scoped failure callbacks should arrive for that mesh.
If only some ranks fail, the mesh remains usable through slices that
exclude the failed ranks.
Overrides may be declared with either ``def`` or ``async def``. An
``async def`` override is awaited on the actor's asyncio event loop --
the same loop that runs endpoint coroutines -- so it can ``await``
other endpoints or I/O. A sync override runs under ``fake_sync_state``
and cannot call ``asyncio.get_running_loop``. If the override raises,
the exception is treated as a new supervision event chained to the
one being handled, matching the ``__exit__`` convention of context
managers.
This method is documentation-only on ``Actor``; subclasses provide the
real implementation.
"""
...
@_doc_stub
def __cleanup__(self, exc: Exception | None) -> None:
"""Called when the actor stops, normally (via ``ActorMesh.stop()``) or
because of an error. The same ``__cleanup__`` runs in both cases;
``exc`` is ``None`` on a normal stop and carries the exception on an
error stop. It is *not* called on fatal failures such as OOMs, panics,
or signals like ``SIGSEGV``. If it exceeds ``HYPERACTOR_CLEANUP_TIMEOUT``,
it is cancelled and the actor is placed in an error state.
By the time this runs, every mesh this actor owns has already been
stopped recursively, and each owned actor's ``__cleanup__`` has already
run. Owned actor meshes and proc meshes are no longer usable from this
method. For shutdown work that needs an owned mesh, expose a dedicated
endpoint and call it before ``stop()``.
Use ``__cleanup__`` to release resources the actor owns directly: open
files, network connections, background threads, asyncio tasks, and the
like -- not other actors or procs.
Overrides may be declared with either ``def`` or ``async def``; the
async-ness must match the actor's endpoints. Actors with sync endpoints
require a sync ``__cleanup__``; actors with async endpoints require an
async ``__cleanup__``. An ``async def`` override is awaited on the
actor's asyncio event loop -- the same loop that runs endpoint
coroutines -- so it can ``await`` other endpoints or I/O. A sync
override runs under ``fake_sync_state`` and cannot call
``asyncio.get_running_loop``. If the override raises, the exception
becomes a supervision event and will notify the owner.
This method is documentation-only on ``Actor``; subclasses provide the
real implementation.
"""
...
class ActorMesh(MeshTrait, Generic[T]):
"""
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.
"""
def __init__(
self,
Class: Type[T],
name: str,
inner: "ActorMeshProtocol",
shape: Shape,
proc_mesh: "Optional[ProcMesh]",
) -> None:
# Class name of the actor.
self.__name__: str = Class.__name__
# The name user gives when spawning the mesh
self._mesh_name = name
self._class: Type[T] = Class
self._inner: "ActorMeshProtocol" = inner
self._shape = shape
self._proc_mesh = proc_mesh
async_endpoints = []
sync_endpoints = []
async_cleanup = None
for attr_name in dir(self._class):
attr_value = getattr(self._class, attr_name, None)
if isinstance(attr_value, EndpointProperty):
# The ActorMesh builtin methods may clash with user-defined endpoints,
# check for this and raise an explainable error.
existing_attr = getattr(self, attr_name, None)
if not isinstance(existing_attr, NotAnEndpoint):
raise ValueError(
"ActorMesh has an attribute that collides with an endpoint name: "
f"attribute={attr_name}, ActorMesh.{attr_name}={getattr(self, attr_name)}, endpoint={attr_value}"
)
# Convert string method name to appropriate MethodSpecifier
kind = (
MethodSpecifier.ExplicitPort
if attr_value._explicit_response_port
else MethodSpecifier.ReturnsResponse
)
setattr(
self,
attr_name,
self._endpoint(
# pyrefly: ignore [bad-argument-count, bad-argument-type]
kind(attr_name),
attr_value._method,
attr_value._propagator,
),
)
if inspect.iscoroutinefunction(attr_value._method):
async_endpoints.append(attr_name)
else:
sync_endpoints.append(attr_name)
if attr_name == "__cleanup__" and _is_user_override(attr_value):
async_cleanup = inspect.iscoroutinefunction(attr_value)
if sync_endpoints and async_endpoints:
raise ValueError(
f"{self._class} mixes both async and sync endpoints."
"Synchronous endpoints cannot be mixed with async endpoints because they can cause the asyncio loop to deadlock if they wait."
f"sync: {sync_endpoints} async: {async_endpoints}"
)
if sync_endpoints and async_cleanup:
raise ValueError(
f"{self._class} has sync endpoints, but an async __cleanup__. Make sure __cleanup__ is also synchronous."
"Synchronous endpoints cannot be mixed with async endpoints because they can cause the asyncio loop to deadlock if they wait."
f"sync: {sync_endpoints}"
)
# Check for False explicitly because None means there is no override.
if async_endpoints and async_cleanup is False:
raise ValueError(
f"{self._class} has async endpoints, but a synchronous __cleanup__. Make sure __cleanup__ is also async."
"Synchronous endpoints cannot be mixed with async endpoints because they can cause the asyncio loop to deadlock if they wait."
f"async: {async_endpoints}"
)
def __getattr__(self, attr: str) -> NotAnEndpoint:
if attr in dir(self._class):
return NotAnEndpoint(self, attr)
raise AttributeError(attr)
def _endpoint(
self,
name: MethodSpecifier,
impl: Callable[Concatenate[Any, P], Awaitable[R]],
propagator: Propagator,
) -> Any:
return ActorEndpoint(
self._inner,
name,
self._shape,
self._mesh_name,
inspect.signature(impl),
self._proc_mesh,
propagator,
)
def __reduce_ex__(
self, protocol: Any
) -> "Tuple[Type[ActorMesh[T]], Tuple[Any, ...]]":
return ActorMesh, (
self._class,
self._mesh_name,
self._inner,
self._shape,
self._proc_mesh,
)
@property
def _ndslice(self) -> NDSlice:
return self._shape.ndslice
@property
# pyrefly: ignore [bad-override]
def _labels(self) -> Iterable[str]:
return self._shape.labels
def _new_with_shape(self, shape: Shape) -> "ActorMesh[T]":
sliced = self._inner.new_with_region(shape.region)
return ActorMesh(self._class, self._mesh_name, sliced, shape, self._proc_mesh)
def __repr__(self) -> str:
return f"ActorMesh(class={self._class}, shape={self._shape}), inner={type(self._inner)})"
def stop(self, reason: str = "stopped by client") -> "Future[None]":
instance = context().actor_instance._as_rust()
return Future(coro=self._inner.stop(instance, reason))
@property
def initialized(self) -> Future[None]:
return Future(coro=self._inner.initialized())
@property
def _name(self) -> Future[str]:
"""Retrieves the name stored in the ActorMesh internally."""
# Not called "name" to avoid clashing with a common endpoint name.
return Future(coro=self._inner.name())
[docs]class ActorError(Exception):
"""
Deterministic problem with the user's code.
For example, an OOM resulting in trying to allocate too much GPU memory, or violating
some invariant enforced by the various APIs.
"""
[docs] def __init__(
self,
exception: BaseException,
message: str = "A remote actor call has failed.",
) -> None:
self.exception = exception
# 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: str = "".join(actor_mesh_ref_tb)
self.message = message
def __str__(self) -> str:
return f"{self.message}\n {self.exception_formatted}"
[docs]def current_actor_name() -> str:
"""Return the actor id of the currently executing actor as a string."""
return str(context().actor_instance.actor_id)
[docs]def current_rank() -> Point:
"""Return the current message's position within its mesh as a ``Point``."""
return context().message_rank
[docs]def current_size() -> Dict[str, int]:
"""Return the size of each mesh dimension for the current message.
The result maps each dimension label to the number of actors along it.
"""
r = context().message_rank.extent
return {k: r[k] for k in r}
class RootClientActor(Actor):
name: str = "client"
def __supervise__(self, failure: MeshFailure) -> bool:
import os
import socket
import sys
from datetime import datetime
from monarch._src.actor.supervision import UnhandledFaultHookException
from monarch.actor import unhandled_fault_hook # pyre-ignore
try:
unhandled_fault_hook(failure)
except BaseException as e: # noqa: B036 - catch SystemExit from sys.exit; re-raised wrapped
pid = os.getpid()
hostname = socket.gethostname()
report = failure.report()
message = (
f"Unhandled monarch error on the root actor, "
f"hostname={hostname}, PID={pid} at time {datetime.now()}:\n"
f"{report}\n"
)
sys.stderr.write(message)
sys.stderr.flush()
from monarch._rust_bindings.monarch_hyperactor.telemetry import (
instant_event,
)
instant_event(message)
raise UnhandledFaultHookException(repr(e)) from e
return True
@staticmethod
def _pickled_init_args() -> FrozenBuffer:
args = (
ActorInitArgs(RootClientActor, None, None, RootClientActor.name, None, ()),
)
kwargs = {}
state = pickle(
(args, kwargs),
allow_tensor_engine_references=False,
)
return state.buffer()
