Corpus semantic embedding + prove-local weight-hash cache (+ S5 on real KVM, verified) #34
6 changed files with 581 additions and 12 deletions
3
jouleclaw/jouleclaw-rs/Cargo.lock
generated
3
jouleclaw/jouleclaw-rs/Cargo.lock
generated
|
|
@ -3393,13 +3393,16 @@ dependencies = [
|
|||
name = "jouleclaw-corpus"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"jouleclaw-embed-candle",
|
||||
"jouleclaw-energy",
|
||||
"jouleclaw-hf-hub",
|
||||
"jouleclaw-intent",
|
||||
"jouleclaw-prov",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -16,5 +16,22 @@ serde = { workspace = true }
|
|||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
# Real semantic embedder, opt-in. Pulls candle + tokenizers + hf-hub, so it is
|
||||
# behind a feature: the default corpus build stays light and offline. See
|
||||
# `CandleTextEmbedder`. hf-hub + tokio are needed only by the semantic example
|
||||
# (loading the model is async and hits the hub); they ride the same feature so
|
||||
# the default build takes none of them.
|
||||
jouleclaw-embed-candle = { path = "../jouleclaw-embed-candle", optional = true }
|
||||
jouleclaw-hf-hub = { path = "../jouleclaw-hf-hub", optional = true }
|
||||
tokio = { workspace = true, optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
candle = ["dep:jouleclaw-embed-candle", "dep:jouleclaw-hf-hub", "dep:tokio"]
|
||||
|
||||
[[example]]
|
||||
name = "emit_corpus_semantic"
|
||||
required-features = ["candle"]
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -0,0 +1,82 @@
|
|||
//! Emit a **semantically-embedded** corpus from `(receipt, query_text)` pairs.
|
||||
//!
|
||||
//! ```text
|
||||
//! cargo run -p jouleclaw-corpus --features candle --example emit_corpus_semantic -- /tmp/corpus.jsonl
|
||||
//! ```
|
||||
//!
|
||||
//! Unlike `emit_corpus` (which embeds receipt *structure* only), this wires a
|
||||
//! real BERT/MiniLM embedder through [`jouleclaw_corpus::CandleTextEmbedder`],
|
||||
//! so each row carries what the query *meant* — `[semantic ‖ structural]`.
|
||||
//!
|
||||
//! Loading the model hits the HF hub on first run (hf-hub-cached after), which
|
||||
//! is why this example is behind the `candle` feature and asks for the network
|
||||
//! once. The `build_with_text` join — pairing each receipt with the query that
|
||||
//! produced it — is exactly what a real deployment does from its own logs; here
|
||||
//! we stand in a small labelled ledger.
|
||||
|
||||
use jouleclaw_corpus::{CandleTextEmbedder, CorpusBuilder, TextEmbedder, write_jsonl};
|
||||
use jouleclaw_embed_candle::{CandleEmbedder, EmbedderConfig};
|
||||
use jouleclaw_energy::Provenance;
|
||||
use jouleclaw_hf_hub::HfHubClient;
|
||||
use jouleclaw_intent::{IntentVerdict, QueryIntent};
|
||||
use jouleclaw_prov::{CascadeTier, Receipt};
|
||||
|
||||
fn receipt(tier: CascadeTier, uj: u64, intent: QueryIntent) -> Receipt {
|
||||
Receipt {
|
||||
context: None,
|
||||
jc_receipt: "1".into(),
|
||||
id: format!("rcpt-{}-{uj}", tier.wire_tag()),
|
||||
closed_at: "2026-07-20T00:00:00Z".into(),
|
||||
input_hash: "a".repeat(64),
|
||||
tier,
|
||||
joules_uj: uj,
|
||||
energy_provenance: Provenance::HwShunt,
|
||||
tools_touched: Vec::new(),
|
||||
claims: Vec::new(),
|
||||
eoc_stage: None,
|
||||
citations: Vec::new(),
|
||||
tool_receipts: Vec::new(),
|
||||
token_attributions: Vec::new(),
|
||||
model_id: None,
|
||||
intent: Some(IntentVerdict::new(intent, 0.9, "heuristic:v1")),
|
||||
steward: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let out = std::env::args().nth(1).unwrap_or_else(|| "corpus-semantic.jsonl".into());
|
||||
|
||||
// The receipt/query join a real deployment reconstructs from its own logs.
|
||||
// Cheap comprehension closes low; open-ended generation escalates to a model.
|
||||
let ledger = vec![
|
||||
(receipt(CascadeTier::L0Cache, 11, QueryIntent::Comprehension), "what is my current balance?"),
|
||||
(receipt(CascadeTier::L2Embed, 4_050, QueryIntent::Comprehension), "find the closing docs for 123 Oak St"),
|
||||
(receipt(CascadeTier::L3Model, 2_100_000, QueryIntent::Generation), "draft a settlement summary for the buyer"),
|
||||
(receipt(CascadeTier::L4Wire, 60_000_000, QueryIntent::Generation), "write a full risk memo on this portfolio"),
|
||||
(receipt(CascadeTier::L1Lawful, 40, QueryIntent::Process), "convert 3.5 acres to square feet"),
|
||||
];
|
||||
|
||||
eprintln!("loading MiniLM (first run pulls from the hub, then hf-hub-cached)…");
|
||||
let client = HfHubClient::new()?;
|
||||
let inner = CandleEmbedder::load(EmbedderConfig::minilm(), &client).await?;
|
||||
let embedder = CandleTextEmbedder::new(inner)?;
|
||||
eprintln!("embedder ready — semantic dim {}", embedder.dim());
|
||||
|
||||
let items: Vec<(Receipt, String)> =
|
||||
ledger.into_iter().map(|(r, q)| (r, q.to_string())).collect();
|
||||
|
||||
let (rows, stats) = CorpusBuilder::default()
|
||||
.with_monotone_negatives(true)
|
||||
.build_with_text(&items, &embedder);
|
||||
|
||||
write_jsonl(&out, &rows)?;
|
||||
|
||||
println!("receipts seen {}", stats.receipts_seen);
|
||||
println!("receipts used {}", stats.receipts_used);
|
||||
println!("skipped unmapped {}", stats.skipped_unmapped);
|
||||
println!("rows total {}", stats.rows());
|
||||
println!("row width {} (semantic {} + structural)", rows[0].embedding.len(), embedder.dim());
|
||||
println!("wrote {out}");
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -180,6 +180,27 @@ pub trait Embedder {
|
|||
fn dim(&self) -> usize;
|
||||
}
|
||||
|
||||
/// Embeds query *text* into a semantic vector.
|
||||
///
|
||||
/// This is the semantic counterpart to [`Embedder`], and the split is not
|
||||
/// cosmetic: a receipt records a BLAKE3 of its input, never the input itself,
|
||||
/// so an embedder that captures what a query *means* cannot read it off the
|
||||
/// receipt. It needs the text — which only the deployment that logged both the
|
||||
/// receipts and the queries can supply. The [`CorpusBuilder::build_with_text`]
|
||||
/// path takes that `(receipt, text)` join explicitly; the join is the caller's
|
||||
/// responsibility and is not silently verifiable here (the receipt hash is over
|
||||
/// a normalisation this crate does not know).
|
||||
///
|
||||
/// Implement it over any real model. The bundled [`CandleTextEmbedder`] (behind
|
||||
/// the `candle` feature) wraps the fleet's local BERT/BGE embedder.
|
||||
pub trait TextEmbedder {
|
||||
/// Semantic vector for `text`. Must return [`TextEmbedder::dim`] elements
|
||||
/// for every input, so a whole corpus embeds to a consistent width.
|
||||
fn embed_text(&self, text: &str) -> Vec<f32>;
|
||||
/// Dimension of every vector this embedder returns.
|
||||
fn dim(&self) -> usize;
|
||||
}
|
||||
|
||||
/// Number of dimensions [`StructuralEmbedder`] emits.
|
||||
pub const STRUCTURAL_DIM: usize = 15;
|
||||
|
||||
|
|
@ -294,6 +315,7 @@ impl CorpusStats {
|
|||
pub struct CorpusBuilder<E: Embedder> {
|
||||
embedder: E,
|
||||
monotone_negatives: bool,
|
||||
append_structural: bool,
|
||||
}
|
||||
|
||||
impl Default for CorpusBuilder<StructuralEmbedder> {
|
||||
|
|
@ -308,6 +330,7 @@ impl<E: Embedder> CorpusBuilder<E> {
|
|||
Self {
|
||||
embedder,
|
||||
monotone_negatives: false,
|
||||
append_structural: true,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -329,6 +352,19 @@ impl<E: Embedder> CorpusBuilder<E> {
|
|||
self
|
||||
}
|
||||
|
||||
/// On the [`build_with_text`](Self::build_with_text) path, concatenate the
|
||||
/// structural features onto the semantic vector (default: on).
|
||||
///
|
||||
/// Meaning and shape are complementary signals — *what* a query asks and
|
||||
/// *how* it was served — so by default a text row carries both:
|
||||
/// `[semantic ‖ structural]`, width `text.dim() + STRUCTURAL_DIM`. Set this
|
||||
/// off for a pure-semantic vector of width `text.dim()`. Either way the
|
||||
/// width is constant across the corpus, which is what the trainer requires.
|
||||
pub fn with_append_structural(mut self, yes: bool) -> Self {
|
||||
self.append_structural = yes;
|
||||
self
|
||||
}
|
||||
|
||||
/// Resolve the stage a receipt closed at, preferring the runtime's own
|
||||
/// label over the tier mapping.
|
||||
fn stage_of(receipt: &Receipt) -> Option<CorpusStage> {
|
||||
|
|
@ -339,19 +375,16 @@ impl<E: Embedder> CorpusBuilder<E> {
|
|||
.or_else(|| CorpusStage::from_tier(receipt.tier))
|
||||
}
|
||||
|
||||
/// Project one receipt. Returns an empty vec when it cannot be mapped.
|
||||
pub fn rows_from_receipt(&self, receipt: &Receipt) -> Vec<CorpusRow> {
|
||||
let Some(stage) = Self::stage_of(receipt) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let embedding = self.embedder.embed(receipt);
|
||||
/// The positive row plus any monotone-escalation negatives, given an
|
||||
/// already-computed feature vector. Shared by the structural and text
|
||||
/// paths — only the vector differs between them.
|
||||
fn assemble_rows(&self, stage: CorpusStage, embedding: Vec<f32>, cost: u64) -> Vec<CorpusRow> {
|
||||
let mut rows = vec![CorpusRow {
|
||||
embedding: embedding.clone(),
|
||||
stage,
|
||||
success: true,
|
||||
cost_microjoules: receipt.joules_uj,
|
||||
cost_microjoules: cost,
|
||||
}];
|
||||
|
||||
if self.monotone_negatives {
|
||||
for cheaper in [CorpusStage::Cache, CorpusStage::Kv] {
|
||||
if cheaper.rank() < stage.rank() {
|
||||
|
|
@ -359,7 +392,7 @@ impl<E: Embedder> CorpusBuilder<E> {
|
|||
embedding: embedding.clone(),
|
||||
stage: cheaper,
|
||||
success: false,
|
||||
cost_microjoules: receipt.joules_uj,
|
||||
cost_microjoules: cost,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -367,6 +400,38 @@ impl<E: Embedder> CorpusBuilder<E> {
|
|||
rows
|
||||
}
|
||||
|
||||
/// Project one receipt from its structure alone. Empty when unmappable.
|
||||
pub fn rows_from_receipt(&self, receipt: &Receipt) -> Vec<CorpusRow> {
|
||||
let Some(stage) = Self::stage_of(receipt) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let embedding = self.embedder.embed(receipt);
|
||||
self.assemble_rows(stage, embedding, receipt.joules_uj)
|
||||
}
|
||||
|
||||
/// Project one receipt joined to the query text that produced it, embedding
|
||||
/// the text semantically. Empty when the receipt is unmappable.
|
||||
///
|
||||
/// The row vector is `text_embedder.embed_text(text)`, with the structural
|
||||
/// features appended unless [`with_append_structural`](Self::with_append_structural)
|
||||
/// was turned off. This is where a corpus stops encoding only the *shape* of
|
||||
/// a request and starts encoding what it *asked*.
|
||||
pub fn rows_from_receipt_with_text<T: TextEmbedder>(
|
||||
&self,
|
||||
receipt: &Receipt,
|
||||
text: &str,
|
||||
text_embedder: &T,
|
||||
) -> Vec<CorpusRow> {
|
||||
let Some(stage) = Self::stage_of(receipt) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut embedding = text_embedder.embed_text(text);
|
||||
if self.append_structural {
|
||||
embedding.extend(self.embedder.embed(receipt));
|
||||
}
|
||||
self.assemble_rows(stage, embedding, receipt.joules_uj)
|
||||
}
|
||||
|
||||
/// Project a batch of receipts, reporting what was used and skipped.
|
||||
pub fn build(&self, receipts: &[Receipt]) -> (Vec<CorpusRow>, CorpusStats) {
|
||||
let mut out = Vec::new();
|
||||
|
|
@ -392,6 +457,98 @@ impl<E: Embedder> CorpusBuilder<E> {
|
|||
}
|
||||
(out, stats)
|
||||
}
|
||||
|
||||
/// Project `(receipt, query_text)` pairs into a semantically-embedded
|
||||
/// corpus, reporting what was used and skipped.
|
||||
///
|
||||
/// The join is the caller's: only the deployment that logged both the
|
||||
/// receipts and the queries can pair them, because the receipt carries a
|
||||
/// hash of its input, not the input. Rows carry `text_embedder.embed_text`
|
||||
/// of the query (plus structural features unless disabled) — meaning, not
|
||||
/// just shape. `skipped_unmapped` still counts receipts whose closing tier
|
||||
/// has no EOC analogue, exactly as on the structural path.
|
||||
pub fn build_with_text<T: TextEmbedder>(
|
||||
&self,
|
||||
items: &[(Receipt, String)],
|
||||
text_embedder: &T,
|
||||
) -> (Vec<CorpusRow>, CorpusStats) {
|
||||
let mut out = Vec::new();
|
||||
let mut stats = CorpusStats {
|
||||
receipts_seen: items.len(),
|
||||
..Default::default()
|
||||
};
|
||||
for (r, text) in items {
|
||||
let rows = self.rows_from_receipt_with_text(r, text, text_embedder);
|
||||
if rows.is_empty() {
|
||||
stats.skipped_unmapped += 1;
|
||||
continue;
|
||||
}
|
||||
stats.receipts_used += 1;
|
||||
for row in &rows {
|
||||
if row.success {
|
||||
stats.rows_positive += 1;
|
||||
} else {
|
||||
stats.rows_negative += 1;
|
||||
}
|
||||
}
|
||||
out.extend(rows);
|
||||
}
|
||||
(out, stats)
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// Real semantic embedder (candle-backed, opt-in)
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
/// A [`TextEmbedder`] backed by the fleet's local candle BERT/BGE embedder
|
||||
/// ([`jouleclaw_embed_candle::CandleEmbedder`]).
|
||||
///
|
||||
/// Behind the `candle` feature so the default corpus build takes no candle /
|
||||
/// tokenizers / hf-hub dependency. Construction is the caller's: load the
|
||||
/// [`jouleclaw_embed_candle::CandleEmbedder`] (its `load` is async and hits the
|
||||
/// hub on first call, hf-hub-cached after) and hand it here. The dimension is
|
||||
/// probed once at construction so [`TextEmbedder::dim`] is exact.
|
||||
///
|
||||
/// The underlying embedder runs on-device (`Device::Cpu`), no API key — the
|
||||
/// same L2-tier embedding backend the cascade already uses, now feeding the
|
||||
/// corpus that will train the router above it.
|
||||
#[cfg(feature = "candle")]
|
||||
pub struct CandleTextEmbedder {
|
||||
inner: jouleclaw_embed_candle::CandleEmbedder,
|
||||
dim: usize,
|
||||
}
|
||||
|
||||
#[cfg(feature = "candle")]
|
||||
impl CandleTextEmbedder {
|
||||
/// Wrap a loaded [`jouleclaw_embed_candle::CandleEmbedder`], probing the
|
||||
/// output dimension once so `dim()` never lies.
|
||||
pub fn new(
|
||||
inner: jouleclaw_embed_candle::CandleEmbedder,
|
||||
) -> std::result::Result<Self, jouleclaw_embed_candle::EmbedError> {
|
||||
// A non-empty probe: an empty string can pool degenerately in some
|
||||
// models, and all we need is the width.
|
||||
let dim = inner.embed("dimension probe")?.len();
|
||||
Ok(Self { inner, dim })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "candle")]
|
||||
impl TextEmbedder for CandleTextEmbedder {
|
||||
fn embed_text(&self, text: &str) -> Vec<f32> {
|
||||
// The trait is infallible by contract (a corpus row must have a
|
||||
// vector); a tokenizer/model error yields a zero vector of the probed
|
||||
// width rather than a panic or a ragged row. Callers that need to
|
||||
// distinguish an embed failure from a genuinely zero embedding should
|
||||
// use the inner embedder's `embed` directly.
|
||||
self.inner
|
||||
.embed(text)
|
||||
.unwrap_or_else(|_| vec![0.0; self.dim])
|
||||
}
|
||||
|
||||
fn dim(&self) -> usize {
|
||||
self.dim
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
|
@ -604,4 +761,100 @@ mod tests {
|
|||
serde_json::from_str::<CorpusRow>(line).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// ── semantic (text) path ─────────────────────────────────────────
|
||||
|
||||
/// A deterministic stand-in for a real [`TextEmbedder`]. Real models are
|
||||
/// exercised by the `emit_corpus_semantic` example (network-gated); the
|
||||
/// production logic here must be testable offline, so this maps text to a
|
||||
/// fixed-width vector by a stable, content-sensitive rule.
|
||||
struct MockTextEmbedder {
|
||||
dim: usize,
|
||||
}
|
||||
impl TextEmbedder for MockTextEmbedder {
|
||||
fn embed_text(&self, text: &str) -> Vec<f32> {
|
||||
let mut v = vec![0.0_f32; self.dim];
|
||||
for (i, b) in text.bytes().enumerate() {
|
||||
v[i % self.dim] += b as f32 / 255.0;
|
||||
}
|
||||
v
|
||||
}
|
||||
fn dim(&self) -> usize {
|
||||
self.dim
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_row_is_semantic_then_structural_by_default() {
|
||||
let te = MockTextEmbedder { dim: 8 };
|
||||
let r = receipt(CascadeTier::L3Model, 900);
|
||||
let rows = CorpusBuilder::default().rows_from_receipt_with_text(&r, "how do I close a trade?", &te);
|
||||
assert_eq!(rows.len(), 1);
|
||||
// Composite width = semantic ‖ structural.
|
||||
assert_eq!(rows[0].embedding.len(), te.dim() + STRUCTURAL_DIM);
|
||||
assert_eq!(rows[0].stage, CorpusStage::Neural);
|
||||
assert_eq!(rows[0].cost_microjoules, 900);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_row_is_pure_semantic_when_structural_disabled() {
|
||||
let te = MockTextEmbedder { dim: 8 };
|
||||
let r = receipt(CascadeTier::L2Embed, 4000);
|
||||
let rows = CorpusBuilder::default()
|
||||
.with_append_structural(false)
|
||||
.rows_from_receipt_with_text(&r, "what is the tax basis?", &te);
|
||||
assert_eq!(rows[0].embedding.len(), te.dim());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_embedding_actually_reflects_the_query() {
|
||||
// The whole point of the semantic path: different questions embed
|
||||
// differently, where StructuralEmbedder would collapse them.
|
||||
let te = MockTextEmbedder { dim: 8 };
|
||||
let r = receipt(CascadeTier::L3Model, 5);
|
||||
let a = CorpusBuilder::default()
|
||||
.with_append_structural(false)
|
||||
.rows_from_receipt_with_text(&r, "close the position", &te);
|
||||
let b = CorpusBuilder::default()
|
||||
.with_append_structural(false)
|
||||
.rows_from_receipt_with_text(&r, "audit the vendor spend", &te);
|
||||
assert_ne!(a[0].embedding, b[0].embedding);
|
||||
// And the same query is stable.
|
||||
let a2 = CorpusBuilder::default()
|
||||
.with_append_structural(false)
|
||||
.rows_from_receipt_with_text(&r, "close the position", &te);
|
||||
assert_eq!(a[0].embedding, a2[0].embedding);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_path_preserves_skip_and_negative_semantics() {
|
||||
let te = MockTextEmbedder { dim: 4 };
|
||||
let items = vec![
|
||||
(receipt(CascadeTier::L0Cache, 1), "cached hello".to_string()),
|
||||
(receipt(CascadeTier::L1Lawful, 2), "2 + 2".to_string()), // unmappable
|
||||
(receipt(CascadeTier::L3Model, 3), "write me a plan".to_string()),
|
||||
];
|
||||
let (rows, stats) = CorpusBuilder::default()
|
||||
.with_monotone_negatives(true)
|
||||
.build_with_text(&items, &te);
|
||||
assert_eq!(stats.receipts_seen, 3);
|
||||
assert_eq!(stats.receipts_used, 2);
|
||||
assert_eq!(stats.skipped_unmapped, 1); // L1Lawful, same as structural path
|
||||
assert_eq!(rows.len(), stats.rows());
|
||||
// Every row is the composite width, so the corpus is trainer-consistent.
|
||||
assert!(rows.iter().all(|r| r.embedding.len() == te.dim() + STRUCTURAL_DIM));
|
||||
// The L3 row carries its monotone negatives.
|
||||
assert!(rows.iter().any(|r| !r.success));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_corpus_serialises_to_trainer_jsonl() {
|
||||
let te = MockTextEmbedder { dim: 6 };
|
||||
let items = vec![(receipt(CascadeTier::L3Model, 77), "generate a summary".to_string())];
|
||||
let (rows, _) = CorpusBuilder::default().build_with_text(&items, &te);
|
||||
let text = to_jsonl(&rows).unwrap();
|
||||
let back: CorpusRow = serde_json::from_str(text.lines().next().unwrap()).unwrap();
|
||||
assert_eq!(back, rows[0]);
|
||||
assert_eq!(back.stage, CorpusStage::Neural);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -996,6 +996,11 @@ pub struct LocalInferenceInputs {
|
|||
/// Path to the model weights on local disk. Hashed into the proof, so the
|
||||
/// cell attests inference against *these bytes* and not merely "a model".
|
||||
pub weights: String,
|
||||
/// Optional path to a weight-hash cache. `None` (the default) streams the
|
||||
/// full file every run, preserving the strong bytes-bound guarantee. `Some`
|
||||
/// consults a [`WeightHashCache`] keyed on `(size, mtime)` so a multi-GiB
|
||||
/// model is not re-hashed each run — see that type for the trade-off.
|
||||
pub hash_cache: Option<std::path::PathBuf>,
|
||||
}
|
||||
|
||||
/// Stream a BLAKE3 over a file without loading it into memory — weights are
|
||||
|
|
@ -1018,6 +1023,119 @@ fn hash_file(path: &std::path::Path) -> Result<String, ProofError> {
|
|||
Ok(hasher.finalize().to_hex().to_string())
|
||||
}
|
||||
|
||||
/// A path's cheap freshness fingerprint: size in bytes and mtime in
|
||||
/// nanoseconds since the epoch. Read from metadata without touching contents.
|
||||
fn weight_fingerprint(path: &std::path::Path) -> Result<(u64, u128), ProofError> {
|
||||
let md = std::fs::metadata(path)
|
||||
.map_err(|e| ProofError::Conformance(format!("weights stat failed: {e}")))?;
|
||||
let mtime_ns = md
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0);
|
||||
Ok((md.len(), mtime_ns))
|
||||
}
|
||||
|
||||
/// One remembered weight hash, keyed by the file's freshness fingerprint.
|
||||
#[derive(Clone, serde::Serialize, serde::Deserialize)]
|
||||
struct WeightHashEntry {
|
||||
size: u64,
|
||||
mtime_ns: u128,
|
||||
blake3: String,
|
||||
}
|
||||
|
||||
/// A persistent map from weights path → its BLAKE3, so a 13 GiB model is not
|
||||
/// re-hashed on every proof.
|
||||
///
|
||||
/// ## The honest cost of caching a proof input
|
||||
///
|
||||
/// The proof binds to the BLAKE3 of the weight bytes. A cache that skips the
|
||||
/// re-hash must decide the bytes are unchanged *without reading them*, and it
|
||||
/// does so exactly as build systems do: an unchanged `(size, mtime)` is taken
|
||||
/// as evidence the contents are unchanged. That is a freshness heuristic, not a
|
||||
/// proof — an in-place edit that preserves both size and mtime would be missed.
|
||||
///
|
||||
/// So caching is **opt-in**: [`LocalInferenceInputs::hash_cache`] is `None` by
|
||||
/// default and the prover streams the full file every time, preserving the
|
||||
/// strong "bound to exactly these bytes" guarantee. An operator running the
|
||||
/// same weights in a loop sets a cache path and knowingly accepts the
|
||||
/// heuristic. The cache is never a correctness dependency: an unreadable or
|
||||
/// corrupt cache file degrades to a full hash, it never fails the proof.
|
||||
#[derive(Default, serde::Serialize, serde::Deserialize)]
|
||||
struct WeightHashCache {
|
||||
entries: std::collections::BTreeMap<String, WeightHashEntry>,
|
||||
}
|
||||
|
||||
impl WeightHashCache {
|
||||
/// Load the cache from `path`. A missing or unparseable file is an empty
|
||||
/// cache, never an error — the cache is an optimisation, not a dependency.
|
||||
fn load(path: &std::path::Path) -> Self {
|
||||
std::fs::read(path)
|
||||
.ok()
|
||||
.and_then(|b| serde_json::from_slice(&b).ok())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Best-effort persist. A write failure is swallowed: the proof already
|
||||
/// holds the hash it needs.
|
||||
fn save(&self, path: &std::path::Path) {
|
||||
if let Ok(bytes) = serde_json::to_vec_pretty(self) {
|
||||
let _ = std::fs::write(path, bytes);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a resolved weight hash came off the cache or was freshly streamed —
|
||||
/// surfaced so the CLI can be loud about which guarantee applied.
|
||||
enum HashSource {
|
||||
/// Streamed the full file this run.
|
||||
Fresh,
|
||||
/// Reused a cached hash on matching `(size, mtime)`.
|
||||
Cached,
|
||||
}
|
||||
|
||||
/// Resolve the BLAKE3 of `path`, consulting `cache_path` when given.
|
||||
///
|
||||
/// With no cache the file is always fully streamed. With a cache, a matching
|
||||
/// `(size, mtime)` fingerprint returns the stored hash; any miss (absent entry,
|
||||
/// changed size, changed mtime) streams the file and updates the cache. See
|
||||
/// [`WeightHashCache`] for the guarantee this trades.
|
||||
fn resolve_weight_hash(
|
||||
path: &std::path::Path,
|
||||
cache_path: Option<&std::path::Path>,
|
||||
) -> Result<(String, HashSource), ProofError> {
|
||||
let (size, mtime_ns) = weight_fingerprint(path)?;
|
||||
let key = path
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| path.to_path_buf())
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
|
||||
let Some(cache_path) = cache_path else {
|
||||
return Ok((hash_file(path)?, HashSource::Fresh));
|
||||
};
|
||||
|
||||
let mut cache = WeightHashCache::load(cache_path);
|
||||
if let Some(hit) = cache.entries.get(&key) {
|
||||
if hit.size == size && hit.mtime_ns == mtime_ns {
|
||||
return Ok((hit.blake3.clone(), HashSource::Cached));
|
||||
}
|
||||
}
|
||||
|
||||
let blake3 = hash_file(path)?;
|
||||
cache.entries.insert(
|
||||
key,
|
||||
WeightHashEntry {
|
||||
size,
|
||||
mtime_ns,
|
||||
blake3: blake3.clone(),
|
||||
},
|
||||
);
|
||||
cache.save(cache_path);
|
||||
Ok((blake3, HashSource::Fresh))
|
||||
}
|
||||
|
||||
/// Prove **local, offline inference**: run a local engine against local weights
|
||||
/// with no network grant, meter it on this host's own energy counters, and seal
|
||||
/// a cell.
|
||||
|
|
@ -1083,7 +1201,8 @@ pub fn prove_local(
|
|||
let weights_bytes = std::fs::metadata(weights_path)
|
||||
.map(|m| m.len())
|
||||
.unwrap_or(0);
|
||||
let weights_hash = hash_file(weights_path)?;
|
||||
let (weights_hash, hash_source) =
|
||||
resolve_weight_hash(weights_path, inputs.hash_cache.as_deref())?;
|
||||
|
||||
// (2) Deny-by-default egress, asserted rather than assumed.
|
||||
let grants = Grants {
|
||||
|
|
@ -1183,13 +1302,17 @@ pub fn prove_local(
|
|||
.map_err(|e| ProofError::Crypto(format!("attestation: {e}")))?;
|
||||
|
||||
eprintln!(
|
||||
" local: {} ({:.2} GiB, b3:{}…) — {} over {:.1}s, avg {:.1} W [{}]",
|
||||
" local: {} ({:.2} GiB, b3:{}… [{}]) — {} over {:.1}s, avg {:.1} W [{}]",
|
||||
weights_path
|
||||
.file_name()
|
||||
.map(|s| s.to_string_lossy().to_string())
|
||||
.unwrap_or_default(),
|
||||
weights_bytes as f64 / (1u64 << 30) as f64,
|
||||
&weights_hash[..12],
|
||||
match hash_source {
|
||||
HashSource::Fresh => "hashed",
|
||||
HashSource::Cached => "cache hit: trusts size+mtime",
|
||||
},
|
||||
if is_hw { "measured" } else { "ESTIMATED" },
|
||||
measurement.wall.as_secs_f64(),
|
||||
measurement.avg_watts(),
|
||||
|
|
@ -1341,9 +1464,86 @@ mod tests {
|
|||
engine: engine.to_string(),
|
||||
args: vec![],
|
||||
weights: weights.to_string(),
|
||||
hash_cache: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── weight-hash cache ────────────────────────────────────────────
|
||||
|
||||
/// A temp dir unique per process and per call, so tests never collide —
|
||||
/// even under `cargo test` parallelism or two concurrent test binaries.
|
||||
fn unique_dir(tag: &str) -> std::path::PathBuf {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
static N: AtomicU64 = AtomicU64::new(0);
|
||||
let d = std::env::temp_dir().join(format!(
|
||||
"sandbox-proof-{tag}-{}-{}",
|
||||
std::process::id(),
|
||||
N.fetch_add(1, Ordering::Relaxed)
|
||||
));
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hash_cache_hits_on_unchanged_file_and_matches_full_hash() {
|
||||
let dir = unique_dir("hashcache-hit");
|
||||
let w = dir.join("model.bin");
|
||||
std::fs::write(&w, b"some weight bytes").unwrap();
|
||||
let cache = dir.join("cache.json");
|
||||
let _ = std::fs::remove_file(&cache);
|
||||
|
||||
// Miss → streams the file. Must equal the plain full hash.
|
||||
let (h1, s1) = resolve_weight_hash(&w, Some(&cache)).unwrap();
|
||||
assert!(matches!(s1, HashSource::Fresh));
|
||||
assert_eq!(h1, hash_file(&w).unwrap());
|
||||
|
||||
// Second call, file untouched → cache hit, same hash, no re-stream.
|
||||
let (h2, s2) = resolve_weight_hash(&w, Some(&cache)).unwrap();
|
||||
assert!(matches!(s2, HashSource::Cached));
|
||||
assert_eq!(h1, h2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hash_cache_misses_when_contents_change() {
|
||||
let dir = unique_dir("hashcache-change");
|
||||
let w = dir.join("model.bin");
|
||||
let cache = dir.join("cache.json");
|
||||
let _ = std::fs::remove_file(&cache);
|
||||
|
||||
std::fs::write(&w, b"version one").unwrap();
|
||||
let (h1, _) = resolve_weight_hash(&w, Some(&cache)).unwrap();
|
||||
|
||||
// Different length → different fingerprint → re-hash, different digest.
|
||||
std::fs::write(&w, b"version two is longer").unwrap();
|
||||
let (h2, s2) = resolve_weight_hash(&w, Some(&cache)).unwrap();
|
||||
assert!(matches!(s2, HashSource::Fresh));
|
||||
assert_ne!(h1, h2, "a changed file must not return the stale hash");
|
||||
assert_eq!(h2, hash_file(&w).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_cache_always_streams_and_equals_full_hash() {
|
||||
let dir = unique_dir("hashcache-none");
|
||||
let w = dir.join("model.bin");
|
||||
std::fs::write(&w, b"deterministic contents").unwrap();
|
||||
let (h, s) = resolve_weight_hash(&w, None).unwrap();
|
||||
assert!(matches!(s, HashSource::Fresh));
|
||||
assert_eq!(h, hash_file(&w).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupt_cache_file_degrades_to_a_full_hash() {
|
||||
// The cache is an optimisation, never a correctness dependency.
|
||||
let dir = unique_dir("hashcache-corrupt");
|
||||
let w = dir.join("model.bin");
|
||||
std::fs::write(&w, b"real bytes").unwrap();
|
||||
let cache = dir.join("cache.json");
|
||||
std::fs::write(&cache, b"}{ this is not json").unwrap();
|
||||
let (h, s) = resolve_weight_hash(&w, Some(&cache)).unwrap();
|
||||
assert!(matches!(s, HashSource::Fresh));
|
||||
assert_eq!(h, hash_file(&w).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_proof_refuses_absent_weights() {
|
||||
// An air-gap proof whose weights are not actually on this disk is
|
||||
|
|
|
|||
|
|
@ -150,7 +150,21 @@ fn prove_local(path: &str) {
|
|||
.ok()
|
||||
.map(|s| s.split_whitespace().map(String::from).collect())
|
||||
.unwrap_or_else(|| vec![weights.clone()]);
|
||||
let inputs = sandbox_proof::LocalInferenceInputs { engine, args, weights };
|
||||
// Weight-hash cache: on by default at a temp path so a 13 GiB model is not
|
||||
// re-hashed on every run. `SANDBOX_WEIGHTS_HASH_CACHE=off` forces a full
|
||||
// stream every time (the strong bytes-bound guarantee); any other value is
|
||||
// used as the cache file path. The prover prints which one applied.
|
||||
let hash_cache = match std::env::var("SANDBOX_WEIGHTS_HASH_CACHE").as_deref() {
|
||||
Ok("off") | Ok("none") | Ok("") => None,
|
||||
Ok(p) => Some(std::path::PathBuf::from(p)),
|
||||
Err(_) => Some(std::env::temp_dir().join("sandbox-proof-weights-hash.json")),
|
||||
};
|
||||
let inputs = sandbox_proof::LocalInferenceInputs {
|
||||
engine,
|
||||
args,
|
||||
weights,
|
||||
hash_cache,
|
||||
};
|
||||
let cloud = env("SANDBOX_CLOUD", "local");
|
||||
let cell = match sandbox_proof::prove_local(&inputs, &cloud, now_us(), &proof_key()) {
|
||||
Ok(c) => c,
|
||||
|
|
|
|||
Loading…
Reference in a new issue