monarch_rdma/backend/ibverbs/
device_selection.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//! ibverbs-specific device selection: pairs a [`MemoryLocation`] with the
10//! RDMA NIC(s) that have the best PCIe path to it.
11
12use std::sync::LazyLock;
13
14use anyhow::Context;
15use anyhow::Result;
16use dashmap::DashMap;
17
18use super::device::IbvDevice;
19use super::device::IbvDeviceImpl;
20use super::primitives::IbvDeviceInfo;
21use crate::device_selection::MemoryLocation;
22use crate::device_selection::PCIAddress;
23use crate::device_selection::PciPath;
24use crate::device_selection::cpu_path;
25use crate::device_selection::get_cuda_pci_address;
26use crate::device_selection::pci_path;
27
28/// What an [`IbvConfig`](super::primitives::IbvConfig) targets: a memory
29/// location (whose best NIC is auto-selected) or an explicit NIC by name.
30#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
31pub enum IbvDeviceTarget {
32    /// Auto-select the best NIC for a CPU/GPU memory location.
33    MemoryLocation(MemoryLocation),
34    /// Use the NIC with this exact device name (e.g. `"mlx5_0"`).
35    Nic(String),
36}
37
38impl IbvDeviceTarget {
39    /// Target the best NIC for CPU memory on NUMA node `numa`.
40    pub fn cpu(numa: u32) -> Self {
41        Self::MemoryLocation(MemoryLocation::Cpu(Some(numa)))
42    }
43
44    /// Target the best NIC for GPU memory on CUDA ordinal `ordinal`.
45    pub fn gpu(ordinal: u32) -> Self {
46        Self::MemoryLocation(MemoryLocation::Gpu(Some(ordinal)))
47    }
48
49    /// Target the NIC with the given device name.
50    pub fn nic(name: impl Into<String>) -> Self {
51        Self::Nic(name.into())
52    }
53}
54
55/// The PCI address of an RDMA NIC, resolved from its sysfs device link
56/// (`/sys/class/infiniband/<name>/device`).
57pub fn get_pci_address(device: &IbvDeviceInfo) -> Result<PCIAddress> {
58    let link = format!("/sys/class/infiniband/{}/device", device.name());
59    let resolved =
60        std::fs::canonicalize(&link).with_context(|| format!("resolving sysfs link {link}"))?;
61    resolved
62        .file_name()
63        .and_then(|name| name.to_str())
64        .and_then(PCIAddress::parse)
65        .with_context(|| format!("no PCI address in resolved path {resolved:?}"))
66}
67
68/// The NIC(s) of backend `I` with the best path to `location`, ranked by
69/// [`PciPath::is_better_than`] (most local, then highest port-capped
70/// bandwidth) and returning all that tie for best.
71///
72/// A NIC's path bandwidth is the lesser of its PCIe-chain bottleneck and
73/// its RDMA port speed. Results are cached per `(backend, location)` since
74/// the PCI/NUMA topology is fixed for the process lifetime.
75pub fn select_optimal_ibv_devices<I: IbvDeviceImpl>(
76    location: MemoryLocation,
77) -> Vec<IbvDeviceInfo> {
78    static CACHE: LazyLock<DashMap<(&'static str, MemoryLocation), Vec<IbvDeviceInfo>>> =
79        LazyLock::new(DashMap::new);
80    let key = (I::typename(), location);
81    if let Some(cached) = CACHE.get(&key) {
82        return cached.value().clone();
83    }
84    let result = compute_optimal_ibv_devices::<I>(location);
85    CACHE.insert(key, result.clone());
86    result
87}
88
89/// Uncached core of [`select_optimal_ibv_devices`].
90fn compute_optimal_ibv_devices<I: IbvDeviceImpl>(location: MemoryLocation) -> Vec<IbvDeviceInfo> {
91    let mut best: Option<PciPath> = None;
92    let mut devices: Vec<IbvDeviceInfo> = Vec::new();
93    for nic in IbvDevice::<I>::list() {
94        let Some(path) = nic_path(&nic, location) else {
95            continue;
96        };
97        match best {
98            Some(current) if current.is_better_than(&path) => continue,
99            Some(current) if path.is_better_than(&current) => devices.clear(),
100            _ => {}
101        }
102        best = Some(path);
103        devices.push(nic);
104    }
105    devices
106}
107
108/// The best [`PciPath`] from `location` to `nic`, capped by the NIC's RDMA
109/// port speed. `None` when the path can't be computed — e.g. the NIC's PCI
110/// address or a required GPU address can't be resolved.
111fn nic_path(nic: &IbvDeviceInfo, location: MemoryLocation) -> Option<PciPath> {
112    let nic_addr = get_pci_address(nic).ok()?;
113    let base = match location {
114        MemoryLocation::Cpu(numa) => cpu_path(&nic_addr, numa),
115        MemoryLocation::Gpu(Some(ordinal)) => pci_path(&get_cuda_pci_address(ordinal)?, &nic_addr),
116        MemoryLocation::Gpu(None) => best_gpu_path(&nic_addr)?,
117    };
118    Some(cap_by_port_speed(base, nic))
119}
120
121/// The best path from any visible CUDA device to `nic_addr`, or `None` if
122/// no GPU's PCI address resolves.
123fn best_gpu_path(nic_addr: &PCIAddress) -> Option<PciPath> {
124    (0..cuda_device_count())
125        .filter_map(|ordinal| get_cuda_pci_address(ordinal as u32))
126        .map(|gpu_addr| pci_path(&gpu_addr, nic_addr))
127        .reduce(|a, b| if b.is_better_than(&a) { b } else { a })
128}
129
130/// Caps `path`'s bottleneck at the NIC's RDMA port speed. A NIC with no
131/// ACTIVE port reports a port speed of 0 and is dragged to the worst
132/// case, like an unreadable PCIe link.
133fn cap_by_port_speed(path: PciPath, nic: &IbvDeviceInfo) -> PciPath {
134    PciPath {
135        bottleneck_mbytes_per_sec: path
136            .bottleneck_mbytes_per_sec
137            .min(nic.port_speed_mbytes_per_sec()),
138        ..path
139    }
140}
141
142/// Number of NVIDIA GPUs visible to the kernel driver, counted from the
143/// per-GPU directories under `/proc/driver/nvidia/gpus`. This reads kernel
144/// driver state, so unlike `cuDeviceGetCount` it needs no CUDA
145/// initialization; it returns 0 when the NVIDIA driver is absent (no GPU, or
146/// a non-NVIDIA platform).
147///
148/// TODO(slurye): Generalize this (and `get_cuda_pci_address`) to support AMD.
149pub(crate) fn cuda_device_count() -> i32 {
150    std::fs::read_dir("/proc/driver/nvidia/gpus")
151        .map(|entries| entries.flatten().count() as i32)
152        .unwrap_or(0)
153}
154
155/// Resolves an [`IbvDeviceTarget`] to a single NIC of backend `I`: the
156/// named device for [`IbvDeviceTarget::Nic`], or the best NIC for a memory
157/// location. Both arms are scoped to backend `I`, so a name belonging to a
158/// different backend resolves to `None`.
159pub fn resolve_target<I: IbvDeviceImpl>(target: &IbvDeviceTarget) -> Option<IbvDeviceInfo> {
160    match target {
161        IbvDeviceTarget::Nic(name) => IbvDevice::<I>::list()
162            .into_iter()
163            .find(|device| device.name() == name),
164        IbvDeviceTarget::MemoryLocation(location) => select_optimal_ibv_devices::<I>(*location)
165            .into_iter()
166            .next(),
167    }
168}
169
170/// Process-wide CUDA ordinal → optimal NIC map, computed once per backend and
171/// cached. CPU-only workloads pay no initialization cost.
172pub fn get_cuda_device_to_ibv_device<I: IbvDeviceImpl>() -> &'static Vec<Option<IbvDeviceInfo>> {
173    // A function-local `static` in a generic fn is one cell shared across every
174    // `I`, so the cache must be keyed per backend; `I::typename()` selects it.
175    // Each backend's map is `Box::leak`ed once so the cache stores (and returns) a
176    // `&'static`.
177    static CACHE: LazyLock<DashMap<&'static str, &'static Vec<Option<IbvDeviceInfo>>>> =
178        LazyLock::new(DashMap::new);
179    *CACHE.entry(I::typename()).or_insert_with(|| {
180        let result: Vec<Option<IbvDeviceInfo>> = (0..cuda_device_count())
181            .map(|ordinal| {
182                select_optimal_ibv_devices::<I>(MemoryLocation::Gpu(Some(ordinal as u32)))
183                    .into_iter()
184                    .next()
185            })
186            .collect();
187        Box::leak(Box::new(result))
188    })
189}
monarch_rdma/backend/ibverbs/device_selection.rs

monarch_rdma/backend/ibverbs/
device_selection.rs
