xtask/bench/
model.rs

1//! Resolve benchmark inputs through the CLI without linking the Metal engine.
2
3use crate::util;
4use anyhow::{Context, Result, ensure};
5use serde_json::{Value, json};
6use sha2::{Digest, Sha256};
7use std::{
8    ffi::OsStr,
9    path::{Path, PathBuf},
10    process::Command,
11};
12
13/// Pinned index entry and its currently selected prepared artifact.
14pub(super) struct Model {
15    pub reference: String,
16    pub directory: PathBuf,
17    pub precisions: Vec<u8>,
18    /// Original local checkpoint, used to recognize reports made before the index.
19    pub source_directory: Option<PathBuf>,
20}
21
22impl Model {
23    /// Use the same resolver as inference; preparation and inspection are untimed.
24    pub(super) fn inspect(binary: &Path, selector: &OsStr, root: Option<&Path>) -> Result<Self> {
25        let mut command = Command::new(binary);
26
27        command.arg("inspect").arg(selector).arg("--json");
28
29        if let Some(root) = root {
30            command.arg("--root").arg(root);
31        }
32
33        let output = command.output().context("inspecting benchmark model")?;
34
35        ensure!(
36            output.status.success(),
37            "model inspection failed: {}",
38            String::from_utf8_lossy(&output.stderr).trim()
39        );
40
41        Self::from_inspection(
42            &serde_json::from_slice(&output.stdout).context("decoding model inspection")?,
43        )
44    }
45
46    /// Extract the index identity and prepared location from the public inspection view.
47    fn from_inspection(value: &Value) -> Result<Self> {
48        let index = &value["index"];
49        let reference = index["id"]
50            .as_str()
51            .context("inspection missing model ID")?
52            .to_owned();
53        let artifact = index["artifacts"]
54            .as_array()
55            .context("inspection missing artifacts")?
56            .iter()
57            .find(|artifact| artifact["kind"] == "prepared" && artifact["available"] == true)
58            .context(
59                "model must be prepared before benchmarking; run `cherenkov prepare SOURCE`",
60            )?;
61        let directory = artifact["path"]
62            .as_str()
63            .context("prepared artifact missing path")?
64            .into();
65        let precisions = serde_json::from_value(index["precisions"].clone())
66            .context("inspection missing expert precisions")?;
67
68        Ok(Self {
69            reference,
70            directory,
71            precisions,
72            source_directory: index["source"]["path"].as_str().map(PathBuf::from),
73        })
74    }
75
76    /// Upgrade a proven source identity without exposing its path in shared reports.
77    pub(super) fn migrate_signature(&self, signature: &mut Value) {
78        if signature["model_id"].is_string() {
79            return;
80        }
81
82        let matches_id = signature["model"].as_str() == Some(&self.reference);
83        let matches_path = self.source_directory.as_ref().is_some_and(|source| {
84            let digest: String = Sha256::digest(source.as_os_str().as_encoded_bytes())
85                .iter()
86                .map(|byte| format!("{byte:02x}"))
87                .collect();
88
89            signature["model"]
90                .as_str()
91                .is_some_and(|saved| Some(saved) == source.to_str())
92                || signature["model_path_sha256"].as_str() == Some(&digest)
93        });
94
95        if !matches_id && !matches_path {
96            return;
97        }
98
99        signature["model"] = json!("<model>");
100        signature["model_id"] = json!(self.reference);
101
102        signature
103            .as_object_mut()
104            .unwrap()
105            .remove("model_path_sha256");
106
107        if let Some(metadata) = signature["model_metadata_sha256"].as_object_mut()
108            && let Some(manifest) = metadata.remove("packed/manifest.json")
109        {
110            metadata.entry("manifest.json").or_insert(manifest);
111        }
112    }
113
114    /// Hash equivalent metadata in flat artifacts and legacy model/packed layouts.
115    pub(super) fn metadata(&self) -> Result<Value> {
116        let mut metadata = json!({});
117
118        for name in ["config.json", "tokenizer.json"] {
119            let direct = self.directory.join(name);
120            let path = if direct.is_file() {
121                direct
122            } else {
123                self.directory
124                    .parent()
125                    .context("model metadata missing")?
126                    .join(name)
127            };
128            metadata[name] = json!(util::digest(&path)?);
129        }
130
131        metadata["manifest.json"] = json!(util::digest(&self.directory.join("manifest.json"))?);
132
133        Ok(metadata)
134    }
135}
136
137#[cfg(test)]
138#[path = "../../tests/unit/bench_model.rs"]
139mod tests;
xtask/bench/model.rs

xtask/bench/
model.rs
