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