monarch_rdma/
action.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//! Batched RDMA action layer.
10//!
11//! [`RdmaAction`] accumulates [`crate::RdmaOp`]s through `add_read_into_local`
12//! / `add_write_from_local`, validates the per-op sizes and the local
13//! memory ranges as ops are added, and then dispatches the queued ops
14//! across the available backends in parallel on [`RdmaAction::submit`].
15
16use std::collections::HashMap;
17use std::time::Duration;
18
19use hyperactor::context;
20
21use crate::RdmaManagerActor;
22use crate::RdmaManagerMessageClient;
23use crate::RdmaOp;
24use crate::RdmaOpType;
25use crate::backend::RdmaBackendHandle;
26use crate::local_memory::KeepaliveLocalMemory;
27use crate::rdma_components::RdmaRemoteBuffer;
28
29/// A batch of RDMA operations submitted as a single unit.
30///
31/// `RdmaAction` is a builder that accumulates read-into-local and
32/// write-from-local ops, performs eager validation (size check + local
33/// memory range race detection), and then runs them concurrently across
34/// the available backends on [`Self::submit`].
35///
36/// Local-memory race detection treats an `add_read_into_local` claim as a
37/// write to the local range and an `add_write_from_local` claim as a
38/// read; two reads of the same range merge into one claim, anything else
39/// errors. Remote-side ranges are deliberately not tracked.
40pub struct RdmaAction {
41    entries: Vec<RdmaOp>,
42    // Claimed local address ranges keyed by `[start, end)`. The stored
43    // [`RdmaOpType`] doubles as the op-kind tag: `ReadIntoLocal` is a
44    // local write, `WriteFromLocal` is a local read.
45    local_claims: HashMap<(usize, usize), RdmaOpType>,
46}
47
48impl Default for RdmaAction {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl RdmaAction {
55    pub fn new() -> Self {
56        Self {
57            entries: Vec::new(),
58            local_claims: HashMap::new(),
59        }
60    }
61
62    /// True if this op writes to the local address range
63    /// (`ReadIntoLocal` reads remote data *into* local memory).
64    fn is_local_write(op_type: RdmaOpType) -> bool {
65        matches!(op_type, RdmaOpType::ReadIntoLocal)
66    }
67
68    /// Queue a read from `remote` into `local`. Records a local *write*
69    /// claim over the local memory range.
70    pub fn add_read_into_local(
71        &mut self,
72        remote: RdmaRemoteBuffer,
73        local: KeepaliveLocalMemory,
74    ) -> Result<&mut Self, anyhow::Error> {
75        if local.size() < remote.size {
76            anyhow::bail!(
77                "destination local memory size ({}) must be >= remote buffer size ({})",
78                local.size(),
79                remote.size,
80            );
81        }
82        self.record_claim(local.addr(), local.size(), RdmaOpType::ReadIntoLocal)?;
83        self.entries.push(RdmaOp {
84            op_type: RdmaOpType::ReadIntoLocal,
85            local,
86            remote,
87        });
88        Ok(self)
89    }
90
91    /// Queue a write from `local` into `remote`. Records a local *read*
92    /// claim over the local memory range.
93    pub fn add_write_from_local(
94        &mut self,
95        remote: RdmaRemoteBuffer,
96        local: KeepaliveLocalMemory,
97    ) -> Result<&mut Self, anyhow::Error> {
98        if local.size() > remote.size {
99            anyhow::bail!(
100                "source local memory size ({}) must be <= remote buffer size ({})",
101                local.size(),
102                remote.size,
103            );
104        }
105        self.record_claim(local.addr(), local.size(), RdmaOpType::WriteFromLocal)?;
106        self.entries.push(RdmaOp {
107            op_type: RdmaOpType::WriteFromLocal,
108            local,
109            remote,
110        });
111        Ok(self)
112    }
113
114    fn record_claim(
115        &mut self,
116        addr: usize,
117        size: usize,
118        op_type: RdmaOpType,
119    ) -> Result<(), anyhow::Error> {
120        let mut start = addr;
121        let mut end = addr.saturating_add(size);
122        let new_is_write = Self::is_local_write(op_type);
123
124        // In one pass, we merge all entries that overlap with the new claim into a single entry.
125        // We then remove all the merged entries.
126        let mut to_remove: Vec<(usize, usize)> = Vec::new();
127        for (&(s, e), &existing) in self.local_claims.iter() {
128            if end <= s || e <= start {
129                continue;
130            }
131            if new_is_write || Self::is_local_write(existing) {
132                anyhow::bail!(
133                    "RdmaAction data race: existing {:?} claim at [{:#x}, {:#x}) overlaps new {:?} claim at [{:#x}, {:#x})",
134                    existing,
135                    s,
136                    e,
137                    op_type,
138                    start,
139                    end,
140                );
141            }
142            start = start.min(s);
143            end = end.max(e);
144            to_remove.push((s, e));
145        }
146
147        for k in &to_remove {
148            self.local_claims.remove(k);
149        }
150        self.local_claims.insert((start, end), op_type);
151        Ok(())
152    }
153
154    /// Submit all queued ops. Ops are grouped by backend and each group
155    /// is submitted in parallel. Safe to call more than once on the same
156    /// action — the queued ops and overlap claims are left intact.
157    ///
158    /// Takes `&mut self` so the borrow checker prevents two submit
159    /// futures from being alive on the same action simultaneously;
160    /// otherwise the local-range overlap detection would be meaningless
161    /// (two in-flight dispatches over the same claimed local memory).
162    pub async fn submit(
163        &mut self,
164        client: &(impl context::Actor + Send + Sync),
165        timeout: Duration,
166    ) -> Result<(), anyhow::Error> {
167        if self.entries.is_empty() {
168            return Ok(());
169        }
170
171        // This proc's spawned backends, in priority order.
172        let handles = RdmaManagerActor::local_handle(client)
173            .get_backend_handles(client)
174            .await?;
175        let mut buckets: Vec<(RdmaBackendHandle, Vec<RdmaOp>)> =
176            handles.into_iter().map(|h| (h, Vec::new())).collect();
177
178        // Route each op to the first backend the local proc runs that the
179        // remote buffer also advertises.
180        for entry in &self.entries {
181            let op = RdmaOp {
182                op_type: entry.op_type,
183                local: entry.local.clone(),
184                remote: entry.remote.clone(),
185            };
186            match buckets
187                .iter_mut()
188                .find(|(handle, _)| entry.remote.is_compatible_with(handle))
189            {
190                Some((_, ops)) => ops.push(op),
191                None => anyhow::bail!("no compatible RDMA backend for buffer: {:?}", entry.remote),
192            }
193        }
194
195        // Submit each non-empty backend group in parallel, waiting for all
196        // groups to finish and accumulating every failure.
197        let pending = buckets.into_iter().filter(|(_, ops)| !ops.is_empty()).map(
198            |(handle, ops)| async move {
199                let name = handle.backend_name();
200                handle
201                    .submit(client, ops, timeout)
202                    .await
203                    .map_err(|e| format!("({name}) {e}"))
204            },
205        );
206        let errors: Vec<String> = futures::future::join_all(pending)
207            .await
208            .into_iter()
209            .filter_map(Result::err)
210            .collect();
211        if errors.is_empty() {
212            Ok(())
213        } else {
214            anyhow::bail!(
215                "RDMA submit failed on {} backend(s):\n{}",
216                errors.len(),
217                errors.join("\n")
218            )
219        }
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    /// Convenience: WRITE-kind claim (i.e. simulates `add_read_into_local`).
228    fn write_claim(action: &mut RdmaAction, addr: usize, size: usize) -> Result<(), anyhow::Error> {
229        action.record_claim(addr, size, RdmaOpType::ReadIntoLocal)
230    }
231
232    /// Convenience: READ-kind claim (i.e. simulates `add_write_from_local`).
233    fn read_claim(action: &mut RdmaAction, addr: usize, size: usize) -> Result<(), anyhow::Error> {
234        action.record_claim(addr, size, RdmaOpType::WriteFromLocal)
235    }
236
237    #[test]
238    fn disjoint_writes_ok() {
239        let mut a = RdmaAction::new();
240        write_claim(&mut a, 0x1000, 0x100).unwrap();
241        write_claim(&mut a, 0x2000, 0x100).unwrap();
242        assert_eq!(a.local_claims.len(), 2);
243    }
244
245    #[test]
246    fn overlapping_writes_error() {
247        let mut a = RdmaAction::new();
248        write_claim(&mut a, 0x1000, 0x100).unwrap();
249        let err = write_claim(&mut a, 0x1080, 0x100).unwrap_err();
250        assert!(err.to_string().contains("data race"));
251    }
252
253    #[test]
254    fn overlapping_reads_merge() {
255        let mut a = RdmaAction::new();
256        read_claim(&mut a, 0x1000, 0x100).unwrap();
257        read_claim(&mut a, 0x1080, 0x100).unwrap();
258        assert_eq!(a.local_claims.len(), 1);
259        let (&(start, end), kind) = a.local_claims.iter().next().unwrap();
260        assert_eq!(start, 0x1000);
261        assert_eq!(end, 0x1180);
262        assert_eq!(*kind, RdmaOpType::WriteFromLocal);
263    }
264
265    #[test]
266    fn cascading_reads_merge() {
267        // Two disjoint reads, then a third that bridges them — the cascade
268        // must absorb both pre-existing entries, not just the first.
269        let mut a = RdmaAction::new();
270        read_claim(&mut a, 0x1000, 0x100).unwrap();
271        read_claim(&mut a, 0x1200, 0x100).unwrap();
272        assert_eq!(a.local_claims.len(), 2);
273        read_claim(&mut a, 0x1080, 0x200).unwrap();
274        assert_eq!(a.local_claims.len(), 1);
275        let (&(start, end), _) = a.local_claims.iter().next().unwrap();
276        assert_eq!(start, 0x1000);
277        assert_eq!(end, 0x1300);
278    }
279
280    #[test]
281    fn write_vs_read_errors() {
282        let mut a = RdmaAction::new();
283        read_claim(&mut a, 0x1000, 0x100).unwrap();
284        let err = write_claim(&mut a, 0x1080, 0x100).unwrap_err();
285        assert!(err.to_string().contains("data race"));
286    }
287
288    #[test]
289    fn read_vs_write_errors() {
290        let mut a = RdmaAction::new();
291        write_claim(&mut a, 0x1000, 0x100).unwrap();
292        let err = read_claim(&mut a, 0x1080, 0x100).unwrap_err();
293        assert!(err.to_string().contains("data race"));
294    }
295
296    #[test]
297    fn touching_ranges_do_not_overlap() {
298        // `[0,100)` and `[100,200)` are adjacent, not overlapping.
299        let mut a = RdmaAction::new();
300        write_claim(&mut a, 0x1000, 0x100).unwrap();
301        write_claim(&mut a, 0x1100, 0x100).unwrap();
302        assert_eq!(a.local_claims.len(), 2);
303    }
304}
monarch_rdma/action.rs

monarch_rdma/
action.rs
