root_client_mailbox: open ports without a proc (#683)

Summary:
Pull Request resolved: #683

To get ready to make proc_mesh/actor_mesh remote references, we need to be able to open ports on the root client without having to wait for the proc to exist.

This diff creates a global Mailbox hooked up to the global router that can be used on the root client to open ports.

On actors the mailbox is looked up via the monarch context set while the actor is running. On the root proc, we use the new global mailbox.

This makes it possible to create a function to create a port without other arguments. This function is actually a class named Channel[T] because of python typing reasons: there is way to pass a type argument to a function in Python. Called channel because that is the rust equivalent. Channel[T].open() returns what used to be called PortTuple, but is now just a real tuple.
ghstack-source-id: 300749961
exported-using-ghexport

Reviewed By: mariusae

Differential Revision: D79215096

fbshipit-source-id: 5ace368b0813467969cc0f43ef587ab4775cbf50

Author: zdevito

hyperactor/src/channel.rs

@@ -575,6 +575,18 @@ pub async fn serve<M: RemoteMessage>(
     .map(|(addr, inner)| (addr, ChannelRx { inner }))
 }
 
+/// Serve on the local address. The server is turned down
+/// when the returned Rx is dropped.
+pub fn serve_local<M: RemoteMessage>() -> (ChannelAddr, ChannelRx<M>) {
+    let (port, rx) = local::serve::<M>();
+    (
+        ChannelAddr::Local(port),
+        ChannelRx {
+            inner: ChannelRxKind::Local(rx),
+        },
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use std::assert_matches::assert_matches;

hyperactor_mesh/src/proc_mesh.rs

@@ -63,6 +63,7 @@ use crate::proc_mesh::mesh_agent::MeshAgent;
 use crate::proc_mesh::mesh_agent::MeshAgentMessageClient;
 use crate::proc_mesh::mesh_agent::StopActorResult;
 use crate::reference::ProcMeshId;
+use crate::shortuuid::ShortUuid;
 
 pub mod mesh_agent;
 
@@ -79,6 +80,26 @@ pub(crate) fn global_router() -> &'static MailboxRouter {
     GLOBAL_ROUTER.get_or_init(MailboxRouter::new)
 }
 
+/// Global mailbox used by the root client to send messages.
+/// This mailbox allows us to open ports before we know which proc the
+/// messages will be sent to.
+pub fn global_mailbox() -> Mailbox {
+    static GLOBAL_MAILBOX: OnceLock<Mailbox> = OnceLock::new();
+    GLOBAL_MAILBOX
+        .get_or_init(|| {
+            let world_id = WorldId(ShortUuid::generate().to_string());
+            let client_proc_id = ProcId(world_id.clone(), 0);
+            let client_proc = Proc::new(
+                client_proc_id.clone(),
+                BoxedMailboxSender::new(global_router().clone()),
+            );
+            global_router().bind(world_id.clone().into(), client_proc.clone());
+
+            client_proc.attach("client").expect("root mailbox creation")
+        })
+        .clone()
+}
+
 type ActorEventRouter = Arc<DashMap<ActorMeshName, mpsc::UnboundedSender<ActorSupervisionEvent>>>;
 /// A ProcMesh maintains a mesh of procs whose lifecycles are managed by
 /// an allocator.

monarch_hyperactor/src/mailbox.rs

@@ -69,6 +69,12 @@ impl PyMailbox {
 
 #[pymethods]
 impl PyMailbox {
+    #[staticmethod]
+    fn root_client_mailbox() -> PyMailbox {
+        PyMailbox {
+            inner: hyperactor_mesh::proc_mesh::global_mailbox(),
+        }
+    }
     fn open_port<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
         let (handle, receiver) = self.inner.open_port();
         let handle = Py::new(py, PythonPortHandle { inner: handle })?;

python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi

@@ -137,7 +137,8 @@ class Mailbox:
     """
     A mailbox from that can receive messages.
     """
-
+    @staticmethod
+    def root_client_mailbox() -> "Mailbox": ...
     def open_port(self) -> tuple[PortHandle, PortReceiver]:
         """Open a port to receive `PythonMessage` messages."""
         ...

python/monarch/_src/actor/actor_mesh.py

@@ -120,6 +120,13 @@ class MonarchContext:
     def get() -> "MonarchContext":
         return _context.get()
 
+    @staticmethod
+    def current_mailbox() -> "Mailbox":
+        context = _context.get(None)
+        if context is not None:
+            return context.mailbox
+        return Mailbox.root_client_mailbox()
+
 
 _context: contextvars.ContextVar[MonarchContext] = contextvars.ContextVar(
     "monarch.actor_mesh._context"
@@ -344,13 +351,13 @@ def _send(
         shape = self._actor_mesh._shape
         return Extent(shape.labels, shape.ndslice.sizes)
 
-    def _port(self, once: bool = False) -> "PortTuple[R]":
-        p, r = PortTuple.create(self._mailbox, once)
+    def _port(self, once: bool = False) -> "Tuple[Port[R], PortReceiver[R]]":
+        p, r = super()._port(once=once)
         if TYPE_CHECKING:
             assert isinstance(
                 r._receiver, (HyPortReceiver | OncePortReceiver)
             ), "unexpected receiver type"
-        return PortTuple(p, PortReceiver(self._mailbox, self._supervise(r._receiver)))
+        return (p, PortReceiver(self._mailbox, self._supervise(r._receiver)))
 
     def _rref(self, args, kwargs):
         self._check_arguments(args, kwargs)
@@ -526,49 +533,25 @@ def exception(self, obj: Exception) -> None:
 
 T = TypeVar("T")
 
-if TYPE_CHECKING:
-    # Python <= 3.10 cannot inherit from Generic[R] and NamedTuple at the same time.
-    # we only need it for type checking though, so copypasta it until 3.11.
-    class PortTuple(NamedTuple, Generic[R]):
-        sender: "Port[R]"
-        receiver: "PortReceiver[R]"
-
-        @staticmethod
-        def create(mailbox: Mailbox, once: bool = False) -> "PortTuple[Any]":
-            handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
-            port_ref = handle.bind()
-            return PortTuple(
-                Port(port_ref, mailbox, rank=None),
-                PortReceiver(mailbox, receiver),
-            )
-else:
-
-    class PortTuple(NamedTuple):
-        sender: "Port[Any]"
-        receiver: "PortReceiver[Any]"
-
-        @staticmethod
-        def create(mailbox: Mailbox, once: bool = False) -> "PortTuple[Any]":
-            handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
-            port_ref = handle.bind()
-            return PortTuple(
-                Port(port_ref, mailbox, rank=None),
-                PortReceiver(mailbox, receiver),
-            )
-
 
 # advance lower-level API for sending messages. This is intentially
 # not part of the Endpoint API because they way it accepts arguments
 # and handles concerns is different.
-def port(endpoint: Endpoint[P, R], once: bool = False) -> "PortTuple[R]":
-    return endpoint._port(once)
-
+class Channel(Generic[R]):
+    @staticmethod
+    def open(once: bool = False) -> Tuple["Port[R]", "PortReceiver[R]"]:
+        mailbox = MonarchContext.current_mailbox()
+        handle, receiver = mailbox.open_once_port() if once else mailbox.open_port()
+        port_ref = handle.bind()
+        return (
+            Port(port_ref, mailbox, rank=None),
+            PortReceiver(mailbox, receiver),
+        )
 
-def ranked_port(
-    endpoint: Endpoint[P, R], once: bool = False
-) -> Tuple["Port[R]", "RankedPortReceiver[R]"]:
-    p, receiver = port(endpoint, once)
-    return p, RankedPortReceiver[R](receiver._mailbox, receiver._receiver)
+    @staticmethod
+    def open_ranked(once: bool = False) -> Tuple["Port[R]", "RankedPortReceiver[R]"]:
+        send, recv = Channel[R].open()
+        return (send, recv.ranked())
 
 
 class PortReceiver(Generic[R]):
@@ -597,6 +580,9 @@ def _process(self, msg: PythonMessage) -> R:
     def recv(self) -> "Future[R]":
         return Future(coro=self._recv())
 
+    def ranked(self) -> "RankedPortReceiver[R]":
+        return RankedPortReceiver[R](self._mailbox, self._receiver)
+
 
 class RankedPortReceiver(PortReceiver[Tuple[int, R]]):
     def _process(self, msg: PythonMessage) -> Tuple[int, R]:

python/monarch/_src/actor/endpoint.py

@@ -40,7 +40,7 @@
         HyPortReceiver,
         OncePortReceiver,
         Port,
-        PortTuple,
+        PortReceiver,
         ValueMesh,
     )
 
@@ -90,9 +90,10 @@ def _send(
         """
         pass
 
-    @abstractmethod
-    def _port(self, once: bool = False) -> "PortTuple[R]":
-        pass
+    def _port(self, once: bool = False) -> "Tuple[Port[R], PortReceiver[R]]":
+        from monarch._src.actor.actor_mesh import Channel
+
+        return Channel[R].open(once)
 
     @abstractmethod
     def _call_name(self) -> Any:
@@ -115,17 +116,14 @@ def choose(self, *args: P.args, **kwargs: P.kwargs) -> Future[R]:
 
         Load balanced RPC-style entrypoint for request/response messaging.
         """
-        from monarch._src.actor.actor_mesh import port
 
-        p, r = port(self, once=True)
+        p, r = self._port(once=True)
         # pyre-ignore
         self._send(args, kwargs, port=p, selection="choose")
         return r.recv()
 
     def call_one(self, *args: P.args, **kwargs: P.kwargs) -> Future[R]:
-        from monarch._src.actor.actor_mesh import port
-
-        p, r = port(self, once=True)
+        p, r = self._port(once=True)
         # pyre-ignore
         extent = self._send(args, kwargs, port=p, selection="choose")
         if extent.nelements != 1:
@@ -135,9 +133,10 @@ def call_one(self, *args: P.args, **kwargs: P.kwargs) -> Future[R]:
         return r.recv()
 
     def call(self, *args: P.args, **kwargs: P.kwargs) -> "Future[ValueMesh[R]]":
-        from monarch._src.actor.actor_mesh import ranked_port, ValueMesh
+        from monarch._src.actor.actor_mesh import ValueMesh
 
-        p, r = ranked_port(self)
+        p, unranked = self._port()
+        r = unranked.ranked()
         # pyre-ignore
         extent = self._send(args, kwargs, port=p)
 
@@ -166,9 +165,8 @@ def _stream(
         This enables processing results from multiple actors incrementally as
         they become available. Returns an async generator of response values.
         """
-        from monarch._src.actor.actor_mesh import port
 
-        p, r = port(self)
+        p, r = self._port()
         # pyre-ignore
         extent = self._send(args, kwargs, port=p)
         for _ in range(extent.nelements):

python/monarch/_src/actor/proc_mesh.py

@@ -41,7 +41,13 @@
 )
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask, Shared
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
-from monarch._src.actor.actor_mesh import _Actor, _ActorMeshRefImpl, Actor, ActorMeshRef
+from monarch._src.actor.actor_mesh import (
+    _Actor,
+    _ActorMeshRefImpl,
+    Actor,
+    ActorMeshRef,
+    MonarchContext,
+)
 
 from monarch._src.actor.allocator import (
     AllocateMixin,
@@ -317,7 +323,7 @@ async def _spawn_nonblocking_on(
         service = ActorMeshRef(
             Class,
             _ActorMeshRefImpl.from_hyperactor_mesh(pm.client, actor_mesh, self),
-            pm.client,
+            MonarchContext.current_mailbox(),
         )
         # useful to have this separate, because eventually we can reconstitute ActorMeshRef objects across pickling by
         # doing `ActorMeshRef(Class, actor_handle)` but not calling _create.

python/monarch/actor/__init__.py

@@ -13,11 +13,11 @@
     Actor,
     ActorError,
     as_endpoint,
+    Channel,
     current_actor_name,
     current_rank,
     current_size,
     Point,
-    port,
     send,
     ValueMesh,
 )
@@ -45,7 +45,7 @@
     "Point",
     "proc_mesh",
     "ProcMesh",
-    "port",
+    "Channel",
     "send",
     "sim_proc_mesh",
     "ValueMesh",

python/monarch/common/remote.py

@@ -28,9 +28,8 @@
 import monarch.common.messages as messages
 
 import torch
-from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape
-from monarch._src.actor.actor_mesh import Port, PortTuple
+from monarch._src.actor.actor_mesh import Port
 from monarch._src.actor.endpoint import Extent, Selection
 
 from monarch.common import _coalescing, device_mesh, stream
@@ -135,20 +134,6 @@ def _send(
         client._request_status()
         return Extent(ambient_mesh._labels, ambient_mesh._ndslice.sizes)
 
-    def _port(self, once: bool = False) -> "PortTuple[R]":
-        ambient_mesh = device_mesh._active
-        if ambient_mesh is None:
-            raise ValueError(
-                "FIXME - cannot create a port without an active proc_mesh, because there is not way to create a port without a mailbox"
-            )
-        mesh_controller = getattr(ambient_mesh.client, "_mesh_controller", None)
-        if mesh_controller is None:
-            raise ValueError(
-                "Cannot create raw port objects with an old-style tensor engine controller."
-            )
-        mailbox: Mailbox = mesh_controller._mailbox
-        return PortTuple.create(mailbox, once)
-
     @property
     def _resolvable(self):
         return resolvable_function(self._remote_impl)

python/monarch/mesh_controller.py

@@ -43,7 +43,7 @@
     ActorId,
 )
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask
-from monarch._src.actor.actor_mesh import ActorEndpoint, Port, PortTuple
+from monarch._src.actor.actor_mesh import ActorEndpoint, Channel, Port
 from monarch._src.actor.endpoint import Selection
 from monarch._src.actor.shape import NDSlice
 from monarch.common import device_mesh, messages, stream
@@ -156,7 +156,7 @@ def fetch(
         defs: Tuple["Tensor", ...],
         uses: Tuple["Tensor", ...],
     ) -> "OldFuture":  # the OldFuture is a lie
-        sender, receiver = PortTuple.create(self._mesh_controller._mailbox, once=True)
+        sender, receiver = Channel.open(once=True)
 
         ident = self.new_node(defs, uses, cast("OldFuture", sender))
         process = mesh._process(shard)
@@ -192,7 +192,7 @@ def shutdown(
         atexit.unregister(self._atexit)
         self._shutdown = True
 
-        sender, receiver = PortTuple.create(self._mesh_controller._mailbox, once=True)
+        sender, receiver = Channel.open(once=True)
         assert sender._port_ref is not None
         self._mesh_controller.sync_at_exit(sender._port_ref.port_id)
         receiver.recv().get(timeout=60)