1use super::*;
3
4#[derive(Clone, Debug, Default)]
5pub struct SequentialScheduler;
6
7#[derive(Clone, Debug)]
8pub struct ParallelScheduler {
9 max_workers: usize,
10}
11
12impl ParallelScheduler {
13 pub fn new(max_workers: usize) -> Result<Self> {
14 if max_workers == 0 {
15 return Err(DagMlError::RuntimeValidation(
16 "parallel scheduler max_workers must be at least 1".to_string(),
17 ));
18 }
19 Ok(Self { max_workers })
20 }
21
22 pub fn max_workers(&self) -> usize {
23 self.max_workers
24 }
25}
26
27#[derive(Clone, Debug)]
28pub(crate) struct PhaseScope {
29 pub(crate) phase: Phase,
30 pub(crate) variant_id: Option<VariantId>,
31 pub(crate) variant: Option<VariantExecutionSpec>,
32 pub(crate) fold_id: Option<FoldId>,
33 pub(crate) seed_root: Option<u64>,
34}
35
36#[derive(Clone, Debug)]
37pub(crate) struct ReplayPredictionCacheContract {
38 pub(crate) requirement: BundlePredictionRequirement,
39 pub(crate) cache: BundlePredictionCacheRecord,
40}
41
42pub(crate) struct MaterializedReplayArtifacts {
43 pub(crate) handles: BTreeMap<NodeId, BTreeMap<String, HandleRef>>,
44 pub(crate) inputs: BTreeMap<NodeId, BTreeMap<String, ArtifactInputSpec>>,
45}
46
47fn prediction_output_ports_for_node(plan: &ExecutionPlan, node_id: &NodeId) -> Result<Vec<String>> {
48 let node = plan
49 .graph_plan
50 .graph
51 .nodes
52 .iter()
53 .find(|node| node.id == *node_id)
54 .ok_or_else(|| {
55 DagMlError::RuntimeValidation(format!(
56 "node `{node_id}` is absent from the execution graph"
57 ))
58 })?;
59 let mut ports = node
60 .ports
61 .outputs
62 .iter()
63 .filter(|port| port.kind == PortKind::Prediction)
64 .map(|port| port.name.clone())
65 .collect::<Vec<_>>();
66 ports.sort();
67 Ok(ports)
68}
69
70fn normalize_prediction_result_port(
71 node_id: &NodeId,
72 block_kind: &str,
73 producer_port: &mut Option<String>,
74 prediction_ports: &[String],
75) -> Result<()> {
76 if let Some(port) = producer_port.as_ref() {
77 if port.trim().is_empty() {
78 return Err(DagMlError::RuntimeValidation(format!(
79 "node `{node_id}` emitted {block_kind} with blank producer_port"
80 )));
81 }
82 if !prediction_ports.iter().any(|candidate| candidate == port) {
83 return Err(DagMlError::RuntimeValidation(format!(
84 "node `{node_id}` emitted {block_kind} for undeclared or non-prediction output port `{port}`; declared prediction ports are {:?}",
85 prediction_ports
86 )));
87 }
88 return Ok(());
89 }
90 match prediction_ports {
91 [only] => {
92 *producer_port = Some(only.clone());
93 Ok(())
94 }
95 [] => Err(DagMlError::RuntimeValidation(format!(
96 "node `{node_id}` emitted {block_kind} without producer_port but declares no prediction output port"
97 ))),
98 _ => Err(DagMlError::RuntimeValidation(format!(
99 "node `{node_id}` emitted {block_kind} without producer_port but declares {} prediction output ports {:?}; multi-output controllers must emit producer_port explicitly",
100 prediction_ports.len(),
101 prediction_ports
102 ))),
103 }
104}
105
106pub(crate) fn normalize_result_prediction_ports(
107 plan: &ExecutionPlan,
108 task: &NodeTask,
109 result: &mut NodeResult,
110) -> Result<()> {
111 if result.predictions.is_empty()
112 && result.observation_predictions.is_empty()
113 && result.aggregated_predictions.is_empty()
114 && result.explanations.is_empty()
115 {
116 return Ok(());
117 }
118 let prediction_ports = prediction_output_ports_for_node(plan, &task.node_plan.node_id)?;
119 for block in &mut result.predictions {
120 normalize_prediction_result_port(
121 &task.node_plan.node_id,
122 "prediction block",
123 &mut block.producer_port,
124 &prediction_ports,
125 )?;
126 }
127 for block in &mut result.observation_predictions {
128 normalize_prediction_result_port(
129 &task.node_plan.node_id,
130 "observation prediction block",
131 &mut block.producer_port,
132 &prediction_ports,
133 )?;
134 }
135 for block in &mut result.aggregated_predictions {
136 normalize_prediction_result_port(
137 &task.node_plan.node_id,
138 "aggregated prediction block",
139 &mut block.producer_port,
140 &prediction_ports,
141 )?;
142 }
143 for block in &mut result.explanations {
144 normalize_prediction_result_port(
145 &task.node_plan.node_id,
146 "explanation block",
147 &mut block.producer_port,
148 &prediction_ports,
149 )?;
150 }
151 Ok(())
152}
153
154fn validate_direct_sample_prediction_result(task: &NodeTask, result: &NodeResult) -> Result<()> {
159 if !result.observation_predictions.is_empty() {
160 return Err(DagMlError::RuntimeValidation(format!(
161 "direct terminal PREDICT node `{}` emitted observation-level predictions; relation aggregation is not permitted",
162 task.node_plan.node_id
163 )));
164 }
165 if !result.aggregated_predictions.is_empty() {
166 return Err(DagMlError::RuntimeValidation(format!(
167 "direct terminal PREDICT node `{}` emitted aggregated predictions; relation aggregation is not permitted",
168 task.node_plan.node_id
169 )));
170 }
171 Ok(())
172}
173
174#[derive(Default)]
175pub(crate) struct PhaseScopeResources<'a> {
176 pub(crate) data_provider: Option<&'a dyn RuntimeDataProvider>,
177 pub(crate) fold_set_override: Option<&'a FoldSet>,
181 pub(crate) node_filter: Option<&'a BTreeSet<NodeId>>,
185 pub(crate) suppress_inner_cv: bool,
188 pub(crate) nested_stacking: Option<NestedStackingInput<'a>>,
192 pub(crate) replay_artifact_handles: Option<&'a BTreeMap<NodeId, BTreeMap<String, HandleRef>>>,
193 pub(crate) replay_artifact_inputs:
194 Option<&'a BTreeMap<NodeId, BTreeMap<String, ArtifactInputSpec>>>,
195 pub(crate) replay_bundle_id: Option<&'a BundleId>,
196 pub(crate) data_envelopes: Option<&'a BTreeMap<String, ExternalDataPlanEnvelope>>,
197 pub(crate) prediction_cache_store: Option<&'a dyn RuntimePredictionCacheStore>,
198 pub(crate) prediction_cache_contracts:
199 Option<&'a BTreeMap<String, ReplayPredictionCacheContract>>,
200 pub(crate) direct_sample_prediction_only: bool,
205 pub(crate) artifact_store: Option<&'a mut InMemoryArtifactStore>,
206}
207
208#[derive(Clone, Copy, Debug, Eq, PartialEq)]
209enum HpoCandidateFitCvOutcome {
210 Completed,
211 Pruned,
212}
213
214struct HpoFoldFeedback<'a> {
215 trial_id: i64,
216 selection: &'a RuntimeHpoSelectionTarget,
217 session: &'a mut dyn RuntimeTunerSession,
218}
219
220fn validate_hpo_progressive_fold_topology(plan: &ExecutionPlan) -> Result<()> {
221 if nested_stacking_campaign_plan(plan)?.is_some() {
222 return Err(DagMlError::RuntimeValidation(
223 "runtime HPO progressive pruning does not support nested-stacking FIT_CV; the scheduler cannot attest one report-grade intermediate per outer fold"
224 .to_string(),
225 ));
226 }
227 let fold_set = plan.fold_set.as_ref().ok_or_else(|| {
228 DagMlError::RuntimeValidation(
229 "runtime HPO progressive pruning requires an explicit validated fold set".to_string(),
230 )
231 })?;
232 if fold_set.partition_mode != FoldPartitionMode::Partition {
233 return Err(DagMlError::RuntimeValidation(
234 "runtime HPO progressive pruning requires FoldPartitionMode::Partition; resampled folds cannot attest one stable validation resource step per sample"
235 .to_string(),
236 ));
237 }
238 i32::try_from(fold_set.folds.len()).map_err(|_| {
239 DagMlError::RuntimeValidation(
240 "runtime HPO fold count exceeds the native intermediate step range".to_string(),
241 )
242 })?;
243 Ok(())
244}
245
246impl SequentialScheduler {
247 pub fn execute_hpo_campaign(
254 &self,
255 plan: &ExecutionPlan,
256 controllers: &RuntimeControllerRegistry,
257 data_provider: &dyn RuntimeDataProvider,
258 ctx: &RunContext,
259 hpo: &RuntimeHpoExecutionContext,
260 ) -> Result<RuntimeHpoCampaignResult> {
261 plan.validate()?;
262 hpo.validate_for_plan(plan)?;
263 validate_hpo_progressive_fold_topology(plan)?;
264 let controller = controllers.get(&hpo.controller_id).ok_or_else(|| {
265 DagMlError::RuntimeValidation(format!(
266 "runtime HPO campaign controller `{}` is not registered",
267 hpo.controller_id
268 ))
269 })?;
270 let task = RuntimeHpoCampaignTask {
271 run_id: ctx.run_id.clone(),
272 operation_id: hpo.operation_id.clone(),
273 controller_id: hpo.controller_id.clone(),
274 target_node_id: hpo.target_node_id.clone(),
275 seed: ctx.root_seed,
276 };
277 let mut session = controller.create_tuner_session(&task, hpo)?;
278 let history_at_start = session.trial_history_len()?;
279 if history_at_start > hpo.trial_budget_total {
280 return Err(DagMlError::RuntimeValidation(format!(
281 "runtime HPO restored native history ({history_at_start}) exceeds total trial budget ({})",
282 hpo.trial_budget_total
283 )));
284 }
285 let remaining_trials = hpo.trial_budget_total - history_at_start;
286 let mut candidates = Vec::new();
287 let mut proposed_variant_ids = BTreeSet::new();
288 let mut trial_variants = BTreeMap::new();
293 let mut incumbent_variants = hpo.resume_variants.clone();
294 let mut terminal_trials = BTreeMap::new();
295 let mut completed_proposals = Vec::new();
296 let mut completed_reports = Vec::new();
297
298 for _ in 0..remaining_trials {
299 let Some(proposal) = session.ask()? else {
300 break;
301 };
302 if trial_variants
303 .insert(proposal.trial_id, proposal.variant.variant_id.clone())
304 .is_some()
305 {
306 return Err(DagMlError::RuntimeValidation(format!(
307 "runtime HPO session proposed duplicate trial `{}`",
308 proposal.trial_id
309 )));
310 }
311 if incumbent_variants
312 .insert(proposal.trial_id, proposal.variant.variant_id.clone())
313 .is_some()
314 {
315 return Err(DagMlError::RuntimeValidation(format!(
316 "runtime HPO session reused restored trial `{}`",
317 proposal.trial_id
318 )));
319 }
320 if !proposed_variant_ids.insert(proposal.variant.variant_id.clone()) {
321 return Err(DagMlError::RuntimeValidation(format!(
322 "runtime HPO session proposed duplicate variant `{}`",
323 proposal.variant.variant_id
324 )));
325 }
326 let mut candidate_plan = plan.clone();
327 candidate_plan.variants = vec![proposal.variant.clone()];
328 candidate_plan.validate()?;
329 let mut candidate_ctx =
330 RunContext::new(ctx.run_id.clone(), proposal.variant.seed.or(ctx.root_seed));
331 candidate_ctx.variant_id = Some(proposal.variant.variant_id.clone());
332
333 let evaluation = {
334 let mut feedback = HpoFoldFeedback {
335 trial_id: proposal.trial_id,
336 selection: &hpo.selection,
337 session: session.as_mut(),
338 };
339 self.execute_hpo_candidate_fit_cv(
340 &candidate_plan,
341 controllers,
342 data_provider,
343 &mut candidate_ctx,
344 &mut feedback,
345 )
346 };
347 let evaluation = match evaluation {
348 Ok(evaluation) => evaluation,
349 Err(error) => {
350 session.tell(
351 proposal.trial_id,
352 RuntimeHpoTerminal::Failed {
353 failure: RuntimeHpoFailure {
354 code: "DAGML_CV_ERROR".to_string(),
355 message: error.to_string(),
356 retryable: false,
357 },
358 },
359 )?;
360 terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
361 continue;
362 }
363 };
364 if evaluation == HpoCandidateFitCvOutcome::Pruned {
365 terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Pruned);
366 continue;
367 }
368 if let Err(error) = candidate_ctx
369 .collect_cross_fold_validation_scores(plan_oof_partition_mode(&candidate_plan))
370 {
371 session.tell(
372 proposal.trial_id,
373 RuntimeHpoTerminal::Failed {
374 failure: RuntimeHpoFailure {
375 code: "DAGML_SCORE_ERROR".to_string(),
376 message: error.to_string(),
377 retryable: false,
378 },
379 },
380 )?;
381 terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
382 continue;
383 }
384 let report = candidate_ctx
385 .score_collector
386 .iter()
387 .find(|report| {
388 report.producer_node == hpo.selection.producer_node
389 && report.producer_port.as_deref()
390 == Some(hpo.selection.producer_port.as_str())
391 && report.partition == PredictionPartition::Validation
392 && report
393 .fold_id
394 .as_ref()
395 .is_some_and(|fold| fold.as_str() == "avg")
396 })
397 .cloned();
398 let Some(mut report) = report else {
399 session.tell(
400 proposal.trial_id,
401 RuntimeHpoTerminal::Failed {
402 failure: RuntimeHpoFailure {
403 code: "DAGML_SCORE_MISSING".to_string(),
404 message: format!(
405 "runtime HPO trial `{}` emitted no target OOF average",
406 proposal.trial_id
407 ),
408 retryable: false,
409 },
410 },
411 )?;
412 terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
413 continue;
414 };
415 report.variant_id = Some(proposal.variant.variant_id.clone());
416 let score = report
417 .metrics
418 .get(hpo.selection.metric.name())
419 .copied()
420 .filter(|score| score.is_finite());
421 let Some(score) = score else {
422 session.tell(
423 proposal.trial_id,
424 RuntimeHpoTerminal::Failed {
425 failure: RuntimeHpoFailure {
426 code: "DAGML_SCORE_NONFINITE".to_string(),
427 message: format!(
428 "runtime HPO trial `{}` emitted no finite `{}` score",
429 proposal.trial_id,
430 hpo.selection.metric.name()
431 ),
432 retryable: false,
433 },
434 },
435 )?;
436 terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
437 continue;
438 };
439 session.tell(proposal.trial_id, RuntimeHpoTerminal::Completed { score })?;
440 terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Completed);
441 completed_proposals.push(proposal.clone());
442 completed_reports.push(RuntimeHpoCompletedReport {
443 trial_id: proposal.trial_id,
444 variant_id: proposal.variant.variant_id.clone(),
445 report: report.clone(),
446 });
447
448 let mut validation_reports = candidate_ctx
449 .score_collector
450 .iter()
451 .filter(|item| item.partition == PredictionPartition::Validation)
452 .cloned()
453 .collect::<Vec<_>>();
454 for item in &mut validation_reports {
455 item.variant_id = Some(proposal.variant.variant_id.clone());
456 }
457 candidates.push(RuntimeHpoCandidateEvaluation {
458 validation_predictions: capture_variant_validation_predictions(
459 &proposal.variant.variant_id,
460 None,
461 &candidate_ctx,
462 ),
463 lineage: candidate_ctx.lineage.records().cloned().collect(),
464 proposal,
465 score,
466 validation_reports,
467 });
468 }
469
470 let history_at_checkpoint = session.trial_history_len()?;
471 if history_at_checkpoint != hpo.trial_budget_total {
472 return Err(DagMlError::RuntimeValidation(format!(
473 "runtime HPO native history ended at {history_at_checkpoint}, expected total trial budget {}",
474 hpo.trial_budget_total
475 )));
476 }
477
478 let checkpoint = RuntimeHpoCheckpointResult {
479 artifact: session.checkpoint()?,
480 provenance: hpo.provenance.clone(),
481 operation_id: hpo.operation_id.clone(),
482 controller_id: hpo.controller_id.clone(),
483 target_node_id: hpo.target_node_id.clone(),
484 completed_proposals,
485 completed_reports,
486 trial_history_len: history_at_checkpoint,
487 };
488 validate_hpo_checkpoint_result(
489 &checkpoint,
490 hpo,
491 &trial_variants,
492 &terminal_trials,
493 history_at_start,
494 )?;
495 let incumbent = session.incumbent(&incumbent_variants)?.ok_or_else(|| {
496 DagMlError::RuntimeValidation(
497 "native HPO campaign has no completed native incumbent after terminalization"
498 .to_string(),
499 )
500 })?;
501 if incumbent.metric != hpo.selection.metric.name()
502 || incumbent.direction != hpo.selection.direction
503 || incumbent_variants.get(&incumbent.trial_id) != Some(&incumbent.variant_id)
504 || !incumbent.score.is_finite()
505 {
506 return Err(DagMlError::RuntimeValidation(
507 "native HPO incumbent is not bound to this scheduler campaign's metric, direction, trial, and variant"
508 .to_string(),
509 ));
510 }
511 let terminal_trials = session.terminal_trial_snapshots(&incumbent_variants)?;
512 if terminal_trials.len() != history_at_checkpoint as usize
513 || terminal_trials
514 .windows(2)
515 .any(|pair| pair[0].trial.id >= pair[1].trial.id)
516 {
517 return Err(DagMlError::RuntimeValidation(
518 "native HPO terminal ledger is not a complete strictly ordered history".to_string(),
519 ));
520 }
521 Ok(RuntimeHpoCampaignResult {
522 operation_id: hpo.operation_id.clone(),
523 controller_id: hpo.controller_id.clone(),
524 target_node_id: hpo.target_node_id.clone(),
525 candidates,
526 checkpoint,
527 incumbent,
528 terminal_trials,
529 })
530 }
531
532 fn execute_hpo_candidate_fit_cv(
533 &self,
534 plan: &ExecutionPlan,
535 controllers: &RuntimeControllerRegistry,
536 data_provider: &dyn RuntimeDataProvider,
537 ctx: &mut RunContext,
538 feedback: &mut HpoFoldFeedback<'_>,
539 ) -> Result<HpoCandidateFitCvOutcome> {
540 let candidate_plan = plan;
541 ctx.configure_global_oof_aggregation(candidate_plan, data_provider)?;
542 let fold_ids = candidate_plan
543 .fold_set
544 .as_ref()
545 .expect("progressive HPO topology was preflighted")
546 .folds
547 .iter()
548 .map(|fold| Some(fold.fold_id.clone()))
549 .collect::<Vec<_>>();
550 let variant = candidate_plan
551 .variants
552 .first()
553 .expect("candidate plan has exactly one variant");
554 for (step, fold_id) in fold_ids.into_iter().enumerate() {
555 let score_start = ctx.score_collector.len();
556 self.execute_phase_scope(
557 candidate_plan,
558 controllers,
559 ctx,
560 PhaseScope {
561 phase: Phase::FitCv,
562 variant_id: Some(variant.variant_id.clone()),
563 variant: Some(VariantExecutionSpec::from_plan(variant)),
564 fold_id: fold_id.clone(),
565 seed_root: variant.seed.or(ctx.root_seed),
566 },
567 PhaseScopeResources {
568 data_provider: Some(data_provider),
569 ..Default::default()
570 },
571 )?;
572 let fold_reports = ctx.score_collector[score_start..]
573 .iter()
574 .filter(|report| {
575 report.producer_node == feedback.selection.producer_node
576 && report.producer_port.as_deref()
577 == Some(feedback.selection.producer_port.as_str())
578 && report.partition == PredictionPartition::Validation
579 && report.fold_id == fold_id
580 })
581 .collect::<Vec<_>>();
582 let [fold_report] = fold_reports.as_slice() else {
583 return Err(DagMlError::RuntimeValidation(format!(
584 "runtime HPO trial `{}` must emit exactly one target validation report per fold; fold {fold_id:?} emitted {}",
585 feedback.trial_id,
586 fold_reports.len()
587 )));
588 };
589 let score = fold_report
590 .metrics
591 .get(feedback.selection.metric.name())
592 .copied()
593 .filter(|score| score.is_finite())
594 .ok_or_else(|| {
595 DagMlError::RuntimeValidation(format!(
596 "runtime HPO trial `{}` fold {fold_id:?} emitted no finite `{}` intermediate score",
597 feedback.trial_id,
598 feedback.selection.metric.name()
599 ))
600 })?;
601 let step = i32::try_from(step).map_err(|_| {
602 DagMlError::RuntimeValidation(
603 "runtime HPO fold intermediate count exceeds i32".to_string(),
604 )
605 })?;
606 if feedback
607 .session
608 .report_intermediate(RuntimeHpoIntermediate {
609 trial_id: feedback.trial_id,
610 step,
611 score,
612 })?
613 == RuntimeHpoIntermediateOutcome::Pruned
614 {
615 return Ok(HpoCandidateFitCvOutcome::Pruned);
616 }
617 }
618 Ok(HpoCandidateFitCvOutcome::Completed)
619 }
620
621 pub fn execute_phase(
622 &self,
623 plan: &ExecutionPlan,
624 controllers: &RuntimeControllerRegistry,
625 ctx: &mut RunContext,
626 phase: Phase,
627 ) -> Result<Vec<NodeResult>> {
628 plan.validate()?;
629 let variant_id = ctx.variant_id.clone();
630 let seed_root = ctx.root_seed;
631 self.execute_phase_scope(
632 plan,
633 controllers,
634 ctx,
635 PhaseScope {
636 phase,
637 variant_id,
638 variant: None,
639 fold_id: None,
640 seed_root,
641 },
642 PhaseScopeResources::default(),
643 )
644 }
645
646 pub fn execute_phase_with_data_provider(
647 &self,
648 plan: &ExecutionPlan,
649 controllers: &RuntimeControllerRegistry,
650 data_provider: &dyn RuntimeDataProvider,
651 ctx: &mut RunContext,
652 phase: Phase,
653 ) -> Result<Vec<NodeResult>> {
654 plan.validate()?;
655 let variant_id = ctx.variant_id.clone();
656 let seed_root = ctx.root_seed;
657 self.execute_phase_scope(
658 plan,
659 controllers,
660 ctx,
661 PhaseScope {
662 phase,
663 variant_id,
664 variant: None,
665 fold_id: None,
666 seed_root,
667 },
668 PhaseScopeResources {
669 data_provider: Some(data_provider),
670 ..Default::default()
671 },
672 )
673 }
674
675 pub fn execute_campaign_phase(
676 &self,
677 plan: &ExecutionPlan,
678 controllers: &RuntimeControllerRegistry,
679 ctx: &mut RunContext,
680 phase: Phase,
681 ) -> Result<Vec<NodeResult>> {
682 plan.validate()?;
683 if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
684 return Err(DagMlError::RuntimeValidation(
685 "nested stacking FIT_CV requires execute_campaign_phase_with_data_provider so the scheduler can materialize parent-bound inner folds"
686 .to_string(),
687 ));
688 }
689 let mut results = Vec::new();
690 let fold_ids = if phase == Phase::FitCv {
691 plan.fold_set
692 .as_ref()
693 .map(|fold_set| {
694 fold_set
695 .folds
696 .iter()
697 .map(|fold| Some(fold.fold_id.clone()))
698 .collect::<Vec<_>>()
699 })
700 .unwrap_or_else(|| vec![None])
701 } else {
702 vec![None]
703 };
704 for variant in &plan.variants {
705 if ctx
706 .variant_id
707 .as_ref()
708 .is_some_and(|requested| requested != &variant.variant_id)
709 {
710 continue;
711 }
712 for fold_id in &fold_ids {
713 let seed_root = variant.seed.or(ctx.root_seed);
714 results.extend(self.execute_phase_scope(
715 plan,
716 controllers,
717 ctx,
718 PhaseScope {
719 phase,
720 variant_id: Some(variant.variant_id.clone()),
721 variant: Some(VariantExecutionSpec::from_plan(variant)),
722 fold_id: fold_id.clone(),
723 seed_root,
724 },
725 PhaseScopeResources::default(),
726 )?);
727 }
728 }
729 Ok(results)
730 }
731
732 pub fn execute_campaign_phase_with_data_provider(
733 &self,
734 plan: &ExecutionPlan,
735 controllers: &RuntimeControllerRegistry,
736 data_provider: &dyn RuntimeDataProvider,
737 ctx: &mut RunContext,
738 phase: Phase,
739 ) -> Result<Vec<NodeResult>> {
740 plan.validate()?;
741 if phase == Phase::FitCv {
742 ctx.configure_global_oof_aggregation(plan, data_provider)?;
743 if let Some(nested) = nested_stacking_campaign_plan(plan)? {
744 return self.execute_nested_stacking_fit_cv(
745 plan,
746 controllers,
747 data_provider,
748 ctx,
749 &nested,
750 );
751 }
752 }
753 let mut results = Vec::new();
754 let fold_ids = if phase == Phase::FitCv {
755 plan.fold_set
756 .as_ref()
757 .map(|fold_set| {
758 fold_set
759 .folds
760 .iter()
761 .map(|fold| Some(fold.fold_id.clone()))
762 .collect::<Vec<_>>()
763 })
764 .unwrap_or_else(|| vec![None])
765 } else {
766 vec![None]
767 };
768 for variant in &plan.variants {
769 if ctx
770 .variant_id
771 .as_ref()
772 .is_some_and(|requested| requested != &variant.variant_id)
773 {
774 continue;
775 }
776 for fold_id in &fold_ids {
777 let seed_root = variant.seed.or(ctx.root_seed);
778 results.extend(self.execute_phase_scope(
779 plan,
780 controllers,
781 ctx,
782 PhaseScope {
783 phase,
784 variant_id: Some(variant.variant_id.clone()),
785 variant: Some(VariantExecutionSpec::from_plan(variant)),
786 fold_id: fold_id.clone(),
787 seed_root,
788 },
789 PhaseScopeResources {
790 data_provider: Some(data_provider),
791 ..Default::default()
792 },
793 )?);
794 }
795 }
796 Ok(results)
797 }
798
799 pub fn execute_campaign_phase_with_data_provider_and_artifact_store(
800 &self,
801 plan: &ExecutionPlan,
802 controllers: &RuntimeControllerRegistry,
803 data_provider: &dyn RuntimeDataProvider,
804 artifact_store: &mut InMemoryArtifactStore,
805 ctx: &mut RunContext,
806 phase: Phase,
807 ) -> Result<Vec<NodeResult>> {
808 plan.validate()?;
809 if phase == Phase::FitCv {
810 ctx.configure_global_oof_aggregation(plan, data_provider)?;
811 if let Some(nested) = nested_stacking_campaign_plan(plan)? {
812 return self.execute_nested_stacking_fit_cv(
816 plan,
817 controllers,
818 data_provider,
819 ctx,
820 &nested,
821 );
822 }
823 }
824 let mut results = Vec::new();
825 let fold_ids = if phase == Phase::FitCv {
826 plan.fold_set
827 .as_ref()
828 .map(|fold_set| {
829 fold_set
830 .folds
831 .iter()
832 .map(|fold| Some(fold.fold_id.clone()))
833 .collect::<Vec<_>>()
834 })
835 .unwrap_or_else(|| vec![None])
836 } else {
837 vec![None]
838 };
839 for variant in &plan.variants {
840 if ctx
841 .variant_id
842 .as_ref()
843 .is_some_and(|requested| requested != &variant.variant_id)
844 {
845 continue;
846 }
847 for fold_id in &fold_ids {
848 let seed_root = variant.seed.or(ctx.root_seed);
849 results.extend(self.execute_phase_scope(
850 plan,
851 controllers,
852 ctx,
853 PhaseScope {
854 phase,
855 variant_id: Some(variant.variant_id.clone()),
856 variant: Some(VariantExecutionSpec::from_plan(variant)),
857 fold_id: fold_id.clone(),
858 seed_root,
859 },
860 PhaseScopeResources {
861 data_provider: Some(data_provider),
862 artifact_store: Some(&mut *artifact_store),
863 ..Default::default()
864 },
865 )?);
866 }
867 }
868 Ok(results)
869 }
870
871 fn execute_nested_stacking_fit_cv(
880 &self,
881 plan: &ExecutionPlan,
882 controllers: &RuntimeControllerRegistry,
883 data_provider: &dyn RuntimeDataProvider,
884 ctx: &mut RunContext,
885 nested: &NestedStackingCampaignPlan,
886 ) -> Result<Vec<NodeResult>> {
887 let parent_fold_ids = nested
888 .outer_scopes
889 .iter()
890 .map(|outer| outer.outer_fold_id.clone())
891 .collect::<BTreeSet<_>>();
892 if let Some(existing) = &ctx.validation_scoring_fold_ids {
893 if existing != &parent_fold_ids {
894 return Err(DagMlError::RuntimeValidation(
895 "nested stacking cannot reuse a run context with a different report-grade outer fold set"
896 .to_string(),
897 ));
898 }
899 } else {
900 ctx.validation_scoring_fold_ids = Some(parent_fold_ids);
901 }
902 let mut results = Vec::new();
903 for variant in &plan.variants {
904 if ctx
905 .variant_id
906 .as_ref()
907 .is_some_and(|requested| requested != &variant.variant_id)
908 {
909 continue;
910 }
911 let seed_root = variant.seed.or(ctx.root_seed);
912 let variant_id = Some(variant.variant_id.clone());
913 let variant_spec = Some(VariantExecutionSpec::from_plan(variant));
914 for outer in &nested.outer_scopes {
915 for inner_fold in &outer.inner.inner_fold_set.folds {
916 results.extend(self.execute_phase_scope(
917 plan,
918 controllers,
919 ctx,
920 PhaseScope {
921 phase: Phase::FitCv,
922 variant_id: variant_id.clone(),
923 variant: variant_spec.clone(),
924 fold_id: Some(inner_fold.fold_id.clone()),
925 seed_root,
926 },
927 PhaseScopeResources {
928 data_provider: Some(data_provider),
929 fold_set_override: Some(&outer.inner.inner_fold_set),
930 node_filter: Some(&nested.base_node_ids),
931 suppress_inner_cv: true,
932 ..Default::default()
933 },
934 )?);
935 }
936
937 results.extend(self.execute_phase_scope(
940 plan,
941 controllers,
942 ctx,
943 PhaseScope {
944 phase: Phase::FitCv,
945 variant_id: variant_id.clone(),
946 variant: variant_spec.clone(),
947 fold_id: Some(outer.outer_fold_id.clone()),
948 seed_root,
949 },
950 PhaseScopeResources {
951 data_provider: Some(data_provider),
952 node_filter: Some(&nested.base_node_ids),
953 suppress_inner_cv: true,
954 ..Default::default()
955 },
956 )?);
957
958 let meta_only = BTreeSet::from([nested.meta_node_id.clone()]);
959 results.extend(self.execute_phase_scope(
960 plan,
961 controllers,
962 ctx,
963 PhaseScope {
964 phase: Phase::FitCv,
965 variant_id: variant_id.clone(),
966 variant: variant_spec.clone(),
967 fold_id: Some(outer.outer_fold_id.clone()),
968 seed_root,
969 },
970 PhaseScopeResources {
971 data_provider: Some(data_provider),
972 node_filter: Some(&meta_only),
973 suppress_inner_cv: true,
974 nested_stacking: Some(NestedStackingInput {
975 meta_node_id: &nested.meta_node_id,
976 inner: &outer.inner,
977 }),
978 ..Default::default()
979 },
980 )?);
981 }
982 }
983 Ok(results)
984 }
985
986 pub fn execute_bundle_replay(
987 &self,
988 replay: BundleReplayExecution<'_>,
989 ctx: &mut RunContext,
990 ) -> Result<Vec<NodeResult>> {
991 self.execute_bundle_replay_with_prediction_mode(replay, ctx, false)
992 }
993
994 pub(crate) fn execute_direct_sample_bundle_replay(
998 &self,
999 replay: BundleReplayExecution<'_>,
1000 ctx: &mut RunContext,
1001 ) -> Result<Vec<NodeResult>> {
1002 if replay.replay_request.phase != Phase::Predict {
1003 return Err(DagMlError::RuntimeValidation(
1004 "direct sample bundle replay is valid only for PREDICT".to_string(),
1005 ));
1006 }
1007 self.execute_bundle_replay_with_prediction_mode(replay, ctx, true)
1008 }
1009
1010 fn execute_bundle_replay_with_prediction_mode(
1011 &self,
1012 replay: BundleReplayExecution<'_>,
1013 ctx: &mut RunContext,
1014 direct_sample_prediction_only: bool,
1015 ) -> Result<Vec<NodeResult>> {
1016 replay.bundle.validate_against_plan(replay.plan)?;
1017 replay
1018 .replay_request
1019 .validate_for_bundle_with_prediction_cache_store(
1020 replay.bundle,
1021 replay.prediction_cache_store.is_some(),
1022 )?;
1023 replay
1024 .bundle
1025 .validate_replay_envelopes(replay.data_envelopes)?;
1026 let prediction_cache_contracts = if replay.replay_request.phase == Phase::Refit {
1027 Some(replay_prediction_cache_contracts(replay.bundle)?)
1028 } else {
1029 None
1030 };
1031 if replay.replay_request.phase == Phase::Refit {
1032 preload_replay_prediction_cache_store(
1033 replay.bundle,
1034 replay.prediction_cache_store,
1035 ctx,
1036 )?;
1037 }
1038 let replay_artifacts = materialize_replay_artifact_handles(
1039 replay.plan,
1040 replay.bundle,
1041 replay.replay_request,
1042 replay.artifact_store,
1043 ctx,
1044 )?;
1045 let selected_variant = replay
1046 .bundle
1047 .selected_variant_id
1048 .as_ref()
1049 .map(|selected| {
1050 replay
1051 .plan
1052 .variants
1053 .iter()
1054 .find(|variant| &variant.variant_id == selected)
1055 .map(VariantExecutionSpec::from_plan)
1056 .ok_or_else(|| {
1057 DagMlError::RuntimeValidation(format!(
1058 "bundle `{}` selected unknown variant `{selected}`",
1059 replay.bundle.bundle_id
1060 ))
1061 })
1062 })
1063 .transpose()?;
1064 let seed_root = selected_variant
1065 .as_ref()
1066 .and_then(|variant| variant.seed)
1067 .or(ctx.root_seed);
1068
1069 self.execute_phase_scope(
1070 replay.plan,
1071 replay.controllers,
1072 ctx,
1073 PhaseScope {
1074 phase: replay.replay_request.phase,
1075 variant_id: replay.bundle.selected_variant_id.clone(),
1076 variant: selected_variant,
1077 fold_id: None,
1078 seed_root,
1079 },
1080 PhaseScopeResources {
1081 data_provider: Some(replay.data_provider),
1082 replay_artifact_handles: Some(&replay_artifacts.handles),
1083 replay_artifact_inputs: Some(&replay_artifacts.inputs),
1084 replay_bundle_id: Some(&replay.bundle.bundle_id),
1085 data_envelopes: Some(replay.data_envelopes),
1086 prediction_cache_store: replay.prediction_cache_store,
1087 prediction_cache_contracts: prediction_cache_contracts.as_ref(),
1088 direct_sample_prediction_only,
1089 ..Default::default()
1090 },
1091 )
1092 }
1093
1094 fn execute_phase_scope(
1095 &self,
1096 plan: &ExecutionPlan,
1097 controllers: &RuntimeControllerRegistry,
1098 ctx: &mut RunContext,
1099 scope: PhaseScope,
1100 mut resources: PhaseScopeResources<'_>,
1101 ) -> Result<Vec<NodeResult>> {
1102 let _phase_span = crate::observability::phase_span(
1103 ctx.run_id.as_str(),
1104 plan.id.as_str(),
1105 scope.phase.as_str(),
1106 scope.variant_id.as_ref().map(VariantId::as_str),
1107 scope.fold_id.as_ref().map(FoldId::as_str),
1108 )
1109 .entered();
1110 let mut results = Vec::new();
1111 let mut output_handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
1112 let mut output_data_views =
1113 BTreeMap::<NodeId, BTreeMap<String, DataProviderViewSpec>>::new();
1114 let mut input_lineage = BTreeMap::<NodeId, LineageId>::new();
1115
1116 for level in plan.node_parallel_levels_for_phase(scope.phase)? {
1117 for node_id in &level {
1118 if resources
1119 .node_filter
1120 .is_some_and(|allowed| !allowed.contains(node_id))
1121 {
1122 continue;
1123 }
1124 let node_plan = plan
1125 .node_plans
1126 .get(node_id)
1127 .expect("execution plan was validated");
1128 if let Some(reduction) = merge_reduction_mode(plan, node_plan) {
1135 if let Some(mut result) =
1136 reassemble_branch_merge(plan, node_plan, ctx, &scope, reduction)?
1137 {
1138 let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1139 let task = NodeTask {
1140 inner_fold_set: None,
1141 run_id: ctx.run_id.clone(),
1142 node_plan: task_node_plan.clone(),
1143 phase: scope.phase,
1144 variant_id: scope.variant_id.clone(),
1145 variant: scope.variant.clone(),
1146 fold_id: scope.fold_id.clone(),
1147 branch_path: Vec::new(),
1148 input_handles: BTreeMap::new(),
1149 data_views: BTreeMap::new(),
1150 prediction_inputs: BTreeMap::new(),
1151 artifact_inputs: BTreeMap::new(),
1152 required_loss_attestations: NodeTask::required_loss_attestations_for(
1153 &task_node_plan,
1154 scope.phase,
1155 )?,
1156 fit_influence: FitInfluenceTask::default(),
1157 seed: None,
1158 };
1159 normalize_result_prediction_ports(plan, &task, &mut result)?;
1160 result.validate_for_task(&task)?;
1161 for prediction in &result.predictions {
1162 ctx.prediction_store.append(prediction.clone())?;
1163 }
1164 apply_result_scoring(
1165 &result,
1166 &mut ctx.score_collector,
1167 &mut ctx.regression_target_records,
1168 )?;
1169 ctx.lineage.record(result.lineage.clone())?;
1170 output_handles.insert(node_id.clone(), result.outputs.clone());
1171 input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1172 results.push(result);
1173 }
1174 continue;
1175 }
1176 let controller = controllers.get(&node_plan.controller_id).ok_or_else(|| {
1177 DagMlError::RuntimeValidation(format!(
1178 "runtime controller `{}` is not registered",
1179 node_plan.controller_id
1180 ))
1181 })?;
1182 let collected_inputs = collect_input_handles(
1183 plan,
1184 node_plan,
1185 &output_handles,
1186 &output_data_views,
1187 &resources,
1188 ctx,
1189 &scope,
1190 )?;
1191 if collected_inputs.skip_node {
1192 continue;
1193 }
1194 let mut input_handles = collected_inputs.handles;
1195 let mut prediction_inputs = collected_inputs.prediction_inputs;
1196 if let Some(nested) = resources.nested_stacking.as_ref() {
1197 replace_nested_stacking_fit_cv_inputs(
1198 plan,
1199 node_plan,
1200 ctx,
1201 &scope,
1202 nested,
1203 &mut input_handles,
1204 &mut prediction_inputs,
1205 )?;
1206 }
1207 let mut artifact_inputs = BTreeMap::new();
1208 if let Some(node_artifact_handles) = resources
1209 .replay_artifact_handles
1210 .and_then(|handles| handles.get(node_id))
1211 {
1212 for (key, handle) in node_artifact_handles {
1213 if input_handles.insert(key.clone(), handle.clone()).is_some() {
1214 return Err(DagMlError::RuntimeValidation(format!(
1215 "node `{node_id}` received duplicate replay artifact input `{key}`"
1216 )));
1217 }
1218 }
1219 }
1220 if let Some(node_artifact_inputs) = resources
1221 .replay_artifact_inputs
1222 .and_then(|inputs| inputs.get(node_id))
1223 {
1224 for (key, spec) in node_artifact_inputs {
1225 if artifact_inputs.insert(key.clone(), spec.clone()).is_some() {
1226 return Err(DagMlError::RuntimeValidation(format!(
1227 "node `{node_id}` received duplicate replay artifact metadata `{key}`"
1228 )));
1229 }
1230 }
1231 }
1232 let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1233 let inner_fold_set = (!resources.suppress_inner_cv)
1234 .then(|| {
1235 inner_fold_set_for_scope(
1236 &plan.campaign,
1237 plan.fold_set.as_ref(),
1238 node_plan,
1239 &scope,
1240 )
1241 })
1242 .transpose()?
1243 .flatten();
1244 let fit_influence = fit_influence_task_for_node(
1245 plan,
1246 &task_node_plan,
1247 &collected_inputs.data_views,
1248 )?;
1249 let task = NodeTask {
1250 inner_fold_set,
1251 run_id: ctx.run_id.clone(),
1252 node_plan: task_node_plan.clone(),
1253 phase: scope.phase,
1254 variant_id: scope.variant_id.clone(),
1255 variant: scope.variant.clone(),
1256 fold_id: scope.fold_id.clone(),
1257 branch_path: Vec::new(),
1258 input_handles,
1259 data_views: collected_inputs.data_views,
1260 prediction_inputs,
1261 artifact_inputs,
1262 required_loss_attestations: NodeTask::required_loss_attestations_for(
1263 &task_node_plan,
1264 scope.phase,
1265 )?,
1266 fit_influence,
1267 seed: derive_task_seed(
1268 scope.seed_root,
1269 scope.variant_id.as_ref(),
1270 scope.fold_id.as_ref(),
1271 &task_node_plan,
1272 scope.phase,
1273 ),
1274 };
1275 let _node_span = crate::observability::node_span(
1276 task.run_id.as_str(),
1277 plan.id.as_str(),
1278 task.phase.as_str(),
1279 task.node_plan.node_id.as_str(),
1280 task.node_plan.controller_id.as_str(),
1281 )
1282 .entered();
1283 let mut result = if task.node_plan.kind == NodeKind::Tuner {
1284 return Err(DagMlError::RuntimeValidation(format!(
1285 "tuner node `{}` requires execute_hpo_campaign with an explicit RuntimeHpoExecutionContext",
1286 task.node_plan.node_id
1287 )));
1288 } else {
1289 match resources.data_provider {
1290 Some(data_provider) => {
1291 controller.invoke_with_data_provider(&task, data_provider)?
1292 }
1293 None => controller.invoke(&task)?,
1294 }
1295 };
1296 record_fit_influence_diagnostic(&task, &mut result);
1297 normalize_result_prediction_ports(plan, &task, &mut result)?;
1298 result.validate_for_task(&task)?;
1299 if resources.direct_sample_prediction_only {
1300 validate_direct_sample_prediction_result(&task, &result)?;
1301 } else {
1302 apply_result_prediction_aggregation(
1303 plan,
1304 controllers,
1305 &task,
1306 &mut result,
1307 &resources,
1308 )?;
1309 }
1310 if let Some(nested) = resources.nested_stacking.as_ref() {
1311 attach_nested_stacking_input_lineage(&mut result, plan, &task, ctx, nested)?;
1312 } else {
1313 attach_coordinator_input_lineage(
1314 &mut result,
1315 plan,
1316 &task.node_plan.node_id,
1317 &input_lineage,
1318 )?;
1319 }
1320 if let Some(store) = resources.artifact_store.as_deref_mut() {
1321 if scope.phase == Phase::Refit {
1322 store.capture_refit_artifacts(&task, &result)?;
1323 }
1324 }
1325 for prediction in &result.predictions {
1326 ctx.prediction_store.append(prediction.clone())?;
1327 }
1328 for prediction in &result.aggregated_predictions {
1329 ctx.aggregated_prediction_store.append(prediction.clone())?;
1330 }
1331 apply_result_scoring(
1332 &result,
1333 &mut ctx.score_collector,
1334 &mut ctx.regression_target_records,
1335 )?;
1336 ctx.lineage.record(result.lineage.clone())?;
1337 let data_views = derive_output_data_views(plan, &task, &result)?;
1338 output_handles.insert(node_id.clone(), result.outputs.clone());
1339 output_data_views.insert(node_id.clone(), data_views);
1340 input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1341 results.push(result);
1342 }
1343 }
1344
1345 Ok(results)
1346 }
1347}
1348
1349impl ParallelScheduler {
1350 pub fn execute_phase(
1351 &self,
1352 plan: &ExecutionPlan,
1353 controllers: &RuntimeControllerRegistry,
1354 ctx: &mut RunContext,
1355 phase: Phase,
1356 ) -> Result<Vec<NodeResult>> {
1357 plan.validate()?;
1358 let variant_id = ctx.variant_id.clone();
1359 let seed_root = ctx.root_seed;
1360 self.execute_phase_scope(
1361 plan,
1362 controllers,
1363 ctx,
1364 PhaseScope {
1365 phase,
1366 variant_id,
1367 variant: None,
1368 fold_id: None,
1369 seed_root,
1370 },
1371 PhaseScopeResources::default(),
1372 )
1373 }
1374
1375 pub fn execute_phase_with_data_provider(
1376 &self,
1377 plan: &ExecutionPlan,
1378 controllers: &RuntimeControllerRegistry,
1379 data_provider: &dyn RuntimeDataProvider,
1380 ctx: &mut RunContext,
1381 phase: Phase,
1382 ) -> Result<Vec<NodeResult>> {
1383 plan.validate()?;
1384 let variant_id = ctx.variant_id.clone();
1385 let seed_root = ctx.root_seed;
1386 self.execute_phase_scope(
1387 plan,
1388 controllers,
1389 ctx,
1390 PhaseScope {
1391 phase,
1392 variant_id,
1393 variant: None,
1394 fold_id: None,
1395 seed_root,
1396 },
1397 PhaseScopeResources {
1398 data_provider: Some(data_provider),
1399 ..Default::default()
1400 },
1401 )
1402 }
1403
1404 pub fn execute_campaign_phase(
1405 &self,
1406 plan: &ExecutionPlan,
1407 controllers: &RuntimeControllerRegistry,
1408 ctx: &mut RunContext,
1409 phase: Phase,
1410 ) -> Result<Vec<NodeResult>> {
1411 plan.validate()?;
1412 if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1413 return Err(DagMlError::RuntimeValidation(
1414 "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1415 .to_string(),
1416 ));
1417 }
1418 let mut results = Vec::new();
1419 let fold_ids = if phase == Phase::FitCv {
1420 plan.fold_set
1421 .as_ref()
1422 .map(|fold_set| {
1423 fold_set
1424 .folds
1425 .iter()
1426 .map(|fold| Some(fold.fold_id.clone()))
1427 .collect::<Vec<_>>()
1428 })
1429 .unwrap_or_else(|| vec![None])
1430 } else {
1431 vec![None]
1432 };
1433 for variant in &plan.variants {
1434 if ctx
1435 .variant_id
1436 .as_ref()
1437 .is_some_and(|requested| requested != &variant.variant_id)
1438 {
1439 continue;
1440 }
1441 for fold_id in &fold_ids {
1442 let seed_root = variant.seed.or(ctx.root_seed);
1443 results.extend(self.execute_phase_scope(
1444 plan,
1445 controllers,
1446 ctx,
1447 PhaseScope {
1448 phase,
1449 variant_id: Some(variant.variant_id.clone()),
1450 variant: Some(VariantExecutionSpec::from_plan(variant)),
1451 fold_id: fold_id.clone(),
1452 seed_root,
1453 },
1454 PhaseScopeResources::default(),
1455 )?);
1456 }
1457 }
1458 Ok(results)
1459 }
1460
1461 pub fn execute_campaign_phase_with_data_provider(
1462 &self,
1463 plan: &ExecutionPlan,
1464 controllers: &RuntimeControllerRegistry,
1465 data_provider: &dyn RuntimeDataProvider,
1466 ctx: &mut RunContext,
1467 phase: Phase,
1468 ) -> Result<Vec<NodeResult>> {
1469 plan.validate()?;
1470 if phase == Phase::FitCv {
1471 ctx.configure_global_oof_aggregation(plan, data_provider)?;
1472 }
1473 if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1474 return Err(DagMlError::RuntimeValidation(
1475 "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1476 .to_string(),
1477 ));
1478 }
1479 let mut results = Vec::new();
1480 let fold_ids = if phase == Phase::FitCv {
1481 plan.fold_set
1482 .as_ref()
1483 .map(|fold_set| {
1484 fold_set
1485 .folds
1486 .iter()
1487 .map(|fold| Some(fold.fold_id.clone()))
1488 .collect::<Vec<_>>()
1489 })
1490 .unwrap_or_else(|| vec![None])
1491 } else {
1492 vec![None]
1493 };
1494 for variant in &plan.variants {
1495 if ctx
1496 .variant_id
1497 .as_ref()
1498 .is_some_and(|requested| requested != &variant.variant_id)
1499 {
1500 continue;
1501 }
1502 for fold_id in &fold_ids {
1503 let seed_root = variant.seed.or(ctx.root_seed);
1504 results.extend(self.execute_phase_scope(
1505 plan,
1506 controllers,
1507 ctx,
1508 PhaseScope {
1509 phase,
1510 variant_id: Some(variant.variant_id.clone()),
1511 variant: Some(VariantExecutionSpec::from_plan(variant)),
1512 fold_id: fold_id.clone(),
1513 seed_root,
1514 },
1515 PhaseScopeResources {
1516 data_provider: Some(data_provider),
1517 ..Default::default()
1518 },
1519 )?);
1520 }
1521 }
1522 Ok(results)
1523 }
1524
1525 pub fn execute_campaign_phase_with_data_provider_and_artifact_store(
1526 &self,
1527 plan: &ExecutionPlan,
1528 controllers: &RuntimeControllerRegistry,
1529 data_provider: &dyn RuntimeDataProvider,
1530 artifact_store: &mut InMemoryArtifactStore,
1531 ctx: &mut RunContext,
1532 phase: Phase,
1533 ) -> Result<Vec<NodeResult>> {
1534 plan.validate()?;
1535 if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1536 return Err(DagMlError::RuntimeValidation(
1537 "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1538 .to_string(),
1539 ));
1540 }
1541 let mut results = Vec::new();
1542 let fold_ids = if phase == Phase::FitCv {
1543 plan.fold_set
1544 .as_ref()
1545 .map(|fold_set| {
1546 fold_set
1547 .folds
1548 .iter()
1549 .map(|fold| Some(fold.fold_id.clone()))
1550 .collect::<Vec<_>>()
1551 })
1552 .unwrap_or_else(|| vec![None])
1553 } else {
1554 vec![None]
1555 };
1556 for variant in &plan.variants {
1557 if ctx
1558 .variant_id
1559 .as_ref()
1560 .is_some_and(|requested| requested != &variant.variant_id)
1561 {
1562 continue;
1563 }
1564 for fold_id in &fold_ids {
1565 let seed_root = variant.seed.or(ctx.root_seed);
1566 results.extend(self.execute_phase_scope(
1567 plan,
1568 controllers,
1569 ctx,
1570 PhaseScope {
1571 phase,
1572 variant_id: Some(variant.variant_id.clone()),
1573 variant: Some(VariantExecutionSpec::from_plan(variant)),
1574 fold_id: fold_id.clone(),
1575 seed_root,
1576 },
1577 PhaseScopeResources {
1578 data_provider: Some(data_provider),
1579 artifact_store: Some(&mut *artifact_store),
1580 ..Default::default()
1581 },
1582 )?);
1583 }
1584 }
1585 Ok(results)
1586 }
1587
1588 pub fn execute_bundle_replay(
1589 &self,
1590 replay: BundleReplayExecution<'_>,
1591 ctx: &mut RunContext,
1592 ) -> Result<Vec<NodeResult>> {
1593 replay.bundle.validate_against_plan(replay.plan)?;
1594 replay
1595 .replay_request
1596 .validate_for_bundle_with_prediction_cache_store(
1597 replay.bundle,
1598 replay.prediction_cache_store.is_some(),
1599 )?;
1600 replay
1601 .bundle
1602 .validate_replay_envelopes(replay.data_envelopes)?;
1603 let prediction_cache_contracts = if replay.replay_request.phase == Phase::Refit {
1604 Some(replay_prediction_cache_contracts(replay.bundle)?)
1605 } else {
1606 None
1607 };
1608 if replay.replay_request.phase == Phase::Refit {
1609 preload_replay_prediction_cache_store(
1610 replay.bundle,
1611 replay.prediction_cache_store,
1612 ctx,
1613 )?;
1614 }
1615 let replay_artifacts = materialize_replay_artifact_handles(
1616 replay.plan,
1617 replay.bundle,
1618 replay.replay_request,
1619 replay.artifact_store,
1620 ctx,
1621 )?;
1622 let selected_variant = replay
1623 .bundle
1624 .selected_variant_id
1625 .as_ref()
1626 .map(|selected| {
1627 replay
1628 .plan
1629 .variants
1630 .iter()
1631 .find(|variant| &variant.variant_id == selected)
1632 .map(VariantExecutionSpec::from_plan)
1633 .ok_or_else(|| {
1634 DagMlError::RuntimeValidation(format!(
1635 "bundle `{}` selected unknown variant `{selected}`",
1636 replay.bundle.bundle_id
1637 ))
1638 })
1639 })
1640 .transpose()?;
1641 let seed_root = selected_variant
1642 .as_ref()
1643 .and_then(|variant| variant.seed)
1644 .or(ctx.root_seed);
1645
1646 self.execute_phase_scope(
1647 replay.plan,
1648 replay.controllers,
1649 ctx,
1650 PhaseScope {
1651 phase: replay.replay_request.phase,
1652 variant_id: replay.bundle.selected_variant_id.clone(),
1653 variant: selected_variant,
1654 fold_id: None,
1655 seed_root,
1656 },
1657 PhaseScopeResources {
1658 data_provider: Some(replay.data_provider),
1659 replay_artifact_handles: Some(&replay_artifacts.handles),
1660 replay_artifact_inputs: Some(&replay_artifacts.inputs),
1661 replay_bundle_id: Some(&replay.bundle.bundle_id),
1662 data_envelopes: Some(replay.data_envelopes),
1663 prediction_cache_store: replay.prediction_cache_store,
1664 prediction_cache_contracts: prediction_cache_contracts.as_ref(),
1665 ..Default::default()
1666 },
1667 )
1668 }
1669
1670 fn execute_phase_scope(
1671 &self,
1672 plan: &ExecutionPlan,
1673 controllers: &RuntimeControllerRegistry,
1674 ctx: &mut RunContext,
1675 scope: PhaseScope,
1676 mut resources: PhaseScopeResources<'_>,
1677 ) -> Result<Vec<NodeResult>> {
1678 let phase_span = crate::observability::phase_span(
1682 ctx.run_id.as_str(),
1683 plan.id.as_str(),
1684 scope.phase.as_str(),
1685 scope.variant_id.as_ref().map(VariantId::as_str),
1686 scope.fold_id.as_ref().map(FoldId::as_str),
1687 );
1688 let _phase_entered = phase_span.clone().entered();
1689 let plan_id = plan.id.as_str();
1691 plan.validate_parallel_controller_capabilities(self.max_workers, scope.phase)?;
1692 let mut results = Vec::new();
1693 let mut output_handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
1694 let mut output_data_views =
1695 BTreeMap::<NodeId, BTreeMap<String, DataProviderViewSpec>>::new();
1696 let mut input_lineage = BTreeMap::<NodeId, LineageId>::new();
1697
1698 for level in plan.node_parallel_levels_for_phase(scope.phase)? {
1699 let mut prepared = Vec::<PreparedNodeTask>::new();
1700 let mut merge_nodes = Vec::<(NodeId, MergeReduction)>::new();
1707 for node_id in &level {
1708 let node_plan = plan
1709 .node_plans
1710 .get(node_id)
1711 .expect("execution plan was validated");
1712 if let Some(reduction) = merge_reduction_mode(plan, node_plan) {
1713 merge_nodes.push((node_id.clone(), reduction));
1714 continue;
1715 }
1716 let collected_inputs = collect_input_handles(
1717 plan,
1718 node_plan,
1719 &output_handles,
1720 &output_data_views,
1721 &resources,
1722 ctx,
1723 &scope,
1724 )?;
1725 if collected_inputs.skip_node {
1726 continue;
1727 }
1728 let mut input_handles = collected_inputs.handles;
1729 let mut artifact_inputs = BTreeMap::new();
1730 if let Some(node_artifact_handles) = resources
1731 .replay_artifact_handles
1732 .and_then(|handles| handles.get(node_id))
1733 {
1734 for (key, handle) in node_artifact_handles {
1735 if input_handles.insert(key.clone(), handle.clone()).is_some() {
1736 return Err(DagMlError::RuntimeValidation(format!(
1737 "node `{node_id}` received duplicate replay artifact input `{key}`"
1738 )));
1739 }
1740 }
1741 }
1742 if let Some(node_artifact_inputs) = resources
1743 .replay_artifact_inputs
1744 .and_then(|inputs| inputs.get(node_id))
1745 {
1746 for (key, spec) in node_artifact_inputs {
1747 if artifact_inputs.insert(key.clone(), spec.clone()).is_some() {
1748 return Err(DagMlError::RuntimeValidation(format!(
1749 "node `{node_id}` received duplicate replay artifact metadata `{key}`"
1750 )));
1751 }
1752 }
1753 }
1754 let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1755 let inner_fold_set = inner_fold_set_for_scope(
1756 &plan.campaign,
1757 plan.fold_set.as_ref(),
1758 node_plan,
1759 &scope,
1760 )?;
1761 let fit_influence = fit_influence_task_for_node(
1762 plan,
1763 &task_node_plan,
1764 &collected_inputs.data_views,
1765 )?;
1766 prepared.push(PreparedNodeTask {
1767 node_id: node_id.clone(),
1768 task: NodeTask {
1769 inner_fold_set,
1770 run_id: ctx.run_id.clone(),
1771 node_plan: task_node_plan.clone(),
1772 phase: scope.phase,
1773 variant_id: scope.variant_id.clone(),
1774 variant: scope.variant.clone(),
1775 fold_id: scope.fold_id.clone(),
1776 branch_path: Vec::new(),
1777 input_handles,
1778 data_views: collected_inputs.data_views,
1779 prediction_inputs: collected_inputs.prediction_inputs,
1780 artifact_inputs,
1781 required_loss_attestations: NodeTask::required_loss_attestations_for(
1782 &task_node_plan,
1783 scope.phase,
1784 )?,
1785 fit_influence,
1786 seed: derive_task_seed(
1787 scope.seed_root,
1788 scope.variant_id.as_ref(),
1789 scope.fold_id.as_ref(),
1790 &task_node_plan,
1791 scope.phase,
1792 ),
1793 },
1794 });
1795 }
1796
1797 for chunk in prepared.chunks(self.max_workers) {
1798 let chunk_results = std::thread::scope(
1799 |thread_scope| -> Result<Vec<NodeResult>> {
1800 let mut handles = Vec::with_capacity(chunk.len());
1801 for prepared_task in chunk {
1802 let controller = controllers
1803 .get(&prepared_task.task.node_plan.controller_id)
1804 .ok_or_else(|| {
1805 DagMlError::RuntimeValidation(format!(
1806 "runtime controller `{}` is not registered",
1807 prepared_task.task.node_plan.controller_id
1808 ))
1809 })?;
1810 let worker_span = phase_span.clone();
1811 handles.push(thread_scope.spawn(move || {
1812 let _worker_span = worker_span.entered();
1813 let _node_span = crate::observability::node_span(
1814 prepared_task.task.run_id.as_str(),
1815 plan_id,
1816 prepared_task.task.phase.as_str(),
1817 prepared_task.task.node_plan.node_id.as_str(),
1818 prepared_task.task.node_plan.controller_id.as_str(),
1819 )
1820 .entered();
1821 let mut result =
1822 if prepared_task.task.node_plan.kind == NodeKind::Tuner {
1823 return Err(DagMlError::RuntimeValidation(format!(
1824 "tuner node `{}` requires execute_hpo_campaign with an explicit RuntimeHpoExecutionContext",
1825 prepared_task.task.node_plan.node_id
1826 )));
1827 } else {
1828 controller.invoke(&prepared_task.task)?
1834 };
1835 record_fit_influence_diagnostic(&prepared_task.task, &mut result);
1836 normalize_result_prediction_ports(
1837 plan,
1838 &prepared_task.task,
1839 &mut result,
1840 )?;
1841 result.validate_for_task(&prepared_task.task)?;
1842 Ok(result)
1843 }));
1844 }
1845 handles
1846 .into_iter()
1847 .map(|handle| {
1848 handle.join().map_err(|_| {
1849 DagMlError::RuntimeValidation(
1850 "parallel scheduler worker panicked".to_string(),
1851 )
1852 })?
1853 })
1854 .collect()
1855 },
1856 )?;
1857
1858 for (prepared_task, mut result) in chunk.iter().zip(chunk_results) {
1859 apply_result_prediction_aggregation(
1860 plan,
1861 controllers,
1862 &prepared_task.task,
1863 &mut result,
1864 &resources,
1865 )?;
1866 if let Some(nested) = resources.nested_stacking.as_ref() {
1867 attach_nested_stacking_input_lineage(
1868 &mut result,
1869 plan,
1870 &prepared_task.task,
1871 ctx,
1872 nested,
1873 )?;
1874 } else {
1875 attach_coordinator_input_lineage(
1876 &mut result,
1877 plan,
1878 &prepared_task.task.node_plan.node_id,
1879 &input_lineage,
1880 )?;
1881 }
1882 if let Some(store) = resources.artifact_store.as_deref_mut() {
1883 if scope.phase == Phase::Refit {
1884 store.capture_refit_artifacts(&prepared_task.task, &result)?;
1885 }
1886 }
1887 for prediction in &result.predictions {
1888 ctx.prediction_store.append(prediction.clone())?;
1889 }
1890 for prediction in &result.aggregated_predictions {
1891 ctx.aggregated_prediction_store.append(prediction.clone())?;
1892 }
1893 apply_result_scoring(
1894 &result,
1895 &mut ctx.score_collector,
1896 &mut ctx.regression_target_records,
1897 )?;
1898 ctx.lineage.record(result.lineage.clone())?;
1899 let data_views = derive_output_data_views(plan, &prepared_task.task, &result)?;
1900 output_handles.insert(prepared_task.node_id.clone(), result.outputs.clone());
1901 output_data_views.insert(prepared_task.node_id.clone(), data_views);
1902 input_lineage.insert(
1903 prepared_task.node_id.clone(),
1904 result.lineage.record_id.clone(),
1905 );
1906 results.push(result);
1907 }
1908 }
1909
1910 for (node_id, reduction) in &merge_nodes {
1915 let node_plan = plan
1916 .node_plans
1917 .get(node_id)
1918 .expect("execution plan was validated");
1919 if let Some(mut result) =
1920 reassemble_branch_merge(plan, node_plan, ctx, &scope, *reduction)?
1921 {
1922 let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1923 let task = NodeTask {
1924 inner_fold_set: None,
1925 run_id: ctx.run_id.clone(),
1926 node_plan: task_node_plan.clone(),
1927 phase: scope.phase,
1928 variant_id: scope.variant_id.clone(),
1929 variant: scope.variant.clone(),
1930 fold_id: scope.fold_id.clone(),
1931 branch_path: Vec::new(),
1932 input_handles: BTreeMap::new(),
1933 data_views: BTreeMap::new(),
1934 prediction_inputs: BTreeMap::new(),
1935 artifact_inputs: BTreeMap::new(),
1936 required_loss_attestations: NodeTask::required_loss_attestations_for(
1937 &task_node_plan,
1938 scope.phase,
1939 )?,
1940 fit_influence: FitInfluenceTask::default(),
1941 seed: None,
1942 };
1943 normalize_result_prediction_ports(plan, &task, &mut result)?;
1944 result.validate_for_task(&task)?;
1945 for prediction in &result.predictions {
1946 ctx.prediction_store.append(prediction.clone())?;
1947 }
1948 apply_result_scoring(
1949 &result,
1950 &mut ctx.score_collector,
1951 &mut ctx.regression_target_records,
1952 )?;
1953 ctx.lineage.record(result.lineage.clone())?;
1954 output_handles.insert(node_id.clone(), result.outputs.clone());
1955 input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1956 results.push(result);
1957 }
1958 }
1959 }
1960
1961 Ok(results)
1962 }
1963}
1964
1965#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1966enum HpoTrialTerminalState {
1967 Completed,
1968 Pruned,
1969 Failed,
1970}
1971
1972fn validate_hpo_checkpoint_result(
1973 checkpoint: &RuntimeHpoCheckpointResult,
1974 hpo: &RuntimeHpoExecutionContext,
1975 trial_variants: &BTreeMap<i64, VariantId>,
1976 terminal_trials: &BTreeMap<i64, HpoTrialTerminalState>,
1977 history_at_start: u32,
1978) -> Result<()> {
1979 checkpoint.artifact.validate().map_err(|error| {
1980 DagMlError::RuntimeValidation(format!(
1981 "runtime HPO checkpoint artifact is invalid: {error}"
1982 ))
1983 })?;
1984 if checkpoint.operation_id != hpo.operation_id
1985 || checkpoint.controller_id != hpo.controller_id
1986 || checkpoint.target_node_id != hpo.target_node_id
1987 || checkpoint.provenance != hpo.provenance
1988 {
1989 return Err(DagMlError::RuntimeValidation(
1990 "runtime HPO checkpoint provenance does not exactly match its execution context"
1991 .to_string(),
1992 ));
1993 }
1994 let proposed_count = u32::try_from(trial_variants.len()).map_err(|_| {
1995 DagMlError::RuntimeValidation(
1996 "runtime HPO scheduler proposal count does not fit u32".to_string(),
1997 )
1998 })?;
1999 if checkpoint.trial_history_len != hpo.trial_budget_total
2000 || checkpoint.trial_history_len < history_at_start
2001 || checkpoint.trial_history_len - history_at_start != proposed_count
2002 {
2003 return Err(DagMlError::RuntimeValidation(
2004 "runtime HPO checkpoint native history is inconsistent with scheduler-observed trials"
2005 .to_string(),
2006 ));
2007 }
2008 if checkpoint.artifact.binding.controller_id != hpo.controller_id.as_str()
2009 || checkpoint.artifact.binding.controller_id != hpo.study.controller_id
2010 || checkpoint.artifact.binding.study_id != hpo.study.study_id
2011 || checkpoint.artifact.methods_abi != hpo.study.methods_abi
2012 {
2013 return Err(DagMlError::RuntimeValidation(
2014 "runtime HPO checkpoint binding/controller/study does not match the active tuner"
2015 .to_string(),
2016 ));
2017 }
2018 let expected_search_space = hpo.study.search_space.fingerprint().map_err(|error| {
2019 DagMlError::RuntimeValidation(format!(
2020 "runtime HPO cannot fingerprint the configured search space: {error}"
2021 ))
2022 })?;
2023 if checkpoint.artifact.binding.search_space_fingerprint != expected_search_space {
2024 return Err(DagMlError::RuntimeValidation(
2025 "runtime HPO checkpoint search-space binding does not match the active study"
2026 .to_string(),
2027 ));
2028 }
2029
2030 let completed_trial_ids = terminal_trials
2031 .iter()
2032 .filter_map(|(trial_id, state)| {
2033 (*state == HpoTrialTerminalState::Completed).then_some(*trial_id)
2034 })
2035 .collect::<BTreeSet<_>>();
2036 let mut proposal_trial_ids = BTreeSet::new();
2037 for proposal in &checkpoint.completed_proposals {
2038 if !proposal_trial_ids.insert(proposal.trial_id) {
2039 return Err(DagMlError::RuntimeValidation(format!(
2040 "runtime HPO checkpoint has duplicate completed proposal for trial `{}`",
2041 proposal.trial_id
2042 )));
2043 }
2044 if trial_variants.get(&proposal.trial_id) != Some(&proposal.variant.variant_id) {
2045 return Err(DagMlError::RuntimeValidation(format!(
2046 "runtime HPO checkpoint proposal for trial `{}` does not exactly match its scheduler proposal",
2047 proposal.trial_id
2048 )));
2049 }
2050 }
2051 if proposal_trial_ids != completed_trial_ids {
2052 return Err(DagMlError::RuntimeValidation(
2053 "runtime HPO checkpoint proposals must cover exactly the completed trials".to_string(),
2054 ));
2055 }
2056
2057 let mut report_trial_ids = BTreeSet::new();
2058 for completed in &checkpoint.completed_reports {
2059 if !report_trial_ids.insert(completed.trial_id) {
2060 return Err(DagMlError::RuntimeValidation(format!(
2061 "runtime HPO checkpoint has duplicate completed report for trial `{}`",
2062 completed.trial_id
2063 )));
2064 }
2065 if trial_variants.get(&completed.trial_id) != Some(&completed.variant_id)
2066 || !proposal_trial_ids.contains(&completed.trial_id)
2067 {
2068 return Err(DagMlError::RuntimeValidation(format!(
2069 "runtime HPO checkpoint report for trial `{}` does not match a completed proposal",
2070 completed.trial_id
2071 )));
2072 }
2073 let report = &completed.report;
2074 if report.producer_node != hpo.selection.producer_node
2075 || report.producer_port.as_deref() != Some(hpo.selection.producer_port.as_str())
2076 || report.partition != PredictionPartition::Validation
2077 || report
2078 .fold_id
2079 .as_ref()
2080 .is_none_or(|fold| fold.as_str() != "avg")
2081 || report.variant_id.as_ref() != Some(&completed.variant_id)
2082 || !report
2083 .metrics
2084 .get(hpo.selection.metric.name())
2085 .is_some_and(|score| score.is_finite())
2086 {
2087 return Err(DagMlError::RuntimeValidation(format!(
2088 "runtime HPO checkpoint report for trial `{}` is not its one finite target OOF average",
2089 completed.trial_id
2090 )));
2091 }
2092 }
2093 if report_trial_ids != completed_trial_ids {
2094 return Err(DagMlError::RuntimeValidation(
2095 "runtime HPO checkpoint reports must cover exactly one OOF average per completed trial"
2096 .to_string(),
2097 ));
2098 }
2099 Ok(())
2100}
2101
2102pub(crate) struct PreparedNodeTask {
2103 pub(crate) node_id: NodeId,
2104 pub(crate) task: NodeTask,
2105}
2106
2107#[cfg(test)]
2110#[allow(clippy::items_after_test_module)]
2111mod hpo_scheduler_tests {
2112 use std::collections::{BTreeMap, BTreeSet};
2113 use std::sync::{Arc, Mutex};
2114
2115 use sha2::{Digest, Sha256};
2116
2117 use super::*;
2118 use crate::controller::{
2119 ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest,
2120 ControllerRegistry, RngPolicy,
2121 };
2122 use crate::data::InMemoryDataProvider;
2123 use crate::fold::{FoldAssignment, FoldPartitionMode, KFoldSpec, NestedCvSpec};
2124 use crate::graph::{
2125 EdgeContract, EdgeSpec, GraphInterface, GraphSpec, NodeSpec, PortRef, PortSchema, PortSpec,
2126 };
2127 use crate::hpo::{
2128 HpoDirection, HpoMetric, HpoOptimizerConfig, HpoParameter, HpoPruner, HpoSampler,
2129 HpoSearchSpace, HpoStudyBinding, MethodsHpoStudyConfig, N4moptCheckpointArtifact,
2130 N4MOPT_ARTIFACT_KIND, N4MOPT_CHECKPOINT_SCHEMA_VERSION, N4MOPT_FORMAT,
2131 };
2132 #[cfg(feature = "methods-optimizer-local")]
2133 use crate::hpo::{MethodsHpoController, MethodsRuntime};
2134 use crate::metrics::RegressionTargetBlock;
2135 use crate::oof::PredictionBlock;
2136 use crate::plan::{build_execution_plan, SplitInvocation};
2137
2138 struct HpoTestModel {
2139 id: ControllerId,
2140 trace: Arc<Mutex<Vec<String>>>,
2141 fail_variant: Option<VariantId>,
2142 score_by_trial: bool,
2143 }
2144
2145 impl RuntimeController for HpoTestModel {
2146 fn controller_id(&self) -> &ControllerId {
2147 &self.id
2148 }
2149
2150 fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
2151 self.trace.lock().unwrap().push("model_cv".to_string());
2152 if task.variant_id.as_ref() == self.fail_variant.as_ref() {
2153 return Err(DagMlError::RuntimeValidation(
2154 "controlled HPO fold failure".to_string(),
2155 ));
2156 }
2157 let sample_id = match task.fold_id.as_ref().map(FoldId::as_str) {
2158 Some("fold:0") => SampleId::new("sample:one").unwrap(),
2159 Some("fold:1") => SampleId::new("sample:two").unwrap(),
2160 other => {
2161 return Err(DagMlError::RuntimeValidation(format!(
2162 "HPO test model received unexpected fold {other:?}"
2163 )));
2164 }
2165 };
2166 let prediction = if self.score_by_trial {
2167 let trial_id = task
2168 .variant_id
2169 .as_ref()
2170 .and_then(|variant| variant.as_str().strip_prefix("hpo:trial:"))
2171 .and_then(|value| value.parse::<i64>().ok())
2172 .ok_or_else(|| {
2173 DagMlError::RuntimeValidation(
2174 "native HPO test model received no trial variant".to_string(),
2175 )
2176 })?;
2177 let magnitude = (trial_id + 1) as f64;
2178 1.0 + magnitude * magnitude
2179 } else {
2180 1.0
2181 };
2182 Ok(NodeResult {
2183 schema_version: None,
2184 node_id: task.node_plan.node_id.clone(),
2185 outputs: BTreeMap::from([(
2186 "prediction".to_string(),
2187 HandleRef {
2188 handle: 2,
2189 kind: HandleKind::Prediction,
2190 owner_controller: self.id.clone(),
2191 },
2192 )]),
2193 predictions: vec![PredictionBlock {
2194 prediction_id: Some(format!("prediction:{}", task.fold_id.as_ref().unwrap())),
2195 producer_node: task.node_plan.node_id.clone(),
2196 producer_port: None,
2197 partition: PredictionPartition::Validation,
2198 fold_id: task.fold_id.clone(),
2199 sample_ids: vec![sample_id.clone()],
2200 values: vec![vec![prediction]],
2201 target_names: vec!["target".to_string()],
2202 }],
2203 observation_predictions: Vec::new(),
2204 aggregated_predictions: Vec::new(),
2205 explanations: Vec::new(),
2206 shape_deltas: Vec::new(),
2207 artifacts: Vec::new(),
2208 artifact_handles: BTreeMap::new(),
2209 fit_influence_diagnostics: Vec::new(),
2210 regression_targets: vec![RegressionTargetBlock {
2211 level: PredictionLevel::Sample,
2212 unit_ids: vec![PredictionUnitId::Sample(sample_id)],
2213 values: vec![vec![1.0]],
2214 target_names: vec!["target".to_string()],
2215 }],
2216 lineage: LineageRecord {
2217 record_id: LineageId::new(format!(
2218 "lineage:hpo-model:{}",
2219 task.fold_id.as_ref().unwrap()
2220 ))
2221 .unwrap(),
2222 run_id: task.run_id.clone(),
2223 node_id: task.node_plan.node_id.clone(),
2224 phase: task.phase,
2225 controller_id: self.id.clone(),
2226 controller_version: task.node_plan.controller_version.clone(),
2227 variant_id: task.variant_id.clone(),
2228 fold_id: task.fold_id.clone(),
2229 branch_path: Vec::new(),
2230 input_lineage: Vec::new(),
2231 artifact_refs: Vec::new(),
2232 params_fingerprint: task.node_plan.params_fingerprint.clone(),
2233 data_model_shape_fingerprint: None,
2234 aggregation_policy_fingerprint: None,
2235 seed: task.seed,
2236 unsafe_flags: BTreeSet::new(),
2237 metrics: BTreeMap::new(),
2238 loss_attestations: Vec::new(),
2239 early_stopping_records: Vec::new(),
2240 },
2241 })
2242 }
2243 }
2244
2245 struct HpoTestTuner {
2246 id: ControllerId,
2247 trace: Arc<Mutex<Vec<String>>>,
2248 history_len: u32,
2249 proposal_count: u32,
2250 prune_at: Option<(i64, i32)>,
2251 }
2252
2253 struct HpoTestSession {
2254 proposals: Vec<RuntimeHpoProposal>,
2255 trace: Arc<Mutex<Vec<String>>>,
2256 checkpoint: N4moptCheckpointArtifact,
2257 history_len: u32,
2258 completed: BTreeMap<i64, f64>,
2259 failed: BTreeMap<i64, RuntimeHpoFailure>,
2260 prune_at: Option<(i64, i32)>,
2261 pruned: BTreeSet<i64>,
2262 intermediates: BTreeMap<i64, Vec<crate::hpo::HpoIntermediate>>,
2263 }
2264
2265 impl RuntimeController for HpoTestTuner {
2266 fn controller_id(&self) -> &ControllerId {
2267 &self.id
2268 }
2269
2270 fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
2271 Err(DagMlError::RuntimeValidation(format!(
2272 "HPO test tuner `{}` was dispatched through generic invoke",
2273 task.node_plan.node_id
2274 )))
2275 }
2276
2277 fn create_tuner_session(
2278 &self,
2279 task: &RuntimeHpoCampaignTask,
2280 context: &RuntimeHpoExecutionContext,
2281 ) -> Result<Box<dyn RuntimeTunerSession>> {
2282 assert_eq!(task.operation_id, context.operation_id);
2283 self.trace
2284 .lock()
2285 .unwrap()
2286 .push("session_factory".to_string());
2287 let payload = vec![7_u8];
2288 let proposals = (0..self.proposal_count)
2289 .map(|offset| {
2290 let trial_id = i64::from(self.history_len + offset + 1);
2291 let mut variant = context.base_variant.clone();
2292 if self.history_len != 0 || self.proposal_count != 1 {
2293 variant.variant_id = VariantId::new(format!("hpo:trial:{trial_id}"))
2294 .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?;
2295 variant.fingerprint = format!("hpo-test-{trial_id}");
2296 }
2297 Ok(RuntimeHpoProposal { trial_id, variant })
2298 })
2299 .collect::<Result<Vec<_>>>()?;
2300 Ok(Box::new(HpoTestSession {
2301 proposals: proposals.into_iter().rev().collect(),
2302 trace: Arc::clone(&self.trace),
2303 history_len: self.history_len,
2304 completed: BTreeMap::new(),
2305 failed: BTreeMap::new(),
2306 prune_at: self.prune_at,
2307 pruned: BTreeSet::new(),
2308 intermediates: BTreeMap::new(),
2309 checkpoint: N4moptCheckpointArtifact {
2310 schema_version: N4MOPT_CHECKPOINT_SCHEMA_VERSION,
2311 artifact_kind: N4MOPT_ARTIFACT_KIND.to_string(),
2312 format: N4MOPT_FORMAT.to_string(),
2313 abi_major: crate::hpo::METHODS_ABI_MAJOR,
2314 abi_min_minor: crate::hpo::METHODS_N4MOPT_MIN_ABI_MINOR,
2315 binding: HpoStudyBinding {
2316 controller_id: context.study.controller_id.clone(),
2317 study_id: context.study.study_id.clone(),
2318 search_space_fingerprint: context
2319 .study
2320 .search_space
2321 .fingerprint()
2322 .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?,
2323 optimizer_fingerprint: "optimizer:test".to_string(),
2324 },
2325 methods_abi: context.study.methods_abi.clone(),
2326 payload_sha256: format!("{:x}", Sha256::digest(&payload)),
2327 opaque_payload: payload,
2328 },
2329 }))
2330 }
2331 }
2332
2333 impl RuntimeTunerSession for HpoTestSession {
2334 fn trial_history_len(&self) -> Result<u32> {
2335 Ok(self.history_len)
2336 }
2337
2338 fn ask(&mut self) -> Result<Option<RuntimeHpoProposal>> {
2339 self.trace.lock().unwrap().push("ask".to_string());
2340 let proposal = self.proposals.pop();
2341 if proposal.is_some() {
2342 self.history_len += 1;
2343 }
2344 Ok(proposal)
2345 }
2346
2347 fn report_intermediate(
2348 &mut self,
2349 intermediate: RuntimeHpoIntermediate,
2350 ) -> Result<RuntimeHpoIntermediateOutcome> {
2351 assert!(intermediate.score.is_finite());
2352 let should_prune = self.prune_at == Some((intermediate.trial_id, intermediate.step));
2353 self.trace.lock().unwrap().push(format!(
2354 "intermediate:{}:{}",
2355 intermediate.trial_id, intermediate.step
2356 ));
2357 self.intermediates
2358 .entry(intermediate.trial_id)
2359 .or_default()
2360 .push(crate::hpo::HpoIntermediate {
2361 sequence: i64::from(intermediate.step) + 1,
2362 step: intermediate.step,
2363 score: intermediate.score,
2364 should_prune,
2365 });
2366 if should_prune {
2367 self.pruned.insert(intermediate.trial_id);
2368 Ok(RuntimeHpoIntermediateOutcome::Pruned)
2369 } else {
2370 Ok(RuntimeHpoIntermediateOutcome::Continue)
2371 }
2372 }
2373
2374 fn tell(&mut self, trial_id: i64, terminal: RuntimeHpoTerminal) -> Result<()> {
2375 assert!(trial_id > 0);
2376 match terminal {
2377 RuntimeHpoTerminal::Completed { score } if score.is_finite() => {
2378 self.trace.lock().unwrap().push("tell".to_string());
2379 self.completed.insert(trial_id, score);
2380 }
2381 RuntimeHpoTerminal::Failed { failure } => {
2382 self.trace.lock().unwrap().push("tell_failed".to_string());
2383 self.failed.insert(trial_id, failure);
2384 }
2385 other => {
2386 return Err(DagMlError::RuntimeValidation(format!(
2387 "test HPO session received unsupported terminal state {other:?}"
2388 )));
2389 }
2390 }
2391 Ok(())
2392 }
2393
2394 fn checkpoint(&self) -> Result<N4moptCheckpointArtifact> {
2395 self.trace.lock().unwrap().push("checkpoint".to_string());
2396 Ok(self.checkpoint.clone())
2397 }
2398
2399 fn incumbent(
2400 &self,
2401 variants: &BTreeMap<i64, VariantId>,
2402 ) -> Result<Option<RuntimeHpoIncumbent>> {
2403 let Some((&trial_id, &score)) = self
2404 .completed
2405 .iter()
2406 .min_by(|left, right| left.1.total_cmp(right.1).then_with(|| left.0.cmp(right.0)))
2407 else {
2408 return Ok(None);
2409 };
2410 Ok(Some(RuntimeHpoIncumbent {
2411 trial_id,
2412 score,
2413 metric: "rmse".to_string(),
2414 direction: HpoDirection::Minimize,
2415 variant_id: variants.get(&trial_id).cloned().unwrap(),
2416 }))
2417 }
2418
2419 fn terminal_trial_snapshots(
2420 &self,
2421 variants: &BTreeMap<i64, VariantId>,
2422 ) -> Result<Vec<RuntimeHpoTerminalSnapshot>> {
2423 if self.completed.is_empty() {
2424 return Err(DagMlError::RuntimeValidation(
2425 "test HPO session has no completed trial".to_string(),
2426 ));
2427 }
2428 let completed = &self.completed;
2429 let failed = &self.failed;
2430 Ok((1..=i64::from(self.history_len))
2431 .map(|id| {
2432 let score = completed.get(&id).copied();
2433 let is_completed = score.is_some();
2434 let pruned = self.pruned.contains(&id);
2435 let failure = failed
2436 .get(&id)
2437 .map(|failure| crate::hpo::HpoFailure {
2438 code: failure.code.clone(),
2439 message: failure.message.clone(),
2440 retryable: failure.retryable,
2441 })
2442 .or_else(|| {
2443 (!is_completed && !pruned).then(|| crate::hpo::HpoFailure {
2444 code: "RESTORED_TEST_FAILURE".to_string(),
2445 message: "synthetic restored terminal".to_string(),
2446 retryable: false,
2447 })
2448 });
2449 RuntimeHpoTerminalSnapshot {
2450 trial: crate::hpo::HpoTrial {
2451 id,
2452 ask_sequence: id,
2453 terminal_sequence: Some(id),
2454 parameters: BTreeMap::new(),
2455 parameter_order: Vec::new(),
2456 status: if is_completed {
2457 crate::hpo::HpoTrialStatus::Completed
2458 } else if pruned {
2459 crate::hpo::HpoTrialStatus::Pruned
2460 } else {
2461 crate::hpo::HpoTrialStatus::Failed
2462 },
2463 score,
2464 rung: 0,
2465 duration: 0.0,
2466 intermediates: self.intermediates.get(&id).cloned().unwrap_or_default(),
2467 failure,
2468 },
2469 variant_id: variants.get(&id).cloned(),
2470 }
2471 })
2472 .collect())
2473 }
2474 }
2475
2476 fn node(id: &str, kind: NodeKind, outputs: Vec<PortSpec>) -> NodeSpec {
2477 NodeSpec {
2478 id: NodeId::new(id).unwrap(),
2479 kind,
2480 operator: None,
2481 params: BTreeMap::new(),
2482 ports: PortSchema {
2483 inputs: Vec::new(),
2484 outputs,
2485 },
2486 metadata: BTreeMap::new(),
2487 seed_label: None,
2488 }
2489 }
2490
2491 fn prediction_port(name: &str) -> PortSpec {
2492 PortSpec {
2493 name: name.to_string(),
2494 kind: PortKind::Prediction,
2495 representation: None,
2496 cardinality: crate::graph::PortCardinality::One,
2497 unit_level: None,
2498 alignment_key: None,
2499 target_level: None,
2500 description: String::new(),
2501 }
2502 }
2503
2504 fn manifest(id: &str, kind: NodeKind) -> ControllerManifest {
2505 ControllerManifest {
2506 controller_id: ControllerId::new(id).unwrap(),
2507 controller_version: "test".to_string(),
2508 operator_kind: kind,
2509 priority: 0,
2510 supported_phases: BTreeSet::from([Phase::FitCv]),
2511 input_ports: Vec::new(),
2512 output_ports: Vec::new(),
2513 data_requirements: None,
2514 capabilities: BTreeSet::from([
2515 ControllerCapability::Deterministic,
2516 ControllerCapability::EmitsPredictions,
2517 ]),
2518 operator_selectors: Vec::new(),
2519 fit_scope: ControllerFitScope::FoldTrain,
2520 rng_policy: RngPolicy::UsesCoreSeed,
2521 artifact_policy: ArtifactPolicy::Serializable,
2522 }
2523 }
2524
2525 #[test]
2526 fn hpo_campaign_invokes_registered_session_and_routes_oof_feedback() {
2527 let target = NodeId::new("model:score").unwrap();
2528 let graph = GraphSpec {
2529 id: "graph:hpo.scheduler".to_string(),
2530 interface: GraphInterface::default(),
2531 nodes: vec![node(
2532 "model:score",
2533 NodeKind::Model,
2534 vec![PortSpec {
2535 name: "prediction".to_string(),
2536 kind: PortKind::Prediction,
2537 representation: None,
2538 cardinality: crate::graph::PortCardinality::One,
2539 unit_level: None,
2540 alignment_key: None,
2541 target_level: None,
2542 description: String::new(),
2543 }],
2544 )],
2545 edges: Vec::new(),
2546 search_space_fingerprint: None,
2547 metadata: BTreeMap::new(),
2548 };
2549 let fold_set = FoldSet {
2550 id: "folds:hpo".to_string(),
2551 sample_ids: vec![
2552 SampleId::new("sample:one").unwrap(),
2553 SampleId::new("sample:two").unwrap(),
2554 ],
2555 folds: vec![
2556 FoldAssignment {
2557 fold_id: FoldId::new("fold:0").unwrap(),
2558 train_sample_ids: vec![SampleId::new("sample:two").unwrap()],
2559 validation_sample_ids: vec![SampleId::new("sample:one").unwrap()],
2560 metadata: BTreeMap::new(),
2561 },
2562 FoldAssignment {
2563 fold_id: FoldId::new("fold:1").unwrap(),
2564 train_sample_ids: vec![SampleId::new("sample:one").unwrap()],
2565 validation_sample_ids: vec![SampleId::new("sample:two").unwrap()],
2566 metadata: BTreeMap::new(),
2567 },
2568 ],
2569 sample_groups: BTreeMap::new(),
2570 partition_mode: FoldPartitionMode::Partition,
2571 };
2572 let mut registry = ControllerRegistry::new();
2573 registry
2574 .register(manifest("controller:model", NodeKind::Model))
2575 .unwrap();
2576 let plan = build_execution_plan(
2577 "plan:hpo.scheduler",
2578 graph,
2579 CampaignSpec {
2580 inner_cv: None,
2581 id: "campaign:hpo.scheduler".to_string(),
2582 root_seed: Some(13),
2583 leakage_policy: Default::default(),
2584 aggregation_policy: Default::default(),
2585 split_invocation: Some(SplitInvocation {
2586 id: "split:hpo".to_string(),
2587 controller_id: None,
2588 leakage_policy: Default::default(),
2589 params: BTreeMap::new(),
2590 fold_set: Some(fold_set),
2591 }),
2592 generation: Default::default(),
2593 shape_plans: BTreeMap::new(),
2594 data_bindings: BTreeMap::new(),
2595 branch_view_plans: Vec::new(),
2596 metadata: BTreeMap::new(),
2597 },
2598 ®istry,
2599 )
2600 .unwrap();
2601 let trace = Arc::new(Mutex::new(Vec::new()));
2602 let mut controllers = RuntimeControllerRegistry::new();
2603 controllers
2604 .register(Box::new(HpoTestTuner {
2605 id: ControllerId::new("controller:tuner").unwrap(),
2606 trace: Arc::clone(&trace),
2607 history_len: 0,
2608 proposal_count: 1,
2609 prune_at: None,
2610 }))
2611 .unwrap();
2612 controllers
2613 .register(Box::new(HpoTestModel {
2614 id: ControllerId::new("controller:model").unwrap(),
2615 trace: Arc::clone(&trace),
2616 fail_variant: None,
2617 score_by_trial: false,
2618 }))
2619 .unwrap();
2620 let hpo = RuntimeHpoExecutionContext {
2621 operation_id: "hpo:test".to_string(),
2622 controller_id: ControllerId::new("controller:tuner").unwrap(),
2623 target_node_id: target.clone(),
2624 base_variant: plan.variants[0].clone(),
2625 trial_budget_total: 1,
2626 study: MethodsHpoStudyConfig {
2627 controller_id: "controller:tuner".to_string(),
2628 study_id: "study:hpo.scheduler".to_string(),
2629 methods_abi: "test-abi".to_string(),
2630 search_space: HpoSearchSpace {
2631 parameters: vec![HpoParameter::Int {
2632 name: "n_components".to_string(),
2633 low: 1,
2634 high: 1,
2635 step: 1,
2636 log: false,
2637 }],
2638 },
2639 optimizer: HpoOptimizerConfig {
2640 sampler: HpoSampler::Random,
2641 pruner: HpoPruner::None,
2642 direction: HpoDirection::Minimize,
2643 metric: HpoMetric::Rmse,
2644 seed: 13,
2645 n_startup_trials: 1,
2646 max_resource: 0,
2647 reduction_factor: 1,
2648 },
2649 },
2650 parameter_paths: BTreeMap::from([(
2651 "n_components".to_string(),
2652 "n_components".to_string(),
2653 )]),
2654 resume_checkpoint: None,
2655 resume_variants: BTreeMap::new(),
2656 resume_terminal_trials: Vec::new(),
2657 selection: RuntimeHpoSelectionTarget {
2658 producer_node: target,
2659 producer_port: "prediction".to_string(),
2660 metric: RegressionMetricKind::Rmse,
2661 direction: HpoDirection::Minimize,
2662 },
2663 provenance: RuntimeHpoProvenance {
2664 graph_fingerprint: plan.graph_fingerprint.clone(),
2665 campaign_fingerprint: plan.campaign_fingerprint.clone(),
2666 controller_fingerprint: plan.controller_fingerprint.clone(),
2667 data_identities_fingerprint: "identity:test".to_string(),
2668 fold_set_fingerprint: plan
2669 .fold_set
2670 .as_ref()
2671 .map(stable_json_fingerprint)
2672 .transpose()
2673 .unwrap(),
2674 training_influence_fingerprint: "influence:test".to_string(),
2675 relation_fingerprint: "relation:test".to_string(),
2676 },
2677 };
2678 let provider = InMemoryDataProvider::new(ControllerId::new("controller:data").unwrap());
2679 let ctx = RunContext::new(RunId::new("run:hpo.scheduler").unwrap(), Some(13));
2680
2681 let result = SequentialScheduler
2682 .execute_hpo_campaign(&plan, &controllers, &provider, &ctx, &hpo)
2683 .unwrap();
2684
2685 assert_eq!(result.operation_id, "hpo:test");
2686 assert_eq!(result.candidates.len(), 1);
2687 assert_eq!(result.checkpoint.completed_proposals.len(), 1);
2688 assert_eq!(result.checkpoint.completed_reports.len(), 1);
2689 assert_eq!(result.candidates[0].lineage.len(), 2);
2690 assert_eq!(result.incumbent.variant_id, plan.variants[0].variant_id);
2691 let mut selected_ctx = RunContext::new(RunId::new("run:hpo.scheduler").unwrap(), Some(13));
2692 selected_ctx.variant_id = Some(plan.variants[0].variant_id.clone());
2693 let selected_results = SequentialScheduler
2694 .execute_campaign_phase_with_data_provider(
2695 &plan,
2696 &controllers,
2697 &provider,
2698 &mut selected_ctx,
2699 Phase::FitCv,
2700 )
2701 .unwrap();
2702 assert_eq!(selected_results.len(), 2);
2703 assert_eq!(selected_ctx.lineage.len(), 2);
2704 assert_eq!(
2705 trace.lock().unwrap().as_slice(),
2706 [
2707 "session_factory",
2708 "ask",
2709 "model_cv",
2710 "intermediate:1:0",
2711 "model_cv",
2712 "intermediate:1:1",
2713 "tell",
2714 "checkpoint",
2715 "model_cv",
2716 "model_cv"
2717 ]
2718 );
2719
2720 let prune_trace = Arc::new(Mutex::new(Vec::new()));
2725 let mut prune_controllers = RuntimeControllerRegistry::new();
2726 prune_controllers
2727 .register(Box::new(HpoTestTuner {
2728 id: ControllerId::new("controller:tuner").unwrap(),
2729 trace: Arc::clone(&prune_trace),
2730 history_len: 0,
2731 proposal_count: 2,
2732 prune_at: Some((2, 0)),
2733 }))
2734 .unwrap();
2735 prune_controllers
2736 .register(Box::new(HpoTestModel {
2737 id: ControllerId::new("controller:model").unwrap(),
2738 trace: Arc::clone(&prune_trace),
2739 fail_variant: None,
2740 score_by_trial: false,
2741 }))
2742 .unwrap();
2743 let mut pruning_hpo = hpo.clone();
2744 pruning_hpo.trial_budget_total = 2;
2745 let pruned = SequentialScheduler
2746 .execute_hpo_campaign(
2747 &plan,
2748 &prune_controllers,
2749 &provider,
2750 &RunContext::new(RunId::new("run:hpo.pruned").unwrap(), Some(13)),
2751 &pruning_hpo,
2752 )
2753 .unwrap();
2754 assert_eq!(pruned.candidates.len(), 1);
2755 assert_eq!(pruned.checkpoint.completed_proposals.len(), 1);
2756 assert_eq!(pruned.checkpoint.completed_reports.len(), 1);
2757 assert_eq!(pruned.incumbent.trial_id, 1);
2758 assert_eq!(
2759 pruned
2760 .terminal_trials
2761 .iter()
2762 .map(|entry| (entry.trial.id, entry.trial.status))
2763 .collect::<Vec<_>>(),
2764 vec![
2765 (1, crate::hpo::HpoTrialStatus::Completed),
2766 (2, crate::hpo::HpoTrialStatus::Pruned),
2767 ]
2768 );
2769 assert_eq!(
2770 prune_trace
2771 .lock()
2772 .unwrap()
2773 .iter()
2774 .filter(|event| event.as_str() == "model_cv")
2775 .count(),
2776 3,
2777 "the pruned second trial must execute one fold, not both"
2778 );
2779
2780 let failed_trace = Arc::new(Mutex::new(Vec::new()));
2785 let mut failed_controllers = RuntimeControllerRegistry::new();
2786 failed_controllers
2787 .register(Box::new(HpoTestTuner {
2788 id: ControllerId::new("controller:tuner").unwrap(),
2789 trace: Arc::clone(&failed_trace),
2790 history_len: 0,
2791 proposal_count: 2,
2792 prune_at: None,
2793 }))
2794 .unwrap();
2795 failed_controllers
2796 .register(Box::new(HpoTestModel {
2797 id: ControllerId::new("controller:model").unwrap(),
2798 trace: Arc::clone(&failed_trace),
2799 fail_variant: Some(VariantId::new("hpo:trial:2").unwrap()),
2800 score_by_trial: false,
2801 }))
2802 .unwrap();
2803 let failed = SequentialScheduler
2804 .execute_hpo_campaign(
2805 &plan,
2806 &failed_controllers,
2807 &provider,
2808 &RunContext::new(RunId::new("run:hpo.failed").unwrap(), Some(13)),
2809 &pruning_hpo,
2810 )
2811 .unwrap();
2812 assert_eq!(failed.candidates.len(), 1);
2813 assert_eq!(
2814 failed.candidates[0]
2815 .validation_predictions
2816 .predictions
2817 .len(),
2818 2
2819 );
2820 assert_eq!(failed.checkpoint.completed_proposals.len(), 1);
2821 assert_eq!(failed.checkpoint.completed_reports.len(), 1);
2822 assert_eq!(
2823 failed
2824 .terminal_trials
2825 .iter()
2826 .map(|entry| (entry.trial.id, entry.trial.status))
2827 .collect::<Vec<_>>(),
2828 vec![
2829 (1, crate::hpo::HpoTrialStatus::Completed),
2830 (2, crate::hpo::HpoTrialStatus::Failed),
2831 ]
2832 );
2833 let failed_trial = &failed.terminal_trials[1].trial;
2834 assert_eq!(
2835 failed_trial.failure.as_ref().unwrap().code,
2836 "DAGML_CV_ERROR"
2837 );
2838 assert!(failed_trial.intermediates.is_empty());
2839 assert_eq!(
2840 failed_trace
2841 .lock()
2842 .unwrap()
2843 .iter()
2844 .filter(|event| event.as_str() == "model_cv")
2845 .count(),
2846 3,
2847 "the failed second trial must not schedule another fold or output handle"
2848 );
2849 assert!(failed_trace
2850 .lock()
2851 .unwrap()
2852 .iter()
2853 .any(|event| event == "tell_failed"));
2854
2855 let mut resampled_plan = plan.clone();
2858 resampled_plan.fold_set.as_mut().unwrap().partition_mode = FoldPartitionMode::Resampled;
2859 let mut resampled_hpo = hpo.clone();
2860 resampled_hpo.provenance.fold_set_fingerprint =
2861 Some(stable_json_fingerprint(resampled_plan.fold_set.as_ref().unwrap()).unwrap());
2862 let trace_len_before = trace.lock().unwrap().len();
2863 let error = SequentialScheduler
2864 .execute_hpo_campaign(
2865 &resampled_plan,
2866 &controllers,
2867 &provider,
2868 &ctx,
2869 &resampled_hpo,
2870 )
2871 .unwrap_err();
2872 assert!(error
2873 .to_string()
2874 .contains("requires FoldPartitionMode::Partition"));
2875 assert_eq!(trace.lock().unwrap().len(), trace_len_before);
2876
2877 let mut empty_plan = plan.clone();
2880 empty_plan.fold_set.as_mut().unwrap().folds.clear();
2881 let trace_len_before = trace.lock().unwrap().len();
2882 let error = SequentialScheduler
2883 .execute_hpo_campaign(&empty_plan, &controllers, &provider, &ctx, &hpo)
2884 .unwrap_err();
2885 assert!(error.to_string().contains("fold set contains no folds"));
2886 assert_eq!(trace.lock().unwrap().len(), trace_len_before);
2887
2888 let mut no_fold_plan = plan.clone();
2891 no_fold_plan.fold_set = None;
2892 let mut no_fold_hpo = hpo.clone();
2893 no_fold_hpo.provenance.fold_set_fingerprint = None;
2894 let trace_len_before = trace.lock().unwrap().len();
2895 let error = SequentialScheduler
2896 .execute_hpo_campaign(&no_fold_plan, &controllers, &provider, &ctx, &no_fold_hpo)
2897 .unwrap_err();
2898 assert!(error
2899 .to_string()
2900 .contains("requires an explicit validated fold set"));
2901 assert_eq!(trace.lock().unwrap().len(), trace_len_before);
2902
2903 let samples = (1..=6)
2907 .map(|index| SampleId::new(format!("nested:{index}")).unwrap())
2908 .collect::<Vec<_>>();
2909 let outer_folds = KFoldSpec {
2910 n_splits: 3,
2911 shuffle: false,
2912 seed: Some(7),
2913 }
2914 .split("folds:hpo.nested", &samples)
2915 .unwrap();
2916 let base_a = NodeId::new("model:nested.base.a").unwrap();
2917 let base_b = NodeId::new("model:nested.base.b").unwrap();
2918 let meta = NodeId::new("model:nested.meta").unwrap();
2919 let mut meta_node = NodeSpec {
2920 id: meta.clone(),
2921 kind: NodeKind::Model,
2922 operator: None,
2923 params: BTreeMap::new(),
2924 ports: PortSchema {
2925 inputs: vec![prediction_port("a"), prediction_port("b")],
2926 outputs: vec![prediction_port("prediction")],
2927 },
2928 metadata: BTreeMap::new(),
2929 seed_label: None,
2930 };
2931 meta_node.metadata.insert(
2932 NESTED_STACKING_EXECUTION_METADATA_KEY.to_string(),
2933 serde_json::json!(NESTED_STACKING_EXECUTION_V1),
2934 );
2935 let nested_graph = GraphSpec {
2936 id: "graph:hpo.nested".to_string(),
2937 interface: GraphInterface::default(),
2938 nodes: vec![
2939 node(
2940 base_a.as_str(),
2941 NodeKind::Model,
2942 vec![prediction_port("prediction")],
2943 ),
2944 node(
2945 base_b.as_str(),
2946 NodeKind::Model,
2947 vec![prediction_port("prediction")],
2948 ),
2949 meta_node,
2950 ],
2951 edges: vec![
2952 EdgeSpec {
2953 source: PortRef {
2954 node_id: base_a,
2955 port_name: "prediction".to_string(),
2956 },
2957 target: PortRef {
2958 node_id: meta.clone(),
2959 port_name: "a".to_string(),
2960 },
2961 contract: EdgeContract {
2962 requires_oof: true,
2963 requires_fold_alignment: true,
2964 ..EdgeContract::new(PortKind::Prediction, None)
2965 },
2966 },
2967 EdgeSpec {
2968 source: PortRef {
2969 node_id: base_b,
2970 port_name: "prediction".to_string(),
2971 },
2972 target: PortRef {
2973 node_id: meta.clone(),
2974 port_name: "b".to_string(),
2975 },
2976 contract: EdgeContract {
2977 requires_oof: true,
2978 requires_fold_alignment: true,
2979 ..EdgeContract::new(PortKind::Prediction, None)
2980 },
2981 },
2982 ],
2983 search_space_fingerprint: None,
2984 metadata: BTreeMap::new(),
2985 };
2986 let mut nested_campaign = plan.campaign.clone();
2987 nested_campaign.id = "campaign:hpo.nested".to_string();
2988 nested_campaign.inner_cv = Some(NestedCvSpec::KFold(KFoldSpec {
2989 n_splits: 2,
2990 shuffle: false,
2991 seed: Some(11),
2992 }));
2993 nested_campaign.split_invocation.as_mut().unwrap().fold_set = Some(outer_folds);
2994 let mut nested_registry = ControllerRegistry::new();
2995 let mut nested_model_manifest = manifest("controller:model", NodeKind::Model);
2996 nested_model_manifest
2997 .capabilities
2998 .insert(ControllerCapability::ConsumesOofPredictions);
2999 nested_registry.register(nested_model_manifest).unwrap();
3000 let nested_plan = build_execution_plan(
3001 "plan:hpo.nested",
3002 nested_graph,
3003 nested_campaign,
3004 &nested_registry,
3005 )
3006 .unwrap();
3007 let mut nested_hpo = hpo.clone();
3008 nested_hpo.target_node_id = meta.clone();
3009 nested_hpo.base_variant = nested_plan.variants[0].clone();
3010 nested_hpo.selection.producer_node = meta;
3011 nested_hpo.provenance.graph_fingerprint = nested_plan.graph_fingerprint.clone();
3012 nested_hpo.provenance.campaign_fingerprint = nested_plan.campaign_fingerprint.clone();
3013 nested_hpo.provenance.controller_fingerprint = nested_plan.controller_fingerprint.clone();
3014 nested_hpo.provenance.fold_set_fingerprint = nested_plan
3015 .fold_set
3016 .as_ref()
3017 .map(stable_json_fingerprint)
3018 .transpose()
3019 .unwrap();
3020 let trace_len_before = trace.lock().unwrap().len();
3021 let error = SequentialScheduler
3022 .execute_hpo_campaign(&nested_plan, &controllers, &provider, &ctx, &nested_hpo)
3023 .unwrap_err();
3024 assert!(error
3025 .to_string()
3026 .contains("does not support nested-stacking FIT_CV"));
3027 assert_eq!(trace.lock().unwrap().len(), trace_len_before);
3028
3029 let resumed_trace = Arc::new(Mutex::new(Vec::new()));
3033 let mut resumed_controllers = RuntimeControllerRegistry::new();
3034 resumed_controllers
3035 .register(Box::new(HpoTestTuner {
3036 id: ControllerId::new("controller:tuner").unwrap(),
3037 trace: Arc::clone(&resumed_trace),
3038 history_len: 2,
3039 proposal_count: 2,
3040 prune_at: None,
3041 }))
3042 .unwrap();
3043 resumed_controllers
3044 .register(Box::new(HpoTestModel {
3045 id: ControllerId::new("controller:model").unwrap(),
3046 trace: Arc::clone(&resumed_trace),
3047 fail_variant: None,
3048 score_by_trial: false,
3049 }))
3050 .unwrap();
3051 let mut resumed_hpo = hpo.clone();
3052 resumed_hpo.trial_budget_total = 4;
3053 let resumed_ctx = RunContext::new(RunId::new("run:hpo.resumed").unwrap(), Some(13));
3054 let resumed = SequentialScheduler
3055 .execute_hpo_campaign(
3056 &plan,
3057 &resumed_controllers,
3058 &provider,
3059 &resumed_ctx,
3060 &resumed_hpo,
3061 )
3062 .unwrap();
3063 assert_eq!(resumed.candidates.len(), 2);
3064 assert_eq!(resumed.checkpoint.trial_history_len, 4);
3065 assert_eq!(
3066 resumed_trace
3067 .lock()
3068 .unwrap()
3069 .iter()
3070 .filter(|event| event.as_str() == "ask")
3071 .count(),
3072 2
3073 );
3074
3075 let mut over_budget_controllers = RuntimeControllerRegistry::new();
3076 over_budget_controllers
3077 .register(Box::new(HpoTestTuner {
3078 id: ControllerId::new("controller:tuner").unwrap(),
3079 trace: Arc::new(Mutex::new(Vec::new())),
3080 history_len: 5,
3081 proposal_count: 0,
3082 prune_at: None,
3083 }))
3084 .unwrap();
3085 let error = SequentialScheduler
3086 .execute_hpo_campaign(
3087 &plan,
3088 &over_budget_controllers,
3089 &provider,
3090 &resumed_ctx,
3091 &resumed_hpo,
3092 )
3093 .unwrap_err();
3094 assert!(error.to_string().contains("exceeds total trial budget"));
3095
3096 #[cfg(feature = "methods-optimizer-local")]
3097 {
3098 let library_path = std::env::var_os("N4M_LIBRARY_PATH")
3102 .expect("native HPO scheduler test requires N4M_LIBRARY_PATH");
3103 let runtime = MethodsRuntime::configure(library_path).unwrap();
3104 let native_trace = Arc::new(Mutex::new(Vec::new()));
3105 let mut native_controllers = RuntimeControllerRegistry::new();
3106 native_controllers
3107 .register(Box::new(MethodsHpoController::new(
3108 ControllerId::new("controller:tuner").unwrap(),
3109 runtime,
3110 )))
3111 .unwrap();
3112 native_controllers
3113 .register(Box::new(HpoTestModel {
3114 id: ControllerId::new("controller:model").unwrap(),
3115 trace: Arc::clone(&native_trace),
3116 fail_variant: None,
3117 score_by_trial: true,
3118 }))
3119 .unwrap();
3120 let mut native_hpo = hpo.clone();
3121 native_hpo.trial_budget_total = 3;
3122 native_hpo.study.study_id = "study:hpo.scheduler.native".to_string();
3123 native_hpo.study.methods_abi = "n4m-abi-2.2".to_string();
3124 native_hpo.study.optimizer.sampler = HpoSampler::Tpe;
3125 native_hpo.study.optimizer.pruner = HpoPruner::Median;
3126 native_hpo.study.optimizer.n_startup_trials = 2;
3127 native_hpo.study.optimizer.seed = 51;
3128 native_hpo.study.optimizer.reduction_factor = 0;
3129 native_hpo.study.search_space.parameters = vec![HpoParameter::Int {
3130 name: "n_components".to_string(),
3131 low: 1,
3132 high: 3,
3133 step: 1,
3134 log: false,
3135 }];
3136 let native = SequentialScheduler
3137 .execute_hpo_campaign(
3138 &plan,
3139 &native_controllers,
3140 &provider,
3141 &RunContext::new(RunId::new("run:hpo.native").unwrap(), Some(13)),
3142 &native_hpo,
3143 )
3144 .unwrap();
3145 assert_eq!(native.checkpoint.trial_history_len, 3);
3146 assert_eq!(native.checkpoint.artifact.format, N4MOPT_FORMAT);
3147 assert_eq!(native.candidates.len(), 2);
3148 assert_eq!(
3149 native
3150 .terminal_trials
3151 .iter()
3152 .map(|entry| entry.trial.status)
3153 .collect::<Vec<_>>(),
3154 vec![
3155 crate::hpo::HpoTrialStatus::Completed,
3156 crate::hpo::HpoTrialStatus::Completed,
3157 crate::hpo::HpoTrialStatus::Pruned,
3158 ]
3159 );
3160 assert!(native.terminal_trials[2]
3161 .trial
3162 .intermediates
3163 .iter()
3164 .any(|item| item.step == 0 && item.should_prune));
3165 assert_eq!(
3166 native_trace
3167 .lock()
3168 .unwrap()
3169 .iter()
3170 .filter(|event| event.as_str() == "model_cv")
3171 .count(),
3172 5,
3173 "native pruning must stop the third trial after its first fold"
3174 );
3175
3176 let prior_checkpoint = native.checkpoint.artifact.clone();
3177 native_hpo.resume_checkpoint = Some(prior_checkpoint.clone());
3178 native_hpo.resume_variants = native
3179 .terminal_trials
3180 .iter()
3181 .filter_map(|snapshot| {
3182 snapshot
3183 .variant_id
3184 .clone()
3185 .map(|variant| (snapshot.trial.id, variant))
3186 })
3187 .collect();
3188 native_hpo.resume_terminal_trials = native.terminal_trials.clone();
3189 native_hpo.trial_budget_total = 4;
3190 let resumed = SequentialScheduler
3191 .execute_hpo_campaign(
3192 &plan,
3193 &native_controllers,
3194 &provider,
3195 &RunContext::new(RunId::new("run:hpo.native.resume").unwrap(), Some(13)),
3196 &native_hpo,
3197 )
3198 .unwrap();
3199 assert_eq!(resumed.checkpoint.trial_history_len, 4);
3200 assert_eq!(resumed.terminal_trials.len(), 4);
3201 assert_eq!(
3202 &resumed.terminal_trials[..3],
3203 native.terminal_trials.as_slice()
3204 );
3205 assert_ne!(
3206 resumed.checkpoint.artifact.opaque_payload,
3207 prior_checkpoint.opaque_payload
3208 );
3209 }
3210 }
3211}
3212
3213pub(crate) fn attach_coordinator_input_lineage(
3214 result: &mut NodeResult,
3215 plan: &ExecutionPlan,
3216 node_id: &NodeId,
3217 upstream_lineage: &BTreeMap<NodeId, LineageId>,
3218) -> Result<()> {
3219 let inferred = inferred_input_lineage_for_node(plan, node_id, upstream_lineage);
3220 if result.lineage.input_lineage.is_empty() {
3221 result.lineage.input_lineage = inferred;
3222 return Ok(());
3223 }
3224
3225 let declared = result
3226 .lineage
3227 .input_lineage
3228 .iter()
3229 .cloned()
3230 .collect::<BTreeSet<_>>()
3231 .into_iter()
3232 .collect::<Vec<_>>();
3233 if declared != inferred {
3234 return Err(DagMlError::RuntimeValidation(format!(
3235 "lineage for node `{}` declared input lineage {:?}, expected {:?}",
3236 result.node_id, declared, inferred
3237 )));
3238 }
3239 result.lineage.input_lineage = declared;
3240 Ok(())
3241}
3242
3243fn attach_nested_stacking_input_lineage(
3250 result: &mut NodeResult,
3251 plan: &ExecutionPlan,
3252 task: &NodeTask,
3253 ctx: &RunContext,
3254 nested: &NestedStackingInput<'_>,
3255) -> Result<()> {
3256 if task.phase != Phase::FitCv || task.node_plan.node_id != *nested.meta_node_id {
3257 return Ok(());
3258 }
3259 let inner_fold_ids = nested
3260 .inner
3261 .inner_fold_set
3262 .folds
3263 .iter()
3264 .map(|fold| fold.fold_id.clone())
3265 .collect::<BTreeSet<_>>();
3266 let source_nodes = incoming_oof_edges(plan, &task.node_plan)?
3267 .into_iter()
3268 .map(|edge| edge.source.node_id.clone())
3269 .collect::<BTreeSet<_>>();
3270 let expected = source_nodes
3271 .iter()
3272 .flat_map(|node_id| {
3273 inner_fold_ids
3274 .iter()
3275 .cloned()
3276 .map(move |fold_id| (node_id.clone(), fold_id))
3277 })
3278 .collect::<BTreeSet<_>>();
3279 let mut actual = BTreeMap::new();
3280 for record in ctx.lineage.records().filter(|record| {
3281 record.phase == Phase::FitCv
3282 && record.variant_id == task.variant_id
3283 && source_nodes.contains(&record.node_id)
3284 && record
3285 .fold_id
3286 .as_ref()
3287 .is_some_and(|fold_id| inner_fold_ids.contains(fold_id))
3288 }) {
3289 let fold_id = record
3290 .fold_id
3291 .clone()
3292 .expect("inner-fold predicate requires a fold id");
3293 if actual
3294 .insert((record.node_id.clone(), fold_id), record.record_id.clone())
3295 .is_some()
3296 {
3297 return Err(DagMlError::RuntimeValidation(
3298 "nested stacking meta input lineage contains duplicate inner-fold evidence"
3299 .to_string(),
3300 ));
3301 }
3302 }
3303 if actual.keys().cloned().collect::<BTreeSet<_>>() != expected {
3304 return Err(DagMlError::RuntimeValidation(
3305 "nested stacking meta input lineage does not exactly cover inner OOF evidence"
3306 .to_string(),
3307 ));
3308 }
3309 let inferred = actual.into_values().collect::<Vec<_>>();
3310 if result.lineage.input_lineage.is_empty() {
3311 result.lineage.input_lineage = inferred;
3312 return Ok(());
3313 }
3314 let declared = result
3315 .lineage
3316 .input_lineage
3317 .iter()
3318 .cloned()
3319 .collect::<BTreeSet<_>>();
3320 if declared.into_iter().collect::<Vec<_>>() != inferred {
3321 return Err(DagMlError::RuntimeValidation(format!(
3322 "nested stacking meta lineage for node `{}` does not match inner OOF evidence",
3323 task.node_plan.node_id
3324 )));
3325 }
3326 result.lineage.input_lineage = inferred;
3327 Ok(())
3328}
3329
3330pub(crate) fn inferred_input_lineage_for_node(
3331 plan: &ExecutionPlan,
3332 node_id: &NodeId,
3333 upstream_lineage: &BTreeMap<NodeId, LineageId>,
3334) -> Vec<LineageId> {
3335 plan.graph_plan
3336 .graph
3337 .edges
3338 .iter()
3339 .filter(|edge| &edge.target.node_id == node_id && edge.contract.propagates_lineage)
3340 .filter_map(|edge| upstream_lineage.get(&edge.source.node_id).cloned())
3341 .collect::<BTreeSet<_>>()
3342 .into_iter()
3343 .collect()
3344}
3345pub(crate) fn collect_input_handles(
3346 plan: &ExecutionPlan,
3347 node_plan: &NodePlan,
3348 output_handles: &BTreeMap<NodeId, BTreeMap<String, HandleRef>>,
3349 output_data_views: &BTreeMap<NodeId, BTreeMap<String, DataProviderViewSpec>>,
3350 resources: &PhaseScopeResources<'_>,
3351 ctx: &RunContext,
3352 scope: &PhaseScope,
3353) -> Result<CollectedInputs> {
3354 let mut inputs = BTreeMap::new();
3355 let mut data_views = BTreeMap::new();
3356 let mut prediction_inputs = BTreeMap::new();
3357 let training_oof_edges = incoming_training_oof_edges(plan, node_plan, scope)?;
3358 let masked_oof_source_ports = if scope.phase == Phase::Predict {
3364 incoming_oof_edges(plan, node_plan)?
3365 } else {
3366 training_oof_edges.clone()
3367 }
3368 .into_iter()
3369 .map(|edge| (edge.source.node_id.clone(), edge.source.port_name.clone()))
3370 .collect::<BTreeSet<_>>();
3371 let bound_data_inputs = node_plan
3372 .data_bindings
3373 .iter()
3374 .map(|binding| binding.input_name.clone())
3375 .collect::<BTreeSet<_>>();
3376 let declared_source_ports = plan
3380 .graph_plan
3381 .graph
3382 .edges
3383 .iter()
3384 .filter(|edge| edge.target.node_id == node_plan.node_id)
3385 .map(|edge| (edge.source.node_id.clone(), edge.source.port_name.clone()))
3386 .collect::<BTreeSet<_>>();
3387 for upstream in &node_plan.input_nodes {
3388 if let Some(handles) = output_handles.get(upstream) {
3389 for (port, handle) in handles {
3390 if !declared_source_ports.contains(&(upstream.clone(), port.clone())) {
3391 continue;
3392 }
3393 if masked_oof_source_ports.contains(&(upstream.clone(), port.clone())) {
3394 continue;
3395 }
3396 inputs.insert(format!("{upstream}.{port}"), handle.clone());
3397 }
3398 }
3399 }
3400 for edge in plan
3401 .graph_plan
3402 .graph
3403 .edges
3404 .iter()
3405 .filter(|edge| edge.target.node_id == node_plan.node_id)
3406 .filter(|edge| edge.contract.kind == PortKind::Data && !edge.contract.requires_oof)
3407 {
3408 if bound_data_inputs.contains(&edge.target.port_name) {
3409 continue;
3410 }
3411 let Some(handles) = output_handles.get(&edge.source.node_id) else {
3412 continue;
3413 };
3414 let Some(handle) = handles.get(&edge.source.port_name) else {
3415 continue;
3416 };
3417 let key = data_view_key(&edge.target.port_name);
3418 if inputs.insert(key.clone(), handle.clone()).is_some() {
3419 return Err(DagMlError::RuntimeValidation(format!(
3420 "node `{}` received duplicate data edge input `{key}`",
3421 node_plan.node_id
3422 )));
3423 }
3424 if let Some(source_views) = output_data_views.get(&edge.source.node_id) {
3425 if let Some(view) = source_views.get(&edge.source.port_name) {
3426 if data_views.insert(key.clone(), view.clone()).is_some() {
3427 return Err(DagMlError::RuntimeValidation(format!(
3428 "node `{}` received duplicate data edge view `{key}`",
3429 node_plan.node_id
3430 )));
3431 }
3432 }
3433 let source_validation_key = validation_data_view_key(&edge.source.port_name);
3434 if let Some(view) = source_views.get(&source_validation_key) {
3435 let validation_key = format!("{key}:validation");
3436 if data_views
3437 .insert(validation_key.clone(), view.clone())
3438 .is_some()
3439 {
3440 return Err(DagMlError::RuntimeValidation(format!(
3441 "node `{}` received duplicate data edge validation view `{validation_key}`",
3442 node_plan.node_id
3443 )));
3444 }
3445 }
3446 }
3447 }
3448 for edge in training_oof_edges {
3449 let key = format!("{}.{}", edge.source.node_id, edge.source.port_name);
3450 let Some(input) = collect_oof_prediction_input(plan, edge, ctx, scope, resources)? else {
3451 return Ok(CollectedInputs {
3452 handles: BTreeMap::new(),
3453 data_views: BTreeMap::new(),
3454 prediction_inputs: BTreeMap::new(),
3455 skip_node: true,
3456 });
3457 };
3458 if inputs.insert(key.clone(), input.handle).is_some() {
3459 return Err(DagMlError::RuntimeValidation(format!(
3460 "node `{}` received duplicate OOF prediction input `{key}`",
3461 node_plan.node_id
3462 )));
3463 }
3464 if prediction_inputs.insert(key.clone(), input.spec).is_some() {
3465 return Err(DagMlError::RuntimeValidation(format!(
3466 "node `{}` received duplicate OOF prediction spec `{key}`",
3467 node_plan.node_id
3468 )));
3469 }
3470 }
3471 if matches!(scope.phase, Phase::Refit | Phase::Predict) {
3478 let off_fold_suffix = scope.phase.as_str().to_ascii_lowercase();
3479 for edge in incoming_oof_edges(plan, node_plan)? {
3480 let Some(input) = collect_off_fold_prediction_input(plan, edge, ctx, scope)? else {
3481 continue;
3482 };
3483 let key = format!(
3484 "{}.{}:{off_fold_suffix}",
3485 edge.source.node_id, edge.source.port_name
3486 );
3487 if inputs.insert(key.clone(), input.handle).is_some() {
3488 return Err(DagMlError::RuntimeValidation(format!(
3489 "node `{}` received duplicate off-fold prediction input `{key}`",
3490 node_plan.node_id
3491 )));
3492 }
3493 if prediction_inputs.insert(key.clone(), input.spec).is_some() {
3494 return Err(DagMlError::RuntimeValidation(format!(
3495 "node `{}` received duplicate off-fold prediction spec `{key}`",
3496 node_plan.node_id
3497 )));
3498 }
3499 }
3500 }
3501 if !node_plan.data_bindings.is_empty() && resources.data_provider.is_none() {
3502 return Err(DagMlError::RuntimeValidation(format!(
3503 "node `{}` requires {} data binding(s) but no runtime data provider is registered",
3504 node_plan.node_id,
3505 node_plan.data_bindings.len()
3506 )));
3507 }
3508 if let Some(data_provider) = resources.data_provider {
3509 let excluded_samples = if scope.phase == Phase::Predict {
3514 BTreeSet::new()
3515 } else {
3516 coordinator_relations_for_node(node_plan, resources)?
3517 .map(|relations| relations.excluded_sample_ids())
3518 .unwrap_or_default()
3519 };
3520 let scope_fold_set = resources.fold_set_override.or(plan.fold_set.as_ref());
3521 for binding in &node_plan.data_bindings {
3522 let predict_cohort = if scope.phase == Phase::Predict {
3523 data_provider.predict_cohort(binding, scope.phase)?
3524 } else {
3525 None
3526 };
3527 let materialized = data_provider.materialize(&DataMaterializationRequest {
3528 run_id: ctx.run_id.clone(),
3529 node_id: node_plan.node_id.clone(),
3530 input_name: binding.input_name.clone(),
3531 phase: scope.phase,
3532 variant_id: scope.variant_id.clone(),
3533 fold_id: scope.fold_id.clone(),
3534 binding: binding.clone(),
3535 predict_cohort: predict_cohort.clone(),
3536 })?;
3537 let branch_view_for_node = branch_view_from_node_metadata(plan, &node_plan.node_id)?;
3538 let mut view = data_view_for_scope(
3539 binding,
3540 scope_fold_set,
3541 scope,
3542 branch_view_for_node.as_ref(),
3543 &excluded_samples,
3544 )?;
3545 if let Some(cohort) = predict_cohort.as_ref() {
3546 bind_predict_cohort_to_view(&mut view, cohort)?;
3547 }
3548 let key = data_view_key(&binding.input_name);
3549 let view_handle = make_data_view_handle(
3550 data_provider,
3551 ctx,
3552 node_plan,
3553 scope,
3554 binding,
3555 DataViewHandleInput {
3556 data_handle: &materialized,
3557 view: &view,
3558 predict_cohort: predict_cohort.as_ref(),
3559 },
3560 )?;
3561 if data_views.insert(key.clone(), view).is_some() {
3562 return Err(DagMlError::RuntimeValidation(format!(
3563 "node `{}` received duplicate data view `{key}`",
3564 node_plan.node_id
3565 )));
3566 }
3567 if inputs.insert(key.clone(), view_handle).is_some() {
3568 return Err(DagMlError::RuntimeValidation(format!(
3569 "node `{}` received duplicate data input `{key}`",
3570 node_plan.node_id
3571 )));
3572 }
3573
3574 if let Some(validation_view) = validation_data_view_for_scope(
3575 binding,
3576 scope_fold_set,
3577 scope,
3578 branch_view_for_node.as_ref(),
3579 &excluded_samples,
3580 )? {
3581 let validation_key = format!("{key}:validation");
3582 let validation_handle = make_data_view_handle(
3583 data_provider,
3584 ctx,
3585 node_plan,
3586 scope,
3587 binding,
3588 DataViewHandleInput {
3589 data_handle: &materialized,
3590 view: &validation_view,
3591 predict_cohort: None,
3592 },
3593 )?;
3594 if data_views
3595 .insert(validation_key.clone(), validation_view)
3596 .is_some()
3597 {
3598 return Err(DagMlError::RuntimeValidation(format!(
3599 "node `{}` received duplicate validation data view `{validation_key}`",
3600 node_plan.node_id
3601 )));
3602 }
3603 if inputs
3604 .insert(validation_key.clone(), validation_handle)
3605 .is_some()
3606 {
3607 return Err(DagMlError::RuntimeValidation(format!(
3608 "node `{}` received duplicate validation data input `{validation_key}`",
3609 node_plan.node_id
3610 )));
3611 }
3612 }
3613 }
3614 }
3615 Ok(CollectedInputs {
3616 handles: inputs,
3617 data_views,
3618 prediction_inputs,
3619 skip_node: false,
3620 })
3621}
3622pub(crate) fn preload_replay_prediction_cache_store(
3623 bundle: &ExecutionBundle,
3624 prediction_cache_store: Option<&dyn RuntimePredictionCacheStore>,
3625 ctx: &mut RunContext,
3626) -> Result<()> {
3627 if bundle.prediction_requirements.is_empty() {
3628 return Ok(());
3629 }
3630 let store = prediction_cache_store.ok_or_else(|| {
3631 DagMlError::RuntimeValidation(format!(
3632 "bundle `{}` cannot preload OOF prediction caches without a prediction cache store",
3633 bundle.bundle_id
3634 ))
3635 })?;
3636 if !ctx.prediction_store.blocks().is_empty() {
3637 return Err(DagMlError::RuntimeValidation(format!(
3638 "bundle `{}` cannot preload OOF prediction caches into a non-empty prediction store",
3639 bundle.bundle_id
3640 )));
3641 }
3642 let contracts = replay_prediction_cache_contracts(bundle)?;
3643 for contract in contracts.values() {
3644 if contract.requirement.prediction_level == PredictionLevel::Sample {
3645 let blocks = store.load_blocks(&contract.cache.requirement_key)?;
3646 if blocks.iter().any(|block| {
3647 block.producer_node != contract.requirement.producer_node
3648 || block.partition != contract.requirement.partition
3649 }) {
3650 return Err(DagMlError::RuntimeValidation(format!(
3651 "prediction cache store returned blocks outside requirement `{}`",
3652 contract.cache.requirement_key
3653 )));
3654 }
3655 let mut payload = build_prediction_cache_payload(&contract.requirement, &blocks)?;
3656 payload.cache_namespace_fingerprints =
3657 contract.cache.cache_namespace_fingerprints.clone();
3658 validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
3659 for block in &payload.blocks {
3660 ctx.prediction_store.append(block.clone())?;
3661 }
3662 } else {
3663 let blocks = store.load_aggregated_blocks(&contract.cache.requirement_key)?;
3664 if blocks.iter().any(|block| {
3665 block.producer_node != contract.requirement.producer_node
3666 || block.partition != contract.requirement.partition
3667 || block.level != contract.requirement.prediction_level
3668 }) {
3669 return Err(DagMlError::RuntimeValidation(format!(
3670 "prediction cache store returned aggregated blocks outside requirement `{}`",
3671 contract.cache.requirement_key
3672 )));
3673 }
3674 let mut payload =
3675 build_aggregated_prediction_cache_payload(&contract.requirement, &blocks)?;
3676 payload.cache_namespace_fingerprints =
3677 contract.cache.cache_namespace_fingerprints.clone();
3678 validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
3679 }
3680 }
3681 Ok(())
3682}
3683
3684pub(crate) fn replay_prediction_cache_contracts(
3685 bundle: &ExecutionBundle,
3686) -> Result<BTreeMap<String, ReplayPredictionCacheContract>> {
3687 bundle.validate()?;
3688 let requirements = bundle
3689 .prediction_requirements
3690 .iter()
3691 .map(|requirement| (requirement.key(), requirement))
3692 .collect::<BTreeMap<_, _>>();
3693 let mut contracts = BTreeMap::new();
3694 for cache in &bundle.prediction_caches {
3695 let requirement = requirements.get(&cache.requirement_key).ok_or_else(|| {
3696 DagMlError::RuntimeValidation(format!(
3697 "prediction cache `{}` references unknown prediction requirement `{}`",
3698 cache.cache_id, cache.requirement_key
3699 ))
3700 })?;
3701 contracts.insert(
3702 cache.requirement_key.clone(),
3703 ReplayPredictionCacheContract {
3704 requirement: (*requirement).clone(),
3705 cache: cache.clone(),
3706 },
3707 );
3708 }
3709 Ok(contracts)
3710}
3711
3712pub(crate) fn materialize_replay_artifact_handles(
3713 plan: &ExecutionPlan,
3714 bundle: &ExecutionBundle,
3715 replay_request: &ReplayPhaseRequest,
3716 artifact_store: &dyn RuntimeArtifactStore,
3717 ctx: &RunContext,
3718) -> Result<MaterializedReplayArtifacts> {
3719 let mut handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
3720 let mut inputs = BTreeMap::<NodeId, BTreeMap<String, ArtifactInputSpec>>::new();
3721 for artifact in &bundle.refit_artifacts {
3722 artifact.validate()?;
3723 let node_plan = plan.node_plans.get(&artifact.node_id).ok_or_else(|| {
3724 DagMlError::RuntimeValidation(format!(
3725 "bundle `{}` artifact references unknown node `{}`",
3726 bundle.bundle_id, artifact.node_id
3727 ))
3728 })?;
3729 if !node_plan.supported_phases.contains(&replay_request.phase) {
3730 return Err(DagMlError::RuntimeValidation(format!(
3731 "bundle `{}` artifact node `{}` does not support replay phase {:?}",
3732 bundle.bundle_id, artifact.node_id, replay_request.phase
3733 )));
3734 }
3735 let handle = artifact_store.materialize(&ArtifactMaterializationRequest {
3736 run_id: ctx.run_id.clone(),
3737 bundle_id: bundle.bundle_id.clone(),
3738 node_id: artifact.node_id.clone(),
3739 phase: replay_request.phase,
3740 variant_id: bundle.selected_variant_id.clone(),
3741 controller_id: artifact.controller_id.clone(),
3742 artifact: artifact.artifact.clone(),
3743 params_fingerprint: artifact.params_fingerprint.clone(),
3744 training_loss_fingerprint: artifact.training_loss_fingerprint.clone(),
3745 })?;
3746 if !matches!(handle.kind, HandleKind::Model | HandleKind::Artifact) {
3747 return Err(DagMlError::RuntimeValidation(format!(
3748 "artifact `{}` materialized as unsupported handle kind {:?}",
3749 artifact.artifact.id, handle.kind
3750 )));
3751 }
3752 if handle.owner_controller != artifact.controller_id {
3753 return Err(DagMlError::RuntimeValidation(format!(
3754 "artifact `{}` handle owner `{}` does not match controller `{}`",
3755 artifact.artifact.id, handle.owner_controller, artifact.controller_id
3756 )));
3757 }
3758 let key = refit_artifact_input_key(&artifact.artifact.id);
3759 if handles
3760 .entry(artifact.node_id.clone())
3761 .or_default()
3762 .insert(key.clone(), handle)
3763 .is_some()
3764 {
3765 return Err(DagMlError::RuntimeValidation(format!(
3766 "duplicate replay artifact input `{key}` for node `{}`",
3767 artifact.node_id
3768 )));
3769 }
3770 if inputs
3771 .entry(artifact.node_id.clone())
3772 .or_default()
3773 .insert(key.clone(), ArtifactInputSpec::from_refit_record(artifact)?)
3774 .is_some()
3775 {
3776 return Err(DagMlError::RuntimeValidation(format!(
3777 "duplicate replay artifact metadata `{key}` for node `{}`",
3778 artifact.node_id
3779 )));
3780 }
3781 }
3782 Ok(MaterializedReplayArtifacts { handles, inputs })
3783}
3784
3785pub(crate) fn derive_task_seed(
3786 root_seed: Option<u64>,
3787 variant_id: Option<&VariantId>,
3788 fold_id: Option<&FoldId>,
3789 node_plan: &NodePlan,
3790 phase: Phase,
3791) -> Option<u64> {
3792 root_seed.map(|root| {
3793 let mut context = SeedContext::root(root);
3794 if let Some(variant_id) = variant_id {
3795 context = context.child(format!("variant:{variant_id}"));
3796 }
3797 if let Some(fold_id) = fold_id {
3798 context = context.child(format!("fold:{fold_id}"));
3799 }
3800 context
3801 .child(format!("node:{}", node_plan.node_id))
3802 .child(format!("phase:{phase:?}"))
3803 .derive_u64("task")
3804 })
3805}