torchmonarch-nightly 2025.7.1__cp311-cp311-manylinux2014_x86_64.whl → 2025.7.25__cp311-cp311-manylinux2014_x86_64.whl

monarch/_src/actor/endpoint.py

@@ -0,0 +1,270 @@
+# 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-unsafe
+
+import functools
+from abc import ABC, abstractmethod
+from operator import mul
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    cast,
+    Concatenate,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    overload,
+    ParamSpec,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
+    TypeVar,
+)
+
+from monarch._src.actor.future import Future
+from monarch._src.actor.tensor_engine_shim import _cached_propagation, fake_call
+
+if TYPE_CHECKING:
+    from monarch._src.actor.actor_mesh import (
+        HyPortReceiver,
+        OncePortReceiver,
+        Port,
+        PortTuple,
+        ValueMesh,
+    )
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+Selection = Literal["all", "choose"] | int
+
+
+class Extent:
+    def __init__(self, labels: Sequence[str], sizes: Sequence[int]) -> None:
+        self.labels = labels
+        self.sizes = sizes
+
+    @property
+    def nelements(self) -> int:
+        return functools.reduce(mul, self.sizes, 1)
+
+    def __str__(self) -> str:
+        return str(dict(zip(self.labels, self.sizes)))
+
+
+Propagator = Any
+
+
+class Endpoint(ABC, Generic[P, R]):
+    def __init__(self, propagator: Propagator) -> None:
+        self._propagator_arg = propagator
+        self._cache: Optional[dict] = None
+
+    @abstractmethod
+    def _send(
+        self,
+        args: Tuple[Any, ...],
+        kwargs: Dict[str, Any],
+        port: "Optional[Port]" = None,
+        selection: Selection = "all",
+    ) -> Extent:
+        """
+        Implements sending a message to the endpoint. The return value of the endpoint will
+        be sent to port if provided. If port is not provided, the return will be dropped,
+        and any exception will cause the actor to fail.
+
+        The return value is the (multi-dimension) size of the actors that were sent a message.
+        For ActorEndpoints this will be the actor_meshes size. For free-function endpoints,
+        this will be the size of the currently active proc_mesh.
+        """
+        pass
+
+    @abstractmethod
+    def _port(self, once: bool = False) -> "PortTuple[R]":
+        pass
+
+    @abstractmethod
+    def _call_name(self) -> Any:
+        """
+        Something to use in InputChecker to represent calling this thingy.
+        """
+        pass
+
+    def _supervise(self, r: "HyPortReceiver | OncePortReceiver") -> Any:
+        return r
+
+    # the following are all 'adverbs' or different ways to handle the
+    # return values of this endpoint. Adverbs should only ever take *args, **kwargs
+    # of the original call. If we want to add syntax sugar for something that needs additional
+    # arguments, it should be implemented as function indepdendent of endpoint like `send`
+    # and `Accumulator`
+    def choose(self, *args: P.args, **kwargs: P.kwargs) -> Future[R]:
+        """
+        Load balanced sends a message to one chosen actor and awaits a result.
+
+        Load balanced RPC-style entrypoint for request/response messaging.
+        """
+        from monarch._src.actor.actor_mesh import port
+
+        p, r = port(self, once=True)
+        # pyre-ignore
+        self._send(args, kwargs, port=p, selection="choose")
+        return r.recv()
+
+    def call_one(self, *args: P.args, **kwargs: P.kwargs) -> Future[R]:
+        from monarch._src.actor.actor_mesh import port
+
+        p, r = port(self, once=True)
+        # pyre-ignore
+        extent = self._send(args, kwargs, port=p, selection="choose")
+        if extent.nelements != 1:
+            raise ValueError(
+                f"Can only use 'call_one' on a single Actor but this actor has shape {extent}"
+            )
+        return r.recv()
+
+    def call(self, *args: P.args, **kwargs: P.kwargs) -> "Future[ValueMesh[R]]":
+        from monarch._src.actor.actor_mesh import ranked_port, ValueMesh
+
+        p, r = ranked_port(self)
+        # pyre-ignore
+        extent = self._send(args, kwargs, port=p)
+
+        async def process() -> "ValueMesh[R]":
+            from monarch._rust_bindings.monarch_hyperactor.shape import Shape
+            from monarch._src.actor.shape import NDSlice
+
+            results: List[R] = [None] * extent.nelements  # pyre-fixme[9]
+            for _ in range(extent.nelements):
+                rank, value = await r.recv()
+                results[rank] = value
+            call_shape = Shape(
+                extent.labels,
+                NDSlice.new_row_major(extent.sizes),
+            )
+            return ValueMesh(call_shape, results)
+
+        return Future(impl=process, requires_loop=False)
+
+    async def stream(self, *args: P.args, **kwargs: P.kwargs) -> AsyncGenerator[R, R]:
+        """
+        Broadcasts to all actors and yields their responses as a stream / generator.
+
+        This enables processing results from multiple actors incrementally as
+        they become available. Returns an async generator of response values.
+        """
+        from monarch._src.actor.actor_mesh import port
+
+        p, r = port(self)
+        # pyre-ignore
+        extent = self._send(args, kwargs, port=p)
+        for _ in range(extent.nelements):
+            yield await r.recv()
+
+    def broadcast(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        """
+        Fire-and-forget broadcast to all actors without waiting for actors to
+        acknowledge receipt.
+
+        In other words, the return of this method does not guarrantee the
+        delivery of the message.
+        """
+        from monarch._src.actor.actor_mesh import send
+
+        # pyre-ignore
+        send(self, args, kwargs)
+
+    def _propagate(self, args, kwargs, fake_args, fake_kwargs):
+        if self._propagator_arg is None or self._propagator_arg == "cached":
+            if self._cache is None:
+                self._cache = {}
+            return _cached_propagation(self._cache, self._resolvable, args, kwargs)
+        elif self._propagator_arg == "inspect":
+            return None
+        elif self._propagator_arg == "mocked":
+            raise NotImplementedError("mocked propagation")
+        else:
+            return fake_call(self._propagator_arg, *fake_args, **fake_kwargs)
+
+    def _fetch_propagate(self, args, kwargs, fake_args, fake_kwargs):
+        if self._propagator_arg is None:
+            return  # no propgator provided, so we just assume no mutations
+        return self._propagate(args, kwargs, fake_args, fake_kwargs)
+
+    def _pipe_propagate(self, args, kwargs, fake_args, fake_kwargs):
+        if not callable(self._propagator_arg):
+            raise ValueError("Must specify explicit callable for pipe")
+        return self._propagate(args, kwargs, fake_args, fake_kwargs)
+
+
+class EndpointProperty(Generic[P, R]):
+    @overload
+    def __init__(
+        self,
+        method: Callable[Concatenate[Any, P], Awaitable[R]],
+        propagator: Propagator,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self, method: Callable[Concatenate[Any, P], R], propagator: Propagator
+    ) -> None: ...
+
+    def __init__(self, method: Any, propagator: Propagator) -> None:
+        self._method = method
+        self._propagator = propagator
+
+    def __get__(self, instance, owner) -> Endpoint[P, R]:
+        # this is a total lie, but we have to actually
+        # recognize this was defined as an endpoint,
+        # and also lookup the method
+        return cast(Endpoint[P, R], self)
+
+
+# This can't just be Callable because otherwise we are not
+# allowed to use type arguments in the return value.
+class EndpointIfy:
+    @overload
+    def __call__(self, function: Callable[P, Awaitable[R]]) -> Endpoint[P, R]: ...
+    @overload
+    def __call__(self, function: Callable[P, R]) -> Endpoint[P, R]: ...
+
+    def __call__(self, function: Any):
+        pass
+
+
+@overload
+def endpoint(
+    method: Callable[Concatenate[Any, P], Awaitable[R]],
+    *,
+    propagate: Propagator = None,
+) -> EndpointProperty[P, R]: ...
+
+
+@overload
+def endpoint(
+    method: Callable[Concatenate[Any, P], R],
+    *,
+    propagate: Propagator = None,
+) -> EndpointProperty[P, R]: ...
+
+
+@overload
+def endpoint(
+    *,
+    propagate: Propagator = None,
+) -> EndpointIfy: ...
+
+
+def endpoint(method=None, *, propagate=None):
+    if method is None:
+        return functools.partial(endpoint, propagate=propagate)
+    return EndpointProperty(method, propagator=propagate)

monarch/_src/actor/event_loop.py

@@ -0,0 +1,97 @@
+# 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.
+
+"""
+Module for managing the event loop used by Monarch Python actors.
+This provides a way to create a Python-aware thread from Rust that runs the worker event loop.
+"""
+
+import asyncio
+import logging
+import threading
+from typing import Optional
+
+from libfb.py.pyre import none_throws
+
+logger = logging.getLogger(__name__)
+
+_event_loop: Optional[asyncio.AbstractEventLoop] = None
+_lock = threading.Lock()
+_ready = threading.Event()
+
+
+def _initialize_event_loop() -> None:
+    """
+    Internal function to initialize the event loop.
+    This creates a new thread with an event loop that runs forever.
+    """
+    global _event_loop, _ready
+    if _event_loop is not None:
+        return
+
+    # Create a new thread that will host our event loop
+    def event_loop_thread():
+        """Target function for the event loop thread."""
+        global _event_loop, _ready
+        try:
+            # Create a new event loop
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            _event_loop = loop
+            _ready.set()
+
+            logger.debug(
+                f"Python worker event loop thread started: {threading.current_thread().name}"
+            )
+            try:
+                # Run the event loop forever
+                loop.run_forever()
+            finally:
+                # Clean up when the loop stops
+                logger.debug("Python worker event loop stopped, closing...")
+                loop.close()
+        except Exception as e:
+            logger.error(f"Error in event loop thread: {e}")
+            _ready.set()
+            raise
+
+    # Create and start the thread
+    threading.Thread(
+        target=event_loop_thread,
+        name="asyncio-event-loop",
+        daemon=True,  # Make it a daemon thread so it doesn't block program exit
+    ).start()
+
+    _ready.wait()  # Wait for the event loop to be ready
+
+    if _event_loop is None:
+        raise RuntimeError("Failed to initialize event loop")
+
+
+def get_event_loop() -> asyncio.AbstractEventLoop:
+    """
+    Get the Python worker event loop.
+    If no event loop is currently running, this will start a new one.
+
+    Expected to be called from rust code.
+    """
+    global _event_loop
+    if _event_loop is None:
+        with _lock:
+            _initialize_event_loop()
+    return none_throws(_event_loop)
+
+
+def stop_event_loop() -> None:
+    """
+    Stop the event loop gracefully.
+    """
+    global _event_loop
+    if _event_loop is not None:
+        logger.debug("Stopping event loop...")
+        event_loop = none_throws(_event_loop)
+        event_loop.call_soon_threadsafe(event_loop.stop)

monarch/_src/actor/future.py

@@ -0,0 +1,100 @@
+# 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.
+
+import asyncio
+import traceback
+from functools import partial
+from typing import Generator, Generic, Optional, TypeVar
+
+R = TypeVar("R")
+
+
+async def _aincomplete(impl, self):
+    try:
+        return self._set_result(await impl())
+    except Exception as e:
+        self._set_exception(e)
+        raise
+
+
+# Future is our generic mechanism for providing both a synchronous and asynchronous API for
+# Monarch Future objects.
+
+# We treat all code as running in one of two contexts: synchronous (asyncio._get_running_loop() is None)
+# or asynchronous.
+
+# Inside of asynchronous code, clients of our API must use `await` to wait for monarch Futures to prevent
+# blocking the surrounding event loop.
+
+# In synchronous code users must call get() because the call is comming from an non-async function so
+# await is not allowed.
+
+# [avoiding async code duplication]
+# Because we allow for two modes, it is tempting as developers of Monarch to start to write two copies of
+# of code for each mode. However, this results in a lot of confusing code duplication.
+# To avoid this, we utilize the fact that synchronous code is allowed to start/complete an asyncio event loop
+# via asyncio.run in order to complete the `get()` operation. So we can just write the async version and use
+# it to implement the synchronoous version.
+
+# However, starting and running an event loop is somewhat expensive. For simple messages, using an event loop
+# is about 4x slower than just directly waiting on the tokio result. To avoid this slow down we perform an
+# optimization. For any case where the `impl` coroutine of a future calls `await` only on PythonFuture
+# (a Tokio future returning a Python value) objects, we pass requires_loop=False to the Future. In this mode,
+# the future will just run the coroutine manually, and the PythonFuture object will recognize it is being awaited
+# without an event loop (search [avoiding code duplication]) and simply do a blocking wait. By avoiding the event
+# loop machinery, this gives it the same throughput as if we ran it synchronously.
+
+
+class Future(Generic[R]):
+    def __init__(self, *, impl, requires_loop=True):
+        self._aget = partial(_aincomplete, impl)
+        self._requires_loop = requires_loop
+
+    def get(self, timeout: Optional[float] = None) -> R:
+        if asyncio._get_running_loop() is not None:
+            raise RuntimeError("get() cannot be called from within an async context")
+        if timeout is not None:
+            return asyncio.run(asyncio.wait_for(self._aget(self), timeout))
+        if not self._requires_loop:
+            try:
+                coro = self._aget(self)
+                next(coro.__await__())
+                tb_str = "".join(traceback.format_stack(coro.cr_frame))
+                raise RuntimeError(
+                    f"a coroutine paused with a future with requires_loop=False cannot block on a python asyncio.Future. Use requires_loop=True.\n{tb_str}"
+                )
+            except StopIteration as e:
+                return e.value
+        return asyncio.run(self._aget(self))
+
+    def __await__(self) -> Generator[R, None, R]:
+        return self._aget(self).__await__()
+
+    def _set_result(self, result):
+        async def af(self):
+            return result
+
+        self._aget = af
+        return result
+
+    def _set_exception(self, e):
+        async def af(self):
+            raise e
+
+        self._aget = af
+
+    # compatibility with old tensor engine Future objects
+    # hopefully we do not need done(), add_callback because
+    # they are harder to implement right.
+    def result(self, timeout: Optional[float] = None) -> R:
+        return self.get(timeout)
+
+    def exception(self, timeout: Optional[float] = None):
+        try:
+            self.get(timeout)
+            return None
+        except Exception as e:
+            return e

monarch/{pdb_wrapper.py → _src/actor/pdb_wrapper.py}

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
 import bdb
 import inspect
 import io
@@ -15,9 +16,10 @@ from dataclasses import dataclass
 from typing import Dict, TYPE_CHECKING
 
 from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
+from monarch._src.actor.sync_state import fake_sync_state
 
 if TYPE_CHECKING:
-    from monarch.debugger import DebugClient
+    from monarch._src.actor.debugger import DebugClient
 
 
 @dataclass
@@ -45,35 +47,38 @@ class PdbWrapper(pdb.Pdb):
         super().__init__(stdout=WriteWrapper(self), stdin=ReadWrapper.create(self))
         self._first = True
 
-    def setup(self, *args, **kwargs):
-        r = super().setup(*args, **kwargs)
-        if self._first:
-            self._first = False
-            # when we enter the debugger, we want to present the user's stack frame
-            # not the nested one inside session.run. This means that the local
-            # variables are what gets printed, etc. To do this
-            # we first execute up 2 to get to that frame.
-            self.do_up(2)
-        return r
-
-    def set_continue(self) -> None:
-        r = super().set_continue()
-        if not self.breaks:
-            # no more breakpoints so this debugger will not
-            # be used again, and we detach from the controller io.
-            self.client_ref.debugger_session_end.call_one(self.rank).get()
-            # break cycle with itself before we exit
-            self.stdin = sys.stdin
-            self.stdout = sys.stdout
-        return r
-
-    def set_trace(self):
-        self.client_ref.debugger_session_start.call_one(
+    def set_trace(self, frame):
+        self.client_ref.debugger_session_start.broadcast(
             self.rank, self.coords, socket.getfqdn(socket.gethostname()), self.actor_id
-        ).get()
+        )
         if self.header:
             self.message(self.header)
-        super().set_trace()
+        super().set_trace(frame)
+
+    def do_clear(self, arg):
+        if not arg:
+            # Sending `clear` without any argument specified will
+            # request confirmation from the user using the `input` function,
+            # which bypasses our ReadWrapper and causes a hang on the client.
+            # To avoid this, we just clear all breakpoints instead without
+            # confirmation.
+            super().clear_all_breaks()
+        else:
+            super().do_clear(arg)
+
+    def end_debug_session(self):
+        self.client_ref.debugger_session_end.broadcast(self.rank)
+        # Once the debug client actor is notified of the session being over,
+        # we need to prevent any additional requests being sent for the session
+        # by redirecting stdin and stdout.
+        self.stdin = sys.stdin
+        self.stdout = sys.stdout
+
+    def post_mortem(self, exc_tb):
+        self._first = False
+        # See builtin implementation of pdb.post_mortem() for reference.
+        self.reset()
+        self.interaction(None, exc_tb)
 
 
 class ReadWrapper(io.RawIOBase):
@@ -81,16 +86,19 @@ class ReadWrapper(io.RawIOBase):
         self.session = session
 
     def readinto(self, b):
-        response = self.session.client_ref.debugger_read.call_one(
-            self.session.rank, len(b)
-        ).get()
-        if response == "detach":
-            # this gets injected by the worker event loop to
-            # get the worker thread to exit on an Exit command.
-            raise bdb.BdbQuit
-        assert isinstance(response, DebuggerWrite) and len(response.payload) <= len(b)
-        b[: len(response.payload)] = response.payload
-        return len(response.payload)
+        with fake_sync_state():
+            response = self.session.client_ref.debugger_read.call_one(
+                self.session.rank, len(b)
+            ).get()
+            if response == "detach":
+                # this gets injected by the worker event loop to
+                # get the worker thread to exit on an Exit command.
+                raise bdb.BdbQuit
+            assert isinstance(response, DebuggerWrite) and len(response.payload) <= len(
+                b
+            )
+            b[: len(response.payload)] = response.payload
+            return len(response.payload)
 
     def readable(self) -> bool:
         return True
@@ -115,21 +123,14 @@ class WriteWrapper:
             function = f"{inspect.getmodulename(self.session.curframe.f_code.co_filename)}.{self.session.curframe.f_code.co_name}"
             # pyre-ignore
             lineno = self.session.curframe.f_lineno
-        self.session.client_ref.debugger_write.call_one(
+        self.session.client_ref.debugger_write.broadcast(
             self.session.rank,
             DebuggerWrite(
                 s.encode(),
                 function,
                 lineno,
             ),
-        ).get()
+        )
 
     def flush(self):
         pass
-
-
-def remote_breakpointhook(
-    rank: int, coords: Dict[str, int], actor_id: ActorId, client_ref: "DebugClient"
-):
-    ds = PdbWrapper(rank, coords, actor_id, client_ref)
-    ds.set_trace()

monarch/{common/pickle_flatten.py → _src/actor/pickle.py}

@@ -6,10 +6,16 @@
 
 import io
 import pickle
+from contextlib import contextmanager, ExitStack
 from typing import Any, Callable, Iterable, List, Tuple
 
 import cloudpickle
 
+try:
+    import torch  # @manual
+except ImportError:
+    torch = None
+
 
 class _Pickler(cloudpickle.Pickler):
     def __init__(self, filter):
@@ -44,5 +50,23 @@ def flatten(obj: Any, filter: Callable[[Any], bool]) -> Tuple[List[Any], bytes]:
 
 
 def unflatten(data: bytes, values: Iterable[Any]) -> Any:
-    up = _Unpickler(data, values)
-    return up.load()
+    with ExitStack() as stack:
+        if torch is not None:
+            stack.enter_context(load_tensors_on_cpu())
+            stack.enter_context(torch.utils._python_dispatch._disable_current_modes())
+        up = _Unpickler(data, values)
+        return up.load()
+
+
+@contextmanager
+def load_tensors_on_cpu():
+    # Ensure that any tensors load from CPU via monkeypatching how Storages are
+    # loaded.
+    old = torch.storage._load_from_bytes
+    try:
+        torch.storage._load_from_bytes = lambda b: torch.load(
+            io.BytesIO(b), map_location="cpu", weights_only=False
+        )
+        yield
+    finally:
+        torch.storage._load_from_bytes = old