1use std::collections::BTreeSet;
21use std::fmt;
22
23use rustev_contract::canonical::record_canonical_bytes;
24use rustev_contract::ids::ArtifactId;
25use rustev_contract::limits::REPLAY_V1;
26use rustev_contract::output::RawOutput;
27use rustev_contract::remote::{RemoteCode, cut_to};
28use rustev_contract::{Document, DocumentError};
29use serde::{Deserialize, Serialize};
30
31pub const BATCH_QUALIFICATION: &str = "rustev.batch-qualification/1";
32pub const MAX_TEXT_BYTES: usize = 256;
34pub const MAX_PAIRS: usize = 10_000;
36
37#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
42#[serde(deny_unknown_fields)]
43pub struct BatchQualification {
44 pub schema: String,
45 pub adapter: String,
47 pub bindings: Bindings,
48 pub batching: BatchingConfig,
50 pub bound: AgreementBound,
52 pub min_cases: u64,
54 pub stage1: Vec<Stage1Result>,
56 pub stage2: Option<Stage2>,
58 pub verdict: Verdict,
59}
60
61crate::doc::eval_document!(BatchQualification, BATCH_QUALIFICATION, REPLAY_V1);
62
63#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
66#[serde(deny_unknown_fields)]
67pub struct Bindings {
68 pub batched: ArtifactId,
69 pub unbatched: ArtifactId,
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
74#[serde(deny_unknown_fields)]
75pub struct BatchingConfig {
76 pub max_items: u32,
78 pub window_ms: u64,
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
83#[serde(deny_unknown_fields)]
84pub struct AgreementBound {
85 pub criterion: Agreement,
86 pub min_agreeing: Ratio,
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
92#[serde(rename_all = "snake_case")]
93pub enum Agreement {
94 Exact,
96 TopLabel,
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
104#[serde(deny_unknown_fields)]
105pub struct Ratio {
106 pub numerator: u64,
107 pub denominator: u64,
108}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
113#[serde(rename_all = "snake_case")]
114pub enum Stage1Case {
115 BatchBounds,
116 MemberAttribution,
117 MissingMember,
118 MalformedMember,
119 PartialMemberLiability,
120 CancelledBeforeSend,
121 CancelledAfterSend,
122 MemberDeadlines,
123 SettlementRemainder,
124 SettlementUnknownCharge,
125 FailureTransient,
126 FailureOverloaded,
127 FailurePermanent,
128 FailureCancelled,
129 PerItemIsolation,
130}
131
132impl Stage1Case {
133 pub const ALL: [Stage1Case; 15] = [
135 Stage1Case::BatchBounds,
136 Stage1Case::MemberAttribution,
137 Stage1Case::MissingMember,
138 Stage1Case::MalformedMember,
139 Stage1Case::PartialMemberLiability,
140 Stage1Case::CancelledBeforeSend,
141 Stage1Case::CancelledAfterSend,
142 Stage1Case::MemberDeadlines,
143 Stage1Case::SettlementRemainder,
144 Stage1Case::SettlementUnknownCharge,
145 Stage1Case::FailureTransient,
146 Stage1Case::FailureOverloaded,
147 Stage1Case::FailurePermanent,
148 Stage1Case::FailureCancelled,
149 Stage1Case::PerItemIsolation,
150 ];
151}
152
153#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
154#[serde(deny_unknown_fields)]
155pub struct Stage1Result {
156 pub case: Stage1Case,
157 pub outcome: CaseOutcome,
158}
159
160impl Stage1Result {
161 pub fn passed(case: Stage1Case) -> Stage1Result {
162 Stage1Result {
163 case,
164 outcome: CaseOutcome::Passed,
165 }
166 }
167
168 pub fn failed(case: Stage1Case, reason: &str) -> Stage1Result {
170 Stage1Result {
171 case,
172 outcome: CaseOutcome::Failed {
173 reason: cut_to(reason, MAX_TEXT_BYTES),
174 },
175 }
176 }
177}
178
179#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
180#[serde(rename_all = "snake_case", deny_unknown_fields)]
181pub enum CaseOutcome {
182 Passed,
183 Failed { reason: String },
184}
185
186#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
188#[serde(deny_unknown_fields)]
189pub struct Stage2 {
190 pub pairs: Vec<PairedAnswer>,
191 pub spend: Spend,
192}
193
194#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
196#[serde(deny_unknown_fields)]
197pub struct PairedAnswer {
198 pub case: String,
200 pub batched: Answer,
201 pub unbatched: Answer,
202}
203
204#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
205#[serde(rename_all = "snake_case", deny_unknown_fields)]
206pub enum Answer {
207 Output(RawOutput),
208 Failure { code: RemoteCode },
209}
210
211#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
213#[serde(deny_unknown_fields)]
214pub struct Spend {
215 pub observed_units: u64,
216 pub estimated_units: u64,
217 pub unknown_settlements: u64,
219}
220
221#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
222#[serde(rename_all = "snake_case")]
223pub enum Verdict {
224 Qualified,
225 NotQualified,
226 Unknown,
227}
228
229#[derive(Debug, Clone, PartialEq, Eq)]
231pub enum BatchQualificationError {
232 Parse(DocumentError),
233 Schema,
234 Text(&'static str),
236 Batching,
238 Bound,
241 Stage1Order,
243 Stage1Reason(Stage1Case),
245 TooManyPairs(usize),
246 PairCase(usize),
248 Verdict {
250 recorded: Verdict,
251 computed: Verdict,
252 },
253}
254
255impl fmt::Display for BatchQualificationError {
256 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
257 use BatchQualificationError as E;
258 match self {
259 E::Parse(e) => write!(f, "{e}"),
260 E::Schema => write!(f, "not a {BATCH_QUALIFICATION} document"),
261 E::Text(what) => write!(f, "{what} is blank or over {MAX_TEXT_BYTES} bytes"),
262 E::Batching => write!(f, "batching with max_items below 2"),
263 E::Bound => write!(f, "the agreement bound or minimum is not a valid bound"),
264 E::Stage1Order => write!(f, "stage 1 cases are repeated or out of order"),
265 E::Stage1Reason(c) => write!(f, "stage 1 case {c:?} failed without a valid reason"),
266 E::TooManyPairs(n) => write!(f, "{n} pairs exceed {MAX_PAIRS}"),
267 E::PairCase(i) => write!(f, "pair {i} has a blank, long or repeated case name"),
268 E::Verdict { recorded, computed } => {
269 write!(
270 f,
271 "verdict {recorded:?} where the report supports {computed:?}"
272 )
273 }
274 }
275 }
276}
277
278impl std::error::Error for BatchQualificationError {}
279
280fn text_ok(s: &str) -> bool {
281 !s.trim().is_empty() && s.len() <= MAX_TEXT_BYTES
282}
283
284fn top_label(o: &RawOutput) -> Option<&str> {
286 let entries = match o {
287 RawOutput::Label(l) => return Some(l),
288 RawOutput::Distribution(m) | RawOutput::Logits(m) | RawOutput::Scores(m) => m,
289 _ => return None,
290 };
291 if entries.values().any(|v| !v.is_finite()) {
292 return None;
293 }
294 let max = entries.values().copied().reduce(f64::max)?;
295 let mut top = entries.iter().filter(|(_, v)| **v == max);
296 match (top.next(), top.next()) {
297 (Some((k, _)), None) => Some(k),
298 _ => None,
299 }
300}
301
302impl AgreementBound {
303 pub fn agrees(&self, pair: &PairedAnswer) -> Option<bool> {
306 match (&pair.batched, &pair.unbatched) {
307 (Answer::Failure { .. }, Answer::Failure { .. }) => None,
308 (Answer::Output(a), Answer::Output(b)) => Some(match self.criterion {
309 Agreement::Exact => match (record_canonical_bytes(a), record_canonical_bytes(b)) {
310 (Ok(a), Ok(b)) => a == b,
311 _ => false,
312 },
313 Agreement::TopLabel => matches!(
314 (top_label(a), top_label(b)),
315 (Some(a), Some(b)) if a == b
316 ),
317 }),
318 _ => Some(false),
319 }
320 }
321}
322
323impl BatchQualification {
324 #[allow(clippy::too_many_arguments)]
327 pub fn new(
328 adapter: &str,
329 bindings: Bindings,
330 batching: BatchingConfig,
331 bound: AgreementBound,
332 min_cases: u64,
333 mut stage1: Vec<Stage1Result>,
334 stage2: Option<Stage2>,
335 ) -> Result<BatchQualification, BatchQualificationError> {
336 stage1.sort_by_key(|r| r.case);
337 let mut report = BatchQualification {
338 schema: BATCH_QUALIFICATION.into(),
339 adapter: adapter.into(),
340 bindings,
341 batching,
342 bound,
343 min_cases,
344 stage1,
345 stage2,
346 verdict: Verdict::Unknown,
347 };
348 report.check_members()?;
349 report.verdict = report.derived_verdict();
350 Ok(report)
351 }
352
353 pub fn parse(bytes: &[u8]) -> Result<BatchQualification, BatchQualificationError> {
355 let report = <Self as Document>::parse(bytes).map_err(BatchQualificationError::Parse)?;
356 report.check()?;
357 Ok(report)
358 }
359
360 pub fn check(&self) -> Result<(), BatchQualificationError> {
363 self.check_members()?;
364 let computed = self.derived_verdict();
365 if self.verdict != computed {
366 return Err(BatchQualificationError::Verdict {
367 recorded: self.verdict,
368 computed,
369 });
370 }
371 Ok(())
372 }
373
374 fn check_members(&self) -> Result<(), BatchQualificationError> {
375 use BatchQualificationError as E;
376 if self.schema != BATCH_QUALIFICATION {
377 return Err(E::Schema);
378 }
379 if !text_ok(&self.adapter) {
380 return Err(E::Text("the adapter name"));
381 }
382 if self.batching.max_items < 2 {
383 return Err(E::Batching);
384 }
385 let r = self.bound.min_agreeing;
386 if r.denominator == 0 || r.numerator > r.denominator || self.min_cases == 0 {
387 return Err(E::Bound);
388 }
389 if self.stage1.windows(2).any(|w| w[0].case >= w[1].case) {
390 return Err(E::Stage1Order);
391 }
392 for result in &self.stage1 {
393 if let CaseOutcome::Failed { reason } = &result.outcome
394 && !text_ok(reason)
395 {
396 return Err(E::Stage1Reason(result.case));
397 }
398 }
399 if let Some(stage2) = &self.stage2 {
400 if stage2.pairs.len() > MAX_PAIRS {
401 return Err(E::TooManyPairs(stage2.pairs.len()));
402 }
403 let mut names = BTreeSet::new();
404 for (i, pair) in stage2.pairs.iter().enumerate() {
405 if !text_ok(&pair.case) || !names.insert(pair.case.as_str()) {
406 return Err(E::PairCase(i));
407 }
408 }
409 }
410 Ok(())
411 }
412
413 pub fn agreement(&self) -> Option<(u64, u64)> {
415 let stage2 = self.stage2.as_ref()?;
416 let (mut agreeing, mut comparable) = (0, 0);
417 for agrees in stage2.pairs.iter().filter_map(|p| self.bound.agrees(p)) {
418 comparable += 1;
419 agreeing += u64::from(agrees);
420 }
421 Some((agreeing, comparable))
422 }
423
424 pub fn derived_verdict(&self) -> Verdict {
427 if self.check_members().is_err() {
428 return Verdict::Unknown;
429 }
430 let failed = self
431 .stage1
432 .iter()
433 .any(|r| matches!(r.outcome, CaseOutcome::Failed { .. }));
434 let complete = self.stage1.len() == Stage1Case::ALL.len() && !failed;
435 let stage2 = self
438 .agreement()
439 .filter(|(_, comparable)| *comparable >= self.min_cases)
440 .map(|(agreeing, comparable)| {
441 let r = self.bound.min_agreeing;
442 u128::from(agreeing) * u128::from(r.denominator)
443 >= u128::from(r.numerator) * u128::from(comparable)
444 });
445 match stage2 {
446 _ if failed => Verdict::NotQualified,
447 Some(false) => Verdict::NotQualified,
448 Some(true) if complete => Verdict::Qualified,
449 _ => Verdict::Unknown,
450 }
451 }
452}
453
454#[cfg(test)]
455mod tests {
456 use std::collections::BTreeMap;
457
458 use super::*;
459
460 const ID: &str = "sha256:1111111111111111111111111111111111111111111111111111111111111111";
461
462 fn bindings() -> Bindings {
463 Bindings {
464 batched: ArtifactId::parse(ID).unwrap(),
465 unbatched: ArtifactId::parse(ID).unwrap(),
466 }
467 }
468
469 fn bound(criterion: Agreement, numerator: u64, denominator: u64) -> AgreementBound {
470 AgreementBound {
471 criterion,
472 min_agreeing: Ratio {
473 numerator,
474 denominator,
475 },
476 }
477 }
478
479 fn all_passed() -> Vec<Stage1Result> {
480 Stage1Case::ALL
481 .iter()
482 .rev()
483 .map(|c| Stage1Result::passed(*c))
484 .collect()
485 }
486
487 fn dist(top: &str) -> Answer {
488 Answer::Output(RawOutput::Distribution(BTreeMap::from([
489 (top.to_string(), 0.75),
490 ("SYNTHETIC-other".to_string(), 0.25),
491 ])))
492 }
493
494 fn pairs(agreeing: usize, disagreeing: usize) -> Stage2 {
495 let mut pairs: Vec<PairedAnswer> = (0..agreeing)
496 .map(|i| PairedAnswer {
497 case: format!("agree-{i}"),
498 batched: dist("a"),
499 unbatched: dist("a"),
500 })
501 .collect();
502 pairs.extend((0..disagreeing).map(|i| PairedAnswer {
503 case: format!("disagree-{i}"),
504 batched: dist("a"),
505 unbatched: dist("b"),
506 }));
507 Stage2 {
508 pairs,
509 spend: Spend {
510 observed_units: 10,
511 estimated_units: 0,
512 unknown_settlements: 1,
513 },
514 }
515 }
516
517 fn report(stage1: Vec<Stage1Result>, stage2: Option<Stage2>) -> BatchQualification {
518 BatchQualification::new(
519 "SYNTHETIC-adapter",
520 bindings(),
521 BatchingConfig {
522 max_items: 3,
523 window_ms: 2_000,
524 },
525 bound(Agreement::TopLabel, 9, 10),
526 10,
527 stage1,
528 stage2,
529 )
530 .unwrap()
531 }
532
533 #[test]
534 fn the_verdict_follows_both_stages() {
535 let r = report(all_passed(), None);
537 assert_eq!(r.verdict, Verdict::Unknown);
538 assert_eq!(
539 r.stage1.iter().map(|r| r.case).collect::<Vec<_>>(),
540 Stage1Case::ALL
541 );
542 assert_eq!(
544 report(all_passed(), Some(pairs(9, 0))).verdict,
545 Verdict::Unknown
546 );
547 assert_eq!(
549 report(all_passed(), Some(pairs(9, 1))).verdict,
550 Verdict::Qualified
551 );
552 assert_eq!(
554 report(all_passed(), Some(pairs(8, 2))).verdict,
555 Verdict::NotQualified
556 );
557 let mut partial = all_passed();
559 partial.pop();
560 assert_eq!(
561 report(partial.clone(), Some(pairs(10, 0))).verdict,
562 Verdict::Unknown
563 );
564 assert_eq!(
565 report(partial, Some(pairs(0, 10))).verdict,
566 Verdict::NotQualified
567 );
568 let mut failed = all_passed();
570 failed[0] = Stage1Result::failed(failed[0].case, "SYNTHETIC reason");
571 assert_eq!(report(failed, None).verdict, Verdict::NotQualified);
572 }
573
574 #[test]
575 fn pairs_compare_by_the_declared_criterion() {
576 let top = bound(Agreement::TopLabel, 1, 1);
577 let exact = bound(Agreement::Exact, 1, 1);
578 let pair = |batched, unbatched| PairedAnswer {
579 case: "c".into(),
580 batched,
581 unbatched,
582 };
583 let shifted = Answer::Output(RawOutput::Distribution(BTreeMap::from([
584 ("a".to_string(), 0.7),
585 ("SYNTHETIC-other".to_string(), 0.3),
586 ])));
587 assert_eq!(top.agrees(&pair(dist("a"), shifted.clone())), Some(true));
588 assert_eq!(exact.agrees(&pair(dist("a"), shifted)), Some(false));
589 assert_eq!(exact.agrees(&pair(dist("a"), dist("a"))), Some(true));
590 let tie = Answer::Output(RawOutput::Scores(BTreeMap::from([
591 ("a".to_string(), 0.5),
592 ("b".to_string(), 0.5),
593 ])));
594 assert_eq!(top.agrees(&pair(tie.clone(), tie)), Some(false));
595 let label = |l: &str| Answer::Output(RawOutput::Label(l.into()));
596 assert_eq!(top.agrees(&pair(label("a"), dist("a"))), Some(true));
597 let failure = || Answer::Failure {
598 code: RemoteCode::RateLimited,
599 };
600 assert_eq!(top.agrees(&pair(failure(), dist("a"))), Some(false));
601 assert_eq!(top.agrees(&pair(failure(), failure())), None);
602 let mut stage2 = pairs(10, 0);
604 stage2.pairs[0].batched = failure();
605 stage2.pairs[0].unbatched = failure();
606 let r = report(all_passed(), Some(stage2));
607 assert_eq!(r.agreement(), Some((9, 9)));
608 assert_eq!(r.verdict, Verdict::Unknown);
609 }
610
611 #[test]
612 fn a_report_round_trips_and_its_verdict_is_checked() {
613 let r = report(all_passed(), Some(pairs(10, 0)));
614 let bytes = r.record_canonical().unwrap();
615 assert_eq!(BatchQualification::parse(&bytes).unwrap(), r);
616 assert!(r.record_digest().is_ok());
617 let mut forged = report(all_passed(), None);
618 forged.verdict = Verdict::Qualified;
619 let bytes = forged.record_canonical().unwrap();
620 assert_eq!(
621 BatchQualification::parse(&bytes),
622 Err(BatchQualificationError::Verdict {
623 recorded: Verdict::Qualified,
624 computed: Verdict::Unknown
625 })
626 );
627 assert!(matches!(
628 BatchQualification::parse(br#"{"schema":"rustev.other/1"}"#),
629 Err(BatchQualificationError::Parse(_))
630 ));
631 }
632
633 #[test]
634 fn malformed_members_are_refused() {
635 use BatchQualificationError as E;
636 let base = report(all_passed(), Some(pairs(1, 0)));
637 let refused = |f: &dyn Fn(&mut BatchQualification), want: E| {
638 let mut r = base.clone();
639 f(&mut r);
640 assert_eq!(r.check(), Err(want.clone()));
641 let rebuilt = BatchQualification::new(
642 &r.adapter,
643 r.bindings.clone(),
644 r.batching,
645 r.bound,
646 r.min_cases,
647 r.stage1.clone(),
648 r.stage2.clone(),
649 );
650 if !matches!(want, E::Schema | E::Stage1Order) {
651 assert_eq!(rebuilt.unwrap_err(), want);
652 }
653 };
654 refused(&|r| r.schema = "rustev.other/1".into(), E::Schema);
655 refused(&|r| r.adapter = " ".into(), E::Text("the adapter name"));
656 refused(
657 &|r| r.adapter = "a".repeat(257),
658 E::Text("the adapter name"),
659 );
660 refused(&|r| r.batching.max_items = 1, E::Batching);
661 refused(&|r| r.bound.min_agreeing.denominator = 0, E::Bound);
662 refused(&|r| r.bound.min_agreeing.numerator = 11, E::Bound);
663 refused(&|r| r.min_cases = 0, E::Bound);
664 refused(&|r| r.stage1.swap(0, 1), E::Stage1Order);
665 refused(&|r| r.stage1[1] = r.stage1[0].clone(), E::Stage1Order);
666 refused(
667 &|r| {
668 r.stage1[2].outcome = CaseOutcome::Failed {
669 reason: String::new(),
670 }
671 },
672 E::Stage1Reason(Stage1Case::MissingMember),
673 );
674 refused(
675 &|r| r.stage2.as_mut().unwrap().pairs[0].case = String::new(),
676 E::PairCase(0),
677 );
678 refused(
679 &|r| {
680 let s = r.stage2.as_mut().unwrap();
681 s.pairs.push(s.pairs[0].clone());
682 },
683 E::PairCase(1),
684 );
685 refused(
686 &|r| {
687 let s = r.stage2.as_mut().unwrap();
688 s.pairs = (0..=MAX_PAIRS)
689 .map(|i| PairedAnswer {
690 case: i.to_string(),
691 ..s.pairs[0].clone()
692 })
693 .collect();
694 },
695 E::TooManyPairs(MAX_PAIRS + 1),
696 );
697 let mut r = report(all_passed(), Some(pairs(10, 0)));
700 assert_eq!(r.verdict, Verdict::Qualified);
701 r.stage1[1] = r.stage1[0].clone();
702 assert_eq!(r.derived_verdict(), Verdict::Unknown);
703 let long = Stage1Result::failed(Stage1Case::BatchBounds, &"é".repeat(200));
705 assert!(
706 matches!(&long.outcome, CaseOutcome::Failed { reason } if reason.len() <= MAX_TEXT_BYTES)
707 );
708 }
709}