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

monarch/mesh_controller.py

@@ -7,40 +7,66 @@
 import atexit
 import logging
 import os
-import time
+
+import pdb  # noqa
 import traceback
 from collections import deque
 from logging import Logger
-from typing import List, NamedTuple, Optional, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    cast,
+    List,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 import torch.utils._python_dispatch
-
-from monarch import NDSlice
-from monarch._rust_bindings.monarch_extension import client, debugger
+from monarch._rust_bindings.monarch_extension import client
 from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monarch/monarch_extension:monarch_extension
     WorldState,
 )
 from monarch._rust_bindings.monarch_extension.mesh_controller import _Controller
+from monarch._rust_bindings.monarch_extension.tensor_worker import Ref
+from monarch._rust_bindings.monarch_hyperactor.actor import (
+    PythonMessage,
+    PythonMessageKind,
+    UnflattenArg,
+)
+from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox
 from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
     ActorId,
 )
+from monarch._src.actor.actor_mesh import ActorEndpoint, Port, PortTuple
+from monarch._src.actor.endpoint import Selection
+from monarch._src.actor.shape import NDSlice
+from monarch.common import device_mesh, messages, stream
+from monarch.common.controller_api import TController
+from monarch.common.function import ResolvableFunction
+from monarch.common.invocation import Seq
+from monarch.common.messages import Referenceable, SendResultOfActorCall
+from monarch.common.stream import StreamRef
+from monarch.common.tensor import dtensor_check, InputChecker, Tensor
+from monarch.common.tree import flatten
+from monarch.tensor_worker_main import _set_trace
 
 if TYPE_CHECKING:
     from monarch._rust_bindings.monarch_hyperactor.proc_mesh import (
         ProcMesh as HyProcMesh,
     )
-    from monarch.proc_mesh import ProcMesh
+    from monarch.actor import ProcMesh
 
 from monarch._rust_bindings.monarch_hyperactor.shape import Point
 
-from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction
 from monarch.common.client import Client
 from monarch.common.controller_api import LogMessage, MessageResult
-from monarch.common.device_mesh import DeviceMesh, no_mesh
+from monarch.common.device_mesh import DeviceMesh
+from monarch.common.future import Future as OldFuture
 from monarch.common.invocation import DeviceException, RemoteException
-from monarch.controller.debugger import read as debugger_read, write as debugger_write
 from monarch.rust_local_mesh import _get_worker_exec_info
-from pyre_extensions import none_throws
 
 logger: Logger = logging.getLogger(__name__)
 
@@ -48,6 +74,7 @@ logger: Logger = logging.getLogger(__name__)
 class Controller(_Controller):
     def __init__(self, workers: "HyProcMesh") -> None:
         super().__init__()
