Big debugging update (#456)

Summary:
Pull Request resolved: #456

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.

Reviewed By: zdevito

Differential Revision: D77568423

fbshipit-source-id: edc05ae7e7ea7f943de2751e01a6486dce7d1903

Author: Sam Lurye

python/monarch/_src/actor/actor_mesh.py

@@ -12,7 +12,6 @@
 import inspect
 import logging
 import random
-import sys
 import traceback
 
 from dataclasses import dataclass
@@ -53,15 +52,12 @@
 from monarch._rust_bindings.monarch_hyperactor.telemetry import enter_span, exit_span
 from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator
 from monarch._src.actor.future import Future
-from monarch._src.actor.pdb_wrapper import remote_breakpointhook
+from monarch._src.actor.pdb_wrapper import PdbWrapper
 
 from monarch._src.actor.pickle import flatten, unpickle
 
 from monarch._src.actor.shape import MeshTrait, NDSlice
 
-if TYPE_CHECKING:
-    from monarch._src.actor.debugger import DebugClient
-
 logger: logging.Logger = logging.getLogger(__name__)
 
 Allocator = ProcessAllocator | LocalAllocator
@@ -97,6 +93,23 @@ def get() -> "MonarchContext":
 )
 
 
+@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")
@@ -538,6 +551,8 @@ async def handle_cast(
             )
             _context.set(ctx)
 
+            DebugContext.set(DebugContext())
+
             args, kwargs = unpickle(message.message, mailbox)
 
             if message.method == "__init__":
@@ -574,9 +589,10 @@ async def instrumented():
                     )
                     try:
                         result = await the_method(self.instance, *args, **kwargs)
+                        self._maybe_exit_debugger()
                     except Exception as e:
                         logging.critical(
-                            "Unahndled exception in actor endpoint",
+                            "Unhandled exception in actor endpoint",
                             exc_info=e,
                         )
                         raise e
@@ -589,11 +605,13 @@ async def instrumented():
                     the_method.__module__, message.method, str(ctx.mailbox.actor_id)
                 )
                 result = the_method(self.instance, *args, **kwargs)
+                self._maybe_exit_debugger()
                 exit_span()
 
             if port is not None:
                 port.send("result", result)
         except Exception as e:
+            self._post_mortem_debug(e.__traceback__)
             traceback.print_exc()
             s = ActorError(e)
 
@@ -604,6 +622,7 @@ async def instrumented():
             else:
                 raise s from None
         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]
@@ -614,6 +633,29 @@ async def instrumented():
                 pass
             raise
 
+    def _maybe_exit_debugger(self, do_continue=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) -> None:
+        from monarch._src.actor.debugger import DebugManager
+
+        if (pdb_wrapper := DebugContext.get().pdb_wrapper) is not None:
+            ctx = MonarchContext.get()
+            pdb_wrapper = PdbWrapper(
+                ctx.point.rank,
+                ctx.point.shape.coordinates(ctx.point.rank),
+                ctx.mailbox.actor_id,
+                DebugManager.ref().get_debug_client.call_one().get(),
+            )
+            DebugContext.set(DebugContext(pdb_wrapper))
+            pdb_wrapper.post_mortem(exc_tb)
+            self._maybe_exit_debugger(do_continue=False)
+
 
 def _is_mailbox(x: object) -> bool:
     return isinstance(x, Mailbox)
@@ -648,19 +690,6 @@ def _new_with_shape(self, shape: Shape) -> "ActorMeshRef":
             "actor implementations are not meshes, but we can't convince the typechecker of it..."
         )
 
-    @endpoint  # pyre-ignore
-    def _set_debug_client(self, client: "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,
-            point.rank,
-            point.shape.coordinates(point.rank),
-            MonarchContext.get().mailbox.actor_id,
-            client,
-        )
-
 
 class ActorMeshRef(MeshTrait, Generic[T]):
     def __init__(

python/monarch/_src/actor/bootstrap_main.py

@@ -54,6 +54,10 @@ def invoke_main():
     except Exception as e:
         logging.warning(f"Failed to set up py-spy: {e}")
 
+    from monarch._src.actor.debugger import remote_breakpointhook
+
+    sys.breakpointhook = remote_breakpointhook
+
     # Start an event loop for PythonActors to use.
     asyncio.run(main())