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