diff --git a/joulehook/joulehook-rs/Cargo.lock b/joulehook/joulehook-rs/Cargo.lock index f2457f0eb..4fc5b9d59 100644 --- a/joulehook/joulehook-rs/Cargo.lock +++ b/joulehook/joulehook-rs/Cargo.lock @@ -553,7 +553,9 @@ dependencies = [ name = "mesh-engine" version = "0.1.0" dependencies = [ + "blake3", "mesh-schema", + "serde", ] [[package]] diff --git a/mesh/SPEC.md b/mesh/SPEC.md index 7b2b13d14..65ea47a09 100644 --- a/mesh/SPEC.md +++ b/mesh/SPEC.md @@ -338,6 +338,19 @@ decision from a warm one. > fields is how a verifier can tell it apart from the stronger one, rather than > having to assume. +A backend that loads its weights from files addresses them as BLAKE3 over the +canonical encoding of the list of files its forward pass reads, in path order, +each as its path relative to the model directory and the BLAKE3 of its bytes. +Every byte is hashed, once, when the backend loads: a digest over a sample, a +header or a size would look like a binding while leaving the weights free to +change beneath it, and is worse than none. The file names are part of the +address because they are part of what loads, and files that do not reach the +forward pass, such as a tokenizer, are not. The reference engines address +`config.json` with `model.safetensors`, or with `model.safetensors.index.json` +and every shard it names; Gemma reports `fp32`, the arithmetic its loader +widens to, and Falcon-H1 `bf16`. A backend that names its weights before +serving MUST report the same digest after (§2). + --- ## §4 Routing — cheapest-sufficient diff --git a/mesh/mesh-rs/Cargo.lock b/mesh/mesh-rs/Cargo.lock index 76be47197..0e44e22ea 100644 --- a/mesh/mesh-rs/Cargo.lock +++ b/mesh/mesh-rs/Cargo.lock @@ -586,7 +586,9 @@ dependencies = [ name = "mesh-engine" version = "0.1.0" dependencies = [ + "blake3", "mesh-schema", + "serde", ] [[package]] diff --git a/mesh/mesh-rs/Cargo.toml b/mesh/mesh-rs/Cargo.toml index dd021985c..b0719e576 100644 --- a/mesh/mesh-rs/Cargo.toml +++ b/mesh/mesh-rs/Cargo.toml @@ -43,6 +43,7 @@ ed25519-dalek = { version = "2", features = ["rand_core"] } serde = { version = "1", features = ["derive"] } serde_json = "1" thiserror = "2" +blake3 = "1" hex = "0.4" [workspace.lints.clippy] diff --git a/mesh/mesh-rs/crates/mesh-engine-gguf/src/lib.rs b/mesh/mesh-rs/crates/mesh-engine-gguf/src/lib.rs index 63d499e66..f947e0d70 100644 --- a/mesh/mesh-rs/crates/mesh-engine-gguf/src/lib.rs +++ b/mesh/mesh-rs/crates/mesh-engine-gguf/src/lib.rs @@ -35,7 +35,11 @@ mod engine { use jouleclaw_loader_gguf::gemma4::Gemma4; use jouleclaw_loader_gguf::ParseError; use mesh_engine::{Served, ServingEngine}; - use mesh_schema::{BackendId, Fidelity, ServeOutcome, ServeRequest}; + use mesh_schema::{BackendId, ContentId, Fidelity, ServeOutcome, ServeRequest}; + + /// The arithmetic this engine's forward pass runs at: the loader widens + /// every tensor to `f32` and decodes in `f32`. + const PRECISION: &str = "fp32"; /// Errors loading a GGUF-backed engine. /// @@ -46,6 +50,9 @@ mod engine { /// The model directory could not be parsed/loaded. #[error("load model: {0:?}")] Load(ParseError), + /// The weights could not be read to address them. + #[error("address the weights: {0}")] + Digest(String), } /// A Mesh backend backed by a GGUF model decoded with the family's pure-Rust @@ -53,6 +60,9 @@ mod engine { /// feature). pub struct GgufEngine { model: Gemma4, + /// The content address of the weights loaded (SPEC §3.6), taken once, + /// at load, over every file the forward pass reads. + digest: ContentId, id: BackendId, recipe: String, fidelity: Fidelity, @@ -75,9 +85,12 @@ mod engine { fidelity: Fidelity, uj_per_token: u64, ) -> Result { + let dir = dir.as_ref(); let model = Gemma4::load(dir).map_err(GgufError::Load)?; + let digest = gemma_weights_digest(dir).map_err(|e| GgufError::Digest(e.to_string()))?; Ok(Self { model, + digest, id: BackendId::new(id), recipe: recipe.into(), fidelity, @@ -86,11 +99,35 @@ mod engine { } } + /// The files the Gemma loader reads: `config.json`, and either one + /// `model.safetensors` or an index together with every shard it names. + pub(crate) fn gemma_weights_digest(dir: &Path) -> std::io::Result { + let index = "model.safetensors.index.json"; + if !dir.join(index).exists() { + return mesh_engine::weights_digest(dir, &["config.json", "model.safetensors"]); + } + let parsed: serde_json::Value = serde_json::from_slice(&std::fs::read(dir.join(index))?) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?; + let shards: Vec = parsed["weight_map"] + .as_object() + .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "index has no weight_map"))? + .values() + .filter_map(|v| v.as_str().map(str::to_owned)) + .collect(); + let mut files = vec!["config.json", index]; + files.extend(shards.iter().map(String::as_str)); + mesh_engine::weights_digest(dir, &files) + } + impl ServingEngine for GgufEngine { fn backend_id(&self) -> BackendId { self.id.clone() } + fn weights_digest(&self) -> Option { + Some(self.digest) + } + fn recipe_id(&self) -> String { self.recipe.clone() } @@ -123,13 +160,10 @@ mod engine { // makes no claim: the gate falls back to the advertised // rating, which is where these engines already stood. delivered_fidelity: None, - // The weights are a real file, so a digest is obtainable — - // but once, at load, never per request. Until this engine does - // that it reports absent rather than something cheaper that - // would look like a binding: the field's presence is a claim to - // re-executability, and a partial digest is worse than none. - weights_digest: None, - precision: None, + // Addressed once, at load, over every byte the forward pass + // reads: the field's presence is a claim to re-executability. + weights_digest: Some(self.digest), + precision: Some(PRECISION.into()), }; Served { tokens, outcome } } @@ -146,7 +180,11 @@ mod falcon { use jouleclaw_loader_gguf::falcon_h1::{FalconH1Config, FalconH1Model}; use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file; use mesh_engine::{Served, ServingEngine}; - use mesh_schema::{BackendId, Fidelity, ServeOutcome, ServeRequest}; + use mesh_schema::{BackendId, ContentId, Fidelity, ServeOutcome, ServeRequest}; + + /// The arithmetic this engine's forward pass runs at: the model binds only + /// `bf16` weights and multiplies them as `bf16`. + const PRECISION: &str = "bf16"; /// Errors loading a Falcon-H1 backend. #[derive(Debug, thiserror::Error)] @@ -159,6 +197,8 @@ mod falcon { Bind(String), #[error("load tokenizer: {0}")] Tokenizer(String), + #[error("address the weights: {0}")] + Digest(String), } /// A Mesh backend backed by a real Falcon-H1 checkpoint plus its tokenizer. @@ -166,6 +206,10 @@ mod falcon { pub struct FalconEngine { model: FalconH1Model, tokenizer: BpeTokenizer, + /// The content address of the weights loaded (SPEC §3.6): `config.json` + /// and `model.safetensors`, the files the forward pass reads. The + /// tokenizer maps text to the ids a caller sends and is not among them. + digest: ContentId, id: BackendId, recipe: String, fidelity: Fidelity, @@ -193,9 +237,12 @@ mod falcon { .map_err(|e| FalconError::Bind(format!("{e:?}")))?; let tokenizer = BpeTokenizer::from_file(dir.join("tokenizer.json")) .map_err(|e| FalconError::Tokenizer(format!("{e:?}")))?; + let digest = mesh_engine::weights_digest(dir, &["config.json", "model.safetensors"]) + .map_err(|e| FalconError::Digest(e.to_string()))?; Ok(Self { model, tokenizer, + digest, id: BackendId::new(id), recipe: recipe.into(), fidelity, @@ -232,6 +279,9 @@ mod falcon { fn backend_id(&self) -> BackendId { self.id.clone() } + fn weights_digest(&self) -> Option { + Some(self.digest) + } fn recipe_id(&self) -> String { self.recipe.clone() } @@ -268,13 +318,10 @@ mod falcon { // makes no claim: the gate falls back to the advertised // rating, which is where these engines already stood. delivered_fidelity: None, - // The weights are a real file, so a digest is obtainable — - // but once, at load, never per request. Until this engine does - // that it reports absent rather than something cheaper that - // would look like a binding: the field's presence is a claim to - // re-executability, and a partial digest is worse than none. - weights_digest: None, - precision: None, + // Addressed once, at load, over every byte the forward pass + // reads: the field's presence is a claim to re-executability. + weights_digest: Some(self.digest), + precision: Some(PRECISION.into()), }; Served { tokens, outcome } } @@ -314,5 +361,71 @@ mod tests { let served = e.serve(&req, &[2u32, 3u32]); assert!(served.outcome.tokens_out > 0); assert_eq!(served.outcome.tokens_accepted, served.outcome.tokens_out); + // The weights are addressed before serving and reported after, the same. + assert!(e.weights_digest().is_some()); + assert_eq!(served.outcome.weights_digest, e.weights_digest()); + assert_eq!(served.outcome.precision.as_deref(), Some("fp32")); + } + + /// A scratch model directory, removed when dropped. + struct ModelDir(std::path::PathBuf); + impl ModelDir { + fn with(tag: &str, files: &[(&str, &[u8])]) -> Self { + let dir = std::env::temp_dir().join(format!("mesh-gguf-digest-{}-{tag}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + for (name, bytes) in files { + std::fs::write(dir.join(name), bytes).unwrap(); + } + ModelDir(dir) + } + } + impl Drop for ModelDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + /// The Gemma engine addresses the files its loader reads: one + /// `model.safetensors`, or an index and every shard the index names, and + /// a shard the index names but the directory lacks is an error, not a + /// smaller address. + #[test] + fn the_gemma_engine_addresses_what_its_loader_reads() { + let single = ModelDir::with("single", &[("config.json", b"{}"), ("model.safetensors", b"w")]); + assert_eq!( + engine::gemma_weights_digest(&single.0).unwrap(), + mesh_engine::weights_digest(&single.0, &["config.json", "model.safetensors"]).unwrap() + ); + + let index = br#"{"weight_map":{"a.weight":"m-1.safetensors","b.weight":"m-2.safetensors","c.weight":"m-1.safetensors"}}"#; + let sharded = ModelDir::with( + "sharded", + &[("config.json", b"{}"), ("model.safetensors.index.json", index), ("m-1.safetensors", b"one"), ("m-2.safetensors", b"two")], + ); + let all = ["config.json", "model.safetensors.index.json", "m-1.safetensors", "m-2.safetensors"]; + assert_eq!(engine::gemma_weights_digest(&sharded.0).unwrap(), mesh_engine::weights_digest(&sharded.0, &all).unwrap()); + + std::fs::remove_file(sharded.0.join("m-2.safetensors")).unwrap(); + assert!(engine::gemma_weights_digest(&sharded.0).is_err(), "a named shard is missing"); + } + + /// On real weights: the address covers every shard actually on disk. Runs + /// only when `MESH_WEIGHTS_DIR` names a model directory; it hashes the + /// weights without loading them. + #[test] + #[ignore = "requires a local model dir via MESH_WEIGHTS_DIR"] + fn addresses_every_shard_of_real_weights() { + let dir = std::path::PathBuf::from(std::env::var("MESH_WEIGHTS_DIR").expect("set MESH_WEIGHTS_DIR")); + let mut files: Vec = std::fs::read_dir(&dir) + .unwrap() + .map(|e| e.unwrap().file_name().into_string().unwrap()) + .filter(|n| n.ends_with(".safetensors") || n == "config.json" || n == "model.safetensors.index.json") + .collect(); + files.sort(); + let names: Vec<&str> = files.iter().map(String::as_str).collect(); + let d = engine::gemma_weights_digest(&dir).unwrap(); + assert_eq!(d, mesh_engine::weights_digest(&dir, &names).unwrap(), "every shard on disk, and nothing else"); + println!("{} files, weights digest {}", names.len(), d.to_hex()); } } diff --git a/mesh/mesh-rs/crates/mesh-engine/Cargo.toml b/mesh/mesh-rs/crates/mesh-engine/Cargo.toml index ae4e2b1d3..73bdee0ac 100644 --- a/mesh/mesh-rs/crates/mesh-engine/Cargo.toml +++ b/mesh/mesh-rs/crates/mesh-engine/Cargo.toml @@ -9,6 +9,8 @@ description = "Mesh — the serving-engine interface a backend implements, plus [dependencies] mesh-schema = { workspace = true } +blake3 = { workspace = true } +serde = { workspace = true } [lints] workspace = true diff --git a/mesh/mesh-rs/crates/mesh-engine/src/lib.rs b/mesh/mesh-rs/crates/mesh-engine/src/lib.rs index 457fd30d5..f3dad7df9 100644 --- a/mesh/mesh-rs/crates/mesh-engine/src/lib.rs +++ b/mesh/mesh-rs/crates/mesh-engine/src/lib.rs @@ -15,6 +15,37 @@ use mesh_schema::{BackendId, ContentId, Fidelity, ServeOutcome, ServeRequest}; +/// The content address of a model's weights (SPEC §3.6), for a backend that +/// loads them from files: BLAKE3 over the canonical list of the files its +/// forward pass reads, each named by its path relative to `dir` and addressed +/// by the BLAKE3 of its bytes, in path order. +/// +/// Naming each file keeps the address honest about structure: the same bytes +/// split across differently named shards address differently, as they load +/// differently. Hashing every byte is the point. A digest over a sample, a +/// header or a file size is cheaper and would look like a binding while +/// leaving the weights free to change underneath it, so this reads them all, +/// once, when the backend loads them. +/// +/// `files` are the paths the forward pass reads, relative to `dir`; the order +/// they are listed in does not matter. +pub fn weights_digest(dir: &std::path::Path, files: &[&str]) -> std::io::Result { + #[derive(serde::Serialize)] + struct WeightsManifest<'a> { + files: Vec<(&'a str, ContentId)>, + } + let mut named: Vec<&str> = files.to_vec(); + named.sort_unstable(); + named.dedup(); + let mut manifest = WeightsManifest { files: Vec::with_capacity(named.len()) }; + for name in named { + let mut hasher = blake3::Hasher::new(); + hasher.update_reader(std::fs::File::open(dir.join(name))?)?; + manifest.files.push((name, ContentId(*hasher.finalize().as_bytes()))); + } + ContentId::of(&manifest).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())) +} + /// What an engine returns from one `serve`: the generated token ids and the /// serving counters. #[derive(Clone, Debug, PartialEq, Eq)] @@ -257,6 +288,55 @@ impl ServingEngine for ReferenceEngine { #[cfg(test)] mod tests { use super::*; + + /// A scratch directory holding `files`, removed when dropped. + struct Scratch(std::path::PathBuf); + impl Scratch { + fn with(tag: &str, files: &[(&str, &[u8])]) -> Self { + let dir = std::env::temp_dir().join(format!("mesh-weights-{}-{tag}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + for (name, bytes) in files { + std::fs::write(dir.join(name), bytes).unwrap(); + } + Scratch(dir) + } + } + impl Drop for Scratch { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + /// The address is a function of every named file's bytes and name, and of + /// nothing else: not the order they are listed in. + #[test] + fn a_weights_digest_addresses_every_byte_and_every_name() { + let base = Scratch::with("base", &[("config.json", b"{\"d\":4}"), ("model.safetensors", b"weights")]); + let files = ["config.json", "model.safetensors"]; + let d = weights_digest(&base.0, &files).unwrap(); + assert_eq!(d, weights_digest(&base.0, &["model.safetensors", "config.json"]).unwrap(), "order-free"); + assert_eq!(d, weights_digest(&base.0, &files).unwrap(), "deterministic"); + + let one_byte = Scratch::with("byte", &[("config.json", b"{\"d\":4}"), ("model.safetensors", b"weightz")]); + assert_ne!(d, weights_digest(&one_byte.0, &files).unwrap(), "one changed byte"); + let renamed = Scratch::with("name", &[("config.json", b"{\"d\":4}"), ("shard.safetensors", b"weights")]); + assert_ne!(d, weights_digest(&renamed.0, &["config.json", "shard.safetensors"]).unwrap(), "a renamed file"); + assert_ne!(d, weights_digest(&base.0, &["model.safetensors"]).unwrap(), "a file left out"); + assert!(weights_digest(&base.0, &["config.json", "absent.safetensors"]).is_err(), "a missing file is an error"); + } + + /// Each file is addressed exactly as `ContentId::of_bytes` would address + /// it, though it is streamed rather than read into memory. + #[test] + fn a_streamed_file_hash_is_its_content_id() { + let big: Vec = (0..300_000u32).map(|i| (i % 251) as u8).collect(); + let dir = Scratch::with("stream", &[("w.bin", &big)]); + let mut hasher = blake3::Hasher::new(); + hasher.update_reader(std::fs::File::open(dir.0.join("w.bin")).unwrap()).unwrap(); + assert_eq!(ContentId(*hasher.finalize().as_bytes()), ContentId::of_bytes(&big)); + } + use mesh_schema::ContentId; fn req(tag: u8, max_tokens: u32) -> ServeRequest {