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                attach_coordinator_input_lineage(
1177                    &mut result,
1178                    plan,
1179                    &task.node_plan.node_id,
1180                    &input_lineage,
1181                )?;
1182                if let Some(store) = resources.artifact_store.as_deref_mut() {
1183                    if scope.phase == Phase::Refit {
1184                        store.capture_refit_artifacts(&task, &result)?;
1185                    }
1186                }
1187                for prediction in &result.predictions {
1188                    ctx.prediction_store.append(prediction.clone())?;
1189                }
1190                for prediction in &result.aggregated_predictions {
1191                    ctx.aggregated_prediction_store.append(prediction.clone())?;
1192                }
1193                apply_result_scoring(
1194                    &result,
1195                    &mut ctx.score_collector,
1196                    &mut ctx.regression_target_records,
1197                )?;
1198                ctx.lineage.record(result.lineage.clone())?;
1199                let data_views = derive_output_data_views(plan, &task, &result)?;
1200                output_handles.insert(node_id.clone(), result.outputs.clone());
1201                output_data_views.insert(node_id.clone(), data_views);
1202                input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1203                results.push(result);
1204            }
1205        }
1206
1207        Ok(results)
1208    }
1209}
1210
1211impl ParallelScheduler {
1212    pub fn execute_phase(
1213        &self,
1214        plan: &ExecutionPlan,
1215        controllers: &RuntimeControllerRegistry,
1216        ctx: &mut RunContext,
1217        phase: Phase,
1218    ) -> Result<Vec<NodeResult>> {
1219        plan.validate()?;
1220        let variant_id = ctx.variant_id.clone();
1221        let seed_root = ctx.root_seed;
1222        self.execute_phase_scope(
1223            plan,
1224            controllers,
1225            ctx,
1226            PhaseScope {
1227                phase,
1228                variant_id,
1229                variant: None,
1230                fold_id: None,
1231                seed_root,
1232            },
1233            PhaseScopeResources::default(),
1234        )
1235    }
1236
1237    pub fn execute_phase_with_data_provider(
1238        &self,
1239        plan: &ExecutionPlan,
1240        controllers: &RuntimeControllerRegistry,
1241        data_provider: &dyn RuntimeDataProvider,
1242        ctx: &mut RunContext,
1243        phase: Phase,
1244    ) -> Result<Vec<NodeResult>> {
1245        plan.validate()?;
1246        let variant_id = ctx.variant_id.clone();
1247        let seed_root = ctx.root_seed;
1248        self.execute_phase_scope(
1249            plan,
1250            controllers,
1251            ctx,
1252            PhaseScope {
1253                phase,
1254                variant_id,
1255                variant: None,
1256                fold_id: None,
1257                seed_root,
1258            },
1259            PhaseScopeResources {
1260                data_provider: Some(data_provider),
1261                ..Default::default()
1262            },
1263        )
1264    }
1265
1266    pub fn execute_campaign_phase(
1267        &self,
1268        plan: &ExecutionPlan,
1269        controllers: &RuntimeControllerRegistry,
1270        ctx: &mut RunContext,
1271        phase: Phase,
1272    ) -> Result<Vec<NodeResult>> {
1273        plan.validate()?;
1274        if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1275            return Err(DagMlError::RuntimeValidation(
1276                "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1277                    .to_string(),
1278            ));
1279        }
1280        let mut results = Vec::new();
1281        let fold_ids = if phase == Phase::FitCv {
1282            plan.fold_set
1283                .as_ref()
1284                .map(|fold_set| {
1285                    fold_set
1286                        .folds
1287                        .iter()
1288                        .map(|fold| Some(fold.fold_id.clone()))
1289                        .collect::<Vec<_>>()
1290                })
1291                .unwrap_or_else(|| vec![None])
1292        } else {
1293            vec![None]
1294        };
1295        for variant in &plan.variants {
1296            if ctx
1297                .variant_id
1298                .as_ref()
1299                .is_some_and(|requested| requested != &variant.variant_id)
1300            {
1301                continue;
1302            }
1303            for fold_id in &fold_ids {
1304                let seed_root = variant.seed.or(ctx.root_seed);
1305                results.extend(self.execute_phase_scope(
1306                    plan,
1307                    controllers,
1308                    ctx,
1309                    PhaseScope {
1310                        phase,
1311                        variant_id: Some(variant.variant_id.clone()),
1312                        variant: Some(VariantExecutionSpec::from_plan(variant)),
1313                        fold_id: fold_id.clone(),
1314                        seed_root,
1315                    },
1316                    PhaseScopeResources::default(),
1317                )?);
1318            }
1319        }
1320        Ok(results)
1321    }
1322
1323    pub fn execute_campaign_phase_with_data_provider(
1324        &self,
1325        plan: &ExecutionPlan,
1326        controllers: &RuntimeControllerRegistry,
1327        data_provider: &dyn RuntimeDataProvider,
1328        ctx: &mut RunContext,
1329        phase: Phase,
1330    ) -> Result<Vec<NodeResult>> {
1331        plan.validate()?;
1332        if phase == Phase::FitCv {
1333            ctx.configure_global_oof_aggregation(plan, data_provider)?;
1334        }
1335        if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1336            return Err(DagMlError::RuntimeValidation(
1337                "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1338                    .to_string(),
1339            ));
1340        }
1341        let mut results = Vec::new();
1342        let fold_ids = if phase == Phase::FitCv {
1343            plan.fold_set
1344                .as_ref()
1345                .map(|fold_set| {
1346                    fold_set
1347                        .folds
1348                        .iter()
1349                        .map(|fold| Some(fold.fold_id.clone()))
1350                        .collect::<Vec<_>>()
1351                })
1352                .unwrap_or_else(|| vec![None])
1353        } else {
1354            vec![None]
1355        };
1356        for variant in &plan.variants {
1357            if ctx
1358                .variant_id
1359                .as_ref()
1360                .is_some_and(|requested| requested != &variant.variant_id)
1361            {
1362                continue;
1363            }
1364            for fold_id in &fold_ids {
1365                let seed_root = variant.seed.or(ctx.root_seed);
1366                results.extend(self.execute_phase_scope(
1367                    plan,
1368                    controllers,
1369                    ctx,
1370                    PhaseScope {
1371                        phase,
1372                        variant_id: Some(variant.variant_id.clone()),
1373                        variant: Some(VariantExecutionSpec::from_plan(variant)),
1374                        fold_id: fold_id.clone(),
1375                        seed_root,
1376                    },
1377                    PhaseScopeResources {
1378                        data_provider: Some(data_provider),
1379                        ..Default::default()
1380                    },
1381                )?);
1382            }
1383        }
1384        Ok(results)
1385    }
1386
1387    pub fn execute_campaign_phase_with_data_provider_and_artifact_store(
1388        &self,
1389        plan: &ExecutionPlan,
1390        controllers: &RuntimeControllerRegistry,
1391        data_provider: &dyn RuntimeDataProvider,
1392        artifact_store: &mut InMemoryArtifactStore,
1393        ctx: &mut RunContext,
1394        phase: Phase,
1395    ) -> Result<Vec<NodeResult>> {
1396        plan.validate()?;
1397        if phase == Phase::FitCv && nested_stacking_campaign_plan(plan)?.is_some() {
1398            return Err(DagMlError::RuntimeValidation(
1399                "nested stacking FIT_CV is scheduler-serial by construction; use SequentialScheduler so inner OOF evidence is retained before outer evaluation"
1400                    .to_string(),
1401            ));
1402        }
1403        let mut results = Vec::new();
1404        let fold_ids = if phase == Phase::FitCv {
1405            plan.fold_set
1406                .as_ref()
1407                .map(|fold_set| {
1408                    fold_set
1409                        .folds
1410                        .iter()
1411                        .map(|fold| Some(fold.fold_id.clone()))
1412                        .collect::<Vec<_>>()
1413                })
1414                .unwrap_or_else(|| vec![None])
1415        } else {
1416            vec![None]
1417        };
1418        for variant in &plan.variants {
1419            if ctx
1420                .variant_id
1421                .as_ref()
1422                .is_some_and(|requested| requested != &variant.variant_id)
1423            {
1424                continue;
1425            }
1426            for fold_id in &fold_ids {
1427                let seed_root = variant.seed.or(ctx.root_seed);
1428                results.extend(self.execute_phase_scope(
1429                    plan,
1430                    controllers,
1431                    ctx,
1432                    PhaseScope {
1433                        phase,
1434                        variant_id: Some(variant.variant_id.clone()),
1435                        variant: Some(VariantExecutionSpec::from_plan(variant)),
1436                        fold_id: fold_id.clone(),
1437                        seed_root,
1438                    },
1439                    PhaseScopeResources {
1440                        data_provider: Some(data_provider),
1441                        artifact_store: Some(&mut *artifact_store),
1442                        ..Default::default()
1443                    },
1444                )?);
1445            }
1446        }
1447        Ok(results)
1448    }
1449
1450    pub fn execute_bundle_replay(
1451        &self,
1452        replay: BundleReplayExecution<'_>,
1453        ctx: &mut RunContext,
1454    ) -> Result<Vec<NodeResult>> {
1455        replay.bundle.validate_against_plan(replay.plan)?;
1456        replay
1457            .replay_request
1458            .validate_for_bundle_with_prediction_cache_store(
1459                replay.bundle,
1460                replay.prediction_cache_store.is_some(),
1461            )?;
1462        replay
1463            .bundle
1464            .validate_replay_envelopes(replay.data_envelopes)?;
1465        let prediction_cache_contracts = if replay.replay_request.phase == Phase::Refit {
1466            Some(replay_prediction_cache_contracts(replay.bundle)?)
1467        } else {
1468            None
1469        };
1470        if replay.replay_request.phase == Phase::Refit {
1471            preload_replay_prediction_cache_store(
1472                replay.bundle,
1473                replay.prediction_cache_store,
1474                ctx,
1475            )?;
1476        }
1477        let replay_artifacts = materialize_replay_artifact_handles(
1478            replay.plan,
1479            replay.bundle,
1480            replay.replay_request,
1481            replay.artifact_store,
1482            ctx,
1483        )?;
1484        let selected_variant = replay
1485            .bundle
1486            .selected_variant_id
1487            .as_ref()
1488            .map(|selected| {
1489                replay
1490                    .plan
1491                    .variants
1492                    .iter()
1493                    .find(|variant| &variant.variant_id == selected)
1494                    .map(VariantExecutionSpec::from_plan)
1495                    .ok_or_else(|| {
1496                        DagMlError::RuntimeValidation(format!(
1497                            "bundle `{}` selected unknown variant `{selected}`",
1498                            replay.bundle.bundle_id
1499                        ))
1500                    })
1501            })
1502            .transpose()?;
1503        let seed_root = selected_variant
1504            .as_ref()
1505            .and_then(|variant| variant.seed)
1506            .or(ctx.root_seed);
1507
1508        self.execute_phase_scope(
1509            replay.plan,
1510            replay.controllers,
1511            ctx,
1512            PhaseScope {
1513                phase: replay.replay_request.phase,
1514                variant_id: replay.bundle.selected_variant_id.clone(),
1515                variant: selected_variant,
1516                fold_id: None,
1517                seed_root,
1518            },
1519            PhaseScopeResources {
1520                data_provider: Some(replay.data_provider),
1521                replay_artifact_handles: Some(&replay_artifacts.handles),
1522                replay_artifact_inputs: Some(&replay_artifacts.inputs),
1523                replay_bundle_id: Some(&replay.bundle.bundle_id),
1524                data_envelopes: Some(replay.data_envelopes),
1525                prediction_cache_store: replay.prediction_cache_store,
1526                prediction_cache_contracts: prediction_cache_contracts.as_ref(),
1527                ..Default::default()
1528            },
1529        )
1530    }
1531
1532    fn execute_phase_scope(
1533        &self,
1534        plan: &ExecutionPlan,
1535        controllers: &RuntimeControllerRegistry,
1536        ctx: &mut RunContext,
1537        scope: PhaseScope,
1538        mut resources: PhaseScopeResources<'_>,
1539    ) -> Result<Vec<NodeResult>> {
1540        // Hold the phase span on the scheduler thread, and clone it into each
1541        // worker so worker-thread telemetry nests under the phase (tracing spans
1542        // are thread-local and do not auto-propagate across `thread::scope`).
1543        let phase_span = crate::observability::phase_span(
1544            ctx.run_id.as_str(),
1545            plan.id.as_str(),
1546            scope.phase.as_str(),
1547            scope.variant_id.as_ref().map(VariantId::as_str),
1548            scope.fold_id.as_ref().map(FoldId::as_str),
1549        );
1550        let _phase_entered = phase_span.clone().entered();
1551        // Borrowed for the `thread::scope` below; workers join before it ends.
1552        let plan_id = plan.id.as_str();
1553        plan.validate_parallel_controller_capabilities(self.max_workers, scope.phase)?;
1554        let mut results = Vec::new();
1555        let mut output_handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
1556        let mut output_data_views =
1557            BTreeMap::<NodeId, BTreeMap<String, DataProviderViewSpec>>::new();
1558        let mut input_lineage = BTreeMap::<NodeId, LineageId>::new();
1559
1560        for level in plan.node_parallel_levels_for_phase(scope.phase)? {
1561            let mut prepared = Vec::<PreparedNodeTask>::new();
1562            // Cross-branch merge nodes (concat or late-fusion) are not controller
1563            // tasks: they read the upstream branch OOF blocks from the prediction
1564            // store and reassemble them on the scheduler thread (no worker), AFTER
1565            // this level's worker tasks have populated the store. They are in a
1566            // later level than their branches, so the store already holds the
1567            // branch OOF by the time we reassemble — see `reassemble_branch_merge`.
1568            let mut merge_nodes = Vec::<(NodeId, MergeReduction)>::new();
1569            for node_id in &level {
1570                let node_plan = plan
1571                    .node_plans
1572                    .get(node_id)
1573                    .expect("execution plan was validated");
1574                if let Some(reduction) = merge_reduction_mode(plan, node_plan) {
1575                    merge_nodes.push((node_id.clone(), reduction));
1576                    continue;
1577                }
1578                let collected_inputs = collect_input_handles(
1579                    plan,
1580                    node_plan,
1581                    &output_handles,
1582                    &output_data_views,
1583                    &resources,
1584                    ctx,
1585                    &scope,
1586                )?;
1587                if collected_inputs.skip_node {
1588                    continue;
1589                }
1590                let mut input_handles = collected_inputs.handles;
1591                let mut artifact_inputs = BTreeMap::new();
1592                if let Some(node_artifact_handles) = resources
1593                    .replay_artifact_handles
1594                    .and_then(|handles| handles.get(node_id))
1595                {
1596                    for (key, handle) in node_artifact_handles {
1597                        if input_handles.insert(key.clone(), handle.clone()).is_some() {
1598                            return Err(DagMlError::RuntimeValidation(format!(
1599                                "node `{node_id}` received duplicate replay artifact input `{key}`"
1600                            )));
1601                        }
1602                    }
1603                }
1604                if let Some(node_artifact_inputs) = resources
1605                    .replay_artifact_inputs
1606                    .and_then(|inputs| inputs.get(node_id))
1607                {
1608                    for (key, spec) in node_artifact_inputs {
1609                        if artifact_inputs.insert(key.clone(), spec.clone()).is_some() {
1610                            return Err(DagMlError::RuntimeValidation(format!(
1611                                "node `{node_id}` received duplicate replay artifact metadata `{key}`"
1612                            )));
1613                        }
1614                    }
1615                }
1616                let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1617                let inner_fold_set = inner_fold_set_for_scope(
1618                    &plan.campaign,
1619                    plan.fold_set.as_ref(),
1620                    node_plan,
1621                    &scope,
1622                )?;
1623                let fit_influence = fit_influence_task_for_node(
1624                    plan,
1625                    &task_node_plan,
1626                    &collected_inputs.data_views,
1627                )?;
1628                prepared.push(PreparedNodeTask {
1629                    node_id: node_id.clone(),
1630                    task: NodeTask {
1631                        inner_fold_set,
1632                        run_id: ctx.run_id.clone(),
1633                        node_plan: task_node_plan.clone(),
1634                        phase: scope.phase,
1635                        variant_id: scope.variant_id.clone(),
1636                        variant: scope.variant.clone(),
1637                        fold_id: scope.fold_id.clone(),
1638                        branch_path: Vec::new(),
1639                        input_handles,
1640                        data_views: collected_inputs.data_views,
1641                        prediction_inputs: collected_inputs.prediction_inputs,
1642                        artifact_inputs,
1643                        required_loss_attestations: NodeTask::required_loss_attestations_for(
1644                            &task_node_plan,
1645                            scope.phase,
1646                        )?,
1647                        fit_influence,
1648                        seed: derive_task_seed(
1649                            scope.seed_root,
1650                            scope.variant_id.as_ref(),
1651                            scope.fold_id.as_ref(),
1652                            &task_node_plan,
1653                            scope.phase,
1654                        ),
1655                    },
1656                });
1657            }
1658
1659            for chunk in prepared.chunks(self.max_workers) {
1660                let chunk_results = std::thread::scope(
1661                    |thread_scope| -> Result<Vec<NodeResult>> {
1662                        let mut handles = Vec::with_capacity(chunk.len());
1663                        for prepared_task in chunk {
1664                            let controller = controllers
1665                                .get(&prepared_task.task.node_plan.controller_id)
1666                                .ok_or_else(|| {
1667                                    DagMlError::RuntimeValidation(format!(
1668                                        "runtime controller `{}` is not registered",
1669                                        prepared_task.task.node_plan.controller_id
1670                                    ))
1671                                })?;
1672                            let worker_span = phase_span.clone();
1673                            handles.push(thread_scope.spawn(move || {
1674                                let _worker_span = worker_span.entered();
1675                                let _node_span = crate::observability::node_span(
1676                                    prepared_task.task.run_id.as_str(),
1677                                    plan_id,
1678                                    prepared_task.task.phase.as_str(),
1679                                    prepared_task.task.node_plan.node_id.as_str(),
1680                                    prepared_task.task.node_plan.controller_id.as_str(),
1681                                )
1682                                .entered();
1683                                let mut result =
1684                                    if prepared_task.task.node_plan.kind == NodeKind::Tuner {
1685                                        return Err(DagMlError::RuntimeValidation(format!(
1686                                            "tuner node `{}` requires execute_hpo_campaign with an explicit RuntimeHpoExecutionContext",
1687                                            prepared_task.task.node_plan.node_id
1688                                        )));
1689                                    } else {
1690                                        // A provider-aware controller may require a
1691                                        // non-Sync host provider.  Parallel native
1692                                        // Methods PLS is deliberately refused by its
1693                                        // HPO preflight; ordinary controllers keep
1694                                        // their opaque-handle invocation here.
1695                                        controller.invoke(&prepared_task.task)?
1696                                    };
1697                                record_fit_influence_diagnostic(&prepared_task.task, &mut result);
1698                                normalize_result_prediction_ports(
1699                                    plan,
1700                                    &prepared_task.task,
1701                                    &mut result,
1702                                )?;
1703                                result.validate_for_task(&prepared_task.task)?;
1704                                Ok(result)
1705                            }));
1706                        }
1707                        handles
1708                            .into_iter()
1709                            .map(|handle| {
1710                                handle.join().map_err(|_| {
1711                                    DagMlError::RuntimeValidation(
1712                                        "parallel scheduler worker panicked".to_string(),
1713                                    )
1714                                })?
1715                            })
1716                            .collect()
1717                    },
1718                )?;
1719
1720                for (prepared_task, mut result) in chunk.iter().zip(chunk_results) {
1721                    apply_result_prediction_aggregation(
1722                        plan,
1723                        controllers,
1724                        &prepared_task.task,
1725                        &mut result,
1726                        &resources,
1727                    )?;
1728                    attach_coordinator_input_lineage(
1729                        &mut result,
1730                        plan,
1731                        &prepared_task.task.node_plan.node_id,
1732                        &input_lineage,
1733                    )?;
1734                    if let Some(store) = resources.artifact_store.as_deref_mut() {
1735                        if scope.phase == Phase::Refit {
1736                            store.capture_refit_artifacts(&prepared_task.task, &result)?;
1737                        }
1738                    }
1739                    for prediction in &result.predictions {
1740                        ctx.prediction_store.append(prediction.clone())?;
1741                    }
1742                    for prediction in &result.aggregated_predictions {
1743                        ctx.aggregated_prediction_store.append(prediction.clone())?;
1744                    }
1745                    apply_result_scoring(
1746                        &result,
1747                        &mut ctx.score_collector,
1748                        &mut ctx.regression_target_records,
1749                    )?;
1750                    ctx.lineage.record(result.lineage.clone())?;
1751                    let data_views = derive_output_data_views(plan, &prepared_task.task, &result)?;
1752                    output_handles.insert(prepared_task.node_id.clone(), result.outputs.clone());
1753                    output_data_views.insert(prepared_task.node_id.clone(), data_views);
1754                    input_lineage.insert(
1755                        prepared_task.node_id.clone(),
1756                        result.lineage.record_id.clone(),
1757                    );
1758                    results.push(result);
1759                }
1760            }
1761
1762            // Reassemble any cross-branch merge nodes in this level now that the
1763            // level's worker tasks have populated the prediction store. Merge nodes
1764            // sit in a later level than the branches they consume, so the upstream
1765            // branch OOF is already present.
1766            for (node_id, reduction) in &merge_nodes {
1767                let node_plan = plan
1768                    .node_plans
1769                    .get(node_id)
1770                    .expect("execution plan was validated");
1771                if let Some(mut result) =
1772                    reassemble_branch_merge(plan, node_plan, ctx, &scope, *reduction)?
1773                {
1774                    let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
1775                    let task = NodeTask {
1776                        inner_fold_set: None,
1777                        run_id: ctx.run_id.clone(),
1778                        node_plan: task_node_plan.clone(),
1779                        phase: scope.phase,
1780                        variant_id: scope.variant_id.clone(),
1781                        variant: scope.variant.clone(),
1782                        fold_id: scope.fold_id.clone(),
1783                        branch_path: Vec::new(),
1784                        input_handles: BTreeMap::new(),
1785                        data_views: BTreeMap::new(),
1786                        prediction_inputs: BTreeMap::new(),
1787                        artifact_inputs: BTreeMap::new(),
1788                        required_loss_attestations: NodeTask::required_loss_attestations_for(
1789                            &task_node_plan,
1790                            scope.phase,
1791                        )?,
1792                        fit_influence: FitInfluenceTask::default(),
1793                        seed: None,
1794                    };
1795                    normalize_result_prediction_ports(plan, &task, &mut result)?;
1796                    result.validate_for_task(&task)?;
1797                    for prediction in &result.predictions {
1798                        ctx.prediction_store.append(prediction.clone())?;
1799                    }
1800                    apply_result_scoring(
1801                        &result,
1802                        &mut ctx.score_collector,
1803                        &mut ctx.regression_target_records,
1804                    )?;
1805                    ctx.lineage.record(result.lineage.clone())?;
1806                    output_handles.insert(node_id.clone(), result.outputs.clone());
1807                    input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1808                    results.push(result);
1809                }
1810            }
1811        }
1812
1813        Ok(results)
1814    }
1815}
1816
1817#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1818enum HpoTrialTerminalState {
1819    Completed,
1820    Pruned,
1821    Failed,
1822}
1823
1824fn validate_hpo_checkpoint_result(
1825    checkpoint: &RuntimeHpoCheckpointResult,
1826    hpo: &RuntimeHpoExecutionContext,
1827    trial_variants: &BTreeMap<i64, VariantId>,
1828    terminal_trials: &BTreeMap<i64, HpoTrialTerminalState>,
1829    history_at_start: u32,
1830) -> Result<()> {
1831    checkpoint.artifact.validate().map_err(|error| {
1832        DagMlError::RuntimeValidation(format!(
1833            "runtime HPO checkpoint artifact is invalid: {error}"
1834        ))
1835    })?;
1836    if checkpoint.operation_id != hpo.operation_id
1837        || checkpoint.controller_id != hpo.controller_id
1838        || checkpoint.target_node_id != hpo.target_node_id
1839        || checkpoint.provenance != hpo.provenance
1840    {
1841        return Err(DagMlError::RuntimeValidation(
1842            "runtime HPO checkpoint provenance does not exactly match its execution context"
1843                .to_string(),
1844        ));
1845    }
1846    let proposed_count = u32::try_from(trial_variants.len()).map_err(|_| {
1847        DagMlError::RuntimeValidation(
1848            "runtime HPO scheduler proposal count does not fit u32".to_string(),
1849        )
1850    })?;
1851    if checkpoint.trial_history_len != hpo.trial_budget_total
1852        || checkpoint.trial_history_len < history_at_start
1853        || checkpoint.trial_history_len - history_at_start != proposed_count
1854    {
1855        return Err(DagMlError::RuntimeValidation(
1856            "runtime HPO checkpoint native history is inconsistent with scheduler-observed trials"
1857                .to_string(),
1858        ));
1859    }
1860    if checkpoint.artifact.binding.controller_id != hpo.controller_id.as_str()
1861        || checkpoint.artifact.binding.controller_id != hpo.study.controller_id
1862        || checkpoint.artifact.binding.study_id != hpo.study.study_id
1863        || checkpoint.artifact.methods_abi != hpo.study.methods_abi
1864    {
1865        return Err(DagMlError::RuntimeValidation(
1866            "runtime HPO checkpoint binding/controller/study does not match the active tuner"
1867                .to_string(),
1868        ));
1869    }
1870    let expected_search_space = hpo.study.search_space.fingerprint().map_err(|error| {
1871        DagMlError::RuntimeValidation(format!(
1872            "runtime HPO cannot fingerprint the configured search space: {error}"
1873        ))
1874    })?;
1875    if checkpoint.artifact.binding.search_space_fingerprint != expected_search_space {
1876        return Err(DagMlError::RuntimeValidation(
1877            "runtime HPO checkpoint search-space binding does not match the active study"
1878                .to_string(),
1879        ));
1880    }
1881
1882    let completed_trial_ids = terminal_trials
1883        .iter()
1884        .filter_map(|(trial_id, state)| {
1885            (*state == HpoTrialTerminalState::Completed).then_some(*trial_id)
1886        })
1887        .collect::<BTreeSet<_>>();
1888    let mut proposal_trial_ids = BTreeSet::new();
1889    for proposal in &checkpoint.completed_proposals {
1890        if !proposal_trial_ids.insert(proposal.trial_id) {
1891            return Err(DagMlError::RuntimeValidation(format!(
1892                "runtime HPO checkpoint has duplicate completed proposal for trial `{}`",
1893                proposal.trial_id
1894            )));
1895        }
1896        if trial_variants.get(&proposal.trial_id) != Some(&proposal.variant.variant_id) {
1897            return Err(DagMlError::RuntimeValidation(format!(
1898                "runtime HPO checkpoint proposal for trial `{}` does not exactly match its scheduler proposal",
1899                proposal.trial_id
1900            )));
1901        }
1902    }
1903    if proposal_trial_ids != completed_trial_ids {
1904        return Err(DagMlError::RuntimeValidation(
1905            "runtime HPO checkpoint proposals must cover exactly the completed trials".to_string(),
1906        ));
1907    }
1908
1909    let mut report_trial_ids = BTreeSet::new();
1910    for completed in &checkpoint.completed_reports {
1911        if !report_trial_ids.insert(completed.trial_id) {
1912            return Err(DagMlError::RuntimeValidation(format!(
1913                "runtime HPO checkpoint has duplicate completed report for trial `{}`",
1914                completed.trial_id
1915            )));
1916        }
1917        if trial_variants.get(&completed.trial_id) != Some(&completed.variant_id)
1918            || !proposal_trial_ids.contains(&completed.trial_id)
1919        {
1920            return Err(DagMlError::RuntimeValidation(format!(
1921                "runtime HPO checkpoint report for trial `{}` does not match a completed proposal",
1922                completed.trial_id
1923            )));
1924        }
1925        let report = &completed.report;
1926        if report.producer_node != hpo.selection.producer_node
1927            || report.producer_port.as_deref() != Some(hpo.selection.producer_port.as_str())
1928            || report.partition != PredictionPartition::Validation
1929            || report
1930                .fold_id
1931                .as_ref()
1932                .is_none_or(|fold| fold.as_str() != "avg")
1933            || report.variant_id.as_ref() != Some(&completed.variant_id)
1934            || !report
1935                .metrics
1936                .get(hpo.selection.metric.name())
1937                .is_some_and(|score| score.is_finite())
1938        {
1939            return Err(DagMlError::RuntimeValidation(format!(
1940                "runtime HPO checkpoint report for trial `{}` is not its one finite target OOF average",
1941                completed.trial_id
1942            )));
1943        }
1944    }
1945    if report_trial_ids != completed_trial_ids {
1946        return Err(DagMlError::RuntimeValidation(
1947            "runtime HPO checkpoint reports must cover exactly one OOF average per completed trial"
1948                .to_string(),
1949        ));
1950    }
1951    Ok(())
1952}
1953
1954pub(crate) struct PreparedNodeTask {
1955    pub(crate) node_id: NodeId,
1956    pub(crate) task: NodeTask,
1957}
1958
1959// This module stays adjacent to the scheduler-owned task preparation it
1960// exercises; the remaining helpers below are shared by both schedulers.
1961#[cfg(test)]
1962#[allow(clippy::items_after_test_module)]
1963mod hpo_scheduler_tests {
1964    use std::collections::{BTreeMap, BTreeSet};
1965    use std::sync::{Arc, Mutex};
1966
1967    use sha2::{Digest, Sha256};
1968
1969    use super::*;
1970    use crate::controller::{
1971        ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest,
1972        ControllerRegistry, RngPolicy,
1973    };
1974    use crate::data::InMemoryDataProvider;
1975    use crate::fold::{FoldAssignment, FoldPartitionMode};
1976    use crate::graph::{GraphInterface, GraphSpec, NodeSpec, PortSchema, PortSpec};
1977    use crate::hpo::{
1978        HpoDirection, HpoMetric, HpoOptimizerConfig, HpoParameter, HpoPruner, HpoSampler,
1979        HpoSearchSpace, HpoStudyBinding, MethodsHpoStudyConfig, N4moptCheckpointArtifact,
1980        N4MOPT_ARTIFACT_KIND, N4MOPT_CHECKPOINT_SCHEMA_VERSION, N4MOPT_FORMAT,
1981    };
1982    use crate::metrics::RegressionTargetBlock;
1983    use crate::oof::PredictionBlock;
1984    use crate::plan::{build_execution_plan, SplitInvocation};
1985
1986    struct HpoTestModel {
1987        id: ControllerId,
1988        trace: Arc<Mutex<Vec<String>>>,
1989    }
1990
1991    impl RuntimeController for HpoTestModel {
1992        fn controller_id(&self) -> &ControllerId {
1993            &self.id
1994        }
1995
1996        fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
1997            self.trace.lock().unwrap().push("model_cv".to_string());
1998            let sample_id = match task.fold_id.as_ref().map(FoldId::as_str) {
1999                Some("fold:0") => SampleId::new("sample:one").unwrap(),
2000                Some("fold:1") => SampleId::new("sample:two").unwrap(),
2001                other => {
2002                    return Err(DagMlError::RuntimeValidation(format!(
2003                        "HPO test model received unexpected fold {other:?}"
2004                    )));
2005                }
2006            };
2007            Ok(NodeResult {
2008                schema_version: None,
2009                node_id: task.node_plan.node_id.clone(),
2010                outputs: BTreeMap::from([(
2011                    "prediction".to_string(),
2012                    HandleRef {
2013                        handle: 2,
2014                        kind: HandleKind::Prediction,
2015                        owner_controller: self.id.clone(),
2016                    },
2017                )]),
2018                predictions: vec![PredictionBlock {
2019                    prediction_id: Some(format!("prediction:{}", task.fold_id.as_ref().unwrap())),
2020                    producer_node: task.node_plan.node_id.clone(),
2021                    producer_port: None,
2022                    partition: PredictionPartition::Validation,
2023                    fold_id: task.fold_id.clone(),
2024                    sample_ids: vec![sample_id.clone()],
2025                    values: vec![vec![1.0]],
2026                    target_names: vec!["target".to_string()],
2027                }],
2028                observation_predictions: Vec::new(),
2029                aggregated_predictions: Vec::new(),
2030                explanations: Vec::new(),
2031                shape_deltas: Vec::new(),
2032                artifacts: Vec::new(),
2033                artifact_handles: BTreeMap::new(),
2034                fit_influence_diagnostics: Vec::new(),
2035                regression_targets: vec![RegressionTargetBlock {
2036                    level: PredictionLevel::Sample,
2037                    unit_ids: vec![PredictionUnitId::Sample(sample_id)],
2038                    values: vec![vec![1.0]],
2039                    target_names: vec!["target".to_string()],
2040                }],
2041                lineage: LineageRecord {
2042                    record_id: LineageId::new(format!(
2043                        "lineage:hpo-model:{}",
2044                        task.fold_id.as_ref().unwrap()
2045                    ))
2046                    .unwrap(),
2047                    run_id: task.run_id.clone(),
2048                    node_id: task.node_plan.node_id.clone(),
2049                    phase: task.phase,
2050                    controller_id: self.id.clone(),
2051                    controller_version: task.node_plan.controller_version.clone(),
2052                    variant_id: task.variant_id.clone(),
2053                    fold_id: task.fold_id.clone(),
2054                    branch_path: Vec::new(),
2055                    input_lineage: Vec::new(),
2056                    artifact_refs: Vec::new(),
2057                    params_fingerprint: task.node_plan.params_fingerprint.clone(),
2058                    data_model_shape_fingerprint: None,
2059                    aggregation_policy_fingerprint: None,
2060                    seed: task.seed,
2061                    unsafe_flags: BTreeSet::new(),
2062                    metrics: BTreeMap::new(),
2063                    loss_attestations: Vec::new(),
2064                    early_stopping_records: Vec::new(),
2065                },
2066            })
2067        }
2068    }
2069
2070    struct HpoTestTuner {
2071        id: ControllerId,
2072        trace: Arc<Mutex<Vec<String>>>,
2073        history_len: u32,
2074        proposal_count: u32,
2075    }
2076
2077    struct HpoTestSession {
2078        proposals: Vec<RuntimeHpoProposal>,
2079        trace: Arc<Mutex<Vec<String>>>,
2080        checkpoint: N4moptCheckpointArtifact,
2081        history_len: u32,
2082        completed: Option<(i64, f64)>,
2083    }
2084
2085    impl RuntimeController for HpoTestTuner {
2086        fn controller_id(&self) -> &ControllerId {
2087            &self.id
2088        }
2089
2090        fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
2091            Err(DagMlError::RuntimeValidation(format!(
2092                "HPO test tuner `{}` was dispatched through generic invoke",
2093                task.node_plan.node_id
2094            )))
2095        }
2096
2097        fn create_tuner_session(
2098            &self,
2099            task: &RuntimeHpoCampaignTask,
2100            context: &RuntimeHpoExecutionContext,
2101        ) -> Result<Box<dyn RuntimeTunerSession>> {
2102            assert_eq!(task.operation_id, context.operation_id);
2103            self.trace
2104                .lock()
2105                .unwrap()
2106                .push("session_factory".to_string());
2107            let payload = vec![7_u8];
2108            let proposals = (0..self.proposal_count)
2109                .map(|offset| {
2110                    let trial_id = i64::from(self.history_len + offset + 1);
2111                    let mut variant = context.base_variant.clone();
2112                    if self.history_len != 0 || self.proposal_count != 1 {
2113                        variant.variant_id = VariantId::new(format!("hpo:trial:{trial_id}"))
2114                            .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?;
2115                        variant.fingerprint = format!("hpo-test-{trial_id}");
2116                    }
2117                    Ok(RuntimeHpoProposal { trial_id, variant })
2118                })
2119                .collect::<Result<Vec<_>>>()?;
2120            Ok(Box::new(HpoTestSession {
2121                proposals: proposals.into_iter().rev().collect(),
2122                trace: Arc::clone(&self.trace),
2123                history_len: self.history_len,
2124                completed: None,
2125                checkpoint: N4moptCheckpointArtifact {
2126                    schema_version: N4MOPT_CHECKPOINT_SCHEMA_VERSION,
2127                    artifact_kind: N4MOPT_ARTIFACT_KIND.to_string(),
2128                    format: N4MOPT_FORMAT.to_string(),
2129                    binding: HpoStudyBinding {
2130                        controller_id: context.study.controller_id.clone(),
2131                        study_id: context.study.study_id.clone(),
2132                        search_space_fingerprint: context
2133                            .study
2134                            .search_space
2135                            .fingerprint()
2136                            .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?,
2137                        optimizer_fingerprint: "optimizer:test".to_string(),
2138                    },
2139                    methods_abi: context.study.methods_abi.clone(),
2140                    payload_sha256: format!("{:x}", Sha256::digest(&payload)),
2141                    opaque_payload: payload,
2142                },
2143            }))
2144        }
2145    }
2146
2147    impl RuntimeTunerSession for HpoTestSession {
2148        fn trial_history_len(&self) -> Result<u32> {
2149            Ok(self.history_len)
2150        }
2151
2152        fn ask(&mut self) -> Result<Option<RuntimeHpoProposal>> {
2153            self.trace.lock().unwrap().push("ask".to_string());
2154            let proposal = self.proposals.pop();
2155            if proposal.is_some() {
2156                self.history_len += 1;
2157            }
2158            Ok(proposal)
2159        }
2160
2161        fn report_intermediate(
2162            &mut self,
2163            intermediate: RuntimeHpoIntermediate,
2164        ) -> Result<RuntimeHpoIntermediateOutcome> {
2165            assert_eq!(intermediate.step, 0);
2166            assert!(intermediate.score.is_finite());
2167            self.trace.lock().unwrap().push("intermediate".to_string());
2168            Ok(RuntimeHpoIntermediateOutcome::Continue)
2169        }
2170
2171        fn tell(&mut self, trial_id: i64, terminal: RuntimeHpoTerminal) -> Result<()> {
2172            assert!(trial_id > 0);
2173            assert!(
2174                matches!(terminal, RuntimeHpoTerminal::Completed { score } if score.is_finite())
2175            );
2176            self.trace.lock().unwrap().push("tell".to_string());
2177            if let RuntimeHpoTerminal::Completed { score } = terminal {
2178                self.completed = Some((trial_id, score));
2179            }
2180            Ok(())
2181        }
2182
2183        fn checkpoint(&self) -> Result<N4moptCheckpointArtifact> {
2184            self.trace.lock().unwrap().push("checkpoint".to_string());
2185            Ok(self.checkpoint.clone())
2186        }
2187
2188        fn incumbent(
2189            &self,
2190            variants: &BTreeMap<i64, VariantId>,
2191        ) -> Result<Option<RuntimeHpoIncumbent>> {
2192            let Some((trial_id, score)) = self.completed else {
2193                return Ok(None);
2194            };
2195            Ok(Some(RuntimeHpoIncumbent {
2196                trial_id,
2197                score,
2198                metric: "rmse".to_string(),
2199                direction: HpoDirection::Minimize,
2200                variant_id: variants.get(&trial_id).cloned().unwrap(),
2201            }))
2202        }
2203
2204        fn terminal_trial_snapshots(
2205            &self,
2206            variants: &BTreeMap<i64, VariantId>,
2207        ) -> Result<Vec<RuntimeHpoTerminalSnapshot>> {
2208            let (trial_id, score) = self.completed.ok_or_else(|| {
2209                DagMlError::RuntimeValidation("test HPO session has no completed trial".to_string())
2210            })?;
2211            Ok((1..=i64::from(self.history_len))
2212                .map(|id| {
2213                    let completed = id == trial_id;
2214                    RuntimeHpoTerminalSnapshot {
2215                        trial: crate::hpo::HpoTrial {
2216                            id,
2217                            ask_sequence: id,
2218                            terminal_sequence: Some(id),
2219                            parameters: BTreeMap::new(),
2220                            parameter_order: Vec::new(),
2221                            status: if completed {
2222                                crate::hpo::HpoTrialStatus::Completed
2223                            } else {
2224                                crate::hpo::HpoTrialStatus::Failed
2225                            },
2226                            score: completed.then_some(score),
2227                            rung: 0,
2228                            duration: 0.0,
2229                            intermediates: Vec::new(),
2230                            failure: (!completed).then(|| crate::hpo::HpoFailure {
2231                                code: "RESTORED_TEST_FAILURE".to_string(),
2232                                message: "synthetic restored terminal".to_string(),
2233                                retryable: false,
2234                            }),
2235                        },
2236                        variant_id: variants.get(&id).cloned(),
2237                    }
2238                })
2239                .collect())
2240        }
2241    }
2242
2243    fn node(id: &str, kind: NodeKind, outputs: Vec<PortSpec>) -> NodeSpec {
2244        NodeSpec {
2245            id: NodeId::new(id).unwrap(),
2246            kind,
2247            operator: None,
2248            params: BTreeMap::new(),
2249            ports: PortSchema {
2250                inputs: Vec::new(),
2251                outputs,
2252            },
2253            metadata: BTreeMap::new(),
2254            seed_label: None,
2255        }
2256    }
2257
2258    fn manifest(id: &str, kind: NodeKind) -> ControllerManifest {
2259        ControllerManifest {
2260            controller_id: ControllerId::new(id).unwrap(),
2261            controller_version: "test".to_string(),
2262            operator_kind: kind,
2263            priority: 0,
2264            supported_phases: BTreeSet::from([Phase::FitCv]),
2265            input_ports: Vec::new(),
2266            output_ports: Vec::new(),
2267            data_requirements: None,
2268            capabilities: BTreeSet::from([
2269                ControllerCapability::Deterministic,
2270                ControllerCapability::EmitsPredictions,
2271            ]),
2272            operator_selectors: Vec::new(),
2273            fit_scope: ControllerFitScope::FoldTrain,
2274            rng_policy: RngPolicy::UsesCoreSeed,
2275            artifact_policy: ArtifactPolicy::Serializable,
2276        }
2277    }
2278
2279    #[test]
2280    fn hpo_campaign_invokes_registered_session_and_routes_oof_feedback() {
2281        let target = NodeId::new("model:score").unwrap();
2282        let graph = GraphSpec {
2283            id: "graph:hpo.scheduler".to_string(),
2284            interface: GraphInterface::default(),
2285            nodes: vec![node(
2286                "model:score",
2287                NodeKind::Model,
2288                vec![PortSpec {
2289                    name: "prediction".to_string(),
2290                    kind: PortKind::Prediction,
2291                    representation: None,
2292                    cardinality: crate::graph::PortCardinality::One,
2293                    unit_level: None,
2294                    alignment_key: None,
2295                    target_level: None,
2296                    description: String::new(),
2297                }],
2298            )],
2299            edges: Vec::new(),
2300            search_space_fingerprint: None,
2301            metadata: BTreeMap::new(),
2302        };
2303        let fold_set = FoldSet {
2304            id: "folds:hpo".to_string(),
2305            sample_ids: vec![
2306                SampleId::new("sample:one").unwrap(),
2307                SampleId::new("sample:two").unwrap(),
2308            ],
2309            folds: vec![
2310                FoldAssignment {
2311                    fold_id: FoldId::new("fold:0").unwrap(),
2312                    train_sample_ids: vec![SampleId::new("sample:two").unwrap()],
2313                    validation_sample_ids: vec![SampleId::new("sample:one").unwrap()],
2314                    metadata: BTreeMap::new(),
2315                },
2316                FoldAssignment {
2317                    fold_id: FoldId::new("fold:1").unwrap(),
2318                    train_sample_ids: vec![SampleId::new("sample:one").unwrap()],
2319                    validation_sample_ids: vec![SampleId::new("sample:two").unwrap()],
2320                    metadata: BTreeMap::new(),
2321                },
2322            ],
2323            sample_groups: BTreeMap::new(),
2324            partition_mode: FoldPartitionMode::Partition,
2325        };
2326        let mut registry = ControllerRegistry::new();
2327        registry
2328            .register(manifest("controller:model", NodeKind::Model))
2329            .unwrap();
2330        let plan = build_execution_plan(
2331            "plan:hpo.scheduler",
2332            graph,
2333            CampaignSpec {
2334                inner_cv: None,
2335                id: "campaign:hpo.scheduler".to_string(),
2336                root_seed: Some(13),
2337                leakage_policy: Default::default(),
2338                aggregation_policy: Default::default(),
2339                split_invocation: Some(SplitInvocation {
2340                    id: "split:hpo".to_string(),
2341                    controller_id: None,
2342                    leakage_policy: Default::default(),
2343                    params: BTreeMap::new(),
2344                    fold_set: Some(fold_set),
2345                }),
2346                generation: Default::default(),
2347                shape_plans: BTreeMap::new(),
2348                data_bindings: BTreeMap::new(),
2349                branch_view_plans: Vec::new(),
2350                metadata: BTreeMap::new(),
2351            },
2352            &registry,
2353        )
2354        .unwrap();
2355        let trace = Arc::new(Mutex::new(Vec::new()));
2356        let mut controllers = RuntimeControllerRegistry::new();
2357        controllers
2358            .register(Box::new(HpoTestTuner {
2359                id: ControllerId::new("controller:tuner").unwrap(),
2360                trace: Arc::clone(&trace),
2361                history_len: 0,
2362                proposal_count: 1,
2363            }))
2364            .unwrap();
2365        controllers
2366            .register(Box::new(HpoTestModel {
2367                id: ControllerId::new("controller:model").unwrap(),
2368                trace: Arc::clone(&trace),
2369            }))
2370            .unwrap();
2371        let hpo = RuntimeHpoExecutionContext {
2372            operation_id: "hpo:test".to_string(),
2373            controller_id: ControllerId::new("controller:tuner").unwrap(),
2374            target_node_id: target.clone(),
2375            base_variant: plan.variants[0].clone(),
2376            trial_budget_total: 1,
2377            study: MethodsHpoStudyConfig {
2378                controller_id: "controller:tuner".to_string(),
2379                study_id: "study:hpo.scheduler".to_string(),
2380                methods_abi: "test-abi".to_string(),
2381                search_space: HpoSearchSpace {
2382                    parameters: vec![HpoParameter::Int {
2383                        name: "n_components".to_string(),
2384                        low: 1,
2385                        high: 1,
2386                        step: 1,
2387                        log: false,
2388                    }],
2389                },
2390                optimizer: HpoOptimizerConfig {
2391                    sampler: HpoSampler::Random,
2392                    pruner: HpoPruner::None,
2393                    direction: HpoDirection::Minimize,
2394                    metric: HpoMetric::Rmse,
2395                    seed: 13,
2396                    n_startup_trials: 1,
2397                    max_resource: 0,
2398                    reduction_factor: 1,
2399                },
2400            },
2401            parameter_paths: BTreeMap::from([(
2402                "n_components".to_string(),
2403                "n_components".to_string(),
2404            )]),
2405            resume_checkpoint: None,
2406            resume_variants: BTreeMap::new(),
2407            resume_terminal_trials: Vec::new(),
2408            selection: RuntimeHpoSelectionTarget {
2409                producer_node: target,
2410                producer_port: "prediction".to_string(),
2411                metric: RegressionMetricKind::Rmse,
2412                direction: HpoDirection::Minimize,
2413            },
2414            provenance: RuntimeHpoProvenance {
2415                graph_fingerprint: plan.graph_fingerprint.clone(),
2416                campaign_fingerprint: plan.campaign_fingerprint.clone(),
2417                controller_fingerprint: plan.controller_fingerprint.clone(),
2418                data_identities_fingerprint: "identity:test".to_string(),
2419                fold_set_fingerprint: plan
2420                    .fold_set
2421                    .as_ref()
2422                    .map(stable_json_fingerprint)
2423                    .transpose()
2424                    .unwrap(),
2425                training_influence_fingerprint: "influence:test".to_string(),
2426                relation_fingerprint: "relation:test".to_string(),
2427            },
2428        };
2429        let provider = InMemoryDataProvider::new(ControllerId::new("controller:data").unwrap());
2430        let ctx = RunContext::new(RunId::new("run:hpo.scheduler").unwrap(), Some(13));
2431
2432        let result = SequentialScheduler
2433            .execute_hpo_campaign(&plan, &controllers, &provider, &ctx, &hpo)
2434            .unwrap();
2435
2436        assert_eq!(result.operation_id, "hpo:test");
2437        assert_eq!(result.candidates.len(), 1);
2438        assert_eq!(result.checkpoint.completed_proposals.len(), 1);
2439        assert_eq!(result.checkpoint.completed_reports.len(), 1);
2440        assert_eq!(result.candidates[0].lineage.len(), 2);
2441        assert_eq!(result.incumbent.variant_id, plan.variants[0].variant_id);
2442        let mut selected_ctx = RunContext::new(RunId::new("run:hpo.scheduler").unwrap(), Some(13));
2443        selected_ctx.variant_id = Some(plan.variants[0].variant_id.clone());
2444        let selected_results = SequentialScheduler
2445            .execute_campaign_phase_with_data_provider(
2446                &plan,
2447                &controllers,
2448                &provider,
2449                &mut selected_ctx,
2450                Phase::FitCv,
2451            )
2452            .unwrap();
2453        assert_eq!(selected_results.len(), 2);
2454        assert_eq!(selected_ctx.lineage.len(), 2);
2455        assert_eq!(
2456            trace.lock().unwrap().as_slice(),
2457            [
2458                "session_factory",
2459                "ask",
2460                "model_cv",
2461                "model_cv",
2462                "intermediate",
2463                "tell",
2464                "checkpoint",
2465                "model_cv",
2466                "model_cv"
2467            ]
2468        );
2469
2470        // The native study may restore failed/pruned history that has no
2471        // completed proposal evidence. Its local count, not the coordinator's
2472        // persisted completed list, determines the remaining global budget.
2473        let resumed_trace = Arc::new(Mutex::new(Vec::new()));
2474        let mut resumed_controllers = RuntimeControllerRegistry::new();
2475        resumed_controllers
2476            .register(Box::new(HpoTestTuner {
2477                id: ControllerId::new("controller:tuner").unwrap(),
2478                trace: Arc::clone(&resumed_trace),
2479                history_len: 2,
2480                proposal_count: 2,
2481            }))
2482            .unwrap();
2483        resumed_controllers
2484            .register(Box::new(HpoTestModel {
2485                id: ControllerId::new("controller:model").unwrap(),
2486                trace: Arc::clone(&resumed_trace),
2487            }))
2488            .unwrap();
2489        let mut resumed_hpo = hpo.clone();
2490        resumed_hpo.trial_budget_total = 4;
2491        let resumed_ctx = RunContext::new(RunId::new("run:hpo.resumed").unwrap(), Some(13));
2492        let resumed = SequentialScheduler
2493            .execute_hpo_campaign(
2494                &plan,
2495                &resumed_controllers,
2496                &provider,
2497                &resumed_ctx,
2498                &resumed_hpo,
2499            )
2500            .unwrap();
2501        assert_eq!(resumed.candidates.len(), 2);
2502        assert_eq!(resumed.checkpoint.trial_history_len, 4);
2503        assert_eq!(
2504            resumed_trace
2505                .lock()
2506                .unwrap()
2507                .iter()
2508                .filter(|event| event.as_str() == "ask")
2509                .count(),
2510            2
2511        );
2512
2513        let mut over_budget_controllers = RuntimeControllerRegistry::new();
2514        over_budget_controllers
2515            .register(Box::new(HpoTestTuner {
2516                id: ControllerId::new("controller:tuner").unwrap(),
2517                trace: Arc::new(Mutex::new(Vec::new())),
2518                history_len: 5,
2519                proposal_count: 0,
2520            }))
2521            .unwrap();
2522        let error = SequentialScheduler
2523            .execute_hpo_campaign(
2524                &plan,
2525                &over_budget_controllers,
2526                &provider,
2527                &resumed_ctx,
2528                &resumed_hpo,
2529            )
2530            .unwrap_err();
2531        assert!(error.to_string().contains("exceeds total trial budget"));
2532    }
2533}
2534
2535pub(crate) fn attach_coordinator_input_lineage(
2536    result: &mut NodeResult,
2537    plan: &ExecutionPlan,
2538    node_id: &NodeId,
2539    upstream_lineage: &BTreeMap<NodeId, LineageId>,
2540) -> Result<()> {
2541    let inferred = inferred_input_lineage_for_node(plan, node_id, upstream_lineage);
2542    if result.lineage.input_lineage.is_empty() {
2543        result.lineage.input_lineage = inferred;
2544        return Ok(());
2545    }
2546
2547    let declared = result
2548        .lineage
2549        .input_lineage
2550        .iter()
2551        .cloned()
2552        .collect::<BTreeSet<_>>()
2553        .into_iter()
2554        .collect::<Vec<_>>();
2555    if declared != inferred {
2556        return Err(DagMlError::RuntimeValidation(format!(
2557            "lineage for node `{}` declared input lineage {:?}, expected {:?}",
2558            result.node_id, declared, inferred
2559        )));
2560    }
2561    result.lineage.input_lineage = declared;
2562    Ok(())
2563}
2564
2565pub(crate) fn inferred_input_lineage_for_node(
2566    plan: &ExecutionPlan,
2567    node_id: &NodeId,
2568    upstream_lineage: &BTreeMap<NodeId, LineageId>,
2569) -> Vec<LineageId> {
2570    plan.graph_plan
2571        .graph
2572        .edges
2573        .iter()
2574        .filter(|edge| &edge.target.node_id == node_id && edge.contract.propagates_lineage)
2575        .filter_map(|edge| upstream_lineage.get(&edge.source.node_id).cloned())
2576        .collect::<BTreeSet<_>>()
2577        .into_iter()
2578        .collect()
2579}
2580pub(crate) fn collect_input_handles(
2581    plan: &ExecutionPlan,
2582    node_plan: &NodePlan,
2583    output_handles: &BTreeMap<NodeId, BTreeMap<String, HandleRef>>,
2584    output_data_views: &BTreeMap<NodeId, BTreeMap<String, DataProviderViewSpec>>,
2585    resources: &PhaseScopeResources<'_>,
2586    ctx: &RunContext,
2587    scope: &PhaseScope,
2588) -> Result<CollectedInputs> {
2589    let mut inputs = BTreeMap::new();
2590    let mut data_views = BTreeMap::new();
2591    let mut prediction_inputs = BTreeMap::new();
2592    let training_oof_edges = incoming_training_oof_edges(plan, node_plan, scope)?;
2593    // An OOF edge replaces exactly one raw producer port. Do not hide sibling
2594    // outputs from the same producer: a meta-node may legally consume both an
2595    // OOF prediction port and an auxiliary non-OOF port. PREDICT has no
2596    // Validation-OOF input, but its raw prediction port must still be masked so
2597    // only the explicit `:predict` off-fold input reaches the controller.
2598    let masked_oof_source_ports = if scope.phase == Phase::Predict {
2599        incoming_oof_edges(plan, node_plan)?
2600    } else {
2601        training_oof_edges.clone()
2602    }
2603    .into_iter()
2604    .map(|edge| (edge.source.node_id.clone(), edge.source.port_name.clone()))
2605    .collect::<BTreeSet<_>>();
2606    let bound_data_inputs = node_plan
2607        .data_bindings
2608        .iter()
2609        .map(|binding| binding.input_name.clone())
2610        .collect::<BTreeSet<_>>();
2611    // Only forward upstream handles for ports this node DECLARES an edge to.
2612    // A controller must never see a handle outside its declared port contract,
2613    // so a sibling consumer of the same producer cannot expose extra ports here.
2614    let declared_source_ports = plan
2615        .graph_plan
2616        .graph
2617        .edges
2618        .iter()
2619        .filter(|edge| edge.target.node_id == node_plan.node_id)
2620        .map(|edge| (edge.source.node_id.clone(), edge.source.port_name.clone()))
2621        .collect::<BTreeSet<_>>();
2622    for upstream in &node_plan.input_nodes {
2623        if let Some(handles) = output_handles.get(upstream) {
2624            for (port, handle) in handles {
2625                if !declared_source_ports.contains(&(upstream.clone(), port.clone())) {
2626                    continue;
2627                }
2628                if masked_oof_source_ports.contains(&(upstream.clone(), port.clone())) {
2629                    continue;
2630                }
2631                inputs.insert(format!("{upstream}.{port}"), handle.clone());
2632            }
2633        }
2634    }
2635    for edge in plan
2636        .graph_plan
2637        .graph
2638        .edges
2639        .iter()
2640        .filter(|edge| edge.target.node_id == node_plan.node_id)
2641        .filter(|edge| edge.contract.kind == PortKind::Data && !edge.contract.requires_oof)
2642    {
2643        if bound_data_inputs.contains(&edge.target.port_name) {
2644            continue;
2645        }
2646        let Some(handles) = output_handles.get(&edge.source.node_id) else {
2647            continue;
2648        };
2649        let Some(handle) = handles.get(&edge.source.port_name) else {
2650            continue;
2651        };
2652        let key = data_view_key(&edge.target.port_name);
2653        if inputs.insert(key.clone(), handle.clone()).is_some() {
2654            return Err(DagMlError::RuntimeValidation(format!(
2655                "node `{}` received duplicate data edge input `{key}`",
2656                node_plan.node_id
2657            )));
2658        }
2659        if let Some(source_views) = output_data_views.get(&edge.source.node_id) {
2660            if let Some(view) = source_views.get(&edge.source.port_name) {
2661                if data_views.insert(key.clone(), view.clone()).is_some() {
2662                    return Err(DagMlError::RuntimeValidation(format!(
2663                        "node `{}` received duplicate data edge view `{key}`",
2664                        node_plan.node_id
2665                    )));
2666                }
2667            }
2668            let source_validation_key = validation_data_view_key(&edge.source.port_name);
2669            if let Some(view) = source_views.get(&source_validation_key) {
2670                let validation_key = format!("{key}:validation");
2671                if data_views
2672                    .insert(validation_key.clone(), view.clone())
2673                    .is_some()
2674                {
2675                    return Err(DagMlError::RuntimeValidation(format!(
2676                        "node `{}` received duplicate data edge validation view `{validation_key}`",
2677                        node_plan.node_id
2678                    )));
2679                }
2680            }
2681        }
2682    }
2683    for edge in training_oof_edges {
2684        let key = format!("{}.{}", edge.source.node_id, edge.source.port_name);
2685        let Some(input) = collect_oof_prediction_input(plan, edge, ctx, scope, resources)? else {
2686            return Ok(CollectedInputs {
2687                handles: BTreeMap::new(),
2688                data_views: BTreeMap::new(),
2689                prediction_inputs: BTreeMap::new(),
2690                skip_node: true,
2691            });
2692        };
2693        if inputs.insert(key.clone(), input.handle).is_some() {
2694            return Err(DagMlError::RuntimeValidation(format!(
2695                "node `{}` received duplicate OOF prediction input `{key}`",
2696                node_plan.node_id
2697            )));
2698        }
2699        if prediction_inputs.insert(key.clone(), input.spec).is_some() {
2700            return Err(DagMlError::RuntimeValidation(format!(
2701                "node `{}` received duplicate OOF prediction spec `{key}`",
2702                node_plan.node_id
2703            )));
2704        }
2705    }
2706    // REFIT / PREDICT: deliver each base producer's off-fold (test / predict)
2707    // predictions to the stacking meta-node as a SEPARATE prediction input (suffixed
2708    // `:test` / `:predict`) so the host meta-model predicts from them. The FIT_CV
2709    // Validation-OOF input above is the meta-features the meta-model trains on; this
2710    // off-fold input is used ONLY for REFIT/PREDICT scoring/prediction, never FIT_CV
2711    // training — keeping the leakage invariant intact.
2712    if matches!(scope.phase, Phase::Refit | Phase::Predict) {
2713        let off_fold_suffix = scope.phase.as_str().to_ascii_lowercase();
2714        for edge in incoming_oof_edges(plan, node_plan)? {
2715            let Some(input) = collect_off_fold_prediction_input(plan, edge, ctx, scope)? else {
2716                continue;
2717            };
2718            let key = format!(
2719                "{}.{}:{off_fold_suffix}",
2720                edge.source.node_id, edge.source.port_name
2721            );
2722            if inputs.insert(key.clone(), input.handle).is_some() {
2723                return Err(DagMlError::RuntimeValidation(format!(
2724                    "node `{}` received duplicate off-fold prediction input `{key}`",
2725                    node_plan.node_id
2726                )));
2727            }
2728            if prediction_inputs.insert(key.clone(), input.spec).is_some() {
2729                return Err(DagMlError::RuntimeValidation(format!(
2730                    "node `{}` received duplicate off-fold prediction spec `{key}`",
2731                    node_plan.node_id
2732                )));
2733            }
2734        }
2735    }
2736    if !node_plan.data_bindings.is_empty() && resources.data_provider.is_none() {
2737        return Err(DagMlError::RuntimeValidation(format!(
2738            "node `{}` requires {} data binding(s) but no runtime data provider is registered",
2739            node_plan.node_id,
2740            node_plan.data_bindings.len()
2741        )));
2742    }
2743    if let Some(data_provider) = resources.data_provider {
2744        // Samples excluded from training (sample-local) are relevant only to
2745        // fitting scopes. A top-level PREDICT must not even resolve the CV
2746        // relation authority: its separately attested cohort below owns the
2747        // complete identity universe for that read.
2748        let excluded_samples = if scope.phase == Phase::Predict {
2749            BTreeSet::new()
2750        } else {
2751            coordinator_relations_for_node(node_plan, resources)?
2752                .map(|relations| relations.excluded_sample_ids())
2753                .unwrap_or_default()
2754        };
2755        let scope_fold_set = resources.fold_set_override.or(plan.fold_set.as_ref());
2756        for binding in &node_plan.data_bindings {
2757            let predict_cohort = if scope.phase == Phase::Predict {
2758                data_provider.predict_cohort(binding, scope.phase)?
2759            } else {
2760                None
2761            };
2762            let materialized = data_provider.materialize(&DataMaterializationRequest {
2763                run_id: ctx.run_id.clone(),
2764                node_id: node_plan.node_id.clone(),
2765                input_name: binding.input_name.clone(),
2766                phase: scope.phase,
2767                variant_id: scope.variant_id.clone(),
2768                fold_id: scope.fold_id.clone(),
2769                binding: binding.clone(),
2770                predict_cohort: predict_cohort.clone(),
2771            })?;
2772            let branch_view_for_node = branch_view_from_node_metadata(plan, &node_plan.node_id)?;
2773            let mut view = data_view_for_scope(
2774                binding,
2775                scope_fold_set,
2776                scope,
2777                branch_view_for_node.as_ref(),
2778                &excluded_samples,
2779            )?;
2780            if let Some(cohort) = predict_cohort.as_ref() {
2781                bind_predict_cohort_to_view(&mut view, cohort)?;
2782            }
2783            let key = data_view_key(&binding.input_name);
2784            let view_handle = make_data_view_handle(
2785                data_provider,
2786                ctx,
2787                node_plan,
2788                scope,
2789                binding,
2790                DataViewHandleInput {
2791                    data_handle: &materialized,
2792                    view: &view,
2793                    predict_cohort: predict_cohort.as_ref(),
2794                },
2795            )?;
2796            if data_views.insert(key.clone(), view).is_some() {
2797                return Err(DagMlError::RuntimeValidation(format!(
2798                    "node `{}` received duplicate data view `{key}`",
2799                    node_plan.node_id
2800                )));
2801            }
2802            if inputs.insert(key.clone(), view_handle).is_some() {
2803                return Err(DagMlError::RuntimeValidation(format!(
2804                    "node `{}` received duplicate data input `{key}`",
2805                    node_plan.node_id
2806                )));
2807            }
2808
2809            if let Some(validation_view) = validation_data_view_for_scope(
2810                binding,
2811                scope_fold_set,
2812                scope,
2813                branch_view_for_node.as_ref(),
2814                &excluded_samples,
2815            )? {
2816                let validation_key = format!("{key}:validation");
2817                let validation_handle = make_data_view_handle(
2818                    data_provider,
2819                    ctx,
2820                    node_plan,
2821                    scope,
2822                    binding,
2823                    DataViewHandleInput {
2824                        data_handle: &materialized,
2825                        view: &validation_view,
2826                        predict_cohort: None,
2827                    },
2828                )?;
2829                if data_views
2830                    .insert(validation_key.clone(), validation_view)
2831                    .is_some()
2832                {
2833                    return Err(DagMlError::RuntimeValidation(format!(
2834                        "node `{}` received duplicate validation data view `{validation_key}`",
2835                        node_plan.node_id
2836                    )));
2837                }
2838                if inputs
2839                    .insert(validation_key.clone(), validation_handle)
2840                    .is_some()
2841                {
2842                    return Err(DagMlError::RuntimeValidation(format!(
2843                        "node `{}` received duplicate validation data input `{validation_key}`",
2844                        node_plan.node_id
2845                    )));
2846                }
2847            }
2848        }
2849    }
2850    Ok(CollectedInputs {
2851        handles: inputs,
2852        data_views,
2853        prediction_inputs,
2854        skip_node: false,
2855    })
2856}
2857pub(crate) fn preload_replay_prediction_cache_store(
2858    bundle: &ExecutionBundle,
2859    prediction_cache_store: Option<&dyn RuntimePredictionCacheStore>,
2860    ctx: &mut RunContext,
2861) -> Result<()> {
2862    if bundle.prediction_requirements.is_empty() {
2863        return Ok(());
2864    }
2865    let store = prediction_cache_store.ok_or_else(|| {
2866        DagMlError::RuntimeValidation(format!(
2867            "bundle `{}` cannot preload OOF prediction caches without a prediction cache store",
2868            bundle.bundle_id
2869        ))
2870    })?;
2871    if !ctx.prediction_store.blocks().is_empty() {
2872        return Err(DagMlError::RuntimeValidation(format!(
2873            "bundle `{}` cannot preload OOF prediction caches into a non-empty prediction store",
2874            bundle.bundle_id
2875        )));
2876    }
2877    let contracts = replay_prediction_cache_contracts(bundle)?;
2878    for contract in contracts.values() {
2879        if contract.requirement.prediction_level == PredictionLevel::Sample {
2880            let blocks = store.load_blocks(&contract.cache.requirement_key)?;
2881            if blocks.iter().any(|block| {
2882                block.producer_node != contract.requirement.producer_node
2883                    || block.partition != contract.requirement.partition
2884            }) {
2885                return Err(DagMlError::RuntimeValidation(format!(
2886                    "prediction cache store returned blocks outside requirement `{}`",
2887                    contract.cache.requirement_key
2888                )));
2889            }
2890            let mut payload = build_prediction_cache_payload(&contract.requirement, &blocks)?;
2891            payload.cache_namespace_fingerprints =
2892                contract.cache.cache_namespace_fingerprints.clone();
2893            validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
2894            for block in &payload.blocks {
2895                ctx.prediction_store.append(block.clone())?;
2896            }
2897        } else {
2898            let blocks = store.load_aggregated_blocks(&contract.cache.requirement_key)?;
2899            if blocks.iter().any(|block| {
2900                block.producer_node != contract.requirement.producer_node
2901                    || block.partition != contract.requirement.partition
2902                    || block.level != contract.requirement.prediction_level
2903            }) {
2904                return Err(DagMlError::RuntimeValidation(format!(
2905                    "prediction cache store returned aggregated blocks outside requirement `{}`",
2906                    contract.cache.requirement_key
2907                )));
2908            }
2909            let mut payload =
2910                build_aggregated_prediction_cache_payload(&contract.requirement, &blocks)?;
2911            payload.cache_namespace_fingerprints =
2912                contract.cache.cache_namespace_fingerprints.clone();
2913            validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
2914        }
2915    }
2916    Ok(())
2917}
2918
2919pub(crate) fn replay_prediction_cache_contracts(
2920    bundle: &ExecutionBundle,
2921) -> Result<BTreeMap<String, ReplayPredictionCacheContract>> {
2922    bundle.validate()?;
2923    let requirements = bundle
2924        .prediction_requirements
2925        .iter()
2926        .map(|requirement| (requirement.key(), requirement))
2927        .collect::<BTreeMap<_, _>>();
2928    let mut contracts = BTreeMap::new();
2929    for cache in &bundle.prediction_caches {
2930        let requirement = requirements.get(&cache.requirement_key).ok_or_else(|| {
2931            DagMlError::RuntimeValidation(format!(
2932                "prediction cache `{}` references unknown prediction requirement `{}`",
2933                cache.cache_id, cache.requirement_key
2934            ))
2935        })?;
2936        contracts.insert(
2937            cache.requirement_key.clone(),
2938            ReplayPredictionCacheContract {
2939                requirement: (*requirement).clone(),
2940                cache: cache.clone(),
2941            },
2942        );
2943    }
2944    Ok(contracts)
2945}
2946
2947pub(crate) fn materialize_replay_artifact_handles(
2948    plan: &ExecutionPlan,
2949    bundle: &ExecutionBundle,
2950    replay_request: &ReplayPhaseRequest,
2951    artifact_store: &dyn RuntimeArtifactStore,
2952    ctx: &RunContext,
2953) -> Result<MaterializedReplayArtifacts> {
2954    let mut handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
2955    let mut inputs = BTreeMap::<NodeId, BTreeMap<String, ArtifactInputSpec>>::new();
2956    for artifact in &bundle.refit_artifacts {
2957        artifact.validate()?;
2958        let node_plan = plan.node_plans.get(&artifact.node_id).ok_or_else(|| {
2959            DagMlError::RuntimeValidation(format!(
2960                "bundle `{}` artifact references unknown node `{}`",
2961                bundle.bundle_id, artifact.node_id
2962            ))
2963        })?;
2964        if !node_plan.supported_phases.contains(&replay_request.phase) {
2965            return Err(DagMlError::RuntimeValidation(format!(
2966                "bundle `{}` artifact node `{}` does not support replay phase {:?}",
2967                bundle.bundle_id, artifact.node_id, replay_request.phase
2968            )));
2969        }
2970        let handle = artifact_store.materialize(&ArtifactMaterializationRequest {
2971            run_id: ctx.run_id.clone(),
2972            bundle_id: bundle.bundle_id.clone(),
2973            node_id: artifact.node_id.clone(),
2974            phase: replay_request.phase,
2975            variant_id: bundle.selected_variant_id.clone(),
2976            controller_id: artifact.controller_id.clone(),
2977            artifact: artifact.artifact.clone(),
2978            params_fingerprint: artifact.params_fingerprint.clone(),
2979            training_loss_fingerprint: artifact.training_loss_fingerprint.clone(),
2980        })?;
2981        if !matches!(handle.kind, HandleKind::Model | HandleKind::Artifact) {
2982            return Err(DagMlError::RuntimeValidation(format!(
2983                "artifact `{}` materialized as unsupported handle kind {:?}",
2984                artifact.artifact.id, handle.kind
2985            )));
2986        }
2987        if handle.owner_controller != artifact.controller_id {
2988            return Err(DagMlError::RuntimeValidation(format!(
2989                "artifact `{}` handle owner `{}` does not match controller `{}`",
2990                artifact.artifact.id, handle.owner_controller, artifact.controller_id
2991            )));
2992        }
2993        let key = refit_artifact_input_key(&artifact.artifact.id);
2994        if handles
2995            .entry(artifact.node_id.clone())
2996            .or_default()
2997            .insert(key.clone(), handle)
2998            .is_some()
2999        {
3000            return Err(DagMlError::RuntimeValidation(format!(
3001                "duplicate replay artifact input `{key}` for node `{}`",
3002                artifact.node_id
3003            )));
3004        }
3005        if inputs
3006            .entry(artifact.node_id.clone())
3007            .or_default()
3008            .insert(key.clone(), ArtifactInputSpec::from_refit_record(artifact)?)
3009            .is_some()
3010        {
3011            return Err(DagMlError::RuntimeValidation(format!(
3012                "duplicate replay artifact metadata `{key}` for node `{}`",
3013                artifact.node_id
3014            )));
3015        }
3016    }
3017    Ok(MaterializedReplayArtifacts { handles, inputs })
3018}
3019
3020pub(crate) fn derive_task_seed(
3021    root_seed: Option<u64>,
3022    variant_id: Option<&VariantId>,
3023    fold_id: Option<&FoldId>,
3024    node_plan: &NodePlan,
3025    phase: Phase,
3026) -> Option<u64> {
3027    root_seed.map(|root| {
3028        let mut context = SeedContext::root(root);
3029        if let Some(variant_id) = variant_id {
3030            context = context.child(format!("variant:{variant_id}"));
3031        }
3032        if let Some(fold_id) = fold_id {
3033            context = context.child(format!("fold:{fold_id}"));
3034        }
3035        context
3036            .child(format!("node:{}", node_plan.node_id))
3037            .child(format!("phase:{phase:?}"))
3038            .derive_u64("task")
3039    })
3040}
dag_ml_core/runtime/scheduler.rs

dag_ml_core/runtime/
scheduler.rs
