[monarch] mesh: support v1 actor meshes in RDMA buffer

Pull Request resolved: #1467

This uses the v0 shim to support v1 meshes in RDMA buffer.
ghstack-source-id: 315017598

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

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

Author: mariusae

monarch_hyperactor/src/v1/proc_mesh.rs

@@ -36,7 +36,7 @@ use crate::v1::actor_mesh::PythonActorMeshImpl;
     name = "ProcMesh",
     module = "monarch._rust_bindings.monarch_hyperactor.v1.proc_mesh"
 )]
-pub(crate) enum PyProcMesh {
+pub enum PyProcMesh {
     Owned(PyProcMeshImpl),
     Ref(PyProcMeshRefImpl),
 }
@@ -50,7 +50,7 @@ impl PyProcMesh {
         Self::Ref(PyProcMeshRefImpl(inner))
     }
 
-    pub(crate) fn mesh_ref(&self) -> Result<ProcMeshRef, anyhow::Error> {
+    pub fn mesh_ref(&self) -> Result<ProcMeshRef, anyhow::Error> {
         match self {
             PyProcMesh::Owned(inner) => Ok(inner.0.borrow()?.clone()),
             PyProcMesh::Ref(inner) => Ok(inner.0.clone()),
@@ -195,7 +195,7 @@ impl PyProcMesh {
     name = "ProcMeshImpl",
     module = "monarch._rust_bindings.monarch_hyperactor.v1.proc_mesh"
 )]
-pub(crate) struct PyProcMeshImpl(SharedCell<ProcMesh>);
+pub struct PyProcMeshImpl(SharedCell<ProcMesh>);
 
 impl PyProcMeshImpl {
     fn __repr__(&self) -> PyResult<String> {
@@ -211,7 +211,7 @@ impl PyProcMeshImpl {
     name = "ProcMeshRefImpl",
     module = "monarch._rust_bindings.monarch_hyperactor.v1.proc_mesh"
 )]
-pub(crate) struct PyProcMeshRefImpl(ProcMeshRef);
+pub struct PyProcMeshRefImpl(ProcMeshRef);
 
 impl PyProcMeshRefImpl {
     fn __repr__(&self) -> PyResult<String> {

monarch_rdma/extension/lib.rs

@@ -18,6 +18,7 @@ use monarch_hyperactor::instance_dispatch;
 use monarch_hyperactor::proc_mesh::PyProcMesh;
 use monarch_hyperactor::pytokio::PyPythonTask;
 use monarch_hyperactor::runtime::signal_safe_block_on;
+use monarch_hyperactor::v1::proc_mesh::PyProcMesh as PyProcMeshV1;
 use monarch_rdma::RdmaBuffer;
 use monarch_rdma::RdmaManagerActor;
 use monarch_rdma::RdmaManagerMessageClient;
@@ -289,31 +290,52 @@ impl PyRdmaManager {
     #[classmethod]
     fn create_rdma_manager_nonblocking(
         _cls: &Bound<'_, PyType>,
-        proc_mesh: &PyProcMesh,
+        proc_mesh: &Bound<'_, PyAny>,
         client: PyInstance,
     ) -> PyResult<PyPythonTask> {
         tracing::debug!("spawning RDMA manager on target proc_mesh nodes");
 
-        let tracked_proc_mesh = proc_mesh.try_inner()?;
+        if let Ok(v0) = proc_mesh.downcast::<PyProcMesh>() {
+            let tracked_proc_mesh = v0.borrow().try_inner()?;
+            PyPythonTask::new(async move {
+                // Spawns the `RdmaManagerActor` on the target proc_mesh.
+                // This allows the `RdmaController` to run on any node while real RDMA operations occur on appropriate hardware.
+                let actor_mesh = instance_dispatch!(client, |cx| {
+                    tracked_proc_mesh
+                        // Pass None to use default config - RdmaManagerActor will use default IbverbsConfig
+                        // TODO - make IbverbsConfig configurable
+                        .spawn::<RdmaManagerActor>(cx, "rdma_manager", &None)
+                        .await
+                        .map_err(|err| PyException::new_err(err.to_string()))?
+                });
 
-        PyPythonTask::new(async move {
-            // Spawns the `RdmaManagerActor` on the target proc_mesh.
-            // This allows the `RdmaController` to run on any node while real RDMA operations occur on appropriate hardware.
-            let actor_mesh = instance_dispatch!(client, |cx| {
-                tracked_proc_mesh
-                    // Pass None to use default config - RdmaManagerActor will use default IbverbsConfig
-                    // TODO - make IbverbsConfig configurable
-                    .spawn::<RdmaManagerActor>(cx, "rdma_manager", &None)
-                    .await
-                    .map_err(|err| PyException::new_err(err.to_string()))?
-            });
+                // Use placeholder device name since actual device is determined on remote node
+                Ok(Some(PyRdmaManager {
+                    inner: actor_mesh,
+                    device: "remote_rdma_device".to_string(),
+                }))
+            })
+        } else {
+            let proc_mesh = proc_mesh.downcast::<PyProcMeshV1>()?.borrow().mesh_ref()?;
+            PyPythonTask::new(async move {
+                let actor_mesh = instance_dispatch!(client, |cx| {
+                    proc_mesh
+                        // Pass None to use default config - RdmaManagerActor will use default IbverbsConfig
+                        // TODO - make IbverbsConfig configurable
+                        .spawn::<RdmaManagerActor>(cx, "rdma_manager", &None)
+                        .await
+                        .map_err(|err| PyException::new_err(err.to_string()))?
+                });
 
-            // Use placeholder device name since actual device is determined on remote node
-            Ok(Some(PyRdmaManager {
-                inner: actor_mesh,
-                device: "remote_rdma_device".to_string(),
-            }))
-        })
+                let actor_mesh = RootActorMesh::from(actor_mesh);
+                let actor_mesh = SharedCell::from(actor_mesh);
+
+                Ok(Some(PyRdmaManager {
+                    inner: actor_mesh,
+                    device: "remote_rdma_device".to_string(),
+                }))
+            })
+        }
     }
 }
 

python/monarch/_src/rdma/rdma.py

@@ -27,7 +27,14 @@
 from monarch._src.actor.actor_mesh import Actor, context
 from monarch._src.actor.endpoint import endpoint
 from monarch._src.actor.future import Future
-from monarch._src.actor.proc_mesh import get_or_spawn_controller, ProcMesh
+from monarch._src.actor.proc_mesh import (
+    get_or_spawn_controller as get_or_spawn_controller_v0,
+    ProcMesh as ProcMeshV0,
+)
+from monarch._src.actor.v1.proc_mesh import (
+    get_or_spawn_controller as get_or_spawn_controller_v1,
+    ProcMesh as ProcMeshV1,
+)
 from pyre_extensions import none_throws
 
 
@@ -54,13 +61,17 @@ def is_rdma_available():
 @functools.cache
 def _ensure_init_rdma_manager() -> Shared[None]:
     async def task() -> None:
-        await (
-            await get_or_spawn_controller("rdma_controller", RdmaController)
-        ).init_rdma_on_mesh.call_one(
-            # FIXME(slurye): Fix this once controller API is working properly
-            # for v1.
-            cast(ProcMesh, none_throws(context().actor_instance.proc_mesh))
-        )
+        proc_mesh = context().actor_instance.proc_mesh
+        if isinstance(proc_mesh, ProcMeshV0):
+            controller = await get_or_spawn_controller_v0(
+                "rdma_controller", RdmaController
+            )
+        else:
+            controller = await get_or_spawn_controller_v1(
+                "rdma_controller", RdmaController
+            )
+
+        await controller.init_rdma_on_mesh.call_one(proc_mesh)
 
     return PythonTask.from_coroutine(task()).spawn()
 
@@ -120,10 +131,10 @@ def _get_addr_and_size(buf: torch.Tensor | memoryview) -> tuple[int, int]:
 
 class RdmaController(Actor):
     def __init__(self) -> None:
-        self._manager_futures: Dict[ProcMesh, Future[_RdmaManager]] = {}
+        self._manager_futures: Dict[ProcMeshV0 | ProcMeshV1, Future[_RdmaManager]] = {}
 
     @endpoint
-    async def init_rdma_on_mesh(self, proc_mesh: ProcMesh) -> None:
+    async def init_rdma_on_mesh(self, proc_mesh: ProcMeshV0 | ProcMeshV1) -> None:
         # Note: RdmaController acts as coordinator and can run on any node
         # The RDMA support check should happen on the target proc_mesh nodes, not on RdmaController's node