Context and Instance python APIs (meta-pytorch#918)

Summary:
Pull Request resolved: meta-pytorch#918

Adding Context/Instance objects as proposed in https://docs.google.com/document/d/11ohUVih8yPmT1lZaV-OpL-Lawj49-RhJJ6tqxlW94S0/edit?tab=t.0 .

This also replaces our global root mailbox with a global root client instance object.

Context is created from the rust Context<Self> object, though it cannot directly contain Context because the lifetime of the rust thing is limited. (Is there some way to store a limited lifetime thing with a longer lifetime and dynamically assert it is still alvie)?
ghstack-source-id: 304344057

Reviewed By: mariusae

Differential Revision: D80505166

fbshipit-source-id: 2b0bfe1610d20199105a3fdc5a1ed37a1444be18

Author: zdevito

hyperactor_mesh/src/proc_mesh.rs

@@ -16,8 +16,10 @@ use async_trait::async_trait;
 use dashmap::DashMap;
 use futures::future::join_all;
 use hyperactor::Actor;
+use hyperactor::ActorHandle;
 use hyperactor::ActorId;
 use hyperactor::ActorRef;
+use hyperactor::Instance;
 use hyperactor::Mailbox;
 use hyperactor::Named;
 use hyperactor::RemoteMessage;
@@ -79,24 +81,24 @@ pub(crate) fn global_router() -> &'static MailboxRouter {
     GLOBAL_ROUTER.get_or_init(MailboxRouter::new)
 }
 
-/// Global mailbox used by the root client to send messages.
+/// Context use by 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::Ranked(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()
+pub fn global_root_client() -> &'static Instance<()> {
+    static GLOBAL_INSTANCE: OnceLock<(Instance<()>, ActorHandle<()>)> = OnceLock::new();
+    let (instance, _) = GLOBAL_INSTANCE.get_or_init(|| {
+        let world_id = WorldId(ShortUuid::generate().to_string());
+        let client_proc_id = ProcId::Ranked(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
+            .instance("client")
+            .expect("root instance create")
+    });
+    instance
 }
 
 type ActorEventRouter = Arc<DashMap<ActorMeshName, mpsc::UnboundedSender<ActorSupervisionEvent>>>;

monarch_hyperactor/src/actor.rs

