Skip to main content

rustev_surface/
calibrate.rs

1//! `calibrate fit | qualify` (spec 006, 3.10), for every host.
2
3use std::collections::BTreeMap;
4
5use rustev_contract::Document;
6use rustev_contract::calibration::CalibrationArtifact;
7use rustev_contract::definition::StepBody;
8use rustev_contract::eval_report::Split;
9use rustev_contract::limits::{DESCRIPTOR_V1, REPLAY_V1};
10use rustev_contract::output::RawOutput;
11use rustev_contract::plan::{Plan, PlanStepDetail};
12use rustev_eval::bundle::BundleInput;
13use rustev_eval::config::EvaluatorConfig;
14use rustev_eval::dataset::DatasetManifest;
15use rustev_eval::fit::{CalibrationFit, FitTarget, Qualification, fit_temperature, qualify};
16use rustev_eval::replay::{ReplayConfig, ReplayOutcome, RetainedValue, reproduce};
17use rustev_eval::resolve::Resolver;
18
19use crate::eval::{check_case_ids, id_of, read_bundles, read_doc, write_all};
20use crate::host::Host;
21use crate::out::{Done, Out, Res, code};
22
23pub const FIT_OUTPUTS: [&str; 2] = ["calibration.json", "calibration-fit.json"];
24
25fn bump(m: &mut BTreeMap<String, u64>, k: &str) {
26    *m.entry(k.to_string()).or_default() += 1;
27}
28
29/// The fit target of `step` in `plan`: its primary binding's artifact, its
30/// task and options, the step id as the calibration's question, and the
31/// plan's distribution tolerance.
32fn target_parts(
33    plan: &Plan,
34    step: &str,
35) -> Option<(rustev_contract::ids::ArtifactId, String, Vec<String>)> {
36    let binding = plan
37        .steps
38        .iter()
39        .find(|s| s.id == step)
40        .and_then(|s| match &s.detail {
41            PlanStepDetail::Semantic(b) => Some(b.artifact.clone()),
42            _ => None,
43        })?;
44    let decl = plan
45        .definition
46        .steps
47        .iter()
48        .find(|s| s.id == step)
49        .and_then(|s| match &s.body {
50            StepBody::Semantic(d) => Some(d),
51            StepBody::Exact(_)
52            | StepBody::Corroborate(_)
53            | StepBody::Acquire(_)
54            | StepBody::Widen(_)
55            | StepBody::Generate(_)
56            | StepBody::Extract(_)
57            | StepBody::AcceptModelDerived(_)
58            | StepBody::Rank(_)
59            | StepBody::Select(_) => None,
60        })?;
61    Some((binding, decl.task.clone(), decl.options.clone()))
62}
63
64/// What `calibrate fit` reads and writes.
65#[derive(Debug, Clone, Copy)]
66pub struct FitArgs<'a> {
67    pub replay: &'a ReplayConfig,
68    pub out: &'a str,
69    pub dataset: &'a str,
70    pub config: &'a str,
71    pub bundles: &'a str,
72    pub step: &'a str,
73}
74
75/// `calibrate fit`: a temperature calibration of one step, fitted on the
76/// calibration split, published in `out`.
77pub fn fit(command: &str, host: &dyn Host, a: FitArgs<'_>, resolver: &dyn Resolver) -> Res {
78    let replay = a.replay;
79    (|| {
80        let out_dir = a.out;
81        for name in FIT_OUTPUTS {
82            host.ensure_absent(&host.join(out_dir, name))
83                .map_err(|e| Done::io(command, e))?;
84        }
85        let dataset_path = a.dataset;
86        let manifest: DatasetManifest = read_doc(command, host, dataset_path, REPLAY_V1.max_bytes)?;
87        manifest
88            .check_structure()
89            .map_err(|e| Done::invalid(command, dataset_path, e.to_string()))?;
90        let config_path = a.config;
91        let config: EvaluatorConfig = read_doc(command, host, config_path, REPLAY_V1.max_bytes)?;
92        check_case_ids(command, dataset_path, manifest.cases.iter())?;
93        let bundles = read_bundles(command, host, a.bundles, manifest.split(Split::Calibration))?;
94        let step = a.step;
95        let mut skipped: BTreeMap<String, u64> = BTreeMap::new();
96        let mut samples: Vec<(String, RawOutput)> = vec![];
97        let mut plan: Option<Plan> = None;
98        for case in manifest.split(Split::Calibration) {
99            let bundle = match bundles.get(&case.id) {
100                None => {
101                    bump(&mut skipped, "bundle-missing");
102                    continue;
103                }
104                Some(BundleInput::Failed(f)) => {
105                    bump(&mut skipped, f.kind.code());
106                    continue;
107                }
108                Some(BundleInput::Loaded(b)) => b,
109            };
110            let r = match reproduce(bundle, resolver, replay) {
111                ReplayOutcome::Reproduced(r) => r,
112                ReplayOutcome::Diverged { .. } => {
113                    bump(&mut skipped, "diverged");
114                    continue;
115                }
116                ReplayOutcome::Incomparable(i) => {
117                    bump(&mut skipped, &i.code());
118                    continue;
119                }
120            };
121            match &plan {
122                None => plan = Some(r.plan().clone()),
123                Some(first) if first != r.plan() => {
124                    return Err(Done::invalid(
125                        command,
126                        dataset_path,
127                        "reproduced cases come from different plans",
128                    ));
129                }
130                Some(_) => {}
131            }
132            let Some((artifact, ..)) = target_parts(r.plan(), step) else {
133                return Err(Done::invalid(
134                    command,
135                    "--step",
136                    format!("{step:?} is not a semantic step bound by the plan"),
137                ));
138            };
139            let Some(supply) = r
140                .supplies()
141                .iter()
142                .find(|x| x.step == step && x.instance.is_empty())
143            else {
144                bump(&mut skipped, "no-output");
145                continue;
146            };
147            match &supply.value {
148                RetainedValue::Reason(_) => bump(&mut skipped, "runtime-failure"),
149                // A fallback's output is never fitted as the primary's.
150                RetainedValue::Output { artifact: a, .. }
151                    if supply.target != 0 || *a != artifact =>
152                {
153                    bump(&mut skipped, "fallback-output")
154                }
155                RetainedValue::Output { output, .. } => {
156                    samples.push((case.id.clone(), output.clone()))
157                }
158            }
159        }
160        let Some(plan) = plan else {
161            return Err(Done::invalid(
162                command,
163                dataset_path,
164                "no calibration-split case reproduced",
165            ));
166        };
167        let Some((artifact, task, options)) = target_parts(&plan, step) else {
168            return Err(Done::invalid(
169                command,
170                "--step",
171                format!("{step:?} is not a semantic step"),
172            ));
173        };
174        let target = FitTarget {
175            artifact: &artifact,
176            task: &task,
177            question: step,
178            options: &options,
179            tolerance: plan.definition.limits.distribution_tolerance,
180        };
181        let (calibration, fit) = fit_temperature(&manifest, &config, &target, &samples)
182            .map_err(|e| Done::invalid(command, dataset_path, e.to_string()))?;
183        let bytes = |r: Result<Vec<u8>, rustev_contract::canonical::CanonicalError>| {
184            r.map_err(|e| Done::invalid(command, out_dir, e.to_string()))
185        };
186        write_all(
187            command,
188            host,
189            out_dir,
190            vec![
191                (FIT_OUTPUTS[0], bytes(calibration.canonical())?),
192                (FIT_OUTPUTS[1], bytes(fit.canonical())?),
193            ],
194        )?;
195        Ok(Done::ok(
196            Out::new(command, "fitted")
197                .set("calibration", &id_of(&calibration))
198                .set("temperature", &fit.temperature)
199                .set("mean_log_loss", &fit.mean_log_loss)
200                .set("fitted", &fit.fitted)
201                .set("excluded", &fit.excluded)
202                .set("skipped", &skipped)
203                .set("out", out_dir),
204        ))
205    })()
206}
207
208/// What `calibrate qualify` reads.
209#[derive(Debug, Clone, Copy)]
210pub struct QualifyArgs<'a> {
211    pub split: Split,
212    pub calibration: &'a str,
213    pub fit: Option<&'a str>,
214    pub dataset: &'a str,
215}
216
217/// `calibrate qualify`: whether a calibration may be used for `split`.
218pub fn qualify_cmd(command: &str, host: &dyn Host, a: QualifyArgs<'_>) -> Res {
219    let split = a.split;
220    (|| {
221        let artifact: CalibrationArtifact =
222            read_doc(command, host, a.calibration, DESCRIPTOR_V1.max_bytes)?;
223        let fit: Option<CalibrationFit> = match a.fit {
224            None => None,
225            Some(path) => Some(read_doc(command, host, path, REPLAY_V1.max_bytes)?),
226        };
227        let manifest: DatasetManifest = read_doc(command, host, a.dataset, REPLAY_V1.max_bytes)?;
228        let out = Out::new(command, "").set("calibration", &id_of(&artifact));
229        Ok(match qualify(&artifact, fit.as_ref(), &manifest, split) {
230            Qualification::Qualified => Done::ok(out.set("status", "qualified")),
231            Qualification::Refused { reason } => Done::new(
232                code::FAIL,
233                out.set("status", "refused_qualification")
234                    .set("reason", &reason),
235            ),
236            Qualification::Unknown { reason } => Done::new(
237                code::UNKNOWN,
238                out.set("status", "unknown").set("reason", &reason),
239            ),
240        })
241    })()
242}