Fix debugger test

D78585722  broke test_debug and somehow the tests don't run on diffs.

This adds the necessary annotations to make it ok for the  debugger code to block the existing event loop while it is doing debugging stuff.

Differential Revision: [D78634155](https://our.internmc.facebook.com/intern/diff/D78634155/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D78634155/)!

ghstack-source-id: 297383325
Pull Request resolved: #590

Author: zdevito

python/monarch/_src/actor/actor_mesh.py

@@ -18,7 +18,6 @@
 import traceback
 
 from abc import ABC, abstractmethod
-from contextlib import contextmanager
 
 from dataclasses import dataclass
 from operator import mul
@@ -80,6 +79,7 @@
 from monarch._src.actor.pickle import flatten, unflatten
 
 from monarch._src.actor.shape import MeshTrait, NDSlice
+from monarch._src.actor.sync_state import fake_sync_state
 
 if TYPE_CHECKING:
     from monarch._src.actor.proc_mesh import ProcMesh
@@ -686,16 +686,6 @@ def _process(self, msg: PythonMessage) -> Tuple[int, R]:
 # We do this by blanking out the running event loop during the call to the synchronous actor function.
 
 
-@contextmanager
-def fake_sync_state():
-    prev_loop = asyncio.events._get_running_loop()
-    asyncio._set_running_loop(None)
-    try:
-        yield
-    finally:
-        asyncio._set_running_loop(prev_loop)
-
-
 class _Actor:
     """
     This is the message handling implementation of a Python actor.
@@ -828,16 +818,17 @@ 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)
+            with fake_sync_state():
+                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:

python/monarch/_src/actor/debugger.py

@@ -24,6 +24,7 @@
     MonarchContext,
 )
 from monarch._src.actor.pdb_wrapper import DebuggerWrite, PdbWrapper
+from monarch._src.actor.sync_state import fake_sync_state
 from tabulate import tabulate
 
 
@@ -549,12 +550,14 @@ def remote_breakpointhook():
             "exists on both your client and worker processes."
         )
 
+    with fake_sync_state():
+        manager = DebugManager.ref().get_debug_client.call_one().get()
     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(),
+        manager,
     )
     DebugContext.set(DebugContext(pdb_wrapper))
     pdb_wrapper.set_trace(frame)

python/monarch/_src/actor/pdb_wrapper.py

@@ -16,6 +16,7 @@
 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._src.actor.debugger import DebugClient
@@ -47,9 +48,9 @@ def __init__(
         self._first = True
 
     def set_trace(self, frame):
-        self.client_ref.debugger_session_start.call_one(
+        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(frame)
@@ -66,7 +67,7 @@ def do_clear(self, arg):
             super().do_clear(arg)
 
     def end_debug_session(self):
-        self.client_ref.debugger_session_end.call_one(self.rank).get()
+        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.
@@ -85,16 +86,19 @@ 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)
+        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
@@ -119,14 +123,14 @@ def write(self, s: str):
             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

python/monarch/_src/actor/sync_state.py

@@ -0,0 +1,18 @@
+# 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
+from contextlib import contextmanager
+
+
+@contextmanager
+def fake_sync_state():
+    prev_loop = asyncio.events._get_running_loop()
+    asyncio._set_running_loop(None)
+    try:
+        yield
+    finally:
+        asyncio._set_running_loop(prev_loop)

python/tests/test_allocator.py

@@ -33,14 +33,13 @@
     ChannelTransport,
 )
 
-from monarch._src.actor.actor_mesh import fake_sync_state
-
 from monarch._src.actor.allocator import (
     ALLOC_LABEL_PROC_MESH_NAME,
     RemoteAllocator,
     StaticRemoteAllocInitializer,
     TorchXRemoteAllocInitializer,
 )
+from monarch._src.actor.sync_state import fake_sync_state
 from monarch.actor import (
     Actor,
     current_rank,