Skip to main content

dag_ml_core/runtime/
scheduler.rs

1// Auto-split from the former monolithic `runtime.rs` (pure refactor).
2use 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
154/// Reject non-direct prediction outputs before the scheduler can aggregate
155/// them.  A terminal external cohort must never fall back to coordinator
156/// relations or a custom aggregation controller simply because a controller
157/// returned observation-level output.
158fn 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    /// Scheduler-owned fold universe for a nested execution scope.  It is
178    /// never inferred from a fold-id string: callers retain the parent-bound
179    /// `NestedFoldSet` and pass only its validated inner set here.
180    pub(crate) fold_set_override: Option<&'a FoldSet>,
181    /// Restrict execution to one dependency-closed subgraph.  Nested stacking
182    /// uses this for its base branches before invoking the meta node; ordinary
183    /// phases leave it empty and keep the full plan topology.
184    pub(crate) node_filter: Option<&'a BTreeSet<NodeId>>,
185    /// An inner base pass must not recursively apply the plan's ordinary
186    /// `inner_cv` policy.  Nested stacking owns that one level explicitly.
187    pub(crate) suppress_inner_cv: bool,
188    /// Explicit inner-OOF/outer-evaluation split for the one declared nested
189    /// stacking meta node.  This is scheduler-private evidence, never a graph
190    /// edge or a controller-selected policy.
191    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    /// Terminal replay is deliberately narrower than generic PREDICT: its
201    /// selected result must be emitted as a direct sample-level block.  This
202    /// prevents the scheduler from reading coordinator relations or invoking
203    /// a custom aggregation controller for an external V2 cohort.
204    pub(crate) direct_sample_prediction_only: bool,
205    pub(crate) artifact_store: Option<&'a mut InMemoryArtifactStore>,
206}
207
208impl SequentialScheduler {
209    /// Run one local tuner session and evaluate every proposal through the
210    /// ordinary FIT_CV scheduler.  The session remains on this thread; only a
211    /// portable [`RuntimeHpoProposal`] and OOF-derived scalar feedback cross
212    /// the controller boundary. SELECT and REFIT deliberately do not occur
213    /// here, so callers can make exactly one selection and one refit after the
214    /// returned report-grade candidate evidence has been audited.
215    pub fn execute_hpo_campaign(
216        &self,
217        plan: &ExecutionPlan,
218        controllers: &RuntimeControllerRegistry,
219        data_provider: &dyn RuntimeDataProvider,
220        ctx: &RunContext,
221        hpo: &RuntimeHpoExecutionContext,
222    ) -> Result<RuntimeHpoCampaignResult> {
223        plan.validate()?;
224        hpo.validate_for_plan(plan)?;
225        let controller = controllers.get(&hpo.controller_id).ok_or_else(|| {
226            DagMlError::RuntimeValidation(format!(
227                "runtime HPO campaign controller `{}` is not registered",
228                hpo.controller_id
229            ))
230        })?;
231        let task = RuntimeHpoCampaignTask {
232            run_id: ctx.run_id.clone(),
233            operation_id: hpo.operation_id.clone(),
234            controller_id: hpo.controller_id.clone(),
235            target_node_id: hpo.target_node_id.clone(),
236            seed: ctx.root_seed,
237        };
238        let mut session = controller.create_tuner_session(&task, hpo)?;
239        let history_at_start = session.trial_history_len()?;
240        if history_at_start > hpo.trial_budget_total {
241            return Err(DagMlError::RuntimeValidation(format!(
242                "runtime HPO restored native history ({history_at_start}) exceeds total trial budget ({})",
243                hpo.trial_budget_total
244            )));
245        }
246        let remaining_trials = hpo.trial_budget_total - history_at_start;
247        let mut candidates = Vec::new();
248        let mut proposed_variant_ids = BTreeSet::new();
249        // Fresh proposals are checkpointed by this call.  The native study can
250        // nevertheless retain an incumbent from a restored terminal trial, so
251        // keep its persisted trial->variant binding separate from the new
252        // checkpoint evidence and extend it as we ask new trials.
253        let mut trial_variants = BTreeMap::new();
254        let mut incumbent_variants = hpo.resume_variants.clone();
255        let mut terminal_trials = BTreeMap::new();
256        let mut completed_proposals = Vec::new();
257        let mut completed_reports = Vec::new();
258
259        for _ in 0..remaining_trials {
260            let Some(proposal) = session.ask()? else {
261                break;
262            };
263            if trial_variants
264                .insert(proposal.trial_id, proposal.variant.variant_id.clone())
265                .is_some()
266            {
267                return Err(DagMlError::RuntimeValidation(format!(
268                    "runtime HPO session proposed duplicate trial `{}`",
269                    proposal.trial_id
270                )));
271            }
272            if incumbent_variants
273                .insert(proposal.trial_id, proposal.variant.variant_id.clone())
274                .is_some()
275            {
276                return Err(DagMlError::RuntimeValidation(format!(
277                    "runtime HPO session reused restored trial `{}`",
278                    proposal.trial_id
279                )));
280            }
281            if !proposed_variant_ids.insert(proposal.variant.variant_id.clone()) {
282                return Err(DagMlError::RuntimeValidation(format!(
283                    "runtime HPO session proposed duplicate variant `{}`",
284                    proposal.variant.variant_id
285                )));
286            }
287            let mut candidate_plan = plan.clone();
288            candidate_plan.variants = vec![proposal.variant.clone()];
289            candidate_plan.validate()?;
290            let mut candidate_ctx =
291                RunContext::new(ctx.run_id.clone(), proposal.variant.seed.or(ctx.root_seed));
292            candidate_ctx.variant_id = Some(proposal.variant.variant_id.clone());
293
294            let evaluation = self.execute_hpo_candidate_fit_cv(
295                &candidate_plan,
296                controllers,
297                data_provider,
298                &mut candidate_ctx,
299            );
300            if let Err(error) = evaluation {
301                session.tell(
302                    proposal.trial_id,
303                    RuntimeHpoTerminal::Failed {
304                        failure: RuntimeHpoFailure {
305                            code: "DAGML_CV_ERROR".to_string(),
306                            message: error.to_string(),
307                            retryable: false,
308                        },
309                    },
310                )?;
311                terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
312                continue;
313            }
314            if let Err(error) = candidate_ctx
315                .collect_cross_fold_validation_scores(plan_oof_partition_mode(&candidate_plan))
316            {
317                session.tell(
318                    proposal.trial_id,
319                    RuntimeHpoTerminal::Failed {
320                        failure: RuntimeHpoFailure {
321                            code: "DAGML_SCORE_ERROR".to_string(),
322                            message: error.to_string(),
323                            retryable: false,
324                        },
325                    },
326                )?;
327                terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
328                continue;
329            }
330            let report = candidate_ctx
331                .score_collector
332                .iter()
333                .find(|report| {
334                    report.producer_node == hpo.selection.producer_node
335                        && report.producer_port.as_deref()
336                            == Some(hpo.selection.producer_port.as_str())
337                        && report.partition == PredictionPartition::Validation
338                        && report
339                            .fold_id
340                            .as_ref()
341                            .is_some_and(|fold| fold.as_str() == "avg")
342                })
343                .cloned();
344            let Some(mut report) = report else {
345                session.tell(
346                    proposal.trial_id,
347                    RuntimeHpoTerminal::Failed {
348                        failure: RuntimeHpoFailure {
349                            code: "DAGML_SCORE_MISSING".to_string(),
350                            message: format!(
351                                "runtime HPO trial `{}` emitted no target OOF average",
352                                proposal.trial_id
353                            ),
354                            retryable: false,
355                        },
356                    },
357                )?;
358                terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
359                continue;
360            };
361            report.variant_id = Some(proposal.variant.variant_id.clone());
362            let score = report
363                .metrics
364                .get(hpo.selection.metric.name())
365                .copied()
366                .filter(|score| score.is_finite());
367            let Some(score) = score else {
368                session.tell(
369                    proposal.trial_id,
370                    RuntimeHpoTerminal::Failed {
371                        failure: RuntimeHpoFailure {
372                            code: "DAGML_SCORE_NONFINITE".to_string(),
373                            message: format!(
374                                "runtime HPO trial `{}` emitted no finite `{}` score",
375                                proposal.trial_id,
376                                hpo.selection.metric.name()
377                            ),
378                            retryable: false,
379                        },
380                    },
381                )?;
382                terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Failed);
383                continue;
384            };
385            let intermediate = RuntimeHpoIntermediate {
386                trial_id: proposal.trial_id,
387                step: 0,
388                score,
389            };
390            if session.report_intermediate(intermediate)? == RuntimeHpoIntermediateOutcome::Pruned {
391                terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Pruned);
392                continue;
393            }
394            session.tell(proposal.trial_id, RuntimeHpoTerminal::Completed { score })?;
395            terminal_trials.insert(proposal.trial_id, HpoTrialTerminalState::Completed);
396            completed_proposals.push(proposal.clone());
397            completed_reports.push(RuntimeHpoCompletedReport {
398                trial_id: proposal.trial_id,
399                variant_id: proposal.variant.variant_id.clone(),
400                report: report.clone(),
401            });
402
403            let mut validation_reports = candidate_ctx
404                .score_collector
405                .iter()
406                .filter(|item| item.partition == PredictionPartition::Validation)
407                .cloned()
408                .collect::<Vec<_>>();
409            for item in &mut validation_reports {
410                item.variant_id = Some(proposal.variant.variant_id.clone());
411            }
412            candidates.push(RuntimeHpoCandidateEvaluation {
413                validation_predictions: capture_variant_validation_predictions(
414                    &proposal.variant.variant_id,
415                    None,
416                    &candidate_ctx,
417                ),
418                lineage: candidate_ctx.lineage.records().cloned().collect(),
419                proposal,
420                score,
421                validation_reports,
422            });
423        }
424
425        let history_at_checkpoint = session.trial_history_len()?;
426        if history_at_checkpoint != hpo.trial_budget_total {
427            return Err(DagMlError::RuntimeValidation(format!(
428                "runtime HPO native history ended at {history_at_checkpoint}, expected total trial budget {}",
429                hpo.trial_budget_total
430            )));
431        }
432
433        let checkpoint = RuntimeHpoCheckpointResult {
434            artifact: session.checkpoint()?,
435            provenance: hpo.provenance.clone(),
436            operation_id: hpo.operation_id.clone(),
437            controller_id: hpo.controller_id.clone(),
438            target_node_id: hpo.target_node_id.clone(),
439            completed_proposals,
440            completed_reports,
441            trial_history_len: history_at_checkpoint,
442        };
443        validate_hpo_checkpoint_result(
444            &checkpoint,
445            hpo,
446            &trial_variants,
447            &terminal_trials,
448            history_at_start,
449        )?;
450        let incumbent = session.incumbent(&incumbent_variants)?.ok_or_else(|| {
451            DagMlError::RuntimeValidation(
452                "native HPO campaign has no completed native incumbent after terminalization"
453                    .to_string(),
454            )
455        })?;
456        if incumbent.metric != hpo.selection.metric.name()
457            || incumbent.direction != hpo.selection.direction
458            || incumbent_variants.get(&incumbent.trial_id) != Some(&incumbent.variant_id)
459            || !incumbent.score.is_finite()
460        {
461            return Err(DagMlError::RuntimeValidation(
462                "native HPO incumbent is not bound to this scheduler campaign's metric, direction, trial, and variant"
463                    .to_string(),
464            ));
465        }
466        let terminal_trials = session.terminal_trial_snapshots(&incumbent_variants)?;
467        if terminal_trials.len() != history_at_checkpoint as usize
468            || terminal_trials
469                .windows(2)
470                .any(|pair| pair[0].trial.id >= pair[1].trial.id)
471        {
472            return Err(DagMlError::RuntimeValidation(
473                "native HPO terminal ledger is not a complete strictly ordered history".to_string(),
474            ));
475        }
476        Ok(RuntimeHpoCampaignResult {
477            operation_id: hpo.operation_id.clone(),
478            controller_id: hpo.controller_id.clone(),
479            target_node_id: hpo.target_node_id.clone(),
480            candidates,
481            checkpoint,
482            incumbent,
483            terminal_trials,
484        })
485    }
486
487    fn execute_hpo_candidate_fit_cv(
488        &self,
489        plan: &ExecutionPlan,
490        controllers: &RuntimeControllerRegistry,
491        data_provider: &dyn RuntimeDataProvider,
492        ctx: &mut RunContext,
493    ) -> Result<Vec<NodeResult>> {
494        if let Some(nested) = nested_stacking_campaign_plan(plan)? {
495            return self.execute_nested_stacking_fit_cv(
496                plan,
497                controllers,
498                data_provider,
499                ctx,
500                &nested,
501            );
502        }
503        let candidate_plan = plan;
504        ctx.configure_global_oof_aggregation(candidate_plan, data_provider)?;
505        let fold_ids = candidate_plan
506            .fold_set
507            .as_ref()
508            .map(|fold_set| {
509                fold_set
510                    .folds
511                    .iter()
512                    .map(|fold| Some(fold.fold_id.clone()))
513                    .collect::<Vec<_>>()
514            })
515            .unwrap_or_else(|| vec![None]);
516        let variant = candidate_plan
517            .variants
518            .first()
519            .expect("candidate plan has exactly one variant");
520        let mut results = Vec::new();
521        for fold_id in fold_ids {
522            results.extend(self.execute_phase_scope(
523                candidate_plan,
524                controllers,
525                ctx,
526                PhaseScope {
527                    phase: Phase::FitCv,
528                    variant_id: Some(variant.variant_id.clone()),
529                    variant: Some(VariantExecutionSpec::from_plan(variant)),
530                    fold_id,
531                    seed_root: variant.seed.or(ctx.root_seed),
532                },
533                PhaseScopeResources {
534                    data_provider: Some(data_provider),
535                    ..Default::default()
536                },
537            )?);
538        }
539        Ok(results)
540    }
541
542    pub fn execute_phase(
543        &self,
544        plan: &ExecutionPlan,
545        controllers: &RuntimeControllerRegistry,
546        ctx: &mut RunContext,
547        phase: Phase,
548    ) -> Result<Vec<NodeResult>> {
549        plan.validate()?;
550        let variant_id = ctx.variant_id.clone();
551        let seed_root = ctx.root_seed;
552        self.execute_phase_scope(
553            plan,
554            controllers,
555            ctx,
556            PhaseScope {
557                phase,
558                variant_id,
559                variant: None,
560                fold_id: None,
561                seed_root,
562            },
563            PhaseScopeResources::default(),
564        )
565    }
566
567    pub fn execute_phase_with_data_provider(
568        &self,
569        plan: &ExecutionPlan,
570        controllers: &RuntimeControllerRegistry,
571        data_provider: &dyn RuntimeDataProvider,
572        ctx: &mut RunContext,
573        phase: Phase,
574    ) -> Result<Vec<NodeResult>> {
575        plan.validate()?;
576        let variant_id = ctx.variant_id.clone();
577        let seed_root = ctx.root_seed;
578        self.execute_phase_scope(
579            plan,
580            controllers,
581            ctx,
582            PhaseScope {
583                phase,
584                variant_id,
585                variant: None,
586                fold_id: None,
587                seed_root,
588            },
589            PhaseScopeResources {
590                data_provider: Some(data_provider),
591                ..Default::default()
592            },
593        )
594    }
595
596    pub fn execute_campaign_phase(
597        &self,
598        plan: &ExecutionPlan,
599        controllers: &RuntimeControllerRegistry,
600        ctx: &mut RunContext,
601        phase: Phase,
602    ) -> Result<Vec<NodeResult>> {
603        plan.validate()?;
604        if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
605            return Err(DagMlError::RuntimeValidation(
606                "nested stacking FIT_CV requires execute_campaign_phase_with_data_provider so the scheduler can materialize parent-bound inner folds"
607                    .to_string(),
608            ));
609        }
610        let mut results = Vec::new();
611        let fold_ids = if phase == Phase::FitCv {
612            plan.fold_set
613                .as_ref()
614                .map(|fold_set| {
615                    fold_set
616                        .folds
617                        .iter()
618                        .map(|fold| Some(fold.fold_id.clone()))
619                        .collect::<Vec<_>>()
620                })
621                .unwrap_or_else(|| vec![None])
622        } else {
623            vec![None]
624        };
625        for variant in &plan.variants {
626            if ctx
627                .variant_id
628                .as_ref()
629                .is_some_and(|requested| requested != &variant.variant_id)
630            {
631                continue;
632            }
633            for fold_id in &fold_ids {
634                let seed_root = variant.seed.or(ctx.root_seed);
635                results.extend(self.execute_phase_scope(
636                    plan,
637                    controllers,
638                    ctx,
639                    PhaseScope {
640                        phase,
641                        variant_id: Some(variant.variant_id.clone()),
642                        variant: Some(VariantExecutionSpec::from_plan(variant)),
643                        fold_id: fold_id.clone(),
644                        seed_root,
645                    },
646                    PhaseScopeResources::default(),
647                )?);
648            }
649        }
650        Ok(results)
651    }
652
653    pub fn execute_campaign_phase_with_data_provider(
654        &self,
655        plan: &ExecutionPlan,
656        controllers: &RuntimeControllerRegistry,
657        data_provider: &dyn RuntimeDataProvider,
658        ctx: &mut RunContext,
659        phase: Phase,
660    ) -> Result<Vec<NodeResult>> {
661        plan.validate()?;
662        if phase == Phase::FitCv {
663            ctx.configure_global_oof_aggregation(plan, data_provider)?;
664            if let Some(nested) = nested_stacking_campaign_plan(plan)? {
665                return self.execute_nested_stacking_fit_cv(
666                    plan,
667                    controllers,
668                    data_provider,
669                    ctx,
670                    &nested,
671                );
672            }
673        }
674        let mut results = Vec::new();
675        let fold_ids = if phase == Phase::FitCv {
676            plan.fold_set
677                .as_ref()
678                .map(|fold_set| {
679                    fold_set
680                        .folds
681                        .iter()
682                        .map(|fold| Some(fold.fold_id.clone()))
683                        .collect::<Vec<_>>()
684                })
685                .unwrap_or_else(|| vec![None])
686        } else {
687            vec![None]
688        };
689        for variant in &plan.variants {
690            if ctx
691                .variant_id
692                .as_ref()
693                .is_some_and(|requested| requested != &variant.variant_id)
694            {
695                continue;
696            }
697            for fold_id in &fold_ids {
698                let seed_root = variant.seed.or(ctx.root_seed);
699                results.extend(self.execute_phase_scope(
700                    plan,
701                    controllers,
702                    ctx,
703                    PhaseScope {
704                        phase,
705                        variant_id: Some(variant.variant_id.clone()),
706                        variant: Some(VariantExecutionSpec::from_plan(variant)),
707                        fold_id: fold_id.clone(),
708                        seed_root,
709                    },
710                    PhaseScopeResources {
711                        data_provider: Some(data_provider),
712                        ..Default::default()
713                    },
714                )?);
715            }
716        }
717        Ok(results)
718    }
719
720    pub fn execute_campaign_phase_with_data_provider_and_artifact_store(
721        &self,
722        plan: &ExecutionPlan,
723        controllers: &RuntimeControllerRegistry,
724        data_provider: &dyn RuntimeDataProvider,
725        artifact_store: &mut InMemoryArtifactStore,
726        ctx: &mut RunContext,
727        phase: Phase,
728    ) -> Result<Vec<NodeResult>> {
729        plan.validate()?;
730        if phase == Phase::FitCv {
731            ctx.configure_global_oof_aggregation(plan, data_provider)?;
732            if let Some(nested) = nested_stacking_campaign_plan(plan)? {
733                // FIT_CV produces no refit artifacts. Keep the data-provider
734                // route canonical rather than silently using an artifact store
735                // that cannot participate in the inner-OOF proof.
736                return self.execute_nested_stacking_fit_cv(
737                    plan,
738                    controllers,
739                    data_provider,
740                    ctx,
741                    &nested,
742                );
743            }
744        }
745        let mut results = Vec::new();
746        let fold_ids = if phase == Phase::FitCv {
747            plan.fold_set
748                .as_ref()
749                .map(|fold_set| {
750                    fold_set
751                        .folds
752                        .iter()
753                        .map(|fold| Some(fold.fold_id.clone()))
754                        .collect::<Vec<_>>()
755                })
756                .unwrap_or_else(|| vec![None])
757        } else {
758            vec![None]
759        };
760        for variant in &plan.variants {
761            if ctx
762                .variant_id
763                .as_ref()
764                .is_some_and(|requested| requested != &variant.variant_id)
765            {
766                continue;
767            }
768            for fold_id in &fold_ids {
769                let seed_root = variant.seed.or(ctx.root_seed);
770                results.extend(self.execute_phase_scope(
771                    plan,
772                    controllers,
773                    ctx,
774                    PhaseScope {
775                        phase,
776                        variant_id: Some(variant.variant_id.clone()),
777                        variant: Some(VariantExecutionSpec::from_plan(variant)),
778                        fold_id: fold_id.clone(),
779                        seed_root,
780                    },
781                    PhaseScopeResources {
782                        data_provider: Some(data_provider),
783                        artifact_store: Some(&mut *artifact_store),
784                        ..Default::default()
785                    },
786                )?);
787            }
788        }
789        Ok(results)
790    }
791
792    /// Execute one explicitly declared nested-stacking FIT_CV campaign.
793    ///
794    /// For every outer fold, base nodes first produce their OOF predictions on
795    /// the parent-bound inner folds (the only rows used to fit the meta-model),
796    /// then independently produce outer-validation predictions (the only rows
797    /// scored by the meta-model).  The meta invocation receives both evidence
798    /// classes under separate keys; it cannot accidentally train on the outer
799    /// validation rows through the generic OOF collector.
800    fn execute_nested_stacking_fit_cv(
801        &self,
802        plan: &ExecutionPlan,
803        controllers: &RuntimeControllerRegistry,
804        data_provider: &dyn RuntimeDataProvider,
805        ctx: &mut RunContext,
806        nested: &NestedStackingCampaignPlan,
807    ) -> Result<Vec<NodeResult>> {
808        let parent_fold_ids = nested
809            .outer_scopes
810            .iter()
811            .map(|outer| outer.outer_fold_id.clone())
812            .collect::<BTreeSet<_>>();
813        if let Some(existing) = &ctx.validation_scoring_fold_ids {
814            if existing != &parent_fold_ids {
815                return Err(DagMlError::RuntimeValidation(
816                    "nested stacking cannot reuse a run context with a different report-grade outer fold set"
817                        .to_string(),
818                ));
819            }
820        } else {
821            ctx.validation_scoring_fold_ids = Some(parent_fold_ids);
822        }
823        let mut results = Vec::new();
824        for variant in &plan.variants {
825            if ctx
826                .variant_id
827                .as_ref()
828                .is_some_and(|requested| requested != &variant.variant_id)
829            {
830                continue;
831            }
832            let seed_root = variant.seed.or(ctx.root_seed);
833            let variant_id = Some(variant.variant_id.clone());
834            let variant_spec = Some(VariantExecutionSpec::from_plan(variant));
835            for outer in &nested.outer_scopes {
836                for inner_fold in &outer.inner.inner_fold_set.folds {
837                    results.extend(self.execute_phase_scope(
838                        plan,
839                        controllers,
840                        ctx,
841                        PhaseScope {
842                            phase: Phase::FitCv,
843                            variant_id: variant_id.clone(),
844                            variant: variant_spec.clone(),
845                            fold_id: Some(inner_fold.fold_id.clone()),
846                            seed_root,
847                        },
848                        PhaseScopeResources {
849                            data_provider: Some(data_provider),
850                            fold_set_override: Some(&outer.inner.inner_fold_set),
851                            node_filter: Some(&nested.base_node_ids),
852                            suppress_inner_cv: true,
853                            ..Default::default()
854                        },
855                    )?);
856                }
857
858                // Materialize outer-validation base features in a distinct
859                // scope.  They stay out of the unsuffixed meta inputs.
860                results.extend(self.execute_phase_scope(
861                    plan,
862                    controllers,
863                    ctx,
864                    PhaseScope {
865                        phase: Phase::FitCv,
866                        variant_id: variant_id.clone(),
867                        variant: variant_spec.clone(),
868                        fold_id: Some(outer.outer_fold_id.clone()),
869                        seed_root,
870                    },
871                    PhaseScopeResources {
872                        data_provider: Some(data_provider),
873                        node_filter: Some(&nested.base_node_ids),
874                        suppress_inner_cv: true,
875                        ..Default::default()
876                    },
877                )?);
878
879                let meta_only = BTreeSet::from([nested.meta_node_id.clone()]);
880                results.extend(self.execute_phase_scope(
881                    plan,
882                    controllers,
883                    ctx,
884                    PhaseScope {
885                        phase: Phase::FitCv,
886                        variant_id: variant_id.clone(),
887                        variant: variant_spec.clone(),
888                        fold_id: Some(outer.outer_fold_id.clone()),
889                        seed_root,
890                    },
891                    PhaseScopeResources {
892                        data_provider: Some(data_provider),
893                        node_filter: Some(&meta_only),
894                        suppress_inner_cv: true,
895                        nested_stacking: Some(NestedStackingInput {
896                            meta_node_id: &nested.meta_node_id,
897                            inner: &outer.inner,
898                        }),
899                        ..Default::default()
900                    },
901                )?);
902            }
903        }
904        Ok(results)
905    }
906
907    pub fn execute_bundle_replay(
908        &self,
909        replay: BundleReplayExecution<'_>,
910        ctx: &mut RunContext,
911    ) -> Result<Vec<NodeResult>> {
912        self.execute_bundle_replay_with_prediction_mode(replay, ctx, false)
913    }
914
915    /// Execute a PREDICT replay whose controller output is required to be a
916    /// direct sample-level block.  This is an internal terminal boundary, not
917    /// a change to the generic replay contract.
918    pub(crate) fn execute_direct_sample_bundle_replay(
919        &self,
920        replay: BundleReplayExecution<'_>,
921        ctx: &mut RunContext,
922    ) -> Result<Vec<NodeResult>> {
923        if replay.replay_request.phase != Phase::Predict {
924            return Err(DagMlError::RuntimeValidation(
925                "direct sample bundle replay is valid only for PREDICT".to_string(),
926            ));
927        }
928        self.execute_bundle_replay_with_prediction_mode(replay, ctx, true)
929    }
930
931    fn execute_bundle_replay_with_prediction_mode(
932        &self,
933        replay: BundleReplayExecution<'_>,
934        ctx: &mut RunContext,
935        direct_sample_prediction_only: bool,
936    ) -> Result<Vec<NodeResult>> {
937        replay.bundle.validate_against_plan(replay.plan)?;
938        replay
939            .replay_request
940            .validate_for_bundle_with_prediction_cache_store(
941                replay.bundle,
942                replay.prediction_cache_store.is_some(),
943            )?;
944        replay
945            .bundle
946            .validate_replay_envelopes(replay.data_envelopes)?;
947        let prediction_cache_contracts = if replay.replay_request.phase == Phase::Refit {
948            Some(replay_prediction_cache_contracts(replay.bundle)?)
949        } else {
950            None
951        };
952        if replay.replay_request.phase == Phase::Refit {
953            preload_replay_prediction_cache_store(
954                replay.bundle,
955                replay.prediction_cache_store,
956                ctx,
957            )?;
958        }
959        let replay_artifacts = materialize_replay_artifact_handles(
960            replay.plan,
961            replay.bundle,
962            replay.replay_request,
963            replay.artifact_store,
964            ctx,
965        )?;
966        let selected_variant = replay
967            .bundle
968            .selected_variant_id
969            .as_ref()
970            .map(|selected| {
971                replay
972                    .plan
973                    .variants
974                    .iter()
975                    .find(|variant| &variant.variant_id == selected)
976                    .map(VariantExecutionSpec::from_plan)
977                    .ok_or_else(|| {
978                        DagMlError::RuntimeValidation(format!(
979                            "bundle `{}` selected unknown variant `{selected}`",
980                            replay.bundle.bundle_id
981                        ))
982                    })
983            })
984            .transpose()?;
985        let seed_root = selected_variant
986            .as_ref()
987            .and_then(|variant| variant.seed)
988            .or(ctx.root_seed);
989
990        self.execute_phase_scope(
991            replay.plan,
992            replay.controllers,
993            ctx,
994            PhaseScope {
995                phase: replay.replay_request.phase,
996                variant_id: replay.bundle.selected_variant_id.clone(),
997                variant: selected_variant,
998                fold_id: None,
999                seed_root,
1000            },
1001            PhaseScopeResources {
1002                data_provider: Some(replay.data_provider),
1003                replay_artifact_handles: Some(&replay_artifacts.handles),
1004                replay_artifact_inputs: Some(&replay_artifacts.inputs),
1005                replay_bundle_id: Some(&replay.bundle.bundle_id),
1006                data_envelopes: Some(replay.data_envelopes),
1007                prediction_cache_store: replay.prediction_cache_store,
1008                prediction_cache_contracts: prediction_cache_contracts.as_ref(),
1009                direct_sample_prediction_only,
1010                ..Default::default()
1011            },
1012        )
1013    }
1014
1015    fn execute_phase_scope(
1016        &self,
1017        plan: &ExecutionPlan,
1018        controllers: &RuntimeControllerRegistry,
1019        ctx: &mut RunContext,
1020        scope: PhaseScope,
1021        mut resources: PhaseScopeResources<'_>,
1022    ) -> Result<Vec<NodeResult>> {
1023        let _phase_span = crate::observability::phase_span(
1024            ctx.run_id.as_str(),
1025            plan.id.as_str(),
1026            scope.phase.as_str(),
1027            scope.variant_id.as_ref().map(VariantId::as_str),
1028            scope.fold_id.as_ref().map(FoldId::as_str),
1029        )
1030        .entered();
1031        let mut results = Vec::new();
1032        let mut output_handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
1033        let mut output_data_views =
1034            BTreeMap::<NodeId, BTreeMap<String, DataProviderViewSpec>>::new();
1035        let mut input_lineage = BTreeMap::<NodeId, LineageId>::new();
1036
1037        for level in plan.node_parallel_levels_for_phase(scope.phase)? {
1038            for node_id in &level {
1039                if resources
1040                    .node_filter
1041                    .is_some_and(|allowed| !allowed.contains(node_id))
1042                {
1043                    continue;
1044                }
1045                let node_plan = plan
1046                    .node_plans
1047                    .get(node_id)
1048                    .expect("execution plan was validated");
1049                // Cross-branch merge reassembly (concat or late-fusion) is a
1050                // scheduler/runtime handler, not a controller call: it reads the
1051                // upstream branch OOF blocks from the prediction store and emits
1052                // one merged per-sample OOF block. Intercept it before the
1053                // controller path (and before the `requires_oof` edge collection,
1054                // which is a stacking contract the branch inputs do not satisfy).
1055                if let Some(reduction) = merge_reduction_mode(plan, node_plan) {
1056                    if let Some(mut result) =
1057                        reassemble_branch_merge(plan, node_plan, ctx, &scope, reduction)?
1058                    {
1059                        let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1060                        let task = NodeTask {
1061                            inner_fold_set: None,
1062                            run_id: ctx.run_id.clone(),
1063                            node_plan: task_node_plan.clone(),
1064                            phase: scope.phase,
1065                            variant_id: scope.variant_id.clone(),
1066                            variant: scope.variant.clone(),
1067                            fold_id: scope.fold_id.clone(),
1068                            branch_path: Vec::new(),
1069                            input_handles: BTreeMap::new(),
1070                            data_views: BTreeMap::new(),
1071                            prediction_inputs: BTreeMap::new(),
1072                            artifact_inputs: BTreeMap::new(),
1073                            required_loss_attestations: NodeTask::required_loss_attestations_for(
1074                                &task_node_plan,
1075                                scope.phase,
1076                            )?,
1077                            fit_influence: FitInfluenceTask::default(),
1078                            seed: None,
1079                        };
1080                        normalize_result_prediction_ports(plan, &task, &mut result)?;
1081                        result.validate_for_task(&task)?;
1082                        for prediction in &result.predictions {
1083                            ctx.prediction_store.append(prediction.clone())?;
1084                        }
1085                        apply_result_scoring(
1086                            &result,
1087                            &mut ctx.score_collector,
1088                            &mut ctx.regression_target_records,
1089                        )?;
1090                        ctx.lineage.record(result.lineage.clone())?;
1091                        output_handles.insert(node_id.clone(), result.outputs.clone());
1092                        input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1093                        results.push(result);
1094                    }
1095                    continue;
1096                }
1097                let controller = controllers.get(&node_plan.controller_id).ok_or_else(|| {
1098                    DagMlError::RuntimeValidation(format!(
1099                        "runtime controller `{}` is not registered",
1100                        node_plan.controller_id
1101                    ))
1102                })?;
1103                let collected_inputs = collect_input_handles(
1104                    plan,
1105                    node_plan,
1106                    &output_handles,
1107                    &output_data_views,
1108                    &resources,
1109                    ctx,
1110                    &scope,
1111                )?;
1112                if collected_inputs.skip_node {
1113                    continue;
1114                }
1115                let mut input_handles = collected_inputs.handles;
1116                let mut prediction_inputs = collected_inputs.prediction_inputs;
1117                if let Some(nested) = resources.nested_stacking.as_ref() {
1118                    replace_nested_stacking_fit_cv_inputs(
1119                        plan,
1120                        node_plan,
1121                        ctx,
1122                        &scope,
1123                        nested,
1124                        &mut input_handles,
1125                        &mut prediction_inputs,
1126                    )?;
1127                }
1128                let mut artifact_inputs = BTreeMap::new();
1129                if let Some(node_artifact_handles) = resources
1130                    .replay_artifact_handles
1131                    .and_then(|handles| handles.get(node_id))
1132                {
1133                    for (key, handle) in node_artifact_handles {
1134                        if input_handles.insert(key.clone(), handle.clone()).is_some() {
1135                            return Err(DagMlError::RuntimeValidation(format!(
1136                                "node `{node_id}` received duplicate replay artifact input `{key}`"
1137                            )));
1138                        }
1139                    }
1140                }
1141                if let Some(node_artifact_inputs) = resources
1142                    .replay_artifact_inputs
1143                    .and_then(|inputs| inputs.get(node_id))
1144                {
1145                    for (key, spec) in node_artifact_inputs {
1146                        if artifact_inputs.insert(key.clone(), spec.clone()).is_some() {
1147                            return Err(DagMlError::RuntimeValidation(format!(
1148                                "node `{node_id}` received duplicate replay artifact metadata `{key}`"
1149                            )));
1150                        }
1151                    }
1152                }
1153                let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1154                let inner_fold_set = (!resources.suppress_inner_cv)
1155                    .then(|| {
1156                        inner_fold_set_for_scope(
1157                            &plan.campaign,
1158                            plan.fold_set.as_ref(),
1159                            node_plan,
1160                            &scope,
1161                        )
1162                    })
1163                    .transpose()?
1164                    .flatten();
1165                let fit_influence = fit_influence_task_for_node(
1166                    plan,
1167                    &task_node_plan,
1168                    &collected_inputs.data_views,
1169                )?;
1170                let task = NodeTask {
1171                    inner_fold_set,
1172                    run_id: ctx.run_id.clone(),
1173                    node_plan: task_node_plan.clone(),
1174                    phase: scope.phase,
1175                    variant_id: scope.variant_id.clone(),
1176                    variant: scope.variant.clone(),
1177                    fold_id: scope.fold_id.clone(),
1178                    branch_path: Vec::new(),
1179                    input_handles,
1180                    data_views: collected_inputs.data_views,
1181                    prediction_inputs,
1182                    artifact_inputs,
1183                    required_loss_attestations: NodeTask::required_loss_attestations_for(
1184                        &task_node_plan,
1185                        scope.phase,
1186                    )?,
1187                    fit_influence,
1188                    seed: derive_task_seed(
1189                        scope.seed_root,
1190                        scope.variant_id.as_ref(),
1191                        scope.fold_id.as_ref(),
1192                        &task_node_plan,
1193                        scope.phase,
1194                    ),
1195                };
1196                let _node_span = crate::observability::node_span(
1197                    task.run_id.as_str(),
1198                    plan.id.as_str(),
1199                    task.phase.as_str(),
1200                    task.node_plan.node_id.as_str(),
1201                    task.node_plan.controller_id.as_str(),
1202                )
1203                .entered();
1204                let mut result = if task.node_plan.kind == NodeKind::Tuner {
1205                    return Err(DagMlError::RuntimeValidation(format!(
1206                        "tuner node `{}` requires execute_hpo_campaign with an explicit RuntimeHpoExecutionContext",
1207                        task.node_plan.node_id
1208                    )));
1209                } else {
1210                    match resources.data_provider {
1211                        Some(data_provider) => {
1212                            controller.invoke_with_data_provider(&task, data_provider)?
1213                        }
1214                        None => controller.invoke(&task)?,
1215                    }
1216                };
1217                record_fit_influence_diagnostic(&task, &mut result);
1218                normalize_result_prediction_ports(plan, &task, &mut result)?;
1219                result.validate_for_task(&task)?;
1220                if resources.direct_sample_prediction_only {
1221                    validate_direct_sample_prediction_result(&task, &result)?;
1222                } else {
1223                    apply_result_prediction_aggregation(
1224                        plan,
1225                        controllers,
1226                        &task,
1227                        &mut result,
1228                        &resources,
1229                    )?;
1230                }
1231                if let Some(nested) = resources.nested_stacking.as_ref() {
1232                    attach_nested_stacking_input_lineage(&mut result, plan, &task, ctx, nested)?;
1233                } else {
1234                    attach_coordinator_input_lineage(
1235                        &mut result,
1236                        plan,
1237                        &task.node_plan.node_id,
1238                        &input_lineage,
1239                    )?;
1240                }
1241                if let Some(store) = resources.artifact_store.as_deref_mut() {
1242                    if scope.phase == Phase::Refit {
1243                        store.capture_refit_artifacts(&task, &result)?;
1244                    }
1245                }
1246                for prediction in &result.predictions {
1247                    ctx.prediction_store.append(prediction.clone())?;
1248                }
1249                for prediction in &result.aggregated_predictions {
1250                    ctx.aggregated_prediction_store.append(prediction.clone())?;
1251                }
1252                apply_result_scoring(
1253                    &result,
1254                    &mut ctx.score_collector,
1255                    &mut ctx.regression_target_records,
1256                )?;
1257                ctx.lineage.record(result.lineage.clone())?;
1258                let data_views = derive_output_data_views(plan, &task, &result)?;
1259                output_handles.insert(node_id.clone(), result.outputs.clone());
1260                output_data_views.insert(node_id.clone(), data_views);
1261                input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1262                results.push(result);
1263            }
1264        }
1265
1266        Ok(results)
1267    }
1268}
1269
1270impl ParallelScheduler {
1271    pub fn execute_phase(
1272        &self,
1273        plan: &ExecutionPlan,
1274        controllers: &RuntimeControllerRegistry,
1275        ctx: &mut RunContext,
1276        phase: Phase,
1277    ) -> Result<Vec<NodeResult>> {
1278        plan.validate()?;
1279        let variant_id = ctx.variant_id.clone();
1280        let seed_root = ctx.root_seed;
1281        self.execute_phase_scope(
1282            plan,
1283            controllers,
1284            ctx,
1285            PhaseScope {
1286                phase,
1287                variant_id,
1288                variant: None,
1289                fold_id: None,
1290                seed_root,
1291            },
1292            PhaseScopeResources::default(),
1293        )
1294    }
1295
1296    pub fn execute_phase_with_data_provider(
1297        &self,
1298        plan: &ExecutionPlan,
1299        controllers: &RuntimeControllerRegistry,
1300        data_provider: &dyn RuntimeDataProvider,
1301        ctx: &mut RunContext,
1302        phase: Phase,
1303    ) -> Result<Vec<NodeResult>> {
1304        plan.validate()?;
1305        let variant_id = ctx.variant_id.clone();
1306        let seed_root = ctx.root_seed;
1307        self.execute_phase_scope(
1308            plan,
1309            controllers,
1310            ctx,
1311            PhaseScope {
1312                phase,
1313                variant_id,
1314                variant: None,
1315                fold_id: None,
1316                seed_root,
1317            },
1318            PhaseScopeResources {
1319                data_provider: Some(data_provider),
1320                ..Default::default()
1321            },
1322        )
1323    }
1324
1325    pub fn execute_campaign_phase(
1326        &self,
1327        plan: &ExecutionPlan,
1328        controllers: &RuntimeControllerRegistry,
1329        ctx: &mut RunContext,
1330        phase: Phase,
1331    ) -> Result<Vec<NodeResult>> {
1332        plan.validate()?;
1333        if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1334            return Err(DagMlError::RuntimeValidation(
1335                "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1336                    .to_string(),
1337            ));
1338        }
1339        let mut results = Vec::new();
1340        let fold_ids = if phase == Phase::FitCv {
1341            plan.fold_set
1342                .as_ref()
1343                .map(|fold_set| {
1344                    fold_set
1345                        .folds
1346                        .iter()
1347                        .map(|fold| Some(fold.fold_id.clone()))
1348                        .collect::<Vec<_>>()
1349                })
1350                .unwrap_or_else(|| vec![None])
1351        } else {
1352            vec![None]
1353        };
1354        for variant in &plan.variants {
1355            if ctx
1356                .variant_id
1357                .as_ref()
1358                .is_some_and(|requested| requested != &variant.variant_id)
1359            {
1360                continue;
1361            }
1362            for fold_id in &fold_ids {
1363                let seed_root = variant.seed.or(ctx.root_seed);
1364                results.extend(self.execute_phase_scope(
1365                    plan,
1366                    controllers,
1367                    ctx,
1368                    PhaseScope {
1369                        phase,
1370                        variant_id: Some(variant.variant_id.clone()),
1371                        variant: Some(VariantExecutionSpec::from_plan(variant)),
1372                        fold_id: fold_id.clone(),
1373                        seed_root,
1374                    },
1375                    PhaseScopeResources::default(),
1376                )?);
1377            }
1378        }
1379        Ok(results)
1380    }
1381
1382    pub fn execute_campaign_phase_with_data_provider(
1383        &self,
1384        plan: &ExecutionPlan,
1385        controllers: &RuntimeControllerRegistry,
1386        data_provider: &dyn RuntimeDataProvider,
1387        ctx: &mut RunContext,
1388        phase: Phase,
1389    ) -> Result<Vec<NodeResult>> {
1390        plan.validate()?;
1391        if phase == Phase::FitCv {
1392            ctx.configure_global_oof_aggregation(plan, data_provider)?;
1393        }
1394        if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1395            return Err(DagMlError::RuntimeValidation(
1396                "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1397                    .to_string(),
1398            ));
1399        }
1400        let mut results = Vec::new();
1401        let fold_ids = if phase == Phase::FitCv {
1402            plan.fold_set
1403                .as_ref()
1404                .map(|fold_set| {
1405                    fold_set
1406                        .folds
1407                        .iter()
1408                        .map(|fold| Some(fold.fold_id.clone()))
1409                        .collect::<Vec<_>>()
1410                })
1411                .unwrap_or_else(|| vec![None])
1412        } else {
1413            vec![None]
1414        };
1415        for variant in &plan.variants {
1416            if ctx
1417                .variant_id
1418                .as_ref()
1419                .is_some_and(|requested| requested != &variant.variant_id)
1420            {
1421                continue;
1422            }
1423            for fold_id in &fold_ids {
1424                let seed_root = variant.seed.or(ctx.root_seed);
1425                results.extend(self.execute_phase_scope(
1426                    plan,
1427                    controllers,
1428                    ctx,
1429                    PhaseScope {
1430                        phase,
1431                        variant_id: Some(variant.variant_id.clone()),
1432                        variant: Some(VariantExecutionSpec::from_plan(variant)),
1433                        fold_id: fold_id.clone(),
1434                        seed_root,
1435                    },
1436                    PhaseScopeResources {
1437                        data_provider: Some(data_provider),
1438                        ..Default::default()
1439                    },
1440                )?);
1441            }
1442        }
1443        Ok(results)
1444    }
1445
1446    pub fn execute_campaign_phase_with_data_provider_and_artifact_store(
1447        &self,
1448        plan: &ExecutionPlan,
1449        controllers: &RuntimeControllerRegistry,
1450        data_provider: &dyn RuntimeDataProvider,
1451        artifact_store: &mut InMemoryArtifactStore,
1452        ctx: &mut RunContext,
1453        phase: Phase,
1454    ) -> Result<Vec<NodeResult>> {
1455        plan.validate()?;
1456        if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1457            return Err(DagMlError::RuntimeValidation(
1458                "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1459                    .to_string(),
1460            ));
1461        }
1462        let mut results = Vec::new();
1463        let fold_ids = if phase == Phase::FitCv {
1464            plan.fold_set
1465                .as_ref()
1466                .map(|fold_set| {
1467                    fold_set
1468                        .folds
1469                        .iter()
1470                        .map(|fold| Some(fold.fold_id.clone()))
1471                        .collect::<Vec<_>>()
1472                })
1473                .unwrap_or_else(|| vec![None])
1474        } else {
1475            vec![None]
1476        };
1477        for variant in &plan.variants {
1478            if ctx
1479                .variant_id
1480                .as_ref()
1481                .is_some_and(|requested| requested != &variant.variant_id)
1482            {
1483                continue;
1484            }
1485            for fold_id in &fold_ids {
1486                let seed_root = variant.seed.or(ctx.root_seed);
1487                results.extend(self.execute_phase_scope(
1488                    plan,
1489                    controllers,
1490                    ctx,
1491                    PhaseScope {
1492                        phase,
1493                        variant_id: Some(variant.variant_id.clone()),
1494                        variant: Some(VariantExecutionSpec::from_plan(variant)),
1495                        fold_id: fold_id.clone(),
1496                        seed_root,
1497                    },
1498                    PhaseScopeResources {
1499                        data_provider: Some(data_provider),
1500                        artifact_store: Some(&mut *artifact_store),
1501                        ..Default::default()
1502                    },
1503                )?);
1504            }
1505        }
1506        Ok(results)
1507    }
1508
1509    pub fn execute_bundle_replay(
1510        &self,
1511        replay: BundleReplayExecution<'_>,
1512        ctx: &mut RunContext,
1513    ) -> Result<Vec<NodeResult>> {
1514        replay.bundle.validate_against_plan(replay.plan)?;
1515        replay
1516            .replay_request
1517            .validate_for_bundle_with_prediction_cache_store(
1518                replay.bundle,
1519                replay.prediction_cache_store.is_some(),
1520            )?;
1521        replay
1522            .bundle
1523            .validate_replay_envelopes(replay.data_envelopes)?;
1524        let prediction_cache_contracts = if replay.replay_request.phase == Phase::Refit {
1525            Some(replay_prediction_cache_contracts(replay.bundle)?)
1526        } else {
1527            None
1528        };
1529        if replay.replay_request.phase == Phase::Refit {
1530            preload_replay_prediction_cache_store(
1531                replay.bundle,
1532                replay.prediction_cache_store,
1533                ctx,
1534            )?;
1535        }
1536        let replay_artifacts = materialize_replay_artifact_handles(
1537            replay.plan,
1538            replay.bundle,
1539            replay.replay_request,
1540            replay.artifact_store,
1541            ctx,
1542        )?;
1543        let selected_variant = replay
1544            .bundle
1545            .selected_variant_id
1546            .as_ref()
1547            .map(|selected| {
1548                replay
1549                    .plan
1550                    .variants
1551                    .iter()
1552                    .find(|variant| &variant.variant_id == selected)
1553                    .map(VariantExecutionSpec::from_plan)
1554                    .ok_or_else(|| {
1555                        DagMlError::RuntimeValidation(format!(
1556                            "bundle `{}` selected unknown variant `{selected}`",
1557                            replay.bundle.bundle_id
1558                        ))
1559                    })
1560            })
1561            .transpose()?;
1562        let seed_root = selected_variant
1563            .as_ref()
1564            .and_then(|variant| variant.seed)
1565            .or(ctx.root_seed);
1566
1567        self.execute_phase_scope(
1568            replay.plan,
1569            replay.controllers,
1570            ctx,
1571            PhaseScope {
1572                phase: replay.replay_request.phase,
1573                variant_id: replay.bundle.selected_variant_id.clone(),
1574                variant: selected_variant,
1575                fold_id: None,
1576                seed_root,
1577            },
1578            PhaseScopeResources {
1579                data_provider: Some(replay.data_provider),
1580                replay_artifact_handles: Some(&replay_artifacts.handles),
1581                replay_artifact_inputs: Some(&replay_artifacts.inputs),
1582                replay_bundle_id: Some(&replay.bundle.bundle_id),
1583                data_envelopes: Some(replay.data_envelopes),
1584                prediction_cache_store: replay.prediction_cache_store,
1585                prediction_cache_contracts: prediction_cache_contracts.as_ref(),
1586                ..Default::default()
1587            },
1588        )
1589    }
1590
1591    fn execute_phase_scope(
1592        &self,
1593        plan: &ExecutionPlan,
1594        controllers: &RuntimeControllerRegistry,
1595        ctx: &mut RunContext,
1596        scope: PhaseScope,
1597        mut resources: PhaseScopeResources<'_>,
1598    ) -> Result<Vec<NodeResult>> {
1599        // Hold the phase span on the scheduler thread, and clone it into each
1600        // worker so worker-thread telemetry nests under the phase (tracing spans
1601        // are thread-local and do not auto-propagate across `thread::scope`).
1602        let phase_span = crate::observability::phase_span(
1603            ctx.run_id.as_str(),
1604            plan.id.as_str(),
1605            scope.phase.as_str(),
1606            scope.variant_id.as_ref().map(VariantId::as_str),
1607            scope.fold_id.as_ref().map(FoldId::as_str),
1608        );
1609        let _phase_entered = phase_span.clone().entered();
1610        // Borrowed for the `thread::scope` below; workers join before it ends.
1611        let plan_id = plan.id.as_str();
1612        plan.validate_parallel_controller_capabilities(self.max_workers, scope.phase)?;
1613        let mut results = Vec::new();
1614        let mut output_handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
1615        let mut output_data_views =
1616            BTreeMap::<NodeId, BTreeMap<String, DataProviderViewSpec>>::new();
1617        let mut input_lineage = BTreeMap::<NodeId, LineageId>::new();
1618
1619        for level in plan.node_parallel_levels_for_phase(scope.phase)? {
1620            let mut prepared = Vec::<PreparedNodeTask>::new();
1621            // Cross-branch merge nodes (concat or late-fusion) are not controller
1622            // tasks: they read the upstream branch OOF blocks from the prediction
1623            // store and reassemble them on the scheduler thread (no worker), AFTER
1624            // this level's worker tasks have populated the store. They are in a
1625            // later level than their branches, so the store already holds the
1626            // branch OOF by the time we reassemble — see `reassemble_branch_merge`.
1627            let mut merge_nodes = Vec::<(NodeId, MergeReduction)>::new();
1628            for node_id in &level {
1629                let node_plan = plan
1630                    .node_plans
1631                    .get(node_id)
1632                    .expect("execution plan was validated");
1633                if let Some(reduction) = merge_reduction_mode(plan, node_plan) {
1634                    merge_nodes.push((node_id.clone(), reduction));
1635                    continue;
1636                }
1637                let collected_inputs = collect_input_handles(
1638                    plan,
1639                    node_plan,
1640                    &output_handles,
1641                    &output_data_views,
1642                    &resources,
1643                    ctx,
1644                    &scope,
1645                )?;
1646                if collected_inputs.skip_node {
1647                    continue;
1648                }
1649                let mut input_handles = collected_inputs.handles;
1650                let mut artifact_inputs = BTreeMap::new();
1651                if let Some(node_artifact_handles) = resources
1652                    .replay_artifact_handles
1653                    .and_then(|handles| handles.get(node_id))
1654                {
1655                    for (key, handle) in node_artifact_handles {
1656                        if input_handles.insert(key.clone(), handle.clone()).is_some() {
1657                            return Err(DagMlError::RuntimeValidation(format!(
1658                                "node `{node_id}` received duplicate replay artifact input `{key}`"
1659                            )));
1660                        }
1661                    }
1662                }
1663                if let Some(node_artifact_inputs) = resources
1664                    .replay_artifact_inputs
1665                    .and_then(|inputs| inputs.get(node_id))
1666                {
1667                    for (key, spec) in node_artifact_inputs {
1668                        if artifact_inputs.insert(key.clone(), spec.clone()).is_some() {
1669                            return Err(DagMlError::RuntimeValidation(format!(
1670                                "node `{node_id}` received duplicate replay artifact metadata `{key}`"
1671                            )));
1672                        }
1673                    }
1674                }
1675                let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1676                let inner_fold_set = inner_fold_set_for_scope(
1677                    &plan.campaign,
1678                    plan.fold_set.as_ref(),
1679                    node_plan,
1680                    &scope,
1681                )?;
1682                let fit_influence = fit_influence_task_for_node(
1683                    plan,
1684                    &task_node_plan,
1685                    &collected_inputs.data_views,
1686                )?;
1687                prepared.push(PreparedNodeTask {
1688                    node_id: node_id.clone(),
1689                    task: NodeTask {
1690                        inner_fold_set,
1691                        run_id: ctx.run_id.clone(),
1692                        node_plan: task_node_plan.clone(),
1693                        phase: scope.phase,
1694                        variant_id: scope.variant_id.clone(),
1695                        variant: scope.variant.clone(),
1696                        fold_id: scope.fold_id.clone(),
1697                        branch_path: Vec::new(),
1698                        input_handles,
1699                        data_views: collected_inputs.data_views,
1700                        prediction_inputs: collected_inputs.prediction_inputs,
1701                        artifact_inputs,
1702                        required_loss_attestations: NodeTask::required_loss_attestations_for(
1703                            &task_node_plan,
1704                            scope.phase,
1705                        )?,
1706                        fit_influence,
1707                        seed: derive_task_seed(
1708                            scope.seed_root,
1709                            scope.variant_id.as_ref(),
1710                            scope.fold_id.as_ref(),
1711                            &task_node_plan,
1712                            scope.phase,
1713                        ),
1714                    },
1715                });
1716            }
1717
1718            for chunk in prepared.chunks(self.max_workers) {
1719                let chunk_results = std::thread::scope(
1720                    |thread_scope| -> Result<Vec<NodeResult>> {
1721                        let mut handles = Vec::with_capacity(chunk.len());
1722                        for prepared_task in chunk {
1723                            let controller = controllers
1724                                .get(&prepared_task.task.node_plan.controller_id)
1725                                .ok_or_else(|| {
1726                                    DagMlError::RuntimeValidation(format!(
1727                                        "runtime controller `{}` is not registered",
1728                                        prepared_task.task.node_plan.controller_id
1729                                    ))
1730                                })?;
1731                            let worker_span = phase_span.clone();
1732                            handles.push(thread_scope.spawn(move || {
1733                                let _worker_span = worker_span.entered();
1734                                let _node_span = crate::observability::node_span(
1735                                    prepared_task.task.run_id.as_str(),
1736                                    plan_id,
1737                                    prepared_task.task.phase.as_str(),
1738                                    prepared_task.task.node_plan.node_id.as_str(),
1739                                    prepared_task.task.node_plan.controller_id.as_str(),
1740                                )
1741                                .entered();
1742                                let mut result =
1743                                    if prepared_task.task.node_plan.kind == NodeKind::Tuner {
1744                                        return Err(DagMlError::RuntimeValidation(format!(
1745                                            "tuner node `{}` requires execute_hpo_campaign with an explicit RuntimeHpoExecutionContext",
1746                                            prepared_task.task.node_plan.node_id
1747                                        )));
1748                                    } else {
1749                                        // A provider-aware controller may require a
1750                                        // non-Sync host provider.  Parallel native
1751                                        // Methods PLS is deliberately refused by its
1752                                        // HPO preflight; ordinary controllers keep
1753                                        // their opaque-handle invocation here.
1754                                        controller.invoke(&prepared_task.task)?
1755                                    };
1756                                record_fit_influence_diagnostic(&prepared_task.task, &mut result);
1757                                normalize_result_prediction_ports(
1758                                    plan,
1759                                    &prepared_task.task,
1760                                    &mut result,
1761                                )?;
1762                                result.validate_for_task(&prepared_task.task)?;
1763                                Ok(result)
1764                            }));
1765                        }
1766                        handles
1767                            .into_iter()
1768                            .map(|handle| {
1769                                handle.join().map_err(|_| {
1770                                    DagMlError::RuntimeValidation(
1771                                        "parallel scheduler worker panicked".to_string(),
1772                                    )
1773                                })?
1774                            })
1775                            .collect()
1776                    },
1777                )?;
1778
1779                for (prepared_task, mut result) in chunk.iter().zip(chunk_results) {
1780                    apply_result_prediction_aggregation(
1781                        plan,
1782                        controllers,
1783                        &prepared_task.task,
1784                        &mut result,
1785                        &resources,
1786                    )?;
1787                    if let Some(nested) = resources.nested_stacking.as_ref() {
1788                        attach_nested_stacking_input_lineage(
1789                            &mut result,
1790                            plan,
1791                            &prepared_task.task,
1792                            ctx,
1793                            nested,
1794                        )?;
1795                    } else {
1796                        attach_coordinator_input_lineage(
1797                            &mut result,
1798                            plan,
1799                            &prepared_task.task.node_plan.node_id,
1800                            &input_lineage,
1801                        )?;
1802                    }
1803                    if let Some(store) = resources.artifact_store.as_deref_mut() {
1804                        if scope.phase == Phase::Refit {
1805                            store.capture_refit_artifacts(&prepared_task.task, &result)?;
1806                        }
1807                    }
1808                    for prediction in &result.predictions {
1809                        ctx.prediction_store.append(prediction.clone())?;
1810                    }
1811                    for prediction in &result.aggregated_predictions {
1812                        ctx.aggregated_prediction_store.append(prediction.clone())?;
1813                    }
1814                    apply_result_scoring(
1815                        &result,
1816                        &mut ctx.score_collector,
1817                        &mut ctx.regression_target_records,
1818                    )?;
1819                    ctx.lineage.record(result.lineage.clone())?;
1820                    let data_views = derive_output_data_views(plan, &prepared_task.task, &result)?;
1821                    output_handles.insert(prepared_task.node_id.clone(), result.outputs.clone());
1822                    output_data_views.insert(prepared_task.node_id.clone(), data_views);
1823                    input_lineage.insert(
1824                        prepared_task.node_id.clone(),
1825                        result.lineage.record_id.clone(),
1826                    );
1827                    results.push(result);
1828                }
1829            }
1830
1831            // Reassemble any cross-branch merge nodes in this level now that the
1832            // level's worker tasks have populated the prediction store. Merge nodes
1833            // sit in a later level than the branches they consume, so the upstream
1834            // branch OOF is already present.
1835            for (node_id, reduction) in &merge_nodes {
1836                let node_plan = plan
1837                    .node_plans
1838                    .get(node_id)
1839                    .expect("execution plan was validated");
1840                if let Some(mut result) =
1841                    reassemble_branch_merge(plan, node_plan, ctx, &scope, *reduction)?
1842                {
1843                    let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1844                    let task = NodeTask {
1845                        inner_fold_set: None,
1846                        run_id: ctx.run_id.clone(),
1847                        node_plan: task_node_plan.clone(),
1848                        phase: scope.phase,
1849                        variant_id: scope.variant_id.clone(),
1850                        variant: scope.variant.clone(),
1851                        fold_id: scope.fold_id.clone(),
1852                        branch_path: Vec::new(),
1853                        input_handles: BTreeMap::new(),
1854                        data_views: BTreeMap::new(),
1855                        prediction_inputs: BTreeMap::new(),
1856                        artifact_inputs: BTreeMap::new(),
1857                        required_loss_attestations: NodeTask::required_loss_attestations_for(
1858                            &task_node_plan,
1859                            scope.phase,
1860                        )?,
1861                        fit_influence: FitInfluenceTask::default(),
1862                        seed: None,
1863                    };
1864                    normalize_result_prediction_ports(plan, &task, &mut result)?;
1865                    result.validate_for_task(&task)?;
1866                    for prediction in &result.predictions {
1867                        ctx.prediction_store.append(prediction.clone())?;
1868                    }
1869                    apply_result_scoring(
1870                        &result,
1871                        &mut ctx.score_collector,
1872                        &mut ctx.regression_target_records,
1873                    )?;
1874                    ctx.lineage.record(result.lineage.clone())?;
1875                    output_handles.insert(node_id.clone(), result.outputs.clone());
1876                    input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1877                    results.push(result);
1878                }
1879            }
1880        }
1881
1882        Ok(results)
1883    }
1884}
1885
1886#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1887enum HpoTrialTerminalState {
1888    Completed,
1889    Pruned,
1890    Failed,
1891}
1892
1893fn validate_hpo_checkpoint_result(
1894    checkpoint: &RuntimeHpoCheckpointResult,
1895    hpo: &RuntimeHpoExecutionContext,
1896    trial_variants: &BTreeMap<i64, VariantId>,
1897    terminal_trials: &BTreeMap<i64, HpoTrialTerminalState>,
1898    history_at_start: u32,
1899) -> Result<()> {
1900    checkpoint.artifact.validate().map_err(|error| {
1901        DagMlError::RuntimeValidation(format!(
1902            "runtime HPO checkpoint artifact is invalid: {error}"
1903        ))
1904    })?;
1905    if checkpoint.operation_id != hpo.operation_id
1906        || checkpoint.controller_id != hpo.controller_id
1907        || checkpoint.target_node_id != hpo.target_node_id
1908        || checkpoint.provenance != hpo.provenance
1909    {
1910        return Err(DagMlError::RuntimeValidation(
1911            "runtime HPO checkpoint provenance does not exactly match its execution context"
1912                .to_string(),
1913        ));
1914    }
1915    let proposed_count = u32::try_from(trial_variants.len()).map_err(|_| {
1916        DagMlError::RuntimeValidation(
1917            "runtime HPO scheduler proposal count does not fit u32".to_string(),
1918        )
1919    })?;
1920    if checkpoint.trial_history_len != hpo.trial_budget_total
1921        || checkpoint.trial_history_len < history_at_start
1922        || checkpoint.trial_history_len - history_at_start != proposed_count
1923    {
1924        return Err(DagMlError::RuntimeValidation(
1925            "runtime HPO checkpoint native history is inconsistent with scheduler-observed trials"
1926                .to_string(),
1927        ));
1928    }
1929    if checkpoint.artifact.binding.controller_id != hpo.controller_id.as_str()
1930        || checkpoint.artifact.binding.controller_id != hpo.study.controller_id
1931        || checkpoint.artifact.binding.study_id != hpo.study.study_id
1932        || checkpoint.artifact.methods_abi != hpo.study.methods_abi
1933    {
1934        return Err(DagMlError::RuntimeValidation(
1935            "runtime HPO checkpoint binding/controller/study does not match the active tuner"
1936                .to_string(),
1937        ));
1938    }
1939    let expected_search_space = hpo.study.search_space.fingerprint().map_err(|error| {
1940        DagMlError::RuntimeValidation(format!(
1941            "runtime HPO cannot fingerprint the configured search space: {error}"
1942        ))
1943    })?;
1944    if checkpoint.artifact.binding.search_space_fingerprint != expected_search_space {
1945        return Err(DagMlError::RuntimeValidation(
1946            "runtime HPO checkpoint search-space binding does not match the active study"
1947                .to_string(),
1948        ));
1949    }
1950
1951    let completed_trial_ids = terminal_trials
1952        .iter()
1953        .filter_map(|(trial_id, state)| {
1954            (*state == HpoTrialTerminalState::Completed).then_some(*trial_id)
1955        })
1956        .collect::<BTreeSet<_>>();
1957    let mut proposal_trial_ids = BTreeSet::new();
1958    for proposal in &checkpoint.completed_proposals {
1959        if !proposal_trial_ids.insert(proposal.trial_id) {
1960            return Err(DagMlError::RuntimeValidation(format!(
1961                "runtime HPO checkpoint has duplicate completed proposal for trial `{}`",
1962                proposal.trial_id
1963            )));
1964        }
1965        if trial_variants.get(&proposal.trial_id) != Some(&proposal.variant.variant_id) {
1966            return Err(DagMlError::RuntimeValidation(format!(
1967                "runtime HPO checkpoint proposal for trial `{}` does not exactly match its scheduler proposal",
1968                proposal.trial_id
1969            )));
1970        }
1971    }
1972    if proposal_trial_ids != completed_trial_ids {
1973        return Err(DagMlError::RuntimeValidation(
1974            "runtime HPO checkpoint proposals must cover exactly the completed trials".to_string(),
1975        ));
1976    }
1977
1978    let mut report_trial_ids = BTreeSet::new();
1979    for completed in &checkpoint.completed_reports {
1980        if !report_trial_ids.insert(completed.trial_id) {
1981            return Err(DagMlError::RuntimeValidation(format!(
1982                "runtime HPO checkpoint has duplicate completed report for trial `{}`",
1983                completed.trial_id
1984            )));
1985        }
1986        if trial_variants.get(&completed.trial_id) != Some(&completed.variant_id)
1987            || !proposal_trial_ids.contains(&completed.trial_id)
1988        {
1989            return Err(DagMlError::RuntimeValidation(format!(
1990                "runtime HPO checkpoint report for trial `{}` does not match a completed proposal",
1991                completed.trial_id
1992            )));
1993        }
1994        let report = &completed.report;
1995        if report.producer_node != hpo.selection.producer_node
1996            || report.producer_port.as_deref() != Some(hpo.selection.producer_port.as_str())
1997            || report.partition != PredictionPartition::Validation
1998            || report
1999                .fold_id
2000                .as_ref()
2001                .is_none_or(|fold| fold.as_str() != "avg")
2002            || report.variant_id.as_ref() != Some(&completed.variant_id)
2003            || !report
2004                .metrics
2005                .get(hpo.selection.metric.name())
2006                .is_some_and(|score| score.is_finite())
2007        {
2008            return Err(DagMlError::RuntimeValidation(format!(
2009                "runtime HPO checkpoint report for trial `{}` is not its one finite target OOF average",
2010                completed.trial_id
2011            )));
2012        }
2013    }
2014    if report_trial_ids != completed_trial_ids {
2015        return Err(DagMlError::RuntimeValidation(
2016            "runtime HPO checkpoint reports must cover exactly one OOF average per completed trial"
2017                .to_string(),
2018        ));
2019    }
2020    Ok(())
2021}
2022
2023pub(crate) struct PreparedNodeTask {
2024    pub(crate) node_id: NodeId,
2025    pub(crate) task: NodeTask,
2026}
2027
2028// This module stays adjacent to the scheduler-owned task preparation it
2029// exercises; the remaining helpers below are shared by both schedulers.
2030#[cfg(test)]
2031#[allow(clippy::items_after_test_module)]
2032mod hpo_scheduler_tests {
2033    use std::collections::{BTreeMap, BTreeSet};
2034    use std::sync::{Arc, Mutex};
2035
2036    use sha2::{Digest, Sha256};
2037
2038    use super::*;
2039    use crate::controller::{
2040        ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest,
2041        ControllerRegistry, RngPolicy,
2042    };
2043    use crate::data::InMemoryDataProvider;
2044    use crate::fold::{FoldAssignment, FoldPartitionMode};
2045    use crate::graph::{GraphInterface, GraphSpec, NodeSpec, PortSchema, PortSpec};
2046    use crate::hpo::{
2047        HpoDirection, HpoMetric, HpoOptimizerConfig, HpoParameter, HpoPruner, HpoSampler,
2048        HpoSearchSpace, HpoStudyBinding, MethodsHpoStudyConfig, N4moptCheckpointArtifact,
2049        N4MOPT_ARTIFACT_KIND, N4MOPT_CHECKPOINT_SCHEMA_VERSION, N4MOPT_FORMAT,
2050    };
2051    use crate::metrics::RegressionTargetBlock;
2052    use crate::oof::PredictionBlock;
2053    use crate::plan::{build_execution_plan, SplitInvocation};
2054
2055    struct HpoTestModel {
2056        id: ControllerId,
2057        trace: Arc<Mutex<Vec<String>>>,
2058    }
2059
2060    impl RuntimeController for HpoTestModel {
2061        fn controller_id(&self) -> &ControllerId {
2062            &self.id
2063        }
2064
2065        fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
2066            self.trace.lock().unwrap().push("model_cv".to_string());
2067            let sample_id = match task.fold_id.as_ref().map(FoldId::as_str) {
2068                Some("fold:0") => SampleId::new("sample:one").unwrap(),
2069                Some("fold:1") => SampleId::new("sample:two").unwrap(),
2070                other => {
2071                    return Err(DagMlError::RuntimeValidation(format!(
2072                        "HPO test model received unexpected fold {other:?}"
2073                    )));
2074                }
2075            };
2076            Ok(NodeResult {
2077                schema_version: None,
2078                node_id: task.node_plan.node_id.clone(),
2079                outputs: BTreeMap::from([(
2080                    "prediction".to_string(),
2081                    HandleRef {
2082                        handle: 2,
2083                        kind: HandleKind::Prediction,
2084                        owner_controller: self.id.clone(),
2085                    },
2086                )]),
2087                predictions: vec![PredictionBlock {
2088                    prediction_id: Some(format!("prediction:{}", task.fold_id.as_ref().unwrap())),
2089                    producer_node: task.node_plan.node_id.clone(),
2090                    producer_port: None,
2091                    partition: PredictionPartition::Validation,
2092                    fold_id: task.fold_id.clone(),
2093                    sample_ids: vec![sample_id.clone()],
2094                    values: vec![vec![1.0]],
2095                    target_names: vec!["target".to_string()],
2096                }],
2097                observation_predictions: Vec::new(),
2098                aggregated_predictions: Vec::new(),
2099                explanations: Vec::new(),
2100                shape_deltas: Vec::new(),
2101                artifacts: Vec::new(),
2102                artifact_handles: BTreeMap::new(),
2103                fit_influence_diagnostics: Vec::new(),
2104                regression_targets: vec![RegressionTargetBlock {
2105                    level: PredictionLevel::Sample,
2106                    unit_ids: vec![PredictionUnitId::Sample(sample_id)],
2107                    values: vec![vec![1.0]],
2108                    target_names: vec!["target".to_string()],
2109                }],
2110                lineage: LineageRecord {
2111                    record_id: LineageId::new(format!(
2112                        "lineage:hpo-model:{}",
2113                        task.fold_id.as_ref().unwrap()
2114                    ))
2115                    .unwrap(),
2116                    run_id: task.run_id.clone(),
2117                    node_id: task.node_plan.node_id.clone(),
2118                    phase: task.phase,
2119                    controller_id: self.id.clone(),
2120                    controller_version: task.node_plan.controller_version.clone(),
2121                    variant_id: task.variant_id.clone(),
2122                    fold_id: task.fold_id.clone(),
2123                    branch_path: Vec::new(),
2124                    input_lineage: Vec::new(),
2125                    artifact_refs: Vec::new(),
2126                    params_fingerprint: task.node_plan.params_fingerprint.clone(),
2127                    data_model_shape_fingerprint: None,
2128                    aggregation_policy_fingerprint: None,
2129                    seed: task.seed,
2130                    unsafe_flags: BTreeSet::new(),
2131                    metrics: BTreeMap::new(),
2132                    loss_attestations: Vec::new(),
2133                    early_stopping_records: Vec::new(),
2134                },
2135            })
2136        }
2137    }
2138
2139    struct HpoTestTuner {
2140        id: ControllerId,
2141        trace: Arc<Mutex<Vec<String>>>,
2142        history_len: u32,
2143        proposal_count: u32,
2144    }
2145
2146    struct HpoTestSession {
2147        proposals: Vec<RuntimeHpoProposal>,
2148        trace: Arc<Mutex<Vec<String>>>,
2149        checkpoint: N4moptCheckpointArtifact,
2150        history_len: u32,
2151        completed: Option<(i64, f64)>,
2152    }
2153
2154    impl RuntimeController for HpoTestTuner {
2155        fn controller_id(&self) -> &ControllerId {
2156            &self.id
2157        }
2158
2159        fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
2160            Err(DagMlError::RuntimeValidation(format!(
2161                "HPO test tuner `{}` was dispatched through generic invoke",
2162                task.node_plan.node_id
2163            )))
2164        }
2165
2166        fn create_tuner_session(
2167            &self,
2168            task: &RuntimeHpoCampaignTask,
2169            context: &RuntimeHpoExecutionContext,
2170        ) -> Result<Box<dyn RuntimeTunerSession>> {
2171            assert_eq!(task.operation_id, context.operation_id);
2172            self.trace
2173                .lock()
2174                .unwrap()
2175                .push("session_factory".to_string());
2176            let payload = vec![7_u8];
2177            let proposals = (0..self.proposal_count)
2178                .map(|offset| {
2179                    let trial_id = i64::from(self.history_len + offset + 1);
2180                    let mut variant = context.base_variant.clone();
2181                    if self.history_len != 0 || self.proposal_count != 1 {
2182                        variant.variant_id = VariantId::new(format!("hpo:trial:{trial_id}"))
2183                            .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?;
2184                        variant.fingerprint = format!("hpo-test-{trial_id}");
2185                    }
2186                    Ok(RuntimeHpoProposal { trial_id, variant })
2187                })
2188                .collect::<Result<Vec<_>>>()?;
2189            Ok(Box::new(HpoTestSession {
2190                proposals: proposals.into_iter().rev().collect(),
2191                trace: Arc::clone(&self.trace),
2192                history_len: self.history_len,
2193                completed: None,
2194                checkpoint: N4moptCheckpointArtifact {
2195                    schema_version: N4MOPT_CHECKPOINT_SCHEMA_VERSION,
2196                    artifact_kind: N4MOPT_ARTIFACT_KIND.to_string(),
2197                    format: N4MOPT_FORMAT.to_string(),
2198                    binding: HpoStudyBinding {
2199                        controller_id: context.study.controller_id.clone(),
2200                        study_id: context.study.study_id.clone(),
2201                        search_space_fingerprint: context
2202                            .study
2203                            .search_space
2204                            .fingerprint()
2205                            .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?,
2206                        optimizer_fingerprint: "optimizer:test".to_string(),
2207                    },
2208                    methods_abi: context.study.methods_abi.clone(),
2209                    payload_sha256: format!("{:x}", Sha256::digest(&payload)),
2210                    opaque_payload: payload,
2211                },
2212            }))
2213        }
2214    }
2215
2216    impl RuntimeTunerSession for HpoTestSession {
2217        fn trial_history_len(&self) -> Result<u32> {
2218            Ok(self.history_len)
2219        }
2220
2221        fn ask(&mut self) -> Result<Option<RuntimeHpoProposal>> {
2222            self.trace.lock().unwrap().push("ask".to_string());
2223            let proposal = self.proposals.pop();
2224            if proposal.is_some() {
2225                self.history_len += 1;
2226            }
2227            Ok(proposal)
2228        }
2229
2230        fn report_intermediate(
2231            &mut self,
2232            intermediate: RuntimeHpoIntermediate,
2233        ) -> Result<RuntimeHpoIntermediateOutcome> {
2234            assert_eq!(intermediate.step, 0);
2235            assert!(intermediate.score.is_finite());
2236            self.trace.lock().unwrap().push("intermediate".to_string());
2237            Ok(RuntimeHpoIntermediateOutcome::Continue)
2238        }
2239
2240        fn tell(&mut self, trial_id: i64, terminal: RuntimeHpoTerminal) -> Result<()> {
2241            assert!(trial_id > 0);
2242            assert!(
2243                matches!(terminal, RuntimeHpoTerminal::Completed { score } if score.is_finite())
2244            );
2245            self.trace.lock().unwrap().push("tell".to_string());
2246            if let RuntimeHpoTerminal::Completed { score } = terminal {
2247                self.completed = Some((trial_id, score));
2248            }
2249            Ok(())
2250        }
2251
2252        fn checkpoint(&self) -> Result<N4moptCheckpointArtifact> {
2253            self.trace.lock().unwrap().push("checkpoint".to_string());
2254            Ok(self.checkpoint.clone())
2255        }
2256
2257        fn incumbent(
2258            &self,
2259            variants: &BTreeMap<i64, VariantId>,
2260        ) -> Result<Option<RuntimeHpoIncumbent>> {
2261            let Some((trial_id, score)) = self.completed else {
2262                return Ok(None);
2263            };
2264            Ok(Some(RuntimeHpoIncumbent {
2265                trial_id,
2266                score,
2267                metric: "rmse".to_string(),
2268                direction: HpoDirection::Minimize,
2269                variant_id: variants.get(&trial_id).cloned().unwrap(),
2270            }))
2271        }
2272
2273        fn terminal_trial_snapshots(
2274            &self,
2275            variants: &BTreeMap<i64, VariantId>,
2276        ) -> Result<Vec<RuntimeHpoTerminalSnapshot>> {
2277            let (trial_id, score) = self.completed.ok_or_else(|| {
2278                DagMlError::RuntimeValidation("test HPO session has no completed trial".to_string())
2279            })?;
2280            Ok((1..=i64::from(self.history_len))
2281                .map(|id| {
2282                    let completed = id == trial_id;
2283                    RuntimeHpoTerminalSnapshot {
2284                        trial: crate::hpo::HpoTrial {
2285                            id,
2286                            ask_sequence: id,
2287                            terminal_sequence: Some(id),
2288                            parameters: BTreeMap::new(),
2289                            parameter_order: Vec::new(),
2290                            status: if completed {
2291                                crate::hpo::HpoTrialStatus::Completed
2292                            } else {
2293                                crate::hpo::HpoTrialStatus::Failed
2294                            },
2295                            score: completed.then_some(score),
2296                            rung: 0,
2297                            duration: 0.0,
2298                            intermediates: Vec::new(),
2299                            failure: (!completed).then(|| crate::hpo::HpoFailure {
2300                                code: "RESTORED_TEST_FAILURE".to_string(),
2301                                message: "synthetic restored terminal".to_string(),
2302                                retryable: false,
2303                            }),
2304                        },
2305                        variant_id: variants.get(&id).cloned(),
2306                    }
2307                })
2308                .collect())
2309        }
2310    }
2311
2312    fn node(id: &str, kind: NodeKind, outputs: Vec<PortSpec>) -> NodeSpec {
2313        NodeSpec {
2314            id: NodeId::new(id).unwrap(),
2315            kind,
2316            operator: None,
2317            params: BTreeMap::new(),
2318            ports: PortSchema {
2319                inputs: Vec::new(),
2320                outputs,
2321            },
2322            metadata: BTreeMap::new(),
2323            seed_label: None,
2324        }
2325    }
2326
2327    fn manifest(id: &str, kind: NodeKind) -> ControllerManifest {
2328        ControllerManifest {
2329            controller_id: ControllerId::new(id).unwrap(),
2330            controller_version: "test".to_string(),
2331            operator_kind: kind,
2332            priority: 0,
2333            supported_phases: BTreeSet::from([Phase::FitCv]),
2334            input_ports: Vec::new(),
2335            output_ports: Vec::new(),
2336            data_requirements: None,
2337            capabilities: BTreeSet::from([
2338                ControllerCapability::Deterministic,
2339                ControllerCapability::EmitsPredictions,
2340            ]),
2341            operator_selectors: Vec::new(),
2342            fit_scope: ControllerFitScope::FoldTrain,
2343            rng_policy: RngPolicy::UsesCoreSeed,
2344            artifact_policy: ArtifactPolicy::Serializable,
2345        }
2346    }
2347
2348    #[test]
2349    fn hpo_campaign_invokes_registered_session_and_routes_oof_feedback() {
2350        let target = NodeId::new("model:score").unwrap();
2351        let graph = GraphSpec {
2352            id: "graph:hpo.scheduler".to_string(),
2353            interface: GraphInterface::default(),
2354            nodes: vec![node(
2355                "model:score",
2356                NodeKind::Model,
2357                vec![PortSpec {
2358                    name: "prediction".to_string(),
2359                    kind: PortKind::Prediction,
2360                    representation: None,
2361                    cardinality: crate::graph::PortCardinality::One,
2362                    unit_level: None,
2363                    alignment_key: None,
2364                    target_level: None,
2365                    description: String::new(),
2366                }],
2367            )],
2368            edges: Vec::new(),
2369            search_space_fingerprint: None,
2370            metadata: BTreeMap::new(),
2371        };
2372        let fold_set = FoldSet {
2373            id: "folds:hpo".to_string(),
2374            sample_ids: vec![
2375                SampleId::new("sample:one").unwrap(),
2376                SampleId::new("sample:two").unwrap(),
2377            ],
2378            folds: vec![
2379                FoldAssignment {
2380                    fold_id: FoldId::new("fold:0").unwrap(),
2381                    train_sample_ids: vec![SampleId::new("sample:two").unwrap()],
2382                    validation_sample_ids: vec![SampleId::new("sample:one").unwrap()],
2383                    metadata: BTreeMap::new(),
2384                },
2385                FoldAssignment {
2386                    fold_id: FoldId::new("fold:1").unwrap(),
2387                    train_sample_ids: vec![SampleId::new("sample:one").unwrap()],
2388                    validation_sample_ids: vec![SampleId::new("sample:two").unwrap()],
2389                    metadata: BTreeMap::new(),
2390                },
2391            ],
2392            sample_groups: BTreeMap::new(),
2393            partition_mode: FoldPartitionMode::Partition,
2394        };
2395        let mut registry = ControllerRegistry::new();
2396        registry
2397            .register(manifest("controller:model", NodeKind::Model))
2398            .unwrap();
2399        let plan = build_execution_plan(
2400            "plan:hpo.scheduler",
2401            graph,
2402            CampaignSpec {
2403                inner_cv: None,
2404                id: "campaign:hpo.scheduler".to_string(),
2405                root_seed: Some(13),
2406                leakage_policy: Default::default(),
2407                aggregation_policy: Default::default(),
2408                split_invocation: Some(SplitInvocation {
2409                    id: "split:hpo".to_string(),
2410                    controller_id: None,
2411                    leakage_policy: Default::default(),
2412                    params: BTreeMap::new(),
2413                    fold_set: Some(fold_set),
2414                }),
2415                generation: Default::default(),
2416                shape_plans: BTreeMap::new(),
2417                data_bindings: BTreeMap::new(),
2418                branch_view_plans: Vec::new(),
2419                metadata: BTreeMap::new(),
2420            },
2421            &registry,
2422        )
2423        .unwrap();
2424        let trace = Arc::new(Mutex::new(Vec::new()));
2425        let mut controllers = RuntimeControllerRegistry::new();
2426        controllers
2427            .register(Box::new(HpoTestTuner {
2428                id: ControllerId::new("controller:tuner").unwrap(),
2429                trace: Arc::clone(&trace),
2430                history_len: 0,
2431                proposal_count: 1,
2432            }))
2433            .unwrap();
2434        controllers
2435            .register(Box::new(HpoTestModel {
2436                id: ControllerId::new("controller:model").unwrap(),
2437                trace: Arc::clone(&trace),
2438            }))
2439            .unwrap();
2440        let hpo = RuntimeHpoExecutionContext {
2441            operation_id: "hpo:test".to_string(),
2442            controller_id: ControllerId::new("controller:tuner").unwrap(),
2443            target_node_id: target.clone(),
2444            base_variant: plan.variants[0].clone(),
2445            trial_budget_total: 1,
2446            study: MethodsHpoStudyConfig {
2447                controller_id: "controller:tuner".to_string(),
2448                study_id: "study:hpo.scheduler".to_string(),
2449                methods_abi: "test-abi".to_string(),
2450                search_space: HpoSearchSpace {
2451                    parameters: vec![HpoParameter::Int {
2452                        name: "n_components".to_string(),
2453                        low: 1,
2454                        high: 1,
2455                        step: 1,
2456                        log: false,
2457                    }],
2458                },
2459                optimizer: HpoOptimizerConfig {
2460                    sampler: HpoSampler::Random,
2461                    pruner: HpoPruner::None,
2462                    direction: HpoDirection::Minimize,
2463                    metric: HpoMetric::Rmse,
2464                    seed: 13,
2465                    n_startup_trials: 1,
2466                    max_resource: 0,
2467                    reduction_factor: 1,
2468                },
2469            },
2470            parameter_paths: BTreeMap::from([(
2471                "n_components".to_string(),
2472                "n_components".to_string(),
2473            )]),
2474            resume_checkpoint: None,
2475            resume_variants: BTreeMap::new(),
2476            resume_terminal_trials: Vec::new(),
2477            selection: RuntimeHpoSelectionTarget {
2478                producer_node: target,
2479                producer_port: "prediction".to_string(),
2480                metric: RegressionMetricKind::Rmse,
2481                direction: HpoDirection::Minimize,
2482            },
2483            provenance: RuntimeHpoProvenance {
2484                graph_fingerprint: plan.graph_fingerprint.clone(),
2485                campaign_fingerprint: plan.campaign_fingerprint.clone(),
2486                controller_fingerprint: plan.controller_fingerprint.clone(),
2487                data_identities_fingerprint: "identity:test".to_string(),
2488                fold_set_fingerprint: plan
2489                    .fold_set
2490                    .as_ref()
2491                    .map(stable_json_fingerprint)
2492                    .transpose()
2493                    .unwrap(),
2494                training_influence_fingerprint: "influence:test".to_string(),
2495                relation_fingerprint: "relation:test".to_string(),
2496            },
2497        };
2498        let provider = InMemoryDataProvider::new(ControllerId::new("controller:data").unwrap());
2499        let ctx = RunContext::new(RunId::new("run:hpo.scheduler").unwrap(), Some(13));
2500
2501        let result = SequentialScheduler
2502            .execute_hpo_campaign(&plan, &controllers, &provider, &ctx, &hpo)
2503            .unwrap();
2504
2505        assert_eq!(result.operation_id, "hpo:test");
2506        assert_eq!(result.candidates.len(), 1);
2507        assert_eq!(result.checkpoint.completed_proposals.len(), 1);
2508        assert_eq!(result.checkpoint.completed_reports.len(), 1);
2509        assert_eq!(result.candidates[0].lineage.len(), 2);
2510        assert_eq!(result.incumbent.variant_id, plan.variants[0].variant_id);
2511        let mut selected_ctx = RunContext::new(RunId::new("run:hpo.scheduler").unwrap(), Some(13));
2512        selected_ctx.variant_id = Some(plan.variants[0].variant_id.clone());
2513        let selected_results = SequentialScheduler
2514            .execute_campaign_phase_with_data_provider(
2515                &plan,
2516                &controllers,
2517                &provider,
2518                &mut selected_ctx,
2519                Phase::FitCv,
2520            )
2521            .unwrap();
2522        assert_eq!(selected_results.len(), 2);
2523        assert_eq!(selected_ctx.lineage.len(), 2);
2524        assert_eq!(
2525            trace.lock().unwrap().as_slice(),
2526            [
2527                "session_factory",
2528                "ask",
2529                "model_cv",
2530                "model_cv",
2531                "intermediate",
2532                "tell",
2533                "checkpoint",
2534                "model_cv",
2535                "model_cv"
2536            ]
2537        );
2538
2539        // The native study may restore failed/pruned history that has no
2540        // completed proposal evidence. Its local count, not the coordinator's
2541        // persisted completed list, determines the remaining global budget.
2542        let resumed_trace = Arc::new(Mutex::new(Vec::new()));
2543        let mut resumed_controllers = RuntimeControllerRegistry::new();
2544        resumed_controllers
2545            .register(Box::new(HpoTestTuner {
2546                id: ControllerId::new("controller:tuner").unwrap(),
2547                trace: Arc::clone(&resumed_trace),
2548                history_len: 2,
2549                proposal_count: 2,
2550            }))
2551            .unwrap();
2552        resumed_controllers
2553            .register(Box::new(HpoTestModel {
2554                id: ControllerId::new("controller:model").unwrap(),
2555                trace: Arc::clone(&resumed_trace),
2556            }))
2557            .unwrap();
2558        let mut resumed_hpo = hpo.clone();
2559        resumed_hpo.trial_budget_total = 4;
2560        let resumed_ctx = RunContext::new(RunId::new("run:hpo.resumed").unwrap(), Some(13));
2561        let resumed = SequentialScheduler
2562            .execute_hpo_campaign(
2563                &plan,
2564                &resumed_controllers,
2565                &provider,
2566                &resumed_ctx,
2567                &resumed_hpo,
2568            )
2569            .unwrap();
2570        assert_eq!(resumed.candidates.len(), 2);
2571        assert_eq!(resumed.checkpoint.trial_history_len, 4);
2572        assert_eq!(
2573            resumed_trace
2574                .lock()
2575                .unwrap()
2576                .iter()
2577                .filter(|event| event.as_str() == "ask")
2578                .count(),
2579            2
2580        );
2581
2582        let mut over_budget_controllers = RuntimeControllerRegistry::new();
2583        over_budget_controllers
2584            .register(Box::new(HpoTestTuner {
2585                id: ControllerId::new("controller:tuner").unwrap(),
2586                trace: Arc::new(Mutex::new(Vec::new())),
2587                history_len: 5,
2588                proposal_count: 0,
2589            }))
2590            .unwrap();
2591        let error = SequentialScheduler
2592            .execute_hpo_campaign(
2593                &plan,
2594                &over_budget_controllers,
2595                &provider,
2596                &resumed_ctx,
2597                &resumed_hpo,
2598            )
2599            .unwrap_err();
2600        assert!(error.to_string().contains("exceeds total trial budget"));
2601    }
2602}
2603
2604pub(crate) fn attach_coordinator_input_lineage(
2605    result: &mut NodeResult,
2606    plan: &ExecutionPlan,
2607    node_id: &NodeId,
2608    upstream_lineage: &BTreeMap<NodeId, LineageId>,
2609) -> Result<()> {
2610    let inferred = inferred_input_lineage_for_node(plan, node_id, upstream_lineage);
2611    if result.lineage.input_lineage.is_empty() {
2612        result.lineage.input_lineage = inferred;
2613        return Ok(());
2614    }
2615
2616    let declared = result
2617        .lineage
2618        .input_lineage
2619        .iter()
2620        .cloned()
2621        .collect::<BTreeSet<_>>()
2622        .into_iter()
2623        .collect::<Vec<_>>();
2624    if declared != inferred {
2625        return Err(DagMlError::RuntimeValidation(format!(
2626            "lineage for node `{}` declared input lineage {:?}, expected {:?}",
2627            result.node_id, declared, inferred
2628        )));
2629    }
2630    result.lineage.input_lineage = declared;
2631    Ok(())
2632}
2633
2634/// The meta invocation in a nested-stacking FIT_CV scope consumes parent-bound
2635/// *inner* OOF blocks, not the outer blocks emitted immediately before it. The
2636/// ordinary per-scope lineage map cannot represent those records because every
2637/// inner fold ran in its own prior scope. Reconstruct the exact dependency set
2638/// from scheduler-owned nested evidence and attach it before recording the
2639/// meta result.
2640fn attach_nested_stacking_input_lineage(
2641    result: &mut NodeResult,
2642    plan: &ExecutionPlan,
2643    task: &NodeTask,
2644    ctx: &RunContext,
2645    nested: &NestedStackingInput<'_>,
2646) -> Result<()> {
2647    if task.phase != Phase::FitCv || task.node_plan.node_id != *nested.meta_node_id {
2648        return Ok(());
2649    }
2650    let inner_fold_ids = nested
2651        .inner
2652        .inner_fold_set
2653        .folds
2654        .iter()
2655        .map(|fold| fold.fold_id.clone())
2656        .collect::<BTreeSet<_>>();
2657    let source_nodes = incoming_oof_edges(plan, &task.node_plan)?
2658        .into_iter()
2659        .map(|edge| edge.source.node_id.clone())
2660        .collect::<BTreeSet<_>>();
2661    let expected = source_nodes
2662        .iter()
2663        .flat_map(|node_id| {
2664            inner_fold_ids
2665                .iter()
2666                .cloned()
2667                .map(move |fold_id| (node_id.clone(), fold_id))
2668        })
2669        .collect::<BTreeSet<_>>();
2670    let mut actual = BTreeMap::new();
2671    for record in ctx.lineage.records().filter(|record| {
2672        record.phase == Phase::FitCv
2673            && record.variant_id == task.variant_id
2674            && source_nodes.contains(&record.node_id)
2675            && record
2676                .fold_id
2677                .as_ref()
2678                .is_some_and(|fold_id| inner_fold_ids.contains(fold_id))
2679    }) {
2680        let fold_id = record
2681            .fold_id
2682            .clone()
2683            .expect("inner-fold predicate requires a fold id");
2684        if actual
2685            .insert((record.node_id.clone(), fold_id), record.record_id.clone())
2686            .is_some()
2687        {
2688            return Err(DagMlError::RuntimeValidation(
2689                "nested stacking meta input lineage contains duplicate inner-fold evidence"
2690                    .to_string(),
2691            ));
2692        }
2693    }
2694    if actual.keys().cloned().collect::<BTreeSet<_>>() != expected {
2695        return Err(DagMlError::RuntimeValidation(
2696            "nested stacking meta input lineage does not exactly cover inner OOF evidence"
2697                .to_string(),
2698        ));
2699    }
2700    let inferred = actual.into_values().collect::<Vec<_>>();
2701    if result.lineage.input_lineage.is_empty() {
2702        result.lineage.input_lineage = inferred;
2703        return Ok(());
2704    }
2705    let declared = result
2706        .lineage
2707        .input_lineage
2708        .iter()
2709        .cloned()
2710        .collect::<BTreeSet<_>>();
2711    if declared.into_iter().collect::<Vec<_>>() != inferred {
2712        return Err(DagMlError::RuntimeValidation(format!(
2713            "nested stacking meta lineage for node `{}` does not match inner OOF evidence",
2714            task.node_plan.node_id
2715        )));
2716    }
2717    result.lineage.input_lineage = inferred;
2718    Ok(())
2719}
2720
2721pub(crate) fn inferred_input_lineage_for_node(
2722    plan: &ExecutionPlan,
2723    node_id: &NodeId,
2724    upstream_lineage: &BTreeMap<NodeId, LineageId>,
2725) -> Vec<LineageId> {
2726    plan.graph_plan
2727        .graph
2728        .edges
2729        .iter()
2730        .filter(|edge| &edge.target.node_id == node_id && edge.contract.propagates_lineage)
2731        .filter_map(|edge| upstream_lineage.get(&edge.source.node_id).cloned())
2732        .collect::<BTreeSet<_>>()
2733        .into_iter()
2734        .collect()
2735}
2736pub(crate) fn collect_input_handles(
2737    plan: &ExecutionPlan,
2738    node_plan: &NodePlan,
2739    output_handles: &BTreeMap<NodeId, BTreeMap<String, HandleRef>>,
2740    output_data_views: &BTreeMap<NodeId, BTreeMap<String, DataProviderViewSpec>>,
2741    resources: &PhaseScopeResources<'_>,
2742    ctx: &RunContext,
2743    scope: &PhaseScope,
2744) -> Result<CollectedInputs> {
2745    let mut inputs = BTreeMap::new();
2746    let mut data_views = BTreeMap::new();
2747    let mut prediction_inputs = BTreeMap::new();
2748    let training_oof_edges = incoming_training_oof_edges(plan, node_plan, scope)?;
2749    // An OOF edge replaces exactly one raw producer port. Do not hide sibling
2750    // outputs from the same producer: a meta-node may legally consume both an
2751    // OOF prediction port and an auxiliary non-OOF port. PREDICT has no
2752    // Validation-OOF input, but its raw prediction port must still be masked so
2753    // only the explicit `:predict` off-fold input reaches the controller.
2754    let masked_oof_source_ports = if scope.phase == Phase::Predict {
2755        incoming_oof_edges(plan, node_plan)?
2756    } else {
2757        training_oof_edges.clone()
2758    }
2759    .into_iter()
2760    .map(|edge| (edge.source.node_id.clone(), edge.source.port_name.clone()))
2761    .collect::<BTreeSet<_>>();
2762    let bound_data_inputs = node_plan
2763        .data_bindings
2764        .iter()
2765        .map(|binding| binding.input_name.clone())
2766        .collect::<BTreeSet<_>>();
2767    // Only forward upstream handles for ports this node DECLARES an edge to.
2768    // A controller must never see a handle outside its declared port contract,
2769    // so a sibling consumer of the same producer cannot expose extra ports here.
2770    let declared_source_ports = plan
2771        .graph_plan
2772        .graph
2773        .edges
2774        .iter()
2775        .filter(|edge| edge.target.node_id == node_plan.node_id)
2776        .map(|edge| (edge.source.node_id.clone(), edge.source.port_name.clone()))
2777        .collect::<BTreeSet<_>>();
2778    for upstream in &node_plan.input_nodes {
2779        if let Some(handles) = output_handles.get(upstream) {
2780            for (port, handle) in handles {
2781                if !declared_source_ports.contains(&(upstream.clone(), port.clone())) {
2782                    continue;
2783                }
2784                if masked_oof_source_ports.contains(&(upstream.clone(), port.clone())) {
2785                    continue;
2786                }
2787                inputs.insert(format!("{upstream}.{port}"), handle.clone());
2788            }
2789        }
2790    }
2791    for edge in plan
2792        .graph_plan
2793        .graph
2794        .edges
2795        .iter()
2796        .filter(|edge| edge.target.node_id == node_plan.node_id)
2797        .filter(|edge| edge.contract.kind == PortKind::Data && !edge.contract.requires_oof)
2798    {
2799        if bound_data_inputs.contains(&edge.target.port_name) {
2800            continue;
2801        }
2802        let Some(handles) = output_handles.get(&edge.source.node_id) else {
2803            continue;
2804        };
2805        let Some(handle) = handles.get(&edge.source.port_name) else {
2806            continue;
2807        };
2808        let key = data_view_key(&edge.target.port_name);
2809        if inputs.insert(key.clone(), handle.clone()).is_some() {
2810            return Err(DagMlError::RuntimeValidation(format!(
2811                "node `{}` received duplicate data edge input `{key}`",
2812                node_plan.node_id
2813            )));
2814        }
2815        if let Some(source_views) = output_data_views.get(&edge.source.node_id) {
2816            if let Some(view) = source_views.get(&edge.source.port_name) {
2817                if data_views.insert(key.clone(), view.clone()).is_some() {
2818                    return Err(DagMlError::RuntimeValidation(format!(
2819                        "node `{}` received duplicate data edge view `{key}`",
2820                        node_plan.node_id
2821                    )));
2822                }
2823            }
2824            let source_validation_key = validation_data_view_key(&edge.source.port_name);
2825            if let Some(view) = source_views.get(&source_validation_key) {
2826                let validation_key = format!("{key}:validation");
2827                if data_views
2828                    .insert(validation_key.clone(), view.clone())
2829                    .is_some()
2830                {
2831                    return Err(DagMlError::RuntimeValidation(format!(
2832                        "node `{}` received duplicate data edge validation view `{validation_key}`",
2833                        node_plan.node_id
2834                    )));
2835                }
2836            }
2837        }
2838    }
2839    for edge in training_oof_edges {
2840        let key = format!("{}.{}", edge.source.node_id, edge.source.port_name);
2841        let Some(input) = collect_oof_prediction_input(plan, edge, ctx, scope, resources)? else {
2842            return Ok(CollectedInputs {
2843                handles: BTreeMap::new(),
2844                data_views: BTreeMap::new(),
2845                prediction_inputs: BTreeMap::new(),
2846                skip_node: true,
2847            });
2848        };
2849        if inputs.insert(key.clone(), input.handle).is_some() {
2850            return Err(DagMlError::RuntimeValidation(format!(
2851                "node `{}` received duplicate OOF prediction input `{key}`",
2852                node_plan.node_id
2853            )));
2854        }
2855        if prediction_inputs.insert(key.clone(), input.spec).is_some() {
2856            return Err(DagMlError::RuntimeValidation(format!(
2857                "node `{}` received duplicate OOF prediction spec `{key}`",
2858                node_plan.node_id
2859            )));
2860        }
2861    }
2862    // REFIT / PREDICT: deliver each base producer's off-fold (test / predict)
2863    // predictions to the stacking meta-node as a SEPARATE prediction input (suffixed
2864    // `:test` / `:predict`) so the host meta-model predicts from them. The FIT_CV
2865    // Validation-OOF input above is the meta-features the meta-model trains on; this
2866    // off-fold input is used ONLY for REFIT/PREDICT scoring/prediction, never FIT_CV
2867    // training — keeping the leakage invariant intact.
2868    if matches!(scope.phase, Phase::Refit | Phase::Predict) {
2869        let off_fold_suffix = scope.phase.as_str().to_ascii_lowercase();
2870        for edge in incoming_oof_edges(plan, node_plan)? {
2871            let Some(input) = collect_off_fold_prediction_input(plan, edge, ctx, scope)? else {
2872                continue;
2873            };
2874            let key = format!(
2875                "{}.{}:{off_fold_suffix}",
2876                edge.source.node_id, edge.source.port_name
2877            );
2878            if inputs.insert(key.clone(), input.handle).is_some() {
2879                return Err(DagMlError::RuntimeValidation(format!(
2880                    "node `{}` received duplicate off-fold prediction input `{key}`",
2881                    node_plan.node_id
2882                )));
2883            }
2884            if prediction_inputs.insert(key.clone(), input.spec).is_some() {
2885                return Err(DagMlError::RuntimeValidation(format!(
2886                    "node `{}` received duplicate off-fold prediction spec `{key}`",
2887                    node_plan.node_id
2888                )));
2889            }
2890        }
2891    }
2892    if !node_plan.data_bindings.is_empty() && resources.data_provider.is_none() {
2893        return Err(DagMlError::RuntimeValidation(format!(
2894            "node `{}` requires {} data binding(s) but no runtime data provider is registered",
2895            node_plan.node_id,
2896            node_plan.data_bindings.len()
2897        )));
2898    }
2899    if let Some(data_provider) = resources.data_provider {
2900        // Samples excluded from training (sample-local) are relevant only to
2901        // fitting scopes. A top-level PREDICT must not even resolve the CV
2902        // relation authority: its separately attested cohort below owns the
2903        // complete identity universe for that read.
2904        let excluded_samples = if scope.phase == Phase::Predict {
2905            BTreeSet::new()
2906        } else {
2907            coordinator_relations_for_node(node_plan, resources)?
2908                .map(|relations| relations.excluded_sample_ids())
2909                .unwrap_or_default()
2910        };
2911        let scope_fold_set = resources.fold_set_override.or(plan.fold_set.as_ref());
2912        for binding in &node_plan.data_bindings {
2913            let predict_cohort = if scope.phase == Phase::Predict {
2914                data_provider.predict_cohort(binding, scope.phase)?
2915            } else {
2916                None
2917            };
2918            let materialized = data_provider.materialize(&DataMaterializationRequest {
2919                run_id: ctx.run_id.clone(),
2920                node_id: node_plan.node_id.clone(),
2921                input_name: binding.input_name.clone(),
2922                phase: scope.phase,
2923                variant_id: scope.variant_id.clone(),
2924                fold_id: scope.fold_id.clone(),
2925                binding: binding.clone(),
2926                predict_cohort: predict_cohort.clone(),
2927            })?;
2928            let branch_view_for_node = branch_view_from_node_metadata(plan, &node_plan.node_id)?;
2929            let mut view = data_view_for_scope(
2930                binding,
2931                scope_fold_set,
2932                scope,
2933                branch_view_for_node.as_ref(),
2934                &excluded_samples,
2935            )?;
2936            if let Some(cohort) = predict_cohort.as_ref() {
2937                bind_predict_cohort_to_view(&mut view, cohort)?;
2938            }
2939            let key = data_view_key(&binding.input_name);
2940            let view_handle = make_data_view_handle(
2941                data_provider,
2942                ctx,
2943                node_plan,
2944                scope,
2945                binding,
2946                DataViewHandleInput {
2947                    data_handle: &materialized,
2948                    view: &view,
2949                    predict_cohort: predict_cohort.as_ref(),
2950                },
2951            )?;
2952            if data_views.insert(key.clone(), view).is_some() {
2953                return Err(DagMlError::RuntimeValidation(format!(
2954                    "node `{}` received duplicate data view `{key}`",
2955                    node_plan.node_id
2956                )));
2957            }
2958            if inputs.insert(key.clone(), view_handle).is_some() {
2959                return Err(DagMlError::RuntimeValidation(format!(
2960                    "node `{}` received duplicate data input `{key}`",
2961                    node_plan.node_id
2962                )));
2963            }
2964
2965            if let Some(validation_view) = validation_data_view_for_scope(
2966                binding,
2967                scope_fold_set,
2968                scope,
2969                branch_view_for_node.as_ref(),
2970                &excluded_samples,
2971            )? {
2972                let validation_key = format!("{key}:validation");
2973                let validation_handle = make_data_view_handle(
2974                    data_provider,
2975                    ctx,
2976                    node_plan,
2977                    scope,
2978                    binding,
2979                    DataViewHandleInput {
2980                        data_handle: &materialized,
2981                        view: &validation_view,
2982                        predict_cohort: None,
2983                    },
2984                )?;
2985                if data_views
2986                    .insert(validation_key.clone(), validation_view)
2987                    .is_some()
2988                {
2989                    return Err(DagMlError::RuntimeValidation(format!(
2990                        "node `{}` received duplicate validation data view `{validation_key}`",
2991                        node_plan.node_id
2992                    )));
2993                }
2994                if inputs
2995                    .insert(validation_key.clone(), validation_handle)
2996                    .is_some()
2997                {
2998                    return Err(DagMlError::RuntimeValidation(format!(
2999                        "node `{}` received duplicate validation data input `{validation_key}`",
3000                        node_plan.node_id
3001                    )));
3002                }
3003            }
3004        }
3005    }
3006    Ok(CollectedInputs {
3007        handles: inputs,
3008        data_views,
3009        prediction_inputs,
3010        skip_node: false,
3011    })
3012}
3013pub(crate) fn preload_replay_prediction_cache_store(
3014    bundle: &ExecutionBundle,
3015    prediction_cache_store: Option<&dyn RuntimePredictionCacheStore>,
3016    ctx: &mut RunContext,
3017) -> Result<()> {
3018    if bundle.prediction_requirements.is_empty() {
3019        return Ok(());
3020    }
3021    let store = prediction_cache_store.ok_or_else(|| {
3022        DagMlError::RuntimeValidation(format!(
3023            "bundle `{}` cannot preload OOF prediction caches without a prediction cache store",
3024            bundle.bundle_id
3025        ))
3026    })?;
3027    if !ctx.prediction_store.blocks().is_empty() {
3028        return Err(DagMlError::RuntimeValidation(format!(
3029            "bundle `{}` cannot preload OOF prediction caches into a non-empty prediction store",
3030            bundle.bundle_id
3031        )));
3032    }
3033    let contracts = replay_prediction_cache_contracts(bundle)?;
3034    for contract in contracts.values() {
3035        if contract.requirement.prediction_level == PredictionLevel::Sample {
3036            let blocks = store.load_blocks(&contract.cache.requirement_key)?;
3037            if blocks.iter().any(|block| {
3038                block.producer_node != contract.requirement.producer_node
3039                    || block.partition != contract.requirement.partition
3040            }) {
3041                return Err(DagMlError::RuntimeValidation(format!(
3042                    "prediction cache store returned blocks outside requirement `{}`",
3043                    contract.cache.requirement_key
3044                )));
3045            }
3046            let mut payload = build_prediction_cache_payload(&contract.requirement, &blocks)?;
3047            payload.cache_namespace_fingerprints =
3048                contract.cache.cache_namespace_fingerprints.clone();
3049            validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
3050            for block in &payload.blocks {
3051                ctx.prediction_store.append(block.clone())?;
3052            }
3053        } else {
3054            let blocks = store.load_aggregated_blocks(&contract.cache.requirement_key)?;
3055            if blocks.iter().any(|block| {
3056                block.producer_node != contract.requirement.producer_node
3057                    || block.partition != contract.requirement.partition
3058                    || block.level != contract.requirement.prediction_level
3059            }) {
3060                return Err(DagMlError::RuntimeValidation(format!(
3061                    "prediction cache store returned aggregated blocks outside requirement `{}`",
3062                    contract.cache.requirement_key
3063                )));
3064            }
3065            let mut payload =
3066                build_aggregated_prediction_cache_payload(&contract.requirement, &blocks)?;
3067            payload.cache_namespace_fingerprints =
3068                contract.cache.cache_namespace_fingerprints.clone();
3069            validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
3070        }
3071    }
3072    Ok(())
3073}
3074
3075pub(crate) fn replay_prediction_cache_contracts(
3076    bundle: &ExecutionBundle,
3077) -> Result<BTreeMap<String, ReplayPredictionCacheContract>> {
3078    bundle.validate()?;
3079    let requirements = bundle
3080        .prediction_requirements
3081        .iter()
3082        .map(|requirement| (requirement.key(), requirement))
3083        .collect::<BTreeMap<_, _>>();
3084    let mut contracts = BTreeMap::new();
3085    for cache in &bundle.prediction_caches {
3086        let requirement = requirements.get(&cache.requirement_key).ok_or_else(|| {
3087            DagMlError::RuntimeValidation(format!(
3088                "prediction cache `{}` references unknown prediction requirement `{}`",
3089                cache.cache_id, cache.requirement_key
3090            ))
3091        })?;
3092        contracts.insert(
3093            cache.requirement_key.clone(),
3094            ReplayPredictionCacheContract {
3095                requirement: (*requirement).clone(),
3096                cache: cache.clone(),
3097            },
3098        );
3099    }
3100    Ok(contracts)
3101}
3102
3103pub(crate) fn materialize_replay_artifact_handles(
3104    plan: &ExecutionPlan,
3105    bundle: &ExecutionBundle,
3106    replay_request: &ReplayPhaseRequest,
3107    artifact_store: &dyn RuntimeArtifactStore,
3108    ctx: &RunContext,
3109) -> Result<MaterializedReplayArtifacts> {
3110    let mut handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
3111    let mut inputs = BTreeMap::<NodeId, BTreeMap<String, ArtifactInputSpec>>::new();
3112    for artifact in &bundle.refit_artifacts {
3113        artifact.validate()?;
3114        let node_plan = plan.node_plans.get(&artifact.node_id).ok_or_else(|| {
3115            DagMlError::RuntimeValidation(format!(
3116                "bundle `{}` artifact references unknown node `{}`",
3117                bundle.bundle_id, artifact.node_id
3118            ))
3119        })?;
3120        if !node_plan.supported_phases.contains(&replay_request.phase) {
3121            return Err(DagMlError::RuntimeValidation(format!(
3122                "bundle `{}` artifact node `{}` does not support replay phase {:?}",
3123                bundle.bundle_id, artifact.node_id, replay_request.phase
3124            )));
3125        }
3126        let handle = artifact_store.materialize(&ArtifactMaterializationRequest {
3127            run_id: ctx.run_id.clone(),
3128            bundle_id: bundle.bundle_id.clone(),
3129            node_id: artifact.node_id.clone(),
3130            phase: replay_request.phase,
3131            variant_id: bundle.selected_variant_id.clone(),
3132            controller_id: artifact.controller_id.clone(),
3133            artifact: artifact.artifact.clone(),
3134            params_fingerprint: artifact.params_fingerprint.clone(),
3135            training_loss_fingerprint: artifact.training_loss_fingerprint.clone(),
3136        })?;
3137        if !matches!(handle.kind, HandleKind::Model | HandleKind::Artifact) {
3138            return Err(DagMlError::RuntimeValidation(format!(
3139                "artifact `{}` materialized as unsupported handle kind {:?}",
3140                artifact.artifact.id, handle.kind
3141            )));
3142        }
3143        if handle.owner_controller != artifact.controller_id {
3144            return Err(DagMlError::RuntimeValidation(format!(
3145                "artifact `{}` handle owner `{}` does not match controller `{}`",
3146                artifact.artifact.id, handle.owner_controller, artifact.controller_id
3147            )));
3148        }
3149        let key = refit_artifact_input_key(&artifact.artifact.id);
3150        if handles
3151            .entry(artifact.node_id.clone())
3152            .or_default()
3153            .insert(key.clone(), handle)
3154            .is_some()
3155        {
3156            return Err(DagMlError::RuntimeValidation(format!(
3157                "duplicate replay artifact input `{key}` for node `{}`",
3158                artifact.node_id
3159            )));
3160        }
3161        if inputs
3162            .entry(artifact.node_id.clone())
3163            .or_default()
3164            .insert(key.clone(), ArtifactInputSpec::from_refit_record(artifact)?)
3165            .is_some()
3166        {
3167            return Err(DagMlError::RuntimeValidation(format!(
3168                "duplicate replay artifact metadata `{key}` for node `{}`",
3169                artifact.node_id
3170            )));
3171        }
3172    }
3173    Ok(MaterializedReplayArtifacts { handles, inputs })
3174}
3175
3176pub(crate) fn derive_task_seed(
3177    root_seed: Option<u64>,
3178    variant_id: Option<&VariantId>,
3179    fold_id: Option<&FoldId>,
3180    node_plan: &NodePlan,
3181    phase: Phase,
3182) -> Option<u64> {
3183    root_seed.map(|root| {
3184        let mut context = SeedContext::root(root);
3185        if let Some(variant_id) = variant_id {
3186            context = context.child(format!("variant:{variant_id}"));
3187        }
3188        if let Some(fold_id) = fold_id {
3189            context = context.child(format!("fold:{fold_id}"));
3190        }
3191        context
3192            .child(format!("node:{}", node_plan.node_id))
3193            .child(format!("phase:{phase:?}"))
3194            .derive_u64("task")
3195    })
3196}