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

monarch_rdma/
backend.rs
