Expose RDMA support through Python APIs

Summary:
This change brings in RDMA support from Rust to the Python APIs, overwriting the existing stand-in RDMABuffer APIs.

There are quite a few changes, so here's a high level overview:

## `monarch_hyperactor` (rust)
Within `PyProcMesh`, adds an optional RDMAMangerActor attribute.

This mimics the prior functionality - before, a `proc_mesh` (in Python) would unconditionally spin up an `RdmaManagerActor`.

Now, a `proc_mesh` will spin up an `RdmaManagerActor` if and only if `tensor_engine` is enabled and `ibverbs` APIs are supported.

A few other pieces must be accessible within the RDMA extension library - specifically `Mailbox` / the send caps capabiliites.

## RDMA extension
Adds an `extension` folder within `monarch_rdma` which contains the actual bindings.

Consistent with other bindings, provides a blocking and non-blocking version.


## RDMA components
Corrects behavior of a few edge cases, i.e. `ibverbs_supported` based on the number of devices, supporting a "loopback" case, wherein two actors share an RDMA buffer but are spawned on the same proc.


## Python
Removes the stand-in RDMAManagerActor

In rdma.py, adds in basic support for RDMABuffer CPU.

While CUDA<>CPU|CUDA is implemented in `monarch_rdma`, this was running into issues with MR registration.

Logging shows that the code path correctly interprets the address as CUDA, but `cuMemGetHandleForAddressRange` returns the handle/fd as `-1` and a null pointer MR.

RDMABuffer will only support CPU for now with GPU support as a follow up, cc dstaay-fb


# Test
Updates tests for anything touching RDMABuffer. Splits out `test_python_actor` into `test_python_actor` and `test_rdma` now that RDMA is utilizing the backend network.

Differential Revision: D76937776

Author: allenwang28

monarch_extension/Cargo.toml

@@ -26,6 +26,7 @@ hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiproces
 libc = "0.2.139"
 monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" }
 monarch_messages = { version = "0.0.0", path = "../monarch_messages", optional = true }
+monarch_rdma_extension = { version = "0.0.0", path = "../monarch_rdma/extension" }
 monarch_simulator_lib = { version = "0.0.0", path = "../monarch_simulator", optional = true }
 monarch_tensor_worker = { version = "0.0.0", path = "../monarch_tensor_worker", optional = true }
 monarch_types = { version = "0.0.0", path = "../monarch_types" }

monarch_extension/src/lib.rs

@@ -120,6 +120,7 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
             module,
             "monarch_extension.mesh_controller",
         )?)?;
+        monarch_rdma_extension::register_python_bindings(&get_or_add_new_module(module, "rdma")?)?;
     }
     simulation_tools::register_python_bindings(&get_or_add_new_module(
         module,
@@ -166,6 +167,7 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
         module,
         "monarch_hyperactor.runtime",
     )?)?;
+
     hyperactor_extension::alloc::register_python_bindings(&get_or_add_new_module(
         module,
         "hyperactor_extension.alloc",

monarch_hyperactor/Cargo.toml

@@ -20,6 +20,7 @@ hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }
 hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiprocess" }
 hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" }
 inventory = "0.3.8"
+monarch_rdma = { version = "0.0.0", path = "../monarch_rdma" }
 monarch_types = { version = "0.0.0", path = "../monarch_types" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
 pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] }

monarch_hyperactor/src/actor_mesh.rs

@@ -62,7 +62,7 @@ impl PythonActorMesh {
     }
 
     #[getter]
-    fn client(&self) -> PyMailbox {
+    pub fn client(&self) -> PyMailbox {
         self.client.clone()
     }
 

monarch_hyperactor/src/mailbox.rs

@@ -52,8 +52,8 @@ use crate::shape::PyShape;
     name = "Mailbox",
     module = "monarch._rust_bindings.monarch_hyperactor.mailbox"
 )]
-pub(super) struct PyMailbox {
-    pub(super) inner: Mailbox,
+pub struct PyMailbox {
+    pub inner: Mailbox,
 }
 
 #[pymethods]

monarch_hyperactor/src/proc_mesh.rs

@@ -28,6 +28,8 @@ use hyperactor_mesh::proc_mesh::SharedSpawnable;
 use hyperactor_mesh::shared_cell::SharedCell;
 use hyperactor_mesh::shared_cell::SharedCellPool;
 use hyperactor_mesh::shared_cell::SharedCellRef;
+use monarch_rdma::IbverbsConfig;
+use monarch_rdma::RdmaManagerActor;
 use monarch_types::PickledPyObject;
 use ndslice::Shape;
 use pyo3::IntoPyObjectExt;
@@ -116,9 +118,14 @@ pub struct PyProcMesh {
     proc_events: SharedCell<Mutex<ProcEvents>>,
     stop_monitor_sender: mpsc::Sender<bool>,
     user_monitor_registered: AtomicBool,
+    pub(super) rdma_manager: Option<SharedCell<RootActorMesh<'static, RdmaManagerActor>>>,
 }
 
-fn allocate_proc_mesh<'py>(py: Python<'py>, alloc: &PyAlloc) -> PyResult<Bound<'py, PyAny>> {
+fn allocate_proc_mesh<'py>(
+    py: Python<'py>,
+    alloc: &PyAlloc,
+    has_tensor_engine: Option<bool>,
+) -> PyResult<Bound<'py, PyAny>> {
     let alloc = match alloc.take() {
         Some(alloc) => alloc,
         None => {
@@ -132,11 +139,15 @@ fn allocate_proc_mesh<'py>(py: Python<'py>, alloc: &PyAlloc) -> PyResult<Bound<'
         let mesh = ProcMesh::allocate(alloc)
             .await
             .map_err(|err| PyException::new_err(err.to_string()))?;
-        Ok(PyProcMesh::monitored(mesh, world_id))
+        PyProcMesh::monitored(mesh, world_id, has_tensor_engine.unwrap_or(false)).await
     })
 }
 
-fn allocate_proc_mesh_blocking<'py>(py: Python<'py>, alloc: &PyAlloc) -> PyResult<PyProcMesh> {
+fn allocate_proc_mesh_blocking<'py>(
+    py: Python<'py>,
+    alloc: &PyAlloc,
+    has_tensor_engine: Option<bool>,
+) -> PyResult<PyProcMesh> {
     let alloc = match alloc.take() {
         Some(alloc) => alloc,
         None => {
@@ -150,14 +161,18 @@ fn allocate_proc_mesh_blocking<'py>(py: Python<'py>, alloc: &PyAlloc) -> PyResul
         let mesh = ProcMesh::allocate(alloc)
             .await
             .map_err(|err| PyException::new_err(err.to_string()))?;
-        Ok(PyProcMesh::monitored(mesh, world_id))
+        PyProcMesh::monitored(mesh, world_id, has_tensor_engine.unwrap_or(false)).await
     })?
 }
 
 impl PyProcMesh {
     /// Create a new [`PyProcMesh`] with a monitor that crashes the
     /// process on any proc failure.
-    fn monitored(mut proc_mesh: ProcMesh, world_id: WorldId) -> Self {
+    async fn monitored(
+        mut proc_mesh: ProcMesh,
+        world_id: WorldId,
+        has_tensor_engine: bool,
+    ) -> Result<Self, PyErr> {
         let (sender, abort_receiver) = mpsc::channel::<bool>(1);
         let proc_events = SharedCell::from(Mutex::new(proc_mesh.events().unwrap()));
         let monitor = tokio::spawn(Self::default_proc_mesh_monitor(
@@ -167,13 +182,34 @@ impl PyProcMesh {
             world_id,
             abort_receiver,
         ));
-        Self {
-            inner: SharedCell::from(TrackedProcMesh::from(proc_mesh)),
+
+        let tracked_proc_mesh = TrackedProcMesh::from(proc_mesh);
+
+        // Create optional RDMA manager
+        let rdma_manager = if has_tensor_engine && monarch_rdma::ibverbs_supported() {
+            // TODO - make this configurable
+            let config = IbverbsConfig::default();
+            tracing::debug!("rdma is enabled, using device {}", config.device);
+            let actor_mesh = tracked_proc_mesh
+                .spawn("rdma_manager", &config)
+                .await
+                .map_err(|err| PyException::new_err(err.to_string()))?;
+            Some(actor_mesh)
+        } else {
+            if has_tensor_engine {
+                tracing::info!("rdma is not enabled on this hardware");
+            }
+            None
+        };
+
+        Ok(Self {
+            inner: SharedCell::from(tracked_proc_mesh),
             keepalive: Keepalive::new(monitor),
             proc_events,
             stop_monitor_sender: sender,
             user_monitor_registered: AtomicBool::new(false),
-        }
+            rdma_manager,
+        })
     }
 
     /// The default monitor of the proc mesh for crashes. If a proc crashes, we print the reason
@@ -224,21 +260,25 @@ impl PyProcMesh {
 #[pymethods]
 impl PyProcMesh {
     #[classmethod]
+    #[pyo3(signature = (alloc, has_tensor_engine=false))]
     fn allocate_nonblocking<'py>(
         _cls: &Bound<'_, PyType>,
         py: Python<'py>,
         alloc: &PyAlloc,
+        has_tensor_engine: bool,
     ) -> PyResult<Bound<'py, PyAny>> {
-        allocate_proc_mesh(py, alloc)
+        allocate_proc_mesh(py, alloc, Some(has_tensor_engine))
     }
 
     #[classmethod]
+    #[pyo3(signature = (alloc, has_tensor_engine=false))]
     fn allocate_blocking<'py>(
         _cls: &Bound<'_, PyType>,
         py: Python<'py>,
         alloc: &PyAlloc,
+        has_tensor_engine: bool,
     ) -> PyResult<PyProcMesh> {
-        allocate_proc_mesh_blocking(py, alloc)
+        allocate_proc_mesh_blocking(py, alloc, Some(has_tensor_engine))
     }
 
     fn spawn_nonblocking<'py>(

monarch_rdma/extension/Cargo.toml

@@ -0,0 +1,20 @@
+# @generated by autocargo from //monarch/monarch_rdma/extension:monarch_rdma_extension
+
+[package]
+name = "monarch_rdma_extension"
+version = "0.0.0"
+authors = ["Meta"]
+edition = "2021"
+license = "BSD-3-Clause"
+
+[lib]
+path = "lib.rs"
+
+[dependencies]
+hyperactor = { version = "0.0.0", path = "../../hyperactor" }
+monarch_hyperactor = { version = "0.0.0", path = "../../monarch_hyperactor" }
+monarch_rdma = { version = "0.0.0", path = ".." }
+pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] }
+pyo3-async-runtimes = { version = "0.24", features = ["attributes", "tokio-runtime"] }
+serde = { version = "1.0.185", features = ["derive", "rc"] }
+serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "unbounded_depth"] }