Big debugging update

Summary:
This diff contains several updates to the monarch actor mesh debugging experience.
- Improved debug input parsing, with the new `cast` command supporting more sophisticated rank selection grammar:
  - `cast ranks(3) pdb_command`: send `pdb_command` to rank 3.
  - `cast ranks(1,3,5) pdb_command`: send `pdb_command` to ranks 1, 3 and 5.
  - `cast ranks(1:10:2) pdb_command`: send `pdb_command` to the ranks in `range(start=1, stop=10, step=2)`.
  - `cast ranks(pp=2, dp=(1,3), tp=2:8) pdb_command`: send `pdb_command` to ranks with `pp` dim 2, `dp` dim 1 or 3, and `tp` dim in `range(2,8)`.
- The debug client is now automatically registered with an actor mesh when that actor mesh is spawned. This means calling `init_debugging(actor_mesh)` is no longer necessary.
- Debugging now works with MAST jobs, by enforcing that breakpoints aren't set in `__main__`, and that the file containing the breakpoint exists on the remote host.
  - The first requirement is due to how `cloudpickle` works -- if an actor endpoint is defined inside `__main__`, `cloudpickle` will serialize it by value instead of by reference. When the code then runs on the remote host, it thinks the location of the code is the user's local `__main__` file, which confuses pdb, because the file doesn't exist at the same path (or may not exist at all) on the remote host.
  - The second requirement is due to important parts of `pdb`'s implementation relying on the ability to search for the file being debugged on the remote host's file system.
- A debugging session for a specific rank is now forcefully exited once the endpoint finishes execution. This contains the debugging experience within user-authored code. It is also necessary for preventing hangs, because if pdb is allowed to continue indefinitely, then control flow will eventually bubble back up to the main asyncio event loop on the worker, at which point everything breaks.
- Hitting a breakpoint now automatically enables post-mortem debugging, so any rank that encounters an exception after hitting a breakpoint will automatically stop at the exception. Attaching the debugger to that rank should then provide an experience like `pdb.post_mortem()`.

## Next steps/gaps I'm aware of (reviewers please read):
- Indexing debug sessions by rank isn't sustainable, because two actor meshes may simultaneously hit breakpoints on the same rank and cause a collision inside the debug client.
- Entering the debug client should happen automatically, rather than requiring the user to do `await debug_client().enter.call_one()`.
- Casting pdb commands should ideally leverage `MeshTrait` rather than reimplementing the selection logic.
- If a mesh was reshaped/renamed so that its dimension names aren't `hosts` and `gpus` anymore, the debugger should reflect the new shape/names.
- The user should be able to enable post-mortem debugging without having to hit a separate breakpoint first.

Differential Revision: D77568423

Author: Sam Lurye

python/monarch/__init__.py

@@ -57,6 +57,7 @@
     from monarch.future import ActorFuture
     from monarch.gradient_generator import grad_function, grad_generator
     from monarch.notebook import mast_mesh, reserve_torchx as mast_reserve
+    from monarch.proc_mesh import debug_client
     from monarch.python_local_mesh import python_local_mesh
     from monarch.rust_backend_mesh import (
         rust_backend_mesh,
@@ -116,6 +117,7 @@
     "LocalAllocator": ("monarch.allocator", "LocalAllocator"),
     "ActorFuture": ("monarch.future", "ActorFuture"),
     "builtins": ("monarch.builtins", "builtins"),
+    "debug_client": ("monarch.proc_mesh", "debug_client"),
 }
 
 
@@ -185,5 +187,6 @@ def __getattr__(name):
     "LocalAllocator",
     "ActorFuture",
     "builtins",
+    "debug_client",
 ]
 assert sorted(__all__) == sorted(_public_api)

python/monarch/actor_mesh.py

@@ -13,6 +13,7 @@
 import io
 import itertools
 import logging
+import os
 import random
 import sys
 import traceback
@@ -61,7 +62,7 @@
 
 from monarch.common.pickle_flatten import flatten, unflatten
 from monarch.common.shape import MeshTrait, NDSlice
