monarch_rdma/

backend.rs

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! RDMA backend implementations.

#[cfg(any(test, feature = "test-utils"))]
pub mod cuda_test_utils;
pub mod ibverbs;
pub mod tcp;

use std::fmt::Debug;
use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use hyperactor::context;
use serde::Deserialize;
use serde::Serialize;
use serde::de::DeserializeOwned;
use typeuri::Named;

use crate::RdmaOp;
use crate::RdmaTransportLevel;
use crate::backend::ibverbs::efa_device::EfaDevice;
use crate::backend::ibverbs::manager_actor::IbvBackend;
use crate::backend::ibverbs::mlx_device::MlxDevice;
use crate::backend::ibverbs::primitives::IbvConfig;
use crate::backend::tcp::manager_actor::TcpBackend;
use crate::local_memory::KeepaliveLocalMemory;
use crate::rdma_components::RdmaRemoteBuffer;

/// Configuration for spawning RDMA backends.
#[derive(Debug, Clone, Default)]
pub struct RdmaConfig {
    /// Configuration for ibverbs-based backends.
    pub(crate) ibv: Option<IbvConfig>,
}

/// A transport backend for RDMA operations.
///
/// Implementors (e.g. [`IbvBackend<I>`], [`TcpBackend`]) are assembled
/// into the [`RdmaBackendHandle`] and [`RdmaRemoteBackends`] registries
/// by [`register_rdma_backends!`].
#[async_trait]
pub trait RdmaBackend: Clone + Debug + Send + Sync + 'static {
    /// Serializable per-buffer context carried on the wire.
    type RemoteBackendContext: Clone + Debug + Serialize + DeserializeOwned + Send + Sync + 'static;

    /// Backend-specific transport details (e.g. a cffi struct with raw
    /// ibverbs handles for GPU-initiated RDMA).
    type TransportInfo;

    /// Whether this backend is available on this host in the current proc.
    /// Returns false if the backend is supported on the host but is
    /// disabled by a config.
    fn available() -> bool;

    /// The transport level this backend provides.
    fn transport_level(&self) -> RdmaTransportLevel;

    /// Low-level transport details for direct control over RDMA
    /// operations (e.g. from a GPU kernel).
    fn transport_info(&self) -> Option<Self::TransportInfo>;

    /// Spawn the backend's actor(s) as children of `cx` and return its handle.
    async fn spawn(cx: &(impl context::Actor + Send + Sync), config: &RdmaConfig) -> Result<Self>
    where
        Self: Sized;

    /// Register `local` for remote access and return its wire context.
    async fn register_remote_buffer(
        &self,
        cx: &(impl context::Actor + Send + Sync),
        remote_buf_id: usize,
        local: KeepaliveLocalMemory,
    ) -> Result<Self::RemoteBackendContext>;

    /// Release a buffer registration by id.
    async fn release_buffer(
        &self,
        cx: &(impl context::Actor + Send + Sync),
        remote_buf_id: usize,
    ) -> Result<()>;

    /// Submit a batch of ops to this backend.
    async fn submit(
        &self,
        cx: &(impl context::Actor + Send + Sync),
        ops: Vec<RdmaOp>,
        timeout: Duration,
    ) -> Result<()>;
}

/// Resolves backend `B`'s context from a buffer's advertised backends.
/// One impl per backend is generated by [`register_rdma_backends!`].
pub(crate) trait ResolveRemoteBackendContext<B: RdmaBackend> {
    fn resolve(&self) -> Option<B::RemoteBackendContext>;
}

/// Derives the per-process RDMA backend registry from a list of
/// `Variant: Handle` pairs, where each `Handle` implements
/// [`RdmaBackend`].
///
/// The expansion defines [`RdmaRemoteBackends`] (a buffer's per-backend
/// wire contexts) with its [`ResolveRemoteBackendContext`] impls and
/// `RdmaRemoteBuffer::resolve_<name>` accessors, [`RdmaBackendHandle`]
/// and its `submit` dispatch, and [`RdmaBackends`] (the proc's spawned
/// backends). The list order is the routing priority.
macro_rules! register_rdma_backends {
    ($($variant:ident: $handle:ty),+ $(,)?) => {
        paste::paste! {
            /// The backends a buffer is reachable through, one slot per backend.
            #[derive(Debug, Clone, Serialize, Deserialize, Named, Default)]
            pub(crate) struct RdmaRemoteBackends {
                $(pub(crate) [<$variant:lower>]: Option<<$handle as RdmaBackend>::RemoteBackendContext>,)+
            }

            $(
                impl ResolveRemoteBackendContext<$handle> for RdmaRemoteBuffer {
                    fn resolve(&self) -> Option<<$handle as RdmaBackend>::RemoteBackendContext> {
                        self.[<resolve_ $variant:lower>]()
                    }
                }
            )+

            impl RdmaRemoteBuffer {
                $(
                    /// Context for this backend, if the buffer advertises it.
                    pub fn [<resolve_ $variant:lower>](
                        &self,
                    ) -> Option<<$handle as RdmaBackend>::RemoteBackendContext> {
                        self.backends.[<$variant:lower>].clone()
                    }
                )+

                /// Whether this buffer advertises a backend compatible with `handle`.
                pub(crate) fn is_compatible_with(&self, handle: &RdmaBackendHandle) -> bool {
                    match handle {
                        $(
                            RdmaBackendHandle::$variant(_) => {
                                self.backends.[<$variant:lower>].is_some()
                            }
                        )+
                    }
                }
            }

            /// The backends spawned on this proc.
            #[derive(Debug, Default)]
            pub(crate) struct RdmaBackends {
                $([<$variant:lower>]: Option<$handle>,)+
            }

            impl RdmaBackends {
                /// Spawn every [`available`](RdmaBackend::available) backend.
                /// A backend that fails to spawn is skipped; bails only if no
                /// backend spawns.
                pub(crate) async fn spawn_available(
                    cx: &(impl context::Actor + Send + Sync),
                    config: &RdmaConfig,
                ) -> Result<Self> {
                    let mut backends = Self::default();
                    let mut errors: Vec<String> = Vec::new();
                    $(
                        if <$handle as RdmaBackend>::available() {
                            match <$handle as RdmaBackend>::spawn(cx, config).await {
                                Ok(handle) => backends.[<$variant:lower>] = Some(handle),
                                Err(e) => errors.push(format!("{}: {e}", stringify!($variant))),
                            }
                        }
                    )+
                    if backends.is_empty() {
                        if errors.is_empty() {
                            anyhow::bail!("no RDMA backend available");
                        }
                        anyhow::bail!(
                            "all available RDMA backends failed to initialize: {}",
                            errors.join("; ")
                        );
                    }
                    if !errors.is_empty() {
                        tracing::warn!("some RDMA backends failed to initialize: {}", errors.join("; "));
                    }
                    Ok(backends)
                }

                fn is_empty(&self) -> bool {
                    $(self.[<$variant:lower>].is_none() &&)+ true
                }

                /// Handles for all spawned backends, in priority order.
                pub(crate) fn handles(&self) -> Vec<RdmaBackendHandle> {
                    let mut handles = Vec::new();
                    $(
                        if let Some(handle) = &self.[<$variant:lower>] {
                            handles.push(RdmaBackendHandle::$variant(handle.clone()));
                        }
                    )+
                    handles
                }

                /// Register `local` with every spawned backend. On the first
                /// failure, release the backends that already registered and
                /// return that error.
                pub(crate) async fn register_all(
                    &self,
                    cx: &(impl context::Actor + Send + Sync),
                    remote_buf_id: usize,
                    local: KeepaliveLocalMemory,
                ) -> Result<RdmaRemoteBackends> {
                    let mut remotes = RdmaRemoteBackends::default();
                    $(
                        if let Some(handle) = &self.[<$variant:lower>] {
                            match <$handle as RdmaBackend>::register_remote_buffer(
                                handle,
                                cx,
                                remote_buf_id,
                                local.clone(),
                            )
                            .await
                            {
                                Ok(context) => remotes.[<$variant:lower>] = Some(context),
                                Err(e) => {
                                    // Release only the backends that already
                                    // registered: this one failed and later
                                    // ones were never reached.
                                    if let Err(release_err) =
                                        self.release_registered(cx, remote_buf_id, &remotes).await
                                    {
                                        tracing::warn!("failed to release remote buffers after registration failure: {release_err}");
                                    }
                                    return Err(e);
                                }
                            }
                        }
                    )+
                    Ok(remotes)
                }

                /// Release `remote_buf_id` from only the backends that
                /// successfully registered (those present in `registered`),
                /// accumulating any failures.
                async fn release_registered(
                    &self,
                    cx: &(impl context::Actor + Send + Sync),
                    remote_buf_id: usize,
                    registered: &RdmaRemoteBackends,
                ) -> Result<()> {
                    let mut errors: Vec<String> = Vec::new();
                    $(
                        if let Some(handle) = &self.[<$variant:lower>]
                            && registered.[<$variant:lower>].is_some()
                        {
                            if let Err(e) =
                                <$handle as RdmaBackend>::release_buffer(handle, cx, remote_buf_id).await
                            {
                                errors.push(format!("({}) {e}", stringify!($variant)));
                            }
                        }
                    )+
                    if errors.is_empty() {
                        Ok(())
                    } else {
                        anyhow::bail!(
                            "RDMA release failed on {} backend(s):\n{}",
                            errors.len(),
                            errors.join("\n")
                        )
                    }
                }

                /// Release `remote_buf_id` from every spawned backend,
                /// accumulating any failures.
                pub(crate) async fn release_all(
                    &self,
                    cx: &(impl context::Actor + Send + Sync),
                    remote_buf_id: usize,
                ) -> Result<()> {
                    let mut errors: Vec<String> = Vec::new();
                    $(
                        if let Some(handle) = &self.[<$variant:lower>] {
                            if let Err(e) =
                                <$handle as RdmaBackend>::release_buffer(handle, cx, remote_buf_id).await
                            {
                                errors.push(format!("({}) {e}", stringify!($variant)));
                            }
                        }
                    )+
                    if errors.is_empty() {
                        Ok(())
                    } else {
                        anyhow::bail!(
                            "RDMA release failed on {} backend(s):\n{}",
                            errors.len(),
                            errors.join("\n")
                        )
                    }
                }
            }
        }

        wirevalue::register_type!(RdmaRemoteBackends);

        /// Handle to a spawned backend.
        #[derive(Debug, Clone)]
        pub enum RdmaBackendHandle {
            $($variant($handle),)+
        }

        impl RdmaBackendHandle {
            /// The backend's name, for diagnostics.
            pub(crate) fn backend_name(&self) -> &'static str {
                match self {
                    $(RdmaBackendHandle::$variant(_) => stringify!($variant),)+
                }
            }

            /// Submit `ops` to this backend.
            pub(crate) async fn submit(
                &self,
                cx: &(impl context::Actor + Send + Sync),
                ops: Vec<RdmaOp>,
                timeout: Duration,
            ) -> Result<()> {
                match self {
                    $(
                        RdmaBackendHandle::$variant(handle) => {
                            <$handle as RdmaBackend>::submit(handle, cx, ops, timeout).await
                        }
                    )+
                }
            }
        }
    };
}

register_rdma_backends! {
    Mlx: IbvBackend<MlxDevice>,
    Efa: IbvBackend<EfaDevice>,
    Tcp: TcpBackend,
}