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(_) => None,
58        })?;
59    Some((binding, decl.task.clone(), decl.options.clone()))
60}
61
62/// What `calibrate fit` reads and writes.
63#[derive(Debug, Clone, Copy)]
64pub struct FitArgs<'a> {
65    pub replay: &'a ReplayConfig,
66    pub out: &'a str,
67    pub dataset: &'a str,
68    pub config: &'a str,
69    pub bundles: &'a str,
70    pub step: &'a str,
71}
72
73/// `calibrate fit`: a temperature calibration of one step, fitted on the
74/// calibration split, published in `out`.
75pub fn fit(command: &str, host: &dyn Host, a: FitArgs<'_>, resolver: &dyn Resolver) -> Res {
76    let replay = a.replay;
77    (|| {
78        let out_dir = a.out;
79        for name in FIT_OUTPUTS {
80            host.ensure_absent(&host.join(out_dir, name))
81                .map_err(|e| Done::io(command, e))?;
82        }
83        let dataset_path = a.dataset;
84        let manifest: DatasetManifest = read_doc(command, host, dataset_path, REPLAY_V1.max_bytes)?;
85        manifest
86            .check_structure()
87            .map_err(|e| Done::invalid(command, dataset_path, e.to_string()))?;
88        let config_path = a.config;
89        let config: EvaluatorConfig = read_doc(command, host, config_path, REPLAY_V1.max_bytes)?;
90        check_case_ids(command, dataset_path, manifest.cases.iter())?;
91        let bundles = read_bundles(command, host, a.bundles, manifest.split(Split::Calibration))?;
92        let step = a.step;
93        let mut skipped: BTreeMap<String, u64> = BTreeMap::new();
94        let mut samples: Vec<(String, RawOutput)> = vec![];
95        let mut plan: Option<Plan> = None;
96        for case in manifest.split(Split::Calibration) {
97            let bundle = match bundles.get(&case.id) {
98                None => {
99                    bump(&mut skipped, "bundle-missing");
100                    continue;
101                }
102                Some(BundleInput::Failed(f)) => {
103                    bump(&mut skipped, f.kind.code());
104                    continue;
105                }
106                Some(BundleInput::Loaded(b)) => b,
107            };
108            let r = match reproduce(bundle, resolver, replay) {
109                ReplayOutcome::Reproduced(r) => r,
110                ReplayOutcome::Diverged { .. } => {
111                    bump(&mut skipped, "diverged");
112                    continue;
113                }
114                ReplayOutcome::Incomparable(i) => {
115                    bump(&mut skipped, &i.code());
116                    continue;
117                }
118            };
119            match &plan {
120                None => plan = Some(r.plan().clone()),
121                Some(first) if first != r.plan() => {
122                    return Err(Done::invalid(
123                        command,
124                        dataset_path,
125                        "reproduced cases come from different plans",
126                    ));
127                }
128                Some(_) => {}
129            }
130            let Some((artifact, ..)) = target_parts(r.plan(), step) else {
131                return Err(Done::invalid(
132                    command,
133                    "--step",
134                    format!("{step:?} is not a semantic step bound by the plan"),
135                ));
136            };
137            let Some(supply) = r
138                .supplies()
139                .iter()
140                .find(|x| x.step == step && x.instance.is_empty())
141            else {
142                bump(&mut skipped, "no-output");
143                continue;
144            };
145            match &supply.value {
146                RetainedValue::Reason(_) => bump(&mut skipped, "runtime-failure"),
147                // A fallback's output is never fitted as the primary's.
148                RetainedValue::Output { artifact: a, .. }
149                    if supply.target != 0 || *a != artifact =>
150                {
151                    bump(&mut skipped, "fallback-output")
152                }
153                RetainedValue::Output { output, .. } => {
154                    samples.push((case.id.clone(), output.clone()))
155                }
156            }
157        }
158        let Some(plan) = plan else {
159            return Err(Done::invalid(
160                command,
161                dataset_path,
162                "no calibration-split case reproduced",
163            ));
164        };
165        let Some((artifact, task, options)) = target_parts(&plan, step) else {
166            return Err(Done::invalid(
167                command,
168                "--step",
169                format!("{step:?} is not a semantic step"),
170            ));
171        };
172        let target = FitTarget {
173            artifact: &artifact,
174            task: &task,
175            question: step,
176            options: &options,
177            tolerance: plan.definition.limits.distribution_tolerance,
178        };
179        let (calibration, fit) = fit_temperature(&manifest, &config, &target, &samples)
180            .map_err(|e| Done::invalid(command, dataset_path, e.to_string()))?;
181        let bytes = |r: Result<Vec<u8>, rustev_contract::canonical::CanonicalError>| {
182            r.map_err(|e| Done::invalid(command, out_dir, e.to_string()))
183        };
184        write_all(
185            command,
186            host,
187            out_dir,
188            vec![
189                (FIT_OUTPUTS[0], bytes(calibration.canonical())?),
190                (FIT_OUTPUTS[1], bytes(fit.canonical())?),
191            ],
192        )?;
193        Ok(Done::ok(
194            Out::new(command, "fitted")
195                .set("calibration", &id_of(&calibration))
196                .set("temperature", &fit.temperature)
197                .set("mean_log_loss", &fit.mean_log_loss)
198                .set("fitted", &fit.fitted)
199                .set("excluded", &fit.excluded)
200                .set("skipped", &skipped)
201                .set("out", out_dir),
202        ))
203    })()
204}
205
206/// What `calibrate qualify` reads.
207#[derive(Debug, Clone, Copy)]
208pub struct QualifyArgs<'a> {
209    pub split: Split,
210    pub calibration: &'a str,
211    pub fit: Option<&'a str>,
212    pub dataset: &'a str,
213}
214
215/// `calibrate qualify`: whether a calibration may be used for `split`.
216pub fn qualify_cmd(command: &str, host: &dyn Host, a: QualifyArgs<'_>) -> Res {
217    let split = a.split;
218    (|| {
219        let artifact: CalibrationArtifact =
220            read_doc(command, host, a.calibration, DESCRIPTOR_V1.max_bytes)?;
221        let fit: Option<CalibrationFit> = match a.fit {
222            None => None,
223            Some(path) => Some(read_doc(command, host, path, REPLAY_V1.max_bytes)?),
224        };
225        let manifest: DatasetManifest = read_doc(command, host, a.dataset, REPLAY_V1.max_bytes)?;
226        let out = Out::new(command, "").set("calibration", &id_of(&artifact));
227        Ok(match qualify(&artifact, fit.as_ref(), &manifest, split) {
228            Qualification::Qualified => Done::ok(out.set("status", "qualified")),
229            Qualification::Refused { reason } => Done::new(
230                code::FAIL,
231                out.set("status", "refused_qualification")
232                    .set("reason", &reason),
233            ),
234            Qualification::Unknown { reason } => Done::new(
235                code::UNKNOWN,
236                out.set("status", "unknown").set("reason", &reason),
237            ),
238        })
239    })()
240}