monarch_rdma/
rdma_manager_actor.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 Manager Actor
10//!
11//! Per-process actor that owns RDMA buffer registrations and delegates
12//! transport-specific work to a NIC backend and [`TcpManagerActor`].
13//!
14//! ## Responsibilities
15//!
16//! - Assigns a unique `remote_buf_id` to each registered local memory handle
17//!   and stores the [`KeepaliveLocalMemory`] for later retrieval.
18//! - Produces [`RdmaRemoteBuffer`] tokens that can be sent to remote peers so
19//!   they can address this buffer over RDMA.
20//! - Delegates MR registration, QP management, and data movement to a NIC
21//!   backend when available, or falls back to the TCP backend
22//!   ([`TcpManagerActor`]).
23//! - Handles remote [`ReleaseBuffer`] requests to clean up registrations.
24
25use std::collections::HashMap;
26use std::sync::OnceLock;
27
28use async_trait::async_trait;
29use hyperactor::Actor;
30use hyperactor::ActorHandle;
31use hyperactor::ActorRef;
32use hyperactor::Context;
33use hyperactor::HandleClient;
34use hyperactor::Handler;
35use hyperactor::Instance;
36use hyperactor::OncePortHandle;
37use hyperactor::OncePortRef;
38use hyperactor::RefClient;
39use hyperactor::RemoteSpawn;
40use hyperactor::context;
41use hyperactor_config::Flattrs;
42use serde::Deserialize;
43use serde::Serialize;
44use typeuri::Named;
45
46use crate::backend::RdmaBackendHandle;
47use crate::backend::RdmaBackends;
48use crate::backend::RdmaConfig;
49use crate::backend::ibverbs::primitives::IbvConfig;
50use crate::backend::tcp::manager_actor::TcpManagerActor;
51use crate::local_memory::KeepaliveLocalMemory;
52use crate::rdma_components::RdmaRemoteBuffer;
53
54/// Helper function to get detailed error messages from RDMAXCEL error codes
55pub fn get_rdmaxcel_error_message(error_code: i32) -> String {
56    unsafe {
57        let c_str = rdmaxcel_sys::rdmaxcel_error_string(error_code);
58        std::ffi::CStr::from_ptr(c_str)
59            .to_string_lossy()
60            .into_owned()
61    }
62}
63
64/// Local-only messages for the [`RdmaManagerActor`].
65///
66/// These messages carry [`KeepaliveLocalMemory`] and are therefore not
67/// serializable -- they can only be sent within the same process.
68#[derive(Handler, HandleClient, Debug)]
69pub enum RdmaManagerMessage {
70    /// Register a local memory handle and return a [`RdmaRemoteBuffer`] that
71    /// remote peers can use to address this buffer over RDMA.
72    RequestBuffer {
73        local: KeepaliveLocalMemory,
74        #[reply]
75        reply: OncePortHandle<RdmaRemoteBuffer>,
76    },
77    /// Look up the local memory handle for a given `remote_buf_id`. Returns
78    /// `None` if the id does not correspond to a registered buffer.
79    RequestLocalMemory {
80        remote_buf_id: usize,
81        #[reply]
82        reply: OncePortHandle<Option<KeepaliveLocalMemory>>,
83    },
84    /// Return in-process handles to all spawned backends, in priority order.
85    GetBackendHandles {
86        #[reply]
87        reply: OncePortHandle<Vec<RdmaBackendHandle>>,
88    },
89}
90
91/// Serializable release message for wire transport.
92///
93/// Used by [`RdmaRemoteBuffer::drop_buffer`] to release a buffer
94/// from a remote process.
95#[derive(Handler, HandleClient, RefClient, Debug, Serialize, Deserialize, Named)]
96pub struct ReleaseBuffer {
97    pub id: usize,
98}
99wirevalue::register_type!(ReleaseBuffer);
100
101/// Serializable query for resolving the [`TcpManagerActor`] ref
102/// from a remote [`RdmaManagerActor`].
103#[derive(Handler, HandleClient, RefClient, Debug, Serialize, Deserialize, Named)]
104pub struct GetTcpActorRef {
105    #[reply]
106    pub reply: OncePortRef<ActorRef<TcpManagerActor>>,
107}
108wirevalue::register_type!(GetTcpActorRef);
109
110#[derive(Debug)]
111#[hyperactor::export(
112    handlers = [
113        GetTcpActorRef,
114        ReleaseBuffer,
115    ],
116)]
117#[hyperactor::spawnable]
118pub struct RdmaManagerActor {
119    next_remote_buf_id: usize,
120    buffers: HashMap<usize, KeepaliveLocalMemory>,
121    params: Option<IbvConfig>,
122    backends: OnceLock<RdmaBackends>,
123}
124
125impl RdmaManagerActor {
126    /// Construct an [`ActorHandle`] for the [`RdmaManagerActor`] co-located
127    /// with the caller.
128    pub fn local_handle(client: &impl context::Actor) -> ActorHandle<Self> {
129        let actor_ref = ActorRef::attest(
130            client
131                .mailbox()
132                .actor_addr()
133                .proc_addr()
134                .actor_addr("rdma_manager"),
135        );
136        actor_ref
137            .downcast_handle(client)
138            .expect("RdmaManagerActor is not in the local process")
139    }
140}
141
142#[async_trait]
143impl RemoteSpawn for RdmaManagerActor {
144    type Params = Option<IbvConfig>;
145
146    async fn new(params: Self::Params, _environment: Flattrs) -> Result<Self, anyhow::Error> {
147        Ok(Self {
148            next_remote_buf_id: 0,
149            buffers: HashMap::new(),
150            params,
151            backends: OnceLock::new(),
152        })
153    }
154}
155
156#[async_trait]
157impl Actor for RdmaManagerActor {
158    async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
159        // Spawn every available backend. `spawn_available` bails when none
160        // is available (e.g. no NIC and TCP fallback disabled).
161        let backends = RdmaBackends::spawn_available(
162            this,
163            &RdmaConfig {
164                ibv: self.params.clone(),
165            },
166        )
167        .await?;
168        self.backends.set(backends).expect("backends set once");
169        Ok(())
170    }
171
172    // This actor is implemented in Rust, but the RDMA registration path may enter
173    // Python and take the GIL. Run its loop on the dedicated rdma runtime rather
174    // than the shared control-plane runtime; see `crate::rdma_runtime`.
175    fn spawn_server_task<F>(future: F) -> tokio::task::JoinHandle<F::Output>
176    where
177        F: std::future::Future + Send + 'static,
178        F::Output: Send + 'static,
179    {
180        crate::rdma_runtime::spawn_on_rdma_runtime(future)
181    }
182}
183
184#[async_trait]
185#[hyperactor::handle(GetTcpActorRef)]
186impl GetTcpActorRefHandler for RdmaManagerActor {
187    async fn get_tcp_actor_ref(
188        &mut self,
189        _cx: &Context<Self>,
190    ) -> Result<ActorRef<TcpManagerActor>, anyhow::Error> {
191        self.backends
192            .get()
193            .expect("backends set in init")
194            .handles()
195            .into_iter()
196            .find_map(|h| match h {
197                RdmaBackendHandle::Tcp(backend) => Some(backend.bind()),
198                _ => None,
199            })
200            .ok_or_else(|| anyhow::anyhow!("TCP backend not available"))
201    }
202}
203
204#[async_trait]
205#[hyperactor::handle(ReleaseBuffer)]
206impl ReleaseBufferHandler for RdmaManagerActor {
207    async fn release_buffer(&mut self, cx: &Context<Self>, id: usize) -> Result<(), anyhow::Error> {
208        self.buffers.remove(&id);
209        self.backends
210            .get()
211            .expect("backends set in init")
212            .release_all(cx, id)
213            .await?;
214        Ok(())
215    }
216}
217
218#[async_trait]
219#[hyperactor::handle(RdmaManagerMessage)]
220impl RdmaManagerMessageHandler for RdmaManagerActor {
221    async fn request_buffer(
222        &mut self,
223        cx: &Context<Self>,
224        local: KeepaliveLocalMemory,
225    ) -> Result<RdmaRemoteBuffer, anyhow::Error> {
226        let remote_buf_id = self.next_remote_buf_id;
227        self.next_remote_buf_id += 1;
228        let size = local.size();
229
230        let backends = self
231            .backends
232            .get()
233            .expect("backends set in init")
234            .register_all(cx, remote_buf_id, local.clone())
235            .await?;
236        self.buffers.insert(remote_buf_id, local);
237
238        Ok(RdmaRemoteBuffer {
239            id: remote_buf_id,
240            size,
241            owner: cx.bind().clone(),
242            backends,
243        })
244    }
245
246    async fn request_local_memory(
247        &mut self,
248        _cx: &Context<Self>,
249        remote_buf_id: usize,
250    ) -> Result<Option<KeepaliveLocalMemory>, anyhow::Error> {
251        Ok(self.buffers.get(&remote_buf_id).cloned())
252    }
253
254    async fn get_backend_handles(
255        &mut self,
256        _cx: &Context<Self>,
257    ) -> Result<Vec<RdmaBackendHandle>, anyhow::Error> {
258        Ok(self.backends.get().expect("backends set in init").handles())
259    }
260}
monarch_rdma/rdma_manager_actor.rs

monarch_rdma/
rdma_manager_actor.rs
