metrale_storage/
rdma_snapshot.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `RdmaSnapshotArena`, the synchronous RDMA transport of the SSM snapshot
4//! spill tier, addressed by byte offset or, in paging mode, by key.
5//!
6//! Owner: storage, SSM snapshot tier.
7//! Invariants:
8//! - Every transfer holds the arena's lock throughout, and returns `Ok` only after
9//!   its RDMA work has completed.
10//! - Without the `cuda` feature and `metrale_rdma_verbs`, both constructors fail.
11//!
12//! Both modes send the v2 paging header with kind `SSM`. Raw mode (`connect`, which
13//! sends `blob_bytes == 0`) gets a private arena and the caller chooses offsets.
14//! Paging mode (`connect_paging`) lets the peer own residency, driven by ALLOC,
15//! COMMIT, GET and REMOVE on the TCP stream. This transport moves host bytes only;
16//! `SsmSnapshotPool::spill_slot` and `fault_in_slot` in metrale-model-engine gather
17//! the state and order the device streams.
18
19// 2026-09-25: Without `cuda` and verbs a stub whose constructors fail takes the real
20// type's place, so callers compile unchanged.
21#[cfg(all(feature = "cuda", metrale_rdma_verbs))]
22pub use imp::RdmaSnapshotArena;
23#[cfg(not(all(feature = "cuda", metrale_rdma_verbs)))]
24pub use stub::RdmaSnapshotArena;
25
26#[cfg(not(all(feature = "cuda", metrale_rdma_verbs)))]
27mod stub {
28    use anyhow::{Result, bail};
29    /// 2026-09-25: Stand-in when RDMA verbs or CUDA are not built. Both constructors
30    /// fail; the other methods panic.
31    pub struct RdmaSnapshotArena;
32    impl RdmaSnapshotArena {
33        pub fn connect(_addr: &str, _arena_bytes: u64, _blob_bytes: usize) -> Result<Self> {
34            bail!("RDMA snapshot tier not built (needs feature `cuda` + metrale_rdma_verbs)")
35        }
36        pub fn connect_paging(_addr: &str, _arena_bytes: u64, _blob_bytes: usize) -> Result<Self> {
37            bail!("RDMA snapshot tier not built (needs feature `cuda` + metrale_rdma_verbs)")
38        }
39        pub fn write(&self, _offset: u64, _bytes: &[u8]) -> Result<()> {
40            unreachable!("stub RdmaSnapshotArena is never constructed")
41        }
42        pub fn read(&self, _offset: u64, _out: &mut [u8]) -> Result<()> {
43            unreachable!("stub RdmaSnapshotArena is never constructed")
44        }
45        pub fn paging_put(&self, _key: u64, _bytes: &[u8]) -> Result<()> {
46            unreachable!("stub RdmaSnapshotArena is never constructed")
47        }
48        pub fn paging_get(&self, _key: u64, _out: &mut [u8]) -> Result<bool> {
49            unreachable!("stub RdmaSnapshotArena is never constructed")
50        }
51        pub fn paging_remove(&self, _key: u64) -> Result<()> {
52            unreachable!("stub RdmaSnapshotArena is never constructed")
53        }
54    }
55}
56
57#[cfg(all(feature = "cuda", metrale_rdma_verbs))]
58mod imp {
59    use std::io::Write;
60    use std::net::TcpStream;
61    use std::sync::Mutex;
62
63    use anyhow::{Result, bail};
64
65    use crate::cuda_min::PinnedBuffer;
66    use metrale_gpu_sys::env::{first_set, first_set_u32};
67    use metrale_gpu_sys::railset::{RailSet, RailSpec};
68    use metrale_gpu_sys::verbs::Verbs;
69
70    /// 2026-09-25: One rail: its QP and one registered `blob_bytes` bounce.
71    struct SnapRail {
72        verbs: Verbs,
73        bounce: PinnedBuffer,
74        lkey: u32,
75        remote_rkey: u32,
76        /// 2026-09-25: This rail's lkey for the shared staging buffer; 0 when staging
77        /// is off. Every rail registers the same buffer, so a chunk lands at its own
78        /// offset whichever rail moves it.
79        staging_lkey: u32,
80    }
81
82    /// 2026-09-25: Mutable transport state, behind one lock because the public
83    /// methods take `&self`.
84    struct ArenaInner {
85        rails: Vec<SnapRail>,
86        /// 2026-09-25: One contiguous `blob_bytes` staging buffer for the striped path
87        /// (`METRALE_SSM_STAGING=1`); `None` moves each blob as one WR through a
88        /// rail's bounce.
89        staging: Option<PinnedBuffer>,
90        rr: usize,
91        next_wr: u64,
92        /// 2026-09-25: Raw mode: held open, idle, for the connection's lifetime.
93        /// Paging mode: the control channel for ALLOC, COMMIT, GET and REMOVE.
94        stream: TcpStream,
95    }
96
97    /// 2026-09-25: RDMA snapshot arena. `write` and `read` move one `blob_bytes`
98    /// blob to or from `base + offset`; the `paging_*` methods address blobs by key.
99    pub struct RdmaSnapshotArena {
100        inner: Mutex<ArenaInner>,
101        remote_base: u64,
102        blob_bytes: usize,
103    }
104
105    // 2026-09-25: SAFETY: every access to the verbs and bounce state goes through
106    // `inner`'s Mutex. `Verbs` and `PinnedBuffer` are `Send`.
107    unsafe impl Send for RdmaSnapshotArena {}
108    unsafe impl Sync for RdmaSnapshotArena {}
109
110    impl RdmaSnapshotArena {
111        /// 2026-09-25: Raw-mode connect to the peer at `addr` for an arena of
112        /// `arena_bytes`, registering `blob_bytes` bounces. Rails as `RdmaKvBackend`:
113        /// rail 0 from `METRALE_EXPERT_RDMA_DEV`/`GID`, rail 1 (only with
114        /// `METRALE_KV_DUAL_RAIL=1`) from `METRALE_KV_RAIL2_DEV`/`GID`.
115        pub fn connect(addr: &str, arena_bytes: u64, blob_bytes: usize) -> Result<Self> {
116            Self::connect_inner(addr, arena_bytes, blob_bytes, false)
117        }
118
119        /// 2026-09-25: Paging-mode connect: the peer owns residency, caching blobs in
120        /// its arena over an NVMe swap file, and the client uses `paging_put`,
121        /// `paging_get` and `paging_remove`. The peer refuses it without a swap
122        /// directory (`RdmaConfig::swap_dir`).
123        pub fn connect_paging(addr: &str, arena_bytes: u64, blob_bytes: usize) -> Result<Self> {
124            Self::connect_inner(addr, arena_bytes, blob_bytes, true)
125        }
126
127        fn connect_inner(
128            addr: &str,
129            arena_bytes: u64,
130            blob_bytes: usize,
131            paging: bool,
132        ) -> Result<Self> {
133            // 2026-09-25: A random 24-bit PSN per rail.
134            let spec =
135                |dev: String, gid: u32| RailSpec::new(dev, gid, rand::random::<u32>() & 0xff_ffff);
136            let rail0 = spec(
137                first_set(&["METRALE_EXPERT_RDMA_DEV"], "roceP2p1s0f1"),
138                first_set_u32(&["METRALE_EXPERT_RDMA_GID"], 3),
139            );
140            let dual = std::env::var("METRALE_KV_DUAL_RAIL").ok().as_deref() == Some("1");
141            let specs: Vec<RailSpec> = if dual {
142                let rail1 = spec(
143                    first_set(&["METRALE_KV_RAIL2_DEV"], "rocep1s0f1"),
144                    first_set_u32(&["METRALE_KV_RAIL2_GID"], 3),
145                );
146                vec![rail0, rail1]
147            } else {
148                vec![rail0]
149            };
150            let n_rails = specs.len();
151
152            let mut stream = TcpStream::connect(addr)
153                .map_err(|e| anyhow::anyhow!("connect snapshot peer {addr}: {e}"))?;
154            stream.set_nodelay(true).ok();
155            // 2026-09-25: Paging mode sends the blob size; raw mode sends
156            // `blob_bytes == 0`, and the peer gives it a private arena.
157            stream.write_all(&crate::snapshot_swap::encode_paging_v2_header(
158                crate::snapshot_swap::PagingKind::SSM,
159                arena_bytes,
160                if paging { blob_bytes as u64 } else { 0 },
161            ))?;
162            let mut rs = RailSet::begin(&mut stream, &specs)?;
163
164            // 2026-09-25: One staging buffer, registered on every rail below.
165            let staging_on = std::env::var("METRALE_SSM_STAGING").ok().as_deref() == Some("1");
166            let staging = if staging_on {
167                Some(PinnedBuffer::new(blob_bytes)?)
168            } else {
169                None
170            };
171
172            // 2026-09-25: Bounce and staging MRs are LOCAL_WRITE only
173            // (`remote_read == false`).
174            let mut parts: Vec<(PinnedBuffer, u32, u32)> = Vec::with_capacity(n_rails);
175            for rail in &mut rs.rails {
176                let bounce = PinnedBuffer::new(blob_bytes)?;
177                // 2026-09-25: SAFETY: the bounce lives as long as the rail, and so the MR.
178                let keys = unsafe { rail.verbs.reg_mr(bounce.ptr, blob_bytes, false)? };
179                // 2026-09-25: SAFETY: `staging` is declared after `rails` in
180                // `ArenaInner`, so it is dropped after them.
181                let staging_lkey = match &staging {
182                    Some(s) => unsafe { rail.verbs.reg_mr(s.ptr, blob_bytes, false)?.lkey },
183                    None => 0,
184                };
185                parts.push((bounce, keys.lkey, staging_lkey));
186            }
187
188            // 2026-09-25: Every rail publishes the same arena base; the last is kept.
189            let server = rs.finish_rw(&mut stream, "snapshot peer")?;
190            let base = server.last().map(|sp| sp.base_addr).unwrap_or(0);
191            let rails: Vec<SnapRail> = rs
192                .into_verbs()
193                .into_iter()
194                .zip(parts)
195                .zip(&server)
196                .map(|((verbs, (bounce, lkey, staging_lkey)), sp)| SnapRail {
197                    verbs,
198                    bounce,
199                    lkey,
200                    remote_rkey: sp.rkey,
201                    staging_lkey,
202                })
203                .collect();
204            tracing::info!(
205                "RdmaSnapshotArena connected to {addr}: {:.1} GiB arena, {n_rails} rail(s), blob {blob_bytes} B",
206                arena_bytes as f64 / (1024.0 * 1024.0 * 1024.0),
207            );
208            Ok(Self {
209                inner: Mutex::new(ArenaInner {
210                    rails,
211                    staging,
212                    rr: 0,
213                    next_wr: 1,
214                    stream,
215                }),
216                remote_base: base,
217                blob_bytes,
218            })
219        }
220
221        #[inline]
222        pub fn blob_bytes(&self) -> usize {
223            self.blob_bytes
224        }
225
226        /// 2026-09-25: RDMA WRITE one `blob_bytes` blob to `base + offset` and wait for
227        /// its completion.
228        pub fn write(&self, offset: u64, bytes: &[u8]) -> Result<()> {
229            if bytes.len() != self.blob_bytes {
230                bail!(
231                    "snapshot write: {} != blob_bytes {}",
232                    bytes.len(),
233                    self.blob_bytes
234                );
235            }
236            let mut g = self.inner.lock().expect("snapshot arena mutex");
237            self.rdma_write_locked(&mut g, self.remote_base + offset, bytes)
238        }
239
240        /// 2026-09-25: RDMA READ one `blob_bytes` blob from `base + offset` into `out`
241        /// and wait for its completion.
242        pub fn read(&self, offset: u64, out: &mut [u8]) -> Result<()> {
243            if out.len() != self.blob_bytes {
244                bail!(
245                    "snapshot read: {} != blob_bytes {}",
246                    out.len(),
247                    self.blob_bytes
248                );
249            }
250            let mut g = self.inner.lock().expect("snapshot arena mutex");
251            self.rdma_read_locked(&mut g, self.remote_base + offset, out)
252        }
253
254        /// 2026-09-25: The next rail, round-robin, and a fresh non-zero wr id.
255        fn rail_and_wr(g: &mut ArenaInner) -> (usize, u64) {
256            let n = g.rails.len();
257            let ri = g.rr % n;
258            g.rr = g.rr.wrapping_add(1);
259            let wr = g.next_wr;
260            g.next_wr = g.next_wr.wrapping_add(1).max(1);
261            (ri, wr)
262        }
263
264        fn rdma_write_locked(&self, g: &mut ArenaInner, raddr: u64, bytes: &[u8]) -> Result<()> {
265            if g.staging.is_some() {
266                return self.rdma_staged(g, raddr, Some(bytes), None);
267            }
268            let (ri, wr) = Self::rail_and_wr(g);
269            let rail = &mut g.rails[ri];
270            // 2026-09-25: SAFETY: the bounce is a live registered MR of `blob_bytes`,
271            // and `bytes` has that length.
272            unsafe {
273                std::ptr::copy_nonoverlapping(
274                    bytes.as_ptr(),
275                    rail.bounce.ptr as *mut u8,
276                    self.blob_bytes,
277                );
278                rail.verbs.post_write(
279                    rail.bounce.ptr,
280                    rail.lkey,
281                    raddr,
282                    rail.remote_rkey,
283                    self.blob_bytes as u32,
284                    wr,
285                )?;
286            }
287            while rail.verbs.poll()? != wr {}
288            Ok(())
289        }
290
291        fn rdma_read_locked(&self, g: &mut ArenaInner, raddr: u64, out: &mut [u8]) -> Result<()> {
292            if g.staging.is_some() {
293                return self.rdma_staged(g, raddr, None, Some(out));
294            }
295            let (ri, wr) = Self::rail_and_wr(g);
296            let rail = &mut g.rails[ri];
297            // 2026-09-25: SAFETY: the bounce is a live registered MR of `blob_bytes`, and
298            // `out` has that length.
299            unsafe {
300                rail.verbs.post_read(
301                    rail.bounce.ptr,
302                    rail.lkey,
303                    raddr,
304                    rail.remote_rkey,
305                    self.blob_bytes as u32,
306                    wr,
307                )?;
308            }
309            while rail.verbs.poll()? != wr {}
310            unsafe {
311                std::ptr::copy_nonoverlapping(
312                    rail.bounce.ptr as *const u8,
313                    out.as_mut_ptr(),
314                    self.blob_bytes,
315                );
316            }
317            Ok(())
318        }
319
320        /// 2026-09-25: Move one blob through the staging buffer in chunks
321        /// (`stripe_plan`), round-robin across rails with at most `staging_depth()` in
322        /// flight per rail. `write_src` set: copy in, then WRITE; `read_dst` set: READ,
323        /// then copy out. Every chunk sits at its blob offset in both the staging
324        /// buffer and the peer arena, so one copy assembles the blob.
325        fn rdma_staged(
326            &self,
327            g: &mut ArenaInner,
328            raddr: u64,
329            write_src: Option<&[u8]>,
330            read_dst: Option<&mut [u8]>,
331        ) -> Result<()> {
332            let staging = g.staging.as_ref().expect("staging present").ptr as *mut u8;
333            let is_read = read_dst.is_some();
334            if let Some(src) = write_src {
335                unsafe { std::ptr::copy_nonoverlapping(src.as_ptr(), staging, self.blob_bytes) };
336            }
337            let n = g.rails.len();
338            let chunk = crate::snapshot_swap::staging_chunk_bytes();
339            let depth = crate::snapshot_swap::staging_depth();
340            let plan = crate::snapshot_swap::stripe_plan(self.blob_bytes, chunk, n);
341            let total: usize = plan.iter().map(|w| w.len()).sum();
342            let mut posted = vec![0usize; n];
343            let mut reaped = vec![0usize; n];
344            let mut done = 0usize;
345            while done < total {
346                for ri in 0..n {
347                    while posted[ri] < plan[ri].len() && (posted[ri] - reaped[ri]) < depth {
348                        let (off, len) = plan[ri][posted[ri]];
349                        let wr = g.next_wr;
350                        g.next_wr = g.next_wr.wrapping_add(1).max(1);
351                        let rail = &mut g.rails[ri];
352                        let local = unsafe { staging.add(off) } as *mut _;
353                        let raddr_chunk = raddr + off as u64;
354                        unsafe {
355                            if is_read {
356                                rail.verbs.post_read(
357                                    local,
358                                    rail.staging_lkey,
359                                    raddr_chunk,
360                                    rail.remote_rkey,
361                                    len as u32,
362                                    wr,
363                                )?;
364                            } else {
365                                rail.verbs.post_write(
366                                    local,
367                                    rail.staging_lkey,
368                                    raddr_chunk,
369                                    rail.remote_rkey,
370                                    len as u32,
371                                    wr,
372                                )?;
373                            }
374                        }
375                        posted[ri] += 1;
376                    }
377                }
378                for ri in 0..n {
379                    if posted[ri] > reaped[ri] {
380                        g.rails[ri].verbs.poll()?;
381                        reaped[ri] += 1;
382                        done += 1;
383                    }
384                }
385            }
386            if let Some(dst) = read_dst {
387                unsafe {
388                    std::ptr::copy_nonoverlapping(
389                        staging as *const u8,
390                        dst.as_mut_ptr(),
391                        self.blob_bytes,
392                    )
393                };
394            }
395            Ok(())
396        }
397
398        /// 2026-09-25: Store `key`'s blob: ALLOC a slot on the control channel, RDMA
399        /// WRITE the blob there, then COMMIT, all under the lock.
400        pub fn paging_put(&self, key: u64, bytes: &[u8]) -> Result<()> {
401            if bytes.len() != self.blob_bytes {
402                bail!(
403                    "paging_put: {} != blob_bytes {}",
404                    bytes.len(),
405                    self.blob_bytes
406                );
407            }
408            let mut g = self.inner.lock().expect("snapshot arena mutex");
409            let off = crate::snapshot_swap::client_alloc(&mut g.stream, key)?;
410            self.rdma_write_locked(&mut g, self.remote_base + off, bytes)?;
411            crate::snapshot_swap::client_commit(&mut g.stream, key)
412        }
413
414        /// 2026-09-25: Read `key`'s blob into `out`; `Ok(false)` when the peer does not
415        /// have it.
416        pub fn paging_get(&self, key: u64, out: &mut [u8]) -> Result<bool> {
417            if out.len() != self.blob_bytes {
418                bail!(
419                    "paging_get: {} != blob_bytes {}",
420                    out.len(),
421                    self.blob_bytes
422                );
423            }
424            let mut g = self.inner.lock().expect("snapshot arena mutex");
425            match crate::snapshot_swap::client_get(&mut g.stream, key)? {
426                Some(off) => {
427                    self.rdma_read_locked(&mut g, self.remote_base + off, out)?;
428                    Ok(true)
429                }
430                None => Ok(false),
431            }
432        }
433
434        /// 2026-09-25: Ask the peer to drop `key`.
435        pub fn paging_remove(&self, key: u64) -> Result<()> {
436            let mut g = self.inner.lock().expect("snapshot arena mutex");
437            crate::snapshot_swap::client_remove(&mut g.stream, key)
438        }
439    }
440}
metrale_storage/rdma_snapshot.rs

metrale_storage/
rdma_snapshot.rs
