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 })
3033 .cloned()
3034 .collect::<Vec<_>>();
3035 if selected.is_empty() {
3036 return Err(DagMlError::RuntimeValidation(format!(
3037 "prediction cache requirement `{}` has no matching prediction blocks",
3038 requirement.key()
3039 )));
3040 }
3041 selected.sort_by(|left, right| {
3042 (
3043 left.fold_id.as_ref().map(ToString::to_string),
3044 left.prediction_id.clone(),
3045 )
3046 .cmp(&(
3047 right.fold_id.as_ref().map(ToString::to_string),
3048 right.prediction_id.clone(),
3049 ))
3050 });
3051 Ok(selected)
3052}
3053
3054fn select_aggregated_prediction_cache_blocks(
3055 requirement: &BundlePredictionRequirement,
3056 blocks: &[AggregatedPredictionBlock],
3057) -> Result<Vec<AggregatedPredictionBlock>> {
3058 requirement.validate()?;
3059 if requirement.prediction_level == PredictionLevel::Sample {
3060 return Err(DagMlError::RuntimeValidation(format!(
3061 "aggregated prediction cache requirement `{}` must use target or group level",
3062 requirement.key()
3063 )));
3064 }
3065 let mut selected = blocks
3066 .iter()
3067 .filter(|block| {
3068 block.producer_node == requirement.producer_node
3069 && block.partition == requirement.partition
3070 && block.level == requirement.prediction_level
3071 })
3072 .cloned()
3073 .collect::<Vec<_>>();
3074 if selected.is_empty() {
3075 return Err(DagMlError::RuntimeValidation(format!(
3076 "aggregated prediction cache requirement `{}` has no matching prediction blocks",
3077 requirement.key()
3078 )));
3079 }
3080 selected.sort_by(|left, right| {
3081 (
3082 left.fold_id.as_ref().map(ToString::to_string),
3083 left.prediction_id.clone(),
3084 )
3085 .cmp(&(
3086 right.fold_id.as_ref().map(ToString::to_string),
3087 right.prediction_id.clone(),
3088 ))
3089 });
3090 Ok(selected)
3091}
3092
3093fn build_prediction_cache_record_from_selected(
3094 requirement: &BundlePredictionRequirement,
3095 selected: &[PredictionBlock],
3096) -> Result<BundlePredictionCacheRecord> {
3097 requirement.validate()?;
3098 if selected.is_empty() {
3099 return Err(DagMlError::RuntimeValidation(format!(
3100 "prediction cache requirement `{}` has no matching prediction blocks",
3101 requirement.key()
3102 )));
3103 }
3104 let mut fold_ids = BTreeSet::new();
3105 let mut sample_ids = BTreeSet::new();
3106 let mut target_names: Option<Vec<String>> = None;
3107 let mut prediction_width: Option<usize> = None;
3108 let mut row_count = 0usize;
3109 let mut block_records = Vec::new();
3110 for block in selected {
3111 if block.producer_node != requirement.producer_node
3112 || block.partition != requirement.partition
3113 {
3114 return Err(DagMlError::RuntimeValidation(format!(
3115 "prediction cache `{}` contains a block outside the requirement scope",
3116 requirement.key()
3117 )));
3118 }
3119 let width = block.validate_shape()?;
3120 if prediction_width.is_some_and(|expected| expected != width) {
3121 return Err(DagMlError::RuntimeValidation(format!(
3122 "prediction cache `{}` has inconsistent prediction width",
3123 requirement.key()
3124 )));
3125 }
3126 prediction_width = Some(width);
3127 let block_target_names = normalized_prediction_targets(block, width);
3128 if target_names
3129 .as_ref()
3130 .is_some_and(|expected| expected != &block_target_names)
3131 {
3132 return Err(DagMlError::RuntimeValidation(format!(
3133 "prediction cache `{}` has inconsistent target names",
3134 requirement.key()
3135 )));
3136 }
3137 target_names = Some(block_target_names);
3138 if let Some(fold_id) = &block.fold_id {
3139 fold_ids.insert(fold_id.clone());
3140 }
3141 sample_ids.extend(block.sample_ids.iter().cloned());
3142 row_count += block.sample_ids.len();
3143 block_records.push(BundlePredictionBlockCacheRecord {
3144 prediction_id: block.prediction_id.clone(),
3145 fold_id: block.fold_id.clone(),
3146 prediction_level: PredictionLevel::Sample,
3147 row_count: block.sample_ids.len(),
3148 unit_ids: Vec::new(),
3149 sample_ids: block.sample_ids.clone(),
3150 content_fingerprint: stable_json_fingerprint(block)?,
3151 });
3152 }
3153
3154 let record = BundlePredictionCacheRecord {
3155 requirement_key: requirement.key(),
3156 cache_id: format!("prediction-cache:{}", requirement.key()),
3157 cache_namespace_fingerprints: Vec::new(),
3158 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
3159 partition: requirement.partition.clone(),
3160 prediction_level: requirement.prediction_level,
3161 fold_ids: fold_ids.into_iter().collect(),
3162 unit_ids: requirement.unit_ids.clone(),
3163 sample_ids: sample_ids.into_iter().collect(),
3164 prediction_width: prediction_width.unwrap_or_default(),
3165 target_names: target_names.unwrap_or_default(),
3166 block_count: block_records.len(),
3167 row_count,
3168 content_fingerprint: stable_json_fingerprint(selected)?,
3169 blocks: block_records,
3170 };
3171 validate_prediction_cache_matches_requirement(&record, requirement)?;
3172 record.validate()?;
3173 Ok(record)
3174}
3175
3176fn build_aggregated_prediction_cache_record_from_selected(
3177 requirement: &BundlePredictionRequirement,
3178 selected: &[AggregatedPredictionBlock],
3179) -> Result<BundlePredictionCacheRecord> {
3180 requirement.validate()?;
3181 if requirement.prediction_level == PredictionLevel::Sample {
3182 return Err(DagMlError::RuntimeValidation(format!(
3183 "aggregated prediction cache requirement `{}` must use target or group level",
3184 requirement.key()
3185 )));
3186 }
3187 if selected.is_empty() {
3188 return Err(DagMlError::RuntimeValidation(format!(
3189 "aggregated prediction cache requirement `{}` has no matching prediction blocks",
3190 requirement.key()
3191 )));
3192 }
3193 let mut fold_ids = BTreeSet::new();
3194 let mut unit_ids = BTreeSet::new();
3195 let mut target_names: Option<Vec<String>> = None;
3196 let mut prediction_width: Option<usize> = None;
3197 let mut row_count = 0usize;
3198 let mut block_records = Vec::new();
3199 for block in selected {
3200 if block.producer_node != requirement.producer_node
3201 || block.partition != requirement.partition
3202 || block.level != requirement.prediction_level
3203 {
3204 return Err(DagMlError::RuntimeValidation(format!(
3205 "aggregated prediction cache `{}` contains a block outside the requirement scope",
3206 requirement.key()
3207 )));
3208 }
3209 let width = block.validate_shape()?;
3210 if prediction_width.is_some_and(|expected| expected != width) {
3211 return Err(DagMlError::RuntimeValidation(format!(
3212 "aggregated prediction cache `{}` has inconsistent prediction width",
3213 requirement.key()
3214 )));
3215 }
3216 prediction_width = Some(width);
3217 let block_target_names = normalized_aggregated_prediction_targets(block, width);
3218 if target_names
3219 .as_ref()
3220 .is_some_and(|expected| expected != &block_target_names)
3221 {
3222 return Err(DagMlError::RuntimeValidation(format!(
3223 "aggregated prediction cache `{}` has inconsistent target names",
3224 requirement.key()
3225 )));
3226 }
3227 target_names = Some(block_target_names);
3228 if let Some(fold_id) = &block.fold_id {
3229 fold_ids.insert(fold_id.clone());
3230 }
3231 unit_ids.extend(block.unit_ids.iter().cloned());
3232 row_count += block.unit_ids.len();
3233 block_records.push(BundlePredictionBlockCacheRecord {
3234 prediction_id: block.prediction_id.clone(),
3235 fold_id: block.fold_id.clone(),
3236 prediction_level: block.level,
3237 row_count: block.unit_ids.len(),
3238 unit_ids: block.unit_ids.clone(),
3239 sample_ids: Vec::new(),
3240 content_fingerprint: stable_json_fingerprint(block)?,
3241 });
3242 }
3243
3244 let record = BundlePredictionCacheRecord {
3245 requirement_key: requirement.key(),
3246 cache_id: format!("prediction-cache:{}", requirement.key()),
3247 cache_namespace_fingerprints: Vec::new(),
3248 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
3249 partition: requirement.partition.clone(),
3250 prediction_level: requirement.prediction_level,
3251 fold_ids: fold_ids.into_iter().collect(),
3252 unit_ids: unit_ids.into_iter().collect(),
3253 sample_ids: Vec::new(),
3254 prediction_width: prediction_width.unwrap_or_default(),
3255 target_names: target_names.unwrap_or_default(),
3256 block_count: block_records.len(),
3257 row_count,
3258 content_fingerprint: stable_json_fingerprint(selected)?,
3259 blocks: block_records,
3260 };
3261 validate_prediction_cache_matches_requirement(&record, requirement)?;
3262 record.validate()?;
3263 Ok(record)
3264}
3265
3266fn validate_prediction_cache_matches_requirement(
3267 cache: &BundlePredictionCacheRecord,
3268 requirement: &BundlePredictionRequirement,
3269) -> Result<()> {
3270 if cache.requirement_key != requirement.key()
3271 || cache.partition != requirement.partition
3272 || cache.prediction_level != requirement.prediction_level
3273 || cache.fold_ids != requirement.fold_ids
3274 || cache.unit_ids != requirement.unit_ids
3275 || cache.sample_ids != requirement.sample_ids
3276 || cache.prediction_width != requirement.prediction_width
3277 || cache.target_names != requirement.target_names
3278 {
3279 return Err(DagMlError::RuntimeValidation(format!(
3280 "prediction cache `{}` does not match requirement `{}`",
3281 cache.cache_id,
3282 requirement.key()
3283 )));
3284 }
3285 Ok(())
3286}
3287
3288fn normalized_prediction_targets(block: &PredictionBlock, width: usize) -> Vec<String> {
3289 if block.target_names.is_empty() {
3290 (0..width).map(|index| format!("p{index}")).collect()
3291 } else {
3292 block.target_names.clone()
3293 }
3294}
3295
3296fn normalized_aggregated_prediction_targets(
3297 block: &AggregatedPredictionBlock,
3298 width: usize,
3299) -> Vec<String> {
3300 if block.target_names.is_empty() {
3301 (0..width).map(|index| format!("p{index}")).collect()
3302 } else {
3303 block.target_names.clone()
3304 }
3305}
3306
3307fn sample_prediction_units(sample_ids: &[SampleId]) -> Vec<PredictionUnitId> {
3308 sample_ids
3309 .iter()
3310 .cloned()
3311 .map(PredictionUnitId::Sample)
3312 .collect()
3313}
3314
3315fn validate_prediction_units(
3316 label: &str,
3317 expected_level: PredictionLevel,
3318 unit_ids: &[PredictionUnitId],
3319) -> Result<()> {
3320 validate_unique_ids(label, unit_ids)?;
3321 for unit_id in unit_ids {
3322 if unit_id.level() != expected_level {
3323 return Err(DagMlError::RuntimeValidation(format!(
3324 "{label} `{unit_id}` does not match prediction level {:?}",
3325 expected_level
3326 )));
3327 }
3328 }
3329 Ok(())
3330}
3331
3332fn validate_fingerprint(label: &str, value: &str) -> Result<()> {
3333 if value.len() != 64 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
3334 return Err(DagMlError::RuntimeValidation(format!(
3335 "{label} fingerprint must be a 64-character hex digest"
3336 )));
3337 }
3338 Ok(())
3339}
3340
3341fn validate_non_empty(label: &str, value: &str) -> Result<()> {
3342 if value.trim().is_empty() {
3343 return Err(DagMlError::RuntimeValidation(format!("{label} is empty")));
3344 }
3345 Ok(())
3346}
3347
3348fn validate_unique_ids<T>(label: &str, values: &[T]) -> Result<()>
3349where
3350 T: Ord + ToString,
3351{
3352 let mut seen = BTreeSet::new();
3353 for value in values {
3354 if !seen.insert(value) {
3355 return Err(DagMlError::RuntimeValidation(format!(
3356 "duplicate {label} `{}`",
3357 value.to_string()
3358 )));
3359 }
3360 }
3361 Ok(())
3362}
3363
3364#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
3365pub struct ReplayPhaseRequest {
3366 pub bundle_id: BundleId,
3367 pub phase: Phase,
3368 #[serde(default)]
3369 pub data_envelope_keys: Vec<String>,
3370}
3371
3372impl ReplayPhaseRequest {
3373 pub fn validate_for_bundle(&self, bundle: &ExecutionBundle) -> Result<()> {
3374 self.validate_for_bundle_with_prediction_cache_store(bundle, false)
3375 }
3376
3377 pub fn validate_for_bundle_with_prediction_cache_store(
3378 &self,
3379 bundle: &ExecutionBundle,
3380 prediction_cache_available: bool,
3381 ) -> Result<()> {
3382 self.validate_for_bundle_internal(bundle, prediction_cache_available)
3383 }
3384
3385 pub fn validate_for_bundle_with_prediction_cache_payloads(
3386 &self,
3387 bundle: &ExecutionBundle,
3388 prediction_cache_payloads: Option<&BundlePredictionCachePayloadSet>,
3389 ) -> Result<()> {
3390 if let Some(payloads) = prediction_cache_payloads {
3391 payloads.validate_against_bundle(bundle)?;
3392 }
3393 self.validate_for_bundle_internal(bundle, prediction_cache_payloads.is_some())
3394 }
3395
3396 fn validate_for_bundle_internal(
3397 &self,
3398 bundle: &ExecutionBundle,
3399 prediction_cache_available: bool,
3400 ) -> Result<()> {
3401 bundle.validate()?;
3402 if self.bundle_id != bundle.bundle_id {
3403 return Err(DagMlError::RuntimeValidation(format!(
3404 "replay request bundle `{}` does not match bundle `{}`",
3405 self.bundle_id, bundle.bundle_id
3406 )));
3407 }
3408 if !matches!(self.phase, Phase::Predict | Phase::Explain | Phase::Refit) {
3409 return Err(DagMlError::RuntimeValidation(format!(
3410 "bundle replay phase {:?} is not supported",
3411 self.phase
3412 )));
3413 }
3414 if self.phase == Phase::Refit && !bundle.prediction_requirements.is_empty() {
3415 if prediction_cache_available {
3416 return self.validate_data_envelope_keys(bundle);
3417 }
3418 return Err(DagMlError::RuntimeValidation(format!(
3419 "bundle `{}` cannot replay REFIT because it depends on {} OOF prediction requirement(s) but stores only prediction cache manifests",
3420 bundle.bundle_id,
3421 bundle.prediction_requirements.len()
3422 )));
3423 }
3424 self.validate_data_envelope_keys(bundle)
3425 }
3426
3427 fn validate_data_envelope_keys(&self, bundle: &ExecutionBundle) -> Result<()> {
3428 let expected = bundle
3429 .data_requirements
3430 .iter()
3431 .map(BundleDataRequirement::key)
3432 .collect::<BTreeSet<_>>();
3433 let mut requested = BTreeSet::new();
3434 for key in &self.data_envelope_keys {
3435 if key.trim().is_empty() {
3436 return Err(DagMlError::RuntimeValidation(
3437 "replay request contains an empty data envelope key".to_string(),
3438 ));
3439 }
3440 if !requested.insert(key.as_str()) {
3441 return Err(DagMlError::RuntimeValidation(format!(
3442 "replay request contains duplicate data envelope key `{key}`"
3443 )));
3444 }
3445 if !expected.contains(key.as_str()) {
3446 return Err(DagMlError::RuntimeValidation(format!(
3447 "replay request references unknown data envelope key `{key}`"
3448 )));
3449 }
3450 }
3451 for requirement in &bundle.data_requirements {
3452 let key = requirement.key();
3453 if !requested.contains(key.as_str()) {
3454 return Err(DagMlError::RuntimeValidation(format!(
3455 "replay request is missing data envelope key `{key}`"
3456 )));
3457 }
3458 }
3459 Ok(())
3460 }
3461}
3462
3463#[cfg(test)]
3464mod tests {
3465 use super::*;
3466 use crate::controller::{ControllerManifest, ControllerRegistry};
3467 use crate::data::{
3468 AggregateRepresentation, RepresentationCardinality, RepresentationCompatibilityOutcome,
3469 RepresentationCompatibilityReport, RepresentationMissingSourcePolicy, RepresentationPlan,
3470 RepresentationReplayManifest,
3471 };
3472 use crate::dsl::{compile_pipeline_dsl_with_generation, PipelineDslSpec};
3473 use crate::graph::GraphSpec;
3474 use crate::ids::{ArtifactId, ControllerId, FoldId, LineageId, RunId, SampleId, TargetId};
3475 use crate::plan::{build_execution_plan, CampaignSpec};
3476 use crate::relation::EntityUnitLevel;
3477 use crate::selection::{
3478 select_candidate, CandidateScore, MetricObjective, SelectionMetric, SelectionPolicy,
3479 };
3480
3481 fn plan() -> ExecutionPlan {
3482 let graph: GraphSpec =
3483 serde_json::from_str(include_str!("../tests/fixtures/package/minimal_graph.json"))
3484 .unwrap();
3485 let campaign: CampaignSpec = serde_json::from_str(include_str!(
3486 "../tests/fixtures/package/campaign_oof_generation.json"
3487 ))
3488 .unwrap();
3489 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3490 "../tests/fixtures/package/controller_manifests.json"
3491 ))
3492 .unwrap();
3493 let mut registry = ControllerRegistry::new();
3494 for manifest in manifests {
3495 registry.register(manifest).unwrap();
3496 }
3497 build_execution_plan("plan:bundle", graph, campaign, ®istry).unwrap()
3498 }
3499
3500 fn methods_hpo_resume_state() -> MethodsHpoResumeState {
3501 let node = NodeId::new("model:hpo").unwrap();
3502 let variant = crate::generation::VariantPlan {
3503 variant_id: VariantId::new("hpo:trial:1").unwrap(),
3504 choices: BTreeMap::new(),
3505 fingerprint: "c".repeat(64),
3506 seed: Some(7),
3507 };
3508 let sample = SampleId::new("sample:hpo.1").unwrap();
3509 let target = RegressionTargetBlock {
3510 level: PredictionLevel::Sample,
3511 unit_ids: vec![PredictionUnitId::Sample(sample.clone())],
3512 values: vec![vec![1.0]],
3513 target_names: vec!["y".to_string()],
3514 };
3515 let report = RegressionMetricReport {
3516 prediction_id: Some("prediction:model:hpo.avg".to_string()),
3517 producer_node: node.clone(),
3518 producer_port: Some("prediction".to_string()),
3519 variant_id: Some(variant.variant_id.clone()),
3520 variant_label: None,
3521 partition: PredictionPartition::Validation,
3522 fold_id: Some(FoldId::new("avg").unwrap()),
3523 level: PredictionLevel::Sample,
3524 row_count: 1,
3525 target_width: 1,
3526 target_names: vec!["y".to_string()],
3527 metrics: BTreeMap::from([("rmse".to_string(), 0.25)]),
3528 };
3529 let payload = vec![1_u8, 2, 3];
3530 MethodsHpoResumeState {
3531 schema_version: METHODS_HPO_RESUME_STATE_SCHEMA_VERSION,
3532 checkpoint: crate::hpo::N4moptCheckpointArtifact {
3533 schema_version: crate::hpo::N4MOPT_CHECKPOINT_SCHEMA_VERSION,
3534 artifact_kind: crate::hpo::N4MOPT_ARTIFACT_KIND.to_string(),
3535 format: crate::hpo::N4MOPT_FORMAT.to_string(),
3536 binding: crate::hpo::HpoStudyBinding {
3537 controller_id: "controller:hpo".to_string(),
3538 study_id: "study:hpo".to_string(),
3539 search_space_fingerprint: "a".repeat(64),
3540 optimizer_fingerprint: "b".repeat(64),
3541 },
3542 methods_abi: "n4m:1".to_string(),
3543 payload_sha256: format!("{:x}", sha2::Sha256::digest(&payload)),
3544 opaque_payload: payload,
3545 },
3546 provenance: MethodsHpoResumeProvenance {
3547 graph_fingerprint: "d".repeat(64),
3548 campaign_fingerprint: "e".repeat(64),
3549 controller_fingerprint: "f".repeat(64),
3550 data_identities_fingerprint: "1".repeat(64),
3551 fold_set_fingerprint: "2".repeat(64),
3552 training_influence_fingerprint: "3".repeat(64),
3553 relation_fingerprint: "4".repeat(64),
3554 selection: MethodsHpoResumeSelection {
3555 selection_id: "selection:hpo".to_string(),
3556 target_node_id: node.clone(),
3557 producer_port: "prediction".to_string(),
3558 metric: "rmse".to_string(),
3559 },
3560 },
3561 operation_id: "hpo:bundle-test".to_string(),
3562 controller_id: ControllerId::new("controller:hpo").unwrap(),
3563 target_node_id: node.clone(),
3564 incumbent: MethodsHpoNativeIncumbent {
3565 trial_id: 1,
3566 score: 0.25,
3567 metric: "rmse".to_string(),
3568 direction: crate::hpo::HpoDirection::Minimize,
3569 variant_id: variant.variant_id.clone(),
3570 },
3571 trial_history_len: 1,
3572 terminal_trials: vec![MethodsHpoTerminalEvidence {
3573 trial: crate::hpo::HpoTrial {
3574 id: 1,
3575 ask_sequence: 1,
3576 terminal_sequence: Some(1),
3577 parameters: BTreeMap::new(),
3578 parameter_order: Vec::new(),
3579 status: crate::hpo::HpoTrialStatus::Completed,
3580 score: Some(0.25),
3581 rung: 0,
3582 duration: 0.0,
3583 intermediates: Vec::new(),
3584 failure: None,
3585 },
3586 variant_id: Some(variant.variant_id.clone()),
3587 }],
3588 completed_proposals: vec![MethodsHpoCompletedProposal {
3589 trial_id: 1,
3590 variant: variant.clone(),
3591 }],
3592 completed_reports: vec![MethodsHpoCompletedReport {
3593 trial_id: 1,
3594 variant_id: variant.variant_id.clone(),
3595 terminal_state: MethodsHpoCompletedState::Completed,
3596 score: 0.25,
3597 report,
3598 }],
3599 candidates: vec![MethodsHpoCandidateEvidence {
3600 trial_id: 1,
3601 variant_id: variant.variant_id.clone(),
3602 variant_label: None,
3603 predictions: vec![PredictionBlock {
3604 prediction_id: Some("prediction:model:hpo.fold0".to_string()),
3605 producer_node: node.clone(),
3606 producer_port: Some("prediction".to_string()),
3607 partition: PredictionPartition::Validation,
3608 fold_id: Some(FoldId::new("fold:0").unwrap()),
3609 sample_ids: vec![sample.clone()],
3610 values: vec![vec![1.25]],
3611 target_names: vec!["y".to_string()],
3612 }],
3613 regression_targets: vec![target.clone()],
3614 oof_average: MethodsHpoOofAverage {
3615 predictions: AggregatedPredictionBlock {
3616 prediction_id: Some("prediction:model:hpo.avg".to_string()),
3617 producer_node: node.clone(),
3618 producer_port: Some("prediction".to_string()),
3619 partition: PredictionPartition::Validation,
3620 fold_id: Some(FoldId::new("avg").unwrap()),
3621 level: PredictionLevel::Sample,
3622 unit_ids: vec![PredictionUnitId::Sample(sample)],
3623 values: vec![vec![1.25]],
3624 target_names: vec!["y".to_string()],
3625 },
3626 y_true: target,
3627 },
3628 lineage: vec![crate::runtime::LineageRecord {
3629 record_id: LineageId::new("lineage:hpo.1").unwrap(),
3630 run_id: RunId::new("run:hpo").unwrap(),
3631 node_id: node,
3632 phase: Phase::FitCv,
3633 controller_id: ControllerId::new("controller:methods.pls").unwrap(),
3634 controller_version: "1".to_string(),
3635 variant_id: Some(variant.variant_id),
3636 fold_id: Some(FoldId::new("fold:0").unwrap()),
3637 branch_path: Vec::new(),
3638 input_lineage: Vec::new(),
3639 artifact_refs: Vec::new(),
3640 params_fingerprint: "5".repeat(64),
3641 data_model_shape_fingerprint: None,
3642 aggregation_policy_fingerprint: None,
3643 seed: Some(7),
3644 unsafe_flags: BTreeSet::new(),
3645 metrics: BTreeMap::new(),
3646 loss_attestations: Vec::new(),
3647 early_stopping_records: Vec::new(),
3648 }],
3649 }],
3650 }
3651 }
3652
3653 fn branch_merge_plan() -> ExecutionPlan {
3654 let graph: GraphSpec = serde_json::from_str(include_str!(
3655 "../tests/fixtures/package/branch_merge_oof_graph.json"
3656 ))
3657 .unwrap();
3658 let campaign: CampaignSpec = serde_json::from_str(include_str!(
3659 "../tests/fixtures/package/campaign_branch_merge_oof.json"
3660 ))
3661 .unwrap();
3662 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3663 "../tests/fixtures/package/controller_manifests.json"
3664 ))
3665 .unwrap();
3666 let mut registry = ControllerRegistry::new();
3667 for manifest in manifests {
3668 registry.register(manifest).unwrap();
3669 }
3670 build_execution_plan("plan:branch.merge.bundle", graph, campaign, ®istry).unwrap()
3671 }
3672
3673 fn separation_concat_merge_plan() -> ExecutionPlan {
3680 let graph: GraphSpec = serde_json::from_str(include_str!(
3681 "../tests/fixtures/package/separation_branch_concat_merge_oof_graph.json"
3682 ))
3683 .unwrap();
3684 let campaign: CampaignSpec = serde_json::from_str(include_str!(
3685 "../tests/fixtures/package/campaign_separation_branch_concat_merge_oof.json"
3686 ))
3687 .unwrap();
3688 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3689 "../tests/fixtures/package/controller_manifests.json"
3690 ))
3691 .unwrap();
3692 let mut registry = ControllerRegistry::new();
3693 for manifest in manifests {
3694 registry.register(manifest).unwrap();
3695 }
3696 build_execution_plan(
3697 "plan:separation.concat.merge.bundle",
3698 graph,
3699 campaign,
3700 ®istry,
3701 )
3702 .unwrap()
3703 }
3704
3705 fn separation_branch_requirement(
3709 producer_node: &str,
3710 partition_samples: &[&str],
3711 partition_folds: &[&str],
3712 ) -> BundlePredictionRequirement {
3713 BundlePredictionRequirement {
3714 producer_node: NodeId::new(producer_node).unwrap(),
3715 source_port: "oof".to_string(),
3716 consumer_node: NodeId::new("merge:sites").unwrap(),
3717 target_port: format!("oof_{producer_node}"),
3718 partition: PredictionPartition::Validation,
3719 prediction_level: PredictionLevel::Sample,
3720 fold_ids: partition_folds
3721 .iter()
3722 .map(|f| FoldId::new(*f).unwrap())
3723 .collect(),
3724 unit_ids: Vec::new(),
3725 sample_ids: partition_samples
3726 .iter()
3727 .map(|s| SampleId::new(*s).unwrap())
3728 .collect(),
3729 prediction_width: 1,
3730 target_names: vec!["y".to_string()],
3731 }
3732 }
3733
3734 fn separation_branch_blocks(
3737 producer_node: &str,
3738 fold0_sample: &str,
3739 fold1_sample: &str,
3740 offset: f64,
3741 ) -> Vec<PredictionBlock> {
3742 let producer_node = NodeId::new(producer_node).unwrap();
3743 vec![
3744 PredictionBlock {
3745 prediction_id: Some(format!("prediction:{producer_node}:fold0")),
3746 producer_node: producer_node.clone(),
3747 producer_port: Some("pred".to_string()),
3748 partition: PredictionPartition::Validation,
3749 fold_id: Some(FoldId::new("fold:0").unwrap()),
3750 sample_ids: vec![SampleId::new(fold0_sample).unwrap()],
3751 values: vec![vec![offset + 0.1]],
3752 target_names: vec!["y".to_string()],
3753 },
3754 PredictionBlock {
3755 prediction_id: Some(format!("prediction:{producer_node}:fold1")),
3756 producer_node,
3757 producer_port: Some("pred".to_string()),
3758 partition: PredictionPartition::Validation,
3759 fold_id: Some(FoldId::new("fold:1").unwrap()),
3760 sample_ids: vec![SampleId::new(fold1_sample).unwrap()],
3761 values: vec![vec![offset + 0.2]],
3762 target_names: vec!["y".to_string()],
3763 },
3764 ]
3765 }
3766
3767 fn executable_dsl_plan() -> ExecutionPlan {
3768 let spec: PipelineDslSpec = serde_json::from_str(include_str!(
3769 "../tests/fixtures/package/pipeline_dsl_branch_merge_executable.json"
3770 ))
3771 .unwrap();
3772 let compiled = compile_pipeline_dsl_with_generation(&spec).unwrap();
3773 let manifests: Vec<ControllerManifest> = serde_json::from_str(include_str!(
3774 "../tests/fixtures/package/controller_manifests.json"
3775 ))
3776 .unwrap();
3777 let mut registry = ControllerRegistry::new();
3778 for manifest in manifests {
3779 registry.register(manifest).unwrap();
3780 }
3781 build_execution_plan(
3782 "plan:dsl.branch.merge.bundle",
3783 compiled.graph,
3784 compiled.campaign_template,
3785 ®istry,
3786 )
3787 .unwrap()
3788 }
3789
3790 fn branch_merge_selection_decisions() -> BTreeMap<String, SelectionDecision> {
3791 serde_json::from_str(include_str!(
3792 "../tests/fixtures/package/bundle/selection_decisions_branch_merge.json"
3793 ))
3794 .unwrap()
3795 }
3796
3797 fn refit_artifact(
3798 plan: &ExecutionPlan,
3799 node_id: &str,
3800 data_requirement_keys: Vec<String>,
3801 prediction_requirement_keys: Vec<String>,
3802 ) -> RefitArtifactRecord {
3803 let node_id = NodeId::new(node_id).unwrap();
3804 let node_plan = plan.node_plans.get(&node_id).unwrap();
3805 RefitArtifactRecord {
3806 node_id: node_plan.node_id.clone(),
3807 controller_id: node_plan.controller_id.clone(),
3808 artifact: ArtifactRef {
3809 id: ArtifactId::new(format!("artifact:{}:refit", node_plan.node_id)).unwrap(),
3810 kind: "mock_model".to_string(),
3811 controller_id: node_plan.controller_id.clone(),
3812 backend: None,
3813 uri: None,
3814 content_fingerprint: None,
3815 size_bytes: Some(128),
3816 plugin: None,
3817 plugin_version: None,
3818 },
3819 params_fingerprint: node_plan.params_fingerprint.clone(),
3820 training_loss_fingerprint: node_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
3821 data_requirement_keys,
3822 prediction_requirement_keys,
3823 }
3824 }
3825
3826 fn branch_merge_samples() -> Vec<SampleId> {
3827 vec![
3828 SampleId::new("sample:1").unwrap(),
3829 SampleId::new("sample:2").unwrap(),
3830 SampleId::new("sample:3").unwrap(),
3831 SampleId::new("sample:4").unwrap(),
3832 ]
3833 }
3834
3835 fn branch_merge_requirement(
3836 producer_node: &str,
3837 target_port: &str,
3838 ) -> BundlePredictionRequirement {
3839 BundlePredictionRequirement {
3840 producer_node: NodeId::new(producer_node).unwrap(),
3841 source_port: "oof".to_string(),
3842 consumer_node: NodeId::new("merge:stack.pred_plus_original.meta:ridge").unwrap(),
3843 target_port: target_port.to_string(),
3844 partition: PredictionPartition::Validation,
3845 prediction_level: PredictionLevel::Sample,
3846 fold_ids: vec![
3847 FoldId::new("fold:0").unwrap(),
3848 FoldId::new("fold:1").unwrap(),
3849 ],
3850 unit_ids: Vec::new(),
3851 sample_ids: branch_merge_samples(),
3852 prediction_width: 1,
3853 target_names: vec!["y".to_string()],
3854 }
3855 }
3856
3857 fn branch_merge_prediction_blocks(producer_node: &str, offset: f64) -> Vec<PredictionBlock> {
3858 let producer_node = NodeId::new(producer_node).unwrap();
3859 let samples = branch_merge_samples();
3860 vec![
3861 PredictionBlock {
3862 prediction_id: Some(format!("prediction:{producer_node}:fold0")),
3863 producer_node: producer_node.clone(),
3864 producer_port: Some("pred".to_string()),
3865 partition: PredictionPartition::Validation,
3866 fold_id: Some(FoldId::new("fold:0").unwrap()),
3867 sample_ids: samples[0..2].to_vec(),
3868 values: vec![vec![offset + 0.1], vec![offset + 0.2]],
3869 target_names: vec!["y".to_string()],
3870 },
3871 PredictionBlock {
3872 prediction_id: Some(format!("prediction:{producer_node}:fold1")),
3873 producer_node,
3874 producer_port: Some("pred".to_string()),
3875 partition: PredictionPartition::Validation,
3876 fold_id: Some(FoldId::new("fold:1").unwrap()),
3877 sample_ids: samples[2..4].to_vec(),
3878 values: vec![vec![offset + 0.3], vec![offset + 0.4]],
3879 target_names: vec!["y".to_string()],
3880 },
3881 ]
3882 }
3883
3884 fn decision() -> SelectionDecision {
3885 select_candidate(
3886 &SelectionPolicy {
3887 id: "select:merge".to_string(),
3888 metric: SelectionMetric {
3889 name: "rmse".to_string(),
3890 objective: MetricObjective::Minimize,
3891 },
3892 required_metric_level: Some(crate::policy::PredictionLevel::Sample),
3893 require_finite: true,
3894 evaluation_scope: None,
3895 refit_slot_plan: None,
3896 stacking_fit_contract: None,
3897 reduction_id: None,
3898 },
3899 &[
3900 CandidateScore {
3901 candidate_id: "model:base".to_string(),
3902 metrics: BTreeMap::from([("rmse".to_string(), 1.0)]),
3903 metadata: BTreeMap::from([(
3904 "metric_level".to_string(),
3905 serde_json::Value::String("sample".to_string()),
3906 )]),
3907 },
3908 CandidateScore {
3909 candidate_id: "model:other".to_string(),
3910 metrics: BTreeMap::from([("rmse".to_string(), 2.0)]),
3911 metadata: BTreeMap::from([(
3912 "metric_level".to_string(),
3913 serde_json::Value::String("sample".to_string()),
3914 )]),
3915 },
3916 ],
3917 )
3918 .unwrap()
3919 }
3920
3921 fn selected_model_base_decision() -> SelectionDecision {
3922 decision()
3923 }
3924
3925 fn model_base_refit_artifact(plan: &ExecutionPlan) -> RefitArtifactRecord {
3926 let model_plan = plan
3927 .node_plans
3928 .get(&NodeId::new("model:base").unwrap())
3929 .unwrap();
3930 RefitArtifactRecord {
3931 node_id: model_plan.node_id.clone(),
3932 controller_id: model_plan.controller_id.clone(),
3933 artifact: ArtifactRef {
3934 id: ArtifactId::new("artifact:model:base:refit").unwrap(),
3935 kind: "sklearn_pickle".to_string(),
3936 controller_id: model_plan.controller_id.clone(),
3937 backend: None,
3938 uri: None,
3939 content_fingerprint: None,
3940 size_bytes: Some(128),
3941 plugin: None,
3942 plugin_version: None,
3943 },
3944 params_fingerprint: model_plan.params_fingerprint.clone(),
3945 training_loss_fingerprint: model_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
3946 data_requirement_keys: vec!["model:base.x".to_string()],
3947 prediction_requirement_keys: Vec::new(),
3948 }
3949 }
3950
3951 #[test]
3952 fn builds_bundle_from_execution_plan() {
3953 let plan = plan();
3954 let artifact = model_base_refit_artifact(&plan);
3955
3956 let bundle = build_execution_bundle(
3957 BundleId::new("bundle:demo").unwrap(),
3958 &plan,
3959 Some(plan.variants[0].variant_id.clone()),
3960 BTreeMap::from([("merge".to_string(), decision())]),
3961 vec![artifact],
3962 )
3963 .unwrap();
3964
3965 bundle.validate_against_plan(&plan).unwrap();
3966 assert_eq!(bundle.data_requirements.len(), 1);
3967 }
3968
3969 #[test]
3970 fn bundle_data_requirements_accept_d7_replay_contracts() {
3971 let plan = plan();
3972 let artifact = model_base_refit_artifact(&plan);
3973 let mut bundle = build_execution_bundle(
3974 BundleId::new("bundle:d7.replay").unwrap(),
3975 &plan,
3976 Some(plan.variants[0].variant_id.clone()),
3977 BTreeMap::from([("merge".to_string(), decision())]),
3978 vec![artifact],
3979 )
3980 .unwrap();
3981 let relation_fingerprint = bundle.data_requirements[0]
3982 .relation_fingerprint
3983 .clone()
3984 .unwrap_or_else(|| "a".repeat(64));
3985 bundle.data_requirements[0].representation_replay_manifest =
3986 Some(RepresentationReplayManifest {
3987 manifest_id: "repr:d7.bundle".to_string(),
3988 representation_plan: RepresentationPlan::Aggregate(AggregateRepresentation {
3989 input_unit_level: EntityUnitLevel::Observation,
3990 output_unit_level: EntityUnitLevel::PhysicalSample,
3991 reducer_id: None,
3992 method: Some("mean".to_string()),
3993 cardinality: RepresentationCardinality::ManyToOne,
3994 }),
3995 combination_plan: None,
3996 output_unit_level: EntityUnitLevel::PhysicalSample,
3997 output_representation: Some("tabular_numeric".to_string()),
3998 relation_fingerprint: Some(relation_fingerprint.clone()),
3999 feature_schema_fingerprint: Some("b".repeat(64)),
4000 final_reduction_id: None,
4001 sample_observation_mapping: Vec::new(),
4002 combo_selection: Vec::new(),
4003 qc_policy_refs: Vec::new(),
4004 outlier_policy_refs: Vec::new(),
4005 missing_source_policy: None,
4006 missing_repetition_policy: None,
4007 prediction_representation: None,
4008 final_output_unit_level: Some(EntityUnitLevel::PhysicalSample),
4009 train_compatibility: None,
4010 predict_compatibility: None,
4011 metadata: BTreeMap::new(),
4012 });
4013 bundle.data_requirements[0].representation_compatibility =
4014 Some(RepresentationCompatibilityReport {
4015 policy: RepresentationMissingSourcePolicy::Strict,
4016 outcome: RepresentationCompatibilityOutcome::Compatible,
4017 fallback_used: None,
4018 warning_severity: None,
4019 affected_source_count: 0,
4020 affected_repetition_count: 0,
4021 affected_sample_count: 0,
4022 train_relation_fingerprint: Some(relation_fingerprint),
4023 predict_relation_fingerprint: None,
4024 train_unit_count: Some(2),
4025 predict_unit_count: Some(2),
4026 fixed_width_required: false,
4027 final_reducer_stabilizes_output: true,
4028 cartesian_combo_count_changed: false,
4029 late_fusion_branch_delta: false,
4030 messages: Vec::new(),
4031 metadata: BTreeMap::new(),
4032 });
4033 bundle.validate_against_plan(&plan).unwrap();
4034
4035 bundle.data_requirements[0]
4036 .representation_replay_manifest
4037 .as_mut()
4038 .unwrap()
4039 .relation_fingerprint = Some("c".repeat(64));
4040 if bundle.data_requirements[0].relation_fingerprint.is_some() {
4041 assert!(bundle.validate().is_err());
4042 }
4043 }
4044
4045 #[test]
4046 fn d9_negative_prediction_cache_refuses_missing_aggregated_unit_ids() {
4047 let cache = BundlePredictionCacheRecord {
4048 requirement_key: "model:base.oof->model:meta.pred".to_string(),
4049 cache_id: "prediction-cache:d9.missing-units".to_string(),
4050 cache_namespace_fingerprints: Vec::new(),
4051 format: BUNDLE_PREDICTION_CACHE_FORMAT.to_string(),
4052 partition: PredictionPartition::Validation,
4053 prediction_level: PredictionLevel::Target,
4054 fold_ids: vec![FoldId::new("fold:0").unwrap()],
4055 unit_ids: Vec::new(),
4056 sample_ids: Vec::new(),
4057 prediction_width: 1,
4058 target_names: vec!["y".to_string()],
4059 block_count: 1,
4060 row_count: 1,
4061 content_fingerprint: "d".repeat(64),
4062 blocks: vec![BundlePredictionBlockCacheRecord {
4063 prediction_id: Some("prediction:d9.target.fold0".to_string()),
4064 fold_id: Some(FoldId::new("fold:0").unwrap()),
4065 prediction_level: PredictionLevel::Target,
4066 row_count: 1,
4067 unit_ids: vec![PredictionUnitId::Target(TargetId::new("target:a").unwrap())],
4068 sample_ids: Vec::new(),
4069 content_fingerprint: "e".repeat(64),
4070 }],
4071 };
4072
4073 let error = cache.validate().unwrap_err().to_string();
4074 assert!(
4075 error.contains("row_count does not match unique unit ids"),
4076 "unexpected D9 missing-unit-id cache error: {error}"
4077 );
4078 }
4079
4080 #[test]
4081 fn refit_artifact_validation_checks_portable_artifact_metadata() {
4082 let plan = plan();
4083 let mut artifact = model_base_refit_artifact(&plan);
4084 artifact.artifact.backend = Some(crate::runtime::ArtifactBackend::Joblib);
4085 artifact.artifact.uri = Some("artifacts/model.joblib".to_string());
4086 artifact.artifact.content_fingerprint = Some("c".repeat(64));
4087 artifact.artifact.plugin = Some("dagml.sklearn".to_string());
4088 artifact.artifact.plugin_version = Some("1.0.0".to_string());
4089 artifact.validate().unwrap();
4090
4091 artifact.artifact.content_fingerprint = Some("short".to_string());
4092 assert!(artifact
4093 .validate()
4094 .unwrap_err()
4095 .to_string()
4096 .contains("artifact content fingerprint"));
4097 }
4098
4099 #[test]
4100 fn bundle_selections_must_match_plan_and_refit_artifacts() {
4101 let plan = plan();
4102 let artifact = model_base_refit_artifact(&plan);
4103 let valid = build_execution_bundle(
4104 BundleId::new("bundle:selected.model").unwrap(),
4105 &plan,
4106 Some(plan.variants[0].variant_id.clone()),
4107 BTreeMap::from([("model".to_string(), selected_model_base_decision())]),
4108 vec![artifact.clone()],
4109 )
4110 .unwrap();
4111 valid.validate_against_plan(&plan).unwrap();
4112
4113 assert!(build_execution_bundle(
4114 BundleId::new("bundle:selected.model.missing.artifact").unwrap(),
4115 &plan,
4116 Some(plan.variants[0].variant_id.clone()),
4117 BTreeMap::from([("model".to_string(), selected_model_base_decision())]),
4118 Vec::new(),
4119 )
4120 .is_err());
4121
4122 let mut missing_level = selected_model_base_decision();
4123 missing_level.metric_level = None;
4124 assert!(build_execution_bundle(
4125 BundleId::new("bundle:selected.missing.level").unwrap(),
4126 &plan,
4127 Some(plan.variants[0].variant_id.clone()),
4128 BTreeMap::from([("model".to_string(), missing_level)]),
4129 vec![artifact.clone()],
4130 )
4131 .is_err());
4132
4133 let mut wrong_level = selected_model_base_decision();
4134 wrong_level.metric_level = Some(crate::policy::PredictionLevel::Target);
4135 assert!(build_execution_bundle(
4136 BundleId::new("bundle:selected.wrong.level").unwrap(),
4137 &plan,
4138 Some(plan.variants[0].variant_id.clone()),
4139 BTreeMap::from([("model".to_string(), wrong_level)]),
4140 vec![artifact.clone()],
4141 )
4142 .is_err());
4143
4144 let mut unknown = selected_model_base_decision();
4145 unknown.selected_candidate_id = "model:missing".to_string();
4146 unknown.ranked_candidates[0].candidate_id = "model:missing".to_string();
4147 assert!(build_execution_bundle(
4148 BundleId::new("bundle:selected.unknown").unwrap(),
4149 &plan,
4150 Some(plan.variants[0].variant_id.clone()),
4151 BTreeMap::from([("model".to_string(), unknown)]),
4152 vec![artifact],
4153 )
4154 .is_err());
4155 }
4156
4157 #[test]
4158 fn bundle_artifact_params_follow_selected_generation_variant() {
4159 let plan = executable_dsl_plan();
4160 let selected_variant = &plan.variants[0];
4161 let node_plan = plan
4162 .node_plans
4163 .get(&NodeId::new("branch:b0.model:ridge").unwrap())
4164 .unwrap();
4165 let effective_params = selected_variant
4166 .effective_params_for_node(&node_plan.node_id, &node_plan.params)
4167 .unwrap();
4168 let effective_fingerprint = stable_json_fingerprint(&effective_params).unwrap();
4169 assert_ne!(effective_fingerprint, node_plan.params_fingerprint);
4170
4171 let artifact = RefitArtifactRecord {
4172 node_id: node_plan.node_id.clone(),
4173 controller_id: node_plan.controller_id.clone(),
4174 artifact: ArtifactRef {
4175 id: ArtifactId::new("artifact:branch:b0.model:ridge:refit").unwrap(),
4176 kind: "mock_model".to_string(),
4177 controller_id: node_plan.controller_id.clone(),
4178 backend: None,
4179 uri: None,
4180 content_fingerprint: None,
4181 size_bytes: Some(128),
4182 plugin: None,
4183 plugin_version: None,
4184 },
4185 params_fingerprint: effective_fingerprint,
4186 training_loss_fingerprint: node_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
4187 data_requirement_keys: vec!["branch:b0.model:ridge.x".to_string()],
4188 prediction_requirement_keys: Vec::new(),
4189 };
4190
4191 build_execution_bundle(
4192 BundleId::new("bundle:dsl.variant.params").unwrap(),
4193 &plan,
4194 Some(selected_variant.variant_id.clone()),
4195 BTreeMap::new(),
4196 vec![artifact.clone()],
4197 )
4198 .unwrap();
4199
4200 let mut stale_artifact = artifact;
4201 stale_artifact.params_fingerprint = node_plan.params_fingerprint.clone();
4202 let error = build_execution_bundle(
4203 BundleId::new("bundle:dsl.variant.params.stale").unwrap(),
4204 &plan,
4205 Some(selected_variant.variant_id.clone()),
4206 BTreeMap::new(),
4207 vec![stale_artifact],
4208 )
4209 .unwrap_err();
4210 assert!(format!("{error}").contains("artifact params"));
4211 }
4212
4213 #[test]
4214 fn branch_merge_bundle_links_selected_refits_and_fold_aligned_oof_caches() {
4215 let plan = branch_merge_plan();
4216 let b0_requirement = branch_merge_requirement("branch:b0.model:ridge", "b0_oof");
4217 let b1_requirement = branch_merge_requirement("branch:b1.model:rf", "b1_oof");
4218 let b0_cache = build_prediction_cache_record(
4219 &b0_requirement,
4220 &branch_merge_prediction_blocks("branch:b0.model:ridge", 0.0),
4221 )
4222 .unwrap();
4223 let b1_cache = build_prediction_cache_record(
4224 &b1_requirement,
4225 &branch_merge_prediction_blocks("branch:b1.model:rf", 1.0),
4226 )
4227 .unwrap();
4228 let b0_artifact = refit_artifact(
4229 &plan,
4230 "branch:b0.model:ridge",
4231 vec!["branch:b0.model:ridge.x".to_string()],
4232 Vec::new(),
4233 );
4234 let b1_artifact = refit_artifact(
4235 &plan,
4236 "branch:b1.model:rf",
4237 vec!["branch:b1.model:rf.x".to_string()],
4238 Vec::new(),
4239 );
4240 let merge_artifact = refit_artifact(
4241 &plan,
4242 "merge:stack.pred_plus_original.meta:ridge",
4243 vec!["merge:stack.pred_plus_original.meta:ridge.x_original".to_string()],
4244 vec![b0_requirement.key(), b1_requirement.key()],
4245 );
4246
4247 let bundle = build_execution_bundle_with_prediction_contracts(
4248 BundleId::new("bundle:branch.merge.selected.refit").unwrap(),
4249 &plan,
4250 Some(plan.variants[0].variant_id.clone()),
4251 branch_merge_selection_decisions(),
4252 vec![
4253 b0_artifact.clone(),
4254 b1_artifact.clone(),
4255 merge_artifact.clone(),
4256 ],
4257 vec![b0_requirement.clone(), b1_requirement.clone()],
4258 vec![b0_cache.clone(), b1_cache.clone()],
4259 )
4260 .unwrap();
4261 bundle.validate_against_plan(&plan).unwrap();
4262 assert_eq!(bundle.selections.len(), 3);
4263 assert_eq!(bundle.prediction_requirements.len(), 2);
4264 assert_eq!(
4265 bundle.refit_artifacts[2].data_requirement_keys,
4266 vec!["merge:stack.pred_plus_original.meta:ridge.x_original"]
4267 );
4268 assert_eq!(
4269 bundle.refit_artifacts[2].prediction_requirement_keys,
4270 vec![
4271 "branch:b0.model:ridge.oof->merge:stack.pred_plus_original.meta:ridge.b0_oof",
4272 "branch:b1.model:rf.oof->merge:stack.pred_plus_original.meta:ridge.b1_oof",
4273 ]
4274 );
4275
4276 assert!(build_execution_bundle_with_prediction_contracts(
4277 BundleId::new("bundle:branch.merge.missing.branch.refit").unwrap(),
4278 &plan,
4279 Some(plan.variants[0].variant_id.clone()),
4280 branch_merge_selection_decisions(),
4281 vec![b0_artifact.clone(), merge_artifact.clone()],
4282 vec![b0_requirement.clone(), b1_requirement.clone()],
4283 vec![b0_cache.clone(), b1_cache.clone()],
4284 )
4285 .is_err());
4286
4287 let mut misaligned_cache = b0_cache;
4288 misaligned_cache.blocks[0].sample_ids = vec![
4289 SampleId::new("sample:1").unwrap(),
4290 SampleId::new("sample:3").unwrap(),
4291 ];
4292 misaligned_cache.blocks[1].sample_ids = vec![
4293 SampleId::new("sample:2").unwrap(),
4294 SampleId::new("sample:4").unwrap(),
4295 ];
4296 let error = build_execution_bundle_with_prediction_contracts(
4297 BundleId::new("bundle:branch.merge.misaligned.oof.cache").unwrap(),
4298 &plan,
4299 Some(plan.variants[0].variant_id.clone()),
4300 branch_merge_selection_decisions(),
4301 vec![b0_artifact, b1_artifact, merge_artifact],
4302 vec![b0_requirement, b1_requirement],
4303 vec![misaligned_cache, b1_cache],
4304 )
4305 .unwrap_err()
4306 .to_string();
4307 assert!(
4308 error.contains("does not match validation samples"),
4309 "unexpected fold-alignment error: {error}"
4310 );
4311 }
4312
4313 #[test]
4318 fn separation_concat_merge_bundle_assembles_and_is_scored() {
4319 let plan = separation_concat_merge_plan();
4320 let a_requirement = separation_branch_requirement(
4323 "branch:site__A.model:pls",
4324 &["sample:1", "sample:3"],
4325 &["fold:0", "fold:1"],
4326 );
4327 let b_requirement = separation_branch_requirement(
4328 "branch:site__B.model:pls",
4329 &["sample:2", "sample:4"],
4330 &["fold:0", "fold:1"],
4331 );
4332 let a_cache = build_prediction_cache_record(
4333 &a_requirement,
4334 &separation_branch_blocks("branch:site__A.model:pls", "sample:1", "sample:3", 0.0),
4335 )
4336 .unwrap();
4337 let b_cache = build_prediction_cache_record(
4338 &b_requirement,
4339 &separation_branch_blocks("branch:site__B.model:pls", "sample:2", "sample:4", 1.0),
4340 )
4341 .unwrap();
4342 let a_artifact = refit_artifact(
4343 &plan,
4344 "branch:site__A.model:pls",
4345 vec!["branch:site__A.model:pls.x".to_string()],
4346 Vec::new(),
4347 );
4348 let b_artifact = refit_artifact(
4349 &plan,
4350 "branch:site__B.model:pls",
4351 vec!["branch:site__B.model:pls.x".to_string()],
4352 Vec::new(),
4353 );
4354
4355 let mut bundle = build_execution_bundle_with_prediction_contracts(
4356 BundleId::new("bundle:separation.concat.merge").unwrap(),
4357 &plan,
4358 Some(plan.variants[0].variant_id.clone()),
4359 BTreeMap::new(),
4360 vec![a_artifact, b_artifact],
4361 vec![a_requirement, b_requirement],
4362 vec![a_cache, b_cache],
4363 )
4364 .expect("separation-branch concat-merge bundle must assemble");
4365
4366 bundle
4370 .validate_against_plan(&plan)
4371 .expect("partition-covering branch inputs must validate as a group");
4372 assert_eq!(bundle.prediction_requirements.len(), 2);
4373
4374 let scores = ScoreSet {
4378 schema_version: crate::metrics::SCORE_SET_SCHEMA_VERSION,
4379 plan_id: plan.id.clone(),
4380 selection_metric: Some("rmse".to_string()),
4381 reports: vec![crate::metrics::RegressionMetricReport {
4382 prediction_id: Some("prediction:merge:sites:avg".to_string()),
4383 producer_node: NodeId::new("merge:sites").unwrap(),
4384 producer_port: Some("pred".to_string()),
4385 variant_id: Some(plan.variants[0].variant_id.clone()),
4386 variant_label: None,
4387 partition: PredictionPartition::Validation,
4388 fold_id: Some(FoldId::new("avg").unwrap()),
4389 level: PredictionLevel::Sample,
4390 row_count: 4,
4391 target_width: 1,
4392 target_names: vec!["y".to_string()],
4393 metrics: BTreeMap::from([("rmse".to_string(), 1.5)]),
4394 }],
4395 };
4396 bundle.scores = Some(scores);
4397 bundle
4398 .validate_against_plan(&plan)
4399 .expect("bundle with merge-producer scores must validate");
4400 let cv_best = bundle
4401 .scores
4402 .as_ref()
4403 .unwrap()
4404 .reports
4405 .iter()
4406 .find(|report| {
4407 report.producer_node.as_str() == "merge:sites"
4408 && report.fold_id.as_ref().map(FoldId::as_str) == Some("avg")
4409 })
4410 .expect("merge producer must have a cross-fold (avg) score");
4411 assert_eq!(cv_best.metrics.get("rmse"), Some(&1.5));
4412 }
4413
4414 #[test]
4418 fn separation_concat_merge_rejects_overlapping_partitions() {
4419 let plan = separation_concat_merge_plan();
4420 let a_requirement = separation_branch_requirement(
4423 "branch:site__A.model:pls",
4424 &["sample:1", "sample:3"],
4425 &["fold:0", "fold:1"],
4426 );
4427 let b_requirement = separation_branch_requirement(
4428 "branch:site__B.model:pls",
4429 &["sample:2", "sample:3"],
4430 &["fold:0", "fold:1"],
4431 );
4432 let a_cache = build_prediction_cache_record(
4433 &a_requirement,
4434 &separation_branch_blocks("branch:site__A.model:pls", "sample:1", "sample:3", 0.0),
4435 )
4436 .unwrap();
4437 let b_cache = build_prediction_cache_record(
4438 &b_requirement,
4439 &separation_branch_blocks("branch:site__B.model:pls", "sample:2", "sample:3", 1.0),
4440 )
4441 .unwrap();
4442 let a_artifact = refit_artifact(
4443 &plan,
4444 "branch:site__A.model:pls",
4445 vec!["branch:site__A.model:pls.x".to_string()],
4446 Vec::new(),
4447 );
4448 let b_artifact = refit_artifact(
4449 &plan,
4450 "branch:site__B.model:pls",
4451 vec!["branch:site__B.model:pls.x".to_string()],
4452 Vec::new(),
4453 );
4454
4455 let error = build_execution_bundle_with_prediction_contracts(
4456 BundleId::new("bundle:separation.concat.merge.overlap").unwrap(),
4457 &plan,
4458 Some(plan.variants[0].variant_id.clone()),
4459 BTreeMap::new(),
4460 vec![a_artifact, b_artifact],
4461 vec![a_requirement, b_requirement],
4462 vec![a_cache, b_cache],
4463 )
4464 .unwrap_err()
4465 .to_string();
4466 assert!(
4467 error.contains("overlapping branch predictions"),
4468 "overlap must be rejected, got: {error}"
4469 );
4470 }
4471
4472 #[test]
4476 fn separation_concat_merge_rejects_incomplete_coverage() {
4477 let plan = separation_concat_merge_plan();
4478 let a_requirement =
4481 separation_branch_requirement("branch:site__A.model:pls", &["sample:1"], &["fold:0"]);
4482 let b_requirement = separation_branch_requirement(
4483 "branch:site__B.model:pls",
4484 &["sample:2", "sample:4"],
4485 &["fold:0", "fold:1"],
4486 );
4487 let a_cache = build_prediction_cache_record(
4488 &a_requirement,
4489 &[PredictionBlock {
4490 prediction_id: Some("prediction:a:fold0".to_string()),
4491 producer_node: NodeId::new("branch:site__A.model:pls").unwrap(),
4492 producer_port: None,
4493 partition: PredictionPartition::Validation,
4494 fold_id: Some(FoldId::new("fold:0").unwrap()),
4495 sample_ids: vec![SampleId::new("sample:1").unwrap()],
4496 values: vec![vec![0.1]],
4497 target_names: vec!["y".to_string()],
4498 }],
4499 )
4500 .unwrap();
4501 let b_cache = build_prediction_cache_record(
4502 &b_requirement,
4503 &separation_branch_blocks("branch:site__B.model:pls", "sample:2", "sample:4", 1.0),
4504 )
4505 .unwrap();
4506 let a_artifact = refit_artifact(
4507 &plan,
4508 "branch:site__A.model:pls",
4509 vec!["branch:site__A.model:pls.x".to_string()],
4510 Vec::new(),
4511 );
4512 let b_artifact = refit_artifact(
4513 &plan,
4514 "branch:site__B.model:pls",
4515 vec!["branch:site__B.model:pls.x".to_string()],
4516 Vec::new(),
4517 );
4518
4519 let error = build_execution_bundle_with_prediction_contracts(
4520 BundleId::new("bundle:separation.concat.merge.gap").unwrap(),
4521 &plan,
4522 Some(plan.variants[0].variant_id.clone()),
4523 BTreeMap::new(),
4524 vec![a_artifact, b_artifact],
4525 vec![a_requirement, b_requirement],
4526 vec![a_cache, b_cache],
4527 )
4528 .unwrap_err()
4529 .to_string();
4530 assert!(
4531 error.contains("do not cover"),
4532 "an OOF gap must be rejected, got: {error}"
4533 );
4534 }
4535
4536 #[test]
4541 fn separation_concat_merge_rejects_missing_branch_edge() {
4542 let plan = separation_concat_merge_plan();
4543 let a_requirement = separation_branch_requirement(
4547 "branch:site__A.model:pls",
4548 &["sample:1", "sample:2", "sample:3", "sample:4"],
4549 &["fold:0", "fold:1"],
4550 );
4551 let a_artifact = refit_artifact(
4552 &plan,
4553 "branch:site__A.model:pls",
4554 vec!["branch:site__A.model:pls.x".to_string()],
4555 Vec::new(),
4556 );
4557 let b_artifact = refit_artifact(
4558 &plan,
4559 "branch:site__B.model:pls",
4560 vec!["branch:site__B.model:pls.x".to_string()],
4561 Vec::new(),
4562 );
4563
4564 let error = build_execution_bundle_with_prediction_contracts(
4565 BundleId::new("bundle:separation.concat.merge.missing.branch").unwrap(),
4566 &plan,
4567 Some(plan.variants[0].variant_id.clone()),
4568 BTreeMap::new(),
4569 vec![a_artifact, b_artifact],
4570 vec![a_requirement],
4571 Vec::new(),
4572 )
4573 .unwrap_err()
4574 .to_string();
4575 assert!(
4576 error.contains("do not match the plan's incoming OOF edges"),
4577 "a missing branch edge must be rejected, got: {error}"
4578 );
4579 }
4580
4581 #[test]
4586 fn separation_concat_merge_rejects_partial_cache_coverage() {
4587 let plan = separation_concat_merge_plan();
4588 let a_requirement = separation_branch_requirement(
4589 "branch:site__A.model:pls",
4590 &["sample:1", "sample:3"],
4591 &["fold:0", "fold:1"],
4592 );
4593 let b_requirement = separation_branch_requirement(
4594 "branch:site__B.model:pls",
4595 &["sample:2", "sample:4"],
4596 &["fold:0", "fold:1"],
4597 );
4598 let a_cache = build_prediction_cache_record(
4600 &a_requirement,
4601 &separation_branch_blocks("branch:site__A.model:pls", "sample:1", "sample:3", 0.0),
4602 )
4603 .unwrap();
4604 let a_artifact = refit_artifact(
4605 &plan,
4606 "branch:site__A.model:pls",
4607 vec!["branch:site__A.model:pls.x".to_string()],
4608 Vec::new(),
4609 );
4610 let b_artifact = refit_artifact(
4611 &plan,
4612 "branch:site__B.model:pls",
4613 vec!["branch:site__B.model:pls.x".to_string()],
4614 Vec::new(),
4615 );
4616
4617 let error = build_execution_bundle_with_prediction_contracts(
4618 BundleId::new("bundle:separation.concat.merge.partial.cache").unwrap(),
4619 &plan,
4620 Some(plan.variants[0].variant_id.clone()),
4621 BTreeMap::new(),
4622 vec![a_artifact, b_artifact],
4623 vec![a_requirement, b_requirement],
4624 vec![a_cache],
4625 )
4626 .unwrap_err()
4627 .to_string();
4628 assert!(
4629 error.contains("partial prediction-cache coverage"),
4630 "a partial-cache concat group must be rejected, got: {error}"
4631 );
4632 }
4633
4634 #[test]
4635 fn prediction_requirements_are_typed_and_validate_against_oof_edges() {
4636 let plan = branch_merge_plan();
4637 let meta_plan = plan
4638 .node_plans
4639 .get(&NodeId::new("merge:stack.pred_plus_original.meta:ridge").unwrap())
4640 .unwrap();
4641 let producer_node = NodeId::new("branch:b0.model:ridge").unwrap();
4642 let fold0 = FoldId::new("fold:0").unwrap();
4643 let fold1 = FoldId::new("fold:1").unwrap();
4644 let samples = [
4645 SampleId::new("sample:1").unwrap(),
4646 SampleId::new("sample:2").unwrap(),
4647 SampleId::new("sample:3").unwrap(),
4648 SampleId::new("sample:4").unwrap(),
4649 ];
4650 let requirement = BundlePredictionRequirement {
4651 producer_node: producer_node.clone(),
4652 source_port: "oof".to_string(),
4653 consumer_node: meta_plan.node_id.clone(),
4654 target_port: "b0_oof".to_string(),
4655 partition: PredictionPartition::Validation,
4656 prediction_level: PredictionLevel::Sample,
4657 fold_ids: vec![fold0.clone(), fold1.clone()],
4658 unit_ids: Vec::new(),
4659 sample_ids: samples.to_vec(),
4660 prediction_width: 1,
4661 target_names: vec!["y".to_string()],
4662 };
4663 let prediction_blocks = vec![
4664 PredictionBlock {
4665 prediction_id: Some("prediction:branch:b0.fold0".to_string()),
4666 producer_node: producer_node.clone(),
4667 producer_port: Some("oof".to_string()),
4668 partition: PredictionPartition::Validation,
4669 fold_id: Some(fold0),
4670 sample_ids: samples[0..2].to_vec(),
4671 values: vec![vec![0.1], vec![0.2]],
4672 target_names: vec!["y".to_string()],
4673 },
4674 PredictionBlock {
4675 prediction_id: Some("prediction:branch:b0.fold1".to_string()),
4676 producer_node: producer_node.clone(),
4677 producer_port: Some("oof".to_string()),
4678 partition: PredictionPartition::Validation,
4679 fold_id: Some(fold1),
4680 sample_ids: samples[2..4].to_vec(),
4681 values: vec![vec![0.3], vec![0.4]],
4682 target_names: vec!["y".to_string()],
4683 },
4684 ];
4685 let cache = build_prediction_cache_record(&requirement, &prediction_blocks).unwrap();
4686 let payload = build_prediction_cache_payload(&requirement, &prediction_blocks).unwrap();
4687 assert_eq!(cache.prediction_level, PredictionLevel::Sample);
4688 assert_eq!(payload.prediction_level, PredictionLevel::Sample);
4689 assert!(cache
4690 .blocks
4691 .iter()
4692 .all(|block| block.prediction_level == PredictionLevel::Sample));
4693 validate_prediction_cache_payload_matches_record(&payload, &cache).unwrap();
4694 let cache_namespace_fingerprints = vec!["a".repeat(64), "b".repeat(64)];
4695 let mut d10_cache = cache.clone();
4696 d10_cache.cache_namespace_fingerprints = cache_namespace_fingerprints.clone();
4697 d10_cache.validate().unwrap();
4698 let mut d10_payload = payload.clone();
4699 d10_payload.cache_namespace_fingerprints = cache_namespace_fingerprints;
4700 d10_payload.validate().unwrap();
4701 validate_prediction_cache_payload_matches_record(&d10_payload, &d10_cache).unwrap();
4702 for forbidden_partition in [
4703 PredictionPartition::Train,
4704 PredictionPartition::Test,
4705 PredictionPartition::Final,
4706 ] {
4707 let mut non_oof_requirement = requirement.clone();
4708 non_oof_requirement.partition = forbidden_partition.clone();
4709 let requirement_error = non_oof_requirement.validate().unwrap_err().to_string();
4710 assert!(
4711 requirement_error.contains("must use validation OOF predictions"),
4712 "D10 cache namespace must not broaden bundle prediction requirements to {forbidden_partition:?}: {requirement_error}"
4713 );
4714
4715 let mut non_oof_cache = d10_cache.clone();
4716 non_oof_cache.partition = forbidden_partition.clone();
4717 let cache_error = non_oof_cache.validate().unwrap_err().to_string();
4718 assert!(
4719 cache_error.contains("must cache validation OOF predictions"),
4720 "D10 cache namespace must not allow non-OOF cache records for {forbidden_partition:?}: {cache_error}"
4721 );
4722
4723 let mut non_oof_payload = d10_payload.clone();
4724 non_oof_payload.partition = forbidden_partition;
4725 let payload_error = non_oof_payload.validate().unwrap_err().to_string();
4726 assert!(
4727 payload_error.contains("must cache validation OOF predictions"),
4728 "D10 cache namespace must not allow non-OOF cache payloads: {payload_error}"
4729 );
4730 }
4731 let mut short_namespace_cache = d10_cache.clone();
4732 short_namespace_cache.cache_namespace_fingerprints.pop();
4733 assert!(short_namespace_cache
4734 .validate()
4735 .unwrap_err()
4736 .to_string()
4737 .contains("namespace fingerprint count"));
4738 let mut short_namespace_payload = d10_payload;
4739 short_namespace_payload.cache_namespace_fingerprints.pop();
4740 assert!(short_namespace_payload
4741 .validate()
4742 .unwrap_err()
4743 .to_string()
4744 .contains("namespace fingerprint count"));
4745 let mut wrong_level_requirement = requirement.clone();
4746 wrong_level_requirement.prediction_level = PredictionLevel::Target;
4747 assert!(wrong_level_requirement.validate().is_err());
4748 let mut wrong_level_cache = cache.clone();
4749 wrong_level_cache.prediction_level = PredictionLevel::Target;
4750 assert!(wrong_level_cache.validate().is_err());
4751 let mut wrong_level_payload = payload.clone();
4752 wrong_level_payload.prediction_level = PredictionLevel::Target;
4753 assert!(wrong_level_payload.validate().is_err());
4754 let prediction_key = requirement.key();
4755 let artifact = RefitArtifactRecord {
4756 node_id: meta_plan.node_id.clone(),
4757 controller_id: meta_plan.controller_id.clone(),
4758 artifact: ArtifactRef {
4759 id: ArtifactId::new("artifact:merge:stack.pred_plus_original.meta:ridge:refit")
4760 .unwrap(),
4761 kind: "mock_model".to_string(),
4762 controller_id: meta_plan.controller_id.clone(),
4763 backend: None,
4764 uri: None,
4765 content_fingerprint: None,
4766 size_bytes: Some(128),
4767 plugin: None,
4768 plugin_version: None,
4769 },
4770 params_fingerprint: meta_plan.params_fingerprint.clone(),
4771 training_loss_fingerprint: meta_plan.training_loss_fingerprint(Phase::Refit).unwrap(),
4772 data_requirement_keys: vec![
4773 "merge:stack.pred_plus_original.meta:ridge.x_original".to_string()
4774 ],
4775 prediction_requirement_keys: vec![prediction_key],
4776 };
4777
4778 assert!(build_execution_bundle_with_prediction_contracts(
4779 BundleId::new("bundle:d10.cache.without.selected.variant").unwrap(),
4780 &plan,
4781 None,
4782 BTreeMap::new(),
4783 vec![artifact.clone()],
4784 vec![requirement.clone()],
4785 vec![d10_cache],
4786 )
4787 .unwrap_err()
4788 .to_string()
4789 .contains("requires selected_variant_id"));
4790
4791 assert!(build_execution_bundle(
4792 BundleId::new("bundle:missing.prediction.requirement").unwrap(),
4793 &plan,
4794 Some(plan.variants[0].variant_id.clone()),
4795 BTreeMap::new(),
4796 vec![artifact.clone()],
4797 )
4798 .is_err());
4799
4800 assert!(build_execution_bundle_with_prediction_requirements(
4801 BundleId::new("bundle:typed.prediction.requirement.without.cache").unwrap(),
4802 &plan,
4803 Some(plan.variants[0].variant_id.clone()),
4804 BTreeMap::new(),
4805 vec![artifact.clone()],
4806 vec![requirement.clone()],
4807 )
4808 .is_err());
4809
4810 let bundle = build_execution_bundle_with_prediction_contracts(
4811 BundleId::new("bundle:typed.prediction.requirement").unwrap(),
4812 &plan,
4813 Some(plan.variants[0].variant_id.clone()),
4814 BTreeMap::new(),
4815 vec![artifact],
4816 vec![requirement],
4817 vec![cache],
4818 )
4819 .unwrap();
4820 bundle.validate_against_plan(&plan).unwrap();
4821 assert_eq!(bundle.prediction_requirements.len(), 1);
4822 assert_eq!(bundle.prediction_caches.len(), 1);
4823 assert_eq!(
4824 bundle.refit_artifacts[0].prediction_requirement_keys,
4825 vec!["branch:b0.model:ridge.oof->merge:stack.pred_plus_original.meta:ridge.b0_oof"]
4826 );
4827 let payload_set = BundlePredictionCachePayloadSet {
4828 bundle_id: bundle.bundle_id.clone(),
4829 schema_version: PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
4830 caches: vec![payload],
4831 };
4832 payload_set.validate_against_bundle(&bundle).unwrap();
4833 let refit_replay_request = ReplayPhaseRequest {
4834 bundle_id: bundle.bundle_id.clone(),
4835 phase: Phase::Refit,
4836 data_envelope_keys: bundle
4837 .data_requirements
4838 .iter()
4839 .map(BundleDataRequirement::key)
4840 .collect(),
4841 };
4842 refit_replay_request
4843 .validate_for_bundle_with_prediction_cache_payloads(&bundle, Some(&payload_set))
4844 .unwrap();
4845 let mut tampered_payload_set = payload_set.clone();
4846 tampered_payload_set.caches[0].blocks[0].values[0][0] = 99.0;
4847 assert!(tampered_payload_set
4848 .validate_against_bundle(&bundle)
4849 .is_err());
4850 let mut missing_payload_set = payload_set.clone();
4851 missing_payload_set.caches.clear();
4852 assert!(missing_payload_set
4853 .validate_against_bundle(&bundle)
4854 .is_err());
4855 assert!(refit_replay_request.validate_for_bundle(&bundle).is_err());
4856
4857 let mut wrong_data_owner = bundle.clone();
4858 wrong_data_owner.refit_artifacts[0].data_requirement_keys =
4859 vec!["branch:b0.model:ridge.x".to_string()];
4860 assert!(wrong_data_owner.validate().is_err());
4861
4862 let mut wrong_prediction_consumer = bundle;
4863 wrong_prediction_consumer.refit_artifacts[0].node_id =
4864 NodeId::new("branch:b0.model:ridge").unwrap();
4865 wrong_prediction_consumer.refit_artifacts[0]
4866 .data_requirement_keys
4867 .clear();
4868 assert!(wrong_prediction_consumer.validate().is_err());
4869 }
4870
4871 #[test]
4872 fn aggregated_prediction_cache_contracts_preserve_unit_ids() {
4873 let plan = branch_merge_plan();
4874 let producer_node = NodeId::new("branch:b0.model:ridge").unwrap();
4875 let consumer_node = NodeId::new("merge:stack.pred_plus_original.meta:ridge").unwrap();
4876 let fold0 = FoldId::new("fold:0").unwrap();
4877 let fold1 = FoldId::new("fold:1").unwrap();
4878 let target_a = PredictionUnitId::Target(TargetId::new("target:a").unwrap());
4879 let target_b = PredictionUnitId::Target(TargetId::new("target:b").unwrap());
4880 let requirement = BundlePredictionRequirement {
4881 producer_node: producer_node.clone(),
4882 source_port: "oof".to_string(),
4883 consumer_node: consumer_node.clone(),
4884 target_port: "b0_oof".to_string(),
4885 partition: PredictionPartition::Validation,
4886 prediction_level: PredictionLevel::Target,
4887 fold_ids: vec![fold0.clone(), fold1.clone()],
4888 unit_ids: vec![target_a.clone(), target_b.clone()],
4889 sample_ids: Vec::new(),
4890 prediction_width: 1,
4891 target_names: vec!["y".to_string()],
4892 };
4893 let aggregated_blocks = vec![
4894 AggregatedPredictionBlock {
4895 prediction_id: Some("prediction:branch:b0.target.fold0".to_string()),
4896 producer_node: producer_node.clone(),
4897 producer_port: Some("pred".to_string()),
4898 partition: PredictionPartition::Validation,
4899 fold_id: Some(fold0),
4900 level: PredictionLevel::Target,
4901 unit_ids: vec![target_a],
4902 values: vec![vec![0.15]],
4903 target_names: vec!["y".to_string()],
4904 },
4905 AggregatedPredictionBlock {
4906 prediction_id: Some("prediction:branch:b0.target.fold1".to_string()),
4907 producer_node,
4908 producer_port: Some("pred".to_string()),
4909 partition: PredictionPartition::Validation,
4910 fold_id: Some(fold1),
4911 level: PredictionLevel::Target,
4912 unit_ids: vec![target_b],
4913 values: vec![vec![0.35]],
4914 target_names: vec!["y".to_string()],
4915 },
4916 ];
4917
4918 let cache =
4919 build_aggregated_prediction_cache_record(&requirement, &aggregated_blocks).unwrap();
4920 let payload =
4921 build_aggregated_prediction_cache_payload(&requirement, &aggregated_blocks).unwrap();
4922 assert_eq!(cache.prediction_level, PredictionLevel::Target);
4923 assert_eq!(cache.unit_ids, requirement.unit_ids);
4924 assert!(cache.sample_ids.is_empty());
4925 assert!(payload.blocks.is_empty());
4926 assert_eq!(payload.aggregated_blocks.len(), 2);
4927 validate_prediction_cache_payload_matches_record(&payload, &cache).unwrap();
4928
4929 let artifact = refit_artifact(
4930 &plan,
4931 "merge:stack.pred_plus_original.meta:ridge",
4932 vec!["merge:stack.pred_plus_original.meta:ridge.x_original".to_string()],
4933 vec![requirement.key()],
4934 );
4935 let bundle = build_execution_bundle_with_prediction_contracts(
4936 BundleId::new("bundle:target.prediction.requirement").unwrap(),
4937 &plan,
4938 Some(plan.variants[0].variant_id.clone()),
4939 BTreeMap::new(),
4940 vec![artifact],
4941 vec![requirement],
4942 vec![cache],
4943 )
4944 .unwrap();
4945 bundle.validate_against_plan(&plan).unwrap();
4946
4947 let mut tampered_payload = payload;
4948 tampered_payload.aggregated_blocks[0].unit_ids =
4949 vec![PredictionUnitId::Target(TargetId::new("target:z").unwrap())];
4950 assert!(validate_prediction_cache_payload_matches_record(
4951 &tampered_payload,
4952 &bundle.prediction_caches[0]
4953 )
4954 .is_err());
4955 }
4956
4957 #[test]
4958 fn replay_envelopes_must_match_bundle_requirements() {
4959 let plan = plan();
4960 let bundle = build_execution_bundle(
4961 BundleId::new("bundle:demo").unwrap(),
4962 &plan,
4963 None,
4964 BTreeMap::new(),
4965 Vec::new(),
4966 )
4967 .unwrap();
4968 let envelope: ExternalDataPlanEnvelope = serde_json::from_str(include_str!(
4969 "../tests/fixtures/package/data/coordinator_data_plan_envelope_sample12.json"
4970 ))
4971 .unwrap();
4972
4973 bundle
4974 .validate_replay_envelopes(&BTreeMap::from([(
4975 "model:base.x".to_string(),
4976 envelope.clone(),
4977 )]))
4978 .unwrap();
4979
4980 let mut mismatched = envelope;
4981 mismatched.schema_fingerprint = "0".repeat(64);
4982 assert!(bundle
4983 .validate_replay_envelopes(&BTreeMap::from([("model:base.x".to_string(), mismatched,)]))
4984 .is_err());
4985 }
4986
4987 #[test]
4988 fn rejects_unsupported_bundle_schema_version() {
4989 let mut bundle = build_execution_bundle(
4990 BundleId::new("bundle:demo").unwrap(),
4991 &plan(),
4992 None,
4993 BTreeMap::new(),
4994 Vec::new(),
4995 )
4996 .unwrap();
4997 bundle.schema_version = EXECUTION_BUNDLE_SCHEMA_VERSION + 1;
4998
4999 assert!(bundle.validate().is_err());
5000
5001 bundle.schema_version = 0;
5002 assert!(bundle.validate().is_err());
5003 }
5004
5005 #[test]
5006 fn rejects_bundle_with_scores_plan_id_mismatch() {
5007 let plan = plan();
5008 let mut bundle = build_execution_bundle(
5009 BundleId::new("bundle:demo").unwrap(),
5010 &plan,
5011 None,
5012 BTreeMap::new(),
5013 Vec::new(),
5014 )
5015 .unwrap();
5016 bundle.scores = Some(ScoreSet {
5017 schema_version: crate::metrics::SCORE_SET_SCHEMA_VERSION,
5018 plan_id: bundle.plan_id.clone(),
5019 selection_metric: Some("rmse".to_string()),
5020 reports: vec![crate::metrics::RegressionMetricReport {
5021 prediction_id: None,
5022 producer_node: NodeId::new("model:compat.0").unwrap(),
5023 producer_port: Some("pred".to_string()),
5024 variant_id: None,
5025 variant_label: None,
5026 partition: PredictionPartition::Test,
5027 fold_id: Some(FoldId::new("final").unwrap()),
5028 level: PredictionLevel::Sample,
5029 row_count: 4,
5030 target_width: 1,
5031 target_names: vec!["y".to_string()],
5032 metrics: BTreeMap::from([("rmse".to_string(), 1.0)]),
5033 }],
5034 });
5035 bundle.validate().unwrap();
5037 bundle.scores.as_mut().unwrap().plan_id = "plan:wrong".to_string();
5039 let err = bundle.validate().unwrap_err().to_string();
5040 assert!(
5041 err.contains("does not match its embedded scores plan_id"),
5042 "{err}"
5043 );
5044 }
5045
5046 #[test]
5047 fn schema_migration_policy_is_explicit_and_refuses_implicit_migrations() {
5048 let bundle_policy = execution_bundle_schema_migration_policy();
5049 assert_eq!(
5050 bundle_policy.current_version,
5051 EXECUTION_BUNDLE_SCHEMA_VERSION
5052 );
5053 assert_eq!(
5054 bundle_policy.min_readable_version,
5055 MIN_READABLE_EXECUTION_BUNDLE_SCHEMA_VERSION
5056 );
5057 assert_eq!(
5058 bundle_policy.min_writable_version,
5059 MIN_WRITABLE_EXECUTION_BUNDLE_SCHEMA_VERSION
5060 );
5061 assert_eq!(
5062 bundle_policy
5063 .automatic_migrations
5064 .get(&LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION),
5065 Some(&EXECUTION_BUNDLE_SCHEMA_VERSION)
5066 );
5067 bundle_policy
5068 .validate_read_version(LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION, "bundle `legacy`")
5069 .unwrap();
5070 bundle_policy
5071 .validate_read_version(EXECUTION_BUNDLE_SCHEMA_VERSION, "bundle `current`")
5072 .unwrap();
5073 assert!(bundle_policy
5074 .validate_read_version(EXECUTION_BUNDLE_SCHEMA_VERSION + 1, "bundle `future`")
5075 .is_err());
5076 assert!(bundle_policy
5077 .validate_read_version(0, "bundle `zero`")
5078 .is_err());
5079
5080 let mut future_policy = SchemaMigrationPolicy {
5081 artifact: "execution_bundle".to_string(),
5082 current_version: 2,
5083 min_readable_version: 1,
5084 min_writable_version: 2,
5085 automatic_migrations: BTreeMap::new(),
5086 };
5087 assert!(future_policy
5088 .validate_read_version(1, "bundle `old-without-migration`")
5089 .is_err());
5090 future_policy.automatic_migrations.insert(1, 2);
5091 future_policy
5092 .validate_read_version(1, "bundle `old-with-migration`")
5093 .unwrap();
5094 }
5095
5096 #[test]
5097 fn prediction_cache_payload_schema_policy_rejects_unsupported_versions() {
5098 let policy = prediction_cache_payload_schema_migration_policy();
5099 assert_eq!(
5100 policy.current_version,
5101 PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION
5102 );
5103 assert_eq!(
5104 policy
5105 .automatic_migrations
5106 .get(&LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION),
5107 Some(&PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION)
5108 );
5109 policy
5110 .validate_read_version(
5111 LEGACY_PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
5112 "payload `legacy`",
5113 )
5114 .unwrap();
5115
5116 let mut payload_set = BundlePredictionCachePayloadSet {
5117 bundle_id: BundleId::new("bundle:payload.schema").unwrap(),
5118 schema_version: PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION,
5119 caches: Vec::new(),
5120 };
5121 payload_set.validate().unwrap();
5122
5123 payload_set.schema_version = PREDICTION_CACHE_PAYLOAD_SCHEMA_VERSION + 1;
5124 assert!(payload_set.validate().is_err());
5125
5126 payload_set.schema_version = 0;
5127 assert!(payload_set.validate().is_err());
5128 }
5129
5130 #[test]
5131 fn prediction_cache_payload_format_enforces_port_family() {
5132 fn payload_for_block(
5133 format: &str,
5134 producer_port: Option<String>,
5135 ) -> BundlePredictionCachePayload {
5136 let block = PredictionBlock {
5137 prediction_id: Some("prediction:model:base.fold0".to_string()),
5138 producer_node: NodeId::new("model:base").unwrap(),
5139 producer_port,
5140 partition: PredictionPartition::Validation,
5141 fold_id: Some(FoldId::new("fold:0").unwrap()),
5142 sample_ids: vec![SampleId::new("sample:1").unwrap()],
5143 values: vec![vec![1.0]],
5144 target_names: vec!["y".to_string()],
5145 };
5146 let blocks = vec![block];
5147 BundlePredictionCachePayload {
5148 requirement_key: "model:base.oof->model:meta.pred".to_string(),
5149 cache_id: "prediction-cache:model:base.oof->model:meta.pred".to_string(),
5150 cache_namespace_fingerprints: Vec::new(),
5151 format: format.to_string(),
5152 partition: PredictionPartition::Validation,
5153 prediction_level: PredictionLevel::Sample,
5154 block_count: blocks.len(),
5155 row_count: blocks.iter().map(|block| block.sample_ids.len()).sum(),
5156 content_fingerprint: stable_json_fingerprint(&blocks).unwrap(),
5157 blocks,
5158 aggregated_blocks: Vec::new(),
5159 }
5160 }
5161
5162 payload_for_block(LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT, None)
5163 .validate()
5164 .unwrap();
5165 payload_for_block(BUNDLE_PREDICTION_CACHE_FORMAT, Some("oof".to_string()))
5166 .validate()
5167 .unwrap();
5168
5169 let v1_with_port = payload_for_block(
5170 LEGACY_BUNDLE_PREDICTION_CACHE_FORMAT,
5171 Some("oof".to_string()),
5172 )
5173 .validate()
5174 .unwrap_err()
5175 .to_string();
5176 assert!(
5177 v1_with_port.contains("V1") && v1_with_port.contains("producer_port"),
5178 "unexpected V1 port-family error: {v1_with_port}"
5179 );
5180
5181 let v2_without_port = payload_for_block(BUNDLE_PREDICTION_CACHE_FORMAT, None)
5182 .validate()
5183 .unwrap_err()
5184 .to_string();
5185 assert!(
5186 v2_without_port.contains("V2") && v2_without_port.contains("requires producer_port"),
5187 "unexpected V2 port-family error: {v2_without_port}"
5188 );
5189 }
5190
5191 #[test]
5192 fn replay_request_requires_predict_explain_or_refit_phase() {
5193 let bundle = build_execution_bundle(
5194 BundleId::new("bundle:demo").unwrap(),
5195 &plan(),
5196 None,
5197 BTreeMap::new(),
5198 Vec::new(),
5199 )
5200 .unwrap();
5201
5202 ReplayPhaseRequest {
5203 bundle_id: bundle.bundle_id.clone(),
5204 phase: Phase::Predict,
5205 data_envelope_keys: vec!["model:base.x".to_string()],
5206 }
5207 .validate_for_bundle(&bundle)
5208 .unwrap();
5209 ReplayPhaseRequest {
5210 bundle_id: bundle.bundle_id.clone(),
5211 phase: Phase::Refit,
5212 data_envelope_keys: vec!["model:base.x".to_string()],
5213 }
5214 .validate_for_bundle(&bundle)
5215 .unwrap();
5216 assert!(ReplayPhaseRequest {
5217 bundle_id: bundle.bundle_id.clone(),
5218 phase: Phase::FitCv,
5219 data_envelope_keys: vec!["model:base.x".to_string()],
5220 }
5221 .validate_for_bundle(&bundle)
5222 .is_err());
5223 assert!(ReplayPhaseRequest {
5224 bundle_id: bundle.bundle_id.clone(),
5225 phase: Phase::Predict,
5226 data_envelope_keys: vec!["model:base.x".to_string(), "model:base.x".to_string()],
5227 }
5228 .validate_for_bundle(&bundle)
5229 .is_err());
5230 assert!(ReplayPhaseRequest {
5231 bundle_id: bundle.bundle_id.clone(),
5232 phase: Phase::Predict,
5233 data_envelope_keys: vec!["model:base.y".to_string()],
5234 }
5235 .validate_for_bundle(&bundle)
5236 .is_err());
5237 }
5238
5239 #[test]
5240 fn prediction_level_wire_absent_parses_as_sample_but_serialization_stays_explicit() {
5241 let mut requirement = BundlePredictionRequirement {
5242 producer_node: NodeId::new("model:base").unwrap(),
5243 source_port: "oof".to_string(),
5244 consumer_node: NodeId::new("model:meta").unwrap(),
5245 target_port: "x".to_string(),
5246 partition: PredictionPartition::Validation,
5247 prediction_level: PredictionLevel::Sample,
5248 fold_ids: vec![FoldId::new("fold:0").unwrap()],
5249 unit_ids: Vec::new(),
5250 sample_ids: vec![SampleId::new("sample:1").unwrap()],
5251 prediction_width: 1,
5252 target_names: vec!["y".to_string()],
5253 };
5254
5255 let mut sample = serde_json::to_value(&requirement).unwrap();
5256 assert_eq!(sample["prediction_level"], serde_json::json!("sample"));
5257
5258 sample.as_object_mut().unwrap().remove("prediction_level");
5259 let parsed: BundlePredictionRequirement = serde_json::from_value(sample).unwrap();
5260 assert_eq!(parsed.prediction_level, PredictionLevel::Sample);
5261
5262 requirement.prediction_level = PredictionLevel::Group;
5263 let group = serde_json::to_value(&requirement).unwrap();
5264 assert_eq!(group["prediction_level"], serde_json::json!("group"));
5265 }
5266
5267 #[test]
5268 fn methods_hpo_resume_state_round_trips_and_rejects_orphan_or_score_drift() {
5269 let state = methods_hpo_resume_state();
5270 state.validate().unwrap();
5271 let json = serde_json::to_string(&state).unwrap();
5272 let parsed: MethodsHpoResumeState = serde_json::from_str(&json).unwrap();
5273 assert_eq!(parsed, state);
5274
5275 let mut orphan = state.clone();
5276 orphan.completed_reports.clear();
5277 assert!(orphan
5278 .validate()
5279 .unwrap_err()
5280 .to_string()
5281 .contains("exactly cover"));
5282
5283 let mut score_drift = state;
5284 score_drift.completed_reports[0].score += 1.1e-12;
5285 assert!(score_drift
5286 .validate()
5287 .unwrap_err()
5288 .to_string()
5289 .contains("score/intermediate evidence is inconsistent"));
5290
5291 let mut failed = serde_json::to_value(methods_hpo_resume_state()).unwrap();
5292 failed["completed_reports"][0]["terminal_state"] = serde_json::json!("failed");
5293 assert!(serde_json::from_value::<MethodsHpoResumeState>(failed).is_err());
5294
5295 let mut uppercase_binding = methods_hpo_resume_state();
5296 uppercase_binding.checkpoint.binding.optimizer_fingerprint = "A".repeat(64);
5297 assert!(uppercase_binding
5298 .validate()
5299 .unwrap_err()
5300 .to_string()
5301 .contains("lowercase SHA-256"));
5302 }
5303
5304 #[test]
5305 fn methods_hpo_resume_state_rejects_unknown_and_duplicate_json_members() {
5306 let plan = plan();
5307 let bundle = build_execution_bundle(
5308 BundleId::new("bundle:hpo.strict").unwrap(),
5309 &plan,
5310 None,
5311 BTreeMap::new(),
5312 Vec::new(),
5313 )
5314 .unwrap();
5315 let mut raw = serde_json::to_value(bundle).unwrap();
5316 raw.as_object_mut()
5317 .unwrap()
5318 .insert("unknown_hpo_member".to_string(), serde_json::json!(true));
5319 assert!(ExecutionBundle::from_json(&serde_json::to_string(&raw).unwrap()).is_err());
5320
5321 let encoded = serde_json::to_string(&raw).unwrap();
5322 let duplicated = encoded.replacen(
5323 "\"bundle_id\":\"bundle:hpo.strict\"",
5324 "\"bundle_id\":\"bundle:hpo.strict\",\"bundle_id\":\"bundle:hpo.strict\"",
5325 1,
5326 );
5327 let error = ExecutionBundle::from_json(&duplicated)
5328 .unwrap_err()
5329 .to_string();
5330 assert!(error.contains("duplicate JSON object key"), "{error}");
5331
5332 let mut state = serde_json::to_value(methods_hpo_resume_state()).unwrap();
5333 state.as_object_mut().unwrap().insert(
5334 "unexpected_candidate_side_channel".to_string(),
5335 serde_json::json!(true),
5336 );
5337 assert!(serde_json::from_value::<MethodsHpoResumeState>(state).is_err());
5338 }
5339
5340 #[test]
5341 fn execution_bundle_v1_json_rejects_null_conformal_member() {
5342 let plan = plan();
5343 let mut bundle = build_execution_bundle(
5344 BundleId::new("bundle:v1.strict-conformal").unwrap(),
5345 &plan,
5346 None,
5347 BTreeMap::new(),
5348 Vec::new(),
5349 )
5350 .unwrap();
5351 bundle.schema_version = LEGACY_EXECUTION_BUNDLE_SCHEMA_VERSION;
5352 let mut raw = serde_json::to_value(bundle).unwrap();
5353 raw.as_object_mut()
5354 .unwrap()
5355 .insert("conformal_calibration".to_string(), serde_json::Value::Null);
5356 let error = ExecutionBundle::from_json(&serde_json::to_string(&raw).unwrap())
5357 .unwrap_err()
5358 .to_string();
5359 assert!(error.contains("including null"), "{error}");
5360 }
5361}