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 as ops are added,
13//! and then dispatches the queued ops across the available backends in
14//! parallel on [`RdmaAction::submit`].
15
16use std::time::Duration;
17
18use hyperactor::context;
19
20use crate::RdmaManagerActor;
21use crate::RdmaOp;
22use crate::RdmaOpType;
23use crate::backend::RdmaBackendHandle;
24use crate::local_memory::KeepaliveLocalMemory;
25use crate::rdma_components::RdmaRemoteBuffer;
26
27/// A batch of RDMA operations submitted as a single unit.
28///
29/// `RdmaAction` is a builder that accumulates read-into-local and
30/// write-from-local ops, checking each op's sizes as it is added, and then
31/// runs them concurrently across the available backends on
32/// [`Self::submit`].
33///
34/// Overlap between the queued ops' memory ranges is not tracked, on either
35/// the local or the remote side. Two ops in one action that write the same
36/// range race, and it is the caller's job not to queue them.
37pub struct RdmaAction {
38    entries: Vec<RdmaOp>,
39}
40
41impl Default for RdmaAction {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl RdmaAction {
48    pub fn new() -> Self {
49        Self {
50            entries: Vec::new(),
51        }
52    }
53
54    /// Queue a read from `remote` into `local`.
55    pub fn add_read_into_local(
56        &mut self,
57        remote: RdmaRemoteBuffer,
58        local: KeepaliveLocalMemory,
59    ) -> Result<&mut Self, anyhow::Error> {
60        if local.size() < remote.size {
61            anyhow::bail!(
62                "destination local memory size ({}) must be >= remote buffer size ({})",
63                local.size(),
64                remote.size,
65            );
66        }
67        self.entries.push(RdmaOp {
68            op_type: RdmaOpType::ReadIntoLocal,
69            local,
70            remote,
71        });
72        Ok(self)
73    }
74
75    /// Queue a write from `local` into `remote`.
76    pub fn add_write_from_local(
77        &mut self,
78        remote: RdmaRemoteBuffer,
79        local: KeepaliveLocalMemory,
80    ) -> Result<&mut Self, anyhow::Error> {
81        if local.size() > remote.size {
82            anyhow::bail!(
83                "source local memory size ({}) must be <= remote buffer size ({})",
84                local.size(),
85                remote.size,
86            );
87        }
88        self.entries.push(RdmaOp {
89            op_type: RdmaOpType::WriteFromLocal,
90            local,
91            remote,
92        });
93        Ok(self)
94    }
95
96    /// Submit all queued ops. Ops are grouped by backend and each group
97    /// is submitted in parallel. Safe to call more than once on the same
98    /// action — the queued ops are left intact.
99    ///
100    /// Takes `&mut self` so the borrow checker prevents two submit futures
101    /// from being alive on the same action at once, which would put two
102    /// in-flight dispatches over the same local memory.
103    pub async fn submit(
104        &mut self,
105        client: &(impl context::Actor + Send + Sync),
106        timeout: Duration,
107    ) -> Result<(), anyhow::Error> {
108        if self.entries.is_empty() {
109            return Ok(());
110        }
111
112        // This proc's spawned backends, in priority order.
113        let handles = RdmaManagerActor::local_backend_handles(client).await?;
114        let mut buckets: Vec<(RdmaBackendHandle, Vec<RdmaOp>)> =
115            handles.into_iter().map(|h| (h, Vec::new())).collect();
116
117        // Route each op to the first backend the local proc runs that the
118        // remote buffer also advertises.
119        for entry in &self.entries {
120            let op = RdmaOp {
121                op_type: entry.op_type,
122                local: entry.local.clone(),
123                remote: entry.remote.clone(),
124            };
125            match buckets
126                .iter_mut()
127                .find(|(handle, _)| entry.remote.is_compatible_with(handle))
128            {
129                Some((_, ops)) => ops.push(op),
130                None => anyhow::bail!("no compatible RDMA backend for buffer: {:?}", entry.remote),
131            }
132        }
133
134        // Submit each non-empty backend group in parallel, waiting for all
135        // groups to finish and accumulating every failure.
136        let pending = buckets.into_iter().filter(|(_, ops)| !ops.is_empty()).map(
137            |(handle, ops)| async move {
138                let name = handle.backend_name();
139                handle
140                    .submit(client, ops, timeout)
141                    .await
142                    .map_err(|e| format!("({name}) {e}"))
143            },
144        );
145        let errors: Vec<String> = futures::future::join_all(pending)
146            .await
147            .into_iter()
148            .filter_map(Result::err)
149            .collect();
150        if errors.is_empty() {
151            Ok(())
152        } else {
153            anyhow::bail!(
154                "RDMA submit failed on {} backend(s):\n{}",
155                errors.len(),
156                errors.join("\n")
157            )
158        }
159    }
160}
monarch_rdma/action.rs

monarch_rdma/
action.rs
