spark_storage/
weight_peer.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Weight-serving peer + wire protocol — the RDMA weight-staging tier.
4//
5// Generalizes `expert_peer` from (layer, expert) expert records to ALL of a
6// model's safetensors tensors, for FAST MODEL SWAPS. A peer holds one or more
7// staged models' shard files mmap'd + `ibv_reg_mr`'d REMOTE_READ in its RAM; a
8// client (`weight_tier_rdma::RdmaWeightLoader`) requests a model by id/path,
9// reads the peer's MANIFEST, then one-sided RDMA-READs each tensor's bytes
10// straight out of the shard MRs (~24 GB/s dual-rail) instead of the ~2 GB/s USB
11// SSD. Weights are READ-ONLY → one-sided READ, no coherence — the exact
12// expert-tier pattern.
13//
14// It's a CACHE: the FIRST stage of a model into the blade faults its pages in
15// from disk (slow); every later swap reads them out of the peer's warm RAM
16// (fast). Pre-warm the rotation set by connecting once.
17//
18// Wire protocol (little-endian), connection-oriented, server responds to the
19// client's model choice first:
20//   1. Client sends the model request: `[u32 len][len bytes of model id/path]`.
21//   2. Server stages that model (mmap + parse headers, cached across
22//      connections) and sends the manifest: `[u32 len][len bytes of JSON]`
23//      (`WeightManifest` — per-tensor {name,dtype,shape,offset,len,shard}).
24//   3. Client sends `[u8 transport_mode]` (only `MODE_VERBS` is served).
25//   4. Verbs handshake (reused verbatim from `expert_peer`): `[u8 n_rails]`,
26//      then per rail a `VerbsServerParams` whose `layers` vector carries this
27//      model's per-SHARD `(mr_base, rkey)` (shards play the role experts' layer
28//      files do). The client replies with its QP params, the server connects
29//      and idles — the client pulls all tensor bytes one-sided.
30//
31// Per-tensor geometry rides the JSON manifest (like `ExpertIndex`); only the
32// per-shard `(base, rkey)` rides `VerbsServerParams` (like the expert peer's
33// per-layer `(base, rkey)`) — keeping shard counts well under the 4096/8 wire
34// caps and the 512-MR-per-QP shim limit (real models have tens of shards).
35//
36// Module layout (SDD split — the client-facing half is un-gated + verbs-free so LoRA lifts it cleanly):
37//   * `manifest` — the manifest types + address/rail math (un-gated).
38//   * `wire`     — the length-prefixed model-request / manifest codec (un-gated).
39//   * `serve`    — the `atlas-weight-peer` daemon (unix; holds the reg_mr true).
40//   * `shard`    — shard resolution + safetensors parse + the warm RO mmap (unix).
41
42mod manifest;
43#[cfg(unix)]
44mod serve;
45#[cfg(unix)]
46mod shard;
47mod wire;
48
49pub use manifest::{WeightManifest, WeightTensorRecord, rail_for_tensor, tensor_remote_addr};
50#[cfg(unix)]
51pub use serve::{WeightPeerConfig, serve};
52pub use wire::{
53    MODEL_REQUEST_MAX, read_model_request, read_weight_manifest, write_model_request,
54    write_weight_manifest,
55};
spark_storage/weight_peer.rs

spark_storage/
weight_peer.rs
