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