Skip to main content

dag_ml_core/runtime/
scheduler.rs

1// Auto-split from the former monolithic `runtime.rs` (pure refactor).
2use super::*;
3
4#[derive(Clone, Debug, Default)]
5pub struct SequentialScheduler;
6
7#[derive(Clone, Debug)]
8pub struct ParallelScheduler {
9    max_workers: usize,
10}
11
12impl ParallelScheduler {
13    pub fn new(max_workers: usize) -> Result<Self> {
14        if max_workers == 0 {
15            return Err(DagMlError::RuntimeValidation(
16                "parallel scheduler max_workers must be at least 1".to_string(),
17            ));
18        }
19        Ok(Self { max_workers })
20    }
21
22    pub fn max_workers(&self) -> usize {
23        self.max_workers
24    }
25}
26
27#[derive(Clone, Debug)]
28pub(crate) struct PhaseScope {
29    pub(crate) phase: Phase,
30    pub(crate) variant_id: Option<VariantId>,
31    pub(crate) variant: Option<VariantExecutionSpec>,
32    pub(crate) fold_id: Option<FoldId>,
33    pub(crate) seed_root: Option<u64>,
34}
35
36#[derive(Clone, Debug)]
37pub(crate) struct ReplayPredictionCacheContract {
38    pub(crate) requirement: BundlePredictionRequirement,
39    pub(crate) cache: BundlePredictionCacheRecord,
40}
41
42pub(crate) struct MaterializedReplayArtifacts {
43    pub(crate) handles: BTreeMap<NodeId, BTreeMap<String, HandleRef>>,
44    pub(crate) inputs: BTreeMap<NodeId, BTreeMap<String, ArtifactInputSpec>>,
45}
46
47#[derive(Default)]
48pub(crate) struct PhaseScopeResources<'a> {
49    pub(crate) data_provider: Option<&'a dyn RuntimeDataProvider>,
50    pub(crate) replay_artifact_handles: Option<&'a BTreeMap<NodeId, BTreeMap<String, HandleRef>>>,
51    pub(crate) replay_artifact_inputs:
52        Option<&'a BTreeMap<NodeId, BTreeMap<String, ArtifactInputSpec>>>,
53    pub(crate) replay_bundle_id: Option<&'a BundleId>,
54    pub(crate) data_envelopes: Option<&'a BTreeMap<String, ExternalDataPlanEnvelope>>,
55    pub(crate) prediction_cache_store: Option<&'a dyn RuntimePredictionCacheStore>,
56    pub(crate) prediction_cache_contracts:
57        Option<&'a BTreeMap<String, ReplayPredictionCacheContract>>,
58    pub(crate) artifact_store: Option<&'a mut InMemoryArtifactStore>,
59}
60
61impl SequentialScheduler {
62    pub fn execute_phase(
63        &self,
64        plan: &ExecutionPlan,
65        controllers: &RuntimeControllerRegistry,
66        ctx: &mut RunContext,
67        phase: Phase,
68    ) -> Result<Vec<NodeResult>> {
69        plan.validate()?;
70        let variant_id = ctx.variant_id.clone();
71        let seed_root = ctx.root_seed;
72        self.execute_phase_scope(
73            plan,
74            controllers,
75            ctx,
76            PhaseScope {
77                phase,
78                variant_id,
79                variant: None,
80                fold_id: None,
81                seed_root,
82            },
83            PhaseScopeResources::default(),
84        )
85    }
86
87    pub fn execute_phase_with_data_provider(
88        &self,
89        plan: &ExecutionPlan,
90        controllers: &RuntimeControllerRegistry,
91        data_provider: &dyn RuntimeDataProvider,
92        ctx: &mut RunContext,
93        phase: Phase,
94    ) -> Result<Vec<NodeResult>> {
95        plan.validate()?;
96        let variant_id = ctx.variant_id.clone();
97        let seed_root = ctx.root_seed;
98        self.execute_phase_scope(
99            plan,
100            controllers,
101            ctx,
102            PhaseScope {
103                phase,
104                variant_id,
105                variant: None,
106                fold_id: None,
107                seed_root,
108            },
109            PhaseScopeResources {
110                data_provider: Some(data_provider),
111                ..Default::default()
112            },
113        )
114    }
115
116    pub fn execute_campaign_phase(
117        &self,
118        plan: &ExecutionPlan,
119        controllers: &RuntimeControllerRegistry,
120        ctx: &mut RunContext,
121        phase: Phase,
122    ) -> Result<Vec<NodeResult>> {
123        plan.validate()?;
124        let mut results = Vec::new();
125        let fold_ids = if phase == Phase::FitCv {
126            plan.fold_set
127                .as_ref()
128                .map(|fold_set| {
129                    fold_set
130                        .folds
131                        .iter()
132                        .map(|fold| Some(fold.fold_id.clone()))
133                        .collect::<Vec<_>>()
134                })
135                .unwrap_or_else(|| vec![None])
136        } else {
137            vec![None]
138        };
139        for variant in &plan.variants {
140            if ctx
141                .variant_id
142                .as_ref()
143                .is_some_and(|requested| requested != &variant.variant_id)
144            {
145                continue;
146            }
147            for fold_id in &fold_ids {
148                let seed_root = variant.seed.or(ctx.root_seed);
149                results.extend(self.execute_phase_scope(
150                    plan,
151                    controllers,
152                    ctx,
153                    PhaseScope {
154                        phase,
155                        variant_id: Some(variant.variant_id.clone()),
156                        variant: Some(VariantExecutionSpec::from_plan(variant)),
157                        fold_id: fold_id.clone(),
158                        seed_root,
159                    },
160                    PhaseScopeResources::default(),
161                )?);
162            }
163        }
164        Ok(results)
165    }
166
167    pub fn execute_campaign_phase_with_data_provider(
168        &self,
169        plan: &ExecutionPlan,
170        controllers: &RuntimeControllerRegistry,
171        data_provider: &dyn RuntimeDataProvider,
172        ctx: &mut RunContext,
173        phase: Phase,
174    ) -> Result<Vec<NodeResult>> {
175        plan.validate()?;
176        let mut results = Vec::new();
177        let fold_ids = if phase == Phase::FitCv {
178            plan.fold_set
179                .as_ref()
180                .map(|fold_set| {
181                    fold_set
182                        .folds
183                        .iter()
184                        .map(|fold| Some(fold.fold_id.clone()))
185                        .collect::<Vec<_>>()
186                })
187                .unwrap_or_else(|| vec![None])
188        } else {
189            vec![None]
190        };
191        for variant in &plan.variants {
192            if ctx
193                .variant_id
194                .as_ref()
195                .is_some_and(|requested| requested != &variant.variant_id)
196            {
197                continue;
198            }
199            for fold_id in &fold_ids {
200                let seed_root = variant.seed.or(ctx.root_seed);
201                results.extend(self.execute_phase_scope(
202                    plan,
203                    controllers,
204                    ctx,
205                    PhaseScope {
206                        phase,
207                        variant_id: Some(variant.variant_id.clone()),
208                        variant: Some(VariantExecutionSpec::from_plan(variant)),
209                        fold_id: fold_id.clone(),
210                        seed_root,
211                    },
212                    PhaseScopeResources {
213                        data_provider: Some(data_provider),
214                        ..Default::default()
215                    },
216                )?);
217            }
218        }
219        Ok(results)
220    }
221
222    pub fn execute_campaign_phase_with_data_provider_and_artifact_store(
223        &self,
224        plan: &ExecutionPlan,
225        controllers: &RuntimeControllerRegistry,
226        data_provider: &dyn RuntimeDataProvider,
227        artifact_store: &mut InMemoryArtifactStore,
228        ctx: &mut RunContext,
229        phase: Phase,
230    ) -> Result<Vec<NodeResult>> {
231        plan.validate()?;
232        let mut results = Vec::new();
233        let fold_ids = if phase == Phase::FitCv {
234            plan.fold_set
235                .as_ref()
236                .map(|fold_set| {
237                    fold_set
238                        .folds
239                        .iter()
240                        .map(|fold| Some(fold.fold_id.clone()))
241                        .collect::<Vec<_>>()
242                })
243                .unwrap_or_else(|| vec![None])
244        } else {
245            vec![None]
246        };
247        for variant in &plan.variants {
248            if ctx
249                .variant_id
250                .as_ref()
251                .is_some_and(|requested| requested != &variant.variant_id)
252            {
253                continue;
254            }
255            for fold_id in &fold_ids {
256                let seed_root = variant.seed.or(ctx.root_seed);
257                results.extend(self.execute_phase_scope(
258                    plan,
259                    controllers,
260                    ctx,
261                    PhaseScope {
262                        phase,
263                        variant_id: Some(variant.variant_id.clone()),
264                        variant: Some(VariantExecutionSpec::from_plan(variant)),
265                        fold_id: fold_id.clone(),
266                        seed_root,
267                    },
268                    PhaseScopeResources {
269                        data_provider: Some(data_provider),
270                        artifact_store: Some(&mut *artifact_store),
271                        ..Default::default()
272                    },
273                )?);
274            }
275        }
276        Ok(results)
277    }
278
279    pub fn execute_bundle_replay(
280        &self,
281        replay: BundleReplayExecution<'_>,
282        ctx: &mut RunContext,
283    ) -> Result<Vec<NodeResult>> {
284        replay.bundle.validate_against_plan(replay.plan)?;
285        replay
286            .replay_request
287            .validate_for_bundle_with_prediction_cache_store(
288                replay.bundle,
289                replay.prediction_cache_store.is_some(),
290            )?;
291        replay
292            .bundle
293            .validate_replay_envelopes(replay.data_envelopes)?;
294        let prediction_cache_contracts = if replay.replay_request.phase == Phase::Refit {
295            Some(replay_prediction_cache_contracts(replay.bundle)?)
296        } else {
297            None
298        };
299        if replay.replay_request.phase == Phase::Refit {
300            preload_replay_prediction_cache_store(
301                replay.bundle,
302                replay.prediction_cache_store,
303                ctx,
304            )?;
305        }
306        let replay_artifacts = materialize_replay_artifact_handles(
307            replay.plan,
308            replay.bundle,
309            replay.replay_request,
310            replay.artifact_store,
311            ctx,
312        )?;
313        let selected_variant = replay
314            .bundle
315            .selected_variant_id
316            .as_ref()
317            .map(|selected| {
318                replay
319                    .plan
320                    .variants
321                    .iter()
322                    .find(|variant| &variant.variant_id == selected)
323                    .map(VariantExecutionSpec::from_plan)
324                    .ok_or_else(|| {
325                        DagMlError::RuntimeValidation(format!(
326                            "bundle `{}` selected unknown variant `{selected}`",
327                            replay.bundle.bundle_id
328                        ))
329                    })
330            })
331            .transpose()?;
332        let seed_root = selected_variant
333            .as_ref()
334            .and_then(|variant| variant.seed)
335            .or(ctx.root_seed);
336
337        self.execute_phase_scope(
338            replay.plan,
339            replay.controllers,
340            ctx,
341            PhaseScope {
342                phase: replay.replay_request.phase,
343                variant_id: replay.bundle.selected_variant_id.clone(),
344                variant: selected_variant,
345                fold_id: None,
346                seed_root,
347            },
348            PhaseScopeResources {
349                data_provider: Some(replay.data_provider),
350                replay_artifact_handles: Some(&replay_artifacts.handles),
351                replay_artifact_inputs: Some(&replay_artifacts.inputs),
352                replay_bundle_id: Some(&replay.bundle.bundle_id),
353                data_envelopes: Some(replay.data_envelopes),
354                prediction_cache_store: replay.prediction_cache_store,
355                prediction_cache_contracts: prediction_cache_contracts.as_ref(),
356                ..Default::default()
357            },
358        )
359    }
360
361    fn execute_phase_scope(
362        &self,
363        plan: &ExecutionPlan,
364        controllers: &RuntimeControllerRegistry,
365        ctx: &mut RunContext,
366        scope: PhaseScope,
367        mut resources: PhaseScopeResources<'_>,
368    ) -> Result<Vec<NodeResult>> {
369        let _phase_span = crate::observability::phase_span(
370            ctx.run_id.as_str(),
371            plan.id.as_str(),
372            scope.phase.as_str(),
373            scope.variant_id.as_ref().map(VariantId::as_str),
374            scope.fold_id.as_ref().map(FoldId::as_str),
375        )
376        .entered();
377        let mut results = Vec::new();
378        let mut output_handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
379        let mut output_data_views =
380            BTreeMap::<NodeId, BTreeMap<String, DataProviderViewSpec>>::new();
381        let mut input_lineage = BTreeMap::<NodeId, LineageId>::new();
382
383        for level in plan.node_parallel_levels_for_phase(scope.phase)? {
384            for node_id in &level {
385                let node_plan = plan
386                    .node_plans
387                    .get(node_id)
388                    .expect("execution plan was validated");
389                // Cross-branch merge reassembly (concat or late-fusion) is a
390                // scheduler/runtime handler, not a controller call: it reads the
391                // upstream branch OOF blocks from the prediction store and emits
392                // one merged per-sample OOF block. Intercept it before the
393                // controller path (and before the `requires_oof` edge collection,
394                // which is a stacking contract the branch inputs do not satisfy).
395                if let Some(reduction) = merge_reduction_mode(plan, node_plan) {
396                    if let Some(result) =
397                        reassemble_branch_merge(plan, node_plan, ctx, &scope, reduction)?
398                    {
399                        for prediction in &result.predictions {
400                            ctx.prediction_store.append(prediction.clone())?;
401                        }
402                        apply_result_scoring(
403                            &result,
404                            &mut ctx.score_collector,
405                            &mut ctx.regression_target_records,
406                        )?;
407                        ctx.lineage.record(result.lineage.clone())?;
408                        output_handles.insert(node_id.clone(), result.outputs.clone());
409                        input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
410                        results.push(result);
411                    }
412                    continue;
413                }
414                let controller = controllers.get(&node_plan.controller_id).ok_or_else(|| {
415                    DagMlError::RuntimeValidation(format!(
416                        "runtime controller `{}` is not registered",
417                        node_plan.controller_id
418                    ))
419                })?;
420                let collected_inputs = collect_input_handles(
421                    plan,
422                    node_plan,
423                    &output_handles,
424                    &output_data_views,
425                    &resources,
426                    ctx,
427                    &scope,
428                )?;
429                if collected_inputs.skip_node {
430                    continue;
431                }
432                let mut input_handles = collected_inputs.handles;
433                let mut artifact_inputs = BTreeMap::new();
434                if let Some(node_artifact_handles) = resources
435                    .replay_artifact_handles
436                    .and_then(|handles| handles.get(node_id))
437                {
438                    for (key, handle) in node_artifact_handles {
439                        if input_handles.insert(key.clone(), handle.clone()).is_some() {
440                            return Err(DagMlError::RuntimeValidation(format!(
441                                "node `{node_id}` received duplicate replay artifact input `{key}`"
442                            )));
443                        }
444                    }
445                }
446                if let Some(node_artifact_inputs) = resources
447                    .replay_artifact_inputs
448                    .and_then(|inputs| inputs.get(node_id))
449                {
450                    for (key, spec) in node_artifact_inputs {
451                        if artifact_inputs.insert(key.clone(), spec.clone()).is_some() {
452                            return Err(DagMlError::RuntimeValidation(format!(
453                                "node `{node_id}` received duplicate replay artifact metadata `{key}`"
454                            )));
455                        }
456                    }
457                }
458                let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
459                let inner_fold_set = inner_fold_set_for_scope(
460                    &plan.campaign,
461                    plan.fold_set.as_ref(),
462                    node_plan,
463                    &scope,
464                )?;
465                let fit_influence = fit_influence_task_for_node(
466                    plan,
467                    &task_node_plan,
468                    &collected_inputs.data_views,
469                )?;
470                let task = NodeTask {
471                    inner_fold_set,
472                    run_id: ctx.run_id.clone(),
473                    node_plan: task_node_plan.clone(),
474                    phase: scope.phase,
475                    variant_id: scope.variant_id.clone(),
476                    variant: scope.variant.clone(),
477                    fold_id: scope.fold_id.clone(),
478                    branch_path: Vec::new(),
479                    input_handles,
480                    data_views: collected_inputs.data_views,
481                    prediction_inputs: collected_inputs.prediction_inputs,
482                    artifact_inputs,
483                    fit_influence,
484                    seed: derive_task_seed(
485                        scope.seed_root,
486                        scope.variant_id.as_ref(),
487                        scope.fold_id.as_ref(),
488                        &task_node_plan,
489                        scope.phase,
490                    ),
491                };
492                let _node_span = crate::observability::node_span(
493                    task.run_id.as_str(),
494                    plan.id.as_str(),
495                    task.phase.as_str(),
496                    task.node_plan.node_id.as_str(),
497                    task.node_plan.controller_id.as_str(),
498                )
499                .entered();
500                let mut result = controller.invoke(&task)?;
501                record_fit_influence_diagnostic(&task, &mut result);
502                result.validate_for_task(&task)?;
503                apply_result_prediction_aggregation(
504                    plan,
505                    controllers,
506                    &task,
507                    &mut result,
508                    &resources,
509                )?;
510                attach_coordinator_input_lineage(
511                    &mut result,
512                    plan,
513                    &task.node_plan.node_id,
514                    &input_lineage,
515                )?;
516                if let Some(store) = resources.artifact_store.as_deref_mut() {
517                    if scope.phase == Phase::Refit {
518                        store.capture_refit_artifacts(&task, &result)?;
519                    }
520                }
521                for prediction in &result.predictions {
522                    ctx.prediction_store.append(prediction.clone())?;
523                }
524                for prediction in &result.aggregated_predictions {
525                    ctx.aggregated_prediction_store.append(prediction.clone())?;
526                }
527                apply_result_scoring(
528                    &result,
529                    &mut ctx.score_collector,
530                    &mut ctx.regression_target_records,
531                )?;
532                ctx.lineage.record(result.lineage.clone())?;
533                let data_views = derive_output_data_views(plan, &task, &result)?;
534                output_handles.insert(node_id.clone(), result.outputs.clone());
535                output_data_views.insert(node_id.clone(), data_views);
536                input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
537                results.push(result);
538            }
539        }
540
541        Ok(results)
542    }
543}
544
545impl ParallelScheduler {
546    pub fn execute_phase(
547        &self,
548        plan: &ExecutionPlan,
549        controllers: &RuntimeControllerRegistry,
550        ctx: &mut RunContext,
551        phase: Phase,
552    ) -> Result<Vec<NodeResult>> {
553        plan.validate()?;
554        let variant_id = ctx.variant_id.clone();
555        let seed_root = ctx.root_seed;
556        self.execute_phase_scope(
557            plan,
558            controllers,
559            ctx,
560            PhaseScope {
561                phase,
562                variant_id,
563                variant: None,
564                fold_id: None,
565                seed_root,
566            },
567            PhaseScopeResources::default(),
568        )
569    }
570
571    pub fn execute_phase_with_data_provider(
572        &self,
573        plan: &ExecutionPlan,
574        controllers: &RuntimeControllerRegistry,
575        data_provider: &dyn RuntimeDataProvider,
576        ctx: &mut RunContext,
577        phase: Phase,
578    ) -> Result<Vec<NodeResult>> {
579        plan.validate()?;
580        let variant_id = ctx.variant_id.clone();
581        let seed_root = ctx.root_seed;
582        self.execute_phase_scope(
583            plan,
584            controllers,
585            ctx,
586            PhaseScope {
587                phase,
588                variant_id,
589                variant: None,
590                fold_id: None,
591                seed_root,
592            },
593            PhaseScopeResources {
594                data_provider: Some(data_provider),
595                ..Default::default()
596            },
597        )
598    }
599
600    pub fn execute_campaign_phase(
601        &self,
602        plan: &ExecutionPlan,
603        controllers: &RuntimeControllerRegistry,
604        ctx: &mut RunContext,
605        phase: Phase,
606    ) -> Result<Vec<NodeResult>> {
607        plan.validate()?;
608        let mut results = Vec::new();
609        let fold_ids = if phase == Phase::FitCv {
610            plan.fold_set
611                .as_ref()
612                .map(|fold_set| {
613                    fold_set
614                        .folds
615                        .iter()
616                        .map(|fold| Some(fold.fold_id.clone()))
617                        .collect::<Vec<_>>()
618                })
619                .unwrap_or_else(|| vec![None])
620        } else {
621            vec![None]
622        };
623        for variant in &plan.variants {
624            if ctx
625                .variant_id
626                .as_ref()
627                .is_some_and(|requested| requested != &variant.variant_id)
628            {
629                continue;
630            }
631            for fold_id in &fold_ids {
632                let seed_root = variant.seed.or(ctx.root_seed);
633                results.extend(self.execute_phase_scope(
634                    plan,
635                    controllers,
636                    ctx,
637                    PhaseScope {
638                        phase,
639                        variant_id: Some(variant.variant_id.clone()),
640                        variant: Some(VariantExecutionSpec::from_plan(variant)),
641                        fold_id: fold_id.clone(),
642                        seed_root,
643                    },
644                    PhaseScopeResources::default(),
645                )?);
646            }
647        }
648        Ok(results)
649    }
650
651    pub fn execute_campaign_phase_with_data_provider(
652        &self,
653        plan: &ExecutionPlan,
654        controllers: &RuntimeControllerRegistry,
655        data_provider: &dyn RuntimeDataProvider,
656        ctx: &mut RunContext,
657        phase: Phase,
658    ) -> Result<Vec<NodeResult>> {
659        plan.validate()?;
660        let mut results = Vec::new();
661        let fold_ids = if phase == Phase::FitCv {
662            plan.fold_set
663                .as_ref()
664                .map(|fold_set| {
665                    fold_set
666                        .folds
667                        .iter()
668                        .map(|fold| Some(fold.fold_id.clone()))
669                        .collect::<Vec<_>>()
670                })
671                .unwrap_or_else(|| vec![None])
672        } else {
673            vec![None]
674        };
675        for variant in &plan.variants {
676            if ctx
677                .variant_id
678                .as_ref()
679                .is_some_and(|requested| requested != &variant.variant_id)
680            {
681                continue;
682            }
683            for fold_id in &fold_ids {
684                let seed_root = variant.seed.or(ctx.root_seed);
685                results.extend(self.execute_phase_scope(
686                    plan,
687                    controllers,
688                    ctx,
689                    PhaseScope {
690                        phase,
691                        variant_id: Some(variant.variant_id.clone()),
692                        variant: Some(VariantExecutionSpec::from_plan(variant)),
693                        fold_id: fold_id.clone(),
694                        seed_root,
695                    },
696                    PhaseScopeResources {
697                        data_provider: Some(data_provider),
698                        ..Default::default()
699                    },
700                )?);
701            }
702        }
703        Ok(results)
704    }
705
706    pub fn execute_campaign_phase_with_data_provider_and_artifact_store(
707        &self,
708        plan: &ExecutionPlan,
709        controllers: &RuntimeControllerRegistry,
710        data_provider: &dyn RuntimeDataProvider,
711        artifact_store: &mut InMemoryArtifactStore,
712        ctx: &mut RunContext,
713        phase: Phase,
714    ) -> Result<Vec<NodeResult>> {
715        plan.validate()?;
716        let mut results = Vec::new();
717        let fold_ids = if phase == Phase::FitCv {
718            plan.fold_set
719                .as_ref()
720                .map(|fold_set| {
721                    fold_set
722                        .folds
723                        .iter()
724                        .map(|fold| Some(fold.fold_id.clone()))
725                        .collect::<Vec<_>>()
726                })
727                .unwrap_or_else(|| vec![None])
728        } else {
729            vec![None]
730        };
731        for variant in &plan.variants {
732            if ctx
733                .variant_id
734                .as_ref()
735                .is_some_and(|requested| requested != &variant.variant_id)
736            {
737                continue;
738            }
739            for fold_id in &fold_ids {
740                let seed_root = variant.seed.or(ctx.root_seed);
741                results.extend(self.execute_phase_scope(
742                    plan,
743                    controllers,
744                    ctx,
745                    PhaseScope {
746                        phase,
747                        variant_id: Some(variant.variant_id.clone()),
748                        variant: Some(VariantExecutionSpec::from_plan(variant)),
749                        fold_id: fold_id.clone(),
750                        seed_root,
751                    },
752                    PhaseScopeResources {
753                        data_provider: Some(data_provider),
754                        artifact_store: Some(&mut *artifact_store),
755                        ..Default::default()
756                    },
757                )?);
758            }
759        }
760        Ok(results)
761    }
762
763    pub fn execute_bundle_replay(
764        &self,
765        replay: BundleReplayExecution<'_>,
766        ctx: &mut RunContext,
767    ) -> Result<Vec<NodeResult>> {
768        replay.bundle.validate_against_plan(replay.plan)?;
769        replay
770            .replay_request
771            .validate_for_bundle_with_prediction_cache_store(
772                replay.bundle,
773                replay.prediction_cache_store.is_some(),
774            )?;
775        replay
776            .bundle
777            .validate_replay_envelopes(replay.data_envelopes)?;
778        let prediction_cache_contracts = if replay.replay_request.phase == Phase::Refit {
779            Some(replay_prediction_cache_contracts(replay.bundle)?)
780        } else {
781            None
782        };
783        if replay.replay_request.phase == Phase::Refit {
784            preload_replay_prediction_cache_store(
785                replay.bundle,
786                replay.prediction_cache_store,
787                ctx,
788            )?;
789        }
790        let replay_artifacts = materialize_replay_artifact_handles(
791            replay.plan,
792            replay.bundle,
793            replay.replay_request,
794            replay.artifact_store,
795            ctx,
796        )?;
797        let selected_variant = replay
798            .bundle
799            .selected_variant_id
800            .as_ref()
801            .map(|selected| {
802                replay
803                    .plan
804                    .variants
805                    .iter()
806                    .find(|variant| &variant.variant_id == selected)
807                    .map(VariantExecutionSpec::from_plan)
808                    .ok_or_else(|| {
809                        DagMlError::RuntimeValidation(format!(
810                            "bundle `{}` selected unknown variant `{selected}`",
811                            replay.bundle.bundle_id
812                        ))
813                    })
814            })
815            .transpose()?;
816        let seed_root = selected_variant
817            .as_ref()
818            .and_then(|variant| variant.seed)
819            .or(ctx.root_seed);
820
821        self.execute_phase_scope(
822            replay.plan,
823            replay.controllers,
824            ctx,
825            PhaseScope {
826                phase: replay.replay_request.phase,
827                variant_id: replay.bundle.selected_variant_id.clone(),
828                variant: selected_variant,
829                fold_id: None,
830                seed_root,
831            },
832            PhaseScopeResources {
833                data_provider: Some(replay.data_provider),
834                replay_artifact_handles: Some(&replay_artifacts.handles),
835                replay_artifact_inputs: Some(&replay_artifacts.inputs),
836                replay_bundle_id: Some(&replay.bundle.bundle_id),
837                data_envelopes: Some(replay.data_envelopes),
838                prediction_cache_store: replay.prediction_cache_store,
839                prediction_cache_contracts: prediction_cache_contracts.as_ref(),
840                ..Default::default()
841            },
842        )
843    }
844
845    fn execute_phase_scope(
846        &self,
847        plan: &ExecutionPlan,
848        controllers: &RuntimeControllerRegistry,
849        ctx: &mut RunContext,
850        scope: PhaseScope,
851        mut resources: PhaseScopeResources<'_>,
852    ) -> Result<Vec<NodeResult>> {
853        // Hold the phase span on the scheduler thread, and clone it into each
854        // worker so worker-thread telemetry nests under the phase (tracing spans
855        // are thread-local and do not auto-propagate across `thread::scope`).
856        let phase_span = crate::observability::phase_span(
857            ctx.run_id.as_str(),
858            plan.id.as_str(),
859            scope.phase.as_str(),
860            scope.variant_id.as_ref().map(VariantId::as_str),
861            scope.fold_id.as_ref().map(FoldId::as_str),
862        );
863        let _phase_entered = phase_span.clone().entered();
864        // Borrowed for the `thread::scope` below; workers join before it ends.
865        let plan_id = plan.id.as_str();
866        plan.validate_parallel_controller_capabilities(self.max_workers, scope.phase)?;
867        let mut results = Vec::new();
868        let mut output_handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
869        let mut output_data_views =
870            BTreeMap::<NodeId, BTreeMap<String, DataProviderViewSpec>>::new();
871        let mut input_lineage = BTreeMap::<NodeId, LineageId>::new();
872
873        for level in plan.node_parallel_levels_for_phase(scope.phase)? {
874            let mut prepared = Vec::<PreparedNodeTask>::new();
875            // Cross-branch merge nodes (concat or late-fusion) are not controller
876            // tasks: they read the upstream branch OOF blocks from the prediction
877            // store and reassemble them on the scheduler thread (no worker), AFTER
878            // this level's worker tasks have populated the store. They are in a
879            // later level than their branches, so the store already holds the
880            // branch OOF by the time we reassemble — see `reassemble_branch_merge`.
881            let mut merge_nodes = Vec::<(NodeId, MergeReduction)>::new();
882            for node_id in &level {
883                let node_plan = plan
884                    .node_plans
885                    .get(node_id)
886                    .expect("execution plan was validated");
887                if let Some(reduction) = merge_reduction_mode(plan, node_plan) {
888                    merge_nodes.push((node_id.clone(), reduction));
889                    continue;
890                }
891                let collected_inputs = collect_input_handles(
892                    plan,
893                    node_plan,
894                    &output_handles,
895                    &output_data_views,
896                    &resources,
897                    ctx,
898                    &scope,
899                )?;
900                if collected_inputs.skip_node {
901                    continue;
902                }
903                let mut input_handles = collected_inputs.handles;
904                let mut artifact_inputs = BTreeMap::new();
905                if let Some(node_artifact_handles) = resources
906                    .replay_artifact_handles
907                    .and_then(|handles| handles.get(node_id))
908                {
909                    for (key, handle) in node_artifact_handles {
910                        if input_handles.insert(key.clone(), handle.clone()).is_some() {
911                            return Err(DagMlError::RuntimeValidation(format!(
912                                "node `{node_id}` received duplicate replay artifact input `{key}`"
913                            )));
914                        }
915                    }
916                }
917                if let Some(node_artifact_inputs) = resources
918                    .replay_artifact_inputs
919                    .and_then(|inputs| inputs.get(node_id))
920                {
921                    for (key, spec) in node_artifact_inputs {
922                        if artifact_inputs.insert(key.clone(), spec.clone()).is_some() {
923                            return Err(DagMlError::RuntimeValidation(format!(
924                                "node `{node_id}` received duplicate replay artifact metadata `{key}`"
925                            )));
926                        }
927                    }
928                }
929                let task_node_plan = effective_node_plan_for_scope(node_plan, &scope)?;
930                let inner_fold_set = inner_fold_set_for_scope(
931                    &plan.campaign,
932                    plan.fold_set.as_ref(),
933                    node_plan,
934                    &scope,
935                )?;
936                let fit_influence = fit_influence_task_for_node(
937                    plan,
938                    &task_node_plan,
939                    &collected_inputs.data_views,
940                )?;
941                prepared.push(PreparedNodeTask {
942                    node_id: node_id.clone(),
943                    task: NodeTask {
944                        inner_fold_set,
945                        run_id: ctx.run_id.clone(),
946                        node_plan: task_node_plan.clone(),
947                        phase: scope.phase,
948                        variant_id: scope.variant_id.clone(),
949                        variant: scope.variant.clone(),
950                        fold_id: scope.fold_id.clone(),
951                        branch_path: Vec::new(),
952                        input_handles,
953                        data_views: collected_inputs.data_views,
954                        prediction_inputs: collected_inputs.prediction_inputs,
955                        artifact_inputs,
956                        fit_influence,
957                        seed: derive_task_seed(
958                            scope.seed_root,
959                            scope.variant_id.as_ref(),
960                            scope.fold_id.as_ref(),
961                            &task_node_plan,
962                            scope.phase,
963                        ),
964                    },
965                });
966            }
967
968            for chunk in prepared.chunks(self.max_workers) {
969                let chunk_results =
970                    std::thread::scope(|thread_scope| -> Result<Vec<NodeResult>> {
971                        let mut handles = Vec::with_capacity(chunk.len());
972                        for prepared_task in chunk {
973                            let controller = controllers
974                                .get(&prepared_task.task.node_plan.controller_id)
975                                .ok_or_else(|| {
976                                    DagMlError::RuntimeValidation(format!(
977                                        "runtime controller `{}` is not registered",
978                                        prepared_task.task.node_plan.controller_id
979                                    ))
980                                })?;
981                            let worker_span = phase_span.clone();
982                            handles.push(thread_scope.spawn(move || {
983                                let _worker_span = worker_span.entered();
984                                let _node_span = crate::observability::node_span(
985                                    prepared_task.task.run_id.as_str(),
986                                    plan_id,
987                                    prepared_task.task.phase.as_str(),
988                                    prepared_task.task.node_plan.node_id.as_str(),
989                                    prepared_task.task.node_plan.controller_id.as_str(),
990                                )
991                                .entered();
992                                let mut result = controller.invoke(&prepared_task.task)?;
993                                record_fit_influence_diagnostic(&prepared_task.task, &mut result);
994                                result.validate_for_task(&prepared_task.task)?;
995                                Ok(result)
996                            }));
997                        }
998                        handles
999                            .into_iter()
1000                            .map(|handle| {
1001                                handle.join().map_err(|_| {
1002                                    DagMlError::RuntimeValidation(
1003                                        "parallel scheduler worker panicked".to_string(),
1004                                    )
1005                                })?
1006                            })
1007                            .collect()
1008                    })?;
1009
1010                for (prepared_task, mut result) in chunk.iter().zip(chunk_results) {
1011                    apply_result_prediction_aggregation(
1012                        plan,
1013                        controllers,
1014                        &prepared_task.task,
1015                        &mut result,
1016                        &resources,
1017                    )?;
1018                    attach_coordinator_input_lineage(
1019                        &mut result,
1020                        plan,
1021                        &prepared_task.task.node_plan.node_id,
1022                        &input_lineage,
1023                    )?;
1024                    if let Some(store) = resources.artifact_store.as_deref_mut() {
1025                        if scope.phase == Phase::Refit {
1026                            store.capture_refit_artifacts(&prepared_task.task, &result)?;
1027                        }
1028                    }
1029                    for prediction in &result.predictions {
1030                        ctx.prediction_store.append(prediction.clone())?;
1031                    }
1032                    for prediction in &result.aggregated_predictions {
1033                        ctx.aggregated_prediction_store.append(prediction.clone())?;
1034                    }
1035                    apply_result_scoring(
1036                        &result,
1037                        &mut ctx.score_collector,
1038                        &mut ctx.regression_target_records,
1039                    )?;
1040                    ctx.lineage.record(result.lineage.clone())?;
1041                    let data_views = derive_output_data_views(plan, &prepared_task.task, &result)?;
1042                    output_handles.insert(prepared_task.node_id.clone(), result.outputs.clone());
1043                    output_data_views.insert(prepared_task.node_id.clone(), data_views);
1044                    input_lineage.insert(
1045                        prepared_task.node_id.clone(),
1046                        result.lineage.record_id.clone(),
1047                    );
1048                    results.push(result);
1049                }
1050            }
1051
1052            // Reassemble any cross-branch merge nodes in this level now that the
1053            // level's worker tasks have populated the prediction store. Merge nodes
1054            // sit in a later level than the branches they consume, so the upstream
1055            // branch OOF is already present.
1056            for (node_id, reduction) in &merge_nodes {
1057                let node_plan = plan
1058                    .node_plans
1059                    .get(node_id)
1060                    .expect("execution plan was validated");
1061                if let Some(result) =
1062                    reassemble_branch_merge(plan, node_plan, ctx, &scope, *reduction)?
1063                {
1064                    for prediction in &result.predictions {
1065                        ctx.prediction_store.append(prediction.clone())?;
1066                    }
1067                    apply_result_scoring(
1068                        &result,
1069                        &mut ctx.score_collector,
1070                        &mut ctx.regression_target_records,
1071                    )?;
1072                    ctx.lineage.record(result.lineage.clone())?;
1073                    output_handles.insert(node_id.clone(), result.outputs.clone());
1074                    input_lineage.insert(node_id.clone(), result.lineage.record_id.clone());
1075                    results.push(result);
1076                }
1077            }
1078        }
1079
1080        Ok(results)
1081    }
1082}
1083
1084pub(crate) struct PreparedNodeTask {
1085    pub(crate) node_id: NodeId,
1086    pub(crate) task: NodeTask,
1087}
1088
1089pub(crate) fn attach_coordinator_input_lineage(
1090    result: &mut NodeResult,
1091    plan: &ExecutionPlan,
1092    node_id: &NodeId,
1093    upstream_lineage: &BTreeMap<NodeId, LineageId>,
1094) -> Result<()> {
1095    let inferred = inferred_input_lineage_for_node(plan, node_id, upstream_lineage);
1096    if result.lineage.input_lineage.is_empty() {
1097        result.lineage.input_lineage = inferred;
1098        return Ok(());
1099    }
1100
1101    let declared = result
1102        .lineage
1103        .input_lineage
1104        .iter()
1105        .cloned()
1106        .collect::<BTreeSet<_>>()
1107        .into_iter()
1108        .collect::<Vec<_>>();
1109    if declared != inferred {
1110        return Err(DagMlError::RuntimeValidation(format!(
1111            "lineage for node `{}` declared input lineage {:?}, expected {:?}",
1112            result.node_id, declared, inferred
1113        )));
1114    }
1115    result.lineage.input_lineage = declared;
1116    Ok(())
1117}
1118
1119pub(crate) fn inferred_input_lineage_for_node(
1120    plan: &ExecutionPlan,
1121    node_id: &NodeId,
1122    upstream_lineage: &BTreeMap<NodeId, LineageId>,
1123) -> Vec<LineageId> {
1124    plan.graph_plan
1125        .graph
1126        .edges
1127        .iter()
1128        .filter(|edge| &edge.target.node_id == node_id && edge.contract.propagates_lineage)
1129        .filter_map(|edge| upstream_lineage.get(&edge.source.node_id).cloned())
1130        .collect::<BTreeSet<_>>()
1131        .into_iter()
1132        .collect()
1133}
1134pub(crate) fn collect_input_handles(
1135    plan: &ExecutionPlan,
1136    node_plan: &NodePlan,
1137    output_handles: &BTreeMap<NodeId, BTreeMap<String, HandleRef>>,
1138    output_data_views: &BTreeMap<NodeId, BTreeMap<String, DataProviderViewSpec>>,
1139    resources: &PhaseScopeResources<'_>,
1140    ctx: &RunContext,
1141    scope: &PhaseScope,
1142) -> Result<CollectedInputs> {
1143    let mut inputs = BTreeMap::new();
1144    let mut data_views = BTreeMap::new();
1145    let mut prediction_inputs = BTreeMap::new();
1146    let training_oof_edges = incoming_training_oof_edges(plan, node_plan, scope)?;
1147    let training_oof_sources = training_oof_edges
1148        .iter()
1149        .map(|edge| edge.source.node_id.clone())
1150        .collect::<BTreeSet<_>>();
1151    let bound_data_inputs = node_plan
1152        .data_bindings
1153        .iter()
1154        .map(|binding| binding.input_name.clone())
1155        .collect::<BTreeSet<_>>();
1156    // Only forward upstream handles for ports this node DECLARES an edge to.
1157    // A controller must never see a handle outside its declared port contract,
1158    // so a sibling consumer of the same producer cannot expose extra ports here.
1159    let declared_source_ports = plan
1160        .graph_plan
1161        .graph
1162        .edges
1163        .iter()
1164        .filter(|edge| edge.target.node_id == node_plan.node_id)
1165        .map(|edge| (edge.source.node_id.clone(), edge.source.port_name.clone()))
1166        .collect::<BTreeSet<_>>();
1167    for upstream in &node_plan.input_nodes {
1168        if training_oof_sources.contains(upstream) {
1169            continue;
1170        }
1171        if let Some(handles) = output_handles.get(upstream) {
1172            for (port, handle) in handles {
1173                if !declared_source_ports.contains(&(upstream.clone(), port.clone())) {
1174                    continue;
1175                }
1176                inputs.insert(format!("{upstream}.{port}"), handle.clone());
1177            }
1178        }
1179    }
1180    for edge in plan
1181        .graph_plan
1182        .graph
1183        .edges
1184        .iter()
1185        .filter(|edge| edge.target.node_id == node_plan.node_id)
1186        .filter(|edge| edge.contract.kind == PortKind::Data && !edge.contract.requires_oof)
1187    {
1188        if bound_data_inputs.contains(&edge.target.port_name) {
1189            continue;
1190        }
1191        let Some(handles) = output_handles.get(&edge.source.node_id) else {
1192            continue;
1193        };
1194        let Some(handle) = handles.get(&edge.source.port_name) else {
1195            continue;
1196        };
1197        let key = data_view_key(&edge.target.port_name);
1198        if inputs.insert(key.clone(), handle.clone()).is_some() {
1199            return Err(DagMlError::RuntimeValidation(format!(
1200                "node `{}` received duplicate data edge input `{key}`",
1201                node_plan.node_id
1202            )));
1203        }
1204        if let Some(source_views) = output_data_views.get(&edge.source.node_id) {
1205            if let Some(view) = source_views.get(&edge.source.port_name) {
1206                if data_views.insert(key.clone(), view.clone()).is_some() {
1207                    return Err(DagMlError::RuntimeValidation(format!(
1208                        "node `{}` received duplicate data edge view `{key}`",
1209                        node_plan.node_id
1210                    )));
1211                }
1212            }
1213            let source_validation_key = validation_data_view_key(&edge.source.port_name);
1214            if let Some(view) = source_views.get(&source_validation_key) {
1215                let validation_key = format!("{key}:validation");
1216                if data_views
1217                    .insert(validation_key.clone(), view.clone())
1218                    .is_some()
1219                {
1220                    return Err(DagMlError::RuntimeValidation(format!(
1221                        "node `{}` received duplicate data edge validation view `{validation_key}`",
1222                        node_plan.node_id
1223                    )));
1224                }
1225            }
1226        }
1227    }
1228    for edge in training_oof_edges {
1229        let key = format!("{}.{}", edge.source.node_id, edge.source.port_name);
1230        let Some(input) = collect_oof_prediction_input(plan, edge, ctx, scope, resources)? else {
1231            return Ok(CollectedInputs {
1232                handles: BTreeMap::new(),
1233                data_views: BTreeMap::new(),
1234                prediction_inputs: BTreeMap::new(),
1235                skip_node: true,
1236            });
1237        };
1238        if inputs.insert(key.clone(), input.handle).is_some() {
1239            return Err(DagMlError::RuntimeValidation(format!(
1240                "node `{}` received duplicate OOF prediction input `{key}`",
1241                node_plan.node_id
1242            )));
1243        }
1244        if prediction_inputs.insert(key.clone(), input.spec).is_some() {
1245            return Err(DagMlError::RuntimeValidation(format!(
1246                "node `{}` received duplicate OOF prediction spec `{key}`",
1247                node_plan.node_id
1248            )));
1249        }
1250    }
1251    // REFIT / PREDICT: deliver each base producer's off-fold (test / predict)
1252    // predictions to the stacking meta-node as a SEPARATE prediction input (suffixed
1253    // `:test` / `:predict`) so the host meta-model predicts from them. The FIT_CV
1254    // Validation-OOF input above is the meta-features the meta-model trains on; this
1255    // off-fold input is used ONLY for REFIT/PREDICT scoring/prediction, never FIT_CV
1256    // training — keeping the leakage invariant intact.
1257    if matches!(scope.phase, Phase::Refit | Phase::Predict) {
1258        let off_fold_suffix = scope.phase.as_str().to_ascii_lowercase();
1259        for edge in incoming_oof_edges(plan, node_plan)? {
1260            let Some(input) = collect_off_fold_prediction_input(plan, edge, ctx, scope)? else {
1261                continue;
1262            };
1263            let key = format!(
1264                "{}.{}:{off_fold_suffix}",
1265                edge.source.node_id, edge.source.port_name
1266            );
1267            if inputs.insert(key.clone(), input.handle).is_some() {
1268                return Err(DagMlError::RuntimeValidation(format!(
1269                    "node `{}` received duplicate off-fold prediction input `{key}`",
1270                    node_plan.node_id
1271                )));
1272            }
1273            if prediction_inputs.insert(key.clone(), input.spec).is_some() {
1274                return Err(DagMlError::RuntimeValidation(format!(
1275                    "node `{}` received duplicate off-fold prediction spec `{key}`",
1276                    node_plan.node_id
1277                )));
1278            }
1279        }
1280    }
1281    if !node_plan.data_bindings.is_empty() && resources.data_provider.is_none() {
1282        return Err(DagMlError::RuntimeValidation(format!(
1283            "node `{}` requires {} data binding(s) but no runtime data provider is registered",
1284            node_plan.node_id,
1285            node_plan.data_bindings.len()
1286        )));
1287    }
1288    if let Some(data_provider) = resources.data_provider {
1289        // Samples excluded from training (sample-local) for this node, derived
1290        // from its coordinator relations. Used to filter FIT view specs so the
1291        // spec, the materialized view, and fit-influence row_weights agree.
1292        let excluded_samples = coordinator_relations_for_node(node_plan, resources)?
1293            .map(|relations| relations.excluded_sample_ids())
1294            .unwrap_or_default();
1295        for binding in &node_plan.data_bindings {
1296            let materialized = data_provider.materialize(&DataMaterializationRequest {
1297                run_id: ctx.run_id.clone(),
1298                node_id: node_plan.node_id.clone(),
1299                input_name: binding.input_name.clone(),
1300                phase: scope.phase,
1301                variant_id: scope.variant_id.clone(),
1302                fold_id: scope.fold_id.clone(),
1303                binding: binding.clone(),
1304            })?;
1305            let branch_view_for_node = branch_view_from_node_metadata(plan, &node_plan.node_id)?;
1306            let view = data_view_for_scope(
1307                binding,
1308                plan.fold_set.as_ref(),
1309                scope,
1310                branch_view_for_node.as_ref(),
1311                &excluded_samples,
1312            )?;
1313            let key = data_view_key(&binding.input_name);
1314            let view_handle = make_data_view_handle(
1315                data_provider,
1316                ctx,
1317                node_plan,
1318                scope,
1319                binding,
1320                &materialized,
1321                &view,
1322            )?;
1323            if data_views.insert(key.clone(), view).is_some() {
1324                return Err(DagMlError::RuntimeValidation(format!(
1325                    "node `{}` received duplicate data view `{key}`",
1326                    node_plan.node_id
1327                )));
1328            }
1329            if inputs.insert(key.clone(), view_handle).is_some() {
1330                return Err(DagMlError::RuntimeValidation(format!(
1331                    "node `{}` received duplicate data input `{key}`",
1332                    node_plan.node_id
1333                )));
1334            }
1335
1336            if let Some(validation_view) = validation_data_view_for_scope(
1337                binding,
1338                plan.fold_set.as_ref(),
1339                scope,
1340                branch_view_for_node.as_ref(),
1341                &excluded_samples,
1342            )? {
1343                let validation_key = format!("{key}:validation");
1344                let validation_handle = make_data_view_handle(
1345                    data_provider,
1346                    ctx,
1347                    node_plan,
1348                    scope,
1349                    binding,
1350                    &materialized,
1351                    &validation_view,
1352                )?;
1353                if data_views
1354                    .insert(validation_key.clone(), validation_view)
1355                    .is_some()
1356                {
1357                    return Err(DagMlError::RuntimeValidation(format!(
1358                        "node `{}` received duplicate validation data view `{validation_key}`",
1359                        node_plan.node_id
1360                    )));
1361                }
1362                if inputs
1363                    .insert(validation_key.clone(), validation_handle)
1364                    .is_some()
1365                {
1366                    return Err(DagMlError::RuntimeValidation(format!(
1367                        "node `{}` received duplicate validation data input `{validation_key}`",
1368                        node_plan.node_id
1369                    )));
1370                }
1371            }
1372        }
1373    }
1374    Ok(CollectedInputs {
1375        handles: inputs,
1376        data_views,
1377        prediction_inputs,
1378        skip_node: false,
1379    })
1380}
1381pub(crate) fn preload_replay_prediction_cache_store(
1382    bundle: &ExecutionBundle,
1383    prediction_cache_store: Option<&dyn RuntimePredictionCacheStore>,
1384    ctx: &mut RunContext,
1385) -> Result<()> {
1386    if bundle.prediction_requirements.is_empty() {
1387        return Ok(());
1388    }
1389    let store = prediction_cache_store.ok_or_else(|| {
1390        DagMlError::RuntimeValidation(format!(
1391            "bundle `{}` cannot preload OOF prediction caches without a prediction cache store",
1392            bundle.bundle_id
1393        ))
1394    })?;
1395    if !ctx.prediction_store.blocks().is_empty() {
1396        return Err(DagMlError::RuntimeValidation(format!(
1397            "bundle `{}` cannot preload OOF prediction caches into a non-empty prediction store",
1398            bundle.bundle_id
1399        )));
1400    }
1401    let contracts = replay_prediction_cache_contracts(bundle)?;
1402    for contract in contracts.values() {
1403        if contract.requirement.prediction_level == PredictionLevel::Sample {
1404            let blocks = store.load_blocks(&contract.cache.requirement_key)?;
1405            if blocks.iter().any(|block| {
1406                block.producer_node != contract.requirement.producer_node
1407                    || block.partition != contract.requirement.partition
1408            }) {
1409                return Err(DagMlError::RuntimeValidation(format!(
1410                    "prediction cache store returned blocks outside requirement `{}`",
1411                    contract.cache.requirement_key
1412                )));
1413            }
1414            let payload = build_prediction_cache_payload(&contract.requirement, &blocks)?;
1415            validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
1416            for block in &payload.blocks {
1417                ctx.prediction_store.append(block.clone())?;
1418            }
1419        } else {
1420            let blocks = store.load_aggregated_blocks(&contract.cache.requirement_key)?;
1421            if blocks.iter().any(|block| {
1422                block.producer_node != contract.requirement.producer_node
1423                    || block.partition != contract.requirement.partition
1424                    || block.level != contract.requirement.prediction_level
1425            }) {
1426                return Err(DagMlError::RuntimeValidation(format!(
1427                    "prediction cache store returned aggregated blocks outside requirement `{}`",
1428                    contract.cache.requirement_key
1429                )));
1430            }
1431            let payload =
1432                build_aggregated_prediction_cache_payload(&contract.requirement, &blocks)?;
1433            validate_prediction_cache_payload_matches_record(&payload, &contract.cache)?;
1434        }
1435    }
1436    Ok(())
1437}
1438
1439pub(crate) fn replay_prediction_cache_contracts(
1440    bundle: &ExecutionBundle,
1441) -> Result<BTreeMap<String, ReplayPredictionCacheContract>> {
1442    bundle.validate()?;
1443    let requirements = bundle
1444        .prediction_requirements
1445        .iter()
1446        .map(|requirement| (requirement.key(), requirement))
1447        .collect::<BTreeMap<_, _>>();
1448    let mut contracts = BTreeMap::new();
1449    for cache in &bundle.prediction_caches {
1450        let requirement = requirements.get(&cache.requirement_key).ok_or_else(|| {
1451            DagMlError::RuntimeValidation(format!(
1452                "prediction cache `{}` references unknown prediction requirement `{}`",
1453                cache.cache_id, cache.requirement_key
1454            ))
1455        })?;
1456        contracts.insert(
1457            cache.requirement_key.clone(),
1458            ReplayPredictionCacheContract {
1459                requirement: (*requirement).clone(),
1460                cache: cache.clone(),
1461            },
1462        );
1463    }
1464    Ok(contracts)
1465}
1466
1467pub(crate) fn materialize_replay_artifact_handles(
1468    plan: &ExecutionPlan,
1469    bundle: &ExecutionBundle,
1470    replay_request: &ReplayPhaseRequest,
1471    artifact_store: &dyn RuntimeArtifactStore,
1472    ctx: &RunContext,
1473) -> Result<MaterializedReplayArtifacts> {
1474    let mut handles = BTreeMap::<NodeId, BTreeMap<String, HandleRef>>::new();
1475    let mut inputs = BTreeMap::<NodeId, BTreeMap<String, ArtifactInputSpec>>::new();
1476    for artifact in &bundle.refit_artifacts {
1477        artifact.validate()?;
1478        let node_plan = plan.node_plans.get(&artifact.node_id).ok_or_else(|| {
1479            DagMlError::RuntimeValidation(format!(
1480                "bundle `{}` artifact references unknown node `{}`",
1481                bundle.bundle_id, artifact.node_id
1482            ))
1483        })?;
1484        if !node_plan.supported_phases.contains(&replay_request.phase) {
1485            return Err(DagMlError::RuntimeValidation(format!(
1486                "bundle `{}` artifact node `{}` does not support replay phase {:?}",
1487                bundle.bundle_id, artifact.node_id, replay_request.phase
1488            )));
1489        }
1490        let handle = artifact_store.materialize(&ArtifactMaterializationRequest {
1491            run_id: ctx.run_id.clone(),
1492            bundle_id: bundle.bundle_id.clone(),
1493            node_id: artifact.node_id.clone(),
1494            phase: replay_request.phase,
1495            variant_id: bundle.selected_variant_id.clone(),
1496            controller_id: artifact.controller_id.clone(),
1497            artifact: artifact.artifact.clone(),
1498            params_fingerprint: artifact.params_fingerprint.clone(),
1499        })?;
1500        if !matches!(handle.kind, HandleKind::Model | HandleKind::Artifact) {
1501            return Err(DagMlError::RuntimeValidation(format!(
1502                "artifact `{}` materialized as unsupported handle kind {:?}",
1503                artifact.artifact.id, handle.kind
1504            )));
1505        }
1506        if handle.owner_controller != artifact.controller_id {
1507            return Err(DagMlError::RuntimeValidation(format!(
1508                "artifact `{}` handle owner `{}` does not match controller `{}`",
1509                artifact.artifact.id, handle.owner_controller, artifact.controller_id
1510            )));
1511        }
1512        let key = refit_artifact_input_key(&artifact.artifact.id);
1513        if handles
1514            .entry(artifact.node_id.clone())
1515            .or_default()
1516            .insert(key.clone(), handle)
1517            .is_some()
1518        {
1519            return Err(DagMlError::RuntimeValidation(format!(
1520                "duplicate replay artifact input `{key}` for node `{}`",
1521                artifact.node_id
1522            )));
1523        }
1524        if inputs
1525            .entry(artifact.node_id.clone())
1526            .or_default()
1527            .insert(key.clone(), ArtifactInputSpec::from_refit_record(artifact)?)
1528            .is_some()
1529        {
1530            return Err(DagMlError::RuntimeValidation(format!(
1531                "duplicate replay artifact metadata `{key}` for node `{}`",
1532                artifact.node_id
1533            )));
1534        }
1535    }
1536    Ok(MaterializedReplayArtifacts { handles, inputs })
1537}
1538
1539pub(crate) fn derive_task_seed(
1540    root_seed: Option<u64>,
1541    variant_id: Option<&VariantId>,
1542    fold_id: Option<&FoldId>,
1543    node_plan: &NodePlan,
1544    phase: Phase,
1545) -> Option<u64> {
1546    root_seed.map(|root| {
1547        let mut context = SeedContext::root(root);
1548        if let Some(variant_id) = variant_id {
1549            context = context.child(format!("variant:{variant_id}"));
1550        }
1551        if let Some(fold_id) = fold_id {
1552            context = context.child(format!("fold:{fold_id}"));
1553        }
1554        context
1555            .child(format!("node:{}", node_plan.node_id))
1556            .child(format!("phase:{phase:?}"))
1557            .derive_u64("task")
1558    })
1559}