spark_storage/weight_peer/
wire.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Weight-staging wire codec (un-gated, CUDA-free, verbs-free).
4//
5// The length-prefixed framing the daemon and its clients speak: the model
6// request `[u32 len][bytes]` and the manifest `[u32 len][JSON]`. Kept beside
7// `manifest` (which it references) and free of any server dependency so the
8// LoRA client (`weight_lora_rdma`) imports only this + `manifest`.
9
10use anyhow::{Context, Result, bail};
11
12use super::manifest::WeightManifest;
13
14fn read_u32<R: std::io::Read>(r: &mut R) -> Result<u32> {
15    let mut b = [0u8; 4];
16    r.read_exact(&mut b).context("read u32")?;
17    Ok(u32::from_le_bytes(b))
18}
19
20/// Longest model id/path the wire accepts (a corrupt/hostile length must not
21/// trigger a huge allocation).
22pub const MODEL_REQUEST_MAX: usize = 8192;
23
24/// Wire form of the model request: `[u32 len][len bytes UTF-8 id/path]`.
25pub fn write_model_request<W: std::io::Write>(w: &mut W, id: &str) -> Result<()> {
26    let bytes = id.as_bytes();
27    if bytes.is_empty() || bytes.len() > MODEL_REQUEST_MAX {
28        bail!("implausible model request length: {}", bytes.len());
29    }
30    w.write_all(&(bytes.len() as u32).to_le_bytes())?;
31    w.write_all(bytes)?;
32    Ok(())
33}
34
35/// Read a model request written by [`write_model_request`].
36pub fn read_model_request<R: std::io::Read>(r: &mut R) -> Result<String> {
37    let len = read_u32(r)? as usize;
38    if len == 0 || len > MODEL_REQUEST_MAX {
39        bail!("implausible model request length: {len}");
40    }
41    let mut buf = vec![0u8; len];
42    r.read_exact(&mut buf).context("read model request body")?;
43    String::from_utf8(buf).context("model request is not valid UTF-8")
44}
45
46/// Serialize + frame a manifest as `[u32 len][len bytes JSON]`.
47pub fn write_weight_manifest<W: std::io::Write>(w: &mut W, m: &WeightManifest) -> Result<()> {
48    let json = serde_json::to_vec(m).context("serialize weight manifest")?;
49    w.write_all(&(json.len() as u32).to_le_bytes())?;
50    w.write_all(&json)?;
51    Ok(())
52}
53
54/// Read + parse a length-prefixed manifest. Shared by the client tier.
55pub fn read_weight_manifest<R: std::io::Read>(r: &mut R) -> Result<WeightManifest> {
56    let len = read_u32(r)? as usize;
57    if len == 0 || len > 256 * 1024 * 1024 {
58        bail!("implausible weight manifest length: {len}");
59    }
60    let mut buf = vec![0u8; len];
61    r.read_exact(&mut buf)
62        .context("read weight manifest json")?;
63    let m: WeightManifest = serde_json::from_slice(&buf).context("parse weight manifest json")?;
64    if m.version != WeightManifest::VERSION {
65        bail!(
66            "weight manifest version {} != supported {}",
67            m.version,
68            WeightManifest::VERSION
69        );
70    }
71    if m.shard_files.len() != m.shard_lens.len() {
72        bail!(
73            "manifest shard_files ({}) / shard_lens ({}) length mismatch",
74            m.shard_files.len(),
75            m.shard_lens.len()
76        );
77    }
78    Ok(m)
79}
80
81#[cfg(test)]
82#[path = "wire_tests.rs"]
83mod tests;
spark_storage/weight_peer/wire.rs

spark_storage/weight_peer/
wire.rs
