1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::campaign::stable_json_fingerprint;
6use crate::controller::{
7 ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest,
8 ControllerRegistry, RngPolicy,
9};
10use crate::controller_adapter::representation_type_id;
11use crate::data::{
12 BranchViewMode, BranchViewPlan, DataBinding, ExternalDataPlanEnvelope, ModelInputFusionMode,
13 ModelInputPortSpec, ModelInputSpec, RepresentationPlan, SOURCE_INDEX_METADATA_KEY,
14};
15use crate::error::{DagMlError, Result};
16use crate::fold::{FoldSet, NestedCvSpec};
17use crate::generation::{
18 enumerate_variants, generation_spec_fingerprint, GenerationSpec, VariantPlan,
19};
20use crate::graph::{GraphSpec, NodeKind, NodeSpec};
21use crate::ids::{ControllerId, FoldId, NodeId, VariantId};
22use crate::phase::Phase;
23use crate::policy::{AggregationPolicy, DataModelShapePlan, LeakageUnitPolicy};
24
25pub const CAMPAIGN_SPEC_SCHEMA_VERSION: u32 = 1;
26pub const CAMPAIGN_SPEC_SCHEMA_ID: &str =
27 "https://github.com/GBeurier/dag-ml/schemas/campaign_spec.v1.schema.json";
28pub const EXECUTION_PLAN_SCHEMA_VERSION: u32 = 1;
29pub const EXECUTION_PLAN_SCHEMA_ID: &str =
30 "https://github.com/GBeurier/dag-ml/schemas/execution_plan.v1.schema.json";
31
32#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
33pub struct SplitInvocation {
34 pub id: String,
35 #[serde(default)]
36 pub controller_id: Option<ControllerId>,
37 #[serde(default)]
38 pub leakage_policy: LeakageUnitPolicy,
39 #[serde(default)]
40 pub params: BTreeMap<String, serde_json::Value>,
41 #[serde(default)]
42 pub fold_set: Option<FoldSet>,
43}
44
45impl SplitInvocation {
46 pub fn validate(&self) -> Result<()> {
47 if self.id.trim().is_empty() {
48 return Err(DagMlError::CampaignValidation(
49 "split invocation id is empty".to_string(),
50 ));
51 }
52 self.leakage_policy.validate()?;
53 if let Some(fold_set) = &self.fold_set {
54 fold_set.validate()?;
55 }
56 Ok(())
57 }
58}
59
60#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
61pub struct CampaignSpec {
62 pub id: String,
63 pub root_seed: Option<u64>,
64 #[serde(default)]
65 pub leakage_policy: LeakageUnitPolicy,
66 #[serde(default)]
67 pub aggregation_policy: AggregationPolicy,
68 #[serde(default)]
69 pub split_invocation: Option<SplitInvocation>,
70 #[serde(default)]
71 pub generation: GenerationSpec,
72 #[serde(default)]
73 pub shape_plans: BTreeMap<NodeId, DataModelShapePlan>,
74 #[serde(default)]
75 pub data_bindings: BTreeMap<NodeId, Vec<DataBinding>>,
76 #[serde(default, skip_serializing_if = "Vec::is_empty")]
77 pub branch_view_plans: Vec<BranchViewPlan>,
78 #[serde(default, skip_serializing_if = "Option::is_none")]
81 pub inner_cv: Option<NestedCvSpec>,
82 #[serde(default)]
83 pub metadata: BTreeMap<String, serde_json::Value>,
84}
85
86impl CampaignSpec {
87 pub fn validate(&self) -> Result<()> {
88 if self.id.trim().is_empty() {
89 return Err(DagMlError::CampaignValidation(
90 "campaign id is empty".to_string(),
91 ));
92 }
93 self.leakage_policy.validate()?;
94 self.aggregation_policy.validate()?;
95 if let Some(inner_cv) = &self.inner_cv {
96 inner_cv.validate()?;
97 }
98 if let Some(split) = &self.split_invocation {
99 split.validate()?;
100 }
101 self.generation.validate()?;
102 for (node_id, shape_plan) in &self.shape_plans {
103 if node_id != &shape_plan.node_id {
104 return Err(DagMlError::CampaignValidation(format!(
105 "shape plan key `{node_id}` does not match node_id `{}`",
106 shape_plan.node_id
107 )));
108 }
109 shape_plan.validate()?;
110 }
111 for (node_id, bindings) in &self.data_bindings {
112 for binding in bindings {
113 if node_id != &binding.node_id {
114 return Err(DagMlError::CampaignValidation(format!(
115 "data binding key `{node_id}` does not match node_id `{}`",
116 binding.node_id
117 )));
118 }
119 binding.validate()?;
120 }
121 }
122 let mut branch_views = BTreeSet::new();
123 for plan in &self.branch_view_plans {
124 plan.validate()?;
125 if !branch_views.insert(plan.view_id.as_str()) {
126 return Err(DagMlError::CampaignValidation(format!(
127 "campaign `{}` contains duplicate branch view `{}`",
128 self.id, plan.view_id
129 )));
130 }
131 }
132 Ok(())
133 }
134
135 pub fn validate_data_envelope_relations(
136 &self,
137 envelope: &ExternalDataPlanEnvelope,
138 ) -> Result<()> {
139 envelope.validate()?;
140 let Some(relations) = &envelope.coordinator_relations else {
141 return Ok(());
142 };
143 let Some(split) = &self.split_invocation else {
144 return Ok(());
145 };
146 let Some(fold_set) = &split.fold_set else {
147 return Ok(());
148 };
149 relations.validate_against_fold_set(fold_set, &self.leakage_policy)?;
150 relations.validate_against_fold_set(fold_set, &split.leakage_policy)
151 }
152}
153
154#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
155pub struct GraphPlan {
156 pub graph: GraphSpec,
157 pub topological_order: Vec<NodeId>,
158 #[serde(default, skip_serializing_if = "Vec::is_empty")]
159 pub parallel_levels: Vec<Vec<NodeId>>,
160}
161
162impl GraphPlan {
163 pub fn from_graph(graph: GraphSpec) -> Result<Self> {
164 let topological_order = graph.topological_order()?;
165 let parallel_levels = graph.parallel_levels()?;
166 Ok(Self {
167 graph,
168 topological_order,
169 parallel_levels,
170 })
171 }
172
173 pub fn parallel_levels(&self) -> Result<Vec<Vec<NodeId>>> {
174 if self.parallel_levels.is_empty() {
175 return self.graph.parallel_levels();
176 }
177 Ok(self.parallel_levels.clone())
178 }
179}
180
181#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
182pub struct NodePlan {
183 pub node_id: NodeId,
184 pub kind: NodeKind,
185 pub controller_id: ControllerId,
186 pub controller_version: String,
187 pub supported_phases: BTreeSet<Phase>,
188 #[serde(default)]
189 pub controller_capabilities: BTreeSet<ControllerCapability>,
190 pub fit_scope: ControllerFitScope,
191 pub rng_policy: RngPolicy,
192 pub artifact_policy: ArtifactPolicy,
193 pub input_nodes: Vec<NodeId>,
194 pub output_nodes: Vec<NodeId>,
195 pub shape_plan: Option<DataModelShapePlan>,
196 #[serde(default)]
197 pub data_bindings: Vec<DataBinding>,
198 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
199 pub params: BTreeMap<String, serde_json::Value>,
200 #[serde(default, skip_serializing_if = "Option::is_none")]
203 pub inner_cv: Option<NestedCvSpec>,
204 pub params_fingerprint: String,
205}
206
207#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
208pub struct ExecutionPlan {
209 pub id: String,
210 pub graph_plan: GraphPlan,
211 pub campaign: CampaignSpec,
212 pub node_plans: BTreeMap<NodeId, NodePlan>,
213 pub controller_manifests: BTreeMap<ControllerId, ControllerManifest>,
214 pub variants: Vec<VariantPlan>,
215 pub fold_set: Option<FoldSet>,
216 pub graph_fingerprint: String,
217 pub campaign_fingerprint: String,
218 pub controller_fingerprint: String,
219}
220
221#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
222pub struct ExecutionScopePlan {
223 pub scope_id: String,
224 pub phase: Phase,
225 pub variant_id: Option<VariantId>,
226 pub variant_seed: Option<u64>,
227 pub fold_id: Option<FoldId>,
228 pub node_levels: Vec<Vec<NodeId>>,
229}
230
231#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
232pub struct PhaseExecutionSchedule {
233 pub plan_id: String,
234 pub phase: Phase,
235 pub scopes: Vec<ExecutionScopePlan>,
236}
237
238impl ExecutionPlan {
239 pub fn validate(&self) -> Result<()> {
240 self.graph_plan.graph.validate()?;
241 self.campaign.validate()?;
242 if !self.graph_plan.parallel_levels.is_empty()
243 && self.graph_plan.parallel_levels != self.graph_plan.graph.parallel_levels()?
244 {
245 return Err(DagMlError::Planning(
246 "graph plan parallel levels do not match graph".to_string(),
247 ));
248 }
249 if self.node_plans.len() != self.graph_plan.graph.nodes.len() {
250 return Err(DagMlError::Planning(
251 "execution plan node count does not match graph".to_string(),
252 ));
253 }
254 for node_id in &self.graph_plan.topological_order {
255 let plan = self.node_plans.get(node_id).ok_or_else(|| {
256 DagMlError::Planning(format!("missing node plan for `{node_id}`"))
257 })?;
258 let manifest = self
259 .controller_manifests
260 .get(&plan.controller_id)
261 .ok_or_else(|| {
262 DagMlError::Planning(format!(
263 "missing controller manifest `{}` for node `{node_id}`",
264 plan.controller_id
265 ))
266 })?;
267 if manifest.operator_kind != plan.kind {
268 return Err(DagMlError::Planning(format!(
269 "node `{node_id}` planned with incompatible controller `{}`",
270 manifest.controller_id
271 )));
272 }
273 if plan.controller_capabilities != manifest.capabilities {
274 return Err(DagMlError::Planning(format!(
275 "node `{node_id}` controller capabilities do not match manifest `{}`",
276 manifest.controller_id
277 )));
278 }
279 if plan.fit_scope != manifest.fit_scope
280 || plan.rng_policy != manifest.rng_policy
281 || plan.artifact_policy != manifest.artifact_policy
282 {
283 return Err(DagMlError::Planning(format!(
284 "node `{node_id}` controller policy fields do not match manifest `{}`",
285 manifest.controller_id
286 )));
287 }
288 for binding in &plan.data_bindings {
289 if binding.node_id != *node_id {
290 return Err(DagMlError::Planning(format!(
291 "node plan `{node_id}` contains data binding for `{}`",
292 binding.node_id
293 )));
294 }
295 binding.validate()?;
296 }
297 let graph_node = self
298 .graph_plan
299 .graph
300 .nodes
301 .iter()
302 .find(|node| &node.id == node_id)
303 .expect("topological node exists in graph");
304 validate_data_binding_requirements(node_id, plan, manifest, graph_node)?;
305 let actual_params_fingerprint = stable_json_fingerprint(&plan.params)?;
306 if actual_params_fingerprint != plan.params_fingerprint {
307 return Err(DagMlError::Planning(format!(
308 "node plan `{node_id}` params fingerprint does not match params"
309 )));
310 }
311 }
312 for (node_id, plan) in &self.node_plans {
318 if let Some(inner_cv) = &plan.inner_cv {
319 inner_cv.validate().map_err(|error| {
320 DagMlError::Planning(format!(
321 "node plan `{node_id}` has invalid inner_cv: {error}"
322 ))
323 })?;
324 }
325 }
326 self.validate_oof_controller_capabilities()?;
327 if let Some(fold_set) = &self.fold_set {
328 fold_set.validate()?;
329 }
330 if self.variants.is_empty() {
331 return Err(DagMlError::Planning(
332 "execution plan has no variants".to_string(),
333 ));
334 }
335 for variant in &self.variants {
336 variant.validate()?;
337 }
338 Ok(())
339 }
340
341 pub fn validate_parallel_controller_capabilities(
342 &self,
343 max_workers: usize,
344 phase: Phase,
345 ) -> Result<()> {
346 if max_workers <= 1 {
347 return Ok(());
348 }
349 let node_ids = self
350 .node_parallel_levels_for_phase(phase)?
351 .into_iter()
352 .flatten()
353 .collect::<Vec<_>>();
354 for node_id in node_ids {
355 let node_plan = self.node_plans.get(&node_id).ok_or_else(|| {
356 DagMlError::Planning(format!("missing node plan for `{node_id}`"))
357 })?;
358 let manifest = self
359 .controller_manifests
360 .get(&node_plan.controller_id)
361 .ok_or_else(|| {
362 DagMlError::Planning(format!(
363 "missing controller manifest `{}` for node `{}`",
364 node_plan.controller_id, node_plan.node_id
365 ))
366 })?;
367 if !manifest.supports_parallel_invocation() {
368 return Err(DagMlError::Planning(format!(
369 "parallel scheduler with {max_workers} workers requires controller `{}` for node `{}` to declare thread_safe or process_safe",
370 manifest.controller_id, node_plan.node_id
371 )));
372 }
373 }
374 Ok(())
375 }
376
377 fn validate_oof_controller_capabilities(&self) -> Result<()> {
378 for edge in &self.graph_plan.graph.edges {
379 if !edge.contract.requires_oof {
380 continue;
381 }
382 let source_plan = self.node_plans.get(&edge.source.node_id).ok_or_else(|| {
383 DagMlError::Planning(format!(
384 "OOF edge source node `{}` has no node plan",
385 edge.source.node_id
386 ))
387 })?;
388 if !source_plan
389 .controller_capabilities
390 .contains(&ControllerCapability::EmitsPredictions)
391 {
392 return Err(DagMlError::Planning(format!(
393 "OOF edge `{}.{}` -> `{}.{}` requires source controller `{}` to declare emits_predictions",
394 edge.source.node_id,
395 edge.source.port_name,
396 edge.target.node_id,
397 edge.target.port_name,
398 source_plan.controller_id
399 )));
400 }
401 let target_plan = self.node_plans.get(&edge.target.node_id).ok_or_else(|| {
402 DagMlError::Planning(format!(
403 "OOF edge target node `{}` has no node plan",
404 edge.target.node_id
405 ))
406 })?;
407 if !target_plan
408 .controller_capabilities
409 .contains(&ControllerCapability::ConsumesOofPredictions)
410 {
411 return Err(DagMlError::Planning(format!(
412 "OOF edge `{}.{}` -> `{}.{}` requires target controller `{}` to declare consumes_oof_predictions",
413 edge.source.node_id,
414 edge.source.port_name,
415 edge.target.node_id,
416 edge.target.port_name,
417 target_plan.controller_id
418 )));
419 }
420 }
421 Ok(())
422 }
423
424 pub fn node_parallel_levels_for_phase(&self, phase: Phase) -> Result<Vec<Vec<NodeId>>> {
425 let levels = self
426 .graph_plan
427 .parallel_levels()?
428 .into_iter()
429 .map(|level| {
430 level
431 .into_iter()
432 .filter(|node_id| {
433 self.node_plans
434 .get(node_id)
435 .is_some_and(|node_plan| node_plan.supported_phases.contains(&phase))
436 })
437 .collect::<Vec<_>>()
438 })
439 .filter(|level| !level.is_empty())
440 .collect::<Vec<_>>();
441 Ok(levels)
442 }
443
444 pub fn campaign_phase_schedule(&self, phase: Phase) -> Result<PhaseExecutionSchedule> {
445 self.validate()?;
446 let node_levels = self.node_parallel_levels_for_phase(phase)?;
447 let fold_ids = if phase == Phase::FitCv {
448 self.fold_set
449 .as_ref()
450 .map(|fold_set| {
451 fold_set
452 .folds
453 .iter()
454 .map(|fold| Some(fold.fold_id.clone()))
455 .collect::<Vec<_>>()
456 })
457 .unwrap_or_else(|| vec![None])
458 } else {
459 vec![None]
460 };
461 let mut scopes = Vec::new();
462 for variant in &self.variants {
463 for fold_id in &fold_ids {
464 scopes.push(ExecutionScopePlan {
465 scope_id: execution_scope_id(
466 phase,
467 Some(&variant.variant_id),
468 fold_id.as_ref(),
469 ),
470 phase,
471 variant_id: Some(variant.variant_id.clone()),
472 variant_seed: variant.seed,
473 fold_id: fold_id.clone(),
474 node_levels: node_levels.clone(),
475 });
476 }
477 }
478 Ok(PhaseExecutionSchedule {
479 plan_id: self.id.clone(),
480 phase,
481 scopes,
482 })
483 }
484
485 pub fn branch_view_for(&self, branch_id: &str) -> Option<&BranchViewPlan> {
489 branch_view_for_in(&self.campaign.branch_view_plans, branch_id)
490 }
491
492 pub fn branch_view_for_path(&self, branch_path: &[String]) -> Option<&BranchViewPlan> {
499 branch_view_for_path_in(&self.campaign.branch_view_plans, branch_path)
500 }
501}
502
503fn validate_data_binding_requirements(
504 node_id: &NodeId,
505 plan: &NodePlan,
506 manifest: &ControllerManifest,
507 node: &NodeSpec,
508) -> Result<()> {
509 let branch_view = branch_view_plan_from_node_metadata(node)?;
510 let Some(model_input) = manifest.model_input_spec()? else {
511 for binding in &plan.data_bindings {
512 let effective_source_ids = effective_binding_source_ids(binding, branch_view.as_ref())?;
513 if effective_source_ids.len() > 1 {
514 return Err(data_requirement_refusal(
515 "dagml.data_requirement.missing_data_requirements",
516 node_id,
517 binding,
518 manifest,
519 "multisource",
520 &effective_source_ids,
521 "multi-source data binding requires controller data_requirements".to_string(),
522 ));
523 }
524 }
525 return Ok(());
526 };
527 for binding in &plan.data_bindings {
528 let Some(port) = model_input
529 .ports
530 .iter()
531 .find(|port| port.name == binding.input_name)
532 else {
533 return Err(DagMlError::Planning(format!(
534 "node `{node_id}` data binding `{}` is not declared by controller `{}` data_requirements",
535 binding.input_name, manifest.controller_id
536 )));
537 };
538 if !port
539 .accepted_representations
540 .iter()
541 .any(|representation| representation == &binding.output_representation)
542 {
543 return Err(DagMlError::Planning(format!(
544 "node `{node_id}` data binding `{}` output representation `{}` is not accepted by controller `{}` data_requirements port `{}`",
545 binding.input_name,
546 binding.output_representation,
547 manifest.controller_id,
548 port.name
549 )));
550 }
551 if let Some(type_id) = representation_type_id(&binding.output_representation) {
552 if !port
553 .accepted_types
554 .iter()
555 .any(|accepted_type| accepted_type.as_str() == type_id)
556 {
557 return Err(DagMlError::Planning(format!(
558 "node `{node_id}` data binding `{}` output representation `{}` has registered type `{type_id}` but controller `{}` data_requirements port `{}` accepts types {:?}",
559 binding.input_name,
560 binding.output_representation,
561 manifest.controller_id,
562 port.name,
563 port.accepted_types
564 )));
565 }
566 }
567 validate_data_binding_source_shape(
568 node_id,
569 binding,
570 port,
571 &model_input,
572 manifest,
573 branch_view.as_ref(),
574 )?;
575 }
576 Ok(())
577}
578
579fn branch_view_plan_from_node_metadata(node: &NodeSpec) -> Result<Option<BranchViewPlan>> {
580 let Some(value) = node.metadata.get("dsl_branch_view_plan") else {
581 return Ok(None);
582 };
583 let plan: BranchViewPlan = serde_json::from_value(value.clone()).map_err(|error| {
584 DagMlError::Planning(format!(
585 "node `{}` carries malformed `dsl_branch_view_plan` metadata: {error}",
586 node.id
587 ))
588 })?;
589 plan.validate()
590 .map_err(|error| DagMlError::Planning(error.to_string()))?;
591 Ok(Some(plan))
592}
593
594fn validate_data_binding_source_shape(
595 node_id: &NodeId,
596 binding: &DataBinding,
597 port: &ModelInputPortSpec,
598 model_input: &ModelInputSpec,
599 manifest: &ControllerManifest,
600 branch_view: Option<&BranchViewPlan>,
601) -> Result<()> {
602 let effective_source_ids = effective_binding_source_ids(binding, branch_view)?;
603 if effective_source_ids.len() < 2 {
604 return Ok(());
605 }
606 if !port.multi_source {
607 return Err(data_requirement_refusal(
608 "dagml.data_requirement.multi_source_port_not_supported",
609 node_id,
610 binding,
611 manifest,
612 "multisource",
613 &effective_source_ids,
614 format!(
615 "controller `{}` data_requirements port `{}` does not declare multi_source=true",
616 manifest.controller_id, port.name
617 ),
618 ));
619 }
620 let Some(fusion) = &model_input.default_fusion else {
621 return Err(data_requirement_refusal(
622 "dagml.data_requirement.missing_multisource_fusion",
623 node_id,
624 binding,
625 manifest,
626 "multisource",
627 &effective_source_ids,
628 "multi-source data binding requires an explicit default_fusion policy".to_string(),
629 ));
630 };
631 validate_fusion_sources_match_binding(
632 node_id,
633 binding,
634 manifest,
635 fusion.representation_plan.as_ref(),
636 &effective_source_ids,
637 )?;
638 match fusion.mode {
639 ModelInputFusionMode::ConcatenateFeatures => {
640 if binding
641 .metadata
642 .get(SOURCE_INDEX_METADATA_KEY)
643 .and_then(serde_json::Value::as_object)
644 .is_none()
645 {
646 return Err(data_requirement_refusal(
647 "dagml.data_requirement.source_concat_requires_source_index",
648 node_id,
649 binding,
650 manifest,
651 "source_concat",
652 &effective_source_ids,
653 "source-concat feature fusion requires data binding metadata.source_index so feature-axis blocks are explicit".to_string(),
654 ));
655 }
656 }
657 ModelInputFusionMode::DictBySource | ModelInputFusionMode::Custom => {}
658 ModelInputFusionMode::SingleSource | ModelInputFusionMode::StackSamples => {
659 let fusion_mode = fusion_mode_label(fusion.mode);
660 return Err(data_requirement_refusal(
661 "dagml.data_requirement.unsupported_multisource_fusion_mode",
662 node_id,
663 binding,
664 manifest,
665 fusion_mode,
666 &effective_source_ids,
667 format!(
668 "multi-source data binding cannot be planned with default_fusion.mode={fusion_mode}"
669 ),
670 ));
671 }
672 }
673 Ok(())
674}
675
676fn fusion_mode_label(mode: ModelInputFusionMode) -> &'static str {
677 match mode {
678 ModelInputFusionMode::SingleSource => "single_source",
679 ModelInputFusionMode::ConcatenateFeatures => "concatenate_features",
680 ModelInputFusionMode::StackSamples => "stack_samples",
681 ModelInputFusionMode::DictBySource => "dict_by_source",
682 ModelInputFusionMode::Custom => "custom",
683 }
684}
685
686fn effective_binding_source_ids(
687 binding: &DataBinding,
688 branch_view: Option<&BranchViewPlan>,
689) -> Result<Vec<String>> {
690 let Some(branch_view) = branch_view else {
691 return Ok(binding.source_ids.clone());
692 };
693 if branch_view.mode != BranchViewMode::BySource {
694 return Ok(binding.source_ids.clone());
695 }
696 if branch_view.selector.source_ids.len() != 1 {
697 return Err(data_requirement_refusal_for_branch(
698 "dagml.data_requirement.unsupported_by_source_shape",
699 binding,
700 branch_view,
701 "by_source branch views must select exactly one source_id for per-source X-chain fit semantics".to_string(),
702 ));
703 }
704 if !binding.source_ids.is_empty() {
705 let declared = binding.source_ids.iter().collect::<BTreeSet<_>>();
706 for source_id in &branch_view.selector.source_ids {
707 if !declared.contains(source_id) {
708 return Err(data_requirement_refusal_for_branch(
709 "dagml.data_requirement.by_source_selector_outside_binding",
710 binding,
711 branch_view,
712 format!(
713 "by_source branch selector source `{source_id}` is not declared by data binding source_ids"
714 ),
715 ));
716 }
717 }
718 }
719 Ok(branch_view.selector.source_ids.clone())
720}
721
722fn validate_fusion_sources_match_binding(
723 node_id: &NodeId,
724 binding: &DataBinding,
725 manifest: &ControllerManifest,
726 representation_plan: Option<&RepresentationPlan>,
727 effective_source_ids: &[String],
728) -> Result<()> {
729 let Some(representation_plan) = representation_plan else {
730 return Ok(());
731 };
732 let component_sources = representation_plan_component_sources(representation_plan);
733 if component_sources.is_empty() {
734 return Ok(());
735 }
736 let declared = component_sources.iter().cloned().collect::<BTreeSet<_>>();
737 let effective = effective_source_ids.iter().collect::<BTreeSet<_>>();
738 if declared != effective {
739 return Err(data_requirement_refusal(
740 "dagml.data_requirement.representation_sources_mismatch",
741 node_id,
742 binding,
743 manifest,
744 "multisource",
745 effective_source_ids,
746 format!(
747 "default_fusion.representation_plan component_source_ids {:?} do not match binding source_ids {:?}",
748 component_sources, effective_source_ids
749 ),
750 ));
751 }
752 Ok(())
753}
754
755fn representation_plan_component_sources(plan: &RepresentationPlan) -> Vec<&String> {
756 match plan {
757 RepresentationPlan::Aggregate(_) => Vec::new(),
758 RepresentationPlan::CartesianProduct(plan) => {
759 plan.combination_plan.component_source_ids.iter().collect()
760 }
761 RepresentationPlan::MonteCarloCartesian(plan) => {
762 plan.combination_plan.component_source_ids.iter().collect()
763 }
764 RepresentationPlan::StackFixed(plan) => plan.component_source_ids.iter().collect(),
765 RepresentationPlan::StackPaddedMasked(plan) => plan.component_source_ids.iter().collect(),
766 }
767}
768
769fn data_requirement_refusal(
770 code: &'static str,
771 node_id: &NodeId,
772 binding: &DataBinding,
773 manifest: &ControllerManifest,
774 shape: &str,
775 source_ids: &[String],
776 message: String,
777) -> DagMlError {
778 DagMlError::Planning(format!(
779 "data requirement refusal: {}",
780 serde_json::json!({
781 "schema_version": 1,
782 "code": code,
783 "node_id": node_id.to_string(),
784 "input_name": binding.input_name.as_str(),
785 "controller_id": manifest.controller_id.to_string(),
786 "shape": shape,
787 "source_ids": source_ids,
788 "message": message
789 })
790 ))
791}
792
793fn data_requirement_refusal_for_branch(
794 code: &'static str,
795 binding: &DataBinding,
796 branch_view: &BranchViewPlan,
797 message: String,
798) -> DagMlError {
799 DagMlError::Planning(format!(
800 "data requirement refusal: {}",
801 serde_json::json!({
802 "schema_version": 1,
803 "code": code,
804 "node_id": binding.node_id.to_string(),
805 "input_name": binding.input_name.as_str(),
806 "branch_view_id": branch_view.view_id.as_str(),
807 "branch_id": branch_view.branch_id.as_str(),
808 "shape": "by_source",
809 "source_ids": &branch_view.selector.source_ids,
810 "message": message
811 })
812 ))
813}
814
815fn branch_view_for_in<'a>(
816 plans: &'a [BranchViewPlan],
817 branch_id: &str,
818) -> Option<&'a BranchViewPlan> {
819 plans.iter().find(|plan| plan.branch_id == branch_id)
820}
821
822fn branch_view_for_path_in<'a>(
823 plans: &'a [BranchViewPlan],
824 branch_path: &[String],
825) -> Option<&'a BranchViewPlan> {
826 for branch_id in branch_path.iter().rev() {
827 if let Some(plan) = branch_view_for_in(plans, branch_id) {
828 return Some(plan);
829 }
830 }
831 None
832}
833
834fn execution_scope_id(
835 phase: Phase,
836 variant_id: Option<&VariantId>,
837 fold_id: Option<&FoldId>,
838) -> String {
839 format!(
840 "scope:{}:{}:{}",
841 phase_scope_label(phase),
842 variant_id
843 .map(ToString::to_string)
844 .unwrap_or_else(|| "base".to_string()),
845 fold_id
846 .map(ToString::to_string)
847 .unwrap_or_else(|| "nofold".to_string())
848 )
849}
850
851fn phase_scope_label(phase: Phase) -> &'static str {
852 match phase {
853 Phase::Compile => "COMPILE",
854 Phase::Plan => "PLAN",
855 Phase::FitCv => "FIT_CV",
856 Phase::Select => "SELECT",
857 Phase::Refit => "REFIT",
858 Phase::Predict => "PREDICT",
859 Phase::Explain => "EXPLAIN",
860 }
861}
862
863pub fn build_execution_plan(
864 id: impl Into<String>,
865 graph: GraphSpec,
866 campaign: CampaignSpec,
867 registry: &ControllerRegistry,
868) -> Result<ExecutionPlan> {
869 let id = id.into();
870 if id.trim().is_empty() {
871 return Err(DagMlError::Planning(
872 "execution plan id is empty".to_string(),
873 ));
874 }
875 campaign.validate()?;
876 let graph_plan = GraphPlan::from_graph(graph)?;
877 validate_campaign_node_targets(&graph_plan.graph, &campaign)?;
878
879 let mut node_plans = BTreeMap::new();
880 let mut controller_manifests = BTreeMap::new();
881 for node_id in &graph_plan.topological_order {
882 let node = graph_plan
883 .graph
884 .nodes
885 .iter()
886 .find(|node| &node.id == node_id)
887 .expect("topological node exists");
888 let manifest = registry.resolve_for_node(node)?;
889 let params = node.params.clone();
890 let params_fingerprint = stable_json_fingerprint(¶ms)?;
891 let inner_cv = match node.metadata.get("dsl_inner_cv") {
895 Some(value) => {
896 let spec =
897 serde_json::from_value::<NestedCvSpec>(value.clone()).map_err(|error| {
898 DagMlError::Planning(format!(
899 "node `{}` has invalid dsl_inner_cv metadata: {error}",
900 node.id
901 ))
902 })?;
903 spec.validate().map_err(|error| {
906 DagMlError::Planning(format!(
907 "node `{}` has invalid dsl_inner_cv metadata: {error}",
908 node.id
909 ))
910 })?;
911 Some(spec)
912 }
913 None => None,
914 };
915 let shape_plan = campaign.shape_plans.get(&node.id).cloned();
916 let data_bindings = campaign
917 .data_bindings
918 .get(&node.id)
919 .cloned()
920 .unwrap_or_default();
921 node_plans.insert(
922 node.id.clone(),
923 NodePlan {
924 inner_cv,
925 node_id: node.id.clone(),
926 kind: node.kind.clone(),
927 controller_id: manifest.controller_id.clone(),
928 controller_version: manifest.controller_version.clone(),
929 supported_phases: manifest.supported_phases.clone(),
930 controller_capabilities: manifest.capabilities.clone(),
931 fit_scope: manifest.fit_scope,
932 rng_policy: manifest.rng_policy,
933 artifact_policy: manifest.artifact_policy,
934 input_nodes: graph_plan.graph.upstream_nodes(&node.id),
935 output_nodes: graph_plan.graph.downstream_nodes(&node.id),
936 shape_plan,
937 data_bindings,
938 params,
939 params_fingerprint,
940 },
941 );
942 controller_manifests.insert(manifest.controller_id.clone(), manifest);
943 }
944
945 let fold_set = campaign
946 .split_invocation
947 .as_ref()
948 .and_then(|split| split.fold_set.clone());
949 validate_search_space_fingerprint(&graph_plan.graph, &campaign)?;
950 let variants = enumerate_variants(&campaign.generation, campaign.root_seed)?;
951 validate_generation_override_targets(&graph_plan.graph, &variants)?;
952 let graph_fingerprint = stable_json_fingerprint(&graph_plan.graph)?;
953 let campaign_fingerprint = stable_json_fingerprint(&campaign)?;
954 let controller_fingerprint = stable_json_fingerprint(&controller_manifests)?;
955 let plan = ExecutionPlan {
956 id,
957 graph_plan,
958 campaign,
959 node_plans,
960 controller_manifests,
961 variants,
962 fold_set,
963 graph_fingerprint,
964 campaign_fingerprint,
965 controller_fingerprint,
966 };
967 plan.validate()?;
968 Ok(plan)
969}
970
971fn validate_search_space_fingerprint(graph: &GraphSpec, campaign: &CampaignSpec) -> Result<()> {
972 let Some(expected_fingerprint) = &graph.search_space_fingerprint else {
973 return Ok(());
974 };
975 if expected_fingerprint.trim().is_empty() {
976 return Err(DagMlError::Planning(format!(
977 "graph `{}` has empty search_space_fingerprint",
978 graph.id
979 )));
980 }
981 let actual_fingerprint = generation_spec_fingerprint(&campaign.generation)?;
982 if expected_fingerprint != &actual_fingerprint {
983 return Err(DagMlError::Planning(format!(
984 "graph `{}` search_space_fingerprint does not match campaign generation spec",
985 graph.id
986 )));
987 }
988 Ok(())
989}
990
991fn validate_generation_override_targets(graph: &GraphSpec, variants: &[VariantPlan]) -> Result<()> {
992 let node_ids = graph
993 .nodes
994 .iter()
995 .map(|node| node.id.clone())
996 .collect::<BTreeSet<_>>();
997 for variant in variants {
998 for node_id in variant.param_override_targets()? {
999 if !node_ids.contains(&node_id) {
1000 return Err(DagMlError::Planning(format!(
1001 "variant `{}` overrides params for unknown node `{node_id}`",
1002 variant.variant_id
1003 )));
1004 }
1005 }
1006 }
1007 Ok(())
1008}
1009
1010fn validate_campaign_node_targets(graph: &GraphSpec, campaign: &CampaignSpec) -> Result<()> {
1011 let node_ids = graph
1012 .nodes
1013 .iter()
1014 .map(|node| &node.id)
1015 .collect::<BTreeSet<_>>();
1016 for node_id in campaign.shape_plans.keys() {
1017 if !node_ids.contains(node_id) {
1018 return Err(DagMlError::Planning(format!(
1019 "shape plan references unknown node `{node_id}`"
1020 )));
1021 }
1022 }
1023 for node_id in campaign.data_bindings.keys() {
1024 if !node_ids.contains(node_id) {
1025 return Err(DagMlError::Planning(format!(
1026 "data binding references unknown node `{node_id}`"
1027 )));
1028 }
1029 }
1030 Ok(())
1031}
1032
1033pub fn prune_plan_to_active(
1065 plan: &ExecutionPlan,
1066 active_nodes: &BTreeSet<NodeId>,
1067 all_choice_nodes: &BTreeSet<NodeId>,
1068 variant: &VariantPlan,
1069) -> Result<ExecutionPlan> {
1070 plan.validate()?;
1071 variant.validate()?;
1072 for node_id in active_nodes {
1073 if !plan.node_plans.contains_key(node_id) {
1074 return Err(DagMlError::Planning(format!(
1075 "operator-SELECT prune: active node `{node_id}` is not in the union plan"
1076 )));
1077 }
1078 }
1079 if active_nodes.is_empty() {
1080 return Err(DagMlError::Planning(
1081 "operator-SELECT prune: active node set is empty".to_string(),
1082 ));
1083 }
1084
1085 let graph = &plan.graph_plan.graph;
1088 let mut ancestors = BTreeSet::<NodeId>::new();
1089 let mut stack: Vec<NodeId> = active_nodes.iter().cloned().collect();
1090 while let Some(node_id) = stack.pop() {
1091 for upstream in graph.upstream_nodes(&node_id) {
1092 if ancestors.insert(upstream.clone()) {
1093 stack.push(upstream);
1094 }
1095 }
1096 }
1097 let shared_prefix = ancestors
1098 .into_iter()
1099 .filter(|node_id| !all_choice_nodes.contains(node_id))
1100 .collect::<BTreeSet<_>>();
1101
1102 let keep = shared_prefix
1103 .iter()
1104 .chain(active_nodes.iter())
1105 .cloned()
1106 .collect::<BTreeSet<_>>();
1107
1108 let mut pruned_graph = graph.clone();
1111 pruned_graph.nodes.retain(|node| keep.contains(&node.id));
1112 pruned_graph
1113 .edges
1114 .retain(|edge| keep.contains(&edge.source.node_id) && keep.contains(&edge.target.node_id));
1115
1116 let graph_plan = GraphPlan::from_graph(pruned_graph)?;
1117
1118 let mut node_plans = BTreeMap::new();
1121 for (node_id, node_plan) in &plan.node_plans {
1122 if !keep.contains(node_id) {
1123 continue;
1124 }
1125 let mut pruned_node_plan = node_plan.clone();
1126 pruned_node_plan.input_nodes = graph_plan.graph.upstream_nodes(node_id);
1127 pruned_node_plan.output_nodes = graph_plan.graph.downstream_nodes(node_id);
1128 node_plans.insert(node_id.clone(), pruned_node_plan);
1129 }
1130
1131 let graph_fingerprint = stable_json_fingerprint(&graph_plan.graph)?;
1132 let pruned = ExecutionPlan {
1133 id: plan.id.clone(),
1134 graph_plan,
1135 campaign: plan.campaign.clone(),
1136 node_plans,
1137 controller_manifests: plan.controller_manifests.clone(),
1138 variants: vec![variant.clone()],
1139 fold_set: plan.fold_set.clone(),
1140 graph_fingerprint,
1141 campaign_fingerprint: plan.campaign_fingerprint.clone(),
1142 controller_fingerprint: plan.controller_fingerprint.clone(),
1143 };
1144 pruned.validate()?;
1145 validate_active_inputs(&pruned, graph)?;
1146 Ok(pruned)
1147}
1148
1149fn validate_active_inputs(plan: &ExecutionPlan, union_graph: &GraphSpec) -> Result<()> {
1163 let pruned_graph = &plan.graph_plan.graph;
1164 let kept: BTreeSet<&NodeId> = pruned_graph.nodes.iter().map(|node| &node.id).collect();
1165
1166 let mut union_port_sources = BTreeMap::<(NodeId, String), usize>::new();
1167 for edge in &union_graph.edges {
1168 if !kept.contains(&edge.target.node_id) {
1169 continue;
1170 }
1171 *union_port_sources
1172 .entry((edge.target.node_id.clone(), edge.target.port_name.clone()))
1173 .or_insert(0) += 1;
1174 }
1175 let mut pruned_port_sources = BTreeMap::<(NodeId, String), usize>::new();
1176 for edge in &pruned_graph.edges {
1177 *pruned_port_sources
1178 .entry((edge.target.node_id.clone(), edge.target.port_name.clone()))
1179 .or_insert(0) += 1;
1180 }
1181
1182 for (key, union_count) in &union_port_sources {
1183 let pruned_count = pruned_port_sources.get(key).copied().unwrap_or(0);
1184 let (node_id, port_name) = key;
1185 if *union_count >= 1 && pruned_count == 0 {
1186 return Err(DagMlError::Planning(format!(
1187 "operator-SELECT prune left node `{node_id}` required input port `{port_name}` with zero surviving sources (dangling): its producer was pruned away"
1188 )));
1189 }
1190 if pruned_count > 1 {
1191 return Err(DagMlError::Planning(format!(
1192 "operator-SELECT prune left node `{node_id}` input port `{port_name}` fed by {pruned_count} surviving sources (ambiguous)"
1193 )));
1194 }
1195 }
1196 Ok(())
1197}
1198
1199#[cfg(test)]
1200mod tests {
1201 use std::collections::{BTreeMap, BTreeSet};
1202 use std::time::{Duration, Instant};
1203
1204 use super::*;
1205 use crate::controller::{
1206 ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest, RngPolicy,
1207 };
1208 use crate::fold::FoldPartitionMode;
1209
1210 #[test]
1211 fn inner_cv_is_declarable_at_campaign_and_node_level() {
1212 let campaign_json = r#"{"id":"c","root_seed":null,"inner_cv":{"kind":"kfold","n_splits":3,"shuffle":false,"seed":5}}"#;
1214 let campaign: CampaignSpec = serde_json::from_str(campaign_json).unwrap();
1215 campaign.validate().unwrap();
1216 assert!(campaign.inner_cv.is_some());
1217
1218 let node_inner = crate::fold::NestedCvSpec::KFold(crate::fold::KFoldSpec {
1220 n_splits: 4,
1221 shuffle: false,
1222 seed: Some(6),
1223 });
1224 let resolved = crate::fold::resolve_inner_cv(Some(&node_inner), campaign.inner_cv.as_ref());
1225 assert_eq!(resolved, Some(&node_inner));
1226
1227 let bare = r#"{"id":"c","root_seed":null}"#;
1229 let bare_campaign: CampaignSpec = serde_json::from_str(bare).unwrap();
1230 assert!(bare_campaign.inner_cv.is_none());
1231 let reserialized = serde_json::to_string(&bare_campaign).unwrap();
1232 assert!(!reserialized.contains("inner_cv"));
1233
1234 let bad: CampaignSpec = serde_json::from_str(
1237 r#"{"id":"c","root_seed":null,"inner_cv":{"kind":"kfold","n_splits":1,"shuffle":false,"seed":null}}"#,
1238 )
1239 .unwrap();
1240 let error = bad.validate().unwrap_err();
1241 assert!(error.to_string().contains("at least two splits"));
1242 }
1243
1244 #[test]
1245 fn execution_plan_validate_rejects_invalid_node_local_inner_cv() {
1246 let campaign = CampaignSpec {
1250 inner_cv: None,
1251 id: "campaign:plan-validate".to_string(),
1252 root_seed: Some(7),
1253 leakage_policy: LeakageUnitPolicy::default(),
1254 aggregation_policy: AggregationPolicy::default(),
1255 split_invocation: None,
1256 generation: Default::default(),
1257 shape_plans: BTreeMap::new(),
1258 data_bindings: BTreeMap::new(),
1259 branch_view_plans: Vec::new(),
1260 metadata: BTreeMap::new(),
1261 };
1262 let mut plan =
1263 build_execution_plan("plan:validate", graph(), campaign, ®istry()).unwrap();
1264 plan.validate().unwrap();
1265 plan.node_plans
1266 .get_mut(&NodeId::new("model:pls").unwrap())
1267 .unwrap()
1268 .inner_cv = Some(crate::fold::NestedCvSpec::KFold(crate::fold::KFoldSpec {
1269 n_splits: 1,
1270 shuffle: false,
1271 seed: None,
1272 }));
1273 let error = plan.validate().unwrap_err();
1274 assert!(matches!(error, DagMlError::Planning(_)));
1275 assert!(error.to_string().contains("invalid inner_cv"));
1276 assert!(error.to_string().contains("at least two splits"));
1277 }
1278
1279 #[test]
1280 fn build_execution_plan_lowers_dsl_inner_cv_metadata_into_node_plan() {
1281 let mut graph = graph();
1282 graph
1283 .nodes
1284 .iter_mut()
1285 .find(|node| node.id.as_str() == "model:pls")
1286 .unwrap()
1287 .metadata
1288 .insert(
1289 "dsl_inner_cv".to_string(),
1290 serde_json::json!({"kind": "kfold", "n_splits": 3, "shuffle": false, "seed": 9}),
1291 );
1292
1293 let campaign = CampaignSpec {
1294 inner_cv: None,
1295 id: "campaign:inner-cv".to_string(),
1296 root_seed: Some(7),
1297 leakage_policy: LeakageUnitPolicy::default(),
1298 aggregation_policy: AggregationPolicy::default(),
1299 split_invocation: None,
1300 generation: Default::default(),
1301 shape_plans: BTreeMap::new(),
1302 data_bindings: BTreeMap::new(),
1303 branch_view_plans: Vec::new(),
1304 metadata: BTreeMap::new(),
1305 };
1306
1307 let plan = build_execution_plan("plan:inner-cv", graph, campaign, ®istry()).unwrap();
1308 match &plan.node_plans[&NodeId::new("model:pls").unwrap()].inner_cv {
1309 Some(crate::fold::NestedCvSpec::KFold(k)) => {
1310 assert_eq!(k.n_splits, 3);
1311 assert_eq!(k.seed, Some(9));
1312 }
1313 other => panic!("expected lowered KFold inner_cv, got {other:?}"),
1314 }
1315 assert!(plan.node_plans[&NodeId::new("transform:snv").unwrap()]
1316 .inner_cv
1317 .is_none());
1318 }
1319
1320 #[test]
1321 fn build_execution_plan_rejects_malformed_dsl_inner_cv_metadata() {
1322 let mut graph = graph();
1323 graph
1324 .nodes
1325 .iter_mut()
1326 .find(|node| node.id.as_str() == "model:pls")
1327 .unwrap()
1328 .metadata
1329 .insert(
1330 "dsl_inner_cv".to_string(),
1331 serde_json::json!({"kind": "not_a_real_kind"}),
1332 );
1333
1334 let campaign = CampaignSpec {
1335 inner_cv: None,
1336 id: "campaign:inner-cv.bad".to_string(),
1337 root_seed: Some(7),
1338 leakage_policy: LeakageUnitPolicy::default(),
1339 aggregation_policy: AggregationPolicy::default(),
1340 split_invocation: None,
1341 generation: Default::default(),
1342 shape_plans: BTreeMap::new(),
1343 data_bindings: BTreeMap::new(),
1344 branch_view_plans: Vec::new(),
1345 metadata: BTreeMap::new(),
1346 };
1347
1348 let error =
1349 build_execution_plan("plan:inner-cv.bad", graph, campaign, ®istry()).unwrap_err();
1350 assert!(matches!(error, DagMlError::Planning(_)));
1351 assert!(error.to_string().contains("invalid dsl_inner_cv metadata"));
1352 }
1353
1354 #[test]
1355 fn build_execution_plan_rejects_semantically_invalid_dsl_inner_cv() {
1356 let mut graph = graph();
1359 graph
1360 .nodes
1361 .iter_mut()
1362 .find(|node| node.id.as_str() == "model:pls")
1363 .unwrap()
1364 .metadata
1365 .insert(
1366 "dsl_inner_cv".to_string(),
1367 serde_json::json!({"kind": "kfold", "n_splits": 1, "shuffle": false, "seed": null}),
1368 );
1369
1370 let campaign = CampaignSpec {
1371 inner_cv: None,
1372 id: "campaign:inner-cv.nsplits".to_string(),
1373 root_seed: Some(7),
1374 leakage_policy: LeakageUnitPolicy::default(),
1375 aggregation_policy: AggregationPolicy::default(),
1376 split_invocation: None,
1377 generation: Default::default(),
1378 shape_plans: BTreeMap::new(),
1379 data_bindings: BTreeMap::new(),
1380 branch_view_plans: Vec::new(),
1381 metadata: BTreeMap::new(),
1382 };
1383
1384 let error = build_execution_plan("plan:inner-cv.nsplits", graph, campaign, ®istry())
1385 .unwrap_err();
1386 assert!(matches!(error, DagMlError::Planning(_)));
1387 assert!(error.to_string().contains("at least two splits"));
1388 }
1389 use crate::data::{
1390 BranchViewMode, BranchViewPlan, DataBinding, DataViewSelector, SOURCE_INDEX_METADATA_KEY,
1391 };
1392 use crate::generation::{
1393 GenerationChoice, GenerationConstraints, GenerationDimension, GenerationParamOverride,
1394 GenerationStrategy,
1395 };
1396 use crate::graph::{
1397 EdgeContract, EdgeSpec, GraphInterface, NodeSpec, PortCardinality, PortKind, PortRef,
1398 PortSchema, PortSpec,
1399 };
1400 use crate::ids::{ControllerId, FoldId, ObservationId, SampleId, TargetId};
1401 use crate::phase::Phase;
1402 use crate::policy::{DataModelShapePlan, Granularity};
1403 use crate::relation::{SampleRelation, SampleRelationSet};
1404
1405 fn port(name: &str, kind: PortKind) -> PortSpec {
1406 PortSpec {
1407 name: name.to_string(),
1408 kind,
1409 representation: None,
1410 cardinality: PortCardinality::One,
1411 unit_level: None,
1412 alignment_key: None,
1413 target_level: None,
1414 description: String::new(),
1415 }
1416 }
1417
1418 fn node(id: &str, kind: NodeKind, inputs: Vec<PortSpec>, outputs: Vec<PortSpec>) -> NodeSpec {
1419 NodeSpec {
1420 id: NodeId::new(id).unwrap(),
1421 kind,
1422 operator: None,
1423 params: BTreeMap::new(),
1424 ports: PortSchema { inputs, outputs },
1425 metadata: BTreeMap::new(),
1426 seed_label: None,
1427 }
1428 }
1429
1430 fn graph() -> GraphSpec {
1431 GraphSpec {
1432 id: "g".to_string(),
1433 interface: GraphInterface::default(),
1434 nodes: vec![
1435 node(
1436 "transform:snv",
1437 NodeKind::Transform,
1438 vec![],
1439 vec![port("x", PortKind::Data)],
1440 ),
1441 node(
1442 "model:pls",
1443 NodeKind::Model,
1444 vec![port("x", PortKind::Data)],
1445 vec![port("pred", PortKind::Prediction)],
1446 ),
1447 ],
1448 edges: vec![EdgeSpec {
1449 source: PortRef {
1450 node_id: NodeId::new("transform:snv").unwrap(),
1451 port_name: "x".to_string(),
1452 },
1453 target: PortRef {
1454 node_id: NodeId::new("model:pls").unwrap(),
1455 port_name: "x".to_string(),
1456 },
1457 contract: EdgeContract {
1458 requires_oof: false,
1459 requires_fold_alignment: false,
1460 ..EdgeContract::new(PortKind::Data, None)
1461 },
1462 }],
1463 search_space_fingerprint: None,
1464 metadata: BTreeMap::new(),
1465 }
1466 }
1467
1468 fn manifest(id: &str, kind: NodeKind) -> ControllerManifest {
1469 let mut capabilities = BTreeSet::from([
1470 ControllerCapability::Deterministic,
1471 ControllerCapability::ThreadSafe,
1472 ControllerCapability::ProcessSafe,
1473 ]);
1474 if kind == NodeKind::Model {
1475 capabilities.insert(ControllerCapability::EmitsPredictions);
1476 capabilities.insert(ControllerCapability::ConsumesOofPredictions);
1477 }
1478 ControllerManifest {
1479 controller_id: ControllerId::new(id).unwrap(),
1480 controller_version: "0.1.0".to_string(),
1481 operator_kind: kind,
1482 priority: 0,
1483 supported_phases: BTreeSet::from([Phase::FitCv, Phase::Refit, Phase::Predict]),
1484 input_ports: Vec::new(),
1485 output_ports: Vec::new(),
1486 data_requirements: None,
1487 capabilities,
1488 operator_selectors: Vec::new(),
1489 fit_scope: ControllerFitScope::FoldTrain,
1490 rng_policy: RngPolicy::UsesCoreSeed,
1491 artifact_policy: ArtifactPolicy::Serializable,
1492 }
1493 }
1494
1495 fn registry() -> ControllerRegistry {
1496 let mut registry = ControllerRegistry::new();
1497 registry
1498 .register(manifest("controller:transform", NodeKind::Transform))
1499 .unwrap();
1500 registry
1501 .register(manifest("controller:model", NodeKind::Model))
1502 .unwrap();
1503 registry
1504 }
1505
1506 fn registry_with_model_data_requirements(
1507 accepted_representations: &[&str],
1508 accepted_types: &[&str],
1509 ) -> ControllerRegistry {
1510 let mut registry = ControllerRegistry::new();
1511 registry
1512 .register(manifest("controller:transform", NodeKind::Transform))
1513 .unwrap();
1514 let mut model = manifest("controller:model", NodeKind::Model);
1515 model.data_requirements = Some(serde_json::json!({
1516 "schema_version": 1,
1517 "ports": [
1518 {
1519 "name": "x",
1520 "accepted_representations": accepted_representations,
1521 "accepted_types": accepted_types,
1522 "rank": 2,
1523 "multi_source": true,
1524 "optional": false
1525 }
1526 ],
1527 "metadata": {
1528 "source": "plan-test"
1529 }
1530 }));
1531 registry.register(model).unwrap();
1532 registry
1533 }
1534
1535 fn registry_with_model_data_requirements_json(
1536 data_requirements: serde_json::Value,
1537 ) -> ControllerRegistry {
1538 let mut registry = ControllerRegistry::new();
1539 registry
1540 .register(manifest("controller:transform", NodeKind::Transform))
1541 .unwrap();
1542 let mut model = manifest("controller:model", NodeKind::Model);
1543 model.data_requirements = Some(data_requirements);
1544 registry.register(model).unwrap();
1545 registry
1546 }
1547
1548 fn model_data_requirements(
1549 multi_source: bool,
1550 default_fusion: Option<serde_json::Value>,
1551 ) -> serde_json::Value {
1552 let mut spec = serde_json::json!({
1553 "schema_version": 1,
1554 "ports": [
1555 {
1556 "name": "x",
1557 "accepted_representations": ["tabular_numeric"],
1558 "accepted_types": ["table"],
1559 "rank": 2,
1560 "multi_source": multi_source,
1561 "optional": false
1562 }
1563 ],
1564 "metadata": {
1565 "source": "plan-test"
1566 }
1567 });
1568 if let Some(default_fusion) = default_fusion {
1569 spec.as_object_mut()
1570 .unwrap()
1571 .insert("default_fusion".to_string(), default_fusion);
1572 }
1573 spec
1574 }
1575
1576 fn multisource_binding(node_id: &NodeId) -> DataBinding {
1577 let mut binding = data_binding(node_id);
1578 binding.request_id = "nir-chem-source-concat".to_string();
1579 binding.feature_set_id = Some("x_fused".to_string());
1580 binding.source_ids = vec!["nir".to_string(), "chem".to_string()];
1581 binding
1582 }
1583
1584 fn add_source_index(binding: &mut DataBinding) {
1585 binding.metadata.insert(
1586 SOURCE_INDEX_METADATA_KEY.to_string(),
1587 serde_json::json!({
1588 "nir": 0,
1589 "chem": 1
1590 }),
1591 );
1592 }
1593
1594 fn source_concat_fusion() -> serde_json::Value {
1595 serde_json::json!({
1596 "mode": "concatenate_features",
1597 "alignment": "sample_id",
1598 "adapter_id": null,
1599 "params": {
1600 "namespace_columns": true
1601 }
1602 })
1603 }
1604
1605 fn by_source_graph(source_ids: Vec<&str>) -> GraphSpec {
1606 let mut graph = graph();
1607 let branch_view = BranchViewPlan {
1608 view_id: "branch_view:source".to_string(),
1609 branch_id: "branch:source".to_string(),
1610 mode: BranchViewMode::BySource,
1611 selector: DataViewSelector {
1612 source_ids: source_ids.into_iter().map(str::to_string).collect(),
1613 ..Default::default()
1614 },
1615 allow_overlap: false,
1616 metadata: BTreeMap::new(),
1617 };
1618 graph
1619 .nodes
1620 .iter_mut()
1621 .find(|node| node.id.as_str() == "model:pls")
1622 .unwrap()
1623 .metadata
1624 .insert(
1625 "dsl_branch_view_plan".to_string(),
1626 serde_json::to_value(branch_view).unwrap(),
1627 );
1628 graph
1629 }
1630
1631 fn refusal_payload(error: DagMlError) -> serde_json::Value {
1632 let message = error.to_string();
1633 let payload = message
1634 .split_once("data requirement refusal: ")
1635 .unwrap_or_else(|| panic!("missing structured refusal payload in: {message}"))
1636 .1;
1637 serde_json::from_str(payload).unwrap()
1638 }
1639
1640 fn campaign(id: &str) -> CampaignSpec {
1641 CampaignSpec {
1642 id: id.to_string(),
1643 root_seed: Some(7),
1644 leakage_policy: LeakageUnitPolicy::default(),
1645 aggregation_policy: AggregationPolicy::default(),
1646 split_invocation: None,
1647 generation: Default::default(),
1648 shape_plans: BTreeMap::new(),
1649 data_bindings: BTreeMap::new(),
1650 branch_view_plans: Vec::new(),
1651 inner_cv: None,
1652 metadata: BTreeMap::new(),
1653 }
1654 }
1655
1656 #[test]
1657 fn build_execution_plan_consumes_controller_data_requirements_for_bindings() {
1658 let model_id = NodeId::new("model:pls").unwrap();
1659 let mut campaign = campaign("campaign:datareq.ok");
1660 campaign.data_bindings = BTreeMap::from([(
1661 model_id,
1662 vec![data_binding(&NodeId::new("model:pls").unwrap())],
1663 )]);
1664 let plan = build_execution_plan(
1665 "plan:datareq.ok",
1666 graph(),
1667 campaign,
1668 ®istry_with_model_data_requirements(&["tabular_numeric"], &["table"]),
1669 )
1670 .unwrap();
1671 assert_eq!(
1672 plan.node_plans[&NodeId::new("model:pls").unwrap()].data_bindings[0]
1673 .output_representation,
1674 "tabular_numeric"
1675 );
1676 }
1677
1678 #[test]
1679 fn build_execution_plan_rejects_binding_representation_outside_data_requirements() {
1680 let model_id = NodeId::new("model:pls").unwrap();
1681 let mut campaign = campaign("campaign:datareq.representation");
1682 campaign.data_bindings =
1683 BTreeMap::from([(model_id.clone(), vec![data_binding(&model_id)])]);
1684 let error = build_execution_plan(
1685 "plan:datareq.representation",
1686 graph(),
1687 campaign,
1688 ®istry_with_model_data_requirements(&["signal_1d"], &["dense_signal"]),
1689 )
1690 .unwrap_err();
1691 assert!(
1692 error.to_string().contains("output representation"),
1693 "unexpected error: {error}"
1694 );
1695 }
1696
1697 #[test]
1698 fn build_execution_plan_rejects_binding_registered_type_outside_data_requirements() {
1699 let model_id = NodeId::new("model:pls").unwrap();
1700 let mut campaign = campaign("campaign:datareq.type");
1701 campaign.data_bindings =
1702 BTreeMap::from([(model_id.clone(), vec![data_binding(&model_id)])]);
1703 let error = build_execution_plan(
1704 "plan:datareq.type",
1705 graph(),
1706 campaign,
1707 ®istry_with_model_data_requirements(&["tabular_numeric"], &["dense_signal"]),
1708 )
1709 .unwrap_err();
1710 assert!(
1711 error.to_string().contains("registered type `table`"),
1712 "unexpected error: {error}"
1713 );
1714 }
1715
1716 #[test]
1717 fn build_execution_plan_rejects_source_concat_without_source_index() {
1718 let model_id = NodeId::new("model:pls").unwrap();
1719 let mut campaign = campaign("campaign:datareq.source-concat.no-index");
1720 campaign.data_bindings =
1721 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
1722
1723 let error = build_execution_plan(
1724 "plan:datareq.source-concat.no-index",
1725 graph(),
1726 campaign,
1727 ®istry_with_model_data_requirements_json(model_data_requirements(
1728 true,
1729 Some(source_concat_fusion()),
1730 )),
1731 )
1732 .unwrap_err();
1733 let payload = refusal_payload(error);
1734
1735 assert_eq!(
1736 payload["code"],
1737 "dagml.data_requirement.source_concat_requires_source_index"
1738 );
1739 assert_eq!(payload["shape"], "source_concat");
1740 assert_eq!(payload["source_ids"], serde_json::json!(["nir", "chem"]));
1741 }
1742
1743 #[test]
1744 fn build_execution_plan_rejects_multisource_binding_without_data_requirements() {
1745 let model_id = NodeId::new("model:pls").unwrap();
1746 let mut campaign = campaign("campaign:datareq.multisource.no-requirements");
1747 campaign.data_bindings =
1748 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
1749
1750 let error = build_execution_plan(
1751 "plan:datareq.multisource.no-requirements",
1752 graph(),
1753 campaign,
1754 ®istry(),
1755 )
1756 .unwrap_err();
1757 let payload = refusal_payload(error);
1758
1759 assert_eq!(
1760 payload["code"],
1761 "dagml.data_requirement.missing_data_requirements"
1762 );
1763 assert_eq!(payload["source_ids"], serde_json::json!(["nir", "chem"]));
1764 }
1765
1766 #[test]
1767 fn build_execution_plan_accepts_source_concat_with_source_index() {
1768 let model_id = NodeId::new("model:pls").unwrap();
1769 let mut binding = multisource_binding(&model_id);
1770 add_source_index(&mut binding);
1771 let mut campaign = campaign("campaign:datareq.source-concat.index");
1772 campaign.data_bindings = BTreeMap::from([(model_id.clone(), vec![binding])]);
1773
1774 let plan = build_execution_plan(
1775 "plan:datareq.source-concat.index",
1776 graph(),
1777 campaign,
1778 ®istry_with_model_data_requirements_json(model_data_requirements(
1779 true,
1780 Some(source_concat_fusion()),
1781 )),
1782 )
1783 .unwrap();
1784
1785 assert_eq!(
1786 plan.node_plans[&model_id].data_bindings[0].metadata[SOURCE_INDEX_METADATA_KEY],
1787 serde_json::json!({"nir": 0, "chem": 1})
1788 );
1789 }
1790
1791 #[test]
1792 fn by_source_branch_allows_single_source_fit_from_multisource_binding() {
1793 let model_id = NodeId::new("model:pls").unwrap();
1794 let mut campaign = campaign("campaign:datareq.by-source.single");
1795 campaign.data_bindings =
1796 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
1797
1798 let plan = build_execution_plan(
1799 "plan:datareq.by-source.single",
1800 by_source_graph(vec!["nir"]),
1801 campaign,
1802 ®istry_with_model_data_requirements_json(model_data_requirements(false, None)),
1803 )
1804 .unwrap();
1805
1806 assert_eq!(
1807 plan.node_plans[&model_id].data_bindings[0].source_ids.len(),
1808 2
1809 );
1810 }
1811
1812 #[test]
1813 fn by_source_branch_refuses_multi_source_selector_shape() {
1814 let model_id = NodeId::new("model:pls").unwrap();
1815 let mut campaign = campaign("campaign:datareq.by-source.multi");
1816 campaign.data_bindings =
1817 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
1818
1819 let error = build_execution_plan(
1820 "plan:datareq.by-source.multi",
1821 by_source_graph(vec!["nir", "chem"]),
1822 campaign,
1823 ®istry_with_model_data_requirements_json(model_data_requirements(true, None)),
1824 )
1825 .unwrap_err();
1826 let payload = refusal_payload(error);
1827
1828 assert_eq!(
1829 payload["code"],
1830 "dagml.data_requirement.unsupported_by_source_shape"
1831 );
1832 assert_eq!(payload["shape"], "by_source");
1833 assert_eq!(payload["source_ids"], serde_json::json!(["nir", "chem"]));
1834 }
1835
1836 fn large_linear_graph(transform_count: usize) -> GraphSpec {
1837 let mut nodes = Vec::new();
1838 let mut edges = Vec::new();
1839 for node_idx in 0..transform_count {
1840 let node_id = format!("transform:t{node_idx:04}");
1841 nodes.push(node(
1842 &node_id,
1843 NodeKind::Transform,
1844 vec![port("x", PortKind::Data)],
1845 vec![port("x", PortKind::Data)],
1846 ));
1847 if node_idx > 0 {
1848 edges.push(EdgeSpec {
1849 source: PortRef {
1850 node_id: NodeId::new(format!("transform:t{:04}", node_idx - 1)).unwrap(),
1851 port_name: "x".to_string(),
1852 },
1853 target: PortRef {
1854 node_id: NodeId::new(&node_id).unwrap(),
1855 port_name: "x".to_string(),
1856 },
1857 contract: EdgeContract::new(PortKind::Data, None),
1858 });
1859 }
1860 }
1861 nodes.push(node(
1862 "model:final",
1863 NodeKind::Model,
1864 vec![port("x", PortKind::Data)],
1865 vec![port("pred", PortKind::Prediction)],
1866 ));
1867 edges.push(EdgeSpec {
1868 source: PortRef {
1869 node_id: NodeId::new(format!("transform:t{:04}", transform_count - 1)).unwrap(),
1870 port_name: "x".to_string(),
1871 },
1872 target: PortRef {
1873 node_id: NodeId::new("model:final").unwrap(),
1874 port_name: "x".to_string(),
1875 },
1876 contract: EdgeContract::new(PortKind::Data, None),
1877 });
1878
1879 GraphSpec {
1880 id: "g:perf.linear".to_string(),
1881 interface: GraphInterface::default(),
1882 nodes,
1883 edges,
1884 search_space_fingerprint: None,
1885 metadata: BTreeMap::new(),
1886 }
1887 }
1888
1889 fn oof_graph() -> GraphSpec {
1890 GraphSpec {
1891 id: "g:oof.capabilities".to_string(),
1892 interface: GraphInterface::default(),
1893 nodes: vec![
1894 node(
1895 "model:base",
1896 NodeKind::Model,
1897 vec![],
1898 vec![port("pred", PortKind::Prediction)],
1899 ),
1900 node(
1901 "model:meta",
1902 NodeKind::Model,
1903 vec![port("pred", PortKind::Prediction)],
1904 vec![port("pred", PortKind::Prediction)],
1905 ),
1906 ],
1907 edges: vec![EdgeSpec {
1908 source: PortRef {
1909 node_id: NodeId::new("model:base").unwrap(),
1910 port_name: "pred".to_string(),
1911 },
1912 target: PortRef {
1913 node_id: NodeId::new("model:meta").unwrap(),
1914 port_name: "pred".to_string(),
1915 },
1916 contract: EdgeContract {
1917 requires_oof: true,
1918 requires_fold_alignment: true,
1919 ..EdgeContract::new(PortKind::Prediction, None)
1920 },
1921 }],
1922 search_space_fingerprint: None,
1923 metadata: BTreeMap::new(),
1924 }
1925 }
1926
1927 fn data_binding(node_id: &NodeId) -> DataBinding {
1928 DataBinding {
1929 node_id: node_id.clone(),
1930 input_name: "x".to_string(),
1931 request_id: "nir-to-tabular".to_string(),
1932 schema_fingerprint: "f97b37872fa22134b508f98fd8e207e5b776b52594fb8f6f5c3e15bee212246b"
1933 .to_string(),
1934 plan_fingerprint: "7c5431d85574b3f337022fa5d25971d5b5cf445b90331b49938f573ff6901e4d"
1935 .to_string(),
1936 relation_fingerprint: Some(
1937 "a3a7e329df35db9f2883a17b8611b7fae6dcaa031875e3ec2c9be1b9e29cbe10".to_string(),
1938 ),
1939 output_representation: "tabular_numeric".to_string(),
1940 feature_set_id: Some("x".to_string()),
1941 source_ids: vec!["nir".to_string()],
1942 require_relations: true,
1943 view_policy: Default::default(),
1944 metadata: BTreeMap::new(),
1945 }
1946 }
1947
1948 fn levels_as_strings(levels: &[Vec<NodeId>]) -> Vec<Vec<String>> {
1949 levels
1950 .iter()
1951 .map(|level| level.iter().map(ToString::to_string).collect())
1952 .collect()
1953 }
1954
1955 #[test]
1956 fn published_campaign_spec_schema_declares_current_contract() {
1957 let schema: serde_json::Value = serde_json::from_str(include_str!(
1958 "../../../docs/contracts/campaign_spec.schema.json"
1959 ))
1960 .unwrap();
1961
1962 assert_eq!(schema["$id"], CAMPAIGN_SPEC_SCHEMA_ID);
1963 assert!(schema["required"]
1964 .as_array()
1965 .unwrap()
1966 .iter()
1967 .any(|field| field.as_str() == Some("id")));
1968 assert!(schema["$defs"]["split_invocation"]["properties"]
1969 .as_object()
1970 .unwrap()
1971 .contains_key("fold_set"));
1972 assert!(schema["$defs"]["aggregation_policy"]["properties"]
1973 .as_object()
1974 .unwrap()
1975 .contains_key("selection_metric_level"));
1976 assert!(schema["$defs"]["aggregation_policy"]["properties"]
1977 .as_object()
1978 .unwrap()
1979 .contains_key("custom_controller"));
1980 assert!(schema["$defs"]["data_binding"]["properties"]
1981 .as_object()
1982 .unwrap()
1983 .contains_key("view_policy"));
1984 assert!(schema["properties"]
1985 .as_object()
1986 .unwrap()
1987 .contains_key("branch_view_plans"));
1988 assert!(schema["$defs"]["branch_view_plan"]["properties"]
1989 .as_object()
1990 .unwrap()
1991 .contains_key("selector"));
1992 }
1993
1994 #[test]
1995 fn published_execution_plan_schema_declares_current_contract() {
1996 let schema: serde_json::Value = serde_json::from_str(include_str!(
1997 "../../../docs/contracts/execution_plan.schema.json"
1998 ))
1999 .unwrap();
2000
2001 assert_eq!(schema["$id"], EXECUTION_PLAN_SCHEMA_ID);
2002 assert!(schema["required"]
2003 .as_array()
2004 .unwrap()
2005 .iter()
2006 .any(|field| field.as_str() == Some("node_plans")));
2007 assert!(schema["properties"]
2008 .as_object()
2009 .unwrap()
2010 .contains_key("controller_fingerprint"));
2011 assert!(schema["$defs"]["node_plan"]["properties"]
2012 .as_object()
2013 .unwrap()
2014 .contains_key("shape_plan"));
2015 assert!(schema["$defs"]["variant_plan"]["properties"]
2016 .as_object()
2017 .unwrap()
2018 .contains_key("choices"));
2019 }
2020
2021 #[test]
2022 fn published_execution_plan_fixture_validates_current_contract() {
2023 let plan: ExecutionPlan = serde_json::from_str(include_str!(
2024 "../../../examples/fixtures/runtime/execution_plan_branch_merge_executable.json"
2025 ))
2026 .unwrap();
2027
2028 plan.validate().unwrap();
2029 assert_eq!(plan.id, "plan:fixture.execution.branch_merge");
2030 assert_eq!(plan.variants.len(), 2);
2031 assert_eq!(plan.node_plans.len(), plan.graph_plan.graph.nodes.len());
2032 }
2033
2034 #[test]
2035 #[ignore = "perf sanity probe; run with --release --ignored --nocapture"]
2036 fn build_execution_plan_large_linear_graph_under_1500ms() {
2037 let started = Instant::now();
2038 let plan = build_execution_plan(
2039 "plan:perf.linear",
2040 large_linear_graph(400),
2041 campaign("campaign:perf.linear"),
2042 ®istry(),
2043 )
2044 .unwrap();
2045 let elapsed = started.elapsed();
2046
2047 assert_eq!(plan.graph_plan.topological_order.len(), 401);
2048 assert_eq!(plan.node_plans.len(), 401);
2049 assert!(
2050 elapsed <= Duration::from_millis(1_500),
2051 "large execution-plan build took {elapsed:?}"
2052 );
2053 }
2054
2055 #[test]
2056 fn builds_execution_plan_with_shape_and_fold_contracts() {
2057 let model_id = NodeId::new("model:pls").unwrap();
2058 let campaign = CampaignSpec {
2059 inner_cv: None,
2060 id: "campaign:oof".to_string(),
2061 root_seed: Some(7),
2062 leakage_policy: LeakageUnitPolicy::default(),
2063 aggregation_policy: AggregationPolicy::default(),
2064 split_invocation: Some(SplitInvocation {
2065 id: "split:outer".to_string(),
2066 controller_id: None,
2067 leakage_policy: LeakageUnitPolicy::default(),
2068 params: BTreeMap::new(),
2069 fold_set: Some(FoldSet {
2070 id: "outer".to_string(),
2071 sample_ids: vec![SampleId::new("s1").unwrap(), SampleId::new("s2").unwrap()],
2072 folds: vec![
2073 crate::fold::FoldAssignment {
2074 fold_id: FoldId::new("fold0").unwrap(),
2075 train_sample_ids: vec![SampleId::new("s2").unwrap()],
2076 validation_sample_ids: vec![SampleId::new("s1").unwrap()],
2077 metadata: BTreeMap::new(),
2078 },
2079 crate::fold::FoldAssignment {
2080 fold_id: FoldId::new("fold1").unwrap(),
2081 train_sample_ids: vec![SampleId::new("s1").unwrap()],
2082 validation_sample_ids: vec![SampleId::new("s2").unwrap()],
2083 metadata: BTreeMap::new(),
2084 },
2085 ],
2086 sample_groups: BTreeMap::new(),
2087 partition_mode: FoldPartitionMode::Partition,
2088 }),
2089 }),
2090 generation: Default::default(),
2091 shape_plans: BTreeMap::from([(
2092 model_id.clone(),
2093 DataModelShapePlan {
2094 node_id: model_id.clone(),
2095 input_granularity: Granularity::Observation,
2096 ..DataModelShapePlan {
2097 node_id: model_id.clone(),
2098 input_granularity: Granularity::Sample,
2099 target_granularity: Granularity::Sample,
2100 fit_rows: crate::policy::FitBoundary::FoldTrain,
2101 predict_rows: crate::policy::FitBoundary::FoldValidation,
2102 feature_namespace: None,
2103 feature_schema_fingerprint: None,
2104 target_space: "raw".to_string(),
2105 aggregation_policy: AggregationPolicy::default(),
2106 augmentation_policy: crate::policy::AugmentationPolicy::default(),
2107 selection_policy: crate::policy::FeatureSelectionPolicy::default(),
2108 }
2109 },
2110 )]),
2111 data_bindings: BTreeMap::from([(model_id.clone(), vec![data_binding(&model_id)])]),
2112 branch_view_plans: Vec::new(),
2113 metadata: BTreeMap::new(),
2114 };
2115
2116 let plan = build_execution_plan("plan:oof", graph(), campaign, ®istry()).unwrap();
2117
2118 assert_eq!(
2119 plan.graph_plan
2120 .topological_order
2121 .iter()
2122 .map(ToString::to_string)
2123 .collect::<Vec<_>>(),
2124 vec!["transform:snv", "model:pls"]
2125 );
2126 assert_eq!(
2127 levels_as_strings(&plan.graph_plan.parallel_levels),
2128 vec![vec!["transform:snv"], vec!["model:pls"]]
2129 );
2130 assert!(plan.node_plans[&model_id]
2131 .controller_capabilities
2132 .contains(&ControllerCapability::EmitsPredictions));
2133 assert!(plan.fold_set.is_some());
2134 let schedule = plan.campaign_phase_schedule(Phase::FitCv).unwrap();
2135 assert_eq!(schedule.scopes.len(), 2);
2136 assert!(schedule.scopes[0].scope_id.starts_with("scope:FIT_CV:"));
2137 assert!(schedule
2138 .scopes
2139 .iter()
2140 .all(|scope| levels_as_strings(&scope.node_levels)
2141 == vec![vec!["transform:snv"], vec!["model:pls"]]));
2142 assert_eq!(
2143 schedule
2144 .scopes
2145 .iter()
2146 .filter_map(|scope| scope.fold_id.as_ref().map(ToString::to_string))
2147 .collect::<Vec<_>>(),
2148 vec!["fold0", "fold1"]
2149 );
2150 assert_eq!(
2151 plan.node_plans
2152 .get(&model_id)
2153 .unwrap()
2154 .controller_id
2155 .as_str(),
2156 "controller:model"
2157 );
2158 assert_eq!(
2159 plan.node_plans.get(&model_id).unwrap().data_bindings.len(),
2160 1
2161 );
2162
2163 let mut bad_plan = plan.clone();
2164 bad_plan.graph_plan.parallel_levels =
2165 vec![vec![model_id], vec![NodeId::new("transform:snv").unwrap()]];
2166 assert!(bad_plan
2167 .validate()
2168 .unwrap_err()
2169 .to_string()
2170 .contains("parallel levels"));
2171
2172 let bad_envelope = ExternalDataPlanEnvelope {
2173 schema_version: crate::data::EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION,
2174 schema_fingerprint: "f97b37872fa22134b508f98fd8e207e5b776b52594fb8f6f5c3e15bee212246b"
2175 .to_string(),
2176 plan_fingerprint: "7c5431d85574b3f337022fa5d25971d5b5cf445b90331b49938f573ff6901e4d"
2177 .to_string(),
2178 relation_fingerprint: None,
2179 coordinator_relations: Some(SampleRelationSet {
2180 records: vec![{
2181 let mut relation = SampleRelation::new(
2182 ObservationId::new("obs:outside").unwrap(),
2183 SampleId::new("sample:outside").unwrap(),
2184 );
2185 relation.target_id = Some(TargetId::new("target:outside").unwrap());
2186 relation.source_id = Some("nir".to_string());
2187 relation
2188 }],
2189 }),
2190 };
2191 assert!(plan
2192 .campaign
2193 .validate_data_envelope_relations(&bad_envelope)
2194 .unwrap_err()
2195 .to_string()
2196 .contains("outside fold set"));
2197 }
2198
2199 #[test]
2200 fn planning_refuses_shape_plan_for_unknown_node() {
2201 let campaign = CampaignSpec {
2202 inner_cv: None,
2203 id: "campaign:oof".to_string(),
2204 root_seed: Some(7),
2205 leakage_policy: LeakageUnitPolicy::default(),
2206 aggregation_policy: AggregationPolicy::default(),
2207 split_invocation: None,
2208 generation: Default::default(),
2209 shape_plans: BTreeMap::from([(
2210 NodeId::new("model:missing").unwrap(),
2211 DataModelShapePlan {
2212 node_id: NodeId::new("model:missing").unwrap(),
2213 input_granularity: Granularity::Sample,
2214 target_granularity: Granularity::Sample,
2215 fit_rows: crate::policy::FitBoundary::FoldTrain,
2216 predict_rows: crate::policy::FitBoundary::FoldValidation,
2217 feature_namespace: None,
2218 feature_schema_fingerprint: None,
2219 target_space: "raw".to_string(),
2220 aggregation_policy: AggregationPolicy::default(),
2221 augmentation_policy: crate::policy::AugmentationPolicy::default(),
2222 selection_policy: crate::policy::FeatureSelectionPolicy::default(),
2223 },
2224 )]),
2225 data_bindings: BTreeMap::new(),
2226 branch_view_plans: Vec::new(),
2227 metadata: BTreeMap::new(),
2228 };
2229
2230 assert!(build_execution_plan("plan:oof", graph(), campaign, ®istry()).is_err());
2231 }
2232
2233 #[test]
2234 fn planning_refuses_oof_edge_without_controller_capabilities() {
2235 let mut registry = ControllerRegistry::new();
2236 let mut model_manifest = manifest("controller:model", NodeKind::Model);
2237 model_manifest
2238 .capabilities
2239 .remove(&ControllerCapability::ConsumesOofPredictions);
2240 registry.register(model_manifest).unwrap();
2241
2242 let err = build_execution_plan(
2243 "plan:oof.capability",
2244 oof_graph(),
2245 CampaignSpec {
2246 inner_cv: None,
2247 id: "campaign:oof.capability".to_string(),
2248 root_seed: Some(11),
2249 leakage_policy: Default::default(),
2250 aggregation_policy: Default::default(),
2251 split_invocation: None,
2252 generation: Default::default(),
2253 shape_plans: BTreeMap::new(),
2254 data_bindings: BTreeMap::new(),
2255 branch_view_plans: Vec::new(),
2256 metadata: BTreeMap::new(),
2257 },
2258 ®istry,
2259 )
2260 .unwrap_err();
2261
2262 assert!(err.to_string().contains("consumes_oof_predictions"));
2263 }
2264
2265 #[test]
2266 fn parallel_controller_capability_validation_requires_safe_manifest() {
2267 let mut registry = ControllerRegistry::new();
2268 let mut transform_manifest = manifest("controller:transform", NodeKind::Transform);
2269 transform_manifest
2270 .capabilities
2271 .remove(&ControllerCapability::ThreadSafe);
2272 transform_manifest
2273 .capabilities
2274 .remove(&ControllerCapability::ProcessSafe);
2275 registry.register(transform_manifest).unwrap();
2276 registry
2277 .register(manifest("controller:model", NodeKind::Model))
2278 .unwrap();
2279 let plan = build_execution_plan(
2280 "plan:parallel.capability",
2281 graph(),
2282 CampaignSpec {
2283 inner_cv: None,
2284 id: "campaign:parallel.capability".to_string(),
2285 root_seed: Some(11),
2286 leakage_policy: Default::default(),
2287 aggregation_policy: Default::default(),
2288 split_invocation: None,
2289 generation: Default::default(),
2290 shape_plans: BTreeMap::new(),
2291 data_bindings: BTreeMap::new(),
2292 branch_view_plans: Vec::new(),
2293 metadata: BTreeMap::new(),
2294 },
2295 ®istry,
2296 )
2297 .unwrap();
2298
2299 assert!(plan
2300 .validate_parallel_controller_capabilities(1, Phase::FitCv)
2301 .is_ok());
2302 let err = plan
2303 .validate_parallel_controller_capabilities(2, Phase::FitCv)
2304 .unwrap_err();
2305 assert!(err.to_string().contains("thread_safe or process_safe"));
2306 }
2307
2308 #[test]
2309 fn planning_refuses_generation_override_for_unknown_node() {
2310 let campaign = CampaignSpec {
2311 inner_cv: None,
2312 id: "campaign:oof".to_string(),
2313 root_seed: Some(7),
2314 leakage_policy: LeakageUnitPolicy::default(),
2315 aggregation_policy: AggregationPolicy::default(),
2316 split_invocation: None,
2317 generation: GenerationSpec {
2318 strategy: GenerationStrategy::Cartesian,
2319 dimensions: vec![GenerationDimension {
2320 name: "model_family".to_string(),
2321 choices: vec![GenerationChoice {
2322 label: "pls".to_string(),
2323 value: serde_json::json!("pls"),
2324 param_overrides: vec![GenerationParamOverride {
2325 node_id: NodeId::new("model:missing").unwrap(),
2326 params: BTreeMap::from([(
2327 "n_components".to_string(),
2328 serde_json::json!(8),
2329 )]),
2330 }],
2331 active_subsequence: None,
2332 }],
2333 }],
2334 max_variants: Some(1),
2335 constraints: GenerationConstraints::default(),
2336 },
2337 shape_plans: BTreeMap::new(),
2338 data_bindings: BTreeMap::new(),
2339 branch_view_plans: Vec::new(),
2340 metadata: BTreeMap::new(),
2341 };
2342
2343 let error = build_execution_plan("plan:oof", graph(), campaign, ®istry())
2344 .unwrap_err()
2345 .to_string();
2346
2347 assert!(error.contains("overrides params for unknown node"));
2348 }
2349
2350 #[test]
2351 fn planning_validates_declared_search_space_fingerprint() {
2352 let campaign = CampaignSpec {
2353 inner_cv: None,
2354 id: "campaign:search.fingerprint".to_string(),
2355 root_seed: Some(7),
2356 leakage_policy: LeakageUnitPolicy::default(),
2357 aggregation_policy: AggregationPolicy::default(),
2358 split_invocation: None,
2359 generation: GenerationSpec {
2360 strategy: GenerationStrategy::Cartesian,
2361 dimensions: vec![GenerationDimension {
2362 name: "model_family".to_string(),
2363 choices: vec![GenerationChoice {
2364 label: "pls".to_string(),
2365 value: serde_json::json!("pls"),
2366 param_overrides: vec![GenerationParamOverride {
2367 node_id: NodeId::new("model:pls").unwrap(),
2368 params: BTreeMap::from([(
2369 "n_components".to_string(),
2370 serde_json::json!(8),
2371 )]),
2372 }],
2373 active_subsequence: None,
2374 }],
2375 }],
2376 max_variants: Some(1),
2377 constraints: GenerationConstraints::default(),
2378 },
2379 shape_plans: BTreeMap::new(),
2380 data_bindings: BTreeMap::new(),
2381 branch_view_plans: Vec::new(),
2382 metadata: BTreeMap::new(),
2383 };
2384 let mut graph = graph();
2385 graph.search_space_fingerprint =
2386 Some(generation_spec_fingerprint(&campaign.generation).unwrap());
2387
2388 let plan = build_execution_plan(
2389 "plan:search.fingerprint",
2390 graph.clone(),
2391 campaign.clone(),
2392 ®istry(),
2393 )
2394 .unwrap();
2395 assert_eq!(plan.variants.len(), 1);
2396
2397 graph.search_space_fingerprint = Some("sha256:not-the-generation-spec".to_string());
2398 let error = build_execution_plan("plan:search.fingerprint", graph, campaign, ®istry())
2399 .unwrap_err()
2400 .to_string();
2401 assert!(error.contains("search_space_fingerprint"));
2402 }
2403
2404 #[test]
2405 fn branch_view_lookup_helpers_match_by_branch_id_and_innermost_path() {
2406 use crate::data::{BranchViewMode, DataViewSelector};
2407
2408 let outer = BranchViewPlan {
2409 view_id: "branch_view:outer".to_string(),
2410 branch_id: "branch:outer".to_string(),
2411 mode: BranchViewMode::BySource,
2412 selector: DataViewSelector {
2413 source_ids: vec!["nir".to_string()],
2414 ..Default::default()
2415 },
2416 allow_overlap: false,
2417 metadata: BTreeMap::new(),
2418 };
2419 let inner = BranchViewPlan {
2420 view_id: "branch_view:inner".to_string(),
2421 branch_id: "branch:inner".to_string(),
2422 mode: BranchViewMode::Separation,
2423 selector: DataViewSelector {
2424 source_ids: vec!["chem".to_string()],
2425 ..Default::default()
2426 },
2427 allow_overlap: false,
2428 metadata: BTreeMap::new(),
2429 };
2430 let plans = vec![outer.clone(), inner.clone()];
2431
2432 assert_eq!(
2433 super::branch_view_for_in(&plans, "branch:outer"),
2434 Some(&outer)
2435 );
2436 assert_eq!(
2437 super::branch_view_for_in(&plans, "branch:inner"),
2438 Some(&inner)
2439 );
2440 assert_eq!(super::branch_view_for_in(&plans, "branch:missing"), None);
2441
2442 let path = vec!["branch:outer".to_string(), "branch:inner".to_string()];
2443 assert_eq!(super::branch_view_for_path_in(&plans, &path), Some(&inner));
2445
2446 let path_outer_only = vec!["branch:outer".to_string()];
2447 assert_eq!(
2448 super::branch_view_for_path_in(&plans, &path_outer_only),
2449 Some(&outer)
2450 );
2451
2452 let empty_path: Vec<String> = Vec::new();
2453 assert_eq!(super::branch_view_for_path_in(&plans, &empty_path), None);
2454
2455 let path_no_match = vec!["branch:other".to_string()];
2456 assert_eq!(super::branch_view_for_path_in(&plans, &path_no_match), None);
2457 }
2458}