Make SharedSpawnable support any non-borrow deref (#393)

Summary:
Pull Request resolved: #393

This makes `SharedSpawnable` more generic and implements it for any type
that can deref to a `ProcMesh`.

The only side-effect here is that it requires the caller to do the `Arc::clone()`,
rather than implicitly doing it in the impl.

Reviewed By: mariusae

Differential Revision: D77378974

fbshipit-source-id: 7db36cf6b83c29e317b80ad49ca34747aa60bc73

Author: andrewjcg

hyperactor_mesh/src/actor_mesh.rs

@@ -9,7 +9,6 @@
 #![allow(dead_code)] // until used publically
 
 use std::ops::Deref;
-use std::sync::Arc;
 
 use async_trait::async_trait;
 use hyperactor::Actor;
@@ -164,8 +163,8 @@ pub trait ActorMesh: Mesh {
 /// Given a shared ProcMesh, we can obtain a [`ActorMesh<'static, _>`]
 /// for it, useful when lifetime must be managed dynamically.
 enum ProcMeshRef<'a> {
-    /// The reference is shared with an [`Arc`].
-    Shared(Arc<ProcMesh>),
+    /// The reference is shared without requiring a reference.
+    Shared(Box<dyn Deref<Target = ProcMesh> + Sync + Send>),
     /// The reference is borrowed with a parameterized
     /// lifetime.
     Borrowed(&'a ProcMesh),
@@ -206,14 +205,14 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
         }
     }
 
-    pub(crate) fn new_shared(
-        proc_mesh: Arc<ProcMesh>,
+    pub(crate) fn new_shared<D: Deref<Target = ProcMesh> + Send + Sync + 'static>(
+        proc_mesh: D,
         name: String,
         actor_supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
         ranks: Vec<ActorRef<A>>,
     ) -> Self {
         Self {
-            proc_mesh: ProcMeshRef::Shared(proc_mesh),
+            proc_mesh: ProcMeshRef::Shared(Box::new(proc_mesh)),
             name,
             ranks,
             actor_supervision_rx,
@@ -499,6 +498,7 @@ pub(crate) mod test_util {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
 
     use hyperactor::ActorId;
     use hyperactor::PortRef;

hyperactor_mesh/src/comm.rs

@@ -738,6 +738,7 @@ mod tests {
             forward_port: tx.bind(),
         };
         let actor_mesh = proc_mesh
+            .clone()
             .spawn::<TestActor>(dest_actor_name, &params)
             .await
             .unwrap();

hyperactor_mesh/src/proc_mesh.rs

@@ -9,6 +9,7 @@
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::fmt;
+use std::ops::Deref;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -532,7 +533,7 @@ impl ProcEvents {
 #[async_trait]
 pub trait SharedSpawnable {
     async fn spawn<A: Actor + RemoteActor>(
-        &self,
+        self,
         actor_name: &str,
         params: &A::Params,
     ) -> Result<RootActorMesh<'static, A>, anyhow::Error>
@@ -541,9 +542,9 @@ pub trait SharedSpawnable {
 }
 
 #[async_trait]
-impl SharedSpawnable for Arc<ProcMesh> {
+impl<D: Deref<Target = ProcMesh> + Send + Sync + 'static> SharedSpawnable for D {
     async fn spawn<A: Actor + RemoteActor>(
-        &self,
+        self,
         actor_name: &str,
         params: &A::Params,
     ) -> Result<RootActorMesh<'static, A>, anyhow::Error>
@@ -555,11 +556,13 @@ impl SharedSpawnable for Arc<ProcMesh> {
             // Instantiate supervision routing BEFORE spawning the actor mesh.
             self.actor_event_router.insert(actor_name.to_string(), tx);
         }
+        let ranks =
+            ProcMesh::spawn_on_procs::<A>(&self.client, self.agents(), actor_name, params).await?;
         Ok(RootActorMesh::new_shared(
-            Arc::clone(self),
+            self,
             actor_name.to_string(),
             rx,
-            ProcMesh::spawn_on_procs::<A>(&self.client, self.agents(), actor_name, params).await?,
+            ranks,
         ))
     }
 }

monarch_extension/src/mesh_controller.rs

@@ -184,6 +184,7 @@ impl _Controller {
         let workers: anyhow::Result<SharedCell<RootActorMesh<'_, WorkerActor>>> =
             signal_safe_block_on(py, async move {
                 let workers = py_proc_mesh
+                    .clone()
                     .spawn(&format!("tensor_engine_workers_{}", id), &param)
                     .await?;
                 //workers.cast(ndslice::Selection::True, )?;

monarch_hyperactor/src/proc_mesh.rs

@@ -228,12 +228,11 @@ impl PyProcMesh {
         let proc_mesh = Arc::clone(&self.inner);
         let keepalive = self.keepalive.clone();
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            let mailbox = proc_mesh.client().clone();
             let actor_mesh = proc_mesh.spawn(&name, &pickled_type).await?;
             let python_actor_mesh = PythonActorMesh {
                 inner: actor_mesh,
-                client: PyMailbox {
-                    inner: proc_mesh.client().clone(),
-                },
+                client: PyMailbox { inner: mailbox },
                 _keepalive: keepalive,
             };
             Python::with_gil(|py| python_actor_mesh.into_py_any(py))
@@ -250,12 +249,11 @@ impl PyProcMesh {
         let proc_mesh = Arc::clone(&self.inner);
         let keepalive = self.keepalive.clone();
         signal_safe_block_on(py, async move {
+            let mailbox = proc_mesh.client().clone();
             let actor_mesh = proc_mesh.spawn(&name, &pickled_type).await?;
             let python_actor_mesh = PythonActorMesh {
                 inner: actor_mesh,
-                client: PyMailbox {
-                    inner: proc_mesh.client().clone(),
-                },
+                client: PyMailbox { inner: mailbox },
                 _keepalive: keepalive,
             };
             Python::with_gil(|py| python_actor_mesh.into_py_any(py))