+        self._mailbox: Mailbox = workers.client
         # Buffer for messages unrelated to debugging that are received while a
         # debugger session is active.
         self._non_debugger_pending_messages: deque[
@@ -58,19 +85,9 @@ class Controller(_Controller):
     def next_message(
         self, timeout: Optional[float]
     ) -> Optional[LogMessage | MessageResult]:
-        if self._non_debugger_pending_messages:
-            msg = self._non_debugger_pending_messages.popleft()
-        else:
-            msg = self._get_next_message(timeout_msec=int((timeout or 0.0) * 1000.0))
-        if msg is None:
-            return None
-
-        if isinstance(msg, client.WorkerResponse):
-            return _worker_response_to_result(msg)
-        elif isinstance(msg, client.LogMessage):
-            return LogMessage(msg.level, msg.message)
-        elif isinstance(msg, client.DebuggerMessage):
-            self._run_debugger_loop(msg)
+        raise RuntimeError(
+            "internal error: tensor engine does not produce futures that call next_message"
+        )
 
     def send(
         self,
@@ -86,56 +103,6 @@ class Controller(_Controller):
         self._drain_and_stop()
         return []
 
-    def _run_debugger_loop(self, message: client.DebuggerMessage) -> None:
-        if not isinstance(message.action, DebuggerAction.Paused):
-            raise RuntimeError(
-                f"Unexpected debugger message {message} when no debugger session is running"
-            )
-
-        self._pending_debugger_sessions.append(message.debugger_actor_id)
-        while self._pending_debugger_sessions:
-            debugger_actor_id = self._pending_debugger_sessions.popleft()
-            rank = debugger_actor_id.rank
-            proc_id = debugger_actor_id.proc_id
-            debugger_write(
-                f"pdb attached to proc {proc_id} with rank {rank}, debugger actor {debugger_actor_id} \n"
-            )
-
-            self._debugger_attach(debugger_actor_id)
-            while True:
-                # TODO: Add appropriate timeout.
-                msg = self._get_next_message(timeout_msec=None)
-
-                if not isinstance(msg, client.DebuggerMessage):
-                    self._non_debugger_pending_messages.append(msg)
-                    continue
-
-                if msg.debugger_actor_id != debugger_actor_id:
-                    if isinstance(msg.action, DebuggerAction.Paused):
-                        self._pending_debugger_sessions.append(msg.debugger_actor_id)
-                        continue
-                    else:
-                        raise RuntimeError(
-                            f"unexpected debugger message {msg} from rank {msg.debugger_actor_id.rank} "
-                            f"when debugging rank {debugger_actor_id.rank}"
-                        )
-
-                action = msg.action
-                if isinstance(action, DebuggerAction.Detach):
-                    break
-                elif isinstance(action, DebuggerAction.Read):
-                    self._debugger_write(
-                        debugger_actor_id, debugger_read(action.requested_size)
-                    )
-                elif isinstance(action, DebuggerAction.Write):
-                    debugger_write(
-                        debugger.get_bytes_from_write_action(action).decode()
-                    )
-                else:
-                    raise RuntimeError(
-                        f"unexpected debugger message {msg} when debugging rank {debugger_actor_id.rank}"
-                    )
-
     def worker_world_state(self) -> WorldState:
         raise NotImplementedError("worker world state")
 
@@ -145,54 +112,6 @@ class Controller(_Controller):
         pass
 
 
-# TODO: Handling conversion of the response can move to a separate module over time
-# especially as we have structured error messages.
-def _worker_response_to_result(result: client.WorkerResponse) -> MessageResult:
-    if not result.is_exception():
-        # The result of the message needs to be unwrapped on a real device.
-        # Staying as a fake tensor will fail the tensor deserialization.
-        with no_mesh.activate():
-            return MessageResult(result.seq, result.result(), None)
-    exc = none_throws(result.exception())
-    if isinstance(exc, client.Error):
-        worker_frames = [
-            traceback.FrameSummary("<unknown>", None, frame)
-            for frame in exc.backtrace.split("\\n")
-        ]
-        return MessageResult(
-            seq=result.seq,
-            result=None,
-            error=RemoteException(
-                seq=exc.caused_by_seq,
-                exception=RuntimeError(exc.backtrace),
-                controller_frame_index=0,  # TODO: T225205291 fix this once we have recording support in rust
-                controller_frames=None,
-                worker_frames=worker_frames,
-                source_actor_id=exc.actor_id,
-                message=f"Remote function in {exc.actor_id} errored.",
-            ),
-        )
-    elif isinstance(exc, client.Failure):
-        frames = [
-            traceback.FrameSummary("<unknown>", None, frame)
-            for frame in exc.backtrace.split("\n")
-        ]
-        reason = f"Actor {exc.actor_id} crashed on {exc.address}, check the host log for details"
-        logger.error(reason)
-        return MessageResult(
-            seq=0,  # seq is not consumed for DeviceException; it will be directly thrown by the client
-            result=None,
-            error=DeviceException(
-                exception=RuntimeError(reason),
-                frames=frames,
-                source_actor_id=exc.actor_id,
-                message=reason,
-            ),
-        )
-    else:
-        raise RuntimeError(f"Unknown exception type: {type(exc)}")
-
-
 def _initialize_env(worker_point: Point, proc_id: str) -> None:
     worker_rank = worker_point.rank
     try:
@@ -213,12 +132,50 @@ def _initialize_env(worker_point: Point, proc_id: str) -> None:
             "LOCAL_WORLD_SIZE": str(gpus_per_host),
         }
         os.environ.update(process_env)
+        pdb.set_trace = _set_trace
+        # workaround for set_manual_seed somehow not working if cuda is not initialized\
+        if torch.cuda.is_available():
+            torch.cuda.init()
     except Exception:
         traceback.print_exc()
         raise
 
 
 class MeshClient(Client):
+    def fetch(
+        self,
+        mesh: "DeviceMesh",
+        stream: "StreamRef",
+        shard,
+        preprocess_message,
+        args,
+        kwargs,
+        defs: Tuple["Tensor", ...],
+        uses: Tuple["Tensor", ...],
+    ) -> "OldFuture":  # the OldFuture is a lie
+        sender, receiver = PortTuple.create(self._mesh_controller._mailbox, once=True)
+
+        ident = self.new_node(defs, uses, cast("OldFuture", sender))
+        process = mesh._process(shard)
+        self.send(
+            process,
+            messages.SendValue(
+                ident,
+                None,
+                defs,
+                preprocess_message,
+                args,
+                kwargs,
+                stream,
+            ),
+        )
+        # we have to ask for status updates
+        # from workers to be sure they have finished
+        # enough work to count this future as finished,
+        # and all potential errors have been reported
+        self._request_status()
+        return cast("OldFuture", receiver.recv())
+
     def shutdown(
         self,
         destroy_pg: bool = True,
@@ -232,27 +189,43 @@ class MeshClient(Client):
         atexit.unregister(self._atexit)
         self._shutdown = True
 
-        # ensure all pending work is finished.
-        # all errors must be messaged back at this point
-        self.new_node_nocoalesce([], [], None, [])
-        self._request_status()
-
-        ttl = 60
-        start_time = time.time()
-        end_time = start_time + ttl
-        while ttl > 0 and self.last_assigned_seq > self.last_processed_seq:
-            ttl = end_time - time.time()
-            self.handle_next_message(ttl)
-            if self._pending_shutdown_error:
-                raise self._pending_shutdown_error
-
-        if ttl <= 0:
-            raise RuntimeError("shutdown timed out")
-
+        sender, receiver = PortTuple.create(self._mesh_controller._mailbox, once=True)
+        assert sender._port_ref is not None
+        self._mesh_controller.sync_at_exit(sender._port_ref.port_id)
+        receiver.recv().get(timeout=60)
         # we are not expecting anything more now, because we already
         # waited for the responses
         self.inner.drain_and_stop()
 
+    @property
+    def _mesh_controller(self) -> Controller:
+        return cast(Controller, self.inner)
+
+    def new_node_nocoalesce(
+        self,
+        defs: Sequence["Tensor"],
+        uses: Sequence["Tensor"],
+        future: Optional["OldFuture"],
+        tracebacks: List[List[traceback.FrameSummary]],
+    ) -> Seq:
+        seq = self._next_seq()
+        for d in defs:
+            d._seq = seq
+        response_port = None
+        if future is not None:
+            # method annotation is a lie to make Client happy
+            port, slice = cast("Tuple[Port[Any], NDSlice]", future)
+            assert port._port_ref is not None
+            response_port = (port._port_ref.port_id, slice)
+        self._mesh_controller.node(seq, defs, uses, response_port, tracebacks)
+        return seq
+
+    def handle_next_message(self, timeout: Optional[float]) -> bool:
+        """
+        Mesh controller message loop is handled by the tokio event loop.
+        """
+        return False
+
 
 def spawn_tensor_engine(proc_mesh: "ProcMesh") -> DeviceMesh:
     # This argument to Controller
@@ -260,7 +233,7 @@ def spawn_tensor_engine(proc_mesh: "ProcMesh") -> DeviceMesh:
     # report the proc ID instead of the rank it currently does.
     gpus = proc_mesh.sizes.get("gpus", 1)
     backend_ctrl = Controller(proc_mesh._proc_mesh)
-    client = MeshClient(backend_ctrl, proc_mesh.size(), gpus)
+    client = MeshClient(cast("TController", backend_ctrl), proc_mesh.size(), gpus)
     dm = DeviceMesh(
         client,
         NDSlice.new_row_major(list(proc_mesh.sizes.values())),
@@ -268,3 +241,154 @@ def spawn_tensor_engine(proc_mesh: "ProcMesh") -> DeviceMesh:
     )
     dm.exit = lambda: client.shutdown()
     return dm
+
+
+class RemoteException(Exception):
+    def __init__(
+        self,
+        worker_error_string: str,  # this should really be an exception + stacktrace but
+        # worker code needs major refactor to make this possible
+        controller_frames: List[traceback.FrameSummary],
+        rank: int,
+    ):
+        self.worker_error_string = worker_error_string
+        self.controller_frames = controller_frames
+        self.rank = rank
+
+    def __str__(self):
+        try:
+            controller_tb = "".join(traceback.format_list(self.controller_frames))
+            return (
+                f"A remote function has failed asynchronously on rank {self.rank}.\n"
+                f"Traceback of where the remote function was issued on controller (most recent call last):\n{controller_tb}"
+                f"Error as reported from worker:\n{self.worker_error_string}"
+            )
+        except Exception:
+            traceback.print_exc()
+            return "<exception formatting RemoteException>"
+
+
+def _cast_call_method_indirect(
+    endpoint: ActorEndpoint,
+    selection: Selection,
+    client: MeshClient,
+    seq: Seq,
+    args_kwargs_tuple: bytes,
+    refs: Sequence[Any],
+) -> Tuple[str, int]:
+    unflatten_args = [
+        UnflattenArg.PyObject if isinstance(ref, Tensor) else UnflattenArg.Mailbox
+        for ref in refs
+    ]
+    broker_id: Tuple[str, int] = client._mesh_controller.broker_id
+    actor_msg = PythonMessage(
+        PythonMessageKind.CallMethodIndirect(
+            endpoint._name, broker_id, seq, unflatten_args
+        ),
+        args_kwargs_tuple,
+    )
+    endpoint._actor_mesh.cast(actor_msg, selection)
+    return broker_id
+
+
+def actor_send(
+    endpoint: ActorEndpoint,
+    args_kwargs_tuple: bytes,
+    refs: Sequence[Any],
+    port: Optional[Port[Any]],
+    selection: Selection,
+):
+    tensors = [ref for ref in refs if isinstance(ref, Tensor)]
+    # we have some monarch references, we need to ensure their
+    # proc_mesh matches that of the tensors we sent to it
+    chosen_stream = stream._active
+    for t in tensors:
+        if hasattr(t, "stream"):
+            chosen_stream = t.stream
+            break
+    with InputChecker(tensors, lambda x: f"actor_call({x})") as checker:
+        checker.check_mesh_stream_local(device_mesh._active, chosen_stream)
+        # TODO: move propagators into Endpoint abstraction and run the propagator to get the
+        # mutates
+        checker.check_permission(())
+    selected_device_mesh = (
+        endpoint._actor_mesh._proc_mesh and endpoint._actor_mesh._proc_mesh._device_mesh
+    )
+    if selected_device_mesh is not checker.mesh:
+        raise ValueError(
+            f"monarch Tensors sent to an actor must be located on the same process as the actor. However {checker.mesh} is not {selected_device_mesh}."
+            "NYI: better serialization of mesh names to make the mismatch more clear."
+        )
+
+    client = cast(MeshClient, checker.mesh.client)
+
+    stream_ref = chosen_stream._to_ref(client)
+
+    fut = (port, checker.mesh._ndslice) if port is not None else None
+
+    ident = client.new_node([], tensors, cast("OldFuture", fut))
+
+    # To ensure that both the actor and the stream execute in order, we send a message
+    # to each at this point. The message to the worker will be handled on the stream actor where
+    # it will send the 'tensor's to the broker actor locally, along with a response port with the
+    # computed value.
+
+    # The message to the generic actor tells it to first wait on the broker to get the local arguments
+    # from the stream, then it will run the actor method, and send the result to response port.
+
+    broker_id = _cast_call_method_indirect(
+        endpoint, selection, client, ident, args_kwargs_tuple, refs
+    )
+    worker_msg = SendResultOfActorCall(ident, broker_id, tensors, [], stream_ref)
+    client.send(checker.mesh._ndslice, worker_msg)
+    # we have to ask for status updates
+    # from workers to be sure they have finished
+    # enough work to count this future as finished,
+    # and all potential errors have been reported
+    client._request_status()
+
+
+def actor_rref(endpoint, args_kwargs_tuple: bytes, refs: Sequence[Any]):
+    chosen_stream = stream._active
+    fake_result, dtensors, mutates, mesh = dtensor_check(
+        endpoint._propagate,
+        cast(ResolvableFunction, endpoint._name),
+        refs,
+        {},
+        device_mesh._active,
+        chosen_stream,
+    )
+    assert mesh is not None
+
+    fake_result_dtensors, unflatten_result = flatten(
+        fake_result, lambda x: isinstance(x, torch.Tensor)
+    )
+    result_dtensors = tuple(
+        Tensor(fake, mesh, chosen_stream) for fake in fake_result_dtensors
+    )
+    seq = mesh.client.new_node(result_dtensors + mutates, dtensors)
+    assert all(t.ref is not None for t in result_dtensors)
+    assert all(t.ref is not None for t in mutates)
+    result = result_msg = unflatten_result(result_dtensors)
+    if len(result_dtensors) == 0:
+        result_msg = None
+
+    broker_id = _cast_call_method_indirect(
+        endpoint, "all", mesh.client, seq, args_kwargs_tuple, refs
+    )
+    # note the device mesh has to be defined regardles so the remote functions
+    # can invoke mesh.rank("...")
+
+    mesh.define_remotely()
+
+    mesh._send(
+        messages.CallActorMethod(
+            seq,
+            result_msg,
+            broker_id,
+            refs,
+            cast("List[Ref]", mutates),
+            stream._active._to_ref(mesh.client),
+        )
+    )
+    return result

monarch/opaque_module.py

@@ -9,7 +9,7 @@ from typing import List
 import torch
 from monarch.common.function_caching import TensorGroup, TensorGroupPattern
 from monarch.common.opaque_ref import OpaqueRef
-from monarch.common.remote import remote
+from monarch.common.remote import call_on_shard_and_fetch, remote
 from monarch.common.tensor_factory import TensorFactory
 from monarch.common.tree import flatten
 from monarch.opaque_object import _fresh_opaque_ref, OpaqueObject
@@ -144,11 +144,9 @@ class OpaqueModule:
 
     def parameters(self):
         if self._parameters is None:
-            tensor_group_pattern = (
-                remote(_get_parameters_shape)
-                .call_on_shard_and_fetch(self._object)
-                .result()
-            )
+            tensor_group_pattern = call_on_shard_and_fetch(
+                remote(_get_parameters_shape), self._object
+            ).result()
             self._parameters = [
                 p.requires_grad_(True)
                 for p in remote(

monarch/opaque_object.py

@@ -14,7 +14,7 @@ from monarch.common.function import (
 )
 
 from monarch.common.opaque_ref import OpaqueRef
-from monarch.common.remote import remote
+from monarch.common.remote import call_on_shard_and_fetch, remote
 
 
 def _invoke_method(obj: OpaqueRef, method_name: str, *args, **kwargs):
@@ -83,6 +83,6 @@ class OpaqueObject(OpaqueRef):
         return endpoint(self, method_name, *args, **kwargs)
 
     def call_method_on_shard_and_fetch(self, method_name, *args, **kwargs):
-        return remote(_invoke_method).call_on_shard_and_fetch(
-            self, method_name, *args, **kwargs
+        return call_on_shard_and_fetch(
+            remote(_invoke_method), self, method_name, *args, **kwargs
         )