How many controllers would a ControllerController control? (#903)

Summary:
Pull Request resolved: #903

This diff introduces the `get_or_spawn_controller(name, Class, *args, **kwargs)` API, which creates a singleton actor of type `Class` indexed by `name` if no such actor with `name` already exists; otherwise, it returns a reference to the existing singleton actor with that `name`. The singleton actor is created on the root client.

To facilitate this API and manage all of the singleton controllers, this diff adds the `_ControllerController` singleton, which lives on the root client and is spawned the first time a proc mesh is created. Every time a new actor mesh is created, the mesh receives a reference to the `_ControllerController`, which is added to `MonarchContext` and is accessed inside `get_or_spawn_controller`.

With the new `get_or_spawn_controller` functionality, we no longer need to initialize `_RdmaManager` and `DebugManager` (or any manager) inside `proc_mesh.py`. (Actually, `DebugManager` is no longer needed at all, and `DebugClient` has been renamed to `DebugController`).

In order to support the `_RdmaManager` pattern, where a controller singleton spawns a manager actor on a user-specified proc mesh, we need a way to pass proc meshes through actor endpoints. This diff implements the exceedingly hacky `ProcMeshRef`, which just contains the proc mesh id, and can't do anything `ProcMesh` can do. However, when passed to an actor endpoint on the root client process, `ProcMeshRef` can be dereferenced into an actual `ProcMesh` by looking up its id in the `_proc_mesh_registry`, which is a global map from proc mesh id to `weakref(ProcMesh)`.

Also, `rdma/__init__.pyi` was moved into `rdma.pyi` to resolve some pesky lint issues related to importing `monarch._rust_bindings.rdma`.
ghstack-source-id: 303907982

Reviewed By: zdevito

Differential Revision: D80374416

fbshipit-source-id: 58373a58454f9986c10f947cde8444858f46acb7

Author: samlurye

monarch_hyperactor/src/pytokio.rs

@@ -8,7 +8,6 @@
 
 use std::error::Error;
 use std::future::Future;
-use std::ops::Deref;
 use std::pin::Pin;
 
 use hyperactor::clock::Clock;
@@ -20,6 +19,7 @@ use pyo3::exceptions::PyStopIteration;
 use pyo3::exceptions::PyTimeoutError;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use pyo3::types::PyNone;
 use pyo3::types::PyType;
 use tokio::sync::Mutex;
 use tokio::sync::watch;
@@ -182,8 +182,17 @@ impl PyPythonTask {
     }
 
     #[staticmethod]
-    fn from_coroutine(coro: PyObject) -> PyResult<PyPythonTask> {
-        PyPythonTask::new(async {
+    fn from_coroutine(py: Python<'_>, coro: PyObject) -> PyResult<PyPythonTask> {
+        // MonarchContext.get() used inside a PythonTask should inherit the value of
+        // MonarchContext from the context in which the PythonTask was constructed.
+        // We need to do this manually because the value of the contextvar isn't
+        // maintained inside the tokio runtime.
+        let monarch_context = py
+            .import("monarch._src.actor.actor_mesh")?
+            .getattr("MonarchContext")?
+            .call_method0("get")?
+            .unbind();
+        PyPythonTask::new(async move {
             let (coroutine_iterator, none) = Python::with_gil(|py| {
                 coro.into_bound(py)
                     .call_method0("__await__")
@@ -196,12 +205,23 @@ impl PyPythonTask {
             }
             loop {
                 let action: PyResult<Action> = Python::with_gil(|py| {
+                    // We may be executing in a new thread at this point, so we need to set the value
+                    // of MonarchContext.
+                    let _context = py
+                        .import("monarch._src.actor.actor_mesh")?
+                        .getattr("_context")?;
+                    let old_context = _context.call_method1("get", (PyNone::get(py),))?;
+                    _context.call_method1("set", (monarch_context.clone_ref(py),))?;
+
                     let result = match last {
                         Ok(value) => coroutine_iterator.bind(py).call_method1("send", (value,)),
                         Err(pyerr) => coroutine_iterator
                             .bind(py)
                             .call_method1("throw", (pyerr.into_value(py),)),
                     };
+
+                    // Reset MonarchContext so that when this tokio thread yields, it has its original state.
+                    _context.call_method1("set", (old_context,))?;
                     match result {
                         Ok(task) => Ok(Action::Wait(
                             task.extract::<Py<PyPythonTask>>()

python/monarch/_rust_bindings/rdma/__init__.pyi renamed to python/monarch/_rust_bindings/rdma.pyi

@@ -4,7 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, final, Optional
+# pyre-strict
+from typing import Any, final
 
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
 
@@ -51,7 +52,7 @@ class _RdmaBuffer:
         client: Any,
         timeout: int,
     ) -> PythonTask[Any]: ...
-    def __reduce__(self) -> tuple: ...
+    def __reduce__(self) -> tuple[Any, ...]: ...
     def __repr__(self) -> str: ...
     @staticmethod
     def new_from_json(json: str) -> _RdmaBuffer: ...

python/monarch/_src/actor/actor_mesh.py

@@ -12,10 +12,8 @@
 import inspect
 import itertools
 import logging
-import os
 import random
 import traceback
-from abc import ABC, abstractmethod
 
 from dataclasses import dataclass
 from traceback import TracebackException
@@ -98,7 +96,7 @@
     from monarch._rust_bindings.monarch_hyperactor.actor import PortProtocol
     from monarch._rust_bindings.monarch_hyperactor.actor_mesh import ActorMeshProtocol
     from monarch._rust_bindings.monarch_hyperactor.mailbox import PortReceiverBase
-    from monarch._src.actor.proc_mesh import ProcMesh
+    from monarch._src.actor.proc_mesh import _ControllerController, ProcMesh
 
 CallMethod = PythonMessageKind.CallMethod
 
@@ -127,14 +125,18 @@ class MonarchContext:
     proc_id: str
     point: Point
     send_queue: Tuple[Optional["Shared[Any]"], int]
+    controller_controller: Optional["_ControllerController"]
+    proc_mesh: Optional["ProcMesh"]  # actually this is a ProcMeshRef under the hood
 
     @staticmethod
     def get() -> "MonarchContext":
         c = _context.get(None)
         if c is None:
             mb = Mailbox.root_client_mailbox()
             proc_id = mb.actor_id.proc_id
-            c = MonarchContext(mb, proc_id, Point(0, singleton_shape), (None, 0))
+            c = MonarchContext(
+                mb, proc_id, Point(0, singleton_shape), (None, 0), None, None
+            )
             _context.set(c)
         return c
 
@@ -778,7 +780,12 @@ async def handle(
             if ctx is None:
                 # we reuse ctx across the actor so that send_queue is preserved between calls.
                 ctx = self._ctx = MonarchContext(
-                    mailbox, mailbox.actor_id.proc_id, Point(rank, shape), (None, 0)
+                    mailbox,
+                    mailbox.actor_id.proc_id,
+                    Point(rank, shape),
+                    (None, 0),
+                    None,
+                    None,
                 )
             ctx.mailbox = mailbox
             ctx.proc_id = mailbox.actor_id.proc_id
@@ -791,7 +798,10 @@ async def handle(
 
             match method:
                 case MethodSpecifier.Init():
-                    Class, *args = args
+                    Class, proc_mesh, controller_controller, *args = args
+                    ctx.controller_controller = controller_controller
+                    ctx.proc_mesh = proc_mesh
+                    _context.set(ctx)
                     try:
                         self.instance = Class(*args, **kwargs)
                     except Exception as e:
@@ -885,7 +895,7 @@ def _maybe_exit_debugger(self, do_continue=True) -> None:
         DebugContext.set(DebugContext())
 
     def _post_mortem_debug(self, exc_tb) -> None:
-        from monarch._src.actor.debugger import DebugManager
+        from monarch._src.actor.debugger import debug_controller
 
         if (pdb_wrapper := DebugContext.get().pdb_wrapper) is not None:
             with fake_sync_state():
@@ -894,7 +904,7 @@ def _post_mortem_debug(self, exc_tb) -> None:
                     ctx.point.rank,
                     ctx.point.shape.coordinates(ctx.point.rank),
                     ctx.mailbox.actor_id,
-                    DebugManager.ref().get_debug_client.call_one().get(),
+                    debug_controller(),
                 )
                 DebugContext.set(DebugContext(pdb_wrapper))
                 pdb_wrapper.post_mortem(exc_tb)
@@ -1015,6 +1025,7 @@ def _create(
         mailbox: Mailbox,
         shape: Shape,
         proc_mesh: "ProcMesh",
+        controller_controller: Optional["_ControllerController"],
         # args and kwargs are passed to the __init__ method of the user defined
         # python actor object.
         *args: Any,
@@ -1038,7 +1049,7 @@ async def null_func(*_args: Iterable[Any], **_kwargs: Dict[str, Any]) -> None:
             None,
             False,
         )
-        send(ep, (mesh._class, *args), kwargs)
+        send(ep, (mesh._class, proc_mesh, controller_controller, *args), kwargs)
 
         return mesh
 

python/monarch/_src/actor/debugger.py

@@ -14,18 +14,16 @@
 from dataclasses import dataclass
 from typing import cast, Dict, Generator, List, Optional, Tuple, Union
 
-from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
-from monarch._src.actor.actor_mesh import Actor, ActorMesh, DebugContext, MonarchContext
+from monarch._src.actor.actor_mesh import Actor, DebugContext, MonarchContext
 from monarch._src.actor.endpoint import endpoint
 from monarch._src.actor.pdb_wrapper import DebuggerWrite, PdbWrapper
+from monarch._src.actor.proc_mesh import get_or_spawn_controller
 from monarch._src.actor.sync_state import fake_sync_state
 from tabulate import tabulate
 
 
 logger = logging.getLogger(__name__)
 
-_DEBUG_MANAGER_ACTOR_NAME = "debug_manager"
-
 
 async def _debugger_input(prompt=""):
     return await asyncio.to_thread(input, prompt)
@@ -424,7 +422,7 @@ class Cast(DebugCommand):
     command: str
 
 
-class DebugClient(Actor):
+class DebugController(Actor):
     """
     Single actor for both remote debuggers and users to talk to.
 
@@ -563,26 +561,12 @@ async def debugger_write(
         await self.sessions.get(actor_name, rank).debugger_write(write)
 
 
-class DebugManager(Actor):
-    @staticmethod
-    @functools.cache
-    def ref() -> "DebugManager":
-        ctx = MonarchContext.get()
-        return cast(
-            DebugManager,
-            ActorMesh.from_actor_id(
-                DebugManager,
-                ActorId.from_string(f"{ctx.proc_id}.{_DEBUG_MANAGER_ACTOR_NAME}[0]"),
-                ctx.mailbox,
-            ),
-        )
-
-    def __init__(self, debug_client: DebugClient) -> None:
-        self._debug_client = debug_client
-
-    @endpoint
-    def get_debug_client(self) -> DebugClient:
-        return self._debug_client
+# Cached so that we don't have to call out to the root client every time,
+# which may be on a different host.
+@functools.cache
+def debug_controller() -> DebugController:
+    with fake_sync_state():
+        return get_or_spawn_controller("debug_controller", DebugController).get()
 
 
 def remote_breakpointhook():
@@ -602,14 +586,12 @@ 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,
-        manager,
+        debug_controller(),
     )
     DebugContext.set(DebugContext(pdb_wrapper))
     pdb_wrapper.set_trace(frame)

python/monarch/_src/actor/pdb_wrapper.py

@@ -19,7 +19,7 @@
 from monarch._src.actor.sync_state import fake_sync_state
 
 if TYPE_CHECKING:
-    from monarch._src.actor.debugger import DebugClient
+    from monarch._src.actor.debugger import DebugController
 
 
 @dataclass
@@ -35,7 +35,7 @@ def __init__(
         rank: int,
         coords: Dict[str, int],
         actor_id: ActorId,
-        client_ref: "DebugClient",
+        client_ref: "DebugController",
         header: str | None = None,
     ):
         self.rank = rank