[monarch] The root client is just a PythonActor

This diff makes the root client actor just another `PythonActor`.

# Why?

Right now the monarch codebase is peppered with special handling to distinguish between normal python actors and the root client "actor", which has type `()` and is actually just a detached `Instance` with no actor loop; it therefore has no message handlers and can't even process supervision events. As a result, we have to wrap the current context's instance in a special `ContextInstance` enum, and everywhere we want to use it, we either have to use the `instance_dispatch!` macro, or insert code that looks like:
```rust
match instance {
  ContextInstance::PythonActor(ins) => { do something },
  ContextInstance::Client(ins) => { do something else },
}
```
This makes the code more error-prone and harder to understand, with the added complication that the client handling is often not idiomatic w.r.t hyperactor due to the lack of message handlers/actor loop. Some examples:
- [Confusing supervision handling where `owner` might not be defined but `is_owned` is still true and so we need to call into a special `unhandled` function instead of continuing to propagate up the hierarchy](https://fburl.com/code/andy3ggr)
- [The root client can't have child actors due to no supervision event handling, so they have to be spawned directly on the root client proc, and even then, there is no way for the supervision event to reach `monarch.actor.unhandled_fault_hook`](https://fburl.com/code/kqd2iwvc)
- [The root client handles undeliverable messages via a bespoke tokio task/thread](https://fburl.com/code/jjgfy5d5)

Making the root client a normal python actor solves these problems, because:
- We don't need a `ContextInstance` enum anymore -- `PyInstance` *always* contains `Instance<PythonActor>`.
- Supervision events follow a unified path as they bubble up through the hierarchy, and *every* unhandled event reaches `RootClientActor.__supervise__`, defined in python, without special handling.
- The root client can handle undeliverable messages using `RootClientActor._handle_undeliverable_message`, defined in python, without special handling.

# Navigating the code changes (guide for reviewers)

There are a lot of file changes here but only some of them are important. I would recommend reviewing them in the following order:
- `monarch/_src/actor/actor_mesh.py`
  - Defines the `RootClientActor` python class and its behavior.
- `hyperactor/src/proc.rs`
  - Introduces `Proc::actor_instance::<A>(...)`, which returns a detached `A`-typed actor instance/handle, along with its supervision receiver, signal receiver and message receiver.
- `monarch_hyperactor/src/actor.rs`
  - Introduces `PythonActor::bootstrap_client()`, which replaces `global_root_client()` in the root client context. This function starts the root client proc, spawns the `RootClientActor`, starts its actor loop and returns the `Instance<PythonActor>`.
  - The root client actor can now handle `SupervisionFailureMessage` just like every other actor in the hierarchy.
  - Implements `PythonActor::handle_supervision_event` to pass the event to the actor's `SupervisionFailureMessage` handler. This way, **every unhandled supervision event in the system makes its way to `RootClientActor.__supervise__` eventually**.
- `monarch_hyperactor/src/v1/actor_mesh.rs`
  - Deletes the special handling from the actor states monitor like `is_owned` and the explicit `unhandled_fault_hook` call. If `owner` is defined, it forwards the `SupervisionFailureMessage`, or else it does nothing.
  - Fixes (what I think was) a bug in `send_state_change`. A supervision event should only be forwarded as `SupervisionFailureMessage` to `owner` if it represents a failure. With the logic before this diff, stopping an actor mesh from inside an actor endpoint would generate a supervision event that reaches `unhandled_fault_hook` and crashes the root process even if it was a healthy stop.
- `monarch_hyperactor/src/context.rs`
  - Deletes `ContextInstance` and replaces it in `PyInstance` with `Instance<PythonActor>`.
- The rest of the changes are pretty much just cleaning up `instance_dispatch!` calls.

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

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

ghstack-source-id: 325421344
Pull Request resolved: #1985

Author: samlurye

hyperactor/src/actor.rs

@@ -312,8 +312,10 @@ where
 /// with the ID of the actor being served.
 #[derive(Debug)]
 pub struct ActorError {
-    pub(crate) actor_id: Box<ActorId>,
-    pub(crate) kind: Box<ActorErrorKind>,
+    /// The ActorId for the actor that generated this error.
+    pub actor_id: Box<ActorId>,
+    /// The kind of error that occurred.
+    pub kind: Box<ActorErrorKind>,
 }
 
 /// The kinds of actor serving errors.

hyperactor/src/proc.rs

@@ -411,7 +411,9 @@ impl Proc {
             .map_err(|existing| anyhow::anyhow!("coordinator port is already set to {existing}"))
     }
 
-    fn handle_supervision_event(&self, event: ActorSupervisionEvent) {
+    /// Handle a supervision event received by the proc. Attempt to forward it to the
+    /// supervision coordinator port if one is set, otherwise crash the process.
+    pub fn handle_supervision_event(&self, event: ActorSupervisionEvent) {
         let result = match self.state().supervision_coordinator_port.get() {
             Some(port) => port.send(event.clone()).map_err(anyhow::Error::from),
             None => Err(anyhow::anyhow!(
@@ -530,26 +532,46 @@ impl Proc {
         Ok(instance.start(actor, actor_loop_receivers.take().unwrap(), work_rx))
     }
 
-    /// Create and return an actor instance and its corresponding handle. This allows actors to be
-    /// "inverted": the caller can use the returned [`Instance`] to send and receive messages,
-    /// launch child actors, etc. The actor itself does not handle any messages, and supervision events
-    /// are always forwarded to the proc. Otherwise the instance acts as a normal actor, and can be
-    /// referenced and stopped.
+    /// Wrapper for [`Proc::actor_instance::<()>`].
     pub fn instance(&self, name: &str) -> Result<(Instance<()>, ActorHandle<()>), anyhow::Error> {
+        let (instance, handle, ..) = self.actor_instance(name)?;
+
+        Ok((instance, handle))
+    }
+
+    /// Create and return an actor instance, its corresponding handle, its signal port receiver,
+    /// its supervision port receiver, and its message receiver. This allows actors to be
+    /// "inverted": the caller can use the returned [`Instance`] to send and receive messages,
+    /// launch child actors, etc. The actor itself does not handle any messages unless driven by
+    /// the caller. Otherwise the instance acts as a normal actor, and can be referenced and
+    /// stopped.
+    pub fn actor_instance<A: Actor>(
+        &self,
+        name: &str,
+    ) -> Result<
+        (
+            Instance<A>,
+            ActorHandle<A>,
+            PortReceiver<ActorSupervisionEvent>,
+            PortReceiver<Signal>,
+            mpsc::UnboundedReceiver<WorkCell<A>>,
+        ),
+        anyhow::Error,
+    > {
         let actor_id = self.allocate_root_id(name)?;
-        let _ = tracing::debug_span!(
+        let span = tracing::debug_span!(
             "actor_instance",
             actor_name = name,
-            actor_type = std::any::type_name::<()>(),
+            actor_type = std::any::type_name::<A>(),
             actor_id = actor_id.to_string(),
         );
-
-        let (instance, _, _) = Instance::new(self.clone(), actor_id.clone(), true, None);
+        let _guard = span.enter();
+        let (instance, actor_loop_receivers, work_rx) =
+            Instance::new(self.clone(), actor_id.clone(), false, None);
+        let (signal_rx, supervision_rx) = actor_loop_receivers.unwrap();
         let handle = ActorHandle::new(instance.inner.cell.clone(), instance.inner.ports.clone());
-
         instance.change_status(ActorStatus::Client);
-
-        Ok((instance, handle))
+        Ok((instance, handle, supervision_rx, signal_rx, work_rx))
     }
 
     /// Create a child instance. Called from `Instance`.
@@ -874,11 +896,11 @@ impl MailboxSender for WeakProc {
 
 /// Represents a single work item used by the instance to dispatch to
 /// actor handles. Specifically, this enables handler polymorphism.
-struct WorkCell<A: Actor + Send>(
+pub struct WorkCell<A: Actor + Send>(
     Box<
         dyn for<'a> FnOnce(
                 &'a mut A,
-                &'a mut Instance<A>,
+                &'a Instance<A>,
             )
                 -> Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + 'a + Send>>
             + Send
@@ -891,7 +913,7 @@ impl<A: Actor + Send> WorkCell<A> {
     fn new(
         f: impl for<'a> FnOnce(
             &'a mut A,
-            &'a mut Instance<A>,
+            &'a Instance<A>,
         )
             -> Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + 'a + Send>>
         + Send
@@ -902,10 +924,10 @@ impl<A: Actor + Send> WorkCell<A> {
     }
 
     /// Handle the message represented by this work cell.
-    fn handle<'a>(
+    pub fn handle<'a>(
         self,
         actor: &'a mut A,
-        instance: &'a mut Instance<A>,
+        instance: &'a Instance<A>,
     ) -> Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send + 'a>> {
         (self.0)(actor, instance)
     }
@@ -1451,7 +1473,8 @@ impl<A: Actor> Instance<A> {
         Ok(())
     }
 
-    async fn handle_supervision_event(
+    /// Handle a supervision event using the provided actor.
+    pub async fn handle_supervision_event(
         &self,
         actor: &mut A,
         supervision_event: ActorSupervisionEvent,
@@ -1483,7 +1506,7 @@ impl<A: Actor> Instance<A> {
 
     #[hyperactor::instrument(fields(actor_id = self.self_id().to_string(), actor_name = self.self_id().name()))]
     async unsafe fn handle_message<M: Message>(
-        &mut self,
+        &self,
         actor: &mut A,
         type_info: Option<&'static TypeInfo>,
         headers: Attrs,
@@ -1519,8 +1542,8 @@ impl<A: Actor> Instance<A> {
         actor.handle(&context, message).await
     }
 
-    /// Spawn on child on this instance. Currently used only by cap::CanSpawn.
-    pub(crate) fn spawn<C: Actor>(&self, actor: C) -> anyhow::Result<ActorHandle<C>> {
+    /// Spawn on child on this instance.
+    pub fn spawn<C: Actor>(&self, actor: C) -> anyhow::Result<ActorHandle<C>> {
         self.inner.proc.spawn_child(self.inner.cell.clone(), actor)
     }
 
@@ -2041,7 +2064,7 @@ impl<A: Actor> Ports<A> {
                 let port = self.mailbox.open_enqueue_port(move |headers, msg: M| {
                     let seq_info = headers.get(SEQ_INFO).cloned();
 
-                    let work = WorkCell::new(move |actor: &mut A, instance: &mut Instance<A>| {
+                    let work = WorkCell::new(move |actor: &mut A, instance: &Instance<A>| {
                         Box::pin(async move {
                             // SAFETY: we guarantee that the passed type_info is for type M.
                             unsafe {

hyperactor_mesh/src/lib.rs

@@ -29,7 +29,7 @@ mod metrics;
 pub mod proc_mesh;
 pub mod reference;
 pub mod resource;
-mod router;
+pub mod router;
 pub mod shared_cell;
 pub mod shortuuid;
 #[cfg(target_os = "linux")]

monarch_extension/src/code_sync.rs

@@ -18,6 +18,7 @@ use hyperactor::context;
 use hyperactor_mesh::Mesh;
 use hyperactor_mesh::RootActorMesh;
 use hyperactor_mesh::shared_cell::SharedCell;
+use monarch_hyperactor;
 use monarch_hyperactor::code_sync::WorkspaceLocation;
 use monarch_hyperactor::code_sync::manager::CodeSyncManager;
 use monarch_hyperactor::code_sync::manager::CodeSyncManagerParams;
@@ -27,8 +28,6 @@ use monarch_hyperactor::code_sync::manager::WorkspaceConfig;
 use monarch_hyperactor::code_sync::manager::WorkspaceShape;
 use monarch_hyperactor::code_sync::manager::code_sync_mesh;
 use monarch_hyperactor::context::PyInstance;
-use monarch_hyperactor::instance_dispatch;
-use monarch_hyperactor::instance_into_dispatch;
 use monarch_hyperactor::proc_mesh::PyProcMesh;
 use monarch_hyperactor::runtime::signal_safe_block_on;
 use monarch_hyperactor::v1::proc_mesh::PyProcMesh as PyProcMeshV1;
@@ -279,32 +278,34 @@ impl CodeSyncMeshClient {
         if let Ok(v0) = proc_mesh.downcast::<PyProcMesh>() {
             let proc_mesh = v0.borrow().try_inner()?;
             signal_safe_block_on(py, async move {
-                let actor_mesh = instance_dispatch!(client, |cx| {
-                    proc_mesh
-                        .spawn(cx, "code_sync_manager", &CodeSyncManagerParams {})
-                        .await?
-                });
+                let actor_mesh = proc_mesh
+                    .spawn(
+                        client.deref(),
+                        "code_sync_manager",
+                        &CodeSyncManagerParams {},
+                    )
+                    .await?;
                 Ok(Self { actor_mesh })
             })?
         } else {
             let proc_mesh = proc_mesh.downcast::<PyProcMeshV1>()?.borrow().mesh_ref()?;
             signal_safe_block_on(py, async move {
-                let actor_mesh = instance_dispatch!(client, |cx| {
-                    proc_mesh
-                        .spawn_service(cx, "code_sync_manager", &CodeSyncManagerParams {})
-                        .await
-                        .map_err(|err| PyException::new_err(err.to_string()))?
-                });
-                instance_dispatch!(client, |cx| {
-                    actor_mesh
-                        .cast(
-                            cx,
-                            SetActorMeshMessage {
-                                actor_mesh: actor_mesh.deref().clone(),
-                            },
-                        )
-                        .map_err(|err| PyException::new_err(err.to_string()))?
-                });
+                let actor_mesh = proc_mesh
+                    .spawn_service(
+                        client.deref(),
+                        "code_sync_manager",
+                        &CodeSyncManagerParams {},
+                    )
+                    .await
+                    .map_err(|err| PyException::new_err(err.to_string()))?;
+                actor_mesh
+                    .cast(
+                        client.deref(),
+                        SetActorMeshMessage {
+                            actor_mesh: actor_mesh.deref().clone(),
+                        },
+                    )
+                    .map_err(|err| PyException::new_err(err.to_string()))?;
                 Ok(Self {
                     actor_mesh: SharedCell::from(RootActorMesh::from(actor_mesh)),
                 })
@@ -324,19 +325,17 @@ impl CodeSyncMeshClient {
     ) -> PyResult<Bound<'py, PyAny>> {
         let instance = instance.clone();
         let actor_mesh = self.actor_mesh.clone();
-        instance_into_dispatch!(instance, |cx| {
-            monarch_hyperactor::runtime::future_into_py(py, async move {
-                CodeSyncMeshClient::sync_workspace_(
-                    &cx,
-                    actor_mesh,
-                    local,
-                    remote,
-                    method.into(),
-                    auto_reload,
-                )
-                .err_into()
-                .await
-            })
+        monarch_hyperactor::runtime::future_into_py(py, async move {
+            CodeSyncMeshClient::sync_workspace_(
+                instance.deref(),
+                actor_mesh,
+                local,
+                remote,
+                method.into(),
+                auto_reload,
+            )
+            .err_into()
+            .await
         })
     }
 
@@ -354,17 +353,15 @@ impl CodeSyncMeshClient {
             py,
             async move {
                 for workspace in workspaces.into_iter() {
-                    instance_dispatch!(instance, async |cx| {
-                        CodeSyncMeshClient::sync_workspace_(
-                            cx,
-                            actor_mesh.clone(),
-                            workspace.local,
-                            workspace.remote,
-                            workspace.method.into(),
-                            auto_reload,
-                        )
-                        .await
-                    })?;
+                    CodeSyncMeshClient::sync_workspace_(
+                        instance.deref(),
+                        actor_mesh.clone(),
+                        workspace.local,
+                        workspace.remote,
+                        workspace.method.into(),
+                        auto_reload,
+                    )
+                    .await?;
                 }
                 anyhow::Ok(())
             }

monarch_extension/src/logging.rs

@@ -8,6 +8,7 @@
 
 #![allow(unsafe_op_in_unsafe_fn)]
 
+use std::ops::Deref;
 use std::time::Duration;
 
 use hyperactor::ActorHandle;
@@ -22,7 +23,6 @@ use hyperactor_mesh::logging::LogForwardMessage;
 use hyperactor_mesh::selection::Selection;
 use hyperactor_mesh::shared_cell::SharedCell;
 use monarch_hyperactor::context::PyInstance;
-use monarch_hyperactor::instance_dispatch;
 use monarch_hyperactor::logging::LoggerRuntimeActor;
 use monarch_hyperactor::logging::LoggerRuntimeMessage;
 use monarch_hyperactor::proc_mesh::PyProcMesh;
@@ -94,13 +94,10 @@ impl LoggingMeshClient {
                 .client_proc()
                 .spawn("log_client", LogClientActor::default())?;
             let client_actor_ref = client_actor.bind();
-            let forwarder_mesh = instance_dispatch!(instance, |cx| {
-                proc_mesh
-                    .spawn(cx, "log_forwarder", &client_actor_ref)
-                    .await?
-            });
-            let logger_mesh =
-                instance_dispatch!(instance, |cx| { proc_mesh.spawn(cx, "logger", &()).await? });
+            let forwarder_mesh = proc_mesh
+                .spawn(instance.deref(), "log_forwarder", &client_actor_ref)
+                .await?;
+            let logger_mesh = proc_mesh.spawn(instance.deref(), "logger", &()).await?;
 
             // Register flush_internal as a on-stop callback
             let client_actor_for_callback = client_actor.clone();

monarch_extension/src/mesh_controller.rs

@@ -46,7 +46,6 @@ use monarch_hyperactor::actor::PythonMessage;
 use monarch_hyperactor::actor::PythonMessageKind;
 use monarch_hyperactor::buffers::FrozenBuffer;
 use monarch_hyperactor::context::PyInstance;
-use monarch_hyperactor::instance_dispatch;
 use monarch_hyperactor::local_state_broker::LocalStateBrokerActor;
 use monarch_hyperactor::mailbox::PyPortId;
 use monarch_hyperactor::ndslice::PySlice;
@@ -140,17 +139,14 @@ impl _Controller {
         let id = NEXT_ID.fetch_add(1, atomic::Ordering::Relaxed);
         let controller_handle: Arc<Mutex<ActorHandle<MeshControllerActor>>> =
             signal_safe_block_on(py, async move {
-                let controller_handle = instance_dispatch!(client, |instance| {
-                    instance.proc().spawn(
-                        &Name::new("mesh_controller").to_string(),
-                        MeshControllerActor::new(MeshControllerActorParams {
-                            proc_mesh,
-                            id,
-                            rank_map,
-                        })
-                        .await,
-                    )?
-                });
+                let controller_handle = client.spawn(
+                    MeshControllerActor::new(MeshControllerActorParams {
+                        proc_mesh,
+                        id,
+                        rank_map,
+                    })
+                    .await,
+                )?;
                 Ok::<_, anyhow::Error>(Arc::new(Mutex::new(controller_handle)))
             })??;
 
@@ -231,8 +227,7 @@ impl _Controller {
     }
 
     fn _drain_and_stop(&mut self, py: Python<'_>, instance: &PyInstance) -> PyResult<()> {
-        let (stop_worker_port, stop_worker_receiver) =
-            instance_dispatch!(instance, |cx_instance| { cx_instance.open_once_port() });
+        let (stop_worker_port, stop_worker_receiver) = instance.open_once_port();
 
         self.controller_handle
             .blocking_lock()
@@ -817,6 +812,10 @@ impl Actor for MeshControllerActor {
         self.brokers = Some(brokers);
         Ok(())
     }
+
+    fn display_name(&self) -> Option<String> {
+        Some(format!("mesh_controller_{}", self.id))
+    }
 }
 
 impl Debug for MeshControllerActor {