1use 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
29fn 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#[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
73pub 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 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#[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
215pub 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}