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(_)
58 | StepBody::Rank(_)
59 | StepBody::Select(_) => None,
60 })?;
61 Some((binding, decl.task.clone(), decl.options.clone()))
62}
63
64#[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
75pub 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 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#[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
217pub 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}