1use 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 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 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 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 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 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 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 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 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}