torchmonarch-nightly 2025.6.11__cp310-cp310-manylinux2014_x86_64.whl → 2025.6.13__cp310-cp310-manylinux2014_x86_64.whl

monarch/debugger.py

@@ -0,0 +1,377 @@
+# 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 logging
+import sys
+from dataclasses import dataclass
+from typing import Dict, List, Tuple, Union
+
+from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
+from monarch.actor_mesh import Actor, endpoint
+
+from monarch.pdb_wrapper import DebuggerWrite
+
+from monarch.proc_mesh import local_proc_mesh
+from tabulate import tabulate
+
+
+logger = logging.getLogger(__name__)
+
+
+CANCEL_TOKEN = object()
+
+
+async def _debugger_input(prompt=""):
+    return await asyncio.to_thread(input, prompt)
+
+
+def _debugger_output(msg):
+    sys.stdout.write(msg)
+    sys.stdout.flush()
+
+
+@dataclass
+class DebugSessionInfo:
+    rank: int
+    coords: Dict[str, int]
+    hostname: str
+    actor_id: ActorId
+    function: str | None
+    lineno: int | None
+
+
+class DebugSession:
+    """Represents a single session with a remote debugger."""
+
+    def __init__(
+        self, rank: int, coords: Dict[str, int], hostname: str, actor_id: ActorId
+    ):
+        self.rank = rank
+        self.coords = coords
+        self.hostname = hostname
+        self.actor_id = actor_id
+        self._active = False
+        self._message_queue = asyncio.Queue()
+        self._task = None
+        self._pending_send_to_actor = asyncio.Queue()
+        self._outputs_since_last_input = []
+        self._function_lineno = None
+        self._need_read = False
+
+    async def _event_loop(self, line=None, suppress_output=False):
+        if not suppress_output:
+            # If the user had previously attached to this debug session,
+            # then it would have printed various messages from the
+            # message queue. When the user re-attaches, we want to
+            # print out all of the output that was printed since the
+            # last command sent to this session.
+            for output in self._outputs_since_last_input:
+                _debugger_output(output.payload.decode())
+
+        while True:
+            # When the user inputs "detach", it uses up a "read" message
+            # without actually responding to the actor being debugged. We
+            # can't manually reinsert the "read" message into the message queue,
+            # so instead the self._need_read flag indicates there's an additional
+            # "read" that we need to respond to.
+            if self._need_read:
+                self._need_read = False
+                message = "read"
+            else:
+                message = await self._message_queue.get()
+            if message == "detach":
+                # Return to the main outer debug loop.
+                break
+            elif message == "read":
+                break_after = False
+                if line is not None:
+                    break_after = True
+                else:
+                    line = await _debugger_input()
+                if line.strip("\n") == "detach":
+                    self._need_read = True
+                    break
+                else:
+                    self._outputs_since_last_input = []
+                    await self._pending_send_to_actor.put((line + "\n").encode())
+                    line = None
+                    if break_after:
+                        break
+            elif message[0] == "write":
+                output = message[1]
+                # If the user sees this output but then detaches from the session,
+                # its useful to store all outputs since the last input so that
+                # they can be printed again when the user re-attaches.
+                self._outputs_since_last_input.append(output)
+                if not suppress_output:
+                    _debugger_output(output.payload.decode())
+
+        if not suppress_output:
+            print(
+                f"Detaching from debug session for rank {self.rank} ({self.hostname})"
+            )
+
+    def get_info(self):
+        function = lineno = None
+        if self._function_lineno is not None:
+            function, lineno = self._function_lineno
+        return DebugSessionInfo(
+            self.rank, self.coords, self.hostname, self.actor_id, function, lineno
+        )
+
+    async def attach(self, line=None, suppress_output=False):
+        self._active = True
+        if not suppress_output:
+            print(f"Attached to debug session for rank {self.rank} ({self.hostname})")
+        self._task = asyncio.create_task(self._event_loop(line, suppress_output))
+        await self._task
+        if not suppress_output:
+            print(f"Detached from debug session for rank {self.rank} ({self.hostname})")
+        self._active = False
+
+    async def detach(self):
+        if self._active:
+            await self._message_queue.put("detach")
+
+    async def debugger_read(self, size: int) -> DebuggerWrite:
+        await self._message_queue.put("read")
+        input_data = await self._pending_send_to_actor.get()
+        if len(input_data) > size:
+            input_data = input_data[:size]
+        return DebuggerWrite(input_data, None, None)
+
+    async def debugger_write(self, write: DebuggerWrite) -> None:
+        if write.function is not None and write.lineno is not None:
+            self._function_lineno = (write.function, write.lineno)
+        await self._message_queue.put(("write", write))
+
+
+class DebugCommand:
+    @staticmethod
+    def parse(line: str) -> Union["DebugCommand", None]:
+        parts = line.strip("\n").split(" ")
+        if len(parts) == 0:
+            return None
+        command = parts[0]
+        match command:
+            case "attach":
+                return Attach._parse(parts)
+            case "list":
+                return ListCommand()
+            case "quit":
+                return Quit()
+            case "cast":
+                return Cast._parse(parts)
+            case "help":
+                return Help()
+            case "continue":
+                return Continue()
+            case _:
+                print(
+                    f"Unknown command {command}. Expected: attach | list | quit | cast | continue | help"
+                )
+                return None
+
+
+@dataclass
+class Attach(DebugCommand):
+    rank: int
+
+    @classmethod
+    def _parse(cls, parts: List[str]) -> "Attach":
+        if len(parts) != 2:
+            raise ValueError("Invalid attach command. Expected: attach <rank>")
+        try:
+            rank = int(parts[1])
+        except ValueError:
+            raise ValueError(f"Invalid rank {parts[1]}. Expected: int")
+        return cls(rank)
+
+
+class ListCommand(DebugCommand):
+    pass
+
+
+class Quit(DebugCommand):
+    pass
+
+
+class Help(DebugCommand):
+    pass
+
+
+class Continue(DebugCommand):
+    pass
+
+
+@dataclass
+class Cast(DebugCommand):
+    ranks: List[int] | None
+    command: str
+
+    @classmethod
+    def _parse(cls, parts: List[str]) -> "Cast":
+        if len(parts) < 3:
+            raise ValueError(
+                "Invalid cast command. Expected: cast {<r0,r1,...> | *} <command>"
+            )
+        str_ranks = parts[1]
+        command = " ".join(parts[2:])
+        if str_ranks == "*":
+            return cls(None, command)
+        else:
+            str_ranks = str_ranks.split(",")
+            if len(str_ranks) == 0:
+                raise ValueError(
+                    "Invalid rank list for cast. Expected at least one rank."
+                )
+            ranks = []
+            for rank in str_ranks:
+                try:
+                    ranks.append(int(rank))
+                except ValueError:
+                    raise ValueError(f"Invalid rank {rank}. Expected: int")
+            return cls(ranks, command)
+
+
+class DebugClient(Actor):
+    """
+    Single actor for both remote debuggers and users to talk to.
+
+    Handles multiple sessions simultanesouly
+    """
+
+    def __init__(self) -> None:
+        self.sessions = {}  # rank -> DebugSession
+
+    @endpoint
+    async def wait_pending_session(self):
+        while len(self.sessions) == 0:
+            await asyncio.sleep(1)
+
+    @endpoint
+    async def list(self) -> List[Tuple[int, Dict[str, int], str, ActorId, str, int]]:
+        table_data = []
+        for _, session in self.sessions.items():
+            info = session.get_info()
+            table_data.append(
+                (
+                    info.rank,
+                    info.coords,
+                    info.hostname,
+                    info.actor_id,
+                    info.function,
+                    info.lineno,
+                )
+            )
+        table_data = sorted(table_data, key=lambda r: r[0])
+
+        headers = ["Rank", "Coords", "Hostname", "Actor ID", "Function", "Line No."]
+        print(tabulate(table_data, headers=headers, tablefmt="grid"))
+
+        return table_data
+
+    @endpoint
+    async def enter(self) -> None:
+        # pyre-ignore
+        await getattr(self, "list")._method(self)  # noqa
+
+        while True:
+            try:
+                user_input = await _debugger_input("monarch_dbg> ")
+                command = DebugCommand.parse(user_input)
+                if isinstance(command, Help):
+                    print("monarch_dbg commands:")
+                    print("\tattach <rank> - attach to a debug session")
+                    print("\tlist - list all debug sessions")
+                    print("\tquit - exit the debugger, leaving all sessions in place")
+                    print(
+                        "\tcast {<r0,r1,...> | *} <command> - send a command to a comma-separated list of ranks, or all ranks"
+                    )
+                    print(
+                        "\tcontinue - tell all ranks to continue execution, then exit the debugger"
+                    )
+                    print("\thelp - print this help message")
+                elif isinstance(command, Attach):
+                    if command.rank not in self.sessions:
+                        print(f"No debug session for rank {command.rank}")
+                    else:
+                        await self.sessions[command.rank].attach()
+                elif isinstance(command, ListCommand):
+                    await getattr(self, "list")._method(self)  # noqa
+                elif isinstance(command, Continue):
+                    # Make sure all ranks have exited their debug sessions.
+                    # If we sent "quit", it would raise BdbQuit, crashing
+                    # the process, which probably isn't what we want.
+                    while len(self.sessions) > 0:
+                        tasks = []
+                        for rank in self.sessions:
+                            tasks.append(
+                                self.sessions[rank].attach("c", suppress_output=True)
+                            )
+                        await asyncio.gather(*tasks)
+                    return
+                elif isinstance(command, Quit):
+                    return
+                elif isinstance(command, Cast):
+                    if command.ranks is None:
+                        ranks = self.sessions.keys()
+                    else:
+                        ranks = command.ranks
+                    tasks = []
+                    for rank in ranks:
+                        if rank in self.sessions:
+                            tasks.append(
+                                self.sessions[rank].attach(
+                                    command.command,
+                                    suppress_output=True,
+                                )
+                            )
+                        else:
+                            print(f"No debug session for rank {rank}")
+                    await asyncio.gather(*tasks)
+            except Exception as e:
+                print(f"Error processing command: {e}")
+
+    ##########################################################################
+    # Debugger APIs
+    #
+    # These endpoints are called by the remote debuggers to establish sessions
+    # and communicate with them.
+    @endpoint
+    async def debugger_session_start(
+        self, rank: int, coords: Dict[str, int], hostname: str, actor_id: ActorId
+    ) -> None:
+        # Create a session if it doesn't exist
+        if rank not in self.sessions:
+            self.sessions[rank] = DebugSession(rank, coords, hostname, actor_id)
+
+    @endpoint
+    async def debugger_session_end(self, rank: int) -> None:
+        """Detach from the current debug session."""
+        session = self.sessions.pop(rank)
+        await session.detach()
+
+    @endpoint
+    async def debugger_read(self, rank: int, size: int) -> DebuggerWrite | str:
+        """Read from the debug session for the given rank."""
+        session = self.sessions[rank]
+
+        return await session.debugger_read(size)
+
+    @endpoint
+    async def debugger_write(self, rank: int, write: DebuggerWrite) -> None:
+        """Write to the debug session for the given rank."""
+        session = self.sessions[rank]
+        await session.debugger_write(write)
+
+
+async def init_debugging(actor_mesh: Actor) -> DebugClient:
+    debugger_proc_mesh = await local_proc_mesh(gpus=1, hosts=1)
+    debug_client_mesh = await debugger_proc_mesh.spawn("debug_client", DebugClient)
+    await actor_mesh._set_debug_client.call(debug_client_mesh)
+    return debug_client_mesh

