metrale_storage/weight_peer/
manifest.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The manifest a weight peer publishes, and the address and rail
4//! arithmetic its clients use (model-weights `weight_tier_rdma.rs` and
5//! `weight_lora_rdma.rs`). Compiled on every platform and independent of
6//! `serve` and `shard`.
7//!
8//! Owner: storage (weight peer).
9//! Invariants: none beyond the types.
10
11use serde::{Deserialize, Serialize};
12
13/// 2026-09-25: Where one tensor of a staged model lies in its shard file. The
14/// client reads `[shard_base + offset_in_shard, + len)`.
15#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
16pub struct WeightTensorRecord {
17    /// 2026-09-25: The tensor's name in the safetensors header.
18    pub name: String,
19    /// 2026-09-25: The safetensors dtype string (`"BF16"`, `"F8_E4M3"`, ...),
20    /// which the base-weight client maps with `WeightDtype::from_safetensors_str`.
21    pub dtype: String,
22    pub shape: Vec<u64>,
23    /// 2026-09-25: File offset of the tensor's first byte: 8 + header size +
24    /// `data_offsets[0]` (`metrale_core::safetensors::tensor_span`).
25    pub offset_in_shard: u64,
26    /// 2026-09-25: Byte length, `data_offsets[1] - data_offsets[0]`.
27    pub len: u64,
28    /// 2026-09-25: Index into [`WeightManifest::shard_files`].
29    pub shard_index: u32,
30    /// 2026-09-25: True for tensors of `extra_weights.safetensors`, which the
31    /// client never drops with its expert-skip filter (`should_skip_tensor` in
32    /// `weight_tier_rdma.rs`).
33    pub extra: bool,
34}
35
36/// 2026-09-25: A staged model's manifest, sent as length-prefixed JSON after
37/// the client's model request (`write_weight_manifest`). The per-shard
38/// `(base, rkey)` pairs are sent later, in the verbs handshake.
39#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
40pub struct WeightManifest {
41    pub version: u32,
42    /// 2026-09-25: The directory the peer staged, as a string.
43    pub model_id: String,
44    /// 2026-09-25: Shard file names in shard order; the verbs `layers` vector
45    /// lists the shard MRs in the same order.
46    pub shard_files: Vec<String>,
47    /// 2026-09-25: Byte length of each shard file, in shard order.
48    pub shard_lens: Vec<u64>,
49    pub tensors: Vec<WeightTensorRecord>,
50}
51
52impl WeightManifest {
53    pub const VERSION: u32 = 1;
54
55    /// 2026-09-25: The number of shard files, which the client checks against
56    /// the number of MRs the peer publishes per rail.
57    pub fn num_shards(&self) -> usize {
58        self.shard_files.len()
59    }
60
61    /// 2026-09-25: The sum of `shard_lens`, which `serve` reserves on the
62    /// `CommitLedger` once per staged model.
63    pub fn total_shard_bytes(&self) -> u64 {
64        self.shard_lens.iter().sum()
65    }
66}
67
68/// 2026-09-25: The rail for tensor `tensor_index`: `tensor_index % n_rails`,
69/// with `n_rails` taken as at least 1. The client's read loop calls it.
70pub fn rail_for_tensor(tensor_index: usize, n_rails: usize) -> usize {
71    tensor_index % n_rails.max(1)
72}
73
74/// 2026-09-25: The peer address of a tensor's first byte: the shard MR's base
75/// plus `offset_in_shard`, which already counts the size prefix and header.
76pub fn tensor_remote_addr(shard_base: u64, offset_in_shard: u64) -> u64 {
77    shard_base + offset_in_shard
78}
79
80#[cfg(test)]
81#[path = "manifest_tests.rs"]
82mod tests;
metrale_storage/weight_peer/manifest.rs

metrale_storage/weight_peer/
manifest.rs
