spark_storage/weight_peer/
manifest.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Weight-staging manifest types + address math (un-gated, CUDA-free, verbs-free).
4//
5// This half is exactly what the RDMA clients import: `WeightManifest`,
6// `WeightTensorRecord`, `rail_for_tensor`, `tensor_remote_addr`. Keeping it free
7// of any back-edge into the server (`serve`/`shard`) is what makes the LoRA
8// carve-out (`weight_lora_rdma`) a clean lift.
9
10use serde::{Deserialize, Serialize};
11
12/// One tensor's placement inside a staged model, mirroring the safetensors
13/// header exactly: `offset_in_shard` is the ABSOLUTE file offset (8-byte size
14/// prefix + header + the tensor's data-section start), `len` is the raw
15/// contiguous byte count (`data_offsets[1] - data_offsets[0]`). The client
16/// RDMA-READs exactly `[shard_base + offset_in_shard .. + len)`.
17#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
18pub struct WeightTensorRecord {
19    /// HuggingFace tensor name — the `WeightStore` key, verbatim.
20    pub name: String,
21    /// Raw safetensors dtype string (`"BF16"`, `"F8_E4M3"`, `"I8"`, …). The
22    /// client maps it via `WeightDtype::from_safetensors_str` — the same closed
23    /// mapping the disk loaders use.
24    pub dtype: String,
25    pub shape: Vec<u64>,
26    /// Absolute byte offset of the tensor's first byte within its shard file.
27    pub offset_in_shard: u64,
28    /// Tensor byte length (authoritative — do NOT recompute from shape; packed
29    /// NVFP4 lengths differ from `numel * byte_size`).
30    pub len: u64,
31    /// Index into [`WeightManifest::shard_files`].
32    pub shard_index: u32,
33    /// True for tensors from `extra_weights.safetensors` (grafted MTP etc.):
34    /// loaded with NO expert-skip filter, exactly like the disk loaders.
35    pub extra: bool,
36}
37
38/// A staged model's manifest: the geometry the client needs to reconstruct a
39/// byte-identical `WeightStore`. Published as length-prefixed JSON right after
40/// the client's model request. The per-shard `(base, rkey)` MR handles ride the
41/// verbs handshake separately (see the module doc).
42#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
43pub struct WeightManifest {
44    pub version: u32,
45    /// The resolved model id/path the peer staged (echo for the client's log).
46    pub model_id: String,
47    /// Shard file names, shard-indexed. `shard_index` in each tensor record
48    /// indexes this list; the verbs `layers` vector is per-shard in this order.
49    pub shard_files: Vec<String>,
50    /// Byte length of each shard file, shard-indexed (parallels `shard_files`).
51    pub shard_lens: Vec<u64>,
52    pub tensors: Vec<WeightTensorRecord>,
53}
54
55impl WeightManifest {
56    pub const VERSION: u32 = 1;
57
58    /// Number of shard files (== the per-rail `layers` MR count the peer
59    /// publishes and the client validates against).
60    pub fn num_shards(&self) -> usize {
61        self.shard_files.len()
62    }
63
64    /// Total registered bytes across all shards (the whole-file MRs) — the
65    /// figure charged against the blade `CommitLedger` once per staged model.
66    pub fn total_shard_bytes(&self) -> u64 {
67        self.shard_lens.iter().sum()
68    }
69}
70
71/// Rail selection for a tensor under dual-rail striping: tensor `N` is served
72/// over rail `N % n_rails`. Factored out (un-gated) so the striping is unit-
73/// testable off the RDMA path — the client's read loop calls this so the tested
74/// logic and the shipped logic are the same. `n_rails` is clamped to `>= 1`.
75pub fn rail_for_tensor(tensor_index: usize, n_rails: usize) -> usize {
76    tensor_index % n_rails.max(1)
77}
78
79/// Absolute peer virtual address of a tensor's first byte: the shard's whole-
80/// file REMOTE_READ MR base plus the tensor's ABSOLUTE in-shard offset (the
81/// safetensors data-section offset, which already includes the 8-byte size
82/// prefix + header). The client RDMA-READs `[addr .. addr + len)`. Factored out
83/// (un-gated) so the address math is unit-testable off the RDMA path.
84pub fn tensor_remote_addr(shard_base: u64, offset_in_shard: u64) -> u64 {
85    shard_base + offset_in_shard
86}
87
88#[cfg(test)]
89#[path = "manifest_tests.rs"]
90mod tests;
spark_storage/weight_peer/manifest.rs

spark_storage/weight_peer/
manifest.rs