monarch/mesh_controller.py

@@ -4,7 +4,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import atexit
 import logging
+import os
+import time
 import traceback
 from collections import deque
 from logging import Logger
@@ -22,6 +25,8 @@ from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarc
     ActorId,
 )
 from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh as HyProcMesh
+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
@@ -29,6 +34,7 @@ from monarch.common.device_mesh import DeviceMesh, no_mesh
 from monarch.common.invocation import DeviceException, RemoteException
 from monarch.controller.debugger import read as debugger_read, write as debugger_write
 from monarch.proc_mesh import ProcMesh
+from monarch.rust_local_mesh import _get_worker_exec_info
 from pyre_extensions import none_throws
 
 logger: Logger = logging.getLogger(__name__)
@@ -72,18 +78,8 @@ class Controller(_Controller):
     def drain_and_stop(
         self,
     ) -> List[LogMessage | MessageResult | client.DebuggerMessage]:
-        logger.info("rust controller shutting down")
-        results = []
-        for msg in self._drain_and_stop():
-            if isinstance(msg, client.WorkerResponse):
-                results.append(_worker_response_to_result(msg))
-            elif isinstance(msg, client.LogMessage):
-                results.append(LogMessage(msg.level, msg.message))
-            elif isinstance(msg, client.DebuggerMessage):
-                results.append(msg)
-            else:
-                raise RuntimeError(f"Unexpected message type {type(msg)}")
-        return results
+        self._drain_and_stop()
+        return []
 
     def _run_debugger_loop(self, message: client.DebuggerMessage) -> None:
         if not isinstance(message.action, DebuggerAction.Paused):
@@ -158,7 +154,6 @@ def _worker_response_to_result(result: client.WorkerResponse) -> MessageResult:
             traceback.FrameSummary("<unknown>", None, frame)
             for frame in exc.backtrace.split("\\n")
         ]
-        logger.error(f"Worker {exc.actor_id} failed")
         return MessageResult(
             seq=result.seq,
             result=None,
@@ -169,7 +164,7 @@ def _worker_response_to_result(result: client.WorkerResponse) -> MessageResult:
                 controller_frames=None,
                 worker_frames=worker_frames,
                 source_actor_id=exc.actor_id,
-                message=f"Worker {exc.actor_id} failed",
+                message=f"Remote function in {exc.actor_id} errored.",
             ),
         )
     elif isinstance(exc, client.Failure):
@@ -193,13 +188,75 @@ def _worker_response_to_result(result: client.WorkerResponse) -> MessageResult:
         raise RuntimeError(f"Unknown exception type: {type(exc)}")
 
 
+def _initialize_env(worker_point: Point, proc_id: str) -> None:
+    worker_rank = worker_point.rank
+    try:
+        _, worker_env = _get_worker_exec_info()
+        local_rank = worker_point["gpus"]
+        gpus_per_host = worker_point.size("gpus")
+        num_worker_procs = len(worker_point.shape)
+        process_env = {
+            **worker_env,
+            "HYPERACTOR_MANAGED_SUBPROCESS": "1",
+            "CUDA_VISIBLE_DEVICES": str(local_rank),
+            "NCCL_HOSTID": f"{proc_id}_host_{worker_rank // gpus_per_host}",
+            # This is needed to avoid a hard failure in ncclx when we do not
+            # have backend topology info (eg. on RE).
+            "NCCL_IGNORE_TOPO_LOAD_FAILURE": "true",
+            "LOCAL_RANK": str(local_rank),
+            "RANK": str(worker_rank),
+            "WORLD_SIZE": str(num_worker_procs),
+            "LOCAL_WORLD_SIZE": str(gpus_per_host),
+        }
+        os.environ.update(process_env)
+    except Exception:
+        traceback.print_exc()
+        raise
+
+
+class MeshClient(Client):
+    def shutdown(
+        self,
+        destroy_pg: bool = True,
+        error_reason: Optional[RemoteException | DeviceException | Exception] = None,
+    ):
+        # return
+        if self.has_shutdown:
+            return
+        logger.info("shutting down the client gracefully")
+
+        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")
+
+        # we are not expecting anything more now, because we already
+        # waited for the responses
+        self.inner.drain_and_stop()
+
+
 def spawn_tensor_engine(proc_mesh: ProcMesh) -> DeviceMesh:
     # This argument to Controller
     # is currently only used for debug printing. It should be fixed to
     # 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 = Client(backend_ctrl, proc_mesh.size(), gpus)
+    client = MeshClient(backend_ctrl, proc_mesh.size(), gpus)
     dm = DeviceMesh(
         client,
         NDSlice.new_row_major(list(proc_mesh.sizes.values())),

monarch/pdb_wrapper.py

@@ -0,0 +1,135 @@
+# 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 bdb
+import inspect
+import io
+import pdb  # noqa
+import socket
+import sys
+from dataclasses import dataclass
+
+from typing import Dict, TYPE_CHECKING
+
+from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
+
+if TYPE_CHECKING:
+    from monarch.debugger import DebugClient
+
+
+@dataclass
+class DebuggerWrite:
+    payload: bytes
+    function: str | None
+    lineno: int | None
+
+
+class PdbWrapper(pdb.Pdb):
+    def __init__(
+        self,
+        rank: int,
+        coords: Dict[str, int],
+        actor_id: ActorId,
+        client_ref: "DebugClient",
+        header: str | None = None,
+    ):
+        self.rank = rank
+        self.coords = coords
+        self.header = header
+        self.actor_id = actor_id
+        self.client_ref = client_ref
+        # pyre-ignore
+        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(
+            self.rank, self.coords, socket.getfqdn(socket.gethostname()), self.actor_id
+        ).get()
+        if self.header:
+            self.message(self.header)
+        super().set_trace()
+
+
+class ReadWrapper(io.RawIOBase):
+    def __init__(self, session: "PdbWrapper"):
+        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)
+
+    def readable(self) -> bool:
+        return True
+
+    @classmethod
+    def create(cls, session: "PdbWrapper"):
+        return io.TextIOWrapper(io.BufferedReader(cls(session)))
+
+
+class WriteWrapper:
+    def __init__(self, session: "PdbWrapper"):
+        self.session = session
+
+    def writable(self) -> bool:
+        return True
+
+    def write(self, s: str):
+        function = None
+        lineno = None
+        if self.session.curframe is not None:
+            # pyre-ignore
+            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.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/proc_mesh.py

@@ -4,9 +4,11 @@
 # 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 sys
 
-from typing import Any, cast, Optional, Type, TypeVar
+from typing import Any, cast, List, Optional, Type, TypeVar
 
 import monarch
 from monarch import ActorFuture as Future
@@ -18,7 +20,7 @@ from monarch._rust_bindings.hyperactor_extension.alloc import (  # @manual=//mon
 )
 from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox
 from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh as HyProcMesh
-from monarch._rust_bindings.monarch_hyperactor.shape import Shape
+from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
 from monarch.actor_mesh import _Actor, _ActorMeshRefImpl, Actor, ActorMeshRef
 
 from monarch.common._device_utils import _local_device_count
@@ -46,14 +48,16 @@ class ProcMesh(MeshTrait):
     def __init__(self, hy_proc_mesh: HyProcMesh) -> None:
         self._proc_mesh = hy_proc_mesh
         self._mailbox: Mailbox = self._proc_mesh.client
-        self._rdma_manager = self._spawn_blocking("rdma_manager", RDMAManager)
+        self._rdma_manager: RDMAManager = self._spawn_blocking(
+            "rdma_manager", RDMAManager
+        )
 
     @property
-    def _ndslice(self):
+    def _ndslice(self) -> Slice:
         return self._proc_mesh.shape.ndslice
 
     @property
-    def _labels(self):
+    def _labels(self) -> List[str]:
         return self._proc_mesh.shape.labels
 
     def _new_with_shape(self, shape: Shape) -> "ProcMesh":