monarch_rdma/
backend.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 backend implementations.
10
11#[cfg(test)]
12pub(crate) mod cuda_test_utils;
13pub mod ibverbs;
14pub mod tcp;
15
16use std::fmt::Debug;
17use std::sync::Arc;
18use std::time::Duration;
19
20use anyhow::Result;
21use async_trait::async_trait;
22use hyperactor::reference;
23use serde::Deserialize;
24use serde::Serialize;
25use tokio::sync::OnceCell;
26
27use crate::RdmaOp;
28use crate::RdmaTransportLevel;
29
30/// Backend-specific context for a remote buffer.
31///
32/// Each variant holds the information needed to perform RDMA operations
33/// using that backend on a particular buffer.
34///
35/// The [`OnceCell`] is lazily populated at runtime and excluded from
36/// serialization; deserializing produces an empty cell.
37#[derive(Debug, Clone)]
38pub enum RdmaRemoteBackendContext {
39    Ibverbs(
40        reference::ActorRef<ibverbs::manager_actor::IbvManagerActor>,
41        Arc<OnceCell<ibverbs::IbvBuffer>>,
42    ),
43    Tcp(reference::ActorRef<tcp::manager_actor::TcpManagerActor>),
44}
45
46impl Serialize for RdmaRemoteBackendContext {
47    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
48        match self {
49            RdmaRemoteBackendContext::Ibverbs(actor_ref, _) => serializer
50                .serialize_newtype_variant("RdmaRemoteBackendContext", 0, "Ibverbs", actor_ref),
51            RdmaRemoteBackendContext::Tcp(actor_ref) => serializer.serialize_newtype_variant(
52                "RdmaRemoteBackendContext",
53                1,
54                "Tcp",
55                actor_ref,
56            ),
57        }
58    }
59}
60
61impl<'de> Deserialize<'de> for RdmaRemoteBackendContext {
62    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
63        #[derive(Deserialize)]
64        #[serde(rename = "RdmaRemoteBackendContext")]
65        enum Repr {
66            Ibverbs(reference::ActorRef<ibverbs::manager_actor::IbvManagerActor>),
67            Tcp(reference::ActorRef<tcp::manager_actor::TcpManagerActor>),
68        }
69
70        match Repr::deserialize(deserializer)? {
71            Repr::Ibverbs(actor_ref) => Ok(RdmaRemoteBackendContext::Ibverbs(
72                actor_ref,
73                Arc::new(OnceCell::new()),
74            )),
75            Repr::Tcp(actor_ref) => Ok(RdmaRemoteBackendContext::Tcp(actor_ref)),
76        }
77    }
78}
79
80/// Backend for executing RDMA operations over a specific transport.
81///
82/// Each backend manages the transport-specific details of connection
83/// management and data movement. The backend decides internally how to
84/// batch and schedule submitted operations.
85///
86/// Current implementations:
87/// - [`ibverbs::IbvManagerActor`] -- ibverbs NIC transport
88/// - [`tcp::TcpManagerActor`] -- TCP fallback transport
89#[async_trait]
90pub trait RdmaBackend: Send + Debug {
91    /// Backend-specific transport details (e.g., a cffi struct with raw
92    /// ibverbs handles for GPU-initiated RDMA).
93    type TransportInfo;
94
95    /// Submit a batch of RDMA operations.
96    ///
97    /// The backend decides internally how to batch, schedule, and execute
98    /// the operations (e.g., managing QPs and connections as needed).
99    async fn submit(
100        &mut self,
101        cx: &(impl hyperactor::context::Actor + Send + Sync),
102        ops: Vec<RdmaOp>,
103        timeout: Duration,
104    ) -> Result<()>;
105
106    /// The transport level provided by this backend.
107    fn transport_level(&self) -> RdmaTransportLevel;
108
109    /// Low-level backend-specific transport details for direct control
110    /// over RDMA operations (e.g., from a GPU kernel).
111    fn transport_info(&self) -> Option<Self::TransportInfo>;
112}
monarch_rdma/backend.rs

monarch_rdma/
backend.rs
