Lazy NIC device init + integration of NIC autodetection. (#1444)

Summary:
Pull Request resolved: #1444

Rewrites RdmaManager internals wrt:

Before - we select a given RDMA-compliant device on (on init via config) and use that for all network traffic.

Now - RdmaManager is a general resource manager for a given process.   For a given memory virtual addr, we will select the closest RDMA device (pci-distance), and setup the relevant QPs, pd, on the fly.   So a given manager could be sending on traffic on "N" rdma devices.

To make this work, we need to add a 'device_name' field to RdmaBuffer, so its a combination of other_id + other_device; introducing a many to many connection setup between two rdma managers.

Reviewed By: casteryh

Differential Revision: D83918488

fbshipit-source-id: f8b7e5cd1abe29e08afbcc526ab02053c9fa12f5

Author: dstaay-fb

monarch_rdma/examples/cuda_ping_pong/src/cuda_ping_pong.rs

@@ -468,7 +468,12 @@ impl Handler<PerformPingPong> for CudaRdmaActor {
         }
         let qp = self
             .rdma_manager
-            .request_queue_pair(cx, remote_buffer.owner.clone())
+            .request_queue_pair(
+                cx,
+                remote_buffer.owner.clone(),
+                local_buffer.device_name.clone(),
+                remote_buffer.device_name.clone(),
+            )
             .await?;
 
         unsafe {

monarch_rdma/src/device_selection.rs

@@ -439,6 +439,31 @@ pub fn select_optimal_rdma_device(device_hint: Option<&str>) -> Option<RdmaDevic
     }
 }
 
+/// Creates a mapping from CUDA PCI addresses to optimal RDMA devices
+///
+/// This function discovers all available CUDA devices and determines the best
+/// RDMA device for each one using the device selection algorithm.
+///
+/// # Returns
+///
+/// * `HashMap<String, RdmaDevice>` - Map from CUDA PCI address to optimal RDMA device
+pub fn create_cuda_to_rdma_mapping() -> HashMap<String, RdmaDevice> {
+    let mut mapping = HashMap::new();
+
+    // Try to discover CUDA devices (GPU 0-8 should be sufficient for most systems)
+    for gpu_idx in 0..8 {
+        let gpu_idx_str = gpu_idx.to_string();
+        if let Some(cuda_pci_addr) = get_cuda_pci_address(&gpu_idx_str) {
+            let cuda_hint = format!("cuda:{}", gpu_idx);
+            if let Some(rdma_device) = select_optimal_rdma_device(Some(&cuda_hint)) {
+                mapping.insert(cuda_pci_addr, rdma_device);
+            }
+        }
+    }
+
+    mapping
+}
+
 /// Resolves RDMA device using auto-detection logic when needed
 ///
 /// This function applies auto-detection for default devices, but otherwise  

monarch_rdma/src/rdma_components.rs

@@ -101,6 +101,7 @@ pub struct RdmaBuffer {
     pub rkey: u32,
     pub addr: usize,
     pub size: usize,
+    pub device_name: String,
 }
 
 impl RdmaBuffer {
@@ -128,10 +129,18 @@ impl RdmaBuffer {
             remote.owner.actor_id(),
             remote,
         );
-        let remote_owner = remote.owner.clone(); // Clone before the move!
+        let remote_owner = remote.owner.clone();
+
+        let local_device = self.device_name.clone();
+        let remote_device = remote.device_name.clone();
         let mut qp = self
             .owner
-            .request_queue_pair_deprecated(client, remote_owner.clone())
+            .request_queue_pair_deprecated(
+                client,
+                remote_owner.clone(),
+                local_device.clone(),
+                remote_device.clone(),
+            )
             .await?;
 
         qp.put(self.clone(), remote)?;
@@ -141,7 +150,7 @@ impl RdmaBuffer {
 
         // Release the queue pair back to the actor
         self.owner
-            .release_queue_pair_deprecated(client, remote_owner, qp)
+            .release_queue_pair_deprecated(client, remote_owner, local_device, remote_device, qp)
             .await?;
 
         result
@@ -173,9 +182,19 @@ impl RdmaBuffer {
             remote,
         );
         let remote_owner = remote.owner.clone(); // Clone before the move!
+
+        // Extract device name from buffer, fallback to a default if not present
+        let local_device = self.device_name.clone();
+        let remote_device = remote.device_name.clone();
+
         let mut qp = self
             .owner
-            .request_queue_pair_deprecated(client, remote_owner.clone())
+            .request_queue_pair_deprecated(
+                client,
+                remote_owner.clone(),
+                local_device.clone(),
+                remote_device.clone(),
+            )
             .await?;
         qp.get(self.clone(), remote)?;
         let result = self
@@ -184,7 +203,7 @@ impl RdmaBuffer {
 
         // Release the queue pair back to the actor
         self.owner
-            .release_queue_pair_deprecated(client, remote_owner, qp)
+            .release_queue_pair_deprecated(client, remote_owner, local_device, remote_device, qp)
             .await?;
 
         result?;
@@ -877,6 +896,7 @@ impl RdmaQueuePair {
             if (wqe_cnt as u64) < (self.send_wqe_idx - self.send_db_idx) {
                 return Err(anyhow::anyhow!("Overflow of WQE, possible data loss"));
             }
+            self.apply_first_op_delay(self.send_db_idx);
             while self.send_db_idx < self.send_wqe_idx {
                 let offset = (self.send_db_idx % wqe_cnt as u64) * stride as u64;
                 let src_ptr = (base_ptr as *mut u8).wrapping_add(offset as usize);