metrale_storage/weight_peer/
serve.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The weight peer server (unix). Per connection it reads a model
4//! request, stages that model (`shard`: the manifest and read-only mmaps of its
5//! shards), sends the manifest (`wire`), and registers every shard mmap for
6//! remote read on each rail (`reg_mr(.., true)`) so the client reads the
7//! tensors with one-sided RDMA.
8//!
9//! Owner: storage (weight peer).
10//! Invariants: every `StagedModel` holds a `CommitLedger` reservation of its
11//! shard bytes, taken before its shards are mapped and released when it drops.
12
13use anyhow::{Context, Result, bail};
14use std::collections::HashMap;
15use std::io::Read;
16use std::net::{TcpListener, TcpStream, ToSocketAddrs};
17use std::path::{Path, PathBuf};
18use std::sync::{Arc, Mutex};
19
20use super::manifest::WeightManifest;
21use super::shard::{Mmap, build_manifest};
22use super::wire::{read_model_request, write_weight_manifest};
23use crate::expert_peer::MODE_VERBS;
24
25/// 2026-09-25: The peer's rails, its staging ceiling, and which model
26/// directories a client may request.
27#[derive(Clone, Debug)]
28pub struct WeightPeerConfig {
29    /// 2026-09-25: `(device, gid_idx)` per rail; a client asking for `n` rails
30    /// gets the first `n`.
31    pub rails: Vec<(String, u32)>,
32    /// 2026-09-25: Ceiling, in bytes, on the shard bytes of all staged models
33    /// together, 0 for none. Each model is charged once, when it is first
34    /// staged; per-connection registrations are not charged.
35    pub max_blade_bytes: u64,
36    /// 2026-09-25: Directories staged at startup. Unless `allow_any_path` is
37    /// set, a request must name one of them by full path or by final component.
38    pub staged_dirs: Vec<PathBuf>,
39    /// 2026-09-25: When true, a request that names any existing directory is
40    /// staged on demand.
41    pub allow_any_path: bool,
42}
43
44impl Default for WeightPeerConfig {
45    fn default() -> Self {
46        Self {
47            rails: vec![("roceP2p1s0f1".into(), 3)],
48            max_blade_bytes: 0,
49            staged_dirs: Vec::new(),
50            allow_any_path: false,
51        }
52    }
53}
54
55/// 2026-09-25: A staged model: its shard mmaps, which stay mapped across
56/// connections, its manifest, and its ledger reservation, released on drop.
57/// Each connection registers the same mappings on its own rails.
58struct StagedModel {
59    // 2026-09-25: Read only by the `metrale_rdma_verbs` `serve_verbs`.
60    #[cfg_attr(not(metrale_rdma_verbs), allow(dead_code))]
61    shard_mmaps: Vec<Mmap>,
62    manifest: WeightManifest,
63    _reservation: crate::blade_cap::Reservation,
64}
65
66type StagedMap = Arc<Mutex<HashMap<String, Arc<StagedModel>>>>;
67
68/// 2026-09-25: Stage `cfg.staged_dirs`, then accept connections on `addr`, one
69/// thread each, until the listener ends. A directory that fails to stage is
70/// logged and skipped.
71///
72/// # Errors
73/// Binding `addr` fails.
74pub fn serve<A: ToSocketAddrs>(addr: A, cfg: WeightPeerConfig) -> Result<()> {
75    let cfg = Arc::new(cfg);
76    let ledger = Arc::new(crate::blade_cap::CommitLedger::new(cfg.max_blade_bytes));
77    let staged: StagedMap = Arc::new(Mutex::new(HashMap::new()));
78
79    for dir in &cfg.staged_dirs {
80        match stage_model(&staged, &ledger, dir) {
81            Ok(m) => tracing::info!(
82                "weight-peer pre-staged {} ({} shards, {} tensors, {:.1} GiB)",
83                m.manifest.model_id,
84                m.manifest.num_shards(),
85                m.manifest.tensors.len(),
86                m.manifest.total_shard_bytes() as f64 / (1024.0 * 1024.0 * 1024.0),
87            ),
88            Err(e) => tracing::warn!("weight-peer pre-stage {} failed: {e}", dir.display()),
89        }
90    }
91
92    let listener = TcpListener::bind(addr).context("bind weight-peer listener")?;
93    let local = listener.local_addr().ok();
94    tracing::info!(
95        "weight-peer serving on {:?} (verbs rails {:?}, cap {}, allow_any_path {})",
96        local,
97        cfg.rails,
98        if cfg.max_blade_bytes == 0 {
99            "unlimited".to_string()
100        } else {
101            format!(
102                "{:.1} GiB",
103                cfg.max_blade_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
104            )
105        },
106        cfg.allow_any_path,
107    );
108
109    for conn in listener.incoming() {
110        let stream = match conn {
111            Ok(s) => s,
112            Err(e) => {
113                tracing::warn!("weight-peer accept error: {e}");
114                continue;
115            }
116        };
117        let cfg = cfg.clone();
118        let ledger = ledger.clone();
119        let staged = staged.clone();
120        std::thread::spawn(move || {
121            if let Err(e) = handle_conn(stream, &cfg, &ledger, &staged) {
122                tracing::warn!("weight-peer connection ended: {e}");
123            }
124        });
125    }
126    Ok(())
127}
128
129fn handle_conn(
130    mut stream: TcpStream,
131    cfg: &WeightPeerConfig,
132    ledger: &Arc<crate::blade_cap::CommitLedger>,
133    staged: &StagedMap,
134) -> Result<()> {
135    stream.set_nodelay(true).ok();
136
137    let request = read_model_request(&mut stream)?;
138    let dir = resolve_request(cfg, &request)?;
139
140    let model = stage_model(staged, ledger, &dir)?;
141    write_weight_manifest(&mut stream, &model.manifest).context("send manifest")?;
142
143    // 2026-09-25: The client then names its transport; only verbs is served.
144    let mut mode = [0u8; 1];
145    stream
146        .read_exact(&mut mode)
147        .context("read transport mode")?;
148    match mode[0] {
149        MODE_VERBS => serve_verbs(stream, &model, cfg),
150        other => bail!("weight-peer only serves verbs; client asked for mode {other}"),
151    }
152}
153
154/// 2026-09-25: The staged directory a request names (by full path or by final
155/// component), or with `allow_any_path` any existing directory it names.
156fn resolve_request(cfg: &WeightPeerConfig, request: &str) -> Result<PathBuf> {
157    let req = Path::new(request);
158    for d in &cfg.staged_dirs {
159        if d == req
160            || d.file_name().and_then(|n| n.to_str()) == Some(request)
161            || d.to_string_lossy() == request
162        {
163            return Ok(d.clone());
164        }
165    }
166    if cfg.allow_any_path && req.is_dir() {
167        return Ok(req.to_path_buf());
168    }
169    bail!(
170        "model '{request}' is not staged (and allow_any_path is off); \
171         pass it to the peer with --stage <dir>"
172    );
173}
174
175/// 2026-09-25: The staged model for `dir`, staging it first if needed: build
176/// the manifest, reserve its shard bytes on the ledger, then mmap every shard.
177/// Keyed by the directory string, so a second call returns the same model.
178fn stage_model(
179    staged: &StagedMap,
180    ledger: &Arc<crate::blade_cap::CommitLedger>,
181    dir: &Path,
182) -> Result<Arc<StagedModel>> {
183    let key = dir.to_string_lossy().into_owned();
184    {
185        let map = staged.lock().unwrap();
186        if let Some(m) = map.get(&key) {
187            return Ok(m.clone());
188        }
189    }
190
191    let (shard_paths, manifest) = build_manifest(dir, &key)?;
192    // 2026-09-25: The reservation is taken before any shard is mapped, and is
193    // released if a mapping below fails.
194    let reservation = ledger
195        .try_reserve(manifest.total_shard_bytes())
196        .context("weight blade cap")?;
197
198    let mut shard_mmaps = Vec::with_capacity(shard_paths.len());
199    for p in &shard_paths {
200        shard_mmaps.push(Mmap::open_ro(p).with_context(|| format!("mmap {}", p.display()))?);
201    }
202
203    let model = Arc::new(StagedModel {
204        shard_mmaps,
205        manifest,
206        _reservation: reservation,
207    });
208    let mut map = staged.lock().unwrap();
209    // 2026-09-25: If another thread staged the same directory meanwhile, its
210    // model is kept and this one, with its reservation, is dropped.
211    Ok(map.entry(key).or_insert(model).clone())
212}
213
214/// 2026-09-25: Without `metrale_rdma_verbs` a verbs request is refused.
215#[cfg(not(metrale_rdma_verbs))]
216fn serve_verbs(
217    _stream: TcpStream,
218    _model: &Arc<StagedModel>,
219    _cfg: &WeightPeerConfig,
220) -> Result<()> {
221    bail!("client requested verbs transport but this peer was built without rdma-core");
222}
223
224#[cfg(metrale_rdma_verbs)]
225fn serve_verbs(
226    mut stream: TcpStream,
227    model: &Arc<StagedModel>,
228    cfg: &WeightPeerConfig,
229) -> Result<()> {
230    use crate::expert_peer::{STATUS_OK, VerbsClientParams, VerbsServerParams, write_server_rails};
231    use metrale_gpu_sys::verbs::Verbs;
232    use std::io::Write;
233
234    let num_shards = model.shard_mmaps.len();
235
236    let mut b1 = [0u8; 1];
237    stream.read_exact(&mut b1).context("read n_rails")?;
238    let n_rails = b1[0] as usize;
239    if n_rails == 0 || n_rails > cfg.rails.len() {
240        bail!(
241            "client asked for {n_rails} rails; peer has {}",
242            cfg.rails.len()
243        );
244    }
245
246    // 2026-09-25: One QP per rail, each with its own PSN. Nothing is charged
247    // to the ledger here: `stage_model` charged the model once.
248    let pid = std::process::id();
249    let mut rails: Vec<Verbs> = Vec::with_capacity(n_rails);
250    for (i, (dev, gid)) in cfg.rails.iter().take(n_rails).enumerate() {
251        let psn = (0x77_7777 ^ pid ^ ((i as u32) << 20)) & 0xff_ffff;
252        rails.push(Verbs::create(dev, *gid, psn)?);
253    }
254
255    // 2026-09-25: Every shard mmap is registered for remote read on every rail:
256    // one rkey per (rail, shard), all at the mapping's own address.
257    let mut per_rail_shards: Vec<Vec<(u64, u32)>> = (0..n_rails)
258        .map(|_| Vec::with_capacity(num_shards))
259        .collect();
260    for m in &model.shard_mmaps {
261        for (ri, v) in rails.iter_mut().enumerate() {
262            // 2026-09-25: SAFETY: the mapping covers `m.len` bytes at `m.addr` and
263            // belongs to `model`, which outlives `rails` (dropped at the end).
264            let keys = unsafe { v.reg_mr(m.addr as *mut _, m.len, true)? };
265            per_rail_shards[ri].push((m.addr as u64, keys.rkey));
266        }
267    }
268
269    // 2026-09-25: One `VerbsServerParams` per rail; its `layers` holds each
270    // shard's `(base, rkey)` in shard order.
271    let sp: Vec<VerbsServerParams> = rails
272        .iter()
273        .enumerate()
274        .map(|(ri, v)| VerbsServerParams {
275            qpn: v.qpn(),
276            psn: v.psn(),
277            gid: v.gid(),
278            layers: std::mem::take(&mut per_rail_shards[ri]),
279        })
280        .collect();
281    write_server_rails(&mut stream, &sp).context("send verbs server params")?;
282
283    stream.read_exact(&mut b1).context("read client n_rails")?;
284    if b1[0] as usize != n_rails {
285        bail!("client rail count mismatch");
286    }
287    for v in rails.iter_mut() {
288        let cp = VerbsClientParams::read_from(&mut stream).context("read verbs client params")?;
289        v.connect(cp.qpn, cp.psn, &cp.gid)?;
290    }
291    stream
292        .write_all(&[STATUS_OK])
293        .context("send verbs ready ack")?;
294    tracing::info!(
295        "weight-peer verbs client connected to {} ({n_rails} rail(s), {num_shards} shard MRs/rail)",
296        model.manifest.model_id,
297    );
298
299    // 2026-09-25: The client reads the tensors itself; wait for it to close
300    // the connection.
301    let mut sink = [0u8; 8];
302    loop {
303        match stream.read(&mut sink) {
304            Ok(0) => break,
305            Ok(_) => {}
306            Err(_) => break,
307        }
308    }
309    // 2026-09-25: The rails, and with them the MRs, are dropped while `model`
310    // still holds the mappings.
311    drop(rails);
312    Ok(())
313}
metrale_storage/weight_peer/serve.rs

metrale_storage/weight_peer/
serve.rs
