atlas_rdma/
wire.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// The RDMA handshake wire codecs, shared by every Atlas RDMA client and the
4// peer daemons (the daemons re-export these, so client and server speak one
5// codec).
6//
7// ** GOLDEN WIRE FORMAT ** — every byte layout here is a frozen external
8// contract (already-deployed peers depend on it). All integers little-endian. The byte
9// vectors are pinned by `tests/wire_roundtrip.rs` and `tests/transcript_golden.rs`,
10// and the frozen constant values by `frozen_wire_constants` (wire_roundtrip.rs).
11//
12// Un-gated on purpose: pure `std::io` + anyhow, so it compiles and unit-tests
13// on the metal/ATLAS_SKIP_BUILD build with no rdma-core.
14
15use anyhow::{Context, Result, bail};
16
17// ── Shared status / transport-mode bytes ──
18pub const STATUS_OK: u8 = 0;
19pub const STATUS_ERR: u8 = 1;
20/// Two-sided TCP record streaming (the expert record path).
21pub const MODE_TCP: u8 = 0;
22/// One-sided RDMA READ over verbs: the server publishes its store's MRs and
23/// the client READs records directly into its arena.
24pub const MODE_VERBS: u8 = 1;
25
26/// The server's half of the verbs handshake (RO dialect: expert / weight /
27/// LoRA tiers): its QP identity plus, per MoE layer (or per shard), the base
28/// virtual address + rkey of that entry's registered MR.
29/// `remote_addr(layer, expert) = layers[layer].0 + expert * record_stride`.
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct VerbsServerParams {
32    pub qpn: u32,
33    pub psn: u32,
34    pub gid: [u8; 16],
35    /// `(mr_base_addr, rkey)` for each MoE layer (expert tier) or shard
36    /// (weight/LoRA tiers), index-addressed.
37    pub layers: Vec<(u64, u32)>,
38}
39
40/// The client's half: just its QP identity (its arena MR is local-only).
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub struct VerbsClientParams {
43    pub qpn: u32,
44    pub psn: u32,
45    pub gid: [u8; 16],
46}
47
48/// The peer's half of the RW-blade handshake (KV overflow / SSM snapshots):
49/// its QP identity + the single RW MR.
50#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51pub struct CacheServerParams {
52    pub qpn: u32,
53    pub psn: u32,
54    pub gid: [u8; 16],
55    pub base_addr: u64,
56    pub rkey: u32,
57}
58
59/// Anything that carries a remote QP identity a client rail can `connect` to.
60/// Implemented by both server-param dialects so the RailSet handshake tail is
61/// shared without homogenizing the two wire layouts.
62pub trait RemoteQp {
63    fn qp_identity(&self) -> (u32, u32, [u8; 16]);
64}
65
66impl RemoteQp for VerbsServerParams {
67    fn qp_identity(&self) -> (u32, u32, [u8; 16]) {
68        (self.qpn, self.psn, self.gid)
69    }
70}
71
72impl RemoteQp for CacheServerParams {
73    fn qp_identity(&self) -> (u32, u32, [u8; 16]) {
74        (self.qpn, self.psn, self.gid)
75    }
76}
77
78impl VerbsServerParams {
79    /// Wire form: `[u32 qpn][u32 psn][16 gid][u32 n_layers]{[u64 base][u32 rkey]}*`.
80    pub fn write_to<W: std::io::Write>(&self, w: &mut W) -> Result<()> {
81        w.write_all(&self.qpn.to_le_bytes())?;
82        w.write_all(&self.psn.to_le_bytes())?;
83        w.write_all(&self.gid)?;
84        w.write_all(&(self.layers.len() as u32).to_le_bytes())?;
85        for (base, rkey) in &self.layers {
86            w.write_all(&base.to_le_bytes())?;
87            w.write_all(&rkey.to_le_bytes())?;
88        }
89        Ok(())
90    }
91
92    pub fn read_from<R: std::io::Read>(r: &mut R) -> Result<Self> {
93        let qpn = read_u32(r)?;
94        let psn = read_u32(r)?;
95        let mut gid = [0u8; 16];
96        r.read_exact(&mut gid).context("read server gid")?;
97        let n = read_u32(r)? as usize;
98        if n == 0 || n > 4096 {
99            bail!("implausible verbs layer count: {n}");
100        }
101        let mut layers = Vec::with_capacity(n);
102        for _ in 0..n {
103            let mut b8 = [0u8; 8];
104            r.read_exact(&mut b8).context("read mr base")?;
105            let base = u64::from_le_bytes(b8);
106            let rkey = read_u32(r)?;
107            layers.push((base, rkey));
108        }
109        Ok(Self {
110            qpn,
111            psn,
112            gid,
113            layers,
114        })
115    }
116}
117
118impl VerbsClientParams {
119    /// Wire form: `[u32 qpn][u32 psn][16 gid]`.
120    pub fn write_to<W: std::io::Write>(&self, w: &mut W) -> Result<()> {
121        w.write_all(&self.qpn.to_le_bytes())?;
122        w.write_all(&self.psn.to_le_bytes())?;
123        w.write_all(&self.gid)?;
124        Ok(())
125    }
126
127    pub fn read_from<R: std::io::Read>(r: &mut R) -> Result<Self> {
128        let qpn = read_u32(r)?;
129        let psn = read_u32(r)?;
130        let mut gid = [0u8; 16];
131        r.read_exact(&mut gid).context("read client gid")?;
132        Ok(Self { qpn, psn, gid })
133    }
134}
135
136impl CacheServerParams {
137    /// Wire form: `[u32 qpn][u32 psn][16 gid][u64 base_addr][u32 rkey]`.
138    pub fn write_to<W: std::io::Write>(&self, w: &mut W) -> Result<()> {
139        w.write_all(&self.qpn.to_le_bytes())?;
140        w.write_all(&self.psn.to_le_bytes())?;
141        w.write_all(&self.gid)?;
142        w.write_all(&self.base_addr.to_le_bytes())?;
143        w.write_all(&self.rkey.to_le_bytes())?;
144        Ok(())
145    }
146
147    pub fn read_from<R: std::io::Read>(r: &mut R) -> Result<Self> {
148        let mut b4 = [0u8; 4];
149        let mut b8 = [0u8; 8];
150        let mut gid = [0u8; 16];
151        r.read_exact(&mut b4).context("kv qpn")?;
152        let qpn = u32::from_le_bytes(b4);
153        r.read_exact(&mut b4).context("kv psn")?;
154        let psn = u32::from_le_bytes(b4);
155        r.read_exact(&mut gid).context("kv gid")?;
156        r.read_exact(&mut b8).context("kv base")?;
157        let base_addr = u64::from_le_bytes(b8);
158        r.read_exact(&mut b4).context("kv rkey")?;
159        let rkey = u32::from_le_bytes(b4);
160        Ok(Self {
161            qpn,
162            psn,
163            gid,
164            base_addr,
165            rkey,
166        })
167    }
168}
169
170fn read_u32<R: std::io::Read>(r: &mut R) -> Result<u32> {
171    let mut b = [0u8; 4];
172    r.read_exact(&mut b).context("read u32")?;
173    Ok(u32::from_le_bytes(b))
174}
175
176/// Frame N per-rail `VerbsServerParams` for the dual-rail RO tiers: a leading
177/// `[u8 n_rails]` count followed by each rail's params — exactly how the RW
178/// blade frames its per-rail `CacheServerParams`. Single-rail (`n == 1`) is the
179/// default, byte-for-byte the pre-dual-rail path plus the one-byte count prefix.
180pub fn write_server_rails<W: std::io::Write>(w: &mut W, rails: &[VerbsServerParams]) -> Result<()> {
181    if rails.is_empty() || rails.len() > 8 {
182        bail!("implausible server rail count: {}", rails.len());
183    }
184    w.write_all(&[rails.len() as u8])?;
185    for sp in rails {
186        sp.write_to(w)?;
187    }
188    Ok(())
189}
190
191/// Read `want` per-rail `VerbsServerParams` framed by a leading `[u8 n_rails]`.
192/// Bails if the framed count is zero, absurd (> 8), or != `want` — the client
193/// already negotiated `want` rails, so any other count is a protocol error.
194pub fn read_server_rails<R: std::io::Read>(
195    r: &mut R,
196    want: usize,
197) -> Result<Vec<VerbsServerParams>> {
198    let mut b1 = [0u8; 1];
199    r.read_exact(&mut b1).context("read server rail count")?;
200    let n = b1[0] as usize;
201    if n == 0 || n > 8 {
202        bail!("implausible server rail count: {n}");
203    }
204    if n != want {
205        bail!("server framed {n} rails but client negotiated {want}");
206    }
207    let mut rails = Vec::with_capacity(n);
208    for _ in 0..n {
209        rails.push(VerbsServerParams::read_from(r)?);
210    }
211    Ok(rails)
212}
atlas_rdma/wire.rs

atlas_rdma/
wire.rs
