1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4use sha2::Digest;
5
6use crate::aggregation::{AggregatedPredictionBlock, PredictionUnitId};
7use crate::campaign::stable_json_fingerprint;
8use crate::canonical::deserialize_external_contract;
9use crate::conformal_runtime::ConformalCalibrationRef;
10use crate::data::{
11 data_binding_requirement_key, ExternalDataPlanEnvelope, RepresentationCompatibilityReport,
12 RepresentationReplayManifest,
13};
14use crate::error::{DagMlError, Result};
15use crate::ids::{ArtifactId, BundleId, ControllerId, FoldId, NodeId, SampleId, VariantId};
16use crate::metrics::{RegressionMetricReport, RegressionTargetBlock, ScoreSet};
17use crate::oof::{PredictionBlock, PredictionPartition};
18use crate::phase::Phase;
19use crate::plan::ExecutionPlan;
20use crate::policy::PredictionLevel;
21use crate::runtime::{ArtifactBackend, ArtifactRef};
22use crate::selection::SelectionDecision;
23
24pub const METHODS_HPO_RESUME_STATE_SCHEMA_VERSION: u32 = 1;
26
27#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
34#[serde(deny_unknown_fields)]
35pub struct MethodsHpoResumeSelection {
36 pub selection_id: String,
37 pub target_node_id: NodeId,
38 pub producer_port: String,
39 pub metric: String,
40}
41
42#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
44#[serde(deny_unknown_fields)]
45pub struct MethodsHpoResumeProvenance {
46 pub graph_fingerprint: String,
47 pub campaign_fingerprint: String,
48 pub controller_fingerprint: String,
49 pub data_identities_fingerprint: String,
50 pub fold_set_fingerprint: String,
51 pub training_influence_fingerprint: String,
52 pub relation_fingerprint: String,
53 pub selection: MethodsHpoResumeSelection,
54}
55
56#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
59#[serde(deny_unknown_fields)]
60pub struct MethodsHpoCompletedProposal {
61 pub trial_id: i64,
62 pub variant: crate::generation::VariantPlan,
63}
64
65#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
68#[serde(deny_unknown_fields)]
69pub struct MethodsHpoCompletedReport {
70 pub trial_id: i64,
71 pub variant_id: VariantId,
72 pub terminal_state: MethodsHpoCompletedState,
75 pub score: f64,
76 pub report: RegressionMetricReport,
77}
78
79#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
80#[serde(rename_all = "snake_case")]
81pub enum MethodsHpoCompletedState {
82 Completed,
83}
84
85#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
88#[serde(deny_unknown_fields)]
89pub struct MethodsHpoOofAverage {
90 pub predictions: AggregatedPredictionBlock,
91 pub y_true: RegressionTargetBlock,
92}
93
94#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
97#[serde(deny_unknown_fields)]
98pub struct MethodsHpoCandidateEvidence {
99 pub trial_id: i64,
100 pub variant_id: VariantId,
101 #[serde(default, skip_serializing_if = "Option::is_none")]
102 pub variant_label: Option<String>,
103 pub predictions: Vec<PredictionBlock>,
104 pub regression_targets: Vec<RegressionTargetBlock>,
105 pub oof_average: MethodsHpoOofAverage,
106 pub lineage: Vec<crate::runtime::LineageRecord>,
107}
108
109#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
112#[serde(deny_unknown_fields)]
113pub struct MethodsHpoNativeIncumbent {
114 pub trial_id: i64,
115 pub score: f64,
116 pub metric: String,
117 pub direction: crate::hpo::HpoDirection,
118 pub variant_id: VariantId,
119}
120
121#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
122#[serde(deny_unknown_fields)]
123pub struct MethodsHpoTerminalEvidence {
124 pub trial: crate::hpo::HpoTrial,
125 pub variant_id: Option<VariantId>,
126}
127
128#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
135#[serde(deny_unknown_fields)]
136pub struct MethodsHpoResumeState {
137 pub schema_version: u32,
138 pub checkpoint: crate::hpo::N4moptCheckpointArtifact,
139 pub provenance: MethodsHpoResumeProvenance,
140 pub operation_id: String,
142 pub controller_id: ControllerId,
143 pub target_node_id: NodeId,
144 pub incumbent: MethodsHpoNativeIncumbent,
145 pub trial_history_len: u32,
148 pub terminal_trials: Vec<MethodsHpoTerminalEvidence>,
149 pub completed_proposals: Vec<MethodsHpoCompletedProposal>,
150 pub completed_reports: Vec<MethodsHpoCompletedReport>,
151 pub candidates: Vec<MethodsHpoCandidateEvidence>,
152}
153
154impl MethodsHpoResumeState {
155 pub fn from_runtime_checkpoint(
160 checkpoint: crate::runtime::RuntimeHpoCheckpointResult,
161 selection: MethodsHpoResumeSelection,
162 candidates: Vec<crate::runtime::RuntimeHpoCandidateEvaluation>,
163 incumbent: crate::runtime::RuntimeHpoIncumbent,
164 terminal_trials: Vec<crate::runtime::RuntimeHpoTerminalSnapshot>,
165 ) -> Result<Self> {
166 Self::from_runtime_checkpoint_with_previous(
167 checkpoint,
168 selection,
169 candidates,
170 incumbent,
171 terminal_trials,
172 None,
173 )
174 }
175
176 pub(crate) fn from_runtime_checkpoint_with_previous(
182 checkpoint: crate::runtime::RuntimeHpoCheckpointResult,
183 selection: MethodsHpoResumeSelection,
184 candidates: Vec<crate::runtime::RuntimeHpoCandidateEvaluation>,
185 incumbent: crate::runtime::RuntimeHpoIncumbent,
186 terminal_trials: Vec<crate::runtime::RuntimeHpoTerminalSnapshot>,
187 previous: Option<Self>,
188 ) -> Result<Self> {
189 let candidate_scores = candidates
190 .iter()
191 .map(|candidate| (candidate.proposal.trial_id, candidate.score))
192 .collect::<BTreeMap<_, _>>();
193 let completed_proposals = checkpoint
194 .completed_proposals
195 .iter()
196 .map(|proposal| MethodsHpoCompletedProposal {
197 trial_id: proposal.trial_id,
198 variant: proposal.variant.clone(),
199 })
200 .collect();
201 let completed_reports = checkpoint
202 .completed_reports
203 .iter()
204 .map(|completed| {
205 let metric = selection.metric.as_str();
206 if !completed.report.metrics.contains_key(metric) {
207 return Err(DagMlError::RuntimeValidation(format!(
208 "Methods HPO completed trial {} has no `{metric}` OOF metric",
209 completed.trial_id
210 )));
211 }
212 let score = candidate_scores
213 .get(&completed.trial_id)
214 .copied()
215 .ok_or_else(|| {
216 DagMlError::RuntimeValidation(format!(
217 "Methods HPO completed trial {} has no candidate OOF evidence",
218 completed.trial_id
219 ))
220 })?;
221 Ok(MethodsHpoCompletedReport {
222 trial_id: completed.trial_id,
223 variant_id: completed.variant_id.clone(),
224 terminal_state: MethodsHpoCompletedState::Completed,
225 score,
226 report: completed.report.clone(),
227 })
228 })
229 .collect::<Result<Vec<_>>>()?;
230 let candidates = candidates
231 .into_iter()
232 .map(|candidate| {
233 let average = candidate
234 .validation_predictions
235 .oof_average
236 .ok_or_else(|| {
237 DagMlError::RuntimeValidation(format!(
238 "Methods HPO trial {} has no cross-fold OOF average",
239 candidate.proposal.trial_id
240 ))
241 })?;
242 Ok(MethodsHpoCandidateEvidence {
243 trial_id: candidate.proposal.trial_id,
244 variant_id: candidate.proposal.variant.variant_id.clone(),
245 variant_label: candidate.validation_predictions.variant_label,
246 predictions: candidate.validation_predictions.predictions,
247 regression_targets: candidate.validation_predictions.regression_targets,
248 oof_average: MethodsHpoOofAverage {
249 predictions: average.predictions,
250 y_true: average.y_true,
251 },
252 lineage: candidate.lineage,
253 })
254 })
255 .collect::<Result<Vec<_>>>()?;
256 let mut state = Self {
257 schema_version: METHODS_HPO_RESUME_STATE_SCHEMA_VERSION,
258 checkpoint: checkpoint.artifact,
259 provenance: MethodsHpoResumeProvenance {
260 graph_fingerprint: checkpoint.provenance.graph_fingerprint,
261 campaign_fingerprint: checkpoint.provenance.campaign_fingerprint,
262 controller_fingerprint: checkpoint.provenance.controller_fingerprint,
263 data_identities_fingerprint: checkpoint.provenance.data_identities_fingerprint,
264 fold_set_fingerprint: checkpoint.provenance.fold_set_fingerprint.ok_or_else(
265 || {
266 DagMlError::RuntimeValidation(
267 "Methods HPO resume requires an attested fold-set fingerprint"
268 .to_string(),
269 )
270 },
271 )?,
272 training_influence_fingerprint: checkpoint
273 .provenance
274 .training_influence_fingerprint,
275 relation_fingerprint: checkpoint.provenance.relation_fingerprint,
276 selection,
277 },
278 operation_id: checkpoint.operation_id,
279 controller_id: checkpoint.controller_id,
280 target_node_id: checkpoint.target_node_id,
281 incumbent: MethodsHpoNativeIncumbent {
282 trial_id: incumbent.trial_id,
283 score: incumbent.score,
284 metric: incumbent.metric,
285 direction: incumbent.direction,
286 variant_id: incumbent.variant_id,
287 },
288 trial_history_len: checkpoint.trial_history_len,
289 terminal_trials: terminal_trials
290 .into_iter()
291 .map(|snapshot| MethodsHpoTerminalEvidence {
292 trial: snapshot.trial,
293 variant_id: snapshot.variant_id,
294 })
295 .collect(),
296 completed_proposals,
297 completed_reports,
298 candidates,
299 };
300 if let Some(previous) = previous {
301 previous.validate()?;
302 if previous.provenance != state.provenance
303 || previous.operation_id != state.operation_id
304 || previous.controller_id != state.controller_id
305 || previous.target_node_id != state.target_node_id
306 || previous.checkpoint.binding != state.checkpoint.binding
307 {
308 return Err(DagMlError::RuntimeValidation(
309 "native Methods HPO resumed campaign state has incompatible provenance"
310 .to_string(),
311 ));
312 }
313 if previous.trial_history_len > state.trial_history_len
314 || !state
315 .terminal_trials
316 .starts_with(previous.terminal_trials.as_slice())
317 {
318 return Err(DagMlError::RuntimeValidation(
319 "native Methods HPO resumed campaign does not preserve its prior terminal history"
320 .to_string(),
321 ));
322 }
323 state
324 .completed_proposals
325 .extend(previous.completed_proposals);
326 state.completed_reports.extend(previous.completed_reports);
327 state.candidates.extend(previous.candidates);
328 state
329 .completed_proposals
330 .sort_by_key(|proposal| proposal.trial_id);
331 state
332 .completed_reports
333 .sort_by_key(|report| report.trial_id);
334 state.candidates.sort_by_key(|candidate| candidate.trial_id);
335 }
336 state.validate()?;
337 Ok(state)
338 }
339
340 pub fn validate(&self) -> Result<()> {
342 if self.schema_version != METHODS_HPO_RESUME_STATE_SCHEMA_VERSION {
343 return Err(DagMlError::RuntimeValidation(format!(
344 "Methods HPO resume state has unsupported schema_version {}",
345 self.schema_version
346 )));
347 }
348 self.checkpoint.validate().map_err(|error| {
349 DagMlError::RuntimeValidation(format!(
350 "Methods HPO resume checkpoint is invalid: {error}"
351 ))
352 })?;
353 if self.checkpoint.binding.controller_id.trim().is_empty()
354 || self.checkpoint.binding.study_id.trim().is_empty()
355 {
356 return Err(DagMlError::RuntimeValidation(
357 "Methods HPO resume checkpoint binding has an empty controller or study identity"
358 .to_string(),
359 ));
360 }
361 if self.operation_id.trim().is_empty()
362 || self.trial_history_len == 0
363 || self.completed_proposals.len() > self.trial_history_len as usize
364 || self.controller_id.as_str() != self.checkpoint.binding.controller_id
365 || self.target_node_id != self.provenance.selection.target_node_id
366 || self.incumbent.metric != self.provenance.selection.metric
367 || !self.incumbent.score.is_finite()
368 {
369 return Err(DagMlError::RuntimeValidation(
370 "Methods HPO resume campaign operation or incumbent identity is invalid"
371 .to_string(),
372 ));
373 }
374 if self.terminal_trials.len() != self.trial_history_len as usize
375 || self
376 .terminal_trials
377 .windows(2)
378 .any(|pair| pair[0].trial.id >= pair[1].trial.id)
379 || self.terminal_trials.iter().any(|evidence| {
380 !matches!(
381 evidence.trial.status,
382 crate::hpo::HpoTrialStatus::Completed
383 | crate::hpo::HpoTrialStatus::Pruned
384 | crate::hpo::HpoTrialStatus::Failed
385 )
386 })
387 {
388 return Err(DagMlError::RuntimeValidation(
389 "Methods HPO terminal trial evidence must exactly cover a strictly ordered terminal native history"
390 .to_string(),
391 ));
392 }
393 for (label, ids) in [
394 (
395 "completed proposals",
396 self.completed_proposals
397 .iter()
398 .map(|item| item.trial_id)
399 .collect::<Vec<_>>(),
400 ),
401 (
402 "completed reports",
403 self.completed_reports
404 .iter()
405 .map(|item| item.trial_id)
406 .collect::<Vec<_>>(),
407 ),
408 (
409 "candidate evidence",
410 self.candidates
411 .iter()
412 .map(|item| item.trial_id)
413 .collect::<Vec<_>>(),
414 ),
415 ] {
416 if ids.windows(2).any(|pair| pair[0] >= pair[1]) {
417 return Err(DagMlError::RuntimeValidation(format!(
418 "Methods HPO resume {label} must be strictly sorted by trial_id"
419 )));
420 }
421 }
422 validate_resume_sha256(
423 "checkpoint search-space",
424 &self.checkpoint.binding.search_space_fingerprint,
425 )?;
426 validate_resume_sha256(
427 "checkpoint optimizer",
428 &self.checkpoint.binding.optimizer_fingerprint,
429 )?;
430 for (label, value) in [
431 ("graph", self.provenance.graph_fingerprint.as_str()),
432 ("campaign", self.provenance.campaign_fingerprint.as_str()),
433 (
434 "controller",
435 self.provenance.controller_fingerprint.as_str(),
436 ),
437 (
438 "data identities",
439 self.provenance.data_identities_fingerprint.as_str(),
440 ),
441 ("fold set", self.provenance.fold_set_fingerprint.as_str()),
442 (
443 "training influence",
444 self.provenance.training_influence_fingerprint.as_str(),
445 ),
446 ("relation", self.provenance.relation_fingerprint.as_str()),
447 ] {
448 validate_resume_sha256(label, value)?;
449 }
450 for (label, value) in [
451 (
452 "selection id",
453 self.provenance.selection.selection_id.as_str(),
454 ),
455 (
456 "selection producer port",
457 self.provenance.selection.producer_port.as_str(),
458 ),
459 (
460 "selection metric",
461 self.provenance.selection.metric.as_str(),
462 ),
463 ] {
464 if value.trim().is_empty() {
465 return Err(DagMlError::RuntimeValidation(format!(
466 "Methods HPO resume {label} must not be empty"
467 )));
468 }
469 }
470 let mut proposals = BTreeMap::new();
471 for proposal in &self.completed_proposals {
472 proposal.variant.validate()?;
473 if proposals
474 .insert(proposal.trial_id, &proposal.variant)
475 .is_some()
476 {
477 return Err(DagMlError::RuntimeValidation(format!(
478 "Methods HPO resume has duplicate completed proposal for trial {}",
479 proposal.trial_id
480 )));
481 }
482 }
483 if proposals.is_empty() {
484 return Err(DagMlError::RuntimeValidation(
485 "Methods HPO resume requires at least one completed proposal".to_string(),
486 ));
487 }
488 let terminal_by_trial = self
489 .terminal_trials
490 .iter()
491 .map(|evidence| (evidence.trial.id, evidence))
492 .collect::<BTreeMap<_, _>>();
493 let completed_terminal_ids = self
494 .terminal_trials
495 .iter()
496 .filter(|evidence| evidence.trial.status == crate::hpo::HpoTrialStatus::Completed)
497 .map(|evidence| evidence.trial.id)
498 .collect::<BTreeSet<_>>();
499 if completed_terminal_ids != proposals.keys().copied().collect()
500 || proposals.iter().any(|(trial_id, variant)| {
501 terminal_by_trial
502 .get(trial_id)
503 .and_then(|evidence| evidence.variant_id.as_ref())
504 != Some(&variant.variant_id)
505 })
506 {
507 return Err(DagMlError::RuntimeValidation(
508 "Methods HPO completed proposal evidence must exactly match terminal native trials"
509 .to_string(),
510 ));
511 }
512 let mut reports = BTreeSet::new();
513 for completed in &self.completed_reports {
514 let Some(variant) = proposals.get(&completed.trial_id) else {
515 return Err(DagMlError::RuntimeValidation(format!(
516 "Methods HPO resume report for trial {} is orphaned",
517 completed.trial_id
518 )));
519 };
520 if !matches!(
521 completed.terminal_state,
522 MethodsHpoCompletedState::Completed
523 ) || variant.variant_id != completed.variant_id
524 || !completed.score.is_finite()
525 {
526 return Err(DagMlError::RuntimeValidation(format!(
527 "Methods HPO resume report for trial {} does not match its proposal",
528 completed.trial_id
529 )));
530 }
531 let terminal = terminal_by_trial[&completed.trial_id];
532 if terminal.trial.score.map(f64::to_bits) != Some(completed.score.to_bits())
533 || terminal
534 .trial
535 .intermediates
536 .iter()
537 .any(|value| !value.score.is_finite())
538 {
539 return Err(DagMlError::RuntimeValidation(format!(
540 "Methods HPO terminal trial {} score/intermediate evidence is inconsistent",
541 completed.trial_id
542 )));
543 }
544 validate_resume_report(completed, &self.provenance.selection)?;
545 if !reports.insert(completed.trial_id) {
546 return Err(DagMlError::RuntimeValidation(format!(
547 "Methods HPO resume has duplicate completed report for trial {}",
548 completed.trial_id
549 )));
550 }
551 }
552 if reports != proposals.keys().copied().collect() {
553 return Err(DagMlError::RuntimeValidation(
554 "Methods HPO resume completed reports must exactly cover completed proposals"
555 .to_string(),
556 ));
557 }
558 let incumbent = self
559 .completed_reports
560 .iter()
561 .find(|report| report.trial_id == self.incumbent.trial_id)
562 .ok_or_else(|| {
563 DagMlError::RuntimeValidation(
564 "Methods HPO resume incumbent is not a completed scheduler trial".to_string(),
565 )
566 })?;
567 if incumbent.variant_id != self.incumbent.variant_id
568 || incumbent.score.to_bits() != self.incumbent.score.to_bits()
569 {
570 return Err(DagMlError::RuntimeValidation(
571 "Methods HPO resume incumbent does not exactly match its completed report"
572 .to_string(),
573 ));
574 }
575 let mut evidence = BTreeSet::new();
576 for candidate in &self.candidates {
577 let Some(variant) = proposals.get(&candidate.trial_id) else {
578 return Err(DagMlError::RuntimeValidation(format!(
579 "Methods HPO candidate evidence for trial {} is orphaned",
580 candidate.trial_id
581 )));
582 };
583 if variant.variant_id != candidate.variant_id {
584 return Err(DagMlError::RuntimeValidation(format!(
585 "Methods HPO candidate evidence for trial {} does not match its proposal",
586 candidate.trial_id
587 )));
588 }
589 validate_resume_candidate(candidate, &self.provenance.selection)?;
590 if !evidence.insert(candidate.trial_id) {
591 return Err(DagMlError::RuntimeValidation(format!(
592 "Methods HPO resume has duplicate candidate evidence for trial {}",
593 candidate.trial_id
594 )));
595 }
596 }
597 if evidence != proposals.keys().copied().collect() {
598 return Err(DagMlError::RuntimeValidation(
599 "Methods HPO candidate evidence must exactly cover completed proposals".to_string(),
600 ));
601 }
602 Ok(())
603 }
604
605 pub fn validate_against_plan(&self, plan: &ExecutionPlan) -> Result<()> {
607 self.validate()?;
608 if self.provenance.graph_fingerprint != plan.graph_fingerprint
609 || self.provenance.campaign_fingerprint
610 != crate::hpo::campaign_provenance_fingerprint(&plan.campaign)?
611 || self.provenance.controller_fingerprint != plan.controller_fingerprint
612 || plan
613 .fold_set
614 .as_ref()
615 .map(stable_json_fingerprint)
616 .transpose()?
617 .as_deref()
618 != Some(self.provenance.fold_set_fingerprint.as_str())
619 {
620 return Err(DagMlError::RuntimeValidation(
621 "Methods HPO resume provenance does not match execution plan".to_string(),
622 ));
623 }
624 let target = plan.node_plans.get(&self.target_node_id).ok_or_else(|| {
625 DagMlError::RuntimeValidation(
626 "Methods HPO resume target node is absent from execution plan".to_string(),
627 )
628 })?;
629 if self.operation_id.trim().is_empty()
630 || target.kind != crate::graph::NodeKind::Model
631 || self.controller_id.as_str() != self.checkpoint.binding.controller_id
632 || self.target_node_id != self.provenance.selection.target_node_id
633 {
634 return Err(DagMlError::RuntimeValidation(
635 "Methods HPO resume campaign/target/checkpoint binding does not match execution plan"
636 .to_string(),
637 ));
638 }
639 Ok(())
640 }
641}
642
643fn validate_resume_sha256(label: &str, value: &str) -> Result<()> {
644 if value.len() != 64
645 || !value
646 .bytes()
647 .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
648 {
649 return Err(DagMlError::RuntimeValidation(format!(
650 "Methods HPO resume {label} fingerprint must be lowercase SHA-256"
651 )));
652 }
653 Ok(())
654}
655
656fn validate_resume_report(
657 completed: &MethodsHpoCompletedReport,
658 selection: &MethodsHpoResumeSelection,
659) -> Result<()> {
660 completed.report.validate()?;
661 let report = &completed.report;
662 let metric = report.metrics.get(&selection.metric).ok_or_else(|| {
663 DagMlError::RuntimeValidation(format!(
664 "Methods HPO report for trial {} omits selection metric `{}`",
665 completed.trial_id, selection.metric
666 ))
667 })?;
668 if report.variant_id.as_ref() != Some(&completed.variant_id)
669 || report.producer_node != selection.target_node_id
670 || report.producer_port.as_deref() != Some(selection.producer_port.as_str())
671 || report.partition != PredictionPartition::Validation
672 || report.fold_id.as_ref().map(FoldId::as_str) != Some("avg")
673 || report.level != PredictionLevel::Sample
674 || (metric - completed.score).abs() > 1e-12
675 {
676 return Err(DagMlError::RuntimeValidation(format!(
677 "Methods HPO report for trial {} must be the exact sample-level validation OOF average",
678 completed.trial_id
679 )));
680 }
681 Ok(())
682}
683
684fn validate_resume_candidate(
685 candidate: &MethodsHpoCandidateEvidence,
686 selection: &MethodsHpoResumeSelection,
687) -> Result<()> {
688 if candidate.predictions.is_empty()
689 || candidate.predictions.len() != candidate.regression_targets.len()
690 || candidate.lineage.is_empty()
691 {
692 return Err(DagMlError::RuntimeValidation(
693 "Methods HPO candidate evidence requires paired fold OOF predictions, targets and lineage"
694 .to_string(),
695 ));
696 }
697 for (prediction, target) in candidate
698 .predictions
699 .iter()
700 .zip(&candidate.regression_targets)
701 {
702 prediction.validate_shape()?;
703 target.validate_shape()?;
704 if prediction.partition != PredictionPartition::Validation
705 || prediction.producer_node != selection.target_node_id
706 || prediction.producer_port.as_deref() != Some(selection.producer_port.as_str())
707 || prediction.sample_ids
708 != target
709 .unit_ids
710 .iter()
711 .filter_map(|unit| match unit {
712 PredictionUnitId::Sample(id) => Some(id.clone()),
713 _ => None,
714 })
715 .collect::<Vec<_>>()
716 {
717 return Err(DagMlError::RuntimeValidation(
718 "Methods HPO candidate fold evidence is not exact validation OOF data for selection"
719 .to_string(),
720 ));
721 }
722 }
723 candidate.oof_average.predictions.validate_shape()?;
724 candidate.oof_average.y_true.validate_shape()?;
725 if candidate.oof_average.predictions.producer_node != selection.target_node_id
726 || candidate.oof_average.predictions.producer_port.as_deref()
727 != Some(selection.producer_port.as_str())
728 || candidate.oof_average.predictions.partition != PredictionPartition::Validation
729 || candidate
730 .oof_average
731 .predictions
732 .fold_id
733 .as_ref()
734 .map(FoldId::as_str)
735 != Some("avg")
736 {
737 return Err(DagMlError::RuntimeValidation(
738 "Methods HPO candidate OOF average does not match selection output".to_string(),
739 ));
740 }
741 let mut lineage_ids = BTreeSet::new();
742 let mut has_target = false;
743 for record in &candidate.lineage {
744 record.validate()?;
745 if !lineage_ids.insert(record.record_id.clone()) {
746 return Err(DagMlError::RuntimeValidation(
747 "Methods HPO candidate evidence has duplicate lineage record ids".to_string(),
748 ));
749 }
750 has_target |= record.node_id == selection.target_node_id
751 && record.variant_id.as_ref() == Some(&candidate.variant_id);
752 }
753 if !has_target {
754 return Err(DagMlError::RuntimeValidation(
755 "Methods HPO candidate lineage does not include its selected target variant"
756 .to_string(),
757 ));
758 }
759 Ok(())
760}
761
762pub const EXECUTION_BUNDLE_SCHEMA_VERSION: u32 = 2;
763pub const PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION: u32 = 2;
764pub const LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION: u32 = 1;
765pub const LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION: u32 = 1;
766pub const LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT: &str = "dag-ml-json-prediction-blocks-v1";
767pub const BUNDLE_PREDICTION_CACHE_FORMAT: &str = "dag-ml-json-prediction-blocks-v2";
768
769pub const MIN_READABLE_EXECUTION_BUNDLE_SCHEMA_VERSION: u32 = 1;
770pub const MIN_WRITABLE_EXECUTION_BUNDLE_SCHEMA_VERSION: u32 = 2;
771pub const MIN_READABLE_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION: u32 = 1;
772pub const MIN_WRITABLE_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION: u32 = 2;
773
774fn default_execution_bundle_schema_version() -> u32 {
775 LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION
776}
777
778fn default_prediction_cache_payload_schema_version() -> u32 {
779 LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION
780}
781
782fn default_prediction_level() -> PredictionLevel {
783 PredictionLevel::Sample
784}
785
786fn supported_prediction_cache_format(format: &str) -> bool {
787 matches!(
788 format,
789 LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT | BUNDLE_PREDICTION_CACHE_FORMAT
790 )
791}
792
793fn prediction_cache_schema_version_for_format(format: &str, owner: &str) -> Result<u32> {
794 match format {
795 LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT => Ok(LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION),
796 BUNDLE_PREDICTION_CACHE_FORMAT => Ok(PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION),
797 _ => Err(DagMlError::RuntimeValidation(format!(
798 "{owner} uses unsupported cache format `{format}`"
799 ))),
800 }
801}
802
803fn expected_prediction_cache_format_for_schema_version(
804 schema_version: u32,
805 owner: &str,
806) -> Result<&'static str> {
807 match schema_version {
808 LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION => Ok(LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT),
809 EXECUTION_BUNDLE_SCHEMA_VERSION => Ok(BUNDLE_PREDICTION_CACHE_FORMAT),
810 _ => Err(DagMlError::RuntimeValidation(format!(
811 "{owner} uses unsupported cache family schema_version {schema_version}"
812 ))),
813 }
814}
815
816fn validate_prediction_cache_format_for_schema_version(
817 format: &str,
818 schema_version: u32,
819 owner: &str,
820) -> Result<()> {
821 let expected = expected_prediction_cache_format_for_schema_version(schema_version, owner)?;
822 if format != expected {
823 return Err(DagMlError::RuntimeValidation(format!(
824 "{owner} uses cache format `{format}` but schema_version {schema_version} requires `{expected}`"
825 )));
826 }
827 Ok(())
828}
829
830fn validate_prediction_block_port_family(
831 producer_port: &Option<String>,
832 schema_version: u32,
833 owner: &str,
834) -> Result<()> {
835 match (schema_version, producer_port.as_deref()) {
836 (LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION, Some(_)) => Err(
837 DagMlError::RuntimeValidation(format!("{owner} is V1 but carries producer_port")),
838 ),
839 (PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION, Some(port)) if port.trim().is_empty() => {
840 Err(DagMlError::RuntimeValidation(format!(
841 "{owner} is V2 but carries an empty producer_port"
842 )))
843 }
844 (PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION, None) => Err(DagMlError::RuntimeValidation(
845 format!("{owner} is V2 and requires producer_port"),
846 )),
847 _ => Ok(()),
848 }
849}
850
851fn validate_prediction_cache_payload_block_family(
852 payload: &BundlePredictionCachePayload,
853 schema_version: u32,
854) -> Result<()> {
855 for block in &payload.blocks {
856 validate_prediction_block_port_family(
857 &block.producer_port,
858 schema_version,
859 &format!(
860 "prediction cache payload `{}` block for node `{}`",
861 payload.cache_id, block.producer_node
862 ),
863 )?;
864 }
865 for block in &payload.aggregated_blocks {
866 validate_prediction_block_port_family(
867 &block.producer_port,
868 schema_version,
869 &format!(
870 "prediction cache payload `{}` aggregated block for node `{}`",
871 payload.cache_id, block.producer_node
872 ),
873 )?;
874 }
875 Ok(())
876}
877
878#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
879pub struct SchemaMigrationPolicy {
880 pub artifact: String,
881 pub current_version: u32,
882 pub min_readable_version: u32,
883 pub min_writable_version: u32,
884 #[serde(default)]
885 pub automatic_migrations: BTreeMap<u32, u32>,
886}
887
888impl SchemaMigrationPolicy {
889 pub fn validate(&self) -> Result<()> {
890 validate_non_empty("schema migration artifact", &self.artifact)?;
891 if self.current_version == 0
892 || self.min_readable_version == 0
893 || self.min_writable_version == 0
894 {
895 return Err(DagMlError::RuntimeValidation(format!(
896 "schema migration policy `{}` has zero version boundary",
897 self.artifact
898 )));
899 }
900 if self.min_readable_version > self.current_version {
901 return Err(DagMlError::RuntimeValidation(format!(
902 "schema migration policy `{}` min_readable_version exceeds current_version",
903 self.artifact
904 )));
905 }
906 if self.min_writable_version > self.current_version {
907 return Err(DagMlError::RuntimeValidation(format!(
908 "schema migration policy `{}` min_writable_version exceeds current_version",
909 self.artifact
910 )));
911 }
912 for (from, to) in &self.automatic_migrations {
913 if *from == 0 || *to == 0 {
914 return Err(DagMlError::RuntimeValidation(format!(
915 "schema migration policy `{}` contains a zero migration version",
916 self.artifact
917 )));
918 }
919 if from == to {
920 return Err(DagMlError::RuntimeValidation(format!(
921 "schema migration policy `{}` contains a no-op migration {from}->{to}",
922 self.artifact
923 )));
924 }
925 if *to > self.current_version {
926 return Err(DagMlError::RuntimeValidation(format!(
927 "schema migration policy `{}` migrates to unsupported future version {to}",
928 self.artifact
929 )));
930 }
931 }
932 Ok(())
933 }
934
935 pub fn validate_read_version(&self, version: u32, owner: &str) -> Result<()> {
936 self.validate()?;
937 if version < self.min_readable_version {
938 return Err(DagMlError::RuntimeValidation(format!(
939 "{owner} uses schema_version {version}, below minimum readable {} for {}",
940 self.min_readable_version, self.artifact
941 )));
942 }
943 if version > self.current_version {
944 return Err(DagMlError::RuntimeValidation(format!(
945 "{owner} uses future schema_version {version}, current readable {} for {}",
946 self.current_version, self.artifact
947 )));
948 }
949 if version != self.current_version && !self.automatic_migrations.contains_key(&version) {
950 return Err(DagMlError::RuntimeValidation(format!(
951 "{owner} uses schema_version {version}, but {} declares no automatic migration to current version {}",
952 self.artifact, self.current_version
953 )));
954 }
955 Ok(())
956 }
957}
958
959pub fn execution_bundle_schema_migration_policy() -> SchemaMigrationPolicy {
960 SchemaMigrationPolicy {
961 artifact: "execution_bundle".to_string(),
962 current_version: EXECUTION_BUNDLE_SCHEMA_VERSION,
963 min_readable_version: MIN_READABLE_EXECUTION_BUNDLE_SCHEMA_VERSION,
964 min_writable_version: MIN_WRITABLE_EXECUTION_BUNDLE_SCHEMA_VERSION,
965 automatic_migrations: BTreeMap::from([(
966 LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION,
967 EXECUTION_BUNDLE_SCHEMA_VERSION,
968 )]),
969 }
970}
971
972pub fn prediction_cache_payload_schema_migration_policy() -> SchemaMigrationPolicy {
973 SchemaMigrationPolicy {
974 artifact: "prediction_cache_payload".to_string(),
975 current_version: PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
976 min_readable_version: MIN_READABLE_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
977 min_writable_version: MIN_WRITABLE_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
978 automatic_migrations: BTreeMap::from([(
979 LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
980 PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
981 )]),
982 }
983}
984
985#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
986pub struct BundleDataRequirement {
987 pub node_id: NodeId,
988 pub input_name: String,
989 pub schema_fingerprint: String,
990 pub plan_fingerprint: String,
991 #[serde(default)]
992 pub relation_fingerprint: Option<String>,
993 pub output_representation: String,
994 #[serde(default)]
995 pub feature_set_id: Option<String>,
996 #[serde(default, skip_serializing_if = "Option::is_none")]
997 pub representation_replay_manifest: Option<RepresentationReplayManifest>,
998 #[serde(default, skip_serializing_if = "Option::is_none")]
999 pub representation_compatibility: Option<RepresentationCompatibilityReport>,
1000}
1001
1002impl BundleDataRequirement {
1003 pub fn key(&self) -> String {
1004 data_binding_requirement_key(&self.node_id, &self.input_name)
1005 }
1006
1007 fn matches_plan_requirement(&self, expected: &Self) -> bool {
1008 self.node_id == expected.node_id
1009 && self.input_name == expected.input_name
1010 && self.schema_fingerprint == expected.schema_fingerprint
1011 && self.plan_fingerprint == expected.plan_fingerprint
1012 && self.relation_fingerprint == expected.relation_fingerprint
1013 && self.output_representation == expected.output_representation
1014 && self.feature_set_id == expected.feature_set_id
1015 }
1016
1017 pub fn validate(&self) -> Result<()> {
1018 if self.input_name.trim().is_empty() {
1019 return Err(DagMlError::CampaignValidation(format!(
1020 "bundle data requirement for `{}` has empty input_name",
1021 self.node_id
1022 )));
1023 }
1024 validate_fingerprint("schema", &self.schema_fingerprint)?;
1025 validate_fingerprint("plan", &self.plan_fingerprint)?;
1026 if let Some(relation_fingerprint) = &self.relation_fingerprint {
1027 validate_fingerprint("relation", relation_fingerprint)?;
1028 }
1029 if let Some(replay_manifest) = &self.representation_replay_manifest {
1030 replay_manifest.validate()?;
1031 if let (Some(requirement), Some(manifest)) = (
1032 self.relation_fingerprint.as_deref(),
1033 replay_manifest.relation_fingerprint.as_deref(),
1034 ) {
1035 if requirement != manifest {
1036 return Err(DagMlError::CampaignValidation(format!(
1037 "bundle data requirement `{}` relation_fingerprint does not match representation replay manifest",
1038 self.key()
1039 )));
1040 }
1041 }
1042 }
1043 if let Some(report) = &self.representation_compatibility {
1044 report.validate()?;
1045 }
1046 if self.output_representation.trim().is_empty() {
1047 return Err(DagMlError::CampaignValidation(format!(
1048 "bundle data requirement `{}` has empty output representation",
1049 self.key()
1050 )));
1051 }
1052 if let Some(feature_set_id) = &self.feature_set_id {
1053 if feature_set_id.trim().is_empty() {
1054 return Err(DagMlError::CampaignValidation(format!(
1055 "bundle data requirement `{}` has empty feature_set_id",
1056 self.key()
1057 )));
1058 }
1059 }
1060 Ok(())
1061 }
1062}
1063
1064#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1065pub struct BundlePredictionRequirement {
1066 pub producer_node: NodeId,
1067 pub source_port: String,
1068 pub consumer_node: NodeId,
1069 pub target_port: String,
1070 pub partition: PredictionPartition,
1071 #[serde(default = "default_prediction_level")]
1072 pub prediction_level: PredictionLevel,
1073 #[serde(default)]
1074 pub fold_ids: Vec<FoldId>,
1075 #[serde(default, skip_serializing_if = "Vec::is_empty")]
1076 pub unit_ids: Vec<PredictionUnitId>,
1077 #[serde(default)]
1078 pub sample_ids: Vec<SampleId>,
1079 pub prediction_width: usize,
1080 pub target_names: Vec<String>,
1081}
1082
1083impl BundlePredictionRequirement {
1084 pub fn key(&self) -> String {
1085 bundle_prediction_requirement_key(
1086 &self.producer_node,
1087 &self.source_port,
1088 &self.consumer_node,
1089 &self.target_port,
1090 )
1091 }
1092
1093 pub fn validate(&self) -> Result<()> {
1094 validate_non_empty("source_port", &self.source_port)?;
1095 validate_non_empty("target_port", &self.target_port)?;
1096 if self.partition != PredictionPartition::Validation {
1097 return Err(DagMlError::RuntimeValidation(format!(
1098 "bundle prediction requirement `{}` must use validation OOF predictions",
1099 self.key()
1100 )));
1101 }
1102 validate_unique_ids("fold id", &self.fold_ids)?;
1103 validate_prediction_requirement_units(self)?;
1104 if self.prediction_width == 0 {
1105 return Err(DagMlError::RuntimeValidation(format!(
1106 "bundle prediction requirement `{}` has zero prediction width",
1107 self.key()
1108 )));
1109 }
1110 if self.target_names.len() != self.prediction_width {
1111 return Err(DagMlError::RuntimeValidation(format!(
1112 "bundle prediction requirement `{}` target name count does not match prediction width",
1113 self.key()
1114 )));
1115 }
1116 for target_name in &self.target_names {
1117 validate_non_empty("target_name", target_name)?;
1118 }
1119 Ok(())
1120 }
1121}
1122
1123pub fn bundle_prediction_requirement_key(
1124 producer_node: &NodeId,
1125 source_port: &str,
1126 consumer_node: &NodeId,
1127 target_port: &str,
1128) -> String {
1129 format!("{producer_node}.{source_port}->{consumer_node}.{target_port}")
1130}
1131
1132#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1133pub struct BundlePredictionBlockCacheRecord {
1134 #[serde(default)]
1135 pub prediction_id: Option<String>,
1136 #[serde(default)]
1137 pub fold_id: Option<FoldId>,
1138 #[serde(default = "default_prediction_level")]
1139 pub prediction_level: PredictionLevel,
1140 pub row_count: usize,
1141 #[serde(default, skip_serializing_if = "Vec::is_empty")]
1142 pub unit_ids: Vec<PredictionUnitId>,
1143 #[serde(default)]
1144 pub sample_ids: Vec<SampleId>,
1145 pub content_fingerprint: String,
1146}
1147
1148impl BundlePredictionBlockCacheRecord {
1149 pub fn validate(&self) -> Result<()> {
1150 if let Some(prediction_id) = &self.prediction_id {
1151 validate_non_empty("prediction_id", prediction_id)?;
1152 }
1153 if self.row_count == 0 {
1154 return Err(DagMlError::RuntimeValidation(
1155 "prediction block cache record has zero rows".to_string(),
1156 ));
1157 }
1158 validate_prediction_cache_block_record_units(self)?;
1159 validate_fingerprint("prediction block cache content", &self.content_fingerprint)
1160 }
1161}
1162
1163#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1164pub struct BundlePredictionCacheRecord {
1165 pub requirement_key: String,
1166 pub cache_id: String,
1167 #[serde(default, skip_serializing_if = "Vec::is_empty")]
1168 pub cache_namespace_fingerprints: Vec<String>,
1169 pub format: String,
1170 pub partition: PredictionPartition,
1171 #[serde(default = "default_prediction_level")]
1172 pub prediction_level: PredictionLevel,
1173 #[serde(default)]
1174 pub fold_ids: Vec<FoldId>,
1175 #[serde(default, skip_serializing_if = "Vec::is_empty")]
1176 pub unit_ids: Vec<PredictionUnitId>,
1177 #[serde(default)]
1178 pub sample_ids: Vec<SampleId>,
1179 pub prediction_width: usize,
1180 pub target_names: Vec<String>,
1181 pub block_count: usize,
1182 pub row_count: usize,
1183 pub content_fingerprint: String,
1184 #[serde(default)]
1185 pub blocks: Vec<BundlePredictionBlockCacheRecord>,
1186}
1187
1188impl BundlePredictionCacheRecord {
1189 pub fn validate(&self) -> Result<()> {
1190 validate_non_empty("requirement_key", &self.requirement_key)?;
1191 validate_non_empty("cache_id", &self.cache_id)?;
1192 validate_prediction_cache_namespace_fingerprints(
1193 &self.cache_id,
1194 &self.cache_namespace_fingerprints,
1195 )?;
1196 validate_non_empty("format", &self.format)?;
1197 if !supported_prediction_cache_format(&self.format) {
1198 return Err(DagMlError::RuntimeValidation(format!(
1199 "prediction cache `{}` uses unsupported format `{}`",
1200 self.cache_id, self.format
1201 )));
1202 }
1203 if self.partition != PredictionPartition::Validation {
1204 return Err(DagMlError::RuntimeValidation(format!(
1205 "prediction cache `{}` must cache validation OOF predictions",
1206 self.cache_id
1207 )));
1208 }
1209 validate_unique_ids("fold id", &self.fold_ids)?;
1210 validate_prediction_cache_record_units(self)?;
1211 if self.prediction_width == 0 {
1212 return Err(DagMlError::RuntimeValidation(format!(
1213 "prediction cache `{}` has zero prediction width",
1214 self.cache_id
1215 )));
1216 }
1217 if self.target_names.len() != self.prediction_width {
1218 return Err(DagMlError::RuntimeValidation(format!(
1219 "prediction cache `{}` target name count does not match prediction width",
1220 self.cache_id
1221 )));
1222 }
1223 for target_name in &self.target_names {
1224 validate_non_empty("target_name", target_name)?;
1225 }
1226 if self.block_count == 0 || self.block_count != self.blocks.len() {
1227 return Err(DagMlError::RuntimeValidation(format!(
1228 "prediction cache `{}` block_count does not match block records",
1229 self.cache_id
1230 )));
1231 }
1232 if !self.cache_namespace_fingerprints.is_empty()
1233 && self.cache_namespace_fingerprints.len() != self.block_count
1234 {
1235 return Err(DagMlError::RuntimeValidation(format!(
1236 "prediction cache `{}` namespace fingerprint count does not match block_count",
1237 self.cache_id
1238 )));
1239 }
1240 validate_prediction_cache_record_blocks(self)?;
1241 validate_fingerprint("prediction cache content", &self.content_fingerprint)?;
1242 Ok(())
1243 }
1244}
1245
1246fn validate_prediction_requirement_units(requirement: &BundlePredictionRequirement) -> Result<()> {
1247 match requirement.prediction_level {
1248 PredictionLevel::Observation => Err(DagMlError::RuntimeValidation(format!(
1249 "bundle prediction requirement `{}` cannot replay observation-level caches; aggregate to sample first",
1250 requirement.key()
1251 ))),
1252 PredictionLevel::Sample => {
1253 validate_unique_ids("sample id", &requirement.sample_ids)?;
1254 if requirement.sample_ids.is_empty() {
1255 return Err(DagMlError::RuntimeValidation(format!(
1256 "bundle prediction requirement `{}` has no sample ids",
1257 requirement.key()
1258 )));
1259 }
1260 if !requirement.unit_ids.is_empty()
1261 && requirement.unit_ids != sample_prediction_units(&requirement.sample_ids)
1262 {
1263 return Err(DagMlError::RuntimeValidation(format!(
1264 "bundle prediction requirement `{}` sample ids do not match unit ids",
1265 requirement.key()
1266 )));
1267 }
1268 Ok(())
1269 }
1270 PredictionLevel::Target | PredictionLevel::Group => {
1271 if !requirement.sample_ids.is_empty() {
1272 return Err(DagMlError::RuntimeValidation(format!(
1273 "bundle prediction requirement `{}` uses {:?} unit ids but also carries sample ids",
1274 requirement.key(),
1275 requirement.prediction_level
1276 )));
1277 }
1278 validate_prediction_units(
1279 "bundle prediction requirement unit",
1280 requirement.prediction_level,
1281 &requirement.unit_ids,
1282 )?;
1283 if requirement.unit_ids.is_empty() {
1284 return Err(DagMlError::RuntimeValidation(format!(
1285 "bundle prediction requirement `{}` has no unit ids",
1286 requirement.key()
1287 )));
1288 }
1289 Ok(())
1290 }
1291 }
1292}
1293
1294fn validate_prediction_cache_block_record_units(
1295 block: &BundlePredictionBlockCacheRecord,
1296) -> Result<()> {
1297 match block.prediction_level {
1298 PredictionLevel::Observation => Err(DagMlError::RuntimeValidation(
1299 "prediction block cache record cannot use observation-level predictions".to_string(),
1300 )),
1301 PredictionLevel::Sample => {
1302 validate_unique_ids("sample id", &block.sample_ids)?;
1303 if block.row_count != block.sample_ids.len() {
1304 return Err(DagMlError::RuntimeValidation(format!(
1305 "prediction block cache record row_count {} does not match {} sample ids",
1306 block.row_count,
1307 block.sample_ids.len()
1308 )));
1309 }
1310 if !block.unit_ids.is_empty()
1311 && block.unit_ids != sample_prediction_units(&block.sample_ids)
1312 {
1313 return Err(DagMlError::RuntimeValidation(
1314 "prediction block cache record sample ids do not match unit ids".to_string(),
1315 ));
1316 }
1317 Ok(())
1318 }
1319 PredictionLevel::Target | PredictionLevel::Group => {
1320 if !block.sample_ids.is_empty() {
1321 return Err(DagMlError::RuntimeValidation(format!(
1322 "prediction block cache record uses {:?} unit ids but also carries sample ids",
1323 block.prediction_level
1324 )));
1325 }
1326 validate_prediction_units(
1327 "prediction block cache record unit",
1328 block.prediction_level,
1329 &block.unit_ids,
1330 )?;
1331 if block.row_count != block.unit_ids.len() {
1332 return Err(DagMlError::RuntimeValidation(format!(
1333 "prediction block cache record row_count {} does not match {} unit ids",
1334 block.row_count,
1335 block.unit_ids.len()
1336 )));
1337 }
1338 Ok(())
1339 }
1340 }
1341}
1342
1343fn validate_prediction_cache_record_units(cache: &BundlePredictionCacheRecord) -> Result<()> {
1344 match cache.prediction_level {
1345 PredictionLevel::Observation => Err(DagMlError::RuntimeValidation(format!(
1346 "prediction cache `{}` cannot use observation-level predictions",
1347 cache.cache_id
1348 ))),
1349 PredictionLevel::Sample => {
1350 validate_unique_ids("sample id", &cache.sample_ids)?;
1351 if cache.row_count != cache.sample_ids.len() {
1352 return Err(DagMlError::RuntimeValidation(format!(
1353 "prediction cache `{}` row_count does not match unique sample ids",
1354 cache.cache_id
1355 )));
1356 }
1357 if !cache.unit_ids.is_empty()
1358 && cache.unit_ids != sample_prediction_units(&cache.sample_ids)
1359 {
1360 return Err(DagMlError::RuntimeValidation(format!(
1361 "prediction cache `{}` sample ids do not match unit ids",
1362 cache.cache_id
1363 )));
1364 }
1365 Ok(())
1366 }
1367 PredictionLevel::Target | PredictionLevel::Group => {
1368 if !cache.sample_ids.is_empty() {
1369 return Err(DagMlError::RuntimeValidation(format!(
1370 "prediction cache `{}` uses {:?} unit ids but also carries sample ids",
1371 cache.cache_id, cache.prediction_level
1372 )));
1373 }
1374 validate_prediction_units(
1375 "prediction cache unit",
1376 cache.prediction_level,
1377 &cache.unit_ids,
1378 )?;
1379 if cache.row_count != cache.unit_ids.len() {
1380 return Err(DagMlError::RuntimeValidation(format!(
1381 "prediction cache `{}` row_count does not match unique unit ids",
1382 cache.cache_id
1383 )));
1384 }
1385 Ok(())
1386 }
1387 }
1388}
1389
1390fn validate_prediction_cache_record_blocks(cache: &BundlePredictionCacheRecord) -> Result<()> {
1391 let mut row_count = 0usize;
1392 let mut samples = BTreeSet::new();
1393 let mut units = BTreeSet::new();
1394 for block in &cache.blocks {
1395 block.validate()?;
1396 if block.prediction_level != cache.prediction_level {
1397 return Err(DagMlError::RuntimeValidation(format!(
1398 "prediction cache `{}` mixes block prediction levels",
1399 cache.cache_id
1400 )));
1401 }
1402 row_count += block.row_count;
1403 match cache.prediction_level {
1404 PredictionLevel::Sample => {
1405 for sample_id in &block.sample_ids {
1406 if !samples.insert(sample_id.clone()) {
1407 return Err(DagMlError::RuntimeValidation(format!(
1408 "prediction cache `{}` contains duplicate sample `{sample_id}`",
1409 cache.cache_id
1410 )));
1411 }
1412 }
1413 }
1414 PredictionLevel::Target | PredictionLevel::Group => {
1415 for unit_id in &block.unit_ids {
1416 if !units.insert(unit_id.clone()) {
1417 return Err(DagMlError::RuntimeValidation(format!(
1418 "prediction cache `{}` contains duplicate unit `{unit_id}`",
1419 cache.cache_id
1420 )));
1421 }
1422 }
1423 }
1424 PredictionLevel::Observation => {
1425 unreachable!("record unit validation rejects observation")
1426 }
1427 }
1428 }
1429 if cache.row_count == 0 || cache.row_count != row_count {
1430 return Err(DagMlError::RuntimeValidation(format!(
1431 "prediction cache `{}` row_count does not match block records",
1432 cache.cache_id
1433 )));
1434 }
1435 if cache.prediction_level == PredictionLevel::Sample {
1436 let expected = cache.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
1437 if samples != expected {
1438 return Err(DagMlError::RuntimeValidation(format!(
1439 "prediction cache `{}` block samples do not match cache sample ids",
1440 cache.cache_id
1441 )));
1442 }
1443 } else {
1444 let expected = cache.unit_ids.iter().cloned().collect::<BTreeSet<_>>();
1445 if units != expected {
1446 return Err(DagMlError::RuntimeValidation(format!(
1447 "prediction cache `{}` block units do not match cache unit ids",
1448 cache.cache_id
1449 )));
1450 }
1451 }
1452 Ok(())
1453}
1454
1455fn validate_prediction_cache_payload_blocks(
1456 payload: &BundlePredictionCachePayload,
1457) -> Result<usize> {
1458 match payload.prediction_level {
1459 PredictionLevel::Observation => Err(DagMlError::RuntimeValidation(format!(
1460 "prediction cache payload `{}` cannot use observation-level predictions",
1461 payload.cache_id
1462 ))),
1463 PredictionLevel::Sample => validate_sample_prediction_cache_payload_blocks(payload),
1464 PredictionLevel::Target | PredictionLevel::Group => {
1465 validate_aggregated_prediction_cache_payload_blocks(payload)
1466 }
1467 }
1468}
1469
1470fn validate_sample_prediction_cache_payload_blocks(
1471 payload: &BundlePredictionCachePayload,
1472) -> Result<usize> {
1473 let mut row_count = 0usize;
1474 let mut sample_ids = BTreeSet::new();
1475 for block in &payload.blocks {
1476 block.validate_shape()?;
1477 if block.partition != payload.partition {
1478 return Err(DagMlError::RuntimeValidation(format!(
1479 "prediction cache payload `{}` contains a block from partition {:?}",
1480 payload.cache_id, block.partition
1481 )));
1482 }
1483 for sample_id in &block.sample_ids {
1484 if !sample_ids.insert(sample_id) {
1485 return Err(DagMlError::RuntimeValidation(format!(
1486 "prediction cache payload `{}` contains duplicate sample `{}`",
1487 payload.cache_id, sample_id
1488 )));
1489 }
1490 }
1491 row_count += block.sample_ids.len();
1492 }
1493 Ok(row_count)
1494}
1495
1496fn validate_aggregated_prediction_cache_payload_blocks(
1497 payload: &BundlePredictionCachePayload,
1498) -> Result<usize> {
1499 let mut row_count = 0usize;
1500 let mut unit_ids = BTreeSet::new();
1501 for block in &payload.aggregated_blocks {
1502 block.validate_shape()?;
1503 if block.partition != payload.partition {
1504 return Err(DagMlError::RuntimeValidation(format!(
1505 "prediction cache payload `{}` contains an aggregated block from partition {:?}",
1506 payload.cache_id, block.partition
1507 )));
1508 }
1509 if block.level != payload.prediction_level {
1510 return Err(DagMlError::RuntimeValidation(format!(
1511 "prediction cache payload `{}` contains {:?} block inside {:?} payload",
1512 payload.cache_id, block.level, payload.prediction_level
1513 )));
1514 }
1515 for unit_id in &block.unit_ids {
1516 if !unit_ids.insert(unit_id) {
1517 return Err(DagMlError::RuntimeValidation(format!(
1518 "prediction cache payload `{}` contains duplicate unit `{unit_id}`",
1519 payload.cache_id
1520 )));
1521 }
1522 }
1523 row_count += block.unit_ids.len();
1524 }
1525 Ok(row_count)
1526}
1527
1528#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1529pub struct BundlePredictionCachePayload {
1530 pub requirement_key: String,
1531 pub cache_id: String,
1532 #[serde(default, skip_serializing_if = "Vec::is_empty")]
1533 pub cache_namespace_fingerprints: Vec<String>,
1534 pub format: String,
1535 pub partition: PredictionPartition,
1536 #[serde(default = "default_prediction_level")]
1537 pub prediction_level: PredictionLevel,
1538 pub block_count: usize,
1539 pub row_count: usize,
1540 pub content_fingerprint: String,
1541 #[serde(default)]
1542 pub blocks: Vec<PredictionBlock>,
1543 #[serde(default, skip_serializing_if = "Vec::is_empty")]
1544 pub aggregated_blocks: Vec<AggregatedPredictionBlock>,
1545}
1546
1547impl BundlePredictionCachePayload {
1548 pub fn validate(&self) -> Result<()> {
1549 validate_non_empty("requirement_key", &self.requirement_key)?;
1550 validate_non_empty("cache_id", &self.cache_id)?;
1551 validate_prediction_cache_namespace_fingerprints(
1552 &self.cache_id,
1553 &self.cache_namespace_fingerprints,
1554 )?;
1555 validate_non_empty("format", &self.format)?;
1556 if !supported_prediction_cache_format(&self.format) {
1557 return Err(DagMlError::RuntimeValidation(format!(
1558 "prediction cache payload `{}` uses unsupported format `{}`",
1559 self.cache_id, self.format
1560 )));
1561 }
1562 let payload_schema_version = prediction_cache_schema_version_for_format(
1563 &self.format,
1564 &format!("prediction cache payload `{}`", self.cache_id),
1565 )?;
1566 if self.partition != PredictionPartition::Validation {
1567 return Err(DagMlError::RuntimeValidation(format!(
1568 "prediction cache payload `{}` must cache validation OOF predictions",
1569 self.cache_id
1570 )));
1571 }
1572 let expected_block_count = if self.prediction_level == PredictionLevel::Sample {
1573 if !self.aggregated_blocks.is_empty() {
1574 return Err(DagMlError::RuntimeValidation(format!(
1575 "prediction cache payload `{}` mixes sample and aggregated blocks",
1576 self.cache_id
1577 )));
1578 }
1579 self.blocks.len()
1580 } else {
1581 if !self.blocks.is_empty() {
1582 return Err(DagMlError::RuntimeValidation(format!(
1583 "prediction cache payload `{}` mixes aggregated and sample blocks",
1584 self.cache_id
1585 )));
1586 }
1587 self.aggregated_blocks.len()
1588 };
1589 if self.block_count == 0 || self.block_count != expected_block_count {
1590 return Err(DagMlError::RuntimeValidation(format!(
1591 "prediction cache payload `{}` block_count does not match blocks",
1592 self.cache_id
1593 )));
1594 }
1595 if !self.cache_namespace_fingerprints.is_empty()
1596 && self.cache_namespace_fingerprints.len() != self.block_count
1597 {
1598 return Err(DagMlError::RuntimeValidation(format!(
1599 "prediction cache payload `{}` namespace fingerprint count does not match block_count",
1600 self.cache_id
1601 )));
1602 }
1603 let row_count = validate_prediction_cache_payload_blocks(self)?;
1604 if self.row_count == 0 || self.row_count != row_count {
1605 return Err(DagMlError::RuntimeValidation(format!(
1606 "prediction cache payload `{}` row_count does not match blocks",
1607 self.cache_id
1608 )));
1609 }
1610 validate_prediction_cache_payload_block_family(self, payload_schema_version)?;
1611 validate_fingerprint(
1612 "prediction cache payload content",
1613 &self.content_fingerprint,
1614 )?;
1615 let actual_fingerprint = if self.prediction_level == PredictionLevel::Sample {
1616 stable_json_fingerprint(&self.blocks)?
1617 } else {
1618 stable_json_fingerprint(&self.aggregated_blocks)?
1619 };
1620 if actual_fingerprint != self.content_fingerprint {
1621 return Err(DagMlError::RuntimeValidation(format!(
1622 "prediction cache payload `{}` content fingerprint does not match blocks",
1623 self.cache_id
1624 )));
1625 }
1626 Ok(())
1627 }
1628}
1629
1630#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1631pub struct BundlePredictionCachePayloadSet {
1632 pub bundle_id: BundleId,
1633 #[serde(default = "default_prediction_cache_payload_schema_version")]
1634 pub schema_version: u32,
1635 #[serde(default)]
1636 pub caches: Vec<BundlePredictionCachePayload>,
1637}
1638
1639impl BundlePredictionCachePayloadSet {
1640 pub fn validate(&self) -> Result<()> {
1641 prediction_cache_payload_schema_migration_policy().validate_read_version(
1642 self.schema_version,
1643 &format!(
1644 "prediction cache payload set for bundle `{}`",
1645 self.bundle_id
1646 ),
1647 )?;
1648 let mut requirement_keys = BTreeSet::new();
1649 let mut cache_ids = BTreeSet::new();
1650 for payload in &self.caches {
1651 payload.validate()?;
1652 validate_prediction_cache_format_for_schema_version(
1653 &payload.format,
1654 self.schema_version,
1655 &format!(
1656 "prediction cache payload `{}` in set for bundle `{}`",
1657 payload.cache_id, self.bundle_id
1658 ),
1659 )?;
1660 validate_prediction_cache_payload_block_family(payload, self.schema_version)?;
1661 if !requirement_keys.insert(payload.requirement_key.as_str()) {
1662 return Err(DagMlError::RuntimeValidation(format!(
1663 "prediction cache payload set for bundle `{}` has duplicate requirement `{}`",
1664 self.bundle_id, payload.requirement_key
1665 )));
1666 }
1667 if !cache_ids.insert(payload.cache_id.as_str()) {
1668 return Err(DagMlError::RuntimeValidation(format!(
1669 "prediction cache payload set for bundle `{}` has duplicate cache id `{}`",
1670 self.bundle_id, payload.cache_id
1671 )));
1672 }
1673 }
1674 Ok(())
1675 }
1676
1677 pub fn validate_against_bundle(&self, bundle: &ExecutionBundle) -> Result<()> {
1678 self.validate()?;
1679 bundle.validate()?;
1680 if self.bundle_id != bundle.bundle_id {
1681 return Err(DagMlError::RuntimeValidation(format!(
1682 "prediction cache payload set bundle `{}` does not match bundle `{}`",
1683 self.bundle_id, bundle.bundle_id
1684 )));
1685 }
1686 if self.schema_version != bundle.schema_version {
1687 return Err(DagMlError::RuntimeValidation(format!(
1688 "prediction cache payload set for bundle `{}` uses schema_version {} but bundle uses schema_version {}",
1689 self.bundle_id, self.schema_version, bundle.schema_version
1690 )));
1691 }
1692 if self.caches.len() != bundle.prediction_caches.len() {
1693 return Err(DagMlError::RuntimeValidation(format!(
1694 "prediction cache payload set for bundle `{}` has {} payload(s) for {} cache record(s)",
1695 self.bundle_id,
1696 self.caches.len(),
1697 bundle.prediction_caches.len()
1698 )));
1699 }
1700 let records_by_requirement = bundle
1701 .prediction_caches
1702 .iter()
1703 .map(|record| (record.requirement_key.as_str(), record))
1704 .collect::<BTreeMap<_, _>>();
1705 let payloads_by_requirement = self
1706 .caches
1707 .iter()
1708 .map(|payload| (payload.requirement_key.as_str(), payload))
1709 .collect::<BTreeMap<_, _>>();
1710 for (requirement_key, record) in records_by_requirement {
1711 let payload = payloads_by_requirement
1712 .get(requirement_key)
1713 .ok_or_else(|| {
1714 DagMlError::RuntimeValidation(format!(
1715 "prediction cache payload set for bundle `{}` is missing requirement `{}`",
1716 self.bundle_id, requirement_key
1717 ))
1718 })?;
1719 validate_prediction_cache_payload_matches_record(payload, record)?;
1720 }
1721 for requirement_key in payloads_by_requirement.keys() {
1722 if !bundle
1723 .prediction_caches
1724 .iter()
1725 .any(|record| record.requirement_key.as_str() == *requirement_key)
1726 {
1727 return Err(DagMlError::RuntimeValidation(format!(
1728 "prediction cache payload set for bundle `{}` contains unknown requirement `{}`",
1729 self.bundle_id, requirement_key
1730 )));
1731 }
1732 }
1733 Ok(())
1734 }
1735}
1736
1737#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1738pub struct RefitArtifactRecord {
1739 pub node_id: NodeId,
1740 pub controller_id: ControllerId,
1741 pub artifact: ArtifactRef,
1742 pub params_fingerprint: String,
1743 #[serde(default, skip_serializing_if = "Option::is_none")]
1744 pub training_loss_fingerprint: Option<String>,
1745 #[serde(default)]
1746 pub data_requirement_keys: Vec<String>,
1747 #[serde(default)]
1748 pub prediction_requirement_keys: Vec<String>,
1749}
1750
1751impl RefitArtifactRecord {
1752 pub fn validate(&self) -> Result<()> {
1753 self.artifact.validate()?;
1754 if self.artifact.id.as_str().is_empty() {
1755 return Err(DagMlError::RuntimeValidation(format!(
1756 "refit artifact for `{}` has empty artifact id",
1757 self.node_id
1758 )));
1759 }
1760 if self.artifact.kind.trim().is_empty() {
1761 return Err(DagMlError::RuntimeValidation(format!(
1762 "refit artifact `{}` has empty artifact kind",
1763 self.artifact.id
1764 )));
1765 }
1766 if self.artifact.controller_id != self.controller_id {
1767 return Err(DagMlError::RuntimeValidation(format!(
1768 "refit artifact `{}` controller `{}` does not match record controller `{}`",
1769 self.artifact.id, self.artifact.controller_id, self.controller_id
1770 )));
1771 }
1772 validate_fingerprint("params", &self.params_fingerprint)?;
1773 if let Some(fingerprint) = &self.training_loss_fingerprint {
1774 validate_fingerprint("training loss", fingerprint)?;
1775 }
1776 let mut seen_keys = BTreeSet::new();
1777 for key in &self.data_requirement_keys {
1778 if key.trim().is_empty() {
1779 return Err(DagMlError::RuntimeValidation(format!(
1780 "refit artifact `{}` has empty data requirement key",
1781 self.artifact.id
1782 )));
1783 }
1784 if !seen_keys.insert(key.as_str()) {
1785 return Err(DagMlError::RuntimeValidation(format!(
1786 "refit artifact `{}` has duplicate data requirement key `{key}`",
1787 self.artifact.id
1788 )));
1789 }
1790 }
1791 let mut seen_prediction_keys = BTreeSet::new();
1792 for key in &self.prediction_requirement_keys {
1793 if key.trim().is_empty() {
1794 return Err(DagMlError::RuntimeValidation(format!(
1795 "refit artifact `{}` has empty prediction requirement key",
1796 self.artifact.id
1797 )));
1798 }
1799 if !seen_prediction_keys.insert(key.as_str()) {
1800 return Err(DagMlError::RuntimeValidation(format!(
1801 "refit artifact `{}` has duplicate prediction requirement key `{key}`",
1802 self.artifact.id
1803 )));
1804 }
1805 }
1806 Ok(())
1807 }
1808}
1809
1810#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1811#[serde(deny_unknown_fields)]
1812pub struct ExecutionBundle {
1813 pub bundle_id: BundleId,
1814 #[serde(default = "default_execution_bundle_schema_version")]
1815 pub schema_version: u32,
1816 pub plan_id: String,
1817 pub graph_fingerprint: String,
1818 pub campaign_fingerprint: String,
1819 pub controller_fingerprint: String,
1820 #[serde(default)]
1821 pub selected_variant_id: Option<VariantId>,
1822 #[serde(default)]
1823 pub selections: BTreeMap<String, SelectionDecision>,
1824 #[serde(default)]
1825 pub refit_artifacts: Vec<RefitArtifactRecord>,
1826 #[serde(default)]
1827 pub prediction_requirements: Vec<BundlePredictionRequirement>,
1828 #[serde(default)]
1829 pub prediction_caches: Vec<BundlePredictionCacheRecord>,
1830 #[serde(default, skip_serializing_if = "Option::is_none")]
1831 pub methods_hpo_resume_state: Option<MethodsHpoResumeState>,
1832 #[serde(default, skip_serializing_if = "Option::is_none")]
1835 pub conformal_calibration: Option<ConformalCalibrationRef>,
1836 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
1837 pub raw_artifact_payloads: BTreeMap<ArtifactId, Vec<u8>>,
1838 #[serde(default, skip_serializing_if = "Option::is_none")]
1842 pub scores: Option<ScoreSet>,
1843 #[serde(default)]
1844 pub data_requirements: Vec<BundleDataRequirement>,
1845 #[serde(default)]
1846 pub unsafe_flags: BTreeSet<String>,
1847 #[serde(default)]
1848 pub metadata: BTreeMap<String, serde_json::Value>,
1849}
1850
1851impl ExecutionBundle {
1852 pub fn from_json(json: &str) -> Result<Self> {
1854 crate::canonical::parse_typed_json(json).map_err(|error| {
1855 DagMlError::RuntimeValidation(format!(
1856 "execution bundle is not strict TCV1 JSON: {error}"
1857 ))
1858 })?;
1859 let raw: serde_json::Value = serde_json::from_str(json)?;
1860 if raw
1861 .get("schema_version")
1862 .and_then(serde_json::Value::as_u64)
1863 == Some(u64::from(LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION))
1864 && raw
1865 .as_object()
1866 .is_some_and(|object| object.contains_key("conformal_calibration"))
1867 {
1868 return Err(DagMlError::RuntimeValidation(
1869 "execution bundle V1 JSON cannot contain conformal_calibration, including null"
1870 .to_string(),
1871 ));
1872 }
1873 let bundle: Self =
1874 deserialize_external_contract(json, "execution bundle", DagMlError::RuntimeValidation)?;
1875 bundle.validate()?;
1876 Ok(bundle)
1877 }
1878
1879 pub fn validate(&self) -> Result<()> {
1880 execution_bundle_schema_migration_policy()
1881 .validate_read_version(self.schema_version, &format!("bundle `{}`", self.bundle_id))?;
1882 if self.plan_id.trim().is_empty() {
1883 return Err(DagMlError::RuntimeValidation(format!(
1884 "bundle `{}` has empty plan_id",
1885 self.bundle_id
1886 )));
1887 }
1888 validate_fingerprint("graph", &self.graph_fingerprint)?;
1889 validate_fingerprint("campaign", &self.campaign_fingerprint)?;
1890 validate_fingerprint("controller", &self.controller_fingerprint)?;
1891 if let Some(state) = &self.methods_hpo_resume_state {
1892 state.validate().map_err(|error| {
1893 DagMlError::RuntimeValidation(format!(
1894 "bundle `{}` has invalid Methods HPO resume state: {error}",
1895 self.bundle_id
1896 ))
1897 })?;
1898 }
1899 if let Some(calibration) = &self.conformal_calibration {
1900 if self.schema_version == LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION {
1901 return Err(DagMlError::RuntimeValidation(
1902 "execution bundle V1 cannot carry conformal state; migrate to V2".to_string(),
1903 ));
1904 }
1905 calibration.validate()?;
1906 }
1907 let expected_raw_artifacts = self
1908 .refit_artifacts
1909 .iter()
1910 .filter(|record| record.artifact.backend == Some(ArtifactBackend::Raw))
1911 .map(|record| &record.artifact.id)
1912 .collect::<BTreeSet<_>>();
1913 let actual_raw_artifacts = self.raw_artifact_payloads.keys().collect::<BTreeSet<_>>();
1914 if expected_raw_artifacts != actual_raw_artifacts {
1915 return Err(DagMlError::RuntimeValidation(format!(
1916 "bundle `{}` raw artifact payload keys must exactly cover RAW refit artifacts (expected {:?}, got {:?})",
1917 self.bundle_id, expected_raw_artifacts, actual_raw_artifacts
1918 )));
1919 }
1920 for record in &self.refit_artifacts {
1921 if record.artifact.backend == Some(ArtifactBackend::Raw) {
1922 let payload = self
1923 .raw_artifact_payloads
1924 .get(&record.artifact.id)
1925 .expect("exact raw payload coverage was checked");
1926 let expected = record
1927 .artifact
1928 .content_fingerprint
1929 .as_deref()
1930 .ok_or_else(|| {
1931 DagMlError::RuntimeValidation(format!(
1932 "raw artifact `{}` has no content fingerprint",
1933 record.artifact.id
1934 ))
1935 })?;
1936 if format!("{:x}", sha2::Sha256::digest(payload)) != expected
1937 || record.artifact.size_bytes != Some(payload.len() as u64)
1938 {
1939 return Err(DagMlError::RuntimeValidation(format!(
1940 "raw artifact payload `{}` does not match its bundle reference",
1941 record.artifact.id
1942 )));
1943 }
1944 }
1945 }
1946 if let Some(scores) = &self.scores {
1947 scores.validate()?;
1948 if scores.schema_version != self.schema_version {
1949 return Err(DagMlError::RuntimeValidation(format!(
1950 "bundle `{}` uses schema_version {} but embedded scores use schema_version {}",
1951 self.bundle_id, self.schema_version, scores.schema_version
1952 )));
1953 }
1954 if scores.plan_id != self.plan_id {
1955 return Err(DagMlError::RuntimeValidation(format!(
1956 "bundle `{}` plan_id `{}` does not match its embedded scores plan_id `{}`",
1957 self.bundle_id, self.plan_id, scores.plan_id
1958 )));
1959 }
1960 }
1961 for (key, decision) in &self.selections {
1962 if key.trim().is_empty() {
1963 return Err(DagMlError::RuntimeValidation(format!(
1964 "bundle `{}` contains empty selection key",
1965 self.bundle_id
1966 )));
1967 }
1968 decision.validate()?;
1969 }
1970 let mut data_keys = BTreeMap::new();
1971 for requirement in &self.data_requirements {
1972 requirement.validate()?;
1973 let key = requirement.key();
1974 if data_keys.insert(key.clone(), requirement).is_some() {
1975 return Err(DagMlError::RuntimeValidation(format!(
1976 "bundle `{}` has duplicate data requirement `{}`",
1977 self.bundle_id, key
1978 )));
1979 }
1980 }
1981 let mut prediction_keys = BTreeMap::new();
1982 for requirement in &self.prediction_requirements {
1983 requirement.validate()?;
1984 let key = requirement.key();
1985 if prediction_keys.insert(key.clone(), requirement).is_some() {
1986 return Err(DagMlError::RuntimeValidation(format!(
1987 "bundle `{}` has duplicate prediction requirement `{}`",
1988 self.bundle_id, key
1989 )));
1990 }
1991 }
1992 let mut prediction_cache_keys = BTreeMap::new();
1993 for cache in &self.prediction_caches {
1994 cache.validate()?;
1995 if !cache.cache_namespace_fingerprints.is_empty() && self.selected_variant_id.is_none()
1996 {
1997 return Err(DagMlError::RuntimeValidation(format!(
1998 "bundle `{}` prediction cache `{}` is D10-enriched and requires selected_variant_id",
1999 self.bundle_id, cache.cache_id
2000 )));
2001 }
2002 validate_prediction_cache_format_for_schema_version(
2003 &cache.format,
2004 self.schema_version,
2005 &format!(
2006 "prediction cache `{}` in bundle `{}`",
2007 cache.cache_id, self.bundle_id
2008 ),
2009 )?;
2010 let requirement = prediction_keys.get(&cache.requirement_key).ok_or_else(|| {
2011 DagMlError::RuntimeValidation(format!(
2012 "prediction cache `{}` references unknown prediction requirement `{}`",
2013 cache.cache_id, cache.requirement_key
2014 ))
2015 })?;
2016 validate_prediction_cache_matches_requirement(cache, requirement)?;
2017 if prediction_cache_keys
2018 .insert(cache.requirement_key.clone(), cache)
2019 .is_some()
2020 {
2021 return Err(DagMlError::RuntimeValidation(format!(
2022 "bundle `{}` has duplicate prediction cache for requirement `{}`",
2023 self.bundle_id, cache.requirement_key
2024 )));
2025 }
2026 }
2027 for artifact in &self.refit_artifacts {
2028 artifact.validate()?;
2029 for key in &artifact.data_requirement_keys {
2030 match data_keys.get(key) {
2031 Some(requirement) if requirement.node_id == artifact.node_id => {}
2032 Some(requirement) => {
2033 return Err(DagMlError::RuntimeValidation(format!(
2034 "refit artifact `{}` for `{}` references data requirement `{key}` owned by `{}`",
2035 artifact.artifact.id, artifact.node_id, requirement.node_id
2036 )));
2037 }
2038 None => {
2039 return Err(DagMlError::RuntimeValidation(format!(
2040 "refit artifact `{}` references unknown data requirement `{key}`",
2041 artifact.artifact.id
2042 )));
2043 }
2044 }
2045 }
2046 for key in &artifact.prediction_requirement_keys {
2047 match prediction_keys.get(key) {
2048 Some(requirement) if requirement.consumer_node == artifact.node_id => {}
2049 Some(requirement) => {
2050 return Err(DagMlError::RuntimeValidation(format!(
2051 "refit artifact `{}` for `{}` references prediction requirement `{key}` consumed by `{}`",
2052 artifact.artifact.id, artifact.node_id, requirement.consumer_node
2053 )));
2054 }
2055 None => {
2056 return Err(DagMlError::RuntimeValidation(format!(
2057 "refit artifact `{}` references unknown prediction requirement `{key}`",
2058 artifact.artifact.id
2059 )));
2060 }
2061 }
2062 if !prediction_cache_keys.contains_key(key) {
2063 return Err(DagMlError::RuntimeValidation(format!(
2064 "refit artifact `{}` references prediction requirement `{key}` without a prediction cache record",
2065 artifact.artifact.id
2066 )));
2067 }
2068 }
2069 }
2070 for unsafe_flag in &self.unsafe_flags {
2071 if unsafe_flag.trim().is_empty() {
2072 return Err(DagMlError::RuntimeValidation(format!(
2073 "bundle `{}` contains an empty unsafe flag",
2074 self.bundle_id
2075 )));
2076 }
2077 }
2078 Ok(())
2079 }
2080
2081 pub fn validate_against_plan(&self, plan: &ExecutionPlan) -> Result<()> {
2082 self.validate()?;
2083 plan.validate()?;
2084 if self.plan_id != plan.id {
2085 return Err(DagMlError::RuntimeValidation(format!(
2086 "bundle `{}` plan_id `{}` does not match plan `{}`",
2087 self.bundle_id, self.plan_id, plan.id
2088 )));
2089 }
2090 if self.graph_fingerprint != plan.graph_fingerprint
2091 || self.campaign_fingerprint != plan.campaign_fingerprint
2092 || self.controller_fingerprint != plan.controller_fingerprint
2093 {
2094 return Err(DagMlError::RuntimeValidation(format!(
2095 "bundle `{}` fingerprints do not match execution plan",
2096 self.bundle_id
2097 )));
2098 }
2099 let selected_variant = match &self.selected_variant_id {
2100 Some(selected_variant_id) => Some(
2101 plan.variants
2102 .iter()
2103 .find(|variant| &variant.variant_id == selected_variant_id)
2104 .ok_or_else(|| {
2105 DagMlError::RuntimeValidation(format!(
2106 "bundle `{}` selected unknown variant `{selected_variant_id}`",
2107 self.bundle_id
2108 ))
2109 })?,
2110 ),
2111 None => None,
2112 };
2113 self.validate_selections_against_plan(plan)?;
2114 if let Some(state) = &self.methods_hpo_resume_state {
2115 state.validate_against_plan(plan)?;
2116 }
2117 let expected_requirements = collect_data_requirements(plan)?;
2118 let expected_by_key = expected_requirements
2119 .iter()
2120 .map(|requirement| (requirement.key(), requirement))
2121 .collect::<BTreeMap<_, _>>();
2122 if self.data_requirements.len() != expected_by_key.len() {
2123 return Err(DagMlError::RuntimeValidation(format!(
2124 "bundle `{}` data requirement count does not match execution plan",
2125 self.bundle_id
2126 )));
2127 }
2128 for requirement in &self.data_requirements {
2129 let key = requirement.key();
2130 let expected = expected_by_key.get(&key).ok_or_else(|| {
2131 DagMlError::RuntimeValidation(format!(
2132 "bundle `{}` data requirement `{key}` does not exist in execution plan",
2133 self.bundle_id
2134 ))
2135 })?;
2136 if !requirement.matches_plan_requirement(expected) {
2137 return Err(DagMlError::RuntimeValidation(format!(
2138 "bundle `{}` data requirement `{key}` does not match execution plan",
2139 self.bundle_id
2140 )));
2141 }
2142 }
2143 for artifact in &self.refit_artifacts {
2144 let node_plan = plan.node_plans.get(&artifact.node_id).ok_or_else(|| {
2145 DagMlError::RuntimeValidation(format!(
2146 "bundle `{}` artifact references unknown node `{}`",
2147 self.bundle_id, artifact.node_id
2148 ))
2149 })?;
2150 if artifact.controller_id != node_plan.controller_id {
2151 return Err(DagMlError::RuntimeValidation(format!(
2152 "bundle `{}` artifact controller for `{}` does not match plan",
2153 self.bundle_id, artifact.node_id
2154 )));
2155 }
2156 let expected_params_fingerprint =
2157 expected_refit_artifact_params_fingerprint(node_plan, selected_variant)?;
2158 if artifact.params_fingerprint != expected_params_fingerprint {
2159 return Err(DagMlError::RuntimeValidation(format!(
2160 "bundle `{}` artifact params for `{}` do not match plan",
2161 self.bundle_id, artifact.node_id
2162 )));
2163 }
2164 if artifact.training_loss_fingerprint
2165 != node_plan.training_loss_fingerprint(Phase::Refit)?
2166 {
2167 return Err(DagMlError::RuntimeValidation(format!(
2168 "bundle `{}` artifact training loss for `{}` does not match plan",
2169 self.bundle_id, artifact.node_id
2170 )));
2171 }
2172 }
2173 for requirement in &self.prediction_requirements {
2174 let edge = plan
2175 .graph_plan
2176 .graph
2177 .edges
2178 .iter()
2179 .find(|edge| {
2180 edge.source.node_id == requirement.producer_node
2181 && edge.source.port_name == requirement.source_port
2182 && edge.target.node_id == requirement.consumer_node
2183 && edge.target.port_name == requirement.target_port
2184 && edge.contract.requires_oof
2185 })
2186 .ok_or_else(|| {
2187 DagMlError::RuntimeValidation(format!(
2188 "bundle `{}` prediction requirement `{}` does not match an OOF edge in the plan",
2189 self.bundle_id,
2190 requirement.key()
2191 ))
2192 })?;
2193 let cache = self
2194 .prediction_caches
2195 .iter()
2196 .find(|cache| cache.requirement_key == requirement.key());
2197 validate_prediction_requirement_against_plan(self, plan, edge, requirement, cache)?;
2198 }
2199 let cache_by_key = self
2206 .prediction_caches
2207 .iter()
2208 .map(|cache| (cache.requirement_key.clone(), cache))
2209 .collect::<BTreeMap<_, _>>();
2210 let mut concat_merge_groups: BTreeMap<NodeId, Vec<&BundlePredictionRequirement>> =
2211 BTreeMap::new();
2212 for requirement in &self.prediction_requirements {
2213 if is_concat_merge_consumer(plan, &requirement.consumer_node) {
2214 concat_merge_groups
2215 .entry(requirement.consumer_node.clone())
2216 .or_default()
2217 .push(requirement);
2218 }
2219 }
2220 for (consumer_node, requirements) in &concat_merge_groups {
2221 validate_concat_merge_requirement_group(
2222 self,
2223 plan,
2224 consumer_node,
2225 requirements,
2226 &cache_by_key,
2227 )?;
2228 }
2229 Ok(())
2230 }
2231
2232 fn validate_selections_against_plan(&self, plan: &ExecutionPlan) -> Result<()> {
2233 if self.selections.is_empty() {
2234 return Ok(());
2235 }
2236 let artifact_node_ids = self
2237 .refit_artifacts
2238 .iter()
2239 .map(|artifact| artifact.node_id.clone())
2240 .collect::<BTreeSet<_>>();
2241 let required_metric_level = plan.campaign.aggregation_policy.selection_metric_level;
2242 for (selection_key, decision) in &self.selections {
2243 match decision.metric_level {
2244 Some(metric_level) if metric_level == required_metric_level => {}
2245 Some(metric_level) => {
2246 return Err(DagMlError::RuntimeValidation(format!(
2247 "bundle `{}` selection `{selection_key}` metric_level {:?} does not match campaign selection_metric_level {:?}",
2248 self.bundle_id, metric_level, required_metric_level
2249 )));
2250 }
2251 None => {
2252 return Err(DagMlError::RuntimeValidation(format!(
2253 "bundle `{}` selection `{selection_key}` is missing metric_level for campaign selection_metric_level {:?}",
2254 self.bundle_id, required_metric_level
2255 )));
2256 }
2257 }
2258 let selected_candidate_id = decision.selected_candidate_id.as_str();
2259 if let Ok(selected_node_id) = NodeId::new(selected_candidate_id) {
2260 if let Some(node_plan) = plan.node_plans.get(&selected_node_id) {
2261 if node_plan.supported_phases.contains(&Phase::Refit)
2262 && !artifact_node_ids.contains(&node_plan.node_id)
2263 {
2264 return Err(DagMlError::RuntimeValidation(format!(
2265 "bundle `{}` selection `{selection_key}` chose refittable node `{}` without a matching refit artifact",
2266 self.bundle_id, node_plan.node_id
2267 )));
2268 }
2269 continue;
2270 }
2271 }
2272 if VariantId::new(selected_candidate_id).is_ok()
2273 && plan
2274 .variants
2275 .iter()
2276 .any(|variant| variant.variant_id.as_str() == selected_candidate_id)
2277 {
2278 continue;
2279 }
2280 return Err(DagMlError::RuntimeValidation(format!(
2281 "bundle `{}` selection `{selection_key}` chose unknown candidate `{selected_candidate_id}` for plan `{}`",
2282 self.bundle_id, plan.id
2283 )));
2284 }
2285 Ok(())
2286 }
2287
2288 pub fn validate_replay_envelopes(
2289 &self,
2290 envelopes: &BTreeMap<String, ExternalDataPlanEnvelope>,
2291 ) -> Result<()> {
2292 self.validate()?;
2293 for requirement in &self.data_requirements {
2294 let key = requirement.key();
2295 let envelope = envelopes.get(&key).ok_or_else(|| {
2296 DagMlError::RuntimeValidation(format!(
2297 "replay is missing external data envelope for `{key}`"
2298 ))
2299 })?;
2300 envelope.validate()?;
2301 if requirement.schema_fingerprint != envelope.schema_fingerprint
2302 || requirement.plan_fingerprint != envelope.plan_fingerprint
2303 || requirement.relation_fingerprint != envelope.relation_fingerprint
2304 {
2305 return Err(DagMlError::RuntimeValidation(format!(
2306 "replay envelope for `{key}` does not match bundle data requirement"
2307 )));
2308 }
2309 }
2310 Ok(())
2311 }
2312}
2313
2314fn expected_refit_artifact_params_fingerprint(
2315 node_plan: &crate::plan::NodePlan,
2316 selected_variant: Option<&crate::generation::VariantPlan>,
2317) -> Result<String> {
2318 let Some(variant) = selected_variant else {
2319 return Ok(node_plan.params_fingerprint.clone());
2320 };
2321 let effective_params =
2322 variant.effective_params_for_node(&node_plan.node_id, &node_plan.params)?;
2323 stable_json_fingerprint(&effective_params)
2324}
2325
2326fn is_concat_merge_consumer(plan: &ExecutionPlan, consumer_node: &NodeId) -> bool {
2335 let Some(node_plan) = plan.node_plans.get(consumer_node) else {
2336 return false;
2337 };
2338 if node_plan.kind != crate::graph::NodeKind::PredictionJoin {
2339 return false;
2340 }
2341 plan.graph_plan
2342 .graph
2343 .nodes
2344 .iter()
2345 .find(|node| &node.id == consumer_node)
2346 .and_then(|node| node.metadata.get("merge_mode"))
2347 .and_then(serde_json::Value::as_str)
2348 == Some("concat")
2349}
2350
2351fn validate_prediction_requirement_against_plan(
2352 bundle: &ExecutionBundle,
2353 plan: &ExecutionPlan,
2354 edge: &crate::graph::EdgeSpec,
2355 requirement: &BundlePredictionRequirement,
2356 cache: Option<&BundlePredictionCacheRecord>,
2357) -> Result<()> {
2358 if !edge.contract.requires_fold_alignment {
2359 return Ok(());
2360 }
2361 if is_concat_merge_consumer(plan, &requirement.consumer_node) {
2370 return validate_concat_merge_branch_input_requirement(bundle, plan, requirement, cache);
2371 }
2372 let fold_set = plan.fold_set.as_ref().ok_or_else(|| {
2373 DagMlError::RuntimeValidation(format!(
2374 "bundle `{}` prediction requirement `{}` needs fold alignment but plan `{}` has no fold set",
2375 bundle.bundle_id,
2376 requirement.key(),
2377 plan.id
2378 ))
2379 })?;
2380 let expected_fold_ids = fold_set
2381 .folds
2382 .iter()
2383 .map(|fold| fold.fold_id.clone())
2384 .collect::<BTreeSet<_>>();
2385 let requirement_fold_ids = requirement
2386 .fold_ids
2387 .iter()
2388 .cloned()
2389 .collect::<BTreeSet<_>>();
2390 if requirement_fold_ids != expected_fold_ids {
2391 return Err(DagMlError::RuntimeValidation(format!(
2392 "bundle `{}` prediction requirement `{}` fold ids do not match plan fold set",
2393 bundle.bundle_id,
2394 requirement.key()
2395 )));
2396 }
2397 if requirement.prediction_level != PredictionLevel::Sample {
2398 if let Some(cache) = cache {
2399 validate_aggregated_prediction_cache_blocks_match_requirement(
2400 bundle,
2401 requirement,
2402 cache,
2403 fold_set.partition_mode,
2404 )?;
2405 }
2406 return Ok(());
2407 }
2408 let expected_sample_ids = fold_set.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
2409 let requirement_sample_ids = requirement
2410 .sample_ids
2411 .iter()
2412 .cloned()
2413 .collect::<BTreeSet<_>>();
2414 if requirement_sample_ids != expected_sample_ids {
2415 return Err(DagMlError::RuntimeValidation(format!(
2416 "bundle `{}` prediction requirement `{}` sample ids do not match plan fold set",
2417 bundle.bundle_id,
2418 requirement.key()
2419 )));
2420 }
2421 if let Some(cache) = cache {
2422 validate_prediction_cache_blocks_match_fold_set(bundle, requirement, cache, fold_set)?;
2423 }
2424 Ok(())
2425}
2426
2427fn validate_concat_merge_branch_input_requirement(
2439 bundle: &ExecutionBundle,
2440 plan: &ExecutionPlan,
2441 requirement: &BundlePredictionRequirement,
2442 cache: Option<&BundlePredictionCacheRecord>,
2443) -> Result<()> {
2444 let fold_set = plan.fold_set.as_ref().ok_or_else(|| {
2445 DagMlError::RuntimeValidation(format!(
2446 "bundle `{}` prediction requirement `{}` needs fold alignment but plan `{}` has no fold set",
2447 bundle.bundle_id,
2448 requirement.key(),
2449 plan.id
2450 ))
2451 })?;
2452 let universe_fold_ids = fold_set
2453 .folds
2454 .iter()
2455 .map(|fold| fold.fold_id.clone())
2456 .collect::<BTreeSet<_>>();
2457 let requirement_fold_ids = requirement
2458 .fold_ids
2459 .iter()
2460 .cloned()
2461 .collect::<BTreeSet<_>>();
2462 if !requirement_fold_ids.is_subset(&universe_fold_ids) {
2463 return Err(DagMlError::RuntimeValidation(format!(
2464 "bundle `{}` concat-merge prediction requirement `{}` has fold ids outside the plan fold set",
2465 bundle.bundle_id,
2466 requirement.key()
2467 )));
2468 }
2469 if requirement.prediction_level != PredictionLevel::Sample {
2472 return Err(DagMlError::RuntimeValidation(format!(
2473 "bundle `{}` concat-merge prediction requirement `{}` must be sample-level (got {:?})",
2474 bundle.bundle_id,
2475 requirement.key(),
2476 requirement.prediction_level
2477 )));
2478 }
2479 let universe_sample_ids = fold_set.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
2480 let requirement_sample_ids = requirement
2481 .sample_ids
2482 .iter()
2483 .cloned()
2484 .collect::<BTreeSet<_>>();
2485 if !requirement_sample_ids.is_subset(&universe_sample_ids) {
2486 return Err(DagMlError::RuntimeValidation(format!(
2487 "bundle `{}` concat-merge prediction requirement `{}` covers samples outside the plan fold set",
2488 bundle.bundle_id,
2489 requirement.key()
2490 )));
2491 }
2492 if let Some(cache) = cache {
2493 let folds = fold_set
2494 .folds
2495 .iter()
2496 .map(|fold| (&fold.fold_id, fold))
2497 .collect::<BTreeMap<_, _>>();
2498 for block in &cache.blocks {
2499 let fold_id = block.fold_id.as_ref().ok_or_else(|| {
2500 DagMlError::RuntimeValidation(format!(
2501 "bundle `{}` prediction cache `{}` has an OOF block without a fold id",
2502 bundle.bundle_id, cache.cache_id
2503 ))
2504 })?;
2505 let fold = folds.get(fold_id).ok_or_else(|| {
2506 DagMlError::RuntimeValidation(format!(
2507 "bundle `{}` prediction cache `{}` references unknown fold `{fold_id}`",
2508 bundle.bundle_id, cache.cache_id
2509 ))
2510 })?;
2511 let block_samples = block.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
2512 if block_samples.len() != block.sample_ids.len() {
2513 return Err(DagMlError::RuntimeValidation(format!(
2514 "bundle `{}` prediction cache `{}` block for fold `{fold_id}` has a duplicate sample for requirement `{}`",
2515 bundle.bundle_id,
2516 cache.cache_id,
2517 requirement.key()
2518 )));
2519 }
2520 let validation_samples = fold
2521 .validation_sample_ids
2522 .iter()
2523 .cloned()
2524 .collect::<BTreeSet<_>>();
2525 if !block_samples.is_subset(&validation_samples) {
2526 return Err(DagMlError::RuntimeValidation(format!(
2527 "bundle `{}` prediction cache `{}` block for fold `{fold_id}` covers samples outside the fold validation set for requirement `{}`",
2528 bundle.bundle_id,
2529 cache.cache_id,
2530 requirement.key()
2531 )));
2532 }
2533 }
2534 }
2535 Ok(())
2536}
2537
2538fn validate_concat_merge_requirement_group(
2551 bundle: &ExecutionBundle,
2552 plan: &ExecutionPlan,
2553 consumer_node: &NodeId,
2554 requirements: &[&BundlePredictionRequirement],
2555 caches: &BTreeMap<String, &BundlePredictionCacheRecord>,
2556) -> Result<()> {
2557 let fold_set = plan.fold_set.as_ref().ok_or_else(|| {
2558 DagMlError::RuntimeValidation(format!(
2559 "bundle `{}` concat-merge node `{consumer_node}` needs fold alignment but plan `{}` has no fold set",
2560 bundle.bundle_id, plan.id
2561 ))
2562 })?;
2563
2564 let expected_keys = plan
2573 .graph_plan
2574 .graph
2575 .edges
2576 .iter()
2577 .filter(|edge| {
2578 &edge.target.node_id == consumer_node
2579 && edge.contract.requires_oof
2580 && edge.contract.requires_fold_alignment
2581 })
2582 .map(|edge| {
2583 bundle_prediction_requirement_key(
2584 &edge.source.node_id,
2585 &edge.source.port_name,
2586 &edge.target.node_id,
2587 &edge.target.port_name,
2588 )
2589 })
2590 .collect::<BTreeSet<_>>();
2591 let supplied_keys = requirements
2592 .iter()
2593 .map(|req| req.key())
2594 .collect::<BTreeSet<_>>();
2595 if supplied_keys != expected_keys {
2596 let missing: Vec<&str> = expected_keys
2597 .difference(&supplied_keys)
2598 .map(String::as_str)
2599 .collect();
2600 let extra: Vec<&str> = supplied_keys
2601 .difference(&expected_keys)
2602 .map(String::as_str)
2603 .collect();
2604 return Err(DagMlError::RuntimeValidation(format!(
2605 "bundle `{}` concat-merge node `{consumer_node}` branch inputs do not match the plan's incoming OOF edges (missing: [{}]; extra: [{}])",
2606 bundle.bundle_id,
2607 missing.join(", "),
2608 extra.join(", ")
2609 )));
2610 }
2611
2612 let expected_universe = fold_set.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
2614 let mut covered_universe = BTreeSet::new();
2615 for requirement in requirements {
2616 for sample_id in &requirement.sample_ids {
2617 if !covered_universe.insert(sample_id.clone()) {
2618 return Err(DagMlError::RuntimeValidation(format!(
2619 "bundle `{}` concat-merge node `{consumer_node}` received overlapping branch predictions: sample `{sample_id}` is covered by more than one partition",
2620 bundle.bundle_id
2621 )));
2622 }
2623 }
2624 }
2625 if covered_universe != expected_universe {
2626 return Err(DagMlError::RuntimeValidation(format!(
2627 "bundle `{}` concat-merge node `{consumer_node}` branch inputs do not cover the full fold set sample universe (each sample exactly once)",
2628 bundle.bundle_id
2629 )));
2630 }
2631
2632 let cached_count = requirements
2643 .iter()
2644 .filter(|req| caches.contains_key(&req.key()))
2645 .count();
2646 if cached_count != 0 && cached_count != requirements.len() {
2647 return Err(DagMlError::RuntimeValidation(format!(
2648 "bundle `{}` concat-merge node `{consumer_node}` has partial prediction-cache coverage ({cached_count} of {} branch inputs cached): all branch inputs must carry a per-fold OOF cache or none",
2649 bundle.bundle_id,
2650 requirements.len()
2651 )));
2652 }
2653 if cached_count == requirements.len() {
2654 let mut covered_by_fold: BTreeMap<FoldId, BTreeSet<SampleId>> = BTreeMap::new();
2655 for requirement in requirements {
2656 let cache = caches.get(&requirement.key()).expect("checked above");
2657 for block in &cache.blocks {
2658 let Some(fold_id) = block.fold_id.as_ref() else {
2659 continue;
2660 };
2661 let covered = covered_by_fold.entry(fold_id.clone()).or_default();
2662 for sample_id in &block.sample_ids {
2663 if !covered.insert(sample_id.clone()) {
2664 return Err(DagMlError::RuntimeValidation(format!(
2665 "bundle `{}` concat-merge node `{consumer_node}` has overlapping branch predictions in fold `{fold_id}`: sample `{sample_id}` is covered by more than one partition",
2666 bundle.bundle_id
2667 )));
2668 }
2669 }
2670 }
2671 }
2672 for fold in &fold_set.folds {
2673 let expected = fold
2674 .validation_sample_ids
2675 .iter()
2676 .cloned()
2677 .collect::<BTreeSet<_>>();
2678 let covered = covered_by_fold.remove(&fold.fold_id).unwrap_or_default();
2679 if covered != expected {
2680 return Err(DagMlError::RuntimeValidation(format!(
2681 "bundle `{}` concat-merge node `{consumer_node}` branch inputs do not cover fold `{}` validation set (each sample exactly once)",
2682 bundle.bundle_id, fold.fold_id
2683 )));
2684 }
2685 }
2686 }
2687 Ok(())
2688}
2689
2690fn validate_prediction_cache_blocks_match_fold_set(
2691 bundle: &ExecutionBundle,
2692 requirement: &BundlePredictionRequirement,
2693 cache: &BundlePredictionCacheRecord,
2694 fold_set: &crate::fold::FoldSet,
2695) -> Result<()> {
2696 let folds = fold_set
2697 .folds
2698 .iter()
2699 .map(|fold| (&fold.fold_id, fold))
2700 .collect::<BTreeMap<_, _>>();
2701 let expected_fold_ids = fold_set
2702 .folds
2703 .iter()
2704 .map(|fold| fold.fold_id.clone())
2705 .collect::<BTreeSet<_>>();
2706 let mut covered_fold_ids = BTreeSet::new();
2707 let mut covered_sample_ids = BTreeSet::new();
2708 for block in &cache.blocks {
2709 let fold_id = block.fold_id.as_ref().ok_or_else(|| {
2710 DagMlError::RuntimeValidation(format!(
2711 "bundle `{}` prediction cache `{}` has an OOF block without a fold id",
2712 bundle.bundle_id, cache.cache_id
2713 ))
2714 })?;
2715 covered_fold_ids.insert(fold_id.clone());
2716 let fold = folds.get(fold_id).ok_or_else(|| {
2717 DagMlError::RuntimeValidation(format!(
2718 "bundle `{}` prediction cache `{}` references unknown fold `{fold_id}`",
2719 bundle.bundle_id, cache.cache_id
2720 ))
2721 })?;
2722 let block_samples = block.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
2723 let expected_samples = fold
2724 .validation_sample_ids
2725 .iter()
2726 .cloned()
2727 .collect::<BTreeSet<_>>();
2728 if block_samples != expected_samples {
2729 return Err(DagMlError::RuntimeValidation(format!(
2730 "bundle `{}` prediction cache `{}` block for fold `{fold_id}` does not match validation samples for requirement `{}`",
2731 bundle.bundle_id,
2732 cache.cache_id,
2733 requirement.key()
2734 )));
2735 }
2736 for sample_id in block_samples {
2737 if !covered_sample_ids.insert(sample_id.clone())
2742 && fold_set.partition_mode == crate::fold::FoldPartitionMode::Partition
2743 {
2744 return Err(DagMlError::RuntimeValidation(format!(
2745 "bundle `{}` prediction cache `{}` has duplicate OOF sample `{sample_id}`",
2746 bundle.bundle_id, cache.cache_id
2747 )));
2748 }
2749 }
2750 }
2751 if covered_fold_ids != expected_fold_ids {
2752 return Err(DagMlError::RuntimeValidation(format!(
2753 "bundle `{}` prediction cache `{}` does not cover all folds for requirement `{}`",
2754 bundle.bundle_id,
2755 cache.cache_id,
2756 requirement.key()
2757 )));
2758 }
2759 let expected_sample_ids = fold_set.sample_ids.iter().cloned().collect::<BTreeSet<_>>();
2760 if covered_sample_ids != expected_sample_ids {
2761 return Err(DagMlError::RuntimeValidation(format!(
2762 "bundle `{}` prediction cache `{}` does not cover the full OOF sample universe for requirement `{}`",
2763 bundle.bundle_id,
2764 cache.cache_id,
2765 requirement.key()
2766 )));
2767 }
2768 Ok(())
2769}
2770
2771fn validate_aggregated_prediction_cache_blocks_match_requirement(
2772 bundle: &ExecutionBundle,
2773 requirement: &BundlePredictionRequirement,
2774 cache: &BundlePredictionCacheRecord,
2775 partition_mode: crate::fold::FoldPartitionMode,
2776) -> Result<()> {
2777 let mut covered_fold_ids = BTreeSet::new();
2778 let mut covered_unit_ids = BTreeSet::new();
2779 for block in &cache.blocks {
2780 if block.prediction_level != requirement.prediction_level {
2781 return Err(DagMlError::RuntimeValidation(format!(
2782 "bundle `{}` prediction cache `{}` block level does not match requirement `{}`",
2783 bundle.bundle_id,
2784 cache.cache_id,
2785 requirement.key()
2786 )));
2787 }
2788 if let Some(fold_id) = &block.fold_id {
2789 covered_fold_ids.insert(fold_id.clone());
2790 }
2791 for unit_id in &block.unit_ids {
2792 if !covered_unit_ids.insert(unit_id.clone())
2796 && partition_mode == crate::fold::FoldPartitionMode::Partition
2797 {
2798 return Err(DagMlError::RuntimeValidation(format!(
2799 "bundle `{}` prediction cache `{}` has duplicate aggregated unit `{unit_id}`",
2800 bundle.bundle_id, cache.cache_id
2801 )));
2802 }
2803 }
2804 }
2805 let expected_fold_ids = requirement
2806 .fold_ids
2807 .iter()
2808 .cloned()
2809 .collect::<BTreeSet<_>>();
2810 if covered_fold_ids != expected_fold_ids {
2811 return Err(DagMlError::RuntimeValidation(format!(
2812 "bundle `{}` prediction cache `{}` does not cover all folds for aggregated requirement `{}`",
2813 bundle.bundle_id,
2814 cache.cache_id,
2815 requirement.key()
2816 )));
2817 }
2818 let expected_unit_ids = requirement
2819 .unit_ids
2820 .iter()
2821 .cloned()
2822 .collect::<BTreeSet<_>>();
2823 if covered_unit_ids != expected_unit_ids {
2824 return Err(DagMlError::RuntimeValidation(format!(
2825 "bundle `{}` prediction cache `{}` does not cover all units for aggregated requirement `{}`",
2826 bundle.bundle_id,
2827 cache.cache_id,
2828 requirement.key()
2829 )));
2830 }
2831 Ok(())
2832}
2833
2834pub fn build_execution_bundle(
2835 bundle_id: BundleId,
2836 plan: &ExecutionPlan,
2837 selected_variant_id: Option<VariantId>,
2838 selections: BTreeMap<String, SelectionDecision>,
2839 refit_artifacts: Vec<RefitArtifactRecord>,
2840) -> Result<ExecutionBundle> {
2841 build_execution_bundle_with_prediction_requirements(
2842 bundle_id,
2843 plan,
2844 selected_variant_id,
2845 selections,
2846 refit_artifacts,
2847 Vec::new(),
2848 )
2849}
2850
2851pub fn build_execution_bundle_with_prediction_requirements(
2852 bundle_id: BundleId,
2853 plan: &ExecutionPlan,
2854 selected_variant_id: Option<VariantId>,
2855 selections: BTreeMap<String, SelectionDecision>,
2856 refit_artifacts: Vec<RefitArtifactRecord>,
2857 prediction_requirements: Vec<BundlePredictionRequirement>,
2858) -> Result<ExecutionBundle> {
2859 build_execution_bundle_with_prediction_contracts(
2860 bundle_id,
2861 plan,
2862 selected_variant_id,
2863 selections,
2864 refit_artifacts,
2865 prediction_requirements,
2866 Vec::new(),
2867 )
2868}
2869
2870pub fn build_execution_bundle_with_prediction_contracts(
2871 bundle_id: BundleId,
2872 plan: &ExecutionPlan,
2873 selected_variant_id: Option<VariantId>,
2874 selections: BTreeMap<String, SelectionDecision>,
2875 refit_artifacts: Vec<RefitArtifactRecord>,
2876 prediction_requirements: Vec<BundlePredictionRequirement>,
2877 prediction_caches: Vec<BundlePredictionCacheRecord>,
2878) -> Result<ExecutionBundle> {
2879 plan.validate()?;
2880 let bundle = ExecutionBundle {
2881 bundle_id,
2882 schema_version: EXECUTION_BUNDLE_SCHEMA_VERSION,
2883 plan_id: plan.id.clone(),
2884 graph_fingerprint: plan.graph_fingerprint.clone(),
2885 campaign_fingerprint: plan.campaign_fingerprint.clone(),
2886 controller_fingerprint: plan.controller_fingerprint.clone(),
2887 selected_variant_id,
2888 selections,
2889 refit_artifacts,
2890 prediction_requirements,
2891 prediction_caches,
2892 methods_hpo_resume_state: None,
2893 conformal_calibration: None,
2894 raw_artifact_payloads: BTreeMap::new(),
2895 scores: None,
2896 data_requirements: collect_data_requirements(plan)?,
2897 unsafe_flags: BTreeSet::new(),
2898 metadata: BTreeMap::new(),
2899 };
2900 if !bundle
2905 .refit_artifacts
2906 .iter()
2907 .any(|record| record.artifact.backend == Some(ArtifactBackend::Raw))
2908 {
2909 bundle.validate_against_plan(plan)?;
2910 }
2911 Ok(bundle)
2912}
2913
2914fn collect_data_requirements(plan: &ExecutionPlan) -> Result<Vec<BundleDataRequirement>> {
2915 let mut requirements = Vec::new();
2916 for node_plan in plan.node_plans.values() {
2917 for binding in &node_plan.data_bindings {
2918 requirements.push(BundleDataRequirement {
2919 node_id: node_plan.node_id.clone(),
2920 input_name: binding.input_name.clone(),
2921 schema_fingerprint: binding.schema_fingerprint.clone(),
2922 plan_fingerprint: binding.plan_fingerprint.clone(),
2923 relation_fingerprint: binding.relation_fingerprint.clone(),
2924 output_representation: binding.output_representation.clone(),
2925 feature_set_id: binding.feature_set_id.clone(),
2926 representation_replay_manifest: None,
2927 representation_compatibility: None,
2928 });
2929 }
2930 }
2931 requirements.sort_by_key(BundleDataRequirement::key);
2932 for requirement in &requirements {
2933 requirement.validate()?;
2934 }
2935 Ok(requirements)
2936}
2937
2938pub fn build_prediction_cache_record(
2939 requirement: &BundlePredictionRequirement,
2940 blocks: &[PredictionBlock],
2941) -> Result<BundlePredictionCacheRecord> {
2942 let selected = select_prediction_cache_blocks(requirement, blocks)?;
2943 build_prediction_cache_record_from_selected(requirement, &selected)
2944}
2945
2946pub fn build_prediction_cache_payload(
2947 requirement: &BundlePredictionRequirement,
2948 blocks: &[PredictionBlock],
2949) -> Result<BundlePredictionCachePayload> {
2950 let selected = select_prediction_cache_blocks(requirement, blocks)?;
2951 let payload = BundlePredictionCachePayload {
2952 requirement_key: requirement.key(),
2953 cache_id: format!("prediction-cache:{}", requirement.key()),
2954 cache_namespace_fingerprints: Vec::new(),
2955 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
2956 partition: requirement.partition.clone(),
2957 prediction_level: requirement.prediction_level,
2958 block_count: selected.len(),
2959 row_count: selected.iter().map(|block| block.sample_ids.len()).sum(),
2960 content_fingerprint: stable_json_fingerprint(&selected)?,
2961 blocks: selected,
2962 aggregated_blocks: Vec::new(),
2963 };
2964 payload.validate()?;
2965 let record = build_prediction_cache_record(requirement, &payload.blocks)?;
2966 validate_prediction_cache_payload_matches_record(&payload, &record)?;
2967 Ok(payload)
2968}
2969
2970pub fn build_aggregated_prediction_cache_record(
2971 requirement: &BundlePredictionRequirement,
2972 blocks: &[AggregatedPredictionBlock],
2973) -> Result<BundlePredictionCacheRecord> {
2974 let selected = select_aggregated_prediction_cache_blocks(requirement, blocks)?;
2975 build_aggregated_prediction_cache_record_from_selected(requirement, &selected)
2976}
2977
2978pub fn build_aggregated_prediction_cache_payload(
2979 requirement: &BundlePredictionRequirement,
2980 blocks: &[AggregatedPredictionBlock],
2981) -> Result<BundlePredictionCachePayload> {
2982 let selected = select_aggregated_prediction_cache_blocks(requirement, blocks)?;
2983 let payload = BundlePredictionCachePayload {
2984 requirement_key: requirement.key(),
2985 cache_id: format!("prediction-cache:{}", requirement.key()),
2986 cache_namespace_fingerprints: Vec::new(),
2987 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
2988 partition: requirement.partition.clone(),
2989 prediction_level: requirement.prediction_level,
2990 block_count: selected.len(),
2991 row_count: selected.iter().map(|block| block.unit_ids.len()).sum(),
2992 content_fingerprint: stable_json_fingerprint(&selected)?,
2993 blocks: Vec::new(),
2994 aggregated_blocks: selected,
2995 };
2996 payload.validate()?;
2997 let record = build_aggregated_prediction_cache_record(requirement, &payload.aggregated_blocks)?;
2998 validate_prediction_cache_payload_matches_record(&payload, &record)?;
2999 Ok(payload)
3000}
3001
3002pub fn validate_prediction_cache_payload_matches_record(
3003 payload: &BundlePredictionCachePayload,
3004 record: &BundlePredictionCacheRecord,
3005) -> Result<()> {
3006 payload.validate()?;
3007 record.validate()?;
3008 if payload.requirement_key != record.requirement_key
3009 || payload.cache_id != record.cache_id
3010 || payload.cache_namespace_fingerprints != record.cache_namespace_fingerprints
3011 || payload.format != record.format
3012 || payload.partition != record.partition
3013 || payload.prediction_level != record.prediction_level
3014 || payload.block_count != record.block_count
3015 || payload.row_count != record.row_count
3016 || payload.content_fingerprint != record.content_fingerprint
3017 {
3018 return Err(DagMlError::RuntimeValidation(format!(
3019 "prediction cache payload `{}` does not match cache record `{}`",
3020 payload.cache_id, record.cache_id
3021 )));
3022 }
3023 let block_records = if payload.prediction_level == PredictionLevel::Sample {
3024 payload
3025 .blocks
3026 .iter()
3027 .map(|block| {
3028 Ok(BundlePredictionBlockCacheRecord {
3029 prediction_id: block.prediction_id.clone(),
3030 fold_id: block.fold_id.clone(),
3031 prediction_level: PredictionLevel::Sample,
3032 row_count: block.sample_ids.len(),
3033 unit_ids: Vec::new(),
3034 sample_ids: block.sample_ids.clone(),
3035 content_fingerprint: stable_json_fingerprint(block)?,
3036 })
3037 })
3038 .collect::<Result<Vec<_>>>()?
3039 } else {
3040 payload
3041 .aggregated_blocks
3042 .iter()
3043 .map(|block| {
3044 Ok(BundlePredictionBlockCacheRecord {
3045 prediction_id: block.prediction_id.clone(),
3046 fold_id: block.fold_id.clone(),
3047 prediction_level: block.level,
3048 row_count: block.unit_ids.len(),
3049 unit_ids: block.unit_ids.clone(),
3050 sample_ids: Vec::new(),
3051 content_fingerprint: stable_json_fingerprint(block)?,
3052 })
3053 })
3054 .collect::<Result<Vec<_>>>()?
3055 };
3056 if block_records != record.blocks {
3057 return Err(DagMlError::RuntimeValidation(format!(
3058 "prediction cache payload `{}` block fingerprints do not match cache record",
3059 payload.cache_id
3060 )));
3061 }
3062 Ok(())
3063}
3064
3065fn validate_prediction_cache_namespace_fingerprints(
3066 cache_id: &str,
3067 fingerprints: &[String],
3068) -> Result<()> {
3069 let mut seen = BTreeSet::new();
3070 for fingerprint in fingerprints {
3071 validate_fingerprint("prediction cache namespace", fingerprint)?;
3072 if !seen.insert(fingerprint.as_str()) {
3073 return Err(DagMlError::RuntimeValidation(format!(
3074 "prediction cache `{cache_id}` has duplicate cache namespace fingerprint `{fingerprint}`"
3075 )));
3076 }
3077 }
3078 Ok(())
3079}
3080
3081fn select_prediction_cache_blocks(
3082 requirement: &BundlePredictionRequirement,
3083 blocks: &[PredictionBlock],
3084) -> Result<Vec<PredictionBlock>> {
3085 requirement.validate()?;
3086 let mut selected = blocks
3087 .iter()
3088 .filter(|block| {
3089 block.producer_node == requirement.producer_node
3090 && block.partition == requirement.partition
3091 && block
3092 .fold_id
3093 .as_ref()
3094 .is_some_and(|fold_id| requirement.fold_ids.contains(fold_id))
3095 })
3096 .cloned()
3097 .collect::<Vec<_>>();
3098 if selected.is_empty() {
3099 return Err(DagMlError::RuntimeValidation(format!(
3100 "prediction cache requirement `{}` has no matching prediction blocks",
3101 requirement.key()
3102 )));
3103 }
3104 selected.sort_by(|left, right| {
3105 (
3106 left.fold_id.as_ref().map(ToString::to_string),
3107 left.prediction_id.clone(),
3108 )
3109 .cmp(&(
3110 right.fold_id.as_ref().map(ToString::to_string),
3111 right.prediction_id.clone(),
3112 ))
3113 });
3114 Ok(selected)
3115}
3116
3117fn select_aggregated_prediction_cache_blocks(
3118 requirement: &BundlePredictionRequirement,
3119 blocks: &[AggregatedPredictionBlock],
3120) -> Result<Vec<AggregatedPredictionBlock>> {
3121 requirement.validate()?;
3122 if requirement.prediction_level == PredictionLevel::Sample {
3123 return Err(DagMlError::RuntimeValidation(format!(
3124 "aggregated prediction cache requirement `{}` must use target or group level",
3125 requirement.key()
3126 )));
3127 }
3128 let mut selected = blocks
3129 .iter()
3130 .filter(|block| {
3131 block.producer_node == requirement.producer_node
3132 && block.partition == requirement.partition
3133 && block.level == requirement.prediction_level
3134 && block
3135 .fold_id
3136 .as_ref()
3137 .is_some_and(|fold_id| requirement.fold_ids.contains(fold_id))
3138 })
3139 .cloned()
3140 .collect::<Vec<_>>();
3141 if selected.is_empty() {
3142 return Err(DagMlError::RuntimeValidation(format!(
3143 "aggregated prediction cache requirement `{}` has no matching prediction blocks",
3144 requirement.key()
3145 )));
3146 }
3147 selected.sort_by(|left, right| {
3148 (
3149 left.fold_id.as_ref().map(ToString::to_string),
3150 left.prediction_id.clone(),
3151 )
3152 .cmp(&(
3153 right.fold_id.as_ref().map(ToString::to_string),
3154 right.prediction_id.clone(),
3155 ))
3156 });
3157 Ok(selected)
3158}
3159
3160fn build_prediction_cache_record_from_selected(
3161 requirement: &BundlePredictionRequirement,
3162 selected: &[PredictionBlock],
3163) -> Result<BundlePredictionCacheRecord> {
3164 requirement.validate()?;
3165 if selected.is_empty() {
3166 return Err(DagMlError::RuntimeValidation(format!(
3167 "prediction cache requirement `{}` has no matching prediction blocks",
3168 requirement.key()
3169 )));
3170 }
3171 let mut fold_ids = BTreeSet::new();
3172 let mut sample_ids = BTreeSet::new();
3173 let mut target_names: Option<Vec<String>> = None;
3174 let mut prediction_width: Option<usize> = None;
3175 let mut row_count = 0usize;
3176 let mut block_records = Vec::new();
3177 for block in selected {
3178 if block.producer_node != requirement.producer_node
3179 || block.partition != requirement.partition
3180 {
3181 return Err(DagMlError::RuntimeValidation(format!(
3182 "prediction cache `{}` contains a block outside the requirement scope",
3183 requirement.key()
3184 )));
3185 }
3186 let width = block.validate_shape()?;
3187 if prediction_width.is_some_and(|expected| expected != width) {
3188 return Err(DagMlError::RuntimeValidation(format!(
3189 "prediction cache `{}` has inconsistent prediction width",
3190 requirement.key()
3191 )));
3192 }
3193 prediction_width = Some(width);
3194 let block_target_names = normalized_prediction_targets(block, width);
3195 if target_names
3196 .as_ref()
3197 .is_some_and(|expected| expected != &block_target_names)
3198 {
3199 return Err(DagMlError::RuntimeValidation(format!(
3200 "prediction cache `{}` has inconsistent target names",
3201 requirement.key()
3202 )));
3203 }
3204 target_names = Some(block_target_names);
3205 if let Some(fold_id) = &block.fold_id {
3206 fold_ids.insert(fold_id.clone());
3207 }
3208 sample_ids.extend(block.sample_ids.iter().cloned());
3209 row_count += block.sample_ids.len();
3210 block_records.push(BundlePredictionBlockCacheRecord {
3211 prediction_id: block.prediction_id.clone(),
3212 fold_id: block.fold_id.clone(),
3213 prediction_level: PredictionLevel::Sample,
3214 row_count: block.sample_ids.len(),
3215 unit_ids: Vec::new(),
3216 sample_ids: block.sample_ids.clone(),
3217 content_fingerprint: stable_json_fingerprint(block)?,
3218 });
3219 }
3220
3221 let record = BundlePredictionCacheRecord {
3222 requirement_key: requirement.key(),
3223 cache_id: format!("prediction-cache:{}", requirement.key()),
3224 cache_namespace_fingerprints: Vec::new(),
3225 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
3226 partition: requirement.partition.clone(),
3227 prediction_level: requirement.prediction_level,
3228 fold_ids: fold_ids.into_iter().collect(),
3229 unit_ids: requirement.unit_ids.clone(),
3230 sample_ids: sample_ids.into_iter().collect(),
3231 prediction_width: prediction_width.unwrap_or_default(),
3232 target_names: target_names.unwrap_or_default(),
3233 block_count: block_records.len(),
3234 row_count,
3235 content_fingerprint: stable_json_fingerprint(selected)?,
3236 blocks: block_records,
3237 };
3238 validate_prediction_cache_matches_requirement(&record, requirement)?;
3239 record.validate()?;
3240 Ok(record)
3241}
3242
3243fn build_aggregated_prediction_cache_record_from_selected(
3244 requirement: &BundlePredictionRequirement,
3245 selected: &[AggregatedPredictionBlock],
3246) -> Result<BundlePredictionCacheRecord> {
3247 requirement.validate()?;
3248 if requirement.prediction_level == PredictionLevel::Sample {
3249 return Err(DagMlError::RuntimeValidation(format!(
3250 "aggregated prediction cache requirement `{}` must use target or group level",
3251 requirement.key()
3252 )));
3253 }
3254 if selected.is_empty() {
3255 return Err(DagMlError::RuntimeValidation(format!(
3256 "aggregated prediction cache requirement `{}` has no matching prediction blocks",
3257 requirement.key()
3258 )));
3259 }
3260 let mut fold_ids = BTreeSet::new();
3261 let mut unit_ids = BTreeSet::new();
3262 let mut target_names: Option<Vec<String>> = None;
3263 let mut prediction_width: Option<usize> = None;
3264 let mut row_count = 0usize;
3265 let mut block_records = Vec::new();
3266 for block in selected {
3267 if block.producer_node != requirement.producer_node
3268 || block.partition != requirement.partition
3269 || block.level != requirement.prediction_level
3270 {
3271 return Err(DagMlError::RuntimeValidation(format!(
3272 "aggregated prediction cache `{}` contains a block outside the requirement scope",
3273 requirement.key()
3274 )));
3275 }
3276 let width = block.validate_shape()?;
3277 if prediction_width.is_some_and(|expected| expected != width) {
3278 return Err(DagMlError::RuntimeValidation(format!(
3279 "aggregated prediction cache `{}` has inconsistent prediction width",
3280 requirement.key()
3281 )));
3282 }
3283 prediction_width = Some(width);
3284 let block_target_names = normalized_aggregated_prediction_targets(block, width);
3285 if target_names
3286 .as_ref()
3287 .is_some_and(|expected| expected != &block_target_names)
3288 {
3289 return Err(DagMlError::RuntimeValidation(format!(
3290 "aggregated prediction cache `{}` has inconsistent target names",
3291 requirement.key()
3292 )));
3293 }
3294 target_names = Some(block_target_names);
3295 if let Some(fold_id) = &block.fold_id {
3296 fold_ids.insert(fold_id.clone());
3297 }
3298 unit_ids.extend(block.unit_ids.iter().cloned());
3299 row_count += block.unit_ids.len();
3300 block_records.push(BundlePredictionBlockCacheRecord {
3301 prediction_id: block.prediction_id.clone(),
3302 fold_id: block.fold_id.clone(),
3303 prediction_level: block.level,
3304 row_count: block.unit_ids.len(),
3305 unit_ids: block.unit_ids.clone(),
3306 sample_ids: Vec::new(),
3307 content_fingerprint: stable_json_fingerprint(block)?,
3308 });
3309 }
3310
3311 let record = BundlePredictionCacheRecord {
3312 requirement_key: requirement.key(),
3313 cache_id: format!("prediction-cache:{}", requirement.key()),
3314 cache_namespace_fingerprints: Vec::new(),
3315 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
3316 partition: requirement.partition.clone(),
3317 prediction_level: requirement.prediction_level,
3318 fold_ids: fold_ids.into_iter().collect(),
3319 unit_ids: unit_ids.into_iter().collect(),
3320 sample_ids: Vec::new(),
3321 prediction_width: prediction_width.unwrap_or_default(),
3322 target_names: target_names.unwrap_or_default(),
3323 block_count: block_records.len(),
3324 row_count,
3325 content_fingerprint: stable_json_fingerprint(selected)?,
3326 blocks: block_records,
3327 };
3328 validate_prediction_cache_matches_requirement(&record, requirement)?;
3329 record.validate()?;
3330 Ok(record)
3331}
3332
3333fn validate_prediction_cache_matches_requirement(
3334 cache: &BundlePredictionCacheRecord,
3335 requirement: &BundlePredictionRequirement,
3336) -> Result<()> {
3337 if cache.requirement_key != requirement.key()
3338 || cache.partition != requirement.partition
3339 || cache.prediction_level != requirement.prediction_level
3340 || cache.fold_ids != requirement.fold_ids
3341 || cache.unit_ids != requirement.unit_ids
3342 || cache.sample_ids != requirement.sample_ids
3343 || cache.prediction_width != requirement.prediction_width
3344 || cache.target_names != requirement.target_names
3345 {
3346 return Err(DagMlError::RuntimeValidation(format!(
3347 "prediction cache `{}` does not match requirement `{}`",
3348 cache.cache_id,
3349 requirement.key()
3350 )));
3351 }
3352 Ok(())
3353}
3354
3355fn normalized_prediction_targets(block: &PredictionBlock, width: usize) -> Vec<String> {
3356 if block.target_names.is_empty() {
3357 (0..width).map(|index| format!("p{index}")).collect()
3358 } else {
3359 block.target_names.clone()
3360 }
3361}
3362
3363fn normalized_aggregated_prediction_targets(
3364 block: &AggregatedPredictionBlock,
3365 width: usize,
3366) -> Vec<String> {
3367 if block.target_names.is_empty() {
3368 (0..width).map(|index| format!("p{index}")).collect()
3369 } else {
3370 block.target_names.clone()
3371 }
3372}
3373
3374fn sample_prediction_units(sample_ids: &[SampleId]) -> Vec<PredictionUnitId> {
3375 sample_ids
3376 .iter()
3377 .cloned()
3378 .map(PredictionUnitId::Sample)
3379 .collect()
3380}
3381
3382fn validate_prediction_units(
3383 label: &str,
3384 expected_level: PredictionLevel,
3385 unit_ids: &[PredictionUnitId],
3386) -> Result<()> {
3387 validate_unique_ids(label, unit_ids)?;
3388 for unit_id in unit_ids {
3389 if unit_id.level() != expected_level {
3390 return Err(DagMlError::RuntimeValidation(format!(
3391 "{label} `{unit_id}` does not match prediction level {:?}",
3392 expected_level
3393 )));
3394 }
3395 }
3396 Ok(())
3397}
3398
3399fn validate_fingerprint(label: &str, value: &str) -> Result<()> {
3400 if value.len() != 64 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
3401 return Err(DagMlError::RuntimeValidation(format!(
3402 "{label} fingerprint must be a 64-character hex digest"
3403 )));
3404 }
3405 Ok(())
3406}
3407
3408fn validate_non_empty(label: &str, value: &str) -> Result<()> {
3409 if value.trim().is_empty() {
3410 return Err(DagMlError::RuntimeValidation(format!("{label} is empty")));
3411 }
3412 Ok(())
3413}
3414
3415fn validate_unique_ids<T>(label: &str, values: &[T]) -> Result<()>
3416where
3417 T: Ord + ToString,
3418{
3419 let mut seen = BTreeSet::new();
3420 for value in values {
3421 if !seen.insert(value) {
3422 return Err(DagMlError::RuntimeValidation(format!(
3423 "duplicate {label} `{}`",
3424 value.to_string()
3425 )));
3426 }
3427 }
3428 Ok(())
3429}
3430
3431#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
3432pub struct ReplayPhaseRequest {
3433 pub bundle_id: BundleId,
3434 pub phase: Phase,
3435 #[serde(default)]
3436 pub data_envelope_keys: Vec<String>,
3437}
3438
3439impl ReplayPhaseRequest {
3440 pub fn validate_for_bundle(&self, bundle: &ExecutionBundle) -> Result<()> {
3441 self.validate_for_bundle_with_prediction_cache_store(bundle, false)
3442 }
3443
3444 pub fn validate_for_bundle_with_prediction_cache_store(
3445 &self,
3446 bundle: &ExecutionBundle,
3447 prediction_cache_available: bool,
3448 ) -> Result<()> {
3449 self.validate_for_bundle_internal(bundle, prediction_cache_available)
3450 }
3451
3452 pub fn validate_for_bundle_with_prediction_cache_payloads(
3453 &self,
3454 bundle: &ExecutionBundle,
3455 prediction_cache_payloads: Option<&BundlePredictionCachePayloadSet>,
3456 ) -> Result<()> {
3457 if let Some(payloads) = prediction_cache_payloads {
3458 payloads.validate_against_bundle(bundle)?;
3459 }
3460 self.validate_for_bundle_internal(bundle, prediction_cache_payloads.is_some())
3461 }
3462
3463 fn validate_for_bundle_internal(
3464 &self,
3465 bundle: &ExecutionBundle,
3466 prediction_cache_available: bool,
3467 ) -> Result<()> {
3468 bundle.validate()?;
3469 if self.bundle_id != bundle.bundle_id {
3470 return Err(DagMlError::RuntimeValidation(format!(
3471 "replay request bundle `{}` does not match bundle `{}`",
3472 self.bundle_id, bundle.bundle_id
3473 )));
3474 }
3475 if !matches!(self.phase, Phase::Predict | Phase::Explain | Phase::Refit) {
3476 return Err(DagMlError::RuntimeValidation(format!(
3477 "bundle replay phase {:?} is not supported",
3478 self.phase
3479 )));
3480 }
3481 if self.phase == Phase::Refit && !bundle.prediction_requirements.is_empty() {
3482 if prediction_cache_available {
3483 return self.validate_data_envelope_keys(bundle);
3484 }
3485 return Err(DagMlError::RuntimeValidation(format!(
3486 "bundle `{}` cannot replay REFIT because it depends on {} OOF prediction requirement(s) but stores only prediction cache manifests",
3487 bundle.bundle_id,
3488 bundle.prediction_requirements.len()
3489 )));
3490 }
3491 self.validate_data_envelope_keys(bundle)
3492 }
3493
3494 fn validate_data_envelope_keys(&self, bundle: &ExecutionBundle) -> Result<()> {
3495 let expected = bundle
3496 .data_requirements
3497 .iter()
3498 .map(BundleDataRequirement::key)
3499 .collect::<BTreeSet<_>>();
3500 let mut requested = BTreeSet::new();
3501 for key in &self.data_envelope_keys {
3502 if key.trim().is_empty() {
3503 return Err(DagMlError::RuntimeValidation(
3504 "replay request contains an empty data envelope key".to_string(),
3505 ));
3506 }
3507 if !requested.insert(key.as_str()) {
3508 return Err(DagMlError::RuntimeValidation(format!(
3509 "replay request contains duplicate data envelope key `{key}`"
3510 )));
3511 }
3512 if !expected.contains(key.as_str()) {
3513 return Err(DagMlError::RuntimeValidation(format!(
3514 "replay request references unknown data envelope key `{key}`"
3515 )));
3516 }
3517 }
3518 for requirement in &bundle.data_requirements {
3519 let key = requirement.key();
3520 if !requested.contains(key.as_str()) {
3521 return Err(DagMlError::RuntimeValidation(format!(
3522 "replay request is missing data envelope key `{key}`"
3523 )));
3524 }
3525 }
3526 Ok(())
3527 }
3528}
3529
3530#[cfg(test)]
3531mod tests {
3532 use super::*;
3533 use crate::controller::{ControllerManifest, ControllerRegistry};
3534 use crate::data::{
3535 AggregateRepresentation, RepresentationCardinality, RepresentationCompatibilityOutcome,
3536 RepresentationCompatibilityReport, RepresentationMissingSourcePolicy, RepresentationPlan,
3537 RepresentationReplayManifest,
3538 };
3539 use crate::dsl::{compile_pipeline_dsl_with_generation, PipelineDslSpec};
3540 use crate::graph::GraphSpec;
3541 use crate::ids::{ArtifactId, ControllerId, FoldId, LineageId, RunId, SampleId, TargetId};
3542 use crate::plan::{build_execution_plan, CampaignSpec};
3543 use crate::relation::EntityUnitLevel;
3544 use crate::selection::{
3545 select_candidate, CandidateScore, MetricObjective, SelectionMetric, SelectionPolicy,
3546 };
3547
3548 fn plan() -> ExecutionPlan {
3549 let graph: GraphSpec =
3550 serde_json::from_str(include_str!("../tests/fixtures/package/minimal_graph.json"))
3551 .unwrap();
3552 let campaign: CampaignSpec = serde_json::from_str(include_str!(
3553 "../tests/fixtures/package/campaign_oof_generation.json"
3554 ))
3555 .unwrap();
3556 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3557 "../tests/fixtures/package/controller_manifests.json"
3558 ))
3559 .unwrap();
3560 let mut registry = ControllerRegistry::new();
3561 for manifest in manifests {
3562 registry.register(manifest).unwrap();
3563 }
3564 build_execution_plan("plan:bundle", graph, campaign, ®istry).unwrap()
3565 }
3566
3567 fn methods_hpo_resume_state() -> MethodsHpoResumeState {
3568 let node = NodeId::new("model:hpo").unwrap();
3569 let variant = crate::generation::VariantPlan {
3570 variant_id: VariantId::new("hpo:trial:1").unwrap(),
3571 choices: BTreeMap::new(),
3572 fingerprint: "c".repeat(64),
3573 seed: Some(7),
3574 };
3575 let sample = SampleId::new("sample:hpo.1").unwrap();
3576 let target = RegressionTargetBlock {
3577 level: PredictionLevel::Sample,
3578 unit_ids: vec![PredictionUnitId::Sample(sample.clone())],
3579 values: vec![vec![1.0]],
3580 target_names: vec!["y".to_string()],
3581 };
3582 let report = RegressionMetricReport {
3583 prediction_id: Some("prediction:model:hpo.avg".to_string()),
3584 producer_node: node.clone(),
3585 producer_port: Some("prediction".to_string()),
3586 variant_id: Some(variant.variant_id.clone()),
3587 variant_label: None,
3588 partition: PredictionPartition::Validation,
3589 fold_id: Some(FoldId::new("avg").unwrap()),
3590 level: PredictionLevel::Sample,
3591 row_count: 1,
3592 target_width: 1,
3593 target_names: vec!["y".to_string()],
3594 metrics: BTreeMap::from([("rmse".to_string(), 0.25)]),
3595 };
3596 let payload = vec![1_u8, 2, 3];
3597 MethodsHpoResumeState {
3598 schema_version: METHODS_HPO_RESUME_STATE_SCHEMA_VERSION,
3599 checkpoint: crate::hpo::N4moptCheckpointArtifact {
3600 schema_version: crate::hpo::N4MOPT_CHECKPOINT_SCHEMA_VERSION,
3601 artifact_kind: crate::hpo::N4MOPT_ARTIFACT_KIND.to_string(),
3602 format: crate::hpo::N4MOPT_FORMAT.to_string(),
3603 abi_major: crate::hpo::METHODS_ABI_MAJOR,
3604 abi_min_minor: crate::hpo::METHODS_N4MOPT_MIN_ABI_MINOR,
3605 binding: crate::hpo::HpoStudyBinding {
3606 controller_id: "controller:hpo".to_string(),
3607 study_id: "study:hpo".to_string(),
3608 search_space_fingerprint: "a".repeat(64),
3609 optimizer_fingerprint: "b".repeat(64),
3610 },
3611 methods_abi: "n4m:1".to_string(),
3612 payload_sha256: format!("{:x}", sha2::Sha256::digest(&payload)),
3613 opaque_payload: payload,
3614 },
3615 provenance: MethodsHpoResumeProvenance {
3616 graph_fingerprint: "d".repeat(64),
3617 campaign_fingerprint: "e".repeat(64),
3618 controller_fingerprint: "f".repeat(64),
3619 data_identities_fingerprint: "1".repeat(64),
3620 fold_set_fingerprint: "2".repeat(64),
3621 training_influence_fingerprint: "3".repeat(64),
3622 relation_fingerprint: "4".repeat(64),
3623 selection: MethodsHpoResumeSelection {
3624 selection_id: "selection:hpo".to_string(),
3625 target_node_id: node.clone(),
3626 producer_port: "prediction".to_string(),
3627 metric: "rmse".to_string(),
3628 },
3629 },
3630 operation_id: "hpo:bundle-test".to_string(),
3631 controller_id: ControllerId::new("controller:hpo").unwrap(),
3632 target_node_id: node.clone(),
3633 incumbent: MethodsHpoNativeIncumbent {
3634 trial_id: 1,
3635 score: 0.25,
3636 metric: "rmse".to_string(),
3637 direction: crate::hpo::HpoDirection::Minimize,
3638 variant_id: variant.variant_id.clone(),
3639 },
3640 trial_history_len: 1,
3641 terminal_trials: vec![MethodsHpoTerminalEvidence {
3642 trial: crate::hpo::HpoTrial {
3643 id: 1,
3644 ask_sequence: 1,
3645 terminal_sequence: Some(1),
3646 parameters: BTreeMap::new(),
3647 parameter_order: Vec::new(),
3648 status: crate::hpo::HpoTrialStatus::Completed,
3649 score: Some(0.25),
3650 rung: 0,
3651 duration: 0.0,
3652 intermediates: Vec::new(),
3653 failure: None,
3654 },
3655 variant_id: Some(variant.variant_id.clone()),
3656 }],
3657 completed_proposals: vec![MethodsHpoCompletedProposal {
3658 trial_id: 1,
3659 variant: variant.clone(),
3660 }],
3661 completed_reports: vec![MethodsHpoCompletedReport {
3662 trial_id: 1,
3663 variant_id: variant.variant_id.clone(),
3664 terminal_state: MethodsHpoCompletedState::Completed,
3665 score: 0.25,
3666 report,
3667 }],
3668 candidates: vec![MethodsHpoCandidateEvidence {
3669 trial_id: 1,
3670 variant_id: variant.variant_id.clone(),
3671 variant_label: None,
3672 predictions: vec![PredictionBlock {
3673 prediction_id: Some("prediction:model:hpo.fold0".to_string()),
3674 producer_node: node.clone(),
3675 producer_port: Some("prediction".to_string()),
3676 partition: PredictionPartition::Validation,
3677 fold_id: Some(FoldId::new("fold:0").unwrap()),
3678 sample_ids: vec![sample.clone()],
3679 values: vec![vec![1.25]],
3680 target_names: vec!["y".to_string()],
3681 }],
3682 regression_targets: vec![target.clone()],
3683 oof_average: MethodsHpoOofAverage {
3684 predictions: AggregatedPredictionBlock {
3685 prediction_id: Some("prediction:model:hpo.avg".to_string()),
3686 producer_node: node.clone(),
3687 producer_port: Some("prediction".to_string()),
3688 partition: PredictionPartition::Validation,
3689 fold_id: Some(FoldId::new("avg").unwrap()),
3690 level: PredictionLevel::Sample,
3691 unit_ids: vec![PredictionUnitId::Sample(sample)],
3692 values: vec![vec![1.25]],
3693 target_names: vec!["y".to_string()],
3694 },
3695 y_true: target,
3696 },
3697 lineage: vec![crate::runtime::LineageRecord {
3698 record_id: LineageId::new("lineage:hpo.1").unwrap(),
3699 run_id: RunId::new("run:hpo").unwrap(),
3700 node_id: node,
3701 phase: Phase::FitCv,
3702 controller_id: ControllerId::new("controller:methods.pls").unwrap(),
3703 controller_version: "1".to_string(),
3704 variant_id: Some(variant.variant_id),
3705 fold_id: Some(FoldId::new("fold:0").unwrap()),
3706 branch_path: Vec::new(),
3707 input_lineage: Vec::new(),
3708 artifact_refs: Vec::new(),
3709 params_fingerprint: "5".repeat(64),
3710 data_model_shape_fingerprint: None,
3711 aggregation_policy_fingerprint: None,
3712 seed: Some(7),
3713 unsafe_flags: BTreeSet::new(),
3714 metrics: BTreeMap::new(),
3715 loss_attestations: Vec::new(),
3716 early_stopping_records: Vec::new(),
3717 }],
3718 }],
3719 }
3720 }
3721
3722 fn branch_merge_plan() -> ExecutionPlan {
3723 let graph: GraphSpec = serde_json::from_str(include_str!(
3724 "../tests/fixtures/package/branch_merge_oof_graph.json"
3725 ))
3726 .unwrap();
3727 let campaign: CampaignSpec = serde_json::from_str(include_str!(
3728 "../tests/fixtures/package/campaign_branch_merge_oof.json"
3729 ))
3730 .unwrap();
3731 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3732 "../tests/fixtures/package/controller_manifests.json"
3733 ))
3734 .unwrap();
3735 let mut registry = ControllerRegistry::new();
3736 for manifest in manifests {
3737 registry.register(manifest).unwrap();
3738 }
3739 build_execution_plan("plan:branch.merge.bundle", graph, campaign, ®istry).unwrap()
3740 }
3741
3742 fn separation_concat_merge_plan() -> ExecutionPlan {
3749 let graph: GraphSpec = serde_json::from_str(include_str!(
3750 "../tests/fixtures/package/separation_branch_concat_merge_oof_graph.json"
3751 ))
3752 .unwrap();
3753 let campaign: CampaignSpec = serde_json::from_str(include_str!(
3754 "../tests/fixtures/package/campaign_separation_branch_concat_merge_oof.json"
3755 ))
3756 .unwrap();
3757 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3758 "../tests/fixtures/package/controller_manifests.json"
3759 ))
3760 .unwrap();
3761 let mut registry = ControllerRegistry::new();
3762 for manifest in manifests {
3763 registry.register(manifest).unwrap();
3764 }
3765 build_execution_plan(
3766 "plan:separation.concat.merge.bundle",
3767 graph,
3768 campaign,
3769 ®istry,
3770 )
3771 .unwrap()
3772 }
3773
3774 fn separation_branch_requirement(
3778 producer_node: &str,
3779 partition_samples: &[&str],
3780 partition_folds: &[&str],
3781 ) -> BundlePredictionRequirement {
3782 BundlePredictionRequirement {
3783 producer_node: NodeId::new(producer_node).unwrap(),
3784 source_port: "oof".to_string(),
3785 consumer_node: NodeId::new("merge:sites").unwrap(),
3786 target_port: format!("oof_{producer_node}"),
3787 partition: PredictionPartition::Validation,
3788 prediction_level: PredictionLevel::Sample,
3789 fold_ids: partition_folds
3790 .iter()
3791 .map(|f| FoldId::new(*f).unwrap())
3792 .collect(),
3793 unit_ids: Vec::new(),
3794 sample_ids: partition_samples
3795 .iter()
3796 .map(|s| SampleId::new(*s).unwrap())
3797 .collect(),
3798 prediction_width: 1,
3799 target_names: vec!["y".to_string()],
3800 }
3801 }
3802
3803 fn separation_branch_blocks(
3806 producer_node: &str,
3807 fold0_sample: &str,
3808 fold1_sample: &str,
3809 offset: f64,
3810 ) -> Vec<PredictionBlock> {
3811 let producer_node = NodeId::new(producer_node).unwrap();
3812 vec![
3813 PredictionBlock {
3814 prediction_id: Some(format!("prediction:{producer_node}:fold0")),
3815 producer_node: producer_node.clone(),
3816 producer_port: Some("pred".to_string()),
3817 partition: PredictionPartition::Validation,
3818 fold_id: Some(FoldId::new("fold:0").unwrap()),
3819 sample_ids: vec![SampleId::new(fold0_sample).unwrap()],
3820 values: vec![vec![offset + 0.1]],
3821 target_names: vec!["y".to_string()],
3822 },
3823 PredictionBlock {
3824 prediction_id: Some(format!("prediction:{producer_node}:fold1")),
3825 producer_node,
3826 producer_port: Some("pred".to_string()),
3827 partition: PredictionPartition::Validation,
3828 fold_id: Some(FoldId::new("fold:1").unwrap()),
3829 sample_ids: vec![SampleId::new(fold1_sample).unwrap()],
3830 values: vec![vec![offset + 0.2]],
3831 target_names: vec!["y".to_string()],
3832 },
3833 ]
3834 }
3835
3836 fn executable_dsl_plan() -> ExecutionPlan {
3837 let spec: PipelineDslSpec = serde_json::from_str(include_str!(
3838 "../tests/fixtures/package/pipeline_dsl_branch_merge_executable.json"
3839 ))
3840 .unwrap();
3841 let compiled = compile_pipeline_dsl_with_generation(&spec).unwrap();
3842 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3843 "../tests/fixtures/package/controller_manifests.json"
3844 ))
3845 .unwrap();
3846 let mut registry = ControllerRegistry::new();
3847 for manifest in manifests {
3848 registry.register(manifest).unwrap();
3849 }
3850 build_execution_plan(
3851 "plan:dsl.branch.merge.bundle",
3852 compiled.graph,
3853 compiled.campaign_template,
3854 ®istry,
3855 )
3856 .unwrap()
3857 }
3858
3859 fn branch_merge_selection_decisions() -> BTreeMap<String, SelectionDecision> {
3860 serde_json::from_str(include_str!(
3861 "../tests/fixtures/package/bundle/selection_decisions_branch_merge.json"
3862 ))
3863 .unwrap()
3864 }
3865
3866 fn refit_artifact(
3867 plan: &ExecutionPlan,
3868 node_id: &str,
3869 data_requirement_keys: Vec<String>,
3870 prediction_requirement_keys: Vec<String>,
3871 ) -> RefitArtifactRecord {
3872 let node_id = NodeId::new(node_id).unwrap();
3873 let node_plan = plan.node_plans.get(&node_id).unwrap();
3874 RefitArtifactRecord {
3875 node_id: node_plan.node_id.clone(),
3876 controller_id: node_plan.controller_id.clone(),
3877 artifact: ArtifactRef {
3878 id: ArtifactId::new(format!("artifact:{}:refit", node_plan.node_id)).unwrap(),
3879 kind: "mock_model".to_string(),
3880 controller_id: node_plan.controller_id.clone(),
3881 backend: None,
3882 uri: None,
3883 content_fingerprint: None,
3884 size_bytes: Some(128),
3885 plugin: None,
3886 plugin_version: None,
3887 abi_major: None,
3888 abi_min_minor: None,
3889 native_predictor_descriptor: None,
3890 },
3891 params_fingerprint: node_plan.params_fingerprint.clone(),
3892 training_loss_fingerprint: node_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
3893 data_requirement_keys,
3894 prediction_requirement_keys,
3895 }
3896 }
3897
3898 fn branch_merge_samples() -> Vec<SampleId> {
3899 vec![
3900 SampleId::new("sample:1").unwrap(),
3901 SampleId::new("sample:2").unwrap(),
3902 SampleId::new("sample:3").unwrap(),
3903 SampleId::new("sample:4").unwrap(),
3904 ]
3905 }
3906
3907 fn branch_merge_requirement(
3908 producer_node: &str,
3909 target_port: &str,
3910 ) -> BundlePredictionRequirement {
3911 BundlePredictionRequirement {
3912 producer_node: NodeId::new(producer_node).unwrap(),
3913 source_port: "oof".to_string(),
3914 consumer_node: NodeId::new("merge:stack.pred_plus_original.meta:ridge").unwrap(),
3915 target_port: target_port.to_string(),
3916 partition: PredictionPartition::Validation,
3917 prediction_level: PredictionLevel::Sample,
3918 fold_ids: vec![
3919 FoldId::new("fold:0").unwrap(),
3920 FoldId::new("fold:1").unwrap(),
3921 ],
3922 unit_ids: Vec::new(),
3923 sample_ids: branch_merge_samples(),
3924 prediction_width: 1,
3925 target_names: vec!["y".to_string()],
3926 }
3927 }
3928
3929 fn branch_merge_prediction_blocks(producer_node: &str, offset: f64) -> Vec<PredictionBlock> {
3930 let producer_node = NodeId::new(producer_node).unwrap();
3931 let samples = branch_merge_samples();
3932 vec![
3933 PredictionBlock {
3934 prediction_id: Some(format!("prediction:{producer_node}:fold0")),
3935 producer_node: producer_node.clone(),
3936 producer_port: Some("pred".to_string()),
3937 partition: PredictionPartition::Validation,
3938 fold_id: Some(FoldId::new("fold:0").unwrap()),
3939 sample_ids: samples[0..2].to_vec(),
3940 values: vec![vec![offset + 0.1], vec![offset + 0.2]],
3941 target_names: vec!["y".to_string()],
3942 },
3943 PredictionBlock {
3944 prediction_id: Some(format!("prediction:{producer_node}:fold1")),
3945 producer_node,
3946 producer_port: Some("pred".to_string()),
3947 partition: PredictionPartition::Validation,
3948 fold_id: Some(FoldId::new("fold:1").unwrap()),
3949 sample_ids: samples[2..4].to_vec(),
3950 values: vec![vec![offset + 0.3], vec![offset + 0.4]],
3951 target_names: vec!["y".to_string()],
3952 },
3953 ]
3954 }
3955
3956 fn decision() -> SelectionDecision {
3957 select_candidate(
3958 &SelectionPolicy {
3959 id: "select:merge".to_string(),
3960 metric: SelectionMetric {
3961 name: "rmse".to_string(),
3962 objective: MetricObjective::Minimize,
3963 },
3964 required_metric_level: Some(crate::policy::PredictionLevel::Sample),
3965 require_finite: true,
3966 evaluation_scope: None,
3967 refit_slot_plan: None,
3968 stacking_fit_contract: None,
3969 reduction_id: None,
3970 },
3971 &[
3972 CandidateScore {
3973 candidate_id: "model:base".to_string(),
3974 metrics: BTreeMap::from([("rmse".to_string(), 1.0)]),
3975 metadata: BTreeMap::from([(
3976 "metric_level".to_string(),
3977 serde_json::Value::String("sample".to_string()),
3978 )]),
3979 },
3980 CandidateScore {
3981 candidate_id: "model:other".to_string(),
3982 metrics: BTreeMap::from([("rmse".to_string(), 2.0)]),
3983 metadata: BTreeMap::from([(
3984 "metric_level".to_string(),
3985 serde_json::Value::String("sample".to_string()),
3986 )]),
3987 },
3988 ],
3989 )
3990 .unwrap()
3991 }
3992
3993 fn selected_model_base_decision() -> SelectionDecision {
3994 decision()
3995 }
3996
3997 fn model_base_refit_artifact(plan: &ExecutionPlan) -> RefitArtifactRecord {
3998 let model_plan = plan
3999 .node_plans
4000 .get(&NodeId::new("model:base").unwrap())
4001 .unwrap();
4002 RefitArtifactRecord {
4003 node_id: model_plan.node_id.clone(),
4004 controller_id: model_plan.controller_id.clone(),
4005 artifact: ArtifactRef {
4006 id: ArtifactId::new("artifact:model:base:refit").unwrap(),
4007 kind: "sklearn_pickle".to_string(),
4008 controller_id: model_plan.controller_id.clone(),
4009 backend: None,
4010 uri: None,
4011 content_fingerprint: None,
4012 size_bytes: Some(128),
4013 plugin: None,
4014 plugin_version: None,
4015 abi_major: None,
4016 abi_min_minor: None,
4017 native_predictor_descriptor: None,
4018 },
4019 params_fingerprint: model_plan.params_fingerprint.clone(),
4020 training_loss_fingerprint: model_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
4021 data_requirement_keys: vec!["model:base.x".to_string()],
4022 prediction_requirement_keys: Vec::new(),
4023 }
4024 }
4025
4026 #[test]
4027 fn builds_bundle_from_execution_plan() {
4028 let plan = plan();
4029 let artifact = model_base_refit_artifact(&plan);
4030
4031 let bundle = build_execution_bundle(
4032 BundleId::new("bundle:demo").unwrap(),
4033 &plan,
4034 Some(plan.variants[0].variant_id.clone()),
4035 BTreeMap::from([("merge".to_string(), decision())]),
4036 vec![artifact],
4037 )
4038 .unwrap();
4039
4040 bundle.validate_against_plan(&plan).unwrap();
4041 assert_eq!(bundle.data_requirements.len(), 1);
4042 }
4043
4044 #[test]
4045 fn bundle_data_requirements_accept_d7_replay_contracts() {
4046 let plan = plan();
4047 let artifact = model_base_refit_artifact(&plan);
4048 let mut bundle = build_execution_bundle(
4049 BundleId::new("bundle:d7.replay").unwrap(),
4050 &plan,
4051 Some(plan.variants[0].variant_id.clone()),
4052 BTreeMap::from([("merge".to_string(), decision())]),
4053 vec![artifact],
4054 )
4055 .unwrap();
4056 let relation_fingerprint = bundle.data_requirements[0]
4057 .relation_fingerprint
4058 .clone()
4059 .unwrap_or_else(|| "a".repeat(64));
4060 bundle.data_requirements[0].representation_replay_manifest =
4061 Some(RepresentationReplayManifest {
4062 manifest_id: "repr:d7.bundle".to_string(),
4063 representation_plan: RepresentationPlan::Aggregate(AggregateRepresentation {
4064 input_unit_level: EntityUnitLevel::Observation,
4065 output_unit_level: EntityUnitLevel::PhysicalSample,
4066 reducer_id: None,
4067 method: Some("mean".to_string()),
4068 cardinality: RepresentationCardinality::ManyToOne,
4069 }),
4070 combination_plan: None,
4071 output_unit_level: EntityUnitLevel::PhysicalSample,
4072 output_representation: Some("tabular_numeric".to_string()),
4073 relation_fingerprint: Some(relation_fingerprint.clone()),
4074 feature_schema_fingerprint: Some("b".repeat(64)),
4075 final_reduction_id: None,
4076 sample_observation_mapping: Vec::new(),
4077 combo_selection: Vec::new(),
4078 qc_policy_refs: Vec::new(),
4079 outlier_policy_refs: Vec::new(),
4080 missing_source_policy: None,
4081 missing_repetition_policy: None,
4082 prediction_representation: None,
4083 final_output_unit_level: Some(EntityUnitLevel::PhysicalSample),
4084 train_compatibility: None,
4085 predict_compatibility: None,
4086 metadata: BTreeMap::new(),
4087 });
4088 bundle.data_requirements[0].representation_compatibility =
4089 Some(RepresentationCompatibilityReport {
4090 policy: RepresentationMissingSourcePolicy::Strict,
4091 outcome: RepresentationCompatibilityOutcome::Compatible,
4092 fallback_used: None,
4093 warning_severity: None,
4094 affected_source_count: 0,
4095 affected_repetition_count: 0,
4096 affected_sample_count: 0,
4097 train_relation_fingerprint: Some(relation_fingerprint),
4098 predict_relation_fingerprint: None,
4099 train_unit_count: Some(2),
4100 predict_unit_count: Some(2),
4101 fixed_width_required: false,
4102 final_reducer_stabilizes_output: true,
4103 cartesian_combo_count_changed: false,
4104 late_fusion_branch_delta: false,
4105 messages: Vec::new(),
4106 metadata: BTreeMap::new(),
4107 });
4108 bundle.validate_against_plan(&plan).unwrap();
4109
4110 bundle.data_requirements[0]
4111 .representation_replay_manifest
4112 .as_mut()
4113 .unwrap()
4114 .relation_fingerprint = Some("c".repeat(64));
4115 if bundle.data_requirements[0].relation_fingerprint.is_some() {
4116 assert!(bundle.validate().is_err());
4117 }
4118 }
4119
4120 #[test]
4121 fn d9_negative_prediction_cache_refuses_missing_aggregated_unit_ids() {
4122 let cache = BundlePredictionCacheRecord {
4123 requirement_key: "model:base.oof->model:meta.pred".to_string(),
4124 cache_id: "prediction-cache:d9.missing-units".to_string(),
4125 cache_namespace_fingerprints: Vec::new(),
4126 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
4127 partition: PredictionPartition::Validation,
4128 prediction_level: PredictionLevel::Target,
4129 fold_ids: vec![FoldId::new("fold:0").unwrap()],
4130 unit_ids: Vec::new(),
4131 sample_ids: Vec::new(),
4132 prediction_width: 1,
4133 target_names: vec!["y".to_string()],
4134 block_count: 1,
4135 row_count: 1,
4136 content_fingerprint: "d".repeat(64),
4137 blocks: vec![BundlePredictionBlockCacheRecord {
4138 prediction_id: Some("prediction:d9.target.fold0".to_string()),
4139 fold_id: Some(FoldId::new("fold:0").unwrap()),
4140 prediction_level: PredictionLevel::Target,
4141 row_count: 1,
4142 unit_ids: vec![PredictionUnitId::Target(TargetId::new("target:a").unwrap())],
4143 sample_ids: Vec::new(),
4144 content_fingerprint: "e".repeat(64),
4145 }],
4146 };
4147
4148 let error = cache.validate().unwrap_err().to_string();
4149 assert!(
4150 error.contains("row_count does not match unique unit ids"),
4151 "unexpected D9 missing-unit-id cache error: {error}"
4152 );
4153 }
4154
4155 #[test]
4156 fn refit_artifact_validation_checks_portable_artifact_metadata() {
4157 let plan = plan();
4158 let mut artifact = model_base_refit_artifact(&plan);
4159 artifact.artifact.backend = Some(crate::runtime::ArtifactBackend::Joblib);
4160 artifact.artifact.uri = Some("artifacts/model.joblib".to_string());
4161 artifact.artifact.content_fingerprint = Some("c".repeat(64));
4162 artifact.artifact.plugin = Some("dagml.sklearn".to_string());
4163 artifact.artifact.plugin_version = Some("1.0.0".to_string());
4164 artifact.validate().unwrap();
4165
4166 artifact.artifact.content_fingerprint = Some("short".to_string());
4167 assert!(artifact
4168 .validate()
4169 .unwrap_err()
4170 .to_string()
4171 .contains("artifact content fingerprint"));
4172 }
4173
4174 #[test]
4175 fn bundle_selections_must_match_plan_and_refit_artifacts() {
4176 let plan = plan();
4177 let artifact = model_base_refit_artifact(&plan);
4178 let valid = build_execution_bundle(
4179 BundleId::new("bundle:selected.model").unwrap(),
4180 &plan,
4181 Some(plan.variants[0].variant_id.clone()),
4182 BTreeMap::from([("model".to_string(), selected_model_base_decision())]),
4183 vec![artifact.clone()],
4184 )
4185 .unwrap();
4186 valid.validate_against_plan(&plan).unwrap();
4187
4188 assert!(build_execution_bundle(
4189 BundleId::new("bundle:selected.model.missing.artifact").unwrap(),
4190 &plan,
4191 Some(plan.variants[0].variant_id.clone()),
4192 BTreeMap::from([("model".to_string(), selected_model_base_decision())]),
4193 Vec::new(),
4194 )
4195 .is_err());
4196
4197 let mut missing_level = selected_model_base_decision();
4198 missing_level.metric_level = None;
4199 assert!(build_execution_bundle(
4200 BundleId::new("bundle:selected.missing.level").unwrap(),
4201 &plan,
4202 Some(plan.variants[0].variant_id.clone()),
4203 BTreeMap::from([("model".to_string(), missing_level)]),
4204 vec![artifact.clone()],
4205 )
4206 .is_err());
4207
4208 let mut wrong_level = selected_model_base_decision();
4209 wrong_level.metric_level = Some(crate::policy::PredictionLevel::Target);
4210 assert!(build_execution_bundle(
4211 BundleId::new("bundle:selected.wrong.level").unwrap(),
4212 &plan,
4213 Some(plan.variants[0].variant_id.clone()),
4214 BTreeMap::from([("model".to_string(), wrong_level)]),
4215 vec![artifact.clone()],
4216 )
4217 .is_err());
4218
4219 let mut unknown = selected_model_base_decision();
4220 unknown.selected_candidate_id = "model:missing".to_string();
4221 unknown.ranked_candidates[0].candidate_id = "model:missing".to_string();
4222 assert!(build_execution_bundle(
4223 BundleId::new("bundle:selected.unknown").unwrap(),
4224 &plan,
4225 Some(plan.variants[0].variant_id.clone()),
4226 BTreeMap::from([("model".to_string(), unknown)]),
4227 vec![artifact],
4228 )
4229 .is_err());
4230 }
4231
4232 #[test]
4233 fn bundle_artifact_params_follow_selected_generation_variant() {
4234 let plan = executable_dsl_plan();
4235 let selected_variant = &plan.variants[0];
4236 let node_plan = plan
4237 .node_plans
4238 .get(&NodeId::new("branch:b0.model:ridge").unwrap())
4239 .unwrap();
4240 let effective_params = selected_variant
4241 .effective_params_for_node(&node_plan.node_id, &node_plan.params)
4242 .unwrap();
4243 let effective_fingerprint = stable_json_fingerprint(&effective_params).unwrap();
4244 assert_ne!(effective_fingerprint, node_plan.params_fingerprint);
4245
4246 let artifact = RefitArtifactRecord {
4247 node_id: node_plan.node_id.clone(),
4248 controller_id: node_plan.controller_id.clone(),
4249 artifact: ArtifactRef {
4250 id: ArtifactId::new("artifact:branch:b0.model:ridge:refit").unwrap(),
4251 kind: "mock_model".to_string(),
4252 controller_id: node_plan.controller_id.clone(),
4253 backend: None,
4254 uri: None,
4255 content_fingerprint: None,
4256 size_bytes: Some(128),
4257 plugin: None,
4258 plugin_version: None,
4259 abi_major: None,
4260 abi_min_minor: None,
4261 native_predictor_descriptor: None,
4262 },
4263 params_fingerprint: effective_fingerprint,
4264 training_loss_fingerprint: node_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
4265 data_requirement_keys: vec!["branch:b0.model:ridge.x".to_string()],
4266 prediction_requirement_keys: Vec::new(),
4267 };
4268
4269 build_execution_bundle(
4270 BundleId::new("bundle:dsl.variant.params").unwrap(),
4271 &plan,
4272 Some(selected_variant.variant_id.clone()),
4273 BTreeMap::new(),
4274 vec![artifact.clone()],
4275 )
4276 .unwrap();
4277
4278 let mut stale_artifact = artifact;
4279 stale_artifact.params_fingerprint = node_plan.params_fingerprint.clone();
4280 let error = build_execution_bundle(
4281 BundleId::new("bundle:dsl.variant.params.stale").unwrap(),
4282 &plan,
4283 Some(selected_variant.variant_id.clone()),
4284 BTreeMap::new(),
4285 vec![stale_artifact],
4286 )
4287 .unwrap_err();
4288 assert!(format!("{error}").contains("artifact params"));
4289 }
4290
4291 #[test]
4292 fn branch_merge_bundle_links_selected_refits_and_fold_aligned_oof_caches() {
4293 let plan = branch_merge_plan();
4294 let b0_requirement = branch_merge_requirement("branch:b0.model:ridge", "b0_oof");
4295 let b1_requirement = branch_merge_requirement("branch:b1.model:rf", "b1_oof");
4296 let b0_cache = build_prediction_cache_record(
4297 &b0_requirement,
4298 &branch_merge_prediction_blocks("branch:b0.model:ridge", 0.0),
4299 )
4300 .unwrap();
4301 let b1_cache = build_prediction_cache_record(
4302 &b1_requirement,
4303 &branch_merge_prediction_blocks("branch:b1.model:rf", 1.0),
4304 )
4305 .unwrap();
4306 let b0_artifact = refit_artifact(
4307 &plan,
4308 "branch:b0.model:ridge",
4309 vec!["branch:b0.model:ridge.x".to_string()],
4310 Vec::new(),
4311 );
4312 let b1_artifact = refit_artifact(
4313 &plan,
4314 "branch:b1.model:rf",
4315 vec!["branch:b1.model:rf.x".to_string()],
4316 Vec::new(),
4317 );
4318 let merge_artifact = refit_artifact(
4319 &plan,
4320 "merge:stack.pred_plus_original.meta:ridge",
4321 vec!["merge:stack.pred_plus_original.meta:ridge.x_original".to_string()],
4322 vec![b0_requirement.key(), b1_requirement.key()],
4323 );
4324
4325 let bundle = build_execution_bundle_with_prediction_contracts(
4326 BundleId::new("bundle:branch.merge.selected.refit").unwrap(),
4327 &plan,
4328 Some(plan.variants[0].variant_id.clone()),
4329 branch_merge_selection_decisions(),
4330 vec![
4331 b0_artifact.clone(),
4332 b1_artifact.clone(),
4333 merge_artifact.clone(),
4334 ],
4335 vec![b0_requirement.clone(), b1_requirement.clone()],
4336 vec![b0_cache.clone(), b1_cache.clone()],
4337 )
4338 .unwrap();
4339 bundle.validate_against_plan(&plan).unwrap();
4340 assert_eq!(bundle.selections.len(), 3);
4341 assert_eq!(bundle.prediction_requirements.len(), 2);
4342 assert_eq!(
4343 bundle.refit_artifacts[2].data_requirement_keys,
4344 vec!["merge:stack.pred_plus_original.meta:ridge.x_original"]
4345 );
4346 assert_eq!(
4347 bundle.refit_artifacts[2].prediction_requirement_keys,
4348 vec![
4349 "branch:b0.model:ridge.oof->merge:stack.pred_plus_original.meta:ridge.b0_oof",
4350 "branch:b1.model:rf.oof->merge:stack.pred_plus_original.meta:ridge.b1_oof",
4351 ]
4352 );
4353
4354 assert!(build_execution_bundle_with_prediction_contracts(
4355 BundleId::new("bundle:branch.merge.missing.branch.refit").unwrap(),
4356 &plan,
4357 Some(plan.variants[0].variant_id.clone()),
4358 branch_merge_selection_decisions(),
4359 vec![b0_artifact.clone(), merge_artifact.clone()],
4360 vec![b0_requirement.clone(), b1_requirement.clone()],
4361 vec![b0_cache.clone(), b1_cache.clone()],
4362 )
4363 .is_err());
4364
4365 let mut misaligned_cache = b0_cache;
4366 misaligned_cache.blocks[0].sample_ids = vec![
4367 SampleId::new("sample:1").unwrap(),
4368 SampleId::new("sample:3").unwrap(),
4369 ];
4370 misaligned_cache.blocks[1].sample_ids = vec![
4371 SampleId::new("sample:2").unwrap(),
4372 SampleId::new("sample:4").unwrap(),
4373 ];
4374 let error = build_execution_bundle_with_prediction_contracts(
4375 BundleId::new("bundle:branch.merge.misaligned.oof.cache").unwrap(),
4376 &plan,
4377 Some(plan.variants[0].variant_id.clone()),
4378 branch_merge_selection_decisions(),
4379 vec![b0_artifact, b1_artifact, merge_artifact],
4380 vec![b0_requirement, b1_requirement],
4381 vec![misaligned_cache, b1_cache],
4382 )
4383 .unwrap_err()
4384 .to_string();
4385 assert!(
4386 error.contains("does not match validation samples"),
4387 "unexpected fold-alignment error: {error}"
4388 );
4389 }
4390
4391 #[test]
4396 fn separation_concat_merge_bundle_assembles_and_is_scored() {
4397 let plan = separation_concat_merge_plan();
4398 let a_requirement = separation_branch_requirement(
4401 "branch:site__A.model:pls",
4402 &["sample:1", "sample:3"],
4403 &["fold:0", "fold:1"],
4404 );
4405 let b_requirement = separation_branch_requirement(
4406 "branch:site__B.model:pls",
4407 &["sample:2", "sample:4"],
4408 &["fold:0", "fold:1"],
4409 );
4410 let a_cache = build_prediction_cache_record(
4411 &a_requirement,
4412 &separation_branch_blocks("branch:site__A.model:pls", "sample:1", "sample:3", 0.0),
4413 )
4414 .unwrap();
4415 let b_cache = build_prediction_cache_record(
4416 &b_requirement,
4417 &separation_branch_blocks("branch:site__B.model:pls", "sample:2", "sample:4", 1.0),
4418 )
4419 .unwrap();
4420 let a_artifact = refit_artifact(
4421 &plan,
4422 "branch:site__A.model:pls",
4423 vec!["branch:site__A.model:pls.x".to_string()],
4424 Vec::new(),
4425 );
4426 let b_artifact = refit_artifact(
4427 &plan,
4428 "branch:site__B.model:pls",
4429 vec!["branch:site__B.model:pls.x".to_string()],
4430 Vec::new(),
4431 );
4432
4433 let mut bundle = build_execution_bundle_with_prediction_contracts(
4434 BundleId::new("bundle:separation.concat.merge").unwrap(),
4435 &plan,
4436 Some(plan.variants[0].variant_id.clone()),
4437 BTreeMap::new(),
4438 vec![a_artifact, b_artifact],
4439 vec![a_requirement, b_requirement],
4440 vec![a_cache, b_cache],
4441 )
4442 .expect("separation-branch concat-merge bundle must assemble");
4443
4444 bundle
4448 .validate_against_plan(&plan)
4449 .expect("partition-covering branch inputs must validate as a group");
4450 assert_eq!(bundle.prediction_requirements.len(), 2);
4451
4452 let scores = ScoreSet {
4456 schema_version: crate::metrics::SCORE_SET_SCHEMA_VERSION,
4457 plan_id: plan.id.clone(),
4458 selection_metric: Some("rmse".to_string()),
4459 reports: vec![crate::metrics::RegressionMetricReport {
4460 prediction_id: Some("prediction:merge:sites:avg".to_string()),
4461 producer_node: NodeId::new("merge:sites").unwrap(),
4462 producer_port: Some("pred".to_string()),
4463 variant_id: Some(plan.variants[0].variant_id.clone()),
4464 variant_label: None,
4465 partition: PredictionPartition::Validation,
4466 fold_id: Some(FoldId::new("avg").unwrap()),
4467 level: PredictionLevel::Sample,
4468 row_count: 4,
4469 target_width: 1,
4470 target_names: vec!["y".to_string()],
4471 metrics: BTreeMap::from([("rmse".to_string(), 1.5)]),
4472 }],
4473 };
4474 bundle.scores = Some(scores);
4475 bundle
4476 .validate_against_plan(&plan)
4477 .expect("bundle with merge-producer scores must validate");
4478 let cv_best = bundle
4479 .scores
4480 .as_ref()
4481 .unwrap()
4482 .reports
4483 .iter()
4484 .find(|report| {
4485 report.producer_node.as_str() == "merge:sites"
4486 && report.fold_id.as_ref().map(FoldId::as_str) == Some("avg")
4487 })
4488 .expect("merge producer must have a cross-fold (avg) score");
4489 assert_eq!(cv_best.metrics.get("rmse"), Some(&1.5));
4490 }
4491
4492 #[test]
4496 fn separation_concat_merge_rejects_overlapping_partitions() {
4497 let plan = separation_concat_merge_plan();
4498 let a_requirement = separation_branch_requirement(
4501 "branch:site__A.model:pls",
4502 &["sample:1", "sample:3"],
4503 &["fold:0", "fold:1"],
4504 );
4505 let b_requirement = separation_branch_requirement(
4506 "branch:site__B.model:pls",
4507 &["sample:2", "sample:3"],
4508 &["fold:0", "fold:1"],
4509 );
4510 let a_cache = build_prediction_cache_record(
4511 &a_requirement,
4512 &separation_branch_blocks("branch:site__A.model:pls", "sample:1", "sample:3", 0.0),
4513 )
4514 .unwrap();
4515 let b_cache = build_prediction_cache_record(
4516 &b_requirement,
4517 &separation_branch_blocks("branch:site__B.model:pls", "sample:2", "sample:3", 1.0),
4518 )
4519 .unwrap();
4520 let a_artifact = refit_artifact(
4521 &plan,
4522 "branch:site__A.model:pls",
4523 vec!["branch:site__A.model:pls.x".to_string()],
4524 Vec::new(),
4525 );
4526 let b_artifact = refit_artifact(
4527 &plan,
4528 "branch:site__B.model:pls",
4529 vec!["branch:site__B.model:pls.x".to_string()],
4530 Vec::new(),
4531 );
4532
4533 let error = build_execution_bundle_with_prediction_contracts(
4534 BundleId::new("bundle:separation.concat.merge.overlap").unwrap(),
4535 &plan,
4536 Some(plan.variants[0].variant_id.clone()),
4537 BTreeMap::new(),
4538 vec![a_artifact, b_artifact],
4539 vec![a_requirement, b_requirement],
4540 vec![a_cache, b_cache],
4541 )
4542 .unwrap_err()
4543 .to_string();
4544 assert!(
4545 error.contains("overlapping branch predictions"),
4546 "overlap must be rejected, got: {error}"
4547 );
4548 }
4549
4550 #[test]
4554 fn separation_concat_merge_rejects_incomplete_coverage() {
4555 let plan = separation_concat_merge_plan();
4556 let a_requirement =
4559 separation_branch_requirement("branch:site__A.model:pls", &["sample:1"], &["fold:0"]);
4560 let b_requirement = separation_branch_requirement(
4561 "branch:site__B.model:pls",
4562 &["sample:2", "sample:4"],
4563 &["fold:0", "fold:1"],
4564 );
4565 let a_cache = build_prediction_cache_record(
4566 &a_requirement,
4567 &[PredictionBlock {
4568 prediction_id: Some("prediction:a:fold0".to_string()),
4569 producer_node: NodeId::new("branch:site__A.model:pls").unwrap(),
4570 producer_port: None,
4571 partition: PredictionPartition::Validation,
4572 fold_id: Some(FoldId::new("fold:0").unwrap()),
4573 sample_ids: vec![SampleId::new("sample:1").unwrap()],
4574 values: vec![vec![0.1]],
4575 target_names: vec!["y".to_string()],
4576 }],
4577 )
4578 .unwrap();
4579 let b_cache = build_prediction_cache_record(
4580 &b_requirement,
4581 &separation_branch_blocks("branch:site__B.model:pls", "sample:2", "sample:4", 1.0),
4582 )
4583 .unwrap();
4584 let a_artifact = refit_artifact(
4585 &plan,
4586 "branch:site__A.model:pls",
4587 vec!["branch:site__A.model:pls.x".to_string()],
4588 Vec::new(),
4589 );
4590 let b_artifact = refit_artifact(
4591 &plan,
4592 "branch:site__B.model:pls",
4593 vec!["branch:site__B.model:pls.x".to_string()],
4594 Vec::new(),
4595 );
4596
4597 let error = build_execution_bundle_with_prediction_contracts(
4598 BundleId::new("bundle:separation.concat.merge.gap").unwrap(),
4599 &plan,
4600 Some(plan.variants[0].variant_id.clone()),
4601 BTreeMap::new(),
4602 vec![a_artifact, b_artifact],
4603 vec![a_requirement, b_requirement],
4604 vec![a_cache, b_cache],
4605 )
4606 .unwrap_err()
4607 .to_string();
4608 assert!(
4609 error.contains("do not cover"),
4610 "an OOF gap must be rejected, got: {error}"
4611 );
4612 }
4613
4614 #[test]
4619 fn separation_concat_merge_rejects_missing_branch_edge() {
4620 let plan = separation_concat_merge_plan();
4621 let a_requirement = separation_branch_requirement(
4625 "branch:site__A.model:pls",
4626 &["sample:1", "sample:2", "sample:3", "sample:4"],
4627 &["fold:0", "fold:1"],
4628 );
4629 let a_artifact = refit_artifact(
4630 &plan,
4631 "branch:site__A.model:pls",
4632 vec!["branch:site__A.model:pls.x".to_string()],
4633 Vec::new(),
4634 );
4635 let b_artifact = refit_artifact(
4636 &plan,
4637 "branch:site__B.model:pls",
4638 vec!["branch:site__B.model:pls.x".to_string()],
4639 Vec::new(),
4640 );
4641
4642 let error = build_execution_bundle_with_prediction_contracts(
4643 BundleId::new("bundle:separation.concat.merge.missing.branch").unwrap(),
4644 &plan,
4645 Some(plan.variants[0].variant_id.clone()),
4646 BTreeMap::new(),
4647 vec![a_artifact, b_artifact],
4648 vec![a_requirement],
4649 Vec::new(),
4650 )
4651 .unwrap_err()
4652 .to_string();
4653 assert!(
4654 error.contains("do not match the plan's incoming OOF edges"),
4655 "a missing branch edge must be rejected, got: {error}"
4656 );
4657 }
4658
4659 #[test]
4664 fn separation_concat_merge_rejects_partial_cache_coverage() {
4665 let plan = separation_concat_merge_plan();
4666 let a_requirement = separation_branch_requirement(
4667 "branch:site__A.model:pls",
4668 &["sample:1", "sample:3"],
4669 &["fold:0", "fold:1"],
4670 );
4671 let b_requirement = separation_branch_requirement(
4672 "branch:site__B.model:pls",
4673 &["sample:2", "sample:4"],
4674 &["fold:0", "fold:1"],
4675 );
4676 let a_cache = build_prediction_cache_record(
4678 &a_requirement,
4679 &separation_branch_blocks("branch:site__A.model:pls", "sample:1", "sample:3", 0.0),
4680 )
4681 .unwrap();
4682 let a_artifact = refit_artifact(
4683 &plan,
4684 "branch:site__A.model:pls",
4685 vec!["branch:site__A.model:pls.x".to_string()],
4686 Vec::new(),
4687 );
4688 let b_artifact = refit_artifact(
4689 &plan,
4690 "branch:site__B.model:pls",
4691 vec!["branch:site__B.model:pls.x".to_string()],
4692 Vec::new(),
4693 );
4694
4695 let error = build_execution_bundle_with_prediction_contracts(
4696 BundleId::new("bundle:separation.concat.merge.partial.cache").unwrap(),
4697 &plan,
4698 Some(plan.variants[0].variant_id.clone()),
4699 BTreeMap::new(),
4700 vec![a_artifact, b_artifact],
4701 vec![a_requirement, b_requirement],
4702 vec![a_cache],
4703 )
4704 .unwrap_err()
4705 .to_string();
4706 assert!(
4707 error.contains("partial prediction-cache coverage"),
4708 "a partial-cache concat group must be rejected, got: {error}"
4709 );
4710 }
4711
4712 #[test]
4713 fn prediction_requirements_are_typed_and_validate_against_oof_edges() {
4714 let plan = branch_merge_plan();
4715 let meta_plan = plan
4716 .node_plans
4717 .get(&NodeId::new("merge:stack.pred_plus_original.meta:ridge").unwrap())
4718 .unwrap();
4719 let producer_node = NodeId::new("branch:b0.model:ridge").unwrap();
4720 let fold0 = FoldId::new("fold:0").unwrap();
4721 let fold1 = FoldId::new("fold:1").unwrap();
4722 let samples = [
4723 SampleId::new("sample:1").unwrap(),
4724 SampleId::new("sample:2").unwrap(),
4725 SampleId::new("sample:3").unwrap(),
4726 SampleId::new("sample:4").unwrap(),
4727 ];
4728 let requirement = BundlePredictionRequirement {
4729 producer_node: producer_node.clone(),
4730 source_port: "oof".to_string(),
4731 consumer_node: meta_plan.node_id.clone(),
4732 target_port: "b0_oof".to_string(),
4733 partition: PredictionPartition::Validation,
4734 prediction_level: PredictionLevel::Sample,
4735 fold_ids: vec![fold0.clone(), fold1.clone()],
4736 unit_ids: Vec::new(),
4737 sample_ids: samples.to_vec(),
4738 prediction_width: 1,
4739 target_names: vec!["y".to_string()],
4740 };
4741 let prediction_blocks = vec![
4742 PredictionBlock {
4743 prediction_id: Some("prediction:branch:b0.fold0".to_string()),
4744 producer_node: producer_node.clone(),
4745 producer_port: Some("oof".to_string()),
4746 partition: PredictionPartition::Validation,
4747 fold_id: Some(fold0),
4748 sample_ids: samples[0..2].to_vec(),
4749 values: vec![vec![0.1], vec![0.2]],
4750 target_names: vec!["y".to_string()],
4751 },
4752 PredictionBlock {
4753 prediction_id: Some("prediction:branch:b0.fold1".to_string()),
4754 producer_node: producer_node.clone(),
4755 producer_port: Some("oof".to_string()),
4756 partition: PredictionPartition::Validation,
4757 fold_id: Some(fold1),
4758 sample_ids: samples[2..4].to_vec(),
4759 values: vec![vec![0.3], vec![0.4]],
4760 target_names: vec!["y".to_string()],
4761 },
4762 PredictionBlock {
4767 prediction_id: Some("prediction:branch:b0.inner".to_string()),
4768 producer_node: producer_node.clone(),
4769 producer_port: Some("oof".to_string()),
4770 partition: PredictionPartition::Validation,
4771 fold_id: Some(FoldId::new("nested:fold:0").unwrap()),
4772 sample_ids: samples[0..2].to_vec(),
4773 values: vec![vec![9.1], vec![9.2]],
4774 target_names: vec!["y".to_string()],
4775 },
4776 ];
4777 let cache = build_prediction_cache_record(&requirement, &prediction_blocks).unwrap();
4778 let payload = build_prediction_cache_payload(&requirement, &prediction_blocks).unwrap();
4779 assert_eq!(cache.prediction_level, PredictionLevel::Sample);
4780 assert_eq!(payload.prediction_level, PredictionLevel::Sample);
4781 assert!(cache
4782 .blocks
4783 .iter()
4784 .all(|block| block.prediction_level == PredictionLevel::Sample));
4785 assert_eq!(cache.block_count, 2);
4786 assert_eq!(cache.row_count, 4);
4787 assert!(cache.blocks.iter().all(|block| {
4788 block
4789 .fold_id
4790 .as_ref()
4791 .is_some_and(|fold_id| requirement.fold_ids.contains(fold_id))
4792 }));
4793 validate_prediction_cache_payload_matches_record(&payload, &cache).unwrap();
4794 let cache_namespace_fingerprints = vec!["a".repeat(64), "b".repeat(64)];
4795 let mut d10_cache = cache.clone();
4796 d10_cache.cache_namespace_fingerprints = cache_namespace_fingerprints.clone();
4797 d10_cache.validate().unwrap();
4798 let mut d10_payload = payload.clone();
4799 d10_payload.cache_namespace_fingerprints = cache_namespace_fingerprints;
4800 d10_payload.validate().unwrap();
4801 validate_prediction_cache_payload_matches_record(&d10_payload, &d10_cache).unwrap();
4802 for forbidden_partition in [
4803 PredictionPartition::Train,
4804 PredictionPartition::Test,
4805 PredictionPartition::Final,
4806 ] {
4807 let mut non_oof_requirement = requirement.clone();
4808 non_oof_requirement.partition = forbidden_partition.clone();
4809 let requirement_error = non_oof_requirement.validate().unwrap_err().to_string();
4810 assert!(
4811 requirement_error.contains("must use validation OOF predictions"),
4812 "D10 cache namespace must not broaden bundle prediction requirements to {forbidden_partition:?}: {requirement_error}"
4813 );
4814
4815 let mut non_oof_cache = d10_cache.clone();
4816 non_oof_cache.partition = forbidden_partition.clone();
4817 let cache_error = non_oof_cache.validate().unwrap_err().to_string();
4818 assert!(
4819 cache_error.contains("must cache validation OOF predictions"),
4820 "D10 cache namespace must not allow non-OOF cache records for {forbidden_partition:?}: {cache_error}"
4821 );
4822
4823 let mut non_oof_payload = d10_payload.clone();
4824 non_oof_payload.partition = forbidden_partition;
4825 let payload_error = non_oof_payload.validate().unwrap_err().to_string();
4826 assert!(
4827 payload_error.contains("must cache validation OOF predictions"),
4828 "D10 cache namespace must not allow non-OOF cache payloads: {payload_error}"
4829 );
4830 }
4831 let mut short_namespace_cache = d10_cache.clone();
4832 short_namespace_cache.cache_namespace_fingerprints.pop();
4833 assert!(short_namespace_cache
4834 .validate()
4835 .unwrap_err()
4836 .to_string()
4837 .contains("namespace fingerprint count"));
4838 let mut short_namespace_payload = d10_payload;
4839 short_namespace_payload.cache_namespace_fingerprints.pop();
4840 assert!(short_namespace_payload
4841 .validate()
4842 .unwrap_err()
4843 .to_string()
4844 .contains("namespace fingerprint count"));
4845 let mut wrong_level_requirement = requirement.clone();
4846 wrong_level_requirement.prediction_level = PredictionLevel::Target;
4847 assert!(wrong_level_requirement.validate().is_err());
4848 let mut wrong_level_cache = cache.clone();
4849 wrong_level_cache.prediction_level = PredictionLevel::Target;
4850 assert!(wrong_level_cache.validate().is_err());
4851 let mut wrong_level_payload = payload.clone();
4852 wrong_level_payload.prediction_level = PredictionLevel::Target;
4853 assert!(wrong_level_payload.validate().is_err());
4854 let prediction_key = requirement.key();
4855 let artifact = RefitArtifactRecord {
4856 node_id: meta_plan.node_id.clone(),
4857 controller_id: meta_plan.controller_id.clone(),
4858 artifact: ArtifactRef {
4859 id: ArtifactId::new("artifact:merge:stack.pred_plus_original.meta:ridge:refit")
4860 .unwrap(),
4861 kind: "mock_model".to_string(),
4862 controller_id: meta_plan.controller_id.clone(),
4863 backend: None,
4864 uri: None,
4865 content_fingerprint: None,
4866 size_bytes: Some(128),
4867 plugin: None,
4868 plugin_version: None,
4869 abi_major: None,
4870 abi_min_minor: None,
4871 native_predictor_descriptor: None,
4872 },
4873 params_fingerprint: meta_plan.params_fingerprint.clone(),
4874 training_loss_fingerprint: meta_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
4875 data_requirement_keys: vec![
4876 "merge:stack.pred_plus_original.meta:ridge.x_original".to_string()
4877 ],
4878 prediction_requirement_keys: vec![prediction_key],
4879 };
4880
4881 assert!(build_execution_bundle_with_prediction_contracts(
4882 BundleId::new("bundle:d10.cache.without.selected.variant").unwrap(),
4883 &plan,
4884 None,
4885 BTreeMap::new(),
4886 vec![artifact.clone()],
4887 vec![requirement.clone()],
4888 vec![d10_cache],
4889 )
4890 .unwrap_err()
4891 .to_string()
4892 .contains("requires selected_variant_id"));
4893
4894 assert!(build_execution_bundle(
4895 BundleId::new("bundle:missing.prediction.requirement").unwrap(),
4896 &plan,
4897 Some(plan.variants[0].variant_id.clone()),
4898 BTreeMap::new(),
4899 vec![artifact.clone()],
4900 )
4901 .is_err());
4902
4903 assert!(build_execution_bundle_with_prediction_requirements(
4904 BundleId::new("bundle:typed.prediction.requirement.without.cache").unwrap(),
4905 &plan,
4906 Some(plan.variants[0].variant_id.clone()),
4907 BTreeMap::new(),
4908 vec![artifact.clone()],
4909 vec![requirement.clone()],
4910 )
4911 .is_err());
4912
4913 let bundle = build_execution_bundle_with_prediction_contracts(
4914 BundleId::new("bundle:typed.prediction.requirement").unwrap(),
4915 &plan,
4916 Some(plan.variants[0].variant_id.clone()),
4917 BTreeMap::new(),
4918 vec![artifact],
4919 vec![requirement],
4920 vec![cache],
4921 )
4922 .unwrap();
4923 bundle.validate_against_plan(&plan).unwrap();
4924 assert_eq!(bundle.prediction_requirements.len(), 1);
4925 assert_eq!(bundle.prediction_caches.len(), 1);
4926 assert_eq!(
4927 bundle.refit_artifacts[0].prediction_requirement_keys,
4928 vec!["branch:b0.model:ridge.oof->merge:stack.pred_plus_original.meta:ridge.b0_oof"]
4929 );
4930 let payload_set = BundlePredictionCachePayloadSet {
4931 bundle_id: bundle.bundle_id.clone(),
4932 schema_version: PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
4933 caches: vec![payload],
4934 };
4935 payload_set.validate_against_bundle(&bundle).unwrap();
4936 let refit_replay_request = ReplayPhaseRequest {
4937 bundle_id: bundle.bundle_id.clone(),
4938 phase: Phase::Refit,
4939 data_envelope_keys: bundle
4940 .data_requirements
4941 .iter()
4942 .map(BundleDataRequirement::key)
4943 .collect(),
4944 };
4945 refit_replay_request
4946 .validate_for_bundle_with_prediction_cache_payloads(&bundle, Some(&payload_set))
4947 .unwrap();
4948 let mut tampered_payload_set = payload_set.clone();
4949 tampered_payload_set.caches[0].blocks[0].values[0][0] = 99.0;
4950 assert!(tampered_payload_set
4951 .validate_against_bundle(&bundle)
4952 .is_err());
4953 let mut missing_payload_set = payload_set.clone();
4954 missing_payload_set.caches.clear();
4955 assert!(missing_payload_set
4956 .validate_against_bundle(&bundle)
4957 .is_err());
4958 assert!(refit_replay_request.validate_for_bundle(&bundle).is_err());
4959
4960 let mut wrong_data_owner = bundle.clone();
4961 wrong_data_owner.refit_artifacts[0].data_requirement_keys =
4962 vec!["branch:b0.model:ridge.x".to_string()];
4963 assert!(wrong_data_owner.validate().is_err());
4964
4965 let mut wrong_prediction_consumer = bundle;
4966 wrong_prediction_consumer.refit_artifacts[0].node_id =
4967 NodeId::new("branch:b0.model:ridge").unwrap();
4968 wrong_prediction_consumer.refit_artifacts[0]
4969 .data_requirement_keys
4970 .clear();
4971 assert!(wrong_prediction_consumer.validate().is_err());
4972 }
4973
4974 #[test]
4975 fn aggregated_prediction_cache_contracts_preserve_unit_ids() {
4976 let plan = branch_merge_plan();
4977 let producer_node = NodeId::new("branch:b0.model:ridge").unwrap();
4978 let consumer_node = NodeId::new("merge:stack.pred_plus_original.meta:ridge").unwrap();
4979 let fold0 = FoldId::new("fold:0").unwrap();
4980 let fold1 = FoldId::new("fold:1").unwrap();
4981 let target_a = PredictionUnitId::Target(TargetId::new("target:a").unwrap());
4982 let target_b = PredictionUnitId::Target(TargetId::new("target:b").unwrap());
4983 let requirement = BundlePredictionRequirement {
4984 producer_node: producer_node.clone(),
4985 source_port: "oof".to_string(),
4986 consumer_node: consumer_node.clone(),
4987 target_port: "b0_oof".to_string(),
4988 partition: PredictionPartition::Validation,
4989 prediction_level: PredictionLevel::Target,
4990 fold_ids: vec![fold0.clone(), fold1.clone()],
4991 unit_ids: vec![target_a.clone(), target_b.clone()],
4992 sample_ids: Vec::new(),
4993 prediction_width: 1,
4994 target_names: vec!["y".to_string()],
4995 };
4996 let aggregated_blocks = vec![
4997 AggregatedPredictionBlock {
4998 prediction_id: Some("prediction:branch:b0.target.fold0".to_string()),
4999 producer_node: producer_node.clone(),
5000 producer_port: Some("pred".to_string()),
5001 partition: PredictionPartition::Validation,
5002 fold_id: Some(fold0),
5003 level: PredictionLevel::Target,
5004 unit_ids: vec![target_a],
5005 values: vec![vec![0.15]],
5006 target_names: vec!["y".to_string()],
5007 },
5008 AggregatedPredictionBlock {
5009 prediction_id: Some("prediction:branch:b0.target.fold1".to_string()),
5010 producer_node,
5011 producer_port: Some("pred".to_string()),
5012 partition: PredictionPartition::Validation,
5013 fold_id: Some(fold1),
5014 level: PredictionLevel::Target,
5015 unit_ids: vec![target_b],
5016 values: vec![vec![0.35]],
5017 target_names: vec!["y".to_string()],
5018 },
5019 ];
5020
5021 let cache =
5022 build_aggregated_prediction_cache_record(&requirement, &aggregated_blocks).unwrap();
5023 let payload =
5024 build_aggregated_prediction_cache_payload(&requirement, &aggregated_blocks).unwrap();
5025 assert_eq!(cache.prediction_level, PredictionLevel::Target);
5026 assert_eq!(cache.unit_ids, requirement.unit_ids);
5027 assert!(cache.sample_ids.is_empty());
5028 assert!(payload.blocks.is_empty());
5029 assert_eq!(payload.aggregated_blocks.len(), 2);
5030 validate_prediction_cache_payload_matches_record(&payload, &cache).unwrap();
5031
5032 let artifact = refit_artifact(
5033 &plan,
5034 "merge:stack.pred_plus_original.meta:ridge",
5035 vec!["merge:stack.pred_plus_original.meta:ridge.x_original".to_string()],
5036 vec![requirement.key()],
5037 );
5038 let bundle = build_execution_bundle_with_prediction_contracts(
5039 BundleId::new("bundle:target.prediction.requirement").unwrap(),
5040 &plan,
5041 Some(plan.variants[0].variant_id.clone()),
5042 BTreeMap::new(),
5043 vec![artifact],
5044 vec![requirement],
5045 vec![cache],
5046 )
5047 .unwrap();
5048 bundle.validate_against_plan(&plan).unwrap();
5049
5050 let mut tampered_payload = payload;
5051 tampered_payload.aggregated_blocks[0].unit_ids =
5052 vec![PredictionUnitId::Target(TargetId::new("target:z").unwrap())];
5053 assert!(validate_prediction_cache_payload_matches_record(
5054 &tampered_payload,
5055 &bundle.prediction_caches[0]
5056 )
5057 .is_err());
5058 }
5059
5060 #[test]
5061 fn replay_envelopes_must_match_bundle_requirements() {
5062 let plan = plan();
5063 let bundle = build_execution_bundle(
5064 BundleId::new("bundle:demo").unwrap(),
5065 &plan,
5066 None,
5067 BTreeMap::new(),
5068 Vec::new(),
5069 )
5070 .unwrap();
5071 let envelope: ExternalDataPlanEnvelope = serde_json::from_str(include_str!(
5072 "../tests/fixtures/package/data/coordinator_data_plan_envelope_sample12.json"
5073 ))
5074 .unwrap();
5075
5076 bundle
5077 .validate_replay_envelopes(&BTreeMap::from([(
5078 "model:base.x".to_string(),
5079 envelope.clone(),
5080 )]))
5081 .unwrap();
5082
5083 let mut mismatched = envelope;
5084 mismatched.schema_fingerprint = "0".repeat(64);
5085 assert!(bundle
5086 .validate_replay_envelopes(&BTreeMap::from([("model:base.x".to_string(), mismatched,)]))
5087 .is_err());
5088 }
5089
5090 #[test]
5091 fn rejects_unsupported_bundle_schema_version() {
5092 let mut bundle = build_execution_bundle(
5093 BundleId::new("bundle:demo").unwrap(),
5094 &plan(),
5095 None,
5096 BTreeMap::new(),
5097 Vec::new(),
5098 )
5099 .unwrap();
5100 bundle.schema_version = EXECUTION_BUNDLE_SCHEMA_VERSION + 1;
5101
5102 assert!(bundle.validate().is_err());
5103
5104 bundle.schema_version = 0;
5105 assert!(bundle.validate().is_err());
5106 }
5107
5108 #[test]
5109 fn rejects_bundle_with_scores_plan_id_mismatch() {
5110 let plan = plan();
5111 let mut bundle = build_execution_bundle(
5112 BundleId::new("bundle:demo").unwrap(),
5113 &plan,
5114 None,
5115 BTreeMap::new(),
5116 Vec::new(),
5117 )
5118 .unwrap();
5119 bundle.scores = Some(ScoreSet {
5120 schema_version: crate::metrics::SCORE_SET_SCHEMA_VERSION,
5121 plan_id: bundle.plan_id.clone(),
5122 selection_metric: Some("rmse".to_string()),
5123 reports: vec![crate::metrics::RegressionMetricReport {
5124 prediction_id: None,
5125 producer_node: NodeId::new("model:compat.0").unwrap(),
5126 producer_port: Some("pred".to_string()),
5127 variant_id: None,
5128 variant_label: None,
5129 partition: PredictionPartition::Test,
5130 fold_id: Some(FoldId::new("final").unwrap()),
5131 level: PredictionLevel::Sample,
5132 row_count: 4,
5133 target_width: 1,
5134 target_names: vec!["y".to_string()],
5135 metrics: BTreeMap::from([("rmse".to_string(), 1.0)]),
5136 }],
5137 });
5138 bundle.validate().unwrap();
5140 bundle.scores.as_mut().unwrap().plan_id = "plan:wrong".to_string();
5142 let err = bundle.validate().unwrap_err().to_string();
5143 assert!(
5144 err.contains("does not match its embedded scores plan_id"),
5145 "{err}"
5146 );
5147 }
5148
5149 #[test]
5150 fn schema_migration_policy_is_explicit_and_refuses_implicit_migrations() {
5151 let bundle_policy = execution_bundle_schema_migration_policy();
5152 assert_eq!(
5153 bundle_policy.current_version,
5154 EXECUTION_BUNDLE_SCHEMA_VERSION
5155 );
5156 assert_eq!(
5157 bundle_policy.min_readable_version,
5158 MIN_READABLE_EXECUTION_BUNDLE_SCHEMA_VERSION
5159 );
5160 assert_eq!(
5161 bundle_policy.min_writable_version,
5162 MIN_WRITABLE_EXECUTION_BUNDLE_SCHEMA_VERSION
5163 );
5164 assert_eq!(
5165 bundle_policy
5166 .automatic_migrations
5167 .get(&LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION),
5168 Some(&EXECUTION_BUNDLE_SCHEMA_VERSION)
5169 );
5170 bundle_policy
5171 .validate_read_version(LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION, "bundle `legacy`")
5172 .unwrap();
5173 bundle_policy
5174 .validate_read_version(EXECUTION_BUNDLE_SCHEMA_VERSION, "bundle `current`")
5175 .unwrap();
5176 assert!(bundle_policy
5177 .validate_read_version(EXECUTION_BUNDLE_SCHEMA_VERSION + 1, "bundle `future`")
5178 .is_err());
5179 assert!(bundle_policy
5180 .validate_read_version(0, "bundle `zero`")
5181 .is_err());
5182
5183 let mut future_policy = SchemaMigrationPolicy {
5184 artifact: "execution_bundle".to_string(),
5185 current_version: 2,
5186 min_readable_version: 1,
5187 min_writable_version: 2,
5188 automatic_migrations: BTreeMap::new(),
5189 };
5190 assert!(future_policy
5191 .validate_read_version(1, "bundle `old-without-migration`")
5192 .is_err());
5193 future_policy.automatic_migrations.insert(1, 2);
5194 future_policy
5195 .validate_read_version(1, "bundle `old-with-migration`")
5196 .unwrap();
5197 }
5198
5199 #[test]
5200 fn prediction_cache_payload_schema_policy_rejects_unsupported_versions() {
5201 let policy = prediction_cache_payload_schema_migration_policy();
5202 assert_eq!(
5203 policy.current_version,
5204 PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION
5205 );
5206 assert_eq!(
5207 policy
5208 .automatic_migrations
5209 .get(&LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION),
5210 Some(&PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION)
5211 );
5212 policy
5213 .validate_read_version(
5214 LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
5215 "payload `legacy`",
5216 )
5217 .unwrap();
5218
5219 let mut payload_set = BundlePredictionCachePayloadSet {
5220 bundle_id: BundleId::new("bundle:payload.schema").unwrap(),
5221 schema_version: PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
5222 caches: Vec::new(),
5223 };
5224 payload_set.validate().unwrap();
5225
5226 payload_set.schema_version = PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION + 1;
5227 assert!(payload_set.validate().is_err());
5228
5229 payload_set.schema_version = 0;
5230 assert!(payload_set.validate().is_err());
5231 }
5232
5233 #[test]
5234 fn prediction_cache_payload_format_enforces_port_family() {
5235 fn payload_for_block(
5236 format: &str,
5237 producer_port: Option<String>,
5238 ) -> BundlePredictionCachePayload {
5239 let block = PredictionBlock {
5240 prediction_id: Some("prediction:model:base.fold0".to_string()),
5241 producer_node: NodeId::new("model:base").unwrap(),
5242 producer_port,
5243 partition: PredictionPartition::Validation,
5244 fold_id: Some(FoldId::new("fold:0").unwrap()),
5245 sample_ids: vec![SampleId::new("sample:1").unwrap()],
5246 values: vec![vec![1.0]],
5247 target_names: vec!["y".to_string()],
5248 };
5249 let blocks = vec![block];
5250 BundlePredictionCachePayload {
5251 requirement_key: "model:base.oof->model:meta.pred".to_string(),
5252 cache_id: "prediction-cache:model:base.oof->model:meta.pred".to_string(),
5253 cache_namespace_fingerprints: Vec::new(),
5254 format: format.to_string(),
5255 partition: PredictionPartition::Validation,
5256 prediction_level: PredictionLevel::Sample,
5257 block_count: blocks.len(),
5258 row_count: blocks.iter().map(|block| block.sample_ids.len()).sum(),
5259 content_fingerprint: stable_json_fingerprint(&blocks).unwrap(),
5260 blocks,
5261 aggregated_blocks: Vec::new(),
5262 }
5263 }
5264
5265 payload_for_block(LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT, None)
5266 .validate()
5267 .unwrap();
5268 payload_for_block(BUNDLE_PREDICTION_CACHE_FORMAT, Some("oof".to_string()))
5269 .validate()
5270 .unwrap();
5271
5272 let v1_with_port = payload_for_block(
5273 LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT,
5274 Some("oof".to_string()),
5275 )
5276 .validate()
5277 .unwrap_err()
5278 .to_string();
5279 assert!(
5280 v1_with_port.contains("V1") && v1_with_port.contains("producer_port"),
5281 "unexpected V1 port-family error: {v1_with_port}"
5282 );
5283
5284 let v2_without_port = payload_for_block(BUNDLE_PREDICTION_CACHE_FORMAT, None)
5285 .validate()
5286 .unwrap_err()
5287 .to_string();
5288 assert!(
5289 v2_without_port.contains("V2") && v2_without_port.contains("requires producer_port"),
5290 "unexpected V2 port-family error: {v2_without_port}"
5291 );
5292 }
5293
5294 #[test]
5295 fn replay_request_requires_predict_explain_or_refit_phase() {
5296 let bundle = build_execution_bundle(
5297 BundleId::new("bundle:demo").unwrap(),
5298 &plan(),
5299 None,
5300 BTreeMap::new(),
5301 Vec::new(),
5302 )
5303 .unwrap();
5304
5305 ReplayPhaseRequest {
5306 bundle_id: bundle.bundle_id.clone(),
5307 phase: Phase::Predict,
5308 data_envelope_keys: vec!["model:base.x".to_string()],
5309 }
5310 .validate_for_bundle(&bundle)
5311 .unwrap();
5312 ReplayPhaseRequest {
5313 bundle_id: bundle.bundle_id.clone(),
5314 phase: Phase::Refit,
5315 data_envelope_keys: vec!["model:base.x".to_string()],
5316 }
5317 .validate_for_bundle(&bundle)
5318 .unwrap();
5319 assert!(ReplayPhaseRequest {
5320 bundle_id: bundle.bundle_id.clone(),
5321 phase: Phase::FitCv,
5322 data_envelope_keys: vec!["model:base.x".to_string()],
5323 }
5324 .validate_for_bundle(&bundle)
5325 .is_err());
5326 assert!(ReplayPhaseRequest {
5327 bundle_id: bundle.bundle_id.clone(),
5328 phase: Phase::Predict,
5329 data_envelope_keys: vec!["model:base.x".to_string(), "model:base.x".to_string()],
5330 }
5331 .validate_for_bundle(&bundle)
5332 .is_err());
5333 assert!(ReplayPhaseRequest {
5334 bundle_id: bundle.bundle_id.clone(),
5335 phase: Phase::Predict,
5336 data_envelope_keys: vec!["model:base.y".to_string()],
5337 }
5338 .validate_for_bundle(&bundle)
5339 .is_err());
5340 }
5341
5342 #[test]
5343 fn prediction_level_wire_absent_parses_as_sample_but_serialization_stays_explicit() {
5344 let mut requirement = BundlePredictionRequirement {
5345 producer_node: NodeId::new("model:base").unwrap(),
5346 source_port: "oof".to_string(),
5347 consumer_node: NodeId::new("model:meta").unwrap(),
5348 target_port: "x".to_string(),
5349 partition: PredictionPartition::Validation,
5350 prediction_level: PredictionLevel::Sample,
5351 fold_ids: vec![FoldId::new("fold:0").unwrap()],
5352 unit_ids: Vec::new(),
5353 sample_ids: vec![SampleId::new("sample:1").unwrap()],
5354 prediction_width: 1,
5355 target_names: vec!["y".to_string()],
5356 };
5357
5358 let mut sample = serde_json::to_value(&requirement).unwrap();
5359 assert_eq!(sample["prediction_level"], serde_json::json!("sample"));
5360
5361 sample.as_object_mut().unwrap().remove("prediction_level");
5362 let parsed: BundlePredictionRequirement = serde_json::from_value(sample).unwrap();
5363 assert_eq!(parsed.prediction_level, PredictionLevel::Sample);
5364
5365 requirement.prediction_level = PredictionLevel::Group;
5366 let group = serde_json::to_value(&requirement).unwrap();
5367 assert_eq!(group["prediction_level"], serde_json::json!("group"));
5368 }
5369
5370 #[test]
5371 fn methods_hpo_resume_state_round_trips_and_rejects_orphan_or_score_drift() {
5372 let state = methods_hpo_resume_state();
5373 state.validate().unwrap();
5374 let json = serde_json::to_string(&state).unwrap();
5375 let parsed: MethodsHpoResumeState = serde_json::from_str(&json).unwrap();
5376 assert_eq!(parsed, state);
5377
5378 let mut orphan = state.clone();
5379 orphan.completed_reports.clear();
5380 assert!(orphan
5381 .validate()
5382 .unwrap_err()
5383 .to_string()
5384 .contains("exactly cover"));
5385
5386 let mut score_drift = state;
5387 score_drift.completed_reports[0].score += 1.1e-12;
5388 assert!(score_drift
5389 .validate()
5390 .unwrap_err()
5391 .to_string()
5392 .contains("score/intermediate evidence is inconsistent"));
5393
5394 let mut failed = serde_json::to_value(methods_hpo_resume_state()).unwrap();
5395 failed["completed_reports"][0]["terminal_state"] = serde_json::json!("failed");
5396 assert!(serde_json::from_value::<MethodsHpoResumeState>(failed).is_err());
5397
5398 let mut uppercase_binding = methods_hpo_resume_state();
5399 uppercase_binding.checkpoint.binding.optimizer_fingerprint = "A".repeat(64);
5400 assert!(uppercase_binding
5401 .validate()
5402 .unwrap_err()
5403 .to_string()
5404 .contains("lowercase SHA-256"));
5405 }
5406
5407 #[test]
5408 fn methods_hpo_resume_state_rejects_unknown_and_duplicate_json_members() {
5409 let plan = plan();
5410 let bundle = build_execution_bundle(
5411 BundleId::new("bundle:hpo.strict").unwrap(),
5412 &plan,
5413 None,
5414 BTreeMap::new(),
5415 Vec::new(),
5416 )
5417 .unwrap();
5418 let mut raw = serde_json::to_value(bundle).unwrap();
5419 raw.as_object_mut()
5420 .unwrap()
5421 .insert("unknown_hpo_member".to_string(), serde_json::json!(true));
5422 assert!(ExecutionBundle::from_json(&serde_json::to_string(&raw).unwrap()).is_err());
5423
5424 let encoded = serde_json::to_string(&raw).unwrap();
5425 let duplicated = encoded.replacen(
5426 "\"bundle_id\":\"bundle:hpo.strict\"",
5427 "\"bundle_id\":\"bundle:hpo.strict\",\"bundle_id\":\"bundle:hpo.strict\"",
5428 1,
5429 );
5430 let error = ExecutionBundle::from_json(&duplicated)
5431 .unwrap_err()
5432 .to_string();
5433 assert!(error.contains("duplicate JSON object key"), "{error}");
5434
5435 let mut state = serde_json::to_value(methods_hpo_resume_state()).unwrap();
5436 state.as_object_mut().unwrap().insert(
5437 "unexpected_candidate_side_channel".to_string(),
5438 serde_json::json!(true),
5439 );
5440 assert!(serde_json::from_value::<MethodsHpoResumeState>(state).is_err());
5441 }
5442
5443 #[test]
5444 fn execution_bundle_v1_json_rejects_null_conformal_member() {
5445 let plan = plan();
5446 let mut bundle = build_execution_bundle(
5447 BundleId::new("bundle:v1.strict-conformal").unwrap(),
5448 &plan,
5449 None,
5450 BTreeMap::new(),
5451 Vec::new(),
5452 )
5453 .unwrap();
5454 bundle.schema_version = LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION;
5455 let mut raw = serde_json::to_value(bundle).unwrap();
5456 raw.as_object_mut()
5457 .unwrap()
5458 .insert("conformal_calibration".to_string(), serde_json::Value::Null);
5459 let error = ExecutionBundle::from_json(&serde_json::to_string(&raw).unwrap())
5460 .unwrap_err()
5461 .to_string();
5462 assert!(error.contains("including null"), "{error}");
5463 }
5464}