-from monarch.pdb_wrapper import remote_breakpointhook
+from monarch.pdb_wrapper import PdbWrapper
 
 if TYPE_CHECKING:
     from monarch.debugger import DebugClient
@@ -578,9 +579,10 @@ async def instrumented():
                     )
                     try:
                         result = await the_method(self.instance, *args, **kwargs)
+                        self.instance._maybe_exit_debugger()
                     except Exception as e:
                         logging.critical(
-                            "Unahndled exception in actor endpoint",
+                            "Unhandled exception in actor endpoint",
                             exc_info=e,
                         )
                         raise e
@@ -593,11 +595,15 @@ async def instrumented():
                     the_method.__module__, message.method, str(ctx.mailbox.actor_id)
                 )
                 result = the_method(self.instance, *args, **kwargs)
+                # pyre-ignore
+                self.instance._maybe_exit_debugger()
                 exit_span()
 
             if port is not None:
                 port.send("result", result)
         except Exception as e:
+            # pyre-ignore
+            self.instance._post_mortem_debug(e.__traceback__)
             traceback.print_exc()
             s = ActorError(e)
 
@@ -608,6 +614,7 @@ async def instrumented():
             else:
                 raise s from None
         except BaseException as e:
+            self.instance._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]
@@ -674,18 +681,63 @@ def _new_with_shape(self, shape: Shape) -> "ActorMeshRef":
         )
 
     @endpoint  # pyre-ignore
-    def _set_debug_client(self, client: "DebugClient") -> None:
+    def _set_debug_client(self, client: "ActorMeshRef[DebugClient]") -> None:
         point = MonarchContext.get().point
         # For some reason, using a lambda instead of functools.partial
         # confuses the pdb wrapper implementation.
-        sys.breakpointhook = functools.partial(  # pyre-ignore
-            remote_breakpointhook,
+        sys.breakpointhook = functools.partial(
+            self._remote_breakpointhook,
             point.rank,
             point.shape.coordinates(point.rank),
             MonarchContext.get().mailbox.actor_id,
             client,
         )
 
+    def _remote_breakpointhook(
+        self,
+        rank: int,
+        coords: Dict[str, int],
+        actor_id: ActorId,
+        client_ref: "DebugClient",
+    ) -> None:
+        frame = inspect.currentframe()
+        assert frame is not None
+        frame = frame.f_back
+        assert frame is not None
+        file = frame.f_code.co_filename
+        line = frame.f_lineno
+        module = frame.f_globals.get("__name__", "__main__")
+        if module == "__main__" and not os.path.exists(file):
+            raise NotImplementedError(
+                f"Remote debugging not supported for breakpoint at {file}:{line} because "
+                f"it is defined inside __main__, and the file does not exist on the host. "
+                "In this case, cloudpickle serialization does not interact nicely with pdb. "
+                "To debug your code, move it out of __main__ and into a module that "
+                "exists on both your client and worker processes."
+            )
+
+        # pyre-ignore
+        self._pdb_wrapper_args = (rank, coords, actor_id, client_ref)
+        # pyre-ignore
+        self._pdb_wrapper = PdbWrapper(*self._pdb_wrapper_args)
+        self._pdb_wrapper.set_trace(frame)
+
+    def _post_mortem_debug(self, exc_tb) -> None:
+        if hasattr(self, "_pdb_wrapper"):
+            # pyre-ignore
+            self._pdb_wrapper = PdbWrapper(*self._pdb_wrapper_args)
+            self._pdb_wrapper.post_mortem(exc_tb)
+            self._maybe_exit_debugger(do_continue=False)
+
+    def _maybe_exit_debugger(self, do_continue=True) -> None:
+        if hasattr(self, "_pdb_wrapper"):
+            if do_continue:
+                # pyre-ignore
+                self._pdb_wrapper.clear_all_breaks()
+                self._pdb_wrapper.do_continue(None)
+            self._pdb_wrapper.end_debug_session()
+            del self._pdb_wrapper
+
 
 class ActorMeshRef(MeshTrait, Generic[T]):
     def __init__(