metrale_storage/
expert_peer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The expert peer, a process that serves an expert store's
4//! records to `RdmaTier` clients, and its wire format (little-endian):
5//!   1. On accept the peer sends the manifest, `[u32 len][len bytes of JSON]`.
6//!   2. The client sends one transport byte, `MODE_TCP` or `MODE_VERBS`.
7//!   3. TCP: the client sends `[u32 layer][u32 expert]`; the peer answers
8//!      `[STATUS_OK][record_stride bytes]` or `[STATUS_ERR]`, until the client
9//!      hangs up or sends `SHUTDOWN_MARKER` in both fields. Verbs: the peer maps
10//!      and registers every layer file for remote read, publishes its rails, and
11//!      waits for hangup while the client reads records itself.
12//!
13//! The peer uses no CUDA.
14//!
15//! Owner: metrale-storage peers.
16//! Invariants:
17//! - A verbs connection reserves `index.total_bytes()` in the ledger before it
18//!   maps or registers anything.
19
20// 2026-09-25: Used only by unix code, under the same cfg, so a Windows build
21// has no unused imports.
22#[cfg(unix)]
23use anyhow::{Context, Result, bail};
24
25// 2026-09-25: The handshake codecs are metrale-gpu-sys `wire`, re-exported for
26// this module's users; crates/gpu-sys/tests/wire_roundtrip.rs tests them.
27pub use metrale_gpu_sys::wire::{
28    MODE_TCP, MODE_VERBS, STATUS_ERR, STATUS_OK, VerbsClientParams, VerbsServerParams,
29    read_server_rails, write_server_rails,
30};
31
32/// 2026-09-25: In both fields of a TCP request, asks the peer to close the
33/// connection.
34pub const SHUTDOWN_MARKER: u32 = u32::MAX;
35
36/// 2026-09-25: Encode a TCP request, `[u32 layer][u32 expert]` little-endian.
37pub fn encode_request(layer: u32, expert: u32) -> [u8; 8] {
38    let mut b = [0u8; 8];
39    b[0..4].copy_from_slice(&layer.to_le_bytes());
40    b[4..8].copy_from_slice(&expert.to_le_bytes());
41    b
42}
43
44/// 2026-09-25: Decode a TCP request; the inverse of `encode_request`.
45pub fn decode_request(b: &[u8; 8]) -> (u32, u32) {
46    let layer = u32::from_le_bytes([b[0], b[1], b[2], b[3]]);
47    let expert = u32::from_le_bytes([b[4], b[5], b[6], b[7]]);
48    (layer, expert)
49}
50
51#[cfg(unix)]
52pub use server_impl::{RdmaConfig, serve};
53
54#[cfg(unix)]
55mod server_impl {
56    use super::*;
57    use crate::expert::ExpertKey;
58    use crate::expert_pack::ExpertFileReader;
59    use std::io::{Read, Write};
60    use std::net::{TcpListener, TcpStream, ToSocketAddrs};
61    use std::path::{Path, PathBuf};
62    use std::sync::Arc;
63
64    /// 2026-09-25: Settings of the verbs transport, unused by TCP clients. A
65    /// client asking for n rails gets the first n `(device, gid_idx)` entries;
66    /// the peer registers every layer file on each.
67    #[derive(Clone, Debug)]
68    pub struct RdmaConfig {
69        pub rails: Vec<(String, u32)>,
70        /// 2026-09-25: Ceiling on store bytes registered across concurrent verbs
71        /// connections; 0 means none. A connection reserves the store size once,
72        /// whatever its rail count.
73        pub max_blade_bytes: u64,
74    }
75
76    impl Default for RdmaConfig {
77        fn default() -> Self {
78            Self {
79                rails: vec![("roceP2p1s0f1".into(), 3)],
80                max_blade_bytes: 0,
81            }
82        }
83    }
84
85    /// 2026-09-25: Serve the store in `dir` on `addr`, one thread per
86    /// connection. Accept errors are logged and skipped; only opening the store
87    /// or binding returns an error.
88    pub fn serve<A: ToSocketAddrs>(dir: &Path, addr: A, rdma: RdmaConfig) -> Result<()> {
89        let reader = Arc::new(ExpertFileReader::open(dir)?);
90        let manifest = serde_json::to_vec(reader.index())?;
91        let dir: Arc<PathBuf> = Arc::new(dir.to_path_buf());
92        let rdma = Arc::new(rdma);
93        // 2026-09-25: One ledger per `serve`, shared by its connection threads.
94        let ledger = Arc::new(crate::blade_cap::CommitLedger::new(rdma.max_blade_bytes));
95        let listener = TcpListener::bind(addr).context("bind expert-peer listener")?;
96        let local = listener.local_addr().ok();
97        tracing::info!(
98            "expert-peer serving {} ({} layers, {} experts, stride {}) on {:?} \
99             (verbs rails {:?}, cap {})",
100            dir.display(),
101            reader.index().num_moe_layers,
102            reader.index().num_experts,
103            reader.index().record_stride,
104            local,
105            rdma.rails,
106            if rdma.max_blade_bytes == 0 {
107                "unlimited".to_string()
108            } else {
109                format!(
110                    "{:.1} GiB",
111                    rdma.max_blade_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
112                )
113            },
114        );
115        for conn in listener.incoming() {
116            let stream = match conn {
117                Ok(s) => s,
118                Err(e) => {
119                    tracing::warn!("expert-peer accept error: {e}");
120                    continue;
121                }
122            };
123            let reader = reader.clone();
124            let manifest = manifest.clone();
125            let dir = dir.clone();
126            let rdma = rdma.clone();
127            let ledger = ledger.clone();
128            std::thread::spawn(move || {
129                if let Err(e) = handle_conn(stream, &reader, &manifest, &dir, &rdma, &ledger) {
130                    tracing::warn!("expert-peer connection ended: {e}");
131                }
132            });
133        }
134        Ok(())
135    }
136
137    fn handle_conn(
138        mut stream: TcpStream,
139        reader: &ExpertFileReader,
140        manifest: &[u8],
141        dir: &Path,
142        rdma: &RdmaConfig,
143        ledger: &Arc<crate::blade_cap::CommitLedger>,
144    ) -> Result<()> {
145        stream.set_nodelay(true).ok();
146        stream.write_all(&(manifest.len() as u32).to_le_bytes())?;
147        stream.write_all(manifest)?;
148
149        // 2026-09-25: TCP reads records on demand and pins nothing, so only the
150        // verbs transport, which maps and registers the store, is charged.
151        let mut mode = [0u8; 1];
152        stream
153            .read_exact(&mut mode)
154            .context("read transport mode")?;
155        match mode[0] {
156            MODE_TCP => serve_tcp(stream, reader),
157            MODE_VERBS => serve_verbs(stream, reader, dir, rdma, ledger),
158            other => bail!("client requested unknown transport mode {other}"),
159        }
160    }
161
162    /// 2026-09-25: Answer each request with the record, or `STATUS_ERR` when it
163    /// cannot be read, until the client hangs up or sends `SHUTDOWN_MARKER`.
164    fn serve_tcp(mut stream: TcpStream, reader: &ExpertFileReader) -> Result<()> {
165        let stride = reader.index().record_stride as usize;
166        let mut req = [0u8; 8];
167        loop {
168            if stream.read_exact(&mut req).is_err() {
169                break;
170            }
171            let (layer, expert) = decode_request(&req);
172            if layer == SHUTDOWN_MARKER && expert == SHUTDOWN_MARKER {
173                break;
174            }
175            match reader.read_record_raw(ExpertKey::new(layer, expert)) {
176                Ok(rec) => {
177                    debug_assert_eq!(rec.len(), stride);
178                    stream.write_all(&[STATUS_OK])?;
179                    stream.write_all(&rec)?;
180                }
181                Err(e) => {
182                    tracing::warn!("expert-peer read {layer}/{expert}: {e}");
183                    stream.write_all(&[STATUS_ERR])?;
184                }
185            }
186        }
187        Ok(())
188    }
189
190    #[cfg(not(metrale_rdma_verbs))]
191    fn serve_verbs(
192        _stream: TcpStream,
193        _reader: &ExpertFileReader,
194        _dir: &Path,
195        _rdma: &RdmaConfig,
196        _ledger: &Arc<crate::blade_cap::CommitLedger>,
197    ) -> Result<()> {
198        bail!("client requested verbs transport but this peer was built without rdma-core");
199    }
200
201    /// 2026-09-25: Map and register each layer file for remote read, publish
202    /// each rail's QP and per-layer `(base, rkey)`, connect to the client's QPs,
203    /// then wait for hangup; the client reads the records itself.
204    #[cfg(metrale_rdma_verbs)]
205    fn serve_verbs(
206        mut stream: TcpStream,
207        reader: &ExpertFileReader,
208        dir: &Path,
209        rdma: &RdmaConfig,
210        ledger: &Arc<crate::blade_cap::CommitLedger>,
211    ) -> Result<()> {
212        use metrale_gpu_sys::verbs::Verbs;
213
214        let index = reader.index();
215        let num_layers = index.num_moe_layers;
216
217        let mut b1 = [0u8; 1];
218        stream.read_exact(&mut b1).context("read n_rails")?;
219        let n_rails = b1[0] as usize;
220        if n_rails == 0 || n_rails > rdma.rails.len() {
221            bail!(
222                "client asked for {n_rails} rails; peer has {}",
223                rdma.rails.len()
224            );
225        }
226
227        // 2026-09-25: The rails register the same mapped pages, so the store is
228        // charged once, before anything is mapped; the reservation is released
229        // on every return.
230        let _reservation = ledger
231            .try_reserve(index.total_bytes())
232            .context("expert blade cap")?;
233
234        // 2026-09-25: The PSN differs per rail and per process.
235        let pid = std::process::id();
236        let mut rails: Vec<Verbs> = Vec::with_capacity(n_rails);
237        for (i, (dev, gid)) in rdma.rails.iter().take(n_rails).enumerate() {
238            let psn = (0x424242 ^ pid ^ ((i as u32) << 20)) & 0xff_ffff;
239            rails.push(Verbs::create(dev, *gid, psn)?);
240        }
241
242        // 2026-09-25: Each layer file is mapped once and registered on every
243        // rail: one rkey per (rail, layer), one base address per layer.
244        let mut mmaps: Vec<Mmap> = Vec::with_capacity(num_layers as usize);
245        let mut per_rail_layers: Vec<Vec<(u64, u32)>> = (0..n_rails)
246            .map(|_| Vec::with_capacity(num_layers as usize))
247            .collect();
248        for l in 0..num_layers {
249            let path = dir.join(index.file_name(l));
250            let m = Mmap::open_ro(&path).with_context(|| format!("mmap {}", path.display()))?;
251            for (ri, v) in rails.iter_mut().enumerate() {
252                // 2026-09-25: SAFETY: the mapping covers `m.len` bytes at `m.addr`.
253                // The explicit drops at the end release `rails` first; an early
254                // return drops `mmaps` first (reverse declaration order).
255                let keys = unsafe { v.reg_mr(m.addr as *mut _, m.len, true)? };
256                per_rail_layers[ri].push((m.addr as u64, keys.rkey));
257            }
258            mmaps.push(m);
259        }
260
261        let sp: Vec<VerbsServerParams> = rails
262            .iter()
263            .enumerate()
264            .map(|(ri, v)| VerbsServerParams {
265                qpn: v.qpn(),
266                psn: v.psn(),
267                gid: v.gid(),
268                layers: std::mem::take(&mut per_rail_layers[ri]),
269            })
270            .collect();
271        write_server_rails(&mut stream, &sp).context("send verbs server params")?;
272
273        stream.read_exact(&mut b1).context("read client n_rails")?;
274        if b1[0] as usize != n_rails {
275            bail!("client rail count mismatch");
276        }
277        for v in rails.iter_mut() {
278            let cp =
279                VerbsClientParams::read_from(&mut stream).context("read verbs client params")?;
280            v.connect(cp.qpn, cp.psn, &cp.gid)?;
281        }
282        stream
283            .write_all(&[STATUS_OK])
284            .context("send verbs ready ack")?;
285        tracing::info!(
286            "expert-peer verbs client connected ({n_rails} rail(s), {} layer MRs/rail)",
287            num_layers,
288        );
289
290        // 2026-09-25: The client reads every record itself; the peer holds the
291        // registrations until hangup.
292        let mut sink = [0u8; 8];
293        loop {
294            match stream.read(&mut sink) {
295                Ok(0) => break,
296                Ok(_) => {}
297                Err(_) => break,
298            }
299        }
300        // 2026-09-25: Registrations go before the mappings they cover.
301        drop(rails);
302        drop(mmaps);
303        Ok(())
304    }
305
306    /// 2026-09-25: A read-only shared mapping of a whole file, unmapped on drop.
307    #[cfg(metrale_rdma_verbs)]
308    struct Mmap {
309        addr: *mut libc::c_void,
310        len: usize,
311    }
312
313    #[cfg(metrale_rdma_verbs)]
314    impl Mmap {
315        fn open_ro(path: &Path) -> Result<Self> {
316            use std::os::fd::AsRawFd;
317            let f = std::fs::File::open(path)?;
318            let len = f.metadata()?.len() as usize;
319            if len == 0 {
320                bail!("empty layer file {}", path.display());
321            }
322            // 2026-09-25: SAFETY: `f` is an open read-only file of `len` bytes. The
323            // mapping stays valid after the file is closed.
324            let addr = unsafe {
325                libc::mmap(
326                    std::ptr::null_mut(),
327                    len,
328                    libc::PROT_READ,
329                    libc::MAP_SHARED,
330                    f.as_raw_fd(),
331                    0,
332                )
333            };
334            if addr == libc::MAP_FAILED {
335                bail!(
336                    "mmap {} failed: {}",
337                    path.display(),
338                    std::io::Error::last_os_error()
339                );
340            }
341            Ok(Self { addr, len })
342        }
343    }
344
345    #[cfg(metrale_rdma_verbs)]
346    impl Drop for Mmap {
347        fn drop(&mut self) {
348            // 2026-09-25: SAFETY: addr/len come from a successful mmap; drop runs once.
349            unsafe { libc::munmap(self.addr, self.len) };
350        }
351    }
352}
353
354/// 2026-09-25: Read and parse the length-prefixed manifest at the start of a
355/// connection (used by `expert_tier_rdma`). Fails on a length of 0 or above
356/// 16 MiB.
357#[cfg(unix)]
358pub fn read_manifest<R: std::io::Read>(stream: &mut R) -> Result<crate::expert_pack::ExpertIndex> {
359    let mut lenb = [0u8; 4];
360    stream
361        .read_exact(&mut lenb)
362        .context("read manifest length")?;
363    let len = u32::from_le_bytes(lenb) as usize;
364    if len == 0 || len > 16 * 1024 * 1024 {
365        bail!("implausible peer manifest length: {len}");
366    }
367    let mut buf = vec![0u8; len];
368    stream.read_exact(&mut buf).context("read manifest json")?;
369    let index: crate::expert_pack::ExpertIndex =
370        serde_json::from_slice(&buf).context("parse peer manifest")?;
371    Ok(index)
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    #[test]
379    fn request_round_trips() {
380        let b = encode_request(7, 42);
381        assert_eq!(decode_request(&b), (7, 42));
382        let s = encode_request(SHUTDOWN_MARKER, SHUTDOWN_MARKER);
383        assert_eq!(decode_request(&s), (SHUTDOWN_MARKER, SHUTDOWN_MARKER));
384    }
385}
metrale_storage/expert_peer.rs

metrale_storage/
expert_peer.rs
