metrale_storage/
expert_tier_rdma.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `RdmaTier`: fetches expert records from an expert peer
4//! (`expert_peer`) straight into the pinned arena and returns residency
5//! addresses inside it, as `UmaArenaTier` does from local files. Transports:
6//! TCP, where the peer reads each record and sends it, and one-sided RDMA
7//! READ from the peer's registered store (`cfg(metrale_rdma_verbs)` only).
8//! Verbs rail 0 comes from `METRALE_EXPERT_RDMA_DEV` / `METRALE_EXPERT_RDMA_GID`;
9//! `METRALE_EXPERT_DUAL_RAIL=1` adds rail 1 from `METRALE_EXPERT_RAIL2_DEV` /
10//! `METRALE_EXPERT_RAIL2_GID`.
11//!
12//! Owner: metrale-storage experts.
13//! Invariants:
14//! - A fetched record passes `residency_from`'s identity check before its
15//!   addresses are returned.
16//! - The rails' registrations of the arena are dropped before the arena.
17
18use std::io::{Read, Write};
19use std::net::TcpStream;
20
21use anyhow::{Context, Result, bail};
22
23use crate::expert::{ExpertKey, ExpertLayout, ExpertRecordSpec};
24use crate::expert_arena::ExpertArena;
25use crate::expert_peer::{MODE_TCP, STATUS_OK, encode_request, read_manifest};
26use crate::expert_tier::{ArenaSlot, ExpertResidency, ExpertTier, TierKind, residency_from};
27
28/// 2026-09-25: The peer transport. `Verbs` exists only with
29/// `cfg(metrale_rdma_verbs)`; a verbs fetch uses rail `expert % n_rails`.
30enum Transport {
31    Tcp,
32    #[cfg(metrale_rdma_verbs)]
33    Verbs(Vec<Rail>),
34}
35
36/// 2026-09-25: One verbs rail: its connection, the arena's lkey on it, and the
37/// peer's per-layer `(base, rkey)` table for it.
38#[cfg(metrale_rdma_verbs)]
39struct Rail {
40    verbs: metrale_gpu_sys::verbs::Verbs,
41    arena_lkey: u32,
42    layers: Vec<(u64, u32)>,
43}
44
45pub struct RdmaTier {
46    stream: TcpStream,
47    // 2026-09-25: Declared before `arena`, so it drops first: the verbs rails
48    // hold registrations of the arena's pinned pages, which must go before the
49    // pages are freed.
50    transport: Transport,
51    arena: ExpertArena,
52    spec: ExpertRecordSpec,
53    layout: ExpertLayout,
54    healthy: bool,
55}
56
57impl RdmaTier {
58    /// 2026-09-25: Connect to the peer at `addr`, read its manifest, allocate
59    /// the arena with the manifest's record stride, and bring up the transport.
60    pub fn connect(
61        addr: &str,
62        num_slabs: u32,
63        slots_per_slab: u32,
64        use_verbs: bool,
65    ) -> Result<Self> {
66        let mut stream =
67            TcpStream::connect(addr).with_context(|| format!("connect expert peer {addr}"))?;
68        stream.set_nodelay(true).ok();
69        let index = read_manifest(&mut stream)?;
70        let spec = index.spec();
71        let layout = index.layout();
72        let arena = ExpertArena::new(num_slabs, slots_per_slab, layout.record_stride as usize)?;
73
74        let transport = if use_verbs {
75            #[cfg(metrale_rdma_verbs)]
76            {
77                connect_verbs(&mut stream, &arena, index.num_moe_layers)?
78            }
79            // 2026-09-25: Without `cfg(metrale_rdma_verbs)` only TCP exists.
80            #[cfg(not(metrale_rdma_verbs))]
81            {
82                let _ = &arena;
83                bail!(
84                    "--expert-backend rdma-verbs needs a build with rdma-core \
85                     (metrale_rdma_verbs cfg); use --expert-backend rdma (TCP) instead"
86                );
87            }
88        } else {
89            stream
90                .write_all(&[MODE_TCP])
91                .context("send TCP transport mode")?;
92            Transport::Tcp
93        };
94
95        let label = match &transport {
96            Transport::Tcp => "TCP".to_string(),
97            #[cfg(metrale_rdma_verbs)]
98            Transport::Verbs(rails) => {
99                format!("verbs (one-sided RDMA READ, {} rail(s))", rails.len())
100            }
101        };
102        tracing::info!(
103            "RdmaTier[{label}] connected to {addr}: {} layers, {} experts, stride {}",
104            index.num_moe_layers,
105            index.num_experts,
106            layout.record_stride
107        );
108        Ok(Self {
109            stream,
110            arena,
111            spec,
112            layout,
113            transport,
114            healthy: true,
115        })
116    }
117
118    pub fn arena(&self) -> &ExpertArena {
119        &self.arena
120    }
121
122    /// 2026-09-25: TCP fetch: send the request, read a status byte, then read
123    /// `stride` bytes straight into the pinned slot.
124    fn fetch_tcp(&mut self, key: ExpertKey, host: *mut u8, stride: usize) -> Result<()> {
125        if let Err(e) = self
126            .stream
127            .write_all(&encode_request(key.layer, key.expert))
128        {
129            self.healthy = false;
130            return Err(e).with_context(|| format!("peer request {:?}", key));
131        }
132        let mut status = [0u8; 1];
133        if let Err(e) = self.stream.read_exact(&mut status) {
134            self.healthy = false;
135            return Err(e).with_context(|| format!("peer status {:?}", key));
136        }
137        if status[0] != STATUS_OK {
138            bail!("peer returned error status {} for {:?}", status[0], key);
139        }
140        // 2026-09-25: SAFETY: `host` points at a `stride`-byte slot inside the
141        // pinned arena.
142        let dst = unsafe { std::slice::from_raw_parts_mut(host, stride) };
143        if let Err(e) = self.stream.read_exact(dst) {
144            self.healthy = false;
145            return Err(e).with_context(|| format!("peer payload {:?}", key));
146        }
147        Ok(())
148    }
149}
150
151/// 2026-09-25: Bring up the verbs transport through `RailSet`: one rail, or two
152/// with `METRALE_EXPERT_DUAL_RAIL=1`; register the whole arena on each; check
153/// the peer's per-rail layer tables against the manifest; connect the rails and
154/// wait for the peer's ack.
155#[cfg(metrale_rdma_verbs)]
156fn connect_verbs(
157    stream: &mut TcpStream,
158    arena: &ExpertArena,
159    num_layers: u32,
160) -> Result<Transport> {
161    use crate::expert_peer::MODE_VERBS;
162    use metrale_gpu_sys::env::{first_set, first_set_u32};
163    use metrale_gpu_sys::railset::{RailSet, RailSpec};
164
165    stream
166        .write_all(&[MODE_VERBS])
167        .context("send verbs transport mode")?;
168
169    // 2026-09-25: Each rail gets a random 24-bit PSN.
170    let spec = |dev: String, gid: u32| RailSpec::new(dev, gid, rand::random::<u32>() & 0xff_ffff);
171    let rail0 = spec(
172        first_set(&["METRALE_EXPERT_RDMA_DEV"], "roceP2p1s0f1"),
173        first_set_u32(&["METRALE_EXPERT_RDMA_GID"], 3),
174    );
175    let dual = std::env::var("METRALE_EXPERT_DUAL_RAIL").ok().as_deref() == Some("1");
176    let specs: Vec<RailSpec> = if dual {
177        let rail1 = spec(
178            first_set(&["METRALE_EXPERT_RAIL2_DEV"], "rocep1s0f1"),
179            first_set_u32(&["METRALE_EXPERT_RAIL2_GID"], 3),
180        );
181        vec![rail0, rail1]
182    } else {
183        vec![rail0]
184    };
185
186    // 2026-09-25: Each rail registers the whole arena with `remote_read = false`.
187    let mut rs = RailSet::begin(stream, &specs)?;
188    let mut arena_lkeys: Vec<u32> = Vec::with_capacity(rs.n_rails());
189    for rail in &mut rs.rails {
190        // 2026-09-25: SAFETY: the arena outlives every registration (see
191        // `RdmaTier::transport`); base_ptr()/total_bytes() describe exactly its
192        // allocation.
193        let keys = unsafe {
194            rail.verbs
195                .reg_mr(arena.base_ptr(), arena.total_bytes(), false)?
196        };
197        arena_lkeys.push(keys.lkey);
198    }
199
200    // 2026-09-25: Each rail's layer table must match the manifest before any
201    // client parameters are sent.
202    let server = rs
203        .read_server_ro(stream)
204        .context("read verbs server params")?;
205    for sp in &server {
206        if sp.layers.len() != num_layers as usize {
207            bail!(
208                "verbs peer published {} layer MRs but manifest has {num_layers} MoE layers",
209                sp.layers.len()
210            );
211        }
212    }
213
214    rs.complete(stream, &server, "verbs peer")?;
215    let rails: Vec<Rail> = rs
216        .into_verbs()
217        .into_iter()
218        .zip(arena_lkeys)
219        .zip(server)
220        .map(|((verbs, arena_lkey), sp)| Rail {
221            verbs,
222            arena_lkey,
223            layers: sp.layers,
224        })
225        .collect();
226    Ok(Transport::Verbs(rails))
227}
228
229impl ExpertTier for RdmaTier {
230    fn fetch(&mut self, key: ExpertKey, slot: ArenaSlot, _stream: u64) -> Result<ExpertResidency> {
231        let stride = self.layout.record_stride as usize;
232        let host = self.arena.slot_host_ptr(slot.slab, slot.slot)?;
233        let dev_va = self.arena.slot_dev_va(slot.slab, slot.slot)?;
234        // 2026-09-25: A copy, taken before `self.transport` is borrowed mutably.
235        let spec = self.spec;
236
237        match &mut self.transport {
238            Transport::Tcp => {
239                self.fetch_tcp(key, host, stride)?;
240            }
241            #[cfg(metrale_rdma_verbs)]
242            Transport::Verbs(rails) => {
243                let ri = (key.expert as usize) % rails.len();
244                let rail = &mut rails[ri];
245                let (base, rkey) = *rail.layers.get(key.layer as usize).with_context(|| {
246                    format!("verbs: no layer MR for layer {} ({:?})", key.layer, key)
247                })?;
248                let remote_addr = base + (key.expert as u64) * (stride as u64);
249                let wr_id = ((key.layer as u64) << 32) | (key.expert as u64);
250                // 2026-09-25: SAFETY: `host` is a `stride`-byte slot inside the
251                // arena registered under `arena_lkey`; remote_addr/rkey address
252                // the peer's layer on the same rail.
253                let post = unsafe {
254                    rail.verbs.post_read(
255                        host as *mut std::ffi::c_void,
256                        rail.arena_lkey,
257                        remote_addr,
258                        rkey,
259                        stride as u32,
260                        wr_id,
261                    )
262                };
263                if let Err(e) = post {
264                    self.healthy = false;
265                    return Err(e).with_context(|| format!("verbs post_read {:?}", key));
266                }
267                match rail.verbs.poll() {
268                    Ok(got) if got == wr_id => {}
269                    Ok(got) => {
270                        self.healthy = false;
271                        bail!("verbs completion wr_id {got:#x} != expected {wr_id:#x} ({key:?})");
272                    }
273                    Err(e) => {
274                        self.healthy = false;
275                        return Err(e).with_context(|| format!("verbs poll {:?}", key));
276                    }
277                }
278            }
279        }
280
281        // 2026-09-25: SAFETY: the slot now holds the `stride` bytes fetched above.
282        let record = unsafe { std::slice::from_raw_parts(host, stride) };
283        residency_from(&spec, record, dev_va, key)
284    }
285
286    fn kind(&self) -> TierKind {
287        TierKind::Rdma
288    }
289
290    /// 2026-09-25: `false` after a failed send, receive, post or poll, or an
291    /// unexpected completion; a peer error status leaves it `true`.
292    fn healthy(&self) -> bool {
293        self.healthy
294    }
295}
metrale_storage/expert_tier_rdma.rs

metrale_storage/
expert_tier_rdma.rs
