Expose RDMA support through Python APIs (#462)

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

examples/grpo_actor.py

@@ -15,7 +15,7 @@
 import torch.optim as optim
 
 from monarch.actor import Actor, ActorMeshRef, endpoint, proc_mesh
-from monarch.rdma import RDMABuffer
+from monarch.tensor_engine import RDMABuffer
 from torch.distributions import Categorical, kl_divergence
 
 """

mode/tensor_engine

@@ -0,0 +1,2 @@
+--config
+monarch.enable_tensor_engine=true

monarch_extension/Cargo.toml

@@ -14,30 +14,14 @@ doctest = false
 crate-type = ["cdylib"]
 
 [dependencies]
-anyhow = "1.0.98"
-async-trait = "0.1.86"
-bincode = "1.3.3"
-clap = { version = "4.5.38", features = ["derive", "env", "string", "unicode", "wrap_help"] }
-controller = { version = "0.0.0", path = "../controller", optional = true }
-hyperactor = { version = "0.0.0", path = "../hyperactor" }
-hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }
-hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiprocess" }
-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_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" }
+controller = { path = "../controller", optional = true }
+monarch_messages = { path = "../monarch_messages", optional = true }
+monarch_simulator_lib = { path = "../monarch_simulator", optional = true }
+monarch_tensor_worker = { path = "../monarch_tensor_worker", optional = true }
 nccl-sys = { path = "../nccl-sys", optional = true }
-ndslice = { version = "0.0.0", path = "../ndslice" }
-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"] }
-tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
-torch-sys = { version = "0.0.0", path = "../torch-sys", optional = true }
-torch-sys-cuda = { version = "0.0.0", path = "../torch-sys-cuda", optional = true }
-tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
+torch-sys = { path = "../torch-sys", optional = true }
+torch-sys-cuda = { path = "../torch-sys-cuda", optional = true }
 
 [features]
-default = ["tensor_engine"]
-tensor_engine = ["dep:controller", "dep:monarch_messages", "dep:monarch_simulator_lib", "dep:monarch_tensor_worker", "dep:nccl-sys", "dep:torch-sys", "dep:torch-sys-cuda"]
+default = []
+tensor_engine = ["dep:controller", "dep:monarch_messages", "dep:monarch_simulator_lib", "dep:monarch_tensor_worker", "dep:nccl-sys", "dep:torch-sys", "dep:torch-sys-cuda", "fbcode//monarch/monarch_rdma/extension:monarch_rdma_extension"]

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",
     )?)?;
+
     monarch_hyperactor::telemetry::register_python_bindings(&get_or_add_new_module(
         module,
         "monarch_hyperactor.telemetry",

monarch_hyperactor/Cargo.toml

@@ -9,22 +9,15 @@ license = "BSD-3-Clause"
 
 [dependencies]
 anyhow = "1.0.98"
-async-trait = "0.1.86"
-bincode = "1.3.3"
 clap = { version = "4.5.38", features = ["derive", "env", "string", "unicode", "wrap_help"] }
-erased-serde = "0.3.27"
-fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
 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_types = { version = "0.0.0", path = "../monarch_types" }
-ndslice = { version = "0.0.0", path = "../ndslice" }
-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_bytes = "0.11"
-thiserror = "2.0.12"
 tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
 tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
+
+[dev-dependencies]
+ndslice = { version = "0.0.0", path = "../ndslice" }
+
+[features]
+default = []
+tensor_engine = ["dep:monarch_rdma"]

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

@@ -27,6 +27,10 @@ use hyperactor_mesh::proc_mesh::SharedSpawnable;
 use hyperactor_mesh::shared_cell::SharedCell;
 use hyperactor_mesh::shared_cell::SharedCellPool;
 use hyperactor_mesh::shared_cell::SharedCellRef;
+#[cfg(feature = "tensor_engine")]
+use monarch_rdma::IbverbsConfig;
+#[cfg(feature = "tensor_engine")]
+use monarch_rdma::RdmaManagerActor;
 use monarch_types::PickledPyObject;
 use ndslice::Shape;
 use pyo3::IntoPyObjectExt;
@@ -116,6 +120,8 @@ pub struct PyProcMesh {
     proc_events: SharedCell<Mutex<ProcEvents>>,
     stop_monitor_sender: mpsc::Sender<bool>,
     user_monitor_registered: AtomicBool,
+    #[cfg(feature = "tensor_engine")]
+    pub(super) rdma_manager: Option<SharedCell<RootActorMesh<'static, RdmaManagerActor>>>,
 }
 
 fn allocate_proc_mesh<'py>(py: Python<'py>, alloc: &PyAlloc) -> PyResult<Bound<'py, PyAny>> {
@@ -132,7 +138,7 @@ 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).await
     })
 }
 
@@ -150,14 +156,14 @@ 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).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) -> 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 +173,34 @@ impl PyProcMesh {
             world_id,
             abort_receiver,
         ));
-        Self {
-            inner: SharedCell::from(TrackedProcMesh::from(proc_mesh)),
+
+        let tracked_proc_mesh: TrackedProcMesh = TrackedProcMesh::from(proc_mesh);
+
+        // Create optional RDMA manager
+        #[cfg(feature = "tensor_engine")]
+        let rdma_manager = if 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 {
+            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),
-        }
+            #[cfg(feature = "tensor_engine")]
+            rdma_manager,
+        })
     }
 
     /// The default monitor of the proc mesh for crashes. If a proc crashes, we print the reason

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"] }