metrale_storage/weight_peer.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The weight peer: a daemon that holds models' safetensors shards mapped in
4//! its RAM and lets a client read any tensor with one-sided RDMA. The client is
5//! `RdmaWeightLoader` in model-weights (`weight_tier_rdma.rs`).
6//!
7//! Protocol on one TCP connection, lengths little-endian (`wire`, `serve`):
8//! 1. The client sends the model request, `[u32 len][len bytes of id or path]`.
9//! 2. The peer stages that model (once per directory) and sends the manifest,
10//! `[u32 len][len bytes of JSON]` (`WeightManifest`).
11//! 3. The client sends `[u8 transport mode]`; only `MODE_VERBS` is served.
12//! 4. The client sends `[u8 n_rails]`. The peer registers every shard mapping for remote
13//! read on each rail and sends one `VerbsServerParams` per rail, whose `layers` holds
14//! each shard's `(base, rkey)`. The client sends `n_rails` again and its QP parameters
15//! per rail; the peer connects, acknowledges, and waits for the client to close the
16//! connection while it reads the tensors.
17//!
18//! Owner: storage (weight peer).
19//! Invariants: `manifest` and `wire` compile on every platform; `serve` and `shard` only
20//! on unix.
21
22mod manifest;
23#[cfg(unix)]
24mod serve;
25#[cfg(unix)]
26mod shard;
27mod wire;
28
29pub use manifest::{WeightManifest, WeightTensorRecord, rail_for_tensor, tensor_remote_addr};
30#[cfg(unix)]
31pub use serve::{WeightPeerConfig, serve};
32pub use wire::{
33 MODEL_REQUEST_MAX, read_model_request, read_weight_manifest, write_model_request,
34 write_weight_manifest,
35};