metrale_storage/weight_peer/
wire.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The weight peer's length-prefixed framing: the model request
4//! `[u32 len][len bytes UTF-8]` and the manifest `[u32 len][len bytes JSON]`,
5//! lengths little-endian. 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 anyhow::{Context, Result, bail};
12
13use super::manifest::WeightManifest;
14
15fn read_u32<R: std::io::Read>(r: &mut R) -> Result<u32> {
16    let mut b = [0u8; 4];
17    r.read_exact(&mut b).context("read u32")?;
18    Ok(u32::from_le_bytes(b))
19}
20
21/// 2026-09-25: The longest model request, in bytes, either side accepts; a
22/// longer length prefix is refused before any buffer is allocated.
23pub const MODEL_REQUEST_MAX: usize = 8192;
24
25/// 2026-09-25: Write the model request `id`.
26///
27/// # Errors
28/// `id` is empty or longer than [`MODEL_REQUEST_MAX`], or the write fails.
29pub fn write_model_request<W: std::io::Write>(w: &mut W, id: &str) -> Result<()> {
30    let bytes = id.as_bytes();
31    if bytes.is_empty() || bytes.len() > MODEL_REQUEST_MAX {
32        bail!("implausible model request length: {}", bytes.len());
33    }
34    w.write_all(&(bytes.len() as u32).to_le_bytes())?;
35    w.write_all(bytes)?;
36    Ok(())
37}
38
39/// 2026-09-25: Read a model request written by [`write_model_request`].
40///
41/// # Errors
42/// A length of 0 or above [`MODEL_REQUEST_MAX`], a short read, or a body that
43/// is not UTF-8.
44pub fn read_model_request<R: std::io::Read>(r: &mut R) -> Result<String> {
45    let len = read_u32(r)? as usize;
46    if len == 0 || len > MODEL_REQUEST_MAX {
47        bail!("implausible model request length: {len}");
48    }
49    let mut buf = vec![0u8; len];
50    r.read_exact(&mut buf).context("read model request body")?;
51    String::from_utf8(buf).context("model request is not valid UTF-8")
52}
53
54/// 2026-09-25: Serialize `m` to JSON and write it with its length prefix.
55pub fn write_weight_manifest<W: std::io::Write>(w: &mut W, m: &WeightManifest) -> Result<()> {
56    let json = serde_json::to_vec(m).context("serialize weight manifest")?;
57    w.write_all(&(json.len() as u32).to_le_bytes())?;
58    w.write_all(&json)?;
59    Ok(())
60}
61
62/// 2026-09-25: Read and parse a manifest written by [`write_weight_manifest`].
63///
64/// # Errors
65/// A length of 0 or above 256 MiB, a short read, bad JSON, a `version` other
66/// than [`WeightManifest::VERSION`], or `shard_files` and `shard_lens` of
67/// different lengths.
68pub fn read_weight_manifest<R: std::io::Read>(r: &mut R) -> Result<WeightManifest> {
69    let len = read_u32(r)? as usize;
70    if len == 0 || len > 256 * 1024 * 1024 {
71        bail!("implausible weight manifest length: {len}");
72    }
73    let mut buf = vec![0u8; len];
74    r.read_exact(&mut buf)
75        .context("read weight manifest json")?;
76    let m: WeightManifest = serde_json::from_slice(&buf).context("parse weight manifest json")?;
77    if m.version != WeightManifest::VERSION {
78        bail!(
79            "weight manifest version {} != supported {}",
80            m.version,
81            WeightManifest::VERSION
82        );
83    }
84    if m.shard_files.len() != m.shard_lens.len() {
85        bail!(
86            "manifest shard_files ({}) / shard_lens ({}) length mismatch",
87            m.shard_files.len(),
88            m.shard_lens.len()
89        );
90    }
91    Ok(m)
92}
93
94#[cfg(test)]
95#[path = "wire_tests.rs"]
96mod tests;
metrale_storage/weight_peer/wire.rs

metrale_storage/weight_peer/
wire.rs
