monarch_rdma/
rdma_components.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! # RDMA Components
10//!
11//! This module provides the core RDMA building blocks for establishing and managing RDMA connections.
12//!
13//! ## Core Components
14//!
15//! * `IbvDomain` - Manages RDMA resources including context, protection domain, and memory region
16//! * `IbvQueuePair` - Handles communication between endpoints via queue pairs and completion queues
17//!
18//! ## RDMA Overview
19//!
20//! Remote Direct Memory Access (RDMA) allows direct memory access from the memory of one computer
21//! into the memory of another without involving either computer's operating system. This permits
22//! high-throughput, low-latency networking with minimal CPU overhead.
23//!
24//! ## Connection Architecture
25//!
26//! The module manages the following ibverbs primitives:
27//!
28//! 1. **Queue Pairs (QP)**: Each connection has a send queue and a receive queue
29//! 2. **Completion Queues (CQ)**: Events are reported when operations complete
30//! 3. **Memory Regions (MR)**: Memory must be registered with the RDMA device before use
31//! 4. **Protection Domains (PD)**: Provide isolation between different connections
32//!
33//! ## Connection Lifecycle
34//!
35//! 1. Create an `IbvDomain` with `new()`
36//! 2. Create an `IbvQueuePair` from the domain
37//! 3. Exchange connection info with remote peer (application must handle this)
38//! 4. Connect to remote endpoint with `connect()`
39//! 5. Perform RDMA operations (read/write)
40//! 6. Poll for completions
41//! 7. Resources are cleaned up when dropped
42
43/// Maximum size for a single RDMA operation in bytes (1 GiB)
44use std::fs;
45use std::result::Result;
46use std::time::Duration;
47
48use hyperactor::ActorRef;
49use hyperactor::context;
50use serde::Deserialize;
51use serde::Serialize;
52use typeuri::Named;
53
54use crate::RdmaAction;
55use crate::RdmaManagerActor;
56use crate::ReleaseBufferClient;
57use crate::backend::RdmaRemoteBackends;
58use crate::local_memory::KeepaliveLocalMemory;
59
60/// Lightweight handle representing a registered RDMA buffer.
61///
62/// Contains an id for the buffer registration, the buffer size, a reference
63/// to the owning [`RdmaManagerActor`], and backend-specific contexts for
64/// performing RDMA operations.
65#[derive(Debug, Named, Clone, Serialize, Deserialize)]
66pub struct RdmaRemoteBuffer {
67    pub id: usize,
68    pub size: usize,
69    pub owner: ActorRef<RdmaManagerActor>,
70    pub(crate) backends: RdmaRemoteBackends,
71}
72wirevalue::register_type!(RdmaRemoteBuffer);
73
74impl RdmaRemoteBuffer {
75    /// Push data from local memory into this remote buffer (local->remote).
76    pub async fn write_from_local(
77        &self,
78        client: &(impl context::Actor + Send + Sync),
79        local: KeepaliveLocalMemory,
80        timeout: u64,
81    ) -> Result<bool, anyhow::Error> {
82        let mut action = RdmaAction::new();
83        action.add_write_from_local(self.clone(), local)?;
84        action.submit(client, Duration::from_secs(timeout)).await?;
85        Ok(true)
86    }
87
88    /// Pull data from this remote buffer into local memory (remote->local).
89    pub async fn read_into_local(
90        &self,
91        client: &(impl context::Actor + Send + Sync),
92        local: KeepaliveLocalMemory,
93        timeout: u64,
94    ) -> Result<bool, anyhow::Error> {
95        let mut action = RdmaAction::new();
96        action.add_read_into_local(self.clone(), local)?;
97        action.submit(client, Duration::from_secs(timeout)).await?;
98        Ok(true)
99    }
100
101    /// Drop the buffer and release remote handles.
102    pub async fn drop_buffer(&self, client: &impl context::Actor) -> Result<(), anyhow::Error> {
103        tracing::debug!("[buffer] dropping buffer id={}", self.id);
104        self.owner.release_buffer(client, self.id).await?;
105        Ok(())
106    }
107}
108
109/// Utility to validate execution context.
110///
111/// Remote Execution environments do not always have access to the nvidia_peermem module
112/// and/or set the PeerMappingOverride parameter due to security. This function can be
113/// used to validate that the execution context when running operations that need this
114/// functionality (ie. cudaHostRegisterIoMemory).
115///
116/// # Returns
117///
118/// * `Ok(())` if the execution context is valid
119/// * `Err(anyhow::Error)` if the execution context is invalid
120pub async fn validate_execution_context() -> Result<(), anyhow::Error> {
121    // Check for nvidia peermem
122    match fs::read_to_string("/proc/modules") {
123        Ok(contents) => {
124            if !contents.contains("nvidia_peermem") {
125                return Err(anyhow::anyhow!(
126                    "nvidia_peermem module not found in /proc/modules"
127                ));
128            }
129        }
130        Err(e) => {
131            return Err(anyhow::anyhow!(e));
132        }
133    }
134
135    // Test file access to nvidia params
136    match fs::read_to_string("/proc/driver/nvidia/params") {
137        Ok(contents) => {
138            if !contents.contains("PeerMappingOverride=1") {
139                return Err(anyhow::anyhow!(
140                    "PeerMappingOverride=1 not found in /proc/driver/nvidia/params"
141                ));
142            }
143        }
144        Err(e) => {
145            return Err(anyhow::anyhow!(e));
146        }
147    }
148    Ok(())
149}
150
151/// Get all segments that have been registered with MRs for the given PD.
152///
153/// Each protection domain maintains independent segment registrations, so
154/// callers must pass the PD whose lkeys they intend to use.
155pub fn get_registered_cuda_segments(
156    pd: *mut rdmaxcel_sys::ibv_pd,
157) -> Vec<rdmaxcel_sys::rdma_segment_info_t> {
158    unsafe {
159        let segment_count = rdmaxcel_sys::rdma_get_active_segment_count(pd);
160        if segment_count <= 0 {
161            return Vec::new();
162        }
163
164        let mut segments = vec![
165            std::mem::MaybeUninit::<rdmaxcel_sys::rdma_segment_info_t>::zeroed()
166                .assume_init();
167            segment_count as usize
168        ];
169        let actual_count = rdmaxcel_sys::rdma_get_all_registered_segment_info(
170            pd,
171            segments.as_mut_ptr(),
172            segment_count,
173        );
174
175        if actual_count > 0 {
176            segments.truncate(actual_count as usize);
177            segments
178        } else {
179            Vec::new()
180        }
181    }
182}
183
184/// Segment scanner callback type alias for convenience.
185pub type SegmentScannerFn = rdmaxcel_sys::RdmaxcelSegmentScannerFn;
186
187/// Register a segment scanner callback.
188///
189/// The scanner callback is called during RDMA segment registration to discover
190/// CUDA memory segments. The callback should fill the provided buffer with
191/// segment information and return the total count of segments found.
192///
193/// If the returned count exceeds the buffer size, the caller will allocate
194/// a larger buffer and retry.
195///
196/// Pass `None` to unregister the scanner.
197///
198/// # Safety
199///
200/// The provided callback function must be safe to call from C code and must
201/// properly handle the segment buffer.
202pub fn register_segment_scanner(scanner: SegmentScannerFn) {
203    // SAFETY: We are registering a callback function pointer with rdmaxcel.
204    unsafe { rdmaxcel_sys::rdmaxcel_register_segment_scanner(scanner) }
205}
monarch_rdma/rdma_components.rs

monarch_rdma/
rdma_components.rs