@@ -625,15 +625,11 @@ impl Handler<PythonMessage> for PythonActor {
         let (sender, receiver) = oneshot::channel();
 
         let future = Python::with_gil(|py| -> Result<_, SerializablePyErr> {
-            let mailbox = mailbox(py, cx);
-            let (rank, shape) = cx.cast_info();
             let awaitable = self.actor.call_method(
                 py,
                 "handle",
                 (
-                    mailbox,
-                    rank,
-                    PyShape::from(shape),
+                    crate::mailbox::Context::new(py, cx),
                     resolved.method,
                     resolved.bytes,
                     PanicFlag {

monarch_hyperactor/src/mailbox.rs

@@ -9,8 +9,10 @@
 use std::hash::DefaultHasher;
 use std::hash::Hash;
 use std::hash::Hasher;
+use std::ops::Deref;
 use std::sync::Arc;
 
+use hyperactor::ActorId;
 use hyperactor::Mailbox;
 use hyperactor::Named;
 use hyperactor::OncePortHandle;
@@ -33,8 +35,12 @@ use hyperactor::mailbox::monitored_return_handle;
 use hyperactor::message::Bind;
 use hyperactor::message::Bindings;
 use hyperactor::message::Unbind;
+use hyperactor_mesh::comm::multicast::CastInfo;
 use hyperactor_mesh::comm::multicast::set_cast_info_on_headers;
+use hyperactor_mesh::proc_mesh::global_root_client;
 use monarch_types::PickledPyObject;
+use monarch_types::py_global;
+use ndslice::shape::Shape;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyEOFError;
 use pyo3::exceptions::PyRuntimeError;
@@ -51,6 +57,7 @@ use crate::proc::PyActorId;
 use crate::pytokio::PyPythonTask;
 use crate::pytokio::PythonTask;
 use crate::runtime::signal_safe_block_on;
+use crate::shape::PyPoint;
 use crate::shape::PyShape;
 #[derive(Clone, Debug)]
 #[pyclass(
@@ -69,12 +76,6 @@ 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 })?;
@@ -702,6 +703,86 @@ inventory::submit! {
     }
 }
 
+#[pyclass(name = "Instance", module = "monarch._src.actor.actor_mesh")]
+struct Instance {
+    mailbox: Mailbox,
+    actor_id: ActorId,
+    #[pyo3(get, set)]
+    proc_mesh: Option<PyObject>,
+    #[pyo3(get, set, name = "_controller_controller")]
+    controller_controller: Option<PyObject>,
+}
+#[pymethods]
+impl Instance {
+    #[getter]
+    fn _mailbox(&self) -> PyMailbox {
+        PyMailbox {
+            inner: self.mailbox.clone(),
+        }
+    }
+    #[getter]
+    fn actor_id(&self) -> PyActorId {
+        self.actor_id.clone().into()
+    }
+}
+
+impl<A: hyperactor::Actor> From<&hyperactor::proc::Instance<A>> for Instance {
+    fn from(ins: &hyperactor::proc::Instance<A>) -> Self {
+        Instance {
+            mailbox: ins.mailbox_for_py().clone(),
+            actor_id: ins.self_id().clone(),
+            proc_mesh: None,
+            controller_controller: None,
+        }
+    }
+}
+
+#[pyclass(name = "Context", module = "monarch._src.actor.actor_mesh")]
+pub(crate) struct Context {
+    instance: Py<Instance>,
+    rank: usize,
+    shape: Shape,
+}
+
+py_global!(point, "monarch._src.actor.actor_mesh", "Point");
+
+#[pymethods]
+impl Context {
+    #[getter]
+    fn actor_instance(&self) -> &Py<Instance> {
+        &self.instance
+    }
+    #[getter]
+    fn message_rank<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let shape = PyShape::from(self.shape.clone());
+        point(py).call1((self.rank, shape.into_pyobject(py).unwrap().unbind()))
+    }
+    #[staticmethod]
+    fn _root_client_context(py: Python<'_>) -> Context {
+        let instance: Instance = global_root_client().into();
+        Context {
+            instance: instance.into_pyobject(py).unwrap().into(),
+            rank: 0,
+            shape: Shape::unity(),
+        }
+    }
+}
+
+impl Context {
+    pub(crate) fn new<T: hyperactor::actor::Actor>(
+        py: Python<'_>,
+        cx: &hyperactor::proc::Context<T>,
+    ) -> Context {
+        let instance: Instance = cx.deref().into();
+        let (rank, shape) = cx.cast_info();
+        Context {
+            instance: instance.into_pyobject(py).unwrap().into(),
+            rank,
+            shape,
+        }
+    }
+}
+
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     hyperactor_mod.add_class::<PyMailbox>()?;
     hyperactor_mod.add_class::<PyPortId>()?;
@@ -713,5 +794,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PythonOncePortHandle>()?;
     hyperactor_mod.add_class::<PythonOncePortRef>()?;
     hyperactor_mod.add_class::<PythonOncePortReceiver>()?;
+    hyperactor_mod.add_class::<Instance>()?;
+    hyperactor_mod.add_class::<Context>()?;
     Ok(())
 }

monarch_hyperactor/src/pytokio.rs

@@ -183,14 +183,13 @@ impl PyPythonTask {
 
     #[staticmethod]
     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.
+        // context() used inside a PythonTask should inherit the value of
+        // context() 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")?
+            .call_method0("context")?
             .unbind();
         PyPythonTask::new(async move {
             let (coroutine_iterator, none) = Python::with_gil(|py| {
@@ -206,7 +205,7 @@ 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.
+                    // of context().
                     let _context = py
                         .import("monarch._src.actor.actor_mesh")?
                         .getattr("_context")?;
@@ -220,7 +219,7 @@ impl PyPythonTask {
                             .call_method1("throw", (pyerr.into_value(py),)),
                     };
 
-                    // Reset MonarchContext so that when this tokio thread yields, it has its original state.
+                    // Reset context() 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(

python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi

@@ -266,9 +266,7 @@ class PortProtocol(Generic[R], Protocol):
 class Actor(Protocol):
     async def handle(
         self,
-        mailbox: Mailbox,
-        rank: int,
-        shape: Shape,
+        context: Any,
         method: MethodSpecifier,
         message: bytes,
         panic_flag: PanicFlag,

python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi

@@ -137,8 +137,6 @@ 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/_rust_bindings/monarch_hyperactor/proc_mesh.pyi

@@ -6,9 +6,10 @@
 
 # pyre-strict
 
-from typing import AsyncIterator, final, Literal, overload, Type
+from typing import Any, AsyncIterator, final, Literal, overload, Type, TYPE_CHECKING
 
-from monarch._rust_bindings.monarch_hyperactor.actor import Actor
+if TYPE_CHECKING:
+    from monarch._rust_bindings.monarch_hyperactor.actor import Actor
 from monarch._rust_bindings.monarch_hyperactor.actor_mesh import (
     PythonActorMesh,
     PythonActorMeshImpl,
@@ -23,7 +24,7 @@ from monarch._rust_bindings.monarch_hyperactor.shape import Shape
 @final
 class ProcMesh:
     @classmethod
-    def allocate_nonblocking(self, alloc: Alloc) -> PythonTask[ProcMesh]:
+    def allocate_nonblocking(self, alloc: Alloc) -> PythonTask["ProcMesh"]:
         """
         Allocate a process mesh according to the provided alloc.
         Returns when the mesh is fully allocated.
@@ -34,7 +35,9 @@ class ProcMesh:
         ...
 
     def spawn_nonblocking(
-        self, name: str, actor: Type[Actor]
+        self,
+        name: str,
+        actor: Any,
     ) -> PythonTask[PythonActorMesh]:
         """
         Spawn a new actor on this mesh.
@@ -47,7 +50,7 @@ class ProcMesh:
 
     @staticmethod
     def spawn_async(
-        proc_mesh: Shared[ProcMesh], name: str, actor: Type[Actor], emulated: bool
+        proc_mesh: Shared["ProcMesh"], name: str, actor: Type["Actor"], emulated: bool
     ) -> PythonActorMesh: ...
     async def monitor(self) -> ProcMeshMonitor:
         """