1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::campaign::stable_json_fingerprint;
6use crate::canonical::deserialize_external_contract;
7use crate::controller::{
8 ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest,
9 ControllerRegistry, RngPolicy,
10};
11use crate::controller_adapter::representation_type_id;
12use crate::criteria::{
13 CriterionInput, ImplementationCapability, LossCapability, SemanticSpecKind,
14 TrainingLossRoleReference,
15};
16use crate::data::{
17 BranchViewMode, BranchViewPlan, DataBinding, ExternalDataPlanEnvelope, ModelInputFusionMode,
18 ModelInputPortSpec, ModelInputSpec, RepresentationPlan, SOURCE_INDEX_METADATA_KEY,
19};
20use crate::error::{DagMlError, Result};
21use crate::fold::{FoldSet, NestedCvSpec};
22use crate::generation::{
23 enumerate_variants, generation_spec_fingerprint, GenerationSpec, VariantPlan,
24};
25use crate::graph::{GraphSpec, NodeKind, NodeSpec, PortKind};
26use crate::ids::{ControllerId, FoldId, NodeId, VariantId};
27use crate::phase::Phase;
28use crate::policy::{AggregationPolicy, DataModelShapePlan, LeakageUnitPolicy};
29
30pub const CAMPAIGN_SPEC_SCHEMA_VERSION: u32 = 1;
31pub const CAMPAIGN_SPEC_SCHEMA_ID: &str =
32 "https://github.com/GBeurier/dag-ml/schemas/campaign_spec.v1.schema.json";
33pub const EXECUTION_PLAN_SCHEMA_VERSION: u32 = 1;
34pub const EXECUTION_PLAN_SCHEMA_ID: &str =
35 "https://github.com/GBeurier/dag-ml/schemas/execution_plan.v1.schema.json";
36
37#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
38pub struct SplitInvocation {
39 pub id: String,
40 #[serde(default)]
41 pub controller_id: Option<ControllerId>,
42 #[serde(default)]
43 pub leakage_policy: LeakageUnitPolicy,
44 #[serde(default)]
45 pub params: BTreeMap<String, serde_json::Value>,
46 #[serde(default)]
47 pub fold_set: Option<FoldSet>,
48}
49
50impl SplitInvocation {
51 pub fn validate(&self) -> Result<()> {
52 if self.id.trim().is_empty() {
53 return Err(DagMlError::CampaignValidation(
54 "split invocation id is empty".to_string(),
55 ));
56 }
57 self.leakage_policy.validate()?;
58 if let Some(fold_set) = &self.fold_set {
59 fold_set.validate()?;
60 }
61 Ok(())
62 }
63}
64
65#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
66pub struct CampaignSpec {
67 pub id: String,
68 pub root_seed: Option<u64>,
69 #[serde(default)]
70 pub leakage_policy: LeakageUnitPolicy,
71 #[serde(default)]
72 pub aggregation_policy: AggregationPolicy,
73 #[serde(default)]
74 pub split_invocation: Option<SplitInvocation>,
75 #[serde(default)]
76 pub generation: GenerationSpec,
77 #[serde(default)]
78 pub shape_plans: BTreeMap<NodeId, DataModelShapePlan>,
79 #[serde(default)]
80 pub data_bindings: BTreeMap<NodeId, Vec<DataBinding>>,
81 #[serde(default, skip_serializing_if = "Vec::is_empty")]
82 pub branch_view_plans: Vec<BranchViewPlan>,
83 #[serde(default, skip_serializing_if = "Option::is_none")]
86 pub inner_cv: Option<NestedCvSpec>,
87 #[serde(default)]
88 pub metadata: BTreeMap<String, serde_json::Value>,
89}
90
91impl CampaignSpec {
92 pub fn from_json(json: &str) -> Result<Self> {
94 let campaign: Self =
95 deserialize_external_contract(json, "campaign", DagMlError::CampaignValidation)?;
96 campaign.validate()?;
97 Ok(campaign)
98 }
99
100 pub fn validate(&self) -> Result<()> {
101 if self.id.trim().is_empty() {
102 return Err(DagMlError::CampaignValidation(
103 "campaign id is empty".to_string(),
104 ));
105 }
106 self.leakage_policy.validate()?;
107 self.aggregation_policy.validate()?;
108 if let Some(inner_cv) = &self.inner_cv {
109 inner_cv.validate()?;
110 }
111 if let Some(split) = &self.split_invocation {
112 split.validate()?;
113 }
114 self.generation.validate()?;
115 for (node_id, shape_plan) in &self.shape_plans {
116 if node_id != &shape_plan.node_id {
117 return Err(DagMlError::CampaignValidation(format!(
118 "shape plan key `{node_id}` does not match node_id `{}`",
119 shape_plan.node_id
120 )));
121 }
122 shape_plan.validate()?;
123 }
124 for (node_id, bindings) in &self.data_bindings {
125 for binding in bindings {
126 if node_id != &binding.node_id {
127 return Err(DagMlError::CampaignValidation(format!(
128 "data binding key `{node_id}` does not match node_id `{}`",
129 binding.node_id
130 )));
131 }
132 binding.validate()?;
133 }
134 }
135 let mut branch_views = BTreeSet::new();
136 for plan in &self.branch_view_plans {
137 plan.validate()?;
138 if !branch_views.insert(plan.view_id.as_str()) {
139 return Err(DagMlError::CampaignValidation(format!(
140 "campaign `{}` contains duplicate branch view `{}`",
141 self.id, plan.view_id
142 )));
143 }
144 }
145 Ok(())
146 }
147
148 pub fn validate_data_envelope_relations(
149 &self,
150 envelope: &ExternalDataPlanEnvelope,
151 ) -> Result<()> {
152 envelope.validate()?;
153 let Some(relations) = &envelope.coordinator_relations else {
154 return Ok(());
155 };
156 let Some(split) = &self.split_invocation else {
157 return Ok(());
158 };
159 let Some(fold_set) = &split.fold_set else {
160 return Ok(());
161 };
162 relations.validate_against_fold_set(fold_set, &self.leakage_policy)?;
163 relations.validate_against_fold_set(fold_set, &split.leakage_policy)
164 }
165}
166
167#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
168pub struct GraphPlan {
169 pub graph: GraphSpec,
170 pub topological_order: Vec<NodeId>,
171 #[serde(default, skip_serializing_if = "Vec::is_empty")]
172 pub parallel_levels: Vec<Vec<NodeId>>,
173}
174
175impl GraphPlan {
176 pub fn from_graph(graph: GraphSpec) -> Result<Self> {
177 let topological_order = graph.topological_order()?;
178 let parallel_levels = graph.parallel_levels()?;
179 Ok(Self {
180 graph,
181 topological_order,
182 parallel_levels,
183 })
184 }
185
186 pub fn parallel_levels(&self) -> Result<Vec<Vec<NodeId>>> {
187 if self.parallel_levels.is_empty() {
188 return self.graph.parallel_levels();
189 }
190 Ok(self.parallel_levels.clone())
191 }
192}
193
194#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
195pub struct NodePlan {
196 pub node_id: NodeId,
197 pub kind: NodeKind,
198 pub controller_id: ControllerId,
199 pub controller_version: String,
200 pub supported_phases: BTreeSet<Phase>,
201 #[serde(default)]
202 pub controller_capabilities: BTreeSet<ControllerCapability>,
203 #[serde(default, skip_serializing_if = "Vec::is_empty")]
204 pub training_losses: Vec<TrainingLossRoleReference>,
205 pub fit_scope: ControllerFitScope,
206 pub rng_policy: RngPolicy,
207 pub artifact_policy: ArtifactPolicy,
208 pub input_nodes: Vec<NodeId>,
209 pub output_nodes: Vec<NodeId>,
210 pub shape_plan: Option<DataModelShapePlan>,
211 #[serde(default)]
212 pub data_bindings: Vec<DataBinding>,
213 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
214 pub params: BTreeMap<String, serde_json::Value>,
215 #[serde(default, skip_serializing_if = "Option::is_none")]
218 pub inner_cv: Option<NestedCvSpec>,
219 pub params_fingerprint: String,
220}
221
222impl NodePlan {
223 pub fn training_losses_for_phase(
224 &self,
225 phase: Phase,
226 ) -> impl Iterator<Item = &TrainingLossRoleReference> {
227 self.training_losses
228 .iter()
229 .filter(move |role| role.phases.contains(&phase))
230 }
231
232 pub fn training_loss_fingerprint(&self, phase: Phase) -> Result<Option<String>> {
233 let roles = self.training_losses_for_phase(phase).collect::<Vec<_>>();
234 if roles.is_empty() {
235 Ok(None)
236 } else {
237 stable_json_fingerprint(&roles).map(Some)
238 }
239 }
240}
241
242#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
243pub struct ExecutionPlan {
244 pub id: String,
245 pub graph_plan: GraphPlan,
246 pub campaign: CampaignSpec,
247 pub node_plans: BTreeMap<NodeId, NodePlan>,
248 pub controller_manifests: BTreeMap<ControllerId, ControllerManifest>,
249 pub variants: Vec<VariantPlan>,
250 pub fold_set: Option<FoldSet>,
251 pub graph_fingerprint: String,
252 pub campaign_fingerprint: String,
253 pub controller_fingerprint: String,
254}
255
256#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
257pub struct ExecutionScopePlan {
258 pub scope_id: String,
259 pub phase: Phase,
260 pub variant_id: Option<VariantId>,
261 pub variant_seed: Option<u64>,
262 pub fold_id: Option<FoldId>,
263 pub node_levels: Vec<Vec<NodeId>>,
264}
265
266#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
267pub struct PhaseExecutionSchedule {
268 pub plan_id: String,
269 pub phase: Phase,
270 pub scopes: Vec<ExecutionScopePlan>,
271}
272
273impl ExecutionPlan {
274 pub fn from_json(json: &str) -> Result<Self> {
284 let plan: Self =
285 deserialize_external_contract(json, "execution plan", DagMlError::Planning)?;
286 plan.validate()?;
287 Ok(plan)
288 }
289
290 pub fn with_training_losses(mut self, roles: Vec<TrainingLossRoleReference>) -> Result<Self> {
296 self.validate()?;
297 let mut roles_by_node = BTreeMap::<NodeId, Vec<TrainingLossRoleReference>>::new();
298 for role in roles {
299 role.validate()?;
300 if !self.node_plans.contains_key(&role.node_id) {
301 return Err(DagMlError::Planning(format!(
302 "training loss references unknown plan node `{}`",
303 role.node_id
304 )));
305 }
306 roles_by_node
307 .entry(role.node_id.clone())
308 .or_default()
309 .push(role);
310 }
311 for node_roles in roles_by_node.values_mut() {
312 node_roles.sort_by(|left, right| {
313 (&left.output_id, &left.phases).cmp(&(&right.output_id, &right.phases))
314 });
315 }
316 for (node_id, node_plan) in &mut self.node_plans {
317 node_plan.training_losses = roles_by_node.remove(node_id).unwrap_or_default();
318 }
319 self.validate()?;
320 Ok(self)
321 }
322
323 pub fn validate(&self) -> Result<()> {
324 self.graph_plan.graph.validate()?;
325 self.campaign.validate()?;
326 if !self.graph_plan.parallel_levels.is_empty()
329 && self.graph_plan.parallel_levels != self.graph_plan.graph.parallel_levels()?
330 {
331 return Err(DagMlError::Planning(
332 "graph plan parallel levels do not match graph".to_string(),
333 ));
334 }
335
336 for (controller_id, manifest) in &self.controller_manifests {
340 manifest.validate()?;
341 if controller_id != &manifest.controller_id {
342 return Err(DagMlError::Planning(format!(
343 "controller manifest keyed `{controller_id}` declares id `{}`",
344 manifest.controller_id
345 )));
346 }
347 }
348
349 let recomputed_graph_fingerprint = stable_json_fingerprint(&self.graph_plan.graph)?;
360 if recomputed_graph_fingerprint != self.graph_fingerprint {
361 return Err(DagMlError::Planning(
362 "execution plan graph_fingerprint does not match the embedded graph".to_string(),
363 ));
364 }
365 let recomputed_campaign_fingerprint = stable_json_fingerprint(&self.campaign)?;
366 if recomputed_campaign_fingerprint != self.campaign_fingerprint {
367 return Err(DagMlError::Planning(
368 "execution plan campaign_fingerprint does not match the embedded campaign"
369 .to_string(),
370 ));
371 }
372 let recomputed_controller_fingerprint =
373 stable_json_fingerprint(&self.controller_manifests)?;
374 if recomputed_controller_fingerprint != self.controller_fingerprint {
375 return Err(DagMlError::Planning(
376 "execution plan controller_fingerprint does not match the embedded controller manifests"
377 .to_string(),
378 ));
379 }
380
381 let graph_node_ids = self
386 .graph_plan
387 .graph
388 .nodes
389 .iter()
390 .map(|node| node.id.clone())
391 .collect::<BTreeSet<_>>();
392 for (node_id, plan) in &self.node_plans {
393 if node_id != &plan.node_id {
394 return Err(DagMlError::Planning(format!(
395 "node plan keyed `{node_id}` declares node_id `{}`",
396 plan.node_id
397 )));
398 }
399 }
400 let plan_node_ids = self.node_plans.keys().cloned().collect::<BTreeSet<_>>();
401 if plan_node_ids != graph_node_ids {
402 return Err(DagMlError::Planning(
403 "execution plan node_plans do not exactly cover the graph node-id set".to_string(),
404 ));
405 }
406
407 if self.graph_plan.topological_order != self.graph_plan.graph.topological_order()? {
412 return Err(DagMlError::Planning(
413 "execution plan topological_order does not match the graph".to_string(),
414 ));
415 }
416
417 let graph_nodes_by_id = self
418 .graph_plan
419 .graph
420 .nodes
421 .iter()
422 .map(|node| (node.id.clone(), node))
423 .collect::<BTreeMap<_, _>>();
424 for (node_id, plan) in &self.node_plans {
425 let graph_node = graph_nodes_by_id
426 .get(node_id)
427 .expect("node_plans keys equal graph node ids");
428 if plan.kind != graph_node.kind {
433 return Err(DagMlError::Planning(format!(
434 "node plan `{node_id}` kind does not match graph node kind"
435 )));
436 }
437 if plan.input_nodes != self.graph_plan.graph.upstream_nodes(node_id)
438 || plan.output_nodes != self.graph_plan.graph.downstream_nodes(node_id)
439 {
440 return Err(DagMlError::Planning(format!(
441 "node plan `{node_id}` input/output adjacency does not match the graph"
442 )));
443 }
444 let manifest = self
445 .controller_manifests
446 .get(&plan.controller_id)
447 .ok_or_else(|| {
448 DagMlError::Planning(format!(
449 "missing controller manifest `{}` for node `{node_id}`",
450 plan.controller_id
451 ))
452 })?;
453 if manifest.operator_kind != plan.kind
458 || manifest.controller_version != plan.controller_version
459 || manifest.supported_phases != plan.supported_phases
460 || manifest.capabilities != plan.controller_capabilities
461 || manifest.fit_scope != plan.fit_scope
462 || manifest.rng_policy != plan.rng_policy
463 || manifest.artifact_policy != plan.artifact_policy
464 {
465 return Err(DagMlError::Planning(format!(
466 "node `{node_id}` node plan does not match controller manifest `{}`",
467 manifest.controller_id
468 )));
469 }
470 for binding in &plan.data_bindings {
471 if binding.node_id != *node_id {
472 return Err(DagMlError::Planning(format!(
473 "node plan `{node_id}` contains data binding for `{}`",
474 binding.node_id
475 )));
476 }
477 binding.validate()?;
478 }
479 validate_data_binding_requirements(node_id, plan, manifest, graph_node)?;
480 validate_node_training_losses(plan)?;
481 let actual_params_fingerprint = stable_json_fingerprint(&plan.params)?;
482 if actual_params_fingerprint != plan.params_fingerprint {
483 return Err(DagMlError::Planning(format!(
484 "node plan `{node_id}` params fingerprint does not match params"
485 )));
486 }
487 if let Some(inner_cv) = &plan.inner_cv {
491 inner_cv.validate().map_err(|error| {
492 DagMlError::Planning(format!(
493 "node plan `{node_id}` has invalid inner_cv: {error}"
494 ))
495 })?;
496 }
497 }
498 self.validate_oof_controller_capabilities()?;
499 if let Some(fold_set) = &self.fold_set {
500 fold_set.validate()?;
501 }
502 if self.variants.is_empty() {
503 return Err(DagMlError::Planning(
504 "execution plan has no variants".to_string(),
505 ));
506 }
507 for variant in &self.variants {
508 variant.validate()?;
509 }
510 Ok(())
511 }
512
513 pub fn validate_parallel_controller_capabilities(
514 &self,
515 max_workers: usize,
516 phase: Phase,
517 ) -> Result<()> {
518 if max_workers <= 1 {
519 return Ok(());
520 }
521 let node_ids = self
522 .node_parallel_levels_for_phase(phase)?
523 .into_iter()
524 .flatten()
525 .collect::<Vec<_>>();
526 for node_id in node_ids {
527 let node_plan = self.node_plans.get(&node_id).ok_or_else(|| {
528 DagMlError::Planning(format!("missing node plan for `{node_id}`"))
529 })?;
530 let manifest = self
531 .controller_manifests
532 .get(&node_plan.controller_id)
533 .ok_or_else(|| {
534 DagMlError::Planning(format!(
535 "missing controller manifest `{}` for node `{}`",
536 node_plan.controller_id, node_plan.node_id
537 ))
538 })?;
539 if !manifest.supports_parallel_invocation() {
540 return Err(DagMlError::Planning(format!(
541 "parallel scheduler with {max_workers} workers requires controller `{}` for node `{}` to declare thread_safe or process_safe",
542 manifest.controller_id, node_plan.node_id
543 )));
544 }
545 }
546 Ok(())
547 }
548
549 fn validate_oof_controller_capabilities(&self) -> Result<()> {
550 for edge in &self.graph_plan.graph.edges {
551 if edge.contract.kind != PortKind::Prediction {
552 continue;
553 }
554 let target_plan = self.node_plans.get(&edge.target.node_id).ok_or_else(|| {
555 DagMlError::Planning(format!(
556 "prediction edge target node `{}` has no node plan",
557 edge.target.node_id
558 ))
559 })?;
560 let target_fits = matches!(
561 target_plan.fit_scope,
562 ControllerFitScope::FoldTrain | ControllerFitScope::FullTrain
563 ) && target_plan
564 .supported_phases
565 .iter()
566 .any(|phase| matches!(phase, Phase::FitCv | Phase::Refit));
567 if target_fits && !edge.contract.requires_oof {
568 return Err(DagMlError::Planning(format!(
569 "prediction edge `{}.{}` -> `{}.{}` enters fitting controller `{}` and must require OOF",
570 edge.source.node_id,
571 edge.source.port_name,
572 edge.target.node_id,
573 edge.target.port_name,
574 target_plan.controller_id
575 )));
576 }
577 if !edge.contract.requires_oof {
578 continue;
579 }
580 let source_plan = self.node_plans.get(&edge.source.node_id).ok_or_else(|| {
581 DagMlError::Planning(format!(
582 "OOF edge source node `{}` has no node plan",
583 edge.source.node_id
584 ))
585 })?;
586 if !source_plan
587 .controller_capabilities
588 .contains(&ControllerCapability::EmitsPredictions)
589 {
590 return Err(DagMlError::Planning(format!(
591 "OOF edge `{}.{}` -> `{}.{}` requires source controller `{}` to declare emits_predictions",
592 edge.source.node_id,
593 edge.source.port_name,
594 edge.target.node_id,
595 edge.target.port_name,
596 source_plan.controller_id
597 )));
598 }
599 if !target_plan
600 .controller_capabilities
601 .contains(&ControllerCapability::ConsumesOofPredictions)
602 {
603 return Err(DagMlError::Planning(format!(
604 "OOF edge `{}.{}` -> `{}.{}` requires target controller `{}` to declare consumes_oof_predictions",
605 edge.source.node_id,
606 edge.source.port_name,
607 edge.target.node_id,
608 edge.target.port_name,
609 target_plan.controller_id
610 )));
611 }
612 }
613 Ok(())
614 }
615
616 pub fn node_parallel_levels_for_phase(&self, phase: Phase) -> Result<Vec<Vec<NodeId>>> {
617 let levels = self
618 .graph_plan
619 .parallel_levels()?
620 .into_iter()
621 .map(|level| {
622 level
623 .into_iter()
624 .filter(|node_id| {
625 self.node_plans
626 .get(node_id)
627 .is_some_and(|node_plan| node_plan.supported_phases.contains(&phase))
628 })
629 .collect::<Vec<_>>()
630 })
631 .filter(|level| !level.is_empty())
632 .collect::<Vec<_>>();
633 Ok(levels)
634 }
635
636 pub fn campaign_phase_schedule(&self, phase: Phase) -> Result<PhaseExecutionSchedule> {
637 self.validate()?;
638 let node_levels = self.node_parallel_levels_for_phase(phase)?;
639 let fold_ids = if phase == Phase::FitCv {
640 self.fold_set
641 .as_ref()
642 .map(|fold_set| {
643 fold_set
644 .folds
645 .iter()
646 .map(|fold| Some(fold.fold_id.clone()))
647 .collect::<Vec<_>>()
648 })
649 .unwrap_or_else(|| vec![None])
650 } else {
651 vec![None]
652 };
653 let mut scopes = Vec::new();
654 for variant in &self.variants {
655 for fold_id in &fold_ids {
656 scopes.push(ExecutionScopePlan {
657 scope_id: execution_scope_id(
658 phase,
659 Some(&variant.variant_id),
660 fold_id.as_ref(),
661 ),
662 phase,
663 variant_id: Some(variant.variant_id.clone()),
664 variant_seed: variant.seed,
665 fold_id: fold_id.clone(),
666 node_levels: node_levels.clone(),
667 });
668 }
669 }
670 Ok(PhaseExecutionSchedule {
671 plan_id: self.id.clone(),
672 phase,
673 scopes,
674 })
675 }
676
677 pub fn branch_view_for(&self, branch_id: &str) -> Option<&BranchViewPlan> {
681 branch_view_for_in(&self.campaign.branch_view_plans, branch_id)
682 }
683
684 pub fn branch_view_for_path(&self, branch_path: &[String]) -> Option<&BranchViewPlan> {
691 branch_view_for_path_in(&self.campaign.branch_view_plans, branch_path)
692 }
693}
694
695fn validate_data_binding_requirements(
696 node_id: &NodeId,
697 plan: &NodePlan,
698 manifest: &ControllerManifest,
699 node: &NodeSpec,
700) -> Result<()> {
701 let branch_view = branch_view_plan_from_node_metadata(node)?;
702 let Some(model_input) = manifest.model_input_spec()? else {
703 for binding in &plan.data_bindings {
704 let effective_source_ids = effective_binding_source_ids(binding, branch_view.as_ref())?;
705 if effective_source_ids.len() > 1 {
706 return Err(data_requirement_refusal(
707 "dagml.data_requirement.missing_data_requirements",
708 node_id,
709 binding,
710 manifest,
711 "multisource",
712 &effective_source_ids,
713 "multi-source data binding requires controller data_requirements".to_string(),
714 ));
715 }
716 }
717 return Ok(());
718 };
719 for binding in &plan.data_bindings {
720 let Some(port) = model_input
721 .ports
722 .iter()
723 .find(|port| port.name == binding.input_name)
724 else {
725 return Err(DagMlError::Planning(format!(
726 "node `{node_id}` data binding `{}` is not declared by controller `{}` data_requirements",
727 binding.input_name, manifest.controller_id
728 )));
729 };
730 if !port
731 .accepted_representations
732 .iter()
733 .any(|representation| representation == &binding.output_representation)
734 {
735 return Err(DagMlError::Planning(format!(
736 "node `{node_id}` data binding `{}` output representation `{}` is not accepted by controller `{}` data_requirements port `{}`",
737 binding.input_name,
738 binding.output_representation,
739 manifest.controller_id,
740 port.name
741 )));
742 }
743 if let Some(type_id) = representation_type_id(&binding.output_representation) {
744 if !port
745 .accepted_types
746 .iter()
747 .any(|accepted_type| accepted_type.as_str() == type_id)
748 {
749 return Err(DagMlError::Planning(format!(
750 "node `{node_id}` data binding `{}` output representation `{}` has registered type `{type_id}` but controller `{}` data_requirements port `{}` accepts types {:?}",
751 binding.input_name,
752 binding.output_representation,
753 manifest.controller_id,
754 port.name,
755 port.accepted_types
756 )));
757 }
758 }
759 validate_data_binding_source_shape(
760 node_id,
761 binding,
762 port,
763 &model_input,
764 manifest,
765 branch_view.as_ref(),
766 )?;
767 }
768 Ok(())
769}
770
771fn validate_node_training_losses(plan: &NodePlan) -> Result<()> {
772 let mut previous_key: Option<(Option<String>, BTreeSet<Phase>)> = None;
773 let mut occupied_phases = BTreeSet::new();
774 for role in &plan.training_losses {
775 role.validate()?;
776 if role.node_id != plan.node_id {
777 return Err(DagMlError::Planning(format!(
778 "node plan `{}` contains training loss for `{}`",
779 plan.node_id, role.node_id
780 )));
781 }
782 let key = (role.output_id.clone(), role.phases.clone());
783 if previous_key
784 .as_ref()
785 .is_some_and(|previous| previous >= &key)
786 {
787 return Err(DagMlError::Planning(format!(
788 "node plan `{}` training losses must be strictly sorted by output_id and phases",
789 plan.node_id
790 )));
791 }
792 previous_key = Some(key);
793 for phase in &role.phases {
794 if !plan.supported_phases.contains(phase) {
795 return Err(DagMlError::Planning(format!(
796 "node `{}` has a training loss for unsupported phase {phase:?}",
797 plan.node_id
798 )));
799 }
800 if !occupied_phases.insert((role.output_id.clone(), *phase)) {
801 return Err(DagMlError::Planning(format!(
802 "node `{}` has overlapping training losses for output {:?} in phase {phase:?}",
803 plan.node_id, role.output_id
804 )));
805 }
806 }
807 if !plan
808 .controller_capabilities
809 .contains(&ControllerCapability::SupportsConfigurableLoss)
810 {
811 return Err(DagMlError::Planning(format!(
812 "node `{}` configures a training loss but its controller does not support configurable loss",
813 plan.node_id
814 )));
815 }
816 if role.loss.spec.kind == SemanticSpecKind::Custom
817 && !plan
818 .controller_capabilities
819 .contains(&ControllerCapability::SupportsCustomLoss)
820 {
821 return Err(DagMlError::Planning(format!(
822 "node `{}` configures a custom loss but its controller does not support custom loss",
823 plan.node_id
824 )));
825 }
826 if role
827 .loss
828 .spec
829 .capabilities
830 .contains(&LossCapability::Differentiable)
831 && !plan
832 .controller_capabilities
833 .contains(&ControllerCapability::SupportsDifferentiableLoss)
834 {
835 return Err(DagMlError::Planning(format!(
836 "node `{}` configures a differentiable loss but its controller does not support differentiable loss",
837 plan.node_id
838 )));
839 }
840 if role
841 .loss
842 .spec
843 .required_inputs
844 .contains(&CriterionInput::SampleWeight)
845 && !plan
846 .controller_capabilities
847 .contains(&ControllerCapability::SupportsSampleWeights)
848 {
849 return Err(DagMlError::Planning(format!(
850 "node `{}` loss requires sample weights but its controller does not support them",
851 plan.node_id
852 )));
853 }
854 if role
855 .loss
856 .spec
857 .required_inputs
858 .contains(&CriterionInput::MissingMask)
859 && !plan
860 .controller_capabilities
861 .contains(&ControllerCapability::SupportsMissingMasks)
862 {
863 return Err(DagMlError::Planning(format!(
864 "node `{}` loss requires missing masks but its controller does not support them",
865 plan.node_id
866 )));
867 }
868 if role
869 .loss
870 .implementation
871 .capabilities
872 .contains(&ImplementationCapability::NeedsGil)
873 && !plan
874 .controller_capabilities
875 .contains(&ControllerCapability::NeedsPythonGil)
876 {
877 return Err(DagMlError::Planning(format!(
878 "node `{}` loss implementation needs the Python GIL but its controller does not declare it",
879 plan.node_id
880 )));
881 }
882 }
883 Ok(())
884}
885
886fn branch_view_plan_from_node_metadata(node: &NodeSpec) -> Result<Option<BranchViewPlan>> {
887 let Some(value) = node.metadata.get("dsl_branch_view_plan") else {
888 return Ok(None);
889 };
890 let plan: BranchViewPlan = serde_json::from_value(value.clone()).map_err(|error| {
891 DagMlError::Planning(format!(
892 "node `{}` carries malformed `dsl_branch_view_plan` metadata: {error}",
893 node.id
894 ))
895 })?;
896 plan.validate()
897 .map_err(|error| DagMlError::Planning(error.to_string()))?;
898 Ok(Some(plan))
899}
900
901fn validate_data_binding_source_shape(
902 node_id: &NodeId,
903 binding: &DataBinding,
904 port: &ModelInputPortSpec,
905 model_input: &ModelInputSpec,
906 manifest: &ControllerManifest,
907 branch_view: Option<&BranchViewPlan>,
908) -> Result<()> {
909 let effective_source_ids = effective_binding_source_ids(binding, branch_view)?;
910 if effective_source_ids.len() < 2 {
911 return Ok(());
912 }
913 if !port.multi_source {
914 return Err(data_requirement_refusal(
915 "dagml.data_requirement.multi_source_port_not_supported",
916 node_id,
917 binding,
918 manifest,
919 "multisource",
920 &effective_source_ids,
921 format!(
922 "controller `{}` data_requirements port `{}` does not declare multi_source=true",
923 manifest.controller_id, port.name
924 ),
925 ));
926 }
927 let Some(fusion) = &model_input.default_fusion else {
928 return Err(data_requirement_refusal(
929 "dagml.data_requirement.missing_multisource_fusion",
930 node_id,
931 binding,
932 manifest,
933 "multisource",
934 &effective_source_ids,
935 "multi-source data binding requires an explicit default_fusion policy".to_string(),
936 ));
937 };
938 validate_fusion_sources_match_binding(
939 node_id,
940 binding,
941 manifest,
942 fusion.representation_plan.as_ref(),
943 &effective_source_ids,
944 )?;
945 match fusion.mode {
946 ModelInputFusionMode::ConcatenateFeatures => {
947 if binding
948 .metadata
949 .get(SOURCE_INDEX_METADATA_KEY)
950 .and_then(serde_json::Value::as_object)
951 .is_none()
952 {
953 return Err(data_requirement_refusal(
954 "dagml.data_requirement.source_concat_requires_source_index",
955 node_id,
956 binding,
957 manifest,
958 "source_concat",
959 &effective_source_ids,
960 "source-concat feature fusion requires data binding metadata.source_index so feature-axis blocks are explicit".to_string(),
961 ));
962 }
963 }
964 ModelInputFusionMode::DictBySource | ModelInputFusionMode::Custom => {}
965 ModelInputFusionMode::SingleSource | ModelInputFusionMode::StackSamples => {
966 let fusion_mode = fusion_mode_label(fusion.mode);
967 return Err(data_requirement_refusal(
968 "dagml.data_requirement.unsupported_multisource_fusion_mode",
969 node_id,
970 binding,
971 manifest,
972 fusion_mode,
973 &effective_source_ids,
974 format!(
975 "multi-source data binding cannot be planned with default_fusion.mode={fusion_mode}"
976 ),
977 ));
978 }
979 }
980 Ok(())
981}
982
983fn fusion_mode_label(mode: ModelInputFusionMode) -> &'static str {
984 match mode {
985 ModelInputFusionMode::SingleSource => "single_source",
986 ModelInputFusionMode::ConcatenateFeatures => "concatenate_features",
987 ModelInputFusionMode::StackSamples => "stack_samples",
988 ModelInputFusionMode::DictBySource => "dict_by_source",
989 ModelInputFusionMode::Custom => "custom",
990 }
991}
992
993fn effective_binding_source_ids(
994 binding: &DataBinding,
995 branch_view: Option<&BranchViewPlan>,
996) -> Result<Vec<String>> {
997 let Some(branch_view) = branch_view else {
998 return Ok(binding.source_ids.clone());
999 };
1000 if branch_view.mode != BranchViewMode::BySource {
1001 return Ok(binding.source_ids.clone());
1002 }
1003 if branch_view.selector.source_ids.len() != 1 {
1004 return Err(data_requirement_refusal_for_branch(
1005 "dagml.data_requirement.unsupported_by_source_shape",
1006 binding,
1007 branch_view,
1008 "by_source branch views must select exactly one source_id for per-source X-chain fit semantics".to_string(),
1009 ));
1010 }
1011 if !binding.source_ids.is_empty() {
1012 let declared = binding.source_ids.iter().collect::<BTreeSet<_>>();
1013 for source_id in &branch_view.selector.source_ids {
1014 if !declared.contains(source_id) {
1015 return Err(data_requirement_refusal_for_branch(
1016 "dagml.data_requirement.by_source_selector_outside_binding",
1017 binding,
1018 branch_view,
1019 format!(
1020 "by_source branch selector source `{source_id}` is not declared by data binding source_ids"
1021 ),
1022 ));
1023 }
1024 }
1025 }
1026 Ok(branch_view.selector.source_ids.clone())
1027}
1028
1029fn validate_fusion_sources_match_binding(
1030 node_id: &NodeId,
1031 binding: &DataBinding,
1032 manifest: &ControllerManifest,
1033 representation_plan: Option<&RepresentationPlan>,
1034 effective_source_ids: &[String],
1035) -> Result<()> {
1036 let Some(representation_plan) = representation_plan else {
1037 return Ok(());
1038 };
1039 let component_sources = representation_plan_component_sources(representation_plan);
1040 if component_sources.is_empty() {
1041 return Ok(());
1042 }
1043 let declared = component_sources.iter().cloned().collect::<BTreeSet<_>>();
1044 let effective = effective_source_ids.iter().collect::<BTreeSet<_>>();
1045 if declared != effective {
1046 return Err(data_requirement_refusal(
1047 "dagml.data_requirement.representation_sources_mismatch",
1048 node_id,
1049 binding,
1050 manifest,
1051 "multisource",
1052 effective_source_ids,
1053 format!(
1054 "default_fusion.representation_plan component_source_ids {:?} do not match binding source_ids {:?}",
1055 component_sources, effective_source_ids
1056 ),
1057 ));
1058 }
1059 Ok(())
1060}
1061
1062fn representation_plan_component_sources(plan: &RepresentationPlan) -> Vec<&String> {
1063 match plan {
1064 RepresentationPlan::Aggregate(_) => Vec::new(),
1065 RepresentationPlan::CartesianProduct(plan) => {
1066 plan.combination_plan.component_source_ids.iter().collect()
1067 }
1068 RepresentationPlan::MonteCarloCartesian(plan) => {
1069 plan.combination_plan.component_source_ids.iter().collect()
1070 }
1071 RepresentationPlan::StackFixed(plan) => plan.component_source_ids.iter().collect(),
1072 RepresentationPlan::StackPaddedMasked(plan) => plan.component_source_ids.iter().collect(),
1073 }
1074}
1075
1076fn data_requirement_refusal(
1077 code: &'static str,
1078 node_id: &NodeId,
1079 binding: &DataBinding,
1080 manifest: &ControllerManifest,
1081 shape: &str,
1082 source_ids: &[String],
1083 message: String,
1084) -> DagMlError {
1085 DagMlError::Planning(format!(
1086 "data requirement refusal: {}",
1087 serde_json::json!({
1088 "schema_version": 1,
1089 "code": code,
1090 "node_id": node_id.to_string(),
1091 "input_name": binding.input_name.as_str(),
1092 "controller_id": manifest.controller_id.to_string(),
1093 "shape": shape,
1094 "source_ids": source_ids,
1095 "message": message
1096 })
1097 ))
1098}
1099
1100fn data_requirement_refusal_for_branch(
1101 code: &'static str,
1102 binding: &DataBinding,
1103 branch_view: &BranchViewPlan,
1104 message: String,
1105) -> DagMlError {
1106 DagMlError::Planning(format!(
1107 "data requirement refusal: {}",
1108 serde_json::json!({
1109 "schema_version": 1,
1110 "code": code,
1111 "node_id": binding.node_id.to_string(),
1112 "input_name": binding.input_name.as_str(),
1113 "branch_view_id": branch_view.view_id.as_str(),
1114 "branch_id": branch_view.branch_id.as_str(),
1115 "shape": "by_source",
1116 "source_ids": &branch_view.selector.source_ids,
1117 "message": message
1118 })
1119 ))
1120}
1121
1122fn branch_view_for_in<'a>(
1123 plans: &'a [BranchViewPlan],
1124 branch_id: &str,
1125) -> Option<&'a BranchViewPlan> {
1126 plans.iter().find(|plan| plan.branch_id == branch_id)
1127}
1128
1129fn branch_view_for_path_in<'a>(
1130 plans: &'a [BranchViewPlan],
1131 branch_path: &[String],
1132) -> Option<&'a BranchViewPlan> {
1133 for branch_id in branch_path.iter().rev() {
1134 if let Some(plan) = branch_view_for_in(plans, branch_id) {
1135 return Some(plan);
1136 }
1137 }
1138 None
1139}
1140
1141fn execution_scope_id(
1142 phase: Phase,
1143 variant_id: Option<&VariantId>,
1144 fold_id: Option<&FoldId>,
1145) -> String {
1146 format!(
1147 "scope:{}:{}:{}",
1148 phase_scope_label(phase),
1149 variant_id
1150 .map(ToString::to_string)
1151 .unwrap_or_else(|| "base".to_string()),
1152 fold_id
1153 .map(ToString::to_string)
1154 .unwrap_or_else(|| "nofold".to_string())
1155 )
1156}
1157
1158fn phase_scope_label(phase: Phase) -> &'static str {
1159 match phase {
1160 Phase::Compile => "COMPILE",
1161 Phase::Plan => "PLAN",
1162 Phase::FitCv => "FIT_CV",
1163 Phase::Select => "SELECT",
1164 Phase::Refit => "REFIT",
1165 Phase::Predict => "PREDICT",
1166 Phase::Explain => "EXPLAIN",
1167 }
1168}
1169
1170pub fn build_execution_plan(
1171 id: impl Into<String>,
1172 graph: GraphSpec,
1173 campaign: CampaignSpec,
1174 registry: &ControllerRegistry,
1175) -> Result<ExecutionPlan> {
1176 let id = id.into();
1177 if id.trim().is_empty() {
1178 return Err(DagMlError::Planning(
1179 "execution plan id is empty".to_string(),
1180 ));
1181 }
1182 campaign.validate()?;
1183 let graph_plan = GraphPlan::from_graph(graph)?;
1184 validate_campaign_node_targets(&graph_plan.graph, &campaign)?;
1185
1186 let mut node_plans = BTreeMap::new();
1187 let mut controller_manifests = BTreeMap::new();
1188 for node_id in &graph_plan.topological_order {
1189 let node = graph_plan
1190 .graph
1191 .nodes
1192 .iter()
1193 .find(|node| &node.id == node_id)
1194 .expect("topological node exists");
1195 let manifest = registry.resolve_for_node(node)?;
1196 let params = node.params.clone();
1197 let params_fingerprint = stable_json_fingerprint(¶ms)?;
1198 let inner_cv = match node.metadata.get("dsl_inner_cv") {
1202 Some(value) => {
1203 let spec =
1204 serde_json::from_value::<NestedCvSpec>(value.clone()).map_err(|error| {
1205 DagMlError::Planning(format!(
1206 "node `{}` has invalid dsl_inner_cv metadata: {error}",
1207 node.id
1208 ))
1209 })?;
1210 spec.validate().map_err(|error| {
1213 DagMlError::Planning(format!(
1214 "node `{}` has invalid dsl_inner_cv metadata: {error}",
1215 node.id
1216 ))
1217 })?;
1218 Some(spec)
1219 }
1220 None => None,
1221 };
1222 let shape_plan = campaign.shape_plans.get(&node.id).cloned();
1223 let data_bindings = campaign
1224 .data_bindings
1225 .get(&node.id)
1226 .cloned()
1227 .unwrap_or_default();
1228 node_plans.insert(
1229 node.id.clone(),
1230 NodePlan {
1231 inner_cv,
1232 node_id: node.id.clone(),
1233 kind: node.kind.clone(),
1234 controller_id: manifest.controller_id.clone(),
1235 controller_version: manifest.controller_version.clone(),
1236 supported_phases: manifest.supported_phases.clone(),
1237 controller_capabilities: manifest.capabilities.clone(),
1238 training_losses: Vec::new(),
1239 fit_scope: manifest.fit_scope,
1240 rng_policy: manifest.rng_policy,
1241 artifact_policy: manifest.artifact_policy,
1242 input_nodes: graph_plan.graph.upstream_nodes(&node.id),
1243 output_nodes: graph_plan.graph.downstream_nodes(&node.id),
1244 shape_plan,
1245 data_bindings,
1246 params,
1247 params_fingerprint,
1248 },
1249 );
1250 controller_manifests.insert(manifest.controller_id.clone(), manifest);
1251 }
1252
1253 let fold_set = campaign
1254 .split_invocation
1255 .as_ref()
1256 .and_then(|split| split.fold_set.clone());
1257 validate_search_space_fingerprint(&graph_plan.graph, &campaign)?;
1258 let variants = enumerate_variants(&campaign.generation, campaign.root_seed)?;
1259 validate_generation_override_targets(&graph_plan.graph, &variants)?;
1260 let graph_fingerprint = stable_json_fingerprint(&graph_plan.graph)?;
1261 let campaign_fingerprint = stable_json_fingerprint(&campaign)?;
1262 let controller_fingerprint = stable_json_fingerprint(&controller_manifests)?;
1263 let plan = ExecutionPlan {
1264 id,
1265 graph_plan,
1266 campaign,
1267 node_plans,
1268 controller_manifests,
1269 variants,
1270 fold_set,
1271 graph_fingerprint,
1272 campaign_fingerprint,
1273 controller_fingerprint,
1274 };
1275 plan.validate()?;
1276 Ok(plan)
1277}
1278
1279fn validate_search_space_fingerprint(graph: &GraphSpec, campaign: &CampaignSpec) -> Result<()> {
1280 let Some(expected_fingerprint) = &graph.search_space_fingerprint else {
1281 return Ok(());
1282 };
1283 if expected_fingerprint.trim().is_empty() {
1284 return Err(DagMlError::Planning(format!(
1285 "graph `{}` has empty search_space_fingerprint",
1286 graph.id
1287 )));
1288 }
1289 let actual_fingerprint = generation_spec_fingerprint(&campaign.generation)?;
1290 if expected_fingerprint != &actual_fingerprint {
1291 return Err(DagMlError::Planning(format!(
1292 "graph `{}` search_space_fingerprint does not match campaign generation spec",
1293 graph.id
1294 )));
1295 }
1296 Ok(())
1297}
1298
1299fn validate_generation_override_targets(graph: &GraphSpec, variants: &[VariantPlan]) -> Result<()> {
1300 let node_ids = graph
1301 .nodes
1302 .iter()
1303 .map(|node| node.id.clone())
1304 .collect::<BTreeSet<_>>();
1305 for variant in variants {
1306 for node_id in variant.param_override_targets()? {
1307 if !node_ids.contains(&node_id) {
1308 return Err(DagMlError::Planning(format!(
1309 "variant `{}` overrides params for unknown node `{node_id}`",
1310 variant.variant_id
1311 )));
1312 }
1313 }
1314 }
1315 Ok(())
1316}
1317
1318fn validate_campaign_node_targets(graph: &GraphSpec, campaign: &CampaignSpec) -> Result<()> {
1319 let node_ids = graph
1320 .nodes
1321 .iter()
1322 .map(|node| &node.id)
1323 .collect::<BTreeSet<_>>();
1324 for node_id in campaign.shape_plans.keys() {
1325 if !node_ids.contains(node_id) {
1326 return Err(DagMlError::Planning(format!(
1327 "shape plan references unknown node `{node_id}`"
1328 )));
1329 }
1330 }
1331 for node_id in campaign.data_bindings.keys() {
1332 if !node_ids.contains(node_id) {
1333 return Err(DagMlError::Planning(format!(
1334 "data binding references unknown node `{node_id}`"
1335 )));
1336 }
1337 }
1338 Ok(())
1339}
1340
1341pub fn prune_plan_to_active(
1373 plan: &ExecutionPlan,
1374 active_nodes: &BTreeSet<NodeId>,
1375 all_choice_nodes: &BTreeSet<NodeId>,
1376 variant: &VariantPlan,
1377) -> Result<ExecutionPlan> {
1378 plan.validate()?;
1379 variant.validate()?;
1380 for node_id in active_nodes {
1381 if !plan.node_plans.contains_key(node_id) {
1382 return Err(DagMlError::Planning(format!(
1383 "operator-SELECT prune: active node `{node_id}` is not in the union plan"
1384 )));
1385 }
1386 }
1387 if active_nodes.is_empty() {
1388 return Err(DagMlError::Planning(
1389 "operator-SELECT prune: active node set is empty".to_string(),
1390 ));
1391 }
1392
1393 let graph = &plan.graph_plan.graph;
1396 let mut ancestors = BTreeSet::<NodeId>::new();
1397 let mut stack: Vec<NodeId> = active_nodes.iter().cloned().collect();
1398 while let Some(node_id) = stack.pop() {
1399 for upstream in graph.upstream_nodes(&node_id) {
1400 if ancestors.insert(upstream.clone()) {
1401 stack.push(upstream);
1402 }
1403 }
1404 }
1405 let shared_prefix = ancestors
1406 .into_iter()
1407 .filter(|node_id| !all_choice_nodes.contains(node_id))
1408 .collect::<BTreeSet<_>>();
1409
1410 let keep = shared_prefix
1411 .iter()
1412 .chain(active_nodes.iter())
1413 .cloned()
1414 .collect::<BTreeSet<_>>();
1415
1416 let mut pruned_graph = graph.clone();
1419 pruned_graph.nodes.retain(|node| keep.contains(&node.id));
1420 pruned_graph
1421 .edges
1422 .retain(|edge| keep.contains(&edge.source.node_id) && keep.contains(&edge.target.node_id));
1423
1424 let graph_plan = GraphPlan::from_graph(pruned_graph)?;
1425
1426 let mut node_plans = BTreeMap::new();
1429 for (node_id, node_plan) in &plan.node_plans {
1430 if !keep.contains(node_id) {
1431 continue;
1432 }
1433 let mut pruned_node_plan = node_plan.clone();
1434 pruned_node_plan.input_nodes = graph_plan.graph.upstream_nodes(node_id);
1435 pruned_node_plan.output_nodes = graph_plan.graph.downstream_nodes(node_id);
1436 node_plans.insert(node_id.clone(), pruned_node_plan);
1437 }
1438
1439 let graph_fingerprint = stable_json_fingerprint(&graph_plan.graph)?;
1440 let pruned = ExecutionPlan {
1441 id: plan.id.clone(),
1442 graph_plan,
1443 campaign: plan.campaign.clone(),
1444 node_plans,
1445 controller_manifests: plan.controller_manifests.clone(),
1446 variants: vec![variant.clone()],
1447 fold_set: plan.fold_set.clone(),
1448 graph_fingerprint,
1449 campaign_fingerprint: plan.campaign_fingerprint.clone(),
1450 controller_fingerprint: plan.controller_fingerprint.clone(),
1451 };
1452 pruned.validate()?;
1453 validate_active_inputs(&pruned, graph)?;
1454 Ok(pruned)
1455}
1456
1457fn validate_active_inputs(plan: &ExecutionPlan, union_graph: &GraphSpec) -> Result<()> {
1471 let pruned_graph = &plan.graph_plan.graph;
1472 let kept: BTreeSet<&NodeId> = pruned_graph.nodes.iter().map(|node| &node.id).collect();
1473
1474 let mut union_port_sources = BTreeMap::<(NodeId, String), usize>::new();
1475 for edge in &union_graph.edges {
1476 if !kept.contains(&edge.target.node_id) {
1477 continue;
1478 }
1479 *union_port_sources
1480 .entry((edge.target.node_id.clone(), edge.target.port_name.clone()))
1481 .or_insert(0) += 1;
1482 }
1483 let mut pruned_port_sources = BTreeMap::<(NodeId, String), usize>::new();
1484 for edge in &pruned_graph.edges {
1485 *pruned_port_sources
1486 .entry((edge.target.node_id.clone(), edge.target.port_name.clone()))
1487 .or_insert(0) += 1;
1488 }
1489
1490 for (key, union_count) in &union_port_sources {
1491 let pruned_count = pruned_port_sources.get(key).copied().unwrap_or(0);
1492 let (node_id, port_name) = key;
1493 if *union_count >= 1 && pruned_count == 0 {
1494 return Err(DagMlError::Planning(format!(
1495 "operator-SELECT prune left node `{node_id}` required input port `{port_name}` with zero surviving sources (dangling): its producer was pruned away"
1496 )));
1497 }
1498 if pruned_count > 1 {
1499 return Err(DagMlError::Planning(format!(
1500 "operator-SELECT prune left node `{node_id}` input port `{port_name}` fed by {pruned_count} surviving sources (ambiguous)"
1501 )));
1502 }
1503 }
1504 Ok(())
1505}
1506
1507#[cfg(test)]
1508mod tests {
1509 use std::collections::{BTreeMap, BTreeSet};
1510 use std::time::{Duration, Instant};
1511
1512 use super::*;
1513 use crate::controller::{
1514 ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest, RngPolicy,
1515 };
1516 use crate::fold::FoldPartitionMode;
1517
1518 #[test]
1519 fn params_fingerprint_pins_serde_json_binary64_spelling() {
1520 let params = BTreeMap::from([
1521 (
1522 "scope".to_string(),
1523 serde_json::Value::String("train_only".to_string()),
1524 ),
1525 ("std".to_string(), serde_json::Value::from(1e-7_f64)),
1526 ]);
1527 assert_eq!(
1528 stable_json_fingerprint(¶ms).unwrap(),
1529 "3f417903752f65005bc9b69bcd23dfcf3ede2cda010e4f1ead6090a1a407b851"
1530 );
1531
1532 for (value, expected) in [
1535 (1e-7_f64, "1e-7"),
1536 (1e-6_f64, "1e-6"),
1537 (1e-5_f64, "0.00001"),
1538 (1e20_f64, "1e+20"),
1539 (1e21_f64, "1e+21"),
1540 (-0.0_f64, "-0.0"),
1541 (0.1_f64, "0.1"),
1542 (2.0_f64, "2.0"),
1543 (f64::from_bits(1), "5e-324"),
1544 ] {
1545 assert_eq!(serde_json::to_string(&value).unwrap(), expected);
1546 }
1547 }
1548
1549 #[test]
1550 fn inner_cv_is_declarable_at_campaign_and_node_level() {
1551 let campaign_json = r#"{"id":"c","root_seed":null,"inner_cv":{"kind":"kfold","n_splits":3,"shuffle":false,"seed":5}}"#;
1553 let campaign: CampaignSpec = serde_json::from_str(campaign_json).unwrap();
1554 campaign.validate().unwrap();
1555 assert!(campaign.inner_cv.is_some());
1556
1557 let node_inner = crate::fold::NestedCvSpec::KFold(crate::fold::KFoldSpec {
1559 n_splits: 4,
1560 shuffle: false,
1561 seed: Some(6),
1562 });
1563 let resolved = crate::fold::resolve_inner_cv(Some(&node_inner), campaign.inner_cv.as_ref());
1564 assert_eq!(resolved, Some(&node_inner));
1565
1566 let bare = r#"{"id":"c","root_seed":null}"#;
1568 let bare_campaign: CampaignSpec = serde_json::from_str(bare).unwrap();
1569 assert!(bare_campaign.inner_cv.is_none());
1570 let reserialized = serde_json::to_string(&bare_campaign).unwrap();
1571 assert!(!reserialized.contains("inner_cv"));
1572
1573 let bad: CampaignSpec = serde_json::from_str(
1576 r#"{"id":"c","root_seed":null,"inner_cv":{"kind":"kfold","n_splits":1,"shuffle":false,"seed":null}}"#,
1577 )
1578 .unwrap();
1579 let error = bad.validate().unwrap_err();
1580 assert!(error.to_string().contains("at least two splits"));
1581 }
1582
1583 #[test]
1584 fn execution_plan_validate_rejects_invalid_node_local_inner_cv() {
1585 let campaign = CampaignSpec {
1589 inner_cv: None,
1590 id: "campaign:plan-validate".to_string(),
1591 root_seed: Some(7),
1592 leakage_policy: LeakageUnitPolicy::default(),
1593 aggregation_policy: AggregationPolicy::default(),
1594 split_invocation: None,
1595 generation: Default::default(),
1596 shape_plans: BTreeMap::new(),
1597 data_bindings: BTreeMap::new(),
1598 branch_view_plans: Vec::new(),
1599 metadata: BTreeMap::new(),
1600 };
1601 let mut plan =
1602 build_execution_plan("plan:validate", graph(), campaign, ®istry()).unwrap();
1603 plan.validate().unwrap();
1604 plan.node_plans
1605 .get_mut(&NodeId::new("model:pls").unwrap())
1606 .unwrap()
1607 .inner_cv = Some(crate::fold::NestedCvSpec::KFold(crate::fold::KFoldSpec {
1608 n_splits: 1,
1609 shuffle: false,
1610 seed: None,
1611 }));
1612 let error = plan.validate().unwrap_err();
1613 assert!(matches!(error, DagMlError::Planning(_)));
1614 assert!(error.to_string().contains("invalid inner_cv"));
1615 assert!(error.to_string().contains("at least two splits"));
1616 }
1617
1618 fn hardening_plan() -> ExecutionPlan {
1619 let campaign = CampaignSpec {
1620 inner_cv: None,
1621 id: "campaign:harden".to_string(),
1622 root_seed: Some(7),
1623 leakage_policy: LeakageUnitPolicy::default(),
1624 aggregation_policy: AggregationPolicy::default(),
1625 split_invocation: None,
1626 generation: Default::default(),
1627 shape_plans: BTreeMap::new(),
1628 data_bindings: BTreeMap::new(),
1629 branch_view_plans: Vec::new(),
1630 metadata: BTreeMap::new(),
1631 };
1632 build_execution_plan("plan:harden", graph(), campaign, ®istry()).unwrap()
1633 }
1634
1635 #[test]
1636 fn execution_plan_external_reader_rejects_positional_struct_sequences() {
1637 let plan = hardening_plan();
1638 let mut wire = serde_json::to_value(&plan).unwrap();
1639 wire["campaign"]["leakage_policy"] = serde_json::json!([]);
1640
1641 let permissive: ExecutionPlan = serde_json::from_value(wire.clone()).unwrap();
1644 permissive.validate().unwrap();
1645
1646 let error = ExecutionPlan::from_json(&serde_json::to_string(&wire).unwrap()).unwrap_err();
1648 assert!(error.to_string().contains("must use a JSON object"));
1649 assert!(error.to_string().contains("campaign.leakage_policy"));
1650 }
1651
1652 #[test]
1653 fn execution_plan_external_reader_preserves_typed_serde_compatibility() {
1654 let plan = hardening_plan();
1655 let mut wire = serde_json::to_value(&plan).unwrap();
1656 wire["graph_plan"]
1657 .as_object_mut()
1658 .unwrap()
1659 .remove("parallel_levels");
1660 wire["graph_plan"]["graph"]
1661 .as_object_mut()
1662 .unwrap()
1663 .insert("forward_compatible".to_string(), serde_json::json!(true));
1664
1665 let parsed = ExecutionPlan::from_json(&serde_json::to_string(&wire).unwrap()).unwrap();
1666 assert!(parsed.graph_plan.parallel_levels.is_empty());
1667 assert_eq!(parsed.graph_plan.graph, plan.graph_plan.graph);
1668 }
1669
1670 #[test]
1671 fn validate_rejects_manifest_keyed_under_foreign_controller_id() {
1672 let mut plan = hardening_plan();
1673 plan.validate().unwrap();
1674 plan.controller_manifests
1675 .get_mut(&ControllerId::new("controller:model").unwrap())
1676 .unwrap()
1677 .controller_id = ControllerId::new("controller:imposter").unwrap();
1678 let error = plan.validate().unwrap_err();
1679 assert!(matches!(error, DagMlError::Planning(_)));
1680 assert!(error.to_string().contains("declares id"));
1681 }
1682
1683 #[test]
1684 fn validate_rejects_node_plan_keyed_under_foreign_node_id() {
1685 let mut plan = hardening_plan();
1686 plan.validate().unwrap();
1687 plan.node_plans
1688 .get_mut(&NodeId::new("model:pls").unwrap())
1689 .unwrap()
1690 .node_id = NodeId::new("transform:snv").unwrap();
1691 let error = plan.validate().unwrap_err();
1692 assert!(matches!(error, DagMlError::Planning(_)));
1693 assert!(error.to_string().contains("declares node_id"));
1694 }
1695
1696 #[test]
1697 fn validate_rejects_node_plans_not_covering_graph_node_set() {
1698 let mut plan = hardening_plan();
1699 plan.validate().unwrap();
1700 let mut ghost = plan
1703 .node_plans
1704 .remove(&NodeId::new("model:pls").unwrap())
1705 .unwrap();
1706 let bogus = NodeId::new("model:ghost").unwrap();
1707 ghost.node_id = bogus.clone();
1708 plan.node_plans.insert(bogus, ghost);
1709 let error = plan.validate().unwrap_err();
1710 assert!(matches!(error, DagMlError::Planning(_)));
1711 assert!(error
1712 .to_string()
1713 .contains("do not exactly cover the graph node-id set"));
1714 }
1715
1716 #[test]
1717 fn validate_rejects_topological_order_that_omits_a_node() {
1718 let mut plan = hardening_plan();
1719 plan.validate().unwrap();
1720 plan.graph_plan.topological_order = vec![NodeId::new("transform:snv").unwrap()];
1723 let error = plan.validate().unwrap_err();
1724 assert!(matches!(error, DagMlError::Planning(_)));
1725 assert!(error
1726 .to_string()
1727 .contains("topological_order does not match"));
1728 }
1729
1730 #[test]
1731 fn validate_rejects_node_plan_kind_that_differs_from_graph_node() {
1732 let mut plan = hardening_plan();
1733 plan.validate().unwrap();
1734 plan.node_plans
1735 .get_mut(&NodeId::new("model:pls").unwrap())
1736 .unwrap()
1737 .kind = NodeKind::Transform;
1738 let error = plan.validate().unwrap_err();
1739 assert!(matches!(error, DagMlError::Planning(_)));
1740 assert!(error
1741 .to_string()
1742 .contains("kind does not match graph node kind"));
1743 }
1744
1745 #[test]
1746 fn validate_rejects_forged_input_adjacency() {
1747 let mut plan = hardening_plan();
1748 plan.validate().unwrap();
1749 plan.node_plans
1752 .get_mut(&NodeId::new("model:pls").unwrap())
1753 .unwrap()
1754 .input_nodes = vec![];
1755 let error = plan.validate().unwrap_err();
1756 assert!(matches!(error, DagMlError::Planning(_)));
1757 assert!(error
1758 .to_string()
1759 .contains("adjacency does not match the graph"));
1760 }
1761
1762 #[test]
1763 fn validate_rejects_node_plan_supported_phases_that_diverge_from_manifest() {
1764 let mut plan = hardening_plan();
1765 plan.validate().unwrap();
1766 plan.node_plans
1769 .get_mut(&NodeId::new("model:pls").unwrap())
1770 .unwrap()
1771 .supported_phases
1772 .insert(Phase::Explain);
1773 let error = plan.validate().unwrap_err();
1774 assert!(matches!(error, DagMlError::Planning(_)));
1775 assert!(error
1776 .to_string()
1777 .contains("does not match controller manifest"));
1778 }
1779
1780 #[test]
1781 fn validate_rejects_node_plan_controller_version_that_diverges_from_manifest() {
1782 let mut plan = hardening_plan();
1783 plan.validate().unwrap();
1784 plan.node_plans
1785 .get_mut(&NodeId::new("model:pls").unwrap())
1786 .unwrap()
1787 .controller_version = "9.9.9".to_string();
1788 let error = plan.validate().unwrap_err();
1789 assert!(matches!(error, DagMlError::Planning(_)));
1790 assert!(error
1791 .to_string()
1792 .contains("does not match controller manifest"));
1793 }
1794
1795 #[test]
1796 fn validate_rejects_stale_graph_fingerprint_after_graph_mutation() {
1797 let mut plan = hardening_plan();
1798 plan.validate().unwrap();
1799 plan.graph_plan
1804 .graph
1805 .metadata
1806 .insert("tampered".to_string(), serde_json::json!(true));
1807 let error = plan.validate().unwrap_err();
1808 assert!(matches!(error, DagMlError::Planning(_)));
1809 assert!(error
1810 .to_string()
1811 .contains("graph_fingerprint does not match"));
1812 }
1813
1814 #[test]
1815 fn validate_rejects_forged_graph_fingerprint_field() {
1816 let mut plan = hardening_plan();
1817 plan.validate().unwrap();
1818 plan.graph_fingerprint = "sha256:forged".to_string();
1821 let error = plan.validate().unwrap_err();
1822 assert!(matches!(error, DagMlError::Planning(_)));
1823 assert!(error
1824 .to_string()
1825 .contains("graph_fingerprint does not match"));
1826 }
1827
1828 #[test]
1829 fn validate_rejects_stale_campaign_fingerprint_after_campaign_mutation() {
1830 let mut plan = hardening_plan();
1831 plan.validate().unwrap();
1832 plan.campaign
1834 .metadata
1835 .insert("tampered".to_string(), serde_json::json!("x"));
1836 let error = plan.validate().unwrap_err();
1837 assert!(matches!(error, DagMlError::Planning(_)));
1838 assert!(error
1839 .to_string()
1840 .contains("campaign_fingerprint does not match"));
1841 }
1842
1843 #[test]
1844 fn validate_rejects_forged_campaign_fingerprint_field() {
1845 let mut plan = hardening_plan();
1846 plan.validate().unwrap();
1847 plan.campaign_fingerprint = "sha256:forged".to_string();
1848 let error = plan.validate().unwrap_err();
1849 assert!(matches!(error, DagMlError::Planning(_)));
1850 assert!(error
1851 .to_string()
1852 .contains("campaign_fingerprint does not match"));
1853 }
1854
1855 #[test]
1856 fn validate_rejects_stale_controller_fingerprint_after_manifest_mutation() {
1857 let mut plan = hardening_plan();
1858 plan.validate().unwrap();
1859 plan.controller_manifests
1865 .get_mut(&ControllerId::new("controller:model").unwrap())
1866 .unwrap()
1867 .priority = 7;
1868 let error = plan.validate().unwrap_err();
1869 assert!(matches!(error, DagMlError::Planning(_)));
1870 assert!(error
1871 .to_string()
1872 .contains("controller_fingerprint does not match"));
1873 }
1874
1875 #[test]
1876 fn validate_rejects_forged_controller_fingerprint_field() {
1877 let mut plan = hardening_plan();
1878 plan.validate().unwrap();
1879 plan.controller_fingerprint = "sha256:forged".to_string();
1880 let error = plan.validate().unwrap_err();
1881 assert!(matches!(error, DagMlError::Planning(_)));
1882 assert!(error
1883 .to_string()
1884 .contains("controller_fingerprint does not match"));
1885 }
1886
1887 #[test]
1888 fn build_execution_plan_lowers_dsl_inner_cv_metadata_into_node_plan() {
1889 let mut graph = graph();
1890 graph
1891 .nodes
1892 .iter_mut()
1893 .find(|node| node.id.as_str() == "model:pls")
1894 .unwrap()
1895 .metadata
1896 .insert(
1897 "dsl_inner_cv".to_string(),
1898 serde_json::json!({"kind": "kfold", "n_splits": 3, "shuffle": false, "seed": 9}),
1899 );
1900
1901 let campaign = CampaignSpec {
1902 inner_cv: None,
1903 id: "campaign:inner-cv".to_string(),
1904 root_seed: Some(7),
1905 leakage_policy: LeakageUnitPolicy::default(),
1906 aggregation_policy: AggregationPolicy::default(),
1907 split_invocation: None,
1908 generation: Default::default(),
1909 shape_plans: BTreeMap::new(),
1910 data_bindings: BTreeMap::new(),
1911 branch_view_plans: Vec::new(),
1912 metadata: BTreeMap::new(),
1913 };
1914
1915 let plan = build_execution_plan("plan:inner-cv", graph, campaign, ®istry()).unwrap();
1916 match &plan.node_plans[&NodeId::new("model:pls").unwrap()].inner_cv {
1917 Some(crate::fold::NestedCvSpec::KFold(k)) => {
1918 assert_eq!(k.n_splits, 3);
1919 assert_eq!(k.seed, Some(9));
1920 }
1921 other => panic!("expected lowered KFold inner_cv, got {other:?}"),
1922 }
1923 assert!(plan.node_plans[&NodeId::new("transform:snv").unwrap()]
1924 .inner_cv
1925 .is_none());
1926 }
1927
1928 #[test]
1929 fn build_execution_plan_rejects_malformed_dsl_inner_cv_metadata() {
1930 let mut graph = graph();
1931 graph
1932 .nodes
1933 .iter_mut()
1934 .find(|node| node.id.as_str() == "model:pls")
1935 .unwrap()
1936 .metadata
1937 .insert(
1938 "dsl_inner_cv".to_string(),
1939 serde_json::json!({"kind": "not_a_real_kind"}),
1940 );
1941
1942 let campaign = CampaignSpec {
1943 inner_cv: None,
1944 id: "campaign:inner-cv.bad".to_string(),
1945 root_seed: Some(7),
1946 leakage_policy: LeakageUnitPolicy::default(),
1947 aggregation_policy: AggregationPolicy::default(),
1948 split_invocation: None,
1949 generation: Default::default(),
1950 shape_plans: BTreeMap::new(),
1951 data_bindings: BTreeMap::new(),
1952 branch_view_plans: Vec::new(),
1953 metadata: BTreeMap::new(),
1954 };
1955
1956 let error =
1957 build_execution_plan("plan:inner-cv.bad", graph, campaign, ®istry()).unwrap_err();
1958 assert!(matches!(error, DagMlError::Planning(_)));
1959 assert!(error.to_string().contains("invalid dsl_inner_cv metadata"));
1960 }
1961
1962 #[test]
1963 fn build_execution_plan_rejects_semantically_invalid_dsl_inner_cv() {
1964 let mut graph = graph();
1967 graph
1968 .nodes
1969 .iter_mut()
1970 .find(|node| node.id.as_str() == "model:pls")
1971 .unwrap()
1972 .metadata
1973 .insert(
1974 "dsl_inner_cv".to_string(),
1975 serde_json::json!({"kind": "kfold", "n_splits": 1, "shuffle": false, "seed": null}),
1976 );
1977
1978 let campaign = CampaignSpec {
1979 inner_cv: None,
1980 id: "campaign:inner-cv.nsplits".to_string(),
1981 root_seed: Some(7),
1982 leakage_policy: LeakageUnitPolicy::default(),
1983 aggregation_policy: AggregationPolicy::default(),
1984 split_invocation: None,
1985 generation: Default::default(),
1986 shape_plans: BTreeMap::new(),
1987 data_bindings: BTreeMap::new(),
1988 branch_view_plans: Vec::new(),
1989 metadata: BTreeMap::new(),
1990 };
1991
1992 let error = build_execution_plan("plan:inner-cv.nsplits", graph, campaign, ®istry())
1993 .unwrap_err();
1994 assert!(matches!(error, DagMlError::Planning(_)));
1995 assert!(error.to_string().contains("at least two splits"));
1996 }
1997 use crate::data::{
1998 BranchViewMode, BranchViewPlan, DataBinding, DataViewSelector, SOURCE_INDEX_METADATA_KEY,
1999 };
2000 use crate::generation::{
2001 GenerationChoice, GenerationConstraints, GenerationDimension, GenerationParamOverride,
2002 GenerationStrategy,
2003 };
2004 use crate::graph::{
2005 EdgeContract, EdgeSpec, GraphInterface, NodeSpec, PortCardinality, PortKind, PortRef,
2006 PortSchema, PortSpec,
2007 };
2008 use crate::ids::{ControllerId, FoldId, ObservationId, SampleId, TargetId};
2009 use crate::phase::Phase;
2010 use crate::policy::{DataModelShapePlan, Granularity};
2011 use crate::relation::{SampleRelation, SampleRelationSet};
2012
2013 fn port(name: &str, kind: PortKind) -> PortSpec {
2014 PortSpec {
2015 name: name.to_string(),
2016 kind,
2017 representation: None,
2018 cardinality: PortCardinality::One,
2019 unit_level: None,
2020 alignment_key: None,
2021 target_level: None,
2022 description: String::new(),
2023 }
2024 }
2025
2026 fn node(id: &str, kind: NodeKind, inputs: Vec<PortSpec>, outputs: Vec<PortSpec>) -> NodeSpec {
2027 NodeSpec {
2028 id: NodeId::new(id).unwrap(),
2029 kind,
2030 operator: None,
2031 params: BTreeMap::new(),
2032 ports: PortSchema { inputs, outputs },
2033 metadata: BTreeMap::new(),
2034 seed_label: None,
2035 }
2036 }
2037
2038 fn graph() -> GraphSpec {
2039 GraphSpec {
2040 id: "g".to_string(),
2041 interface: GraphInterface::default(),
2042 nodes: vec![
2043 node(
2044 "transform:snv",
2045 NodeKind::Transform,
2046 vec![],
2047 vec![port("x", PortKind::Data)],
2048 ),
2049 node(
2050 "model:pls",
2051 NodeKind::Model,
2052 vec![port("x", PortKind::Data)],
2053 vec![port("pred", PortKind::Prediction)],
2054 ),
2055 ],
2056 edges: vec![EdgeSpec {
2057 source: PortRef {
2058 node_id: NodeId::new("transform:snv").unwrap(),
2059 port_name: "x".to_string(),
2060 },
2061 target: PortRef {
2062 node_id: NodeId::new("model:pls").unwrap(),
2063 port_name: "x".to_string(),
2064 },
2065 contract: EdgeContract {
2066 requires_oof: false,
2067 requires_fold_alignment: false,
2068 ..EdgeContract::new(PortKind::Data, None)
2069 },
2070 }],
2071 search_space_fingerprint: None,
2072 metadata: BTreeMap::new(),
2073 }
2074 }
2075
2076 fn manifest(id: &str, kind: NodeKind) -> ControllerManifest {
2077 let mut capabilities = BTreeSet::from([
2078 ControllerCapability::Deterministic,
2079 ControllerCapability::ThreadSafe,
2080 ControllerCapability::ProcessSafe,
2081 ]);
2082 if kind == NodeKind::Model {
2083 capabilities.insert(ControllerCapability::EmitsPredictions);
2084 capabilities.insert(ControllerCapability::ConsumesOofPredictions);
2085 }
2086 ControllerManifest {
2087 controller_id: ControllerId::new(id).unwrap(),
2088 controller_version: "0.1.0".to_string(),
2089 operator_kind: kind,
2090 priority: 0,
2091 supported_phases: BTreeSet::from([Phase::FitCv, Phase::Refit, Phase::Predict]),
2092 input_ports: Vec::new(),
2093 output_ports: Vec::new(),
2094 data_requirements: None,
2095 capabilities,
2096 operator_selectors: Vec::new(),
2097 fit_scope: ControllerFitScope::FoldTrain,
2098 rng_policy: RngPolicy::UsesCoreSeed,
2099 artifact_policy: ArtifactPolicy::Serializable,
2100 }
2101 }
2102
2103 fn registry() -> ControllerRegistry {
2104 let mut registry = ControllerRegistry::new();
2105 registry
2106 .register(manifest("controller:transform", NodeKind::Transform))
2107 .unwrap();
2108 registry
2109 .register(manifest("controller:model", NodeKind::Model))
2110 .unwrap();
2111 registry
2112 }
2113
2114 fn registry_with_model_data_requirements(
2115 accepted_representations: &[&str],
2116 accepted_types: &[&str],
2117 ) -> ControllerRegistry {
2118 let mut registry = ControllerRegistry::new();
2119 registry
2120 .register(manifest("controller:transform", NodeKind::Transform))
2121 .unwrap();
2122 let mut model = manifest("controller:model", NodeKind::Model);
2123 model.data_requirements = Some(serde_json::json!({
2124 "schema_version": 1,
2125 "ports": [
2126 {
2127 "name": "x",
2128 "accepted_representations": accepted_representations,
2129 "accepted_types": accepted_types,
2130 "rank": 2,
2131 "multi_source": true,
2132 "optional": false
2133 }
2134 ],
2135 "metadata": {
2136 "source": "plan-test"
2137 }
2138 }));
2139 registry.register(model).unwrap();
2140 registry
2141 }
2142
2143 fn registry_with_model_data_requirements_json(
2144 data_requirements: serde_json::Value,
2145 ) -> ControllerRegistry {
2146 let mut registry = ControllerRegistry::new();
2147 registry
2148 .register(manifest("controller:transform", NodeKind::Transform))
2149 .unwrap();
2150 let mut model = manifest("controller:model", NodeKind::Model);
2151 model.data_requirements = Some(data_requirements);
2152 registry.register(model).unwrap();
2153 registry
2154 }
2155
2156 fn model_data_requirements(
2157 multi_source: bool,
2158 default_fusion: Option<serde_json::Value>,
2159 ) -> serde_json::Value {
2160 let mut spec = serde_json::json!({
2161 "schema_version": 1,
2162 "ports": [
2163 {
2164 "name": "x",
2165 "accepted_representations": ["tabular_numeric"],
2166 "accepted_types": ["table"],
2167 "rank": 2,
2168 "multi_source": multi_source,
2169 "optional": false
2170 }
2171 ],
2172 "metadata": {
2173 "source": "plan-test"
2174 }
2175 });
2176 if let Some(default_fusion) = default_fusion {
2177 spec.as_object_mut()
2178 .unwrap()
2179 .insert("default_fusion".to_string(), default_fusion);
2180 }
2181 spec
2182 }
2183
2184 fn multisource_binding(node_id: &NodeId) -> DataBinding {
2185 let mut binding = data_binding(node_id);
2186 binding.request_id = "nir-chem-source-concat".to_string();
2187 binding.feature_set_id = Some("x_fused".to_string());
2188 binding.source_ids = vec!["nir".to_string(), "chem".to_string()];
2189 binding
2190 }
2191
2192 fn add_source_index(binding: &mut DataBinding) {
2193 binding.metadata.insert(
2194 SOURCE_INDEX_METADATA_KEY.to_string(),
2195 serde_json::json!({
2196 "nir": 0,
2197 "chem": 1
2198 }),
2199 );
2200 }
2201
2202 fn source_concat_fusion() -> serde_json::Value {
2203 serde_json::json!({
2204 "mode": "concatenate_features",
2205 "alignment": "sample_id",
2206 "adapter_id": null,
2207 "params": {
2208 "namespace_columns": true
2209 }
2210 })
2211 }
2212
2213 fn by_source_graph(source_ids: Vec<&str>) -> GraphSpec {
2214 let mut graph = graph();
2215 let branch_view = BranchViewPlan {
2216 view_id: "branch_view:source".to_string(),
2217 branch_id: "branch:source".to_string(),
2218 mode: BranchViewMode::BySource,
2219 selector: DataViewSelector {
2220 source_ids: source_ids.into_iter().map(str::to_string).collect(),
2221 ..Default::default()
2222 },
2223 allow_overlap: false,
2224 metadata: BTreeMap::new(),
2225 };
2226 graph
2227 .nodes
2228 .iter_mut()
2229 .find(|node| node.id.as_str() == "model:pls")
2230 .unwrap()
2231 .metadata
2232 .insert(
2233 "dsl_branch_view_plan".to_string(),
2234 serde_json::to_value(branch_view).unwrap(),
2235 );
2236 graph
2237 }
2238
2239 fn refusal_payload(error: DagMlError) -> serde_json::Value {
2240 let message = error.to_string();
2241 let payload = message
2242 .split_once("data requirement refusal: ")
2243 .unwrap_or_else(|| panic!("missing structured refusal payload in: {message}"))
2244 .1;
2245 serde_json::from_str(payload).unwrap()
2246 }
2247
2248 fn campaign(id: &str) -> CampaignSpec {
2249 CampaignSpec {
2250 id: id.to_string(),
2251 root_seed: Some(7),
2252 leakage_policy: LeakageUnitPolicy::default(),
2253 aggregation_policy: AggregationPolicy::default(),
2254 split_invocation: None,
2255 generation: Default::default(),
2256 shape_plans: BTreeMap::new(),
2257 data_bindings: BTreeMap::new(),
2258 branch_view_plans: Vec::new(),
2259 inner_cv: None,
2260 metadata: BTreeMap::new(),
2261 }
2262 }
2263
2264 fn custom_loss_role(node_id: &str, output_id: &str) -> TrainingLossRoleReference {
2265 let fixture: serde_json::Value = serde_json::from_str(include_str!(
2266 "../../../examples/fixtures/criteria/javascript_local_implementations.v1.json"
2267 ))
2268 .unwrap();
2269 let mut role: TrainingLossRoleReference =
2270 serde_json::from_value(fixture["training_loss_role"].clone()).unwrap();
2271 role.node_id = NodeId::new(node_id).unwrap();
2272 role.output_id = Some(output_id.to_string());
2273 role
2274 }
2275
2276 #[test]
2277 fn execution_plan_lowers_training_losses_in_canonical_order() {
2278 let mut loss_registry = ControllerRegistry::new();
2279 loss_registry
2280 .register(manifest("controller:transform", NodeKind::Transform))
2281 .unwrap();
2282 let mut model_manifest = manifest("controller:model", NodeKind::Model);
2283 model_manifest.capabilities.extend([
2284 ControllerCapability::SupportsConfigurableLoss,
2285 ControllerCapability::SupportsCustomLoss,
2286 ControllerCapability::SupportsDifferentiableLoss,
2287 ]);
2288 loss_registry.register(model_manifest).unwrap();
2289
2290 let plan = build_execution_plan(
2291 "plan:training-loss-lowering",
2292 graph(),
2293 campaign("campaign:training-loss-lowering"),
2294 &loss_registry,
2295 )
2296 .unwrap();
2297 let role_b = custom_loss_role("model:pls", "b");
2298 let role_a = custom_loss_role("model:pls", "a");
2299 let bound = plan
2300 .clone()
2301 .with_training_losses(vec![role_b, role_a.clone()])
2302 .unwrap();
2303 let model = bound
2304 .node_plans
2305 .get(&NodeId::new("model:pls").unwrap())
2306 .unwrap();
2307 assert_eq!(
2308 model
2309 .training_losses
2310 .iter()
2311 .map(|role| role.output_id.as_deref())
2312 .collect::<Vec<_>>(),
2313 vec![Some("a"), Some("b")]
2314 );
2315
2316 let cleared = bound.with_training_losses(Vec::new()).unwrap();
2317 assert!(cleared
2318 .node_plans
2319 .values()
2320 .all(|node| node.training_losses.is_empty()));
2321
2322 let mut unknown = role_a.clone();
2323 unknown.node_id = NodeId::new("model:unknown").unwrap();
2324 assert!(plan
2325 .clone()
2326 .with_training_losses(vec![unknown])
2327 .unwrap_err()
2328 .to_string()
2329 .contains("unknown plan node"));
2330
2331 let incapable = build_execution_plan(
2332 "plan:training-loss-incapable",
2333 graph(),
2334 campaign("campaign:training-loss-incapable"),
2335 ®istry(),
2336 )
2337 .unwrap();
2338 assert!(incapable
2339 .with_training_losses(vec![role_a])
2340 .unwrap_err()
2341 .to_string()
2342 .contains("does not support configurable loss"));
2343 }
2344
2345 #[test]
2346 fn build_execution_plan_consumes_controller_data_requirements_for_bindings() {
2347 let model_id = NodeId::new("model:pls").unwrap();
2348 let mut campaign = campaign("campaign:datareq.ok");
2349 campaign.data_bindings = BTreeMap::from([(
2350 model_id,
2351 vec![data_binding(&NodeId::new("model:pls").unwrap())],
2352 )]);
2353 let plan = build_execution_plan(
2354 "plan:datareq.ok",
2355 graph(),
2356 campaign,
2357 ®istry_with_model_data_requirements(&["tabular_numeric"], &["table"]),
2358 )
2359 .unwrap();
2360 assert_eq!(
2361 plan.node_plans[&NodeId::new("model:pls").unwrap()].data_bindings[0]
2362 .output_representation,
2363 "tabular_numeric"
2364 );
2365 }
2366
2367 #[test]
2368 fn build_execution_plan_rejects_binding_representation_outside_data_requirements() {
2369 let model_id = NodeId::new("model:pls").unwrap();
2370 let mut campaign = campaign("campaign:datareq.representation");
2371 campaign.data_bindings =
2372 BTreeMap::from([(model_id.clone(), vec![data_binding(&model_id)])]);
2373 let error = build_execution_plan(
2374 "plan:datareq.representation",
2375 graph(),
2376 campaign,
2377 ®istry_with_model_data_requirements(&["signal_1d"], &["dense_signal"]),
2378 )
2379 .unwrap_err();
2380 assert!(
2381 error.to_string().contains("output representation"),
2382 "unexpected error: {error}"
2383 );
2384 }
2385
2386 #[test]
2387 fn build_execution_plan_rejects_binding_registered_type_outside_data_requirements() {
2388 let model_id = NodeId::new("model:pls").unwrap();
2389 let mut campaign = campaign("campaign:datareq.type");
2390 campaign.data_bindings =
2391 BTreeMap::from([(model_id.clone(), vec![data_binding(&model_id)])]);
2392 let error = build_execution_plan(
2393 "plan:datareq.type",
2394 graph(),
2395 campaign,
2396 ®istry_with_model_data_requirements(&["tabular_numeric"], &["dense_signal"]),
2397 )
2398 .unwrap_err();
2399 assert!(
2400 error.to_string().contains("registered type `table`"),
2401 "unexpected error: {error}"
2402 );
2403 }
2404
2405 #[test]
2406 fn build_execution_plan_rejects_source_concat_without_source_index() {
2407 let model_id = NodeId::new("model:pls").unwrap();
2408 let mut campaign = campaign("campaign:datareq.source-concat.no-index");
2409 campaign.data_bindings =
2410 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
2411
2412 let error = build_execution_plan(
2413 "plan:datareq.source-concat.no-index",
2414 graph(),
2415 campaign,
2416 ®istry_with_model_data_requirements_json(model_data_requirements(
2417 true,
2418 Some(source_concat_fusion()),
2419 )),
2420 )
2421 .unwrap_err();
2422 let payload = refusal_payload(error);
2423
2424 assert_eq!(
2425 payload["code"],
2426 "dagml.data_requirement.source_concat_requires_source_index"
2427 );
2428 assert_eq!(payload["shape"], "source_concat");
2429 assert_eq!(payload["source_ids"], serde_json::json!(["nir", "chem"]));
2430 }
2431
2432 #[test]
2433 fn build_execution_plan_rejects_multisource_binding_without_data_requirements() {
2434 let model_id = NodeId::new("model:pls").unwrap();
2435 let mut campaign = campaign("campaign:datareq.multisource.no-requirements");
2436 campaign.data_bindings =
2437 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
2438
2439 let error = build_execution_plan(
2440 "plan:datareq.multisource.no-requirements",
2441 graph(),
2442 campaign,
2443 ®istry(),
2444 )
2445 .unwrap_err();
2446 let payload = refusal_payload(error);
2447
2448 assert_eq!(
2449 payload["code"],
2450 "dagml.data_requirement.missing_data_requirements"
2451 );
2452 assert_eq!(payload["source_ids"], serde_json::json!(["nir", "chem"]));
2453 }
2454
2455 #[test]
2456 fn build_execution_plan_accepts_source_concat_with_source_index() {
2457 let model_id = NodeId::new("model:pls").unwrap();
2458 let mut binding = multisource_binding(&model_id);
2459 add_source_index(&mut binding);
2460 let mut campaign = campaign("campaign:datareq.source-concat.index");
2461 campaign.data_bindings = BTreeMap::from([(model_id.clone(), vec![binding])]);
2462
2463 let plan = build_execution_plan(
2464 "plan:datareq.source-concat.index",
2465 graph(),
2466 campaign,
2467 ®istry_with_model_data_requirements_json(model_data_requirements(
2468 true,
2469 Some(source_concat_fusion()),
2470 )),
2471 )
2472 .unwrap();
2473
2474 assert_eq!(
2475 plan.node_plans[&model_id].data_bindings[0].metadata[SOURCE_INDEX_METADATA_KEY],
2476 serde_json::json!({"nir": 0, "chem": 1})
2477 );
2478 }
2479
2480 #[test]
2481 fn by_source_branch_allows_single_source_fit_from_multisource_binding() {
2482 let model_id = NodeId::new("model:pls").unwrap();
2483 let mut campaign = campaign("campaign:datareq.by-source.single");
2484 campaign.data_bindings =
2485 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
2486
2487 let plan = build_execution_plan(
2488 "plan:datareq.by-source.single",
2489 by_source_graph(vec!["nir"]),
2490 campaign,
2491 ®istry_with_model_data_requirements_json(model_data_requirements(false, None)),
2492 )
2493 .unwrap();
2494
2495 assert_eq!(
2496 plan.node_plans[&model_id].data_bindings[0].source_ids.len(),
2497 2
2498 );
2499 }
2500
2501 #[test]
2502 fn by_source_branch_refuses_multi_source_selector_shape() {
2503 let model_id = NodeId::new("model:pls").unwrap();
2504 let mut campaign = campaign("campaign:datareq.by-source.multi");
2505 campaign.data_bindings =
2506 BTreeMap::from([(model_id.clone(), vec![multisource_binding(&model_id)])]);
2507
2508 let error = build_execution_plan(
2509 "plan:datareq.by-source.multi",
2510 by_source_graph(vec!["nir", "chem"]),
2511 campaign,
2512 ®istry_with_model_data_requirements_json(model_data_requirements(true, None)),
2513 )
2514 .unwrap_err();
2515 let payload = refusal_payload(error);
2516
2517 assert_eq!(
2518 payload["code"],
2519 "dagml.data_requirement.unsupported_by_source_shape"
2520 );
2521 assert_eq!(payload["shape"], "by_source");
2522 assert_eq!(payload["source_ids"], serde_json::json!(["nir", "chem"]));
2523 }
2524
2525 fn large_linear_graph(transform_count: usize) -> GraphSpec {
2526 let mut nodes = Vec::new();
2527 let mut edges = Vec::new();
2528 for node_idx in 0..transform_count {
2529 let node_id = format!("transform:t{node_idx:04}");
2530 nodes.push(node(
2531 &node_id,
2532 NodeKind::Transform,
2533 vec![port("x", PortKind::Data)],
2534 vec![port("x", PortKind::Data)],
2535 ));
2536 if node_idx > 0 {
2537 edges.push(EdgeSpec {
2538 source: PortRef {
2539 node_id: NodeId::new(format!("transform:t{:04}", node_idx - 1)).unwrap(),
2540 port_name: "x".to_string(),
2541 },
2542 target: PortRef {
2543 node_id: NodeId::new(&node_id).unwrap(),
2544 port_name: "x".to_string(),
2545 },
2546 contract: EdgeContract::new(PortKind::Data, None),
2547 });
2548 }
2549 }
2550 nodes.push(node(
2551 "model:final",
2552 NodeKind::Model,
2553 vec![port("x", PortKind::Data)],
2554 vec![port("pred", PortKind::Prediction)],
2555 ));
2556 edges.push(EdgeSpec {
2557 source: PortRef {
2558 node_id: NodeId::new(format!("transform:t{:04}", transform_count - 1)).unwrap(),
2559 port_name: "x".to_string(),
2560 },
2561 target: PortRef {
2562 node_id: NodeId::new("model:final").unwrap(),
2563 port_name: "x".to_string(),
2564 },
2565 contract: EdgeContract::new(PortKind::Data, None),
2566 });
2567
2568 GraphSpec {
2569 id: "g:perf.linear".to_string(),
2570 interface: GraphInterface::default(),
2571 nodes,
2572 edges,
2573 search_space_fingerprint: None,
2574 metadata: BTreeMap::new(),
2575 }
2576 }
2577
2578 fn oof_graph() -> GraphSpec {
2579 GraphSpec {
2580 id: "g:oof.capabilities".to_string(),
2581 interface: GraphInterface::default(),
2582 nodes: vec![
2583 node(
2584 "model:base",
2585 NodeKind::Model,
2586 vec![],
2587 vec![port("pred", PortKind::Prediction)],
2588 ),
2589 node(
2590 "model:meta",
2591 NodeKind::Model,
2592 vec![port("pred", PortKind::Prediction)],
2593 vec![port("pred", PortKind::Prediction)],
2594 ),
2595 ],
2596 edges: vec![EdgeSpec {
2597 source: PortRef {
2598 node_id: NodeId::new("model:base").unwrap(),
2599 port_name: "pred".to_string(),
2600 },
2601 target: PortRef {
2602 node_id: NodeId::new("model:meta").unwrap(),
2603 port_name: "pred".to_string(),
2604 },
2605 contract: EdgeContract {
2606 requires_oof: true,
2607 requires_fold_alignment: true,
2608 ..EdgeContract::new(PortKind::Prediction, None)
2609 },
2610 }],
2611 search_space_fingerprint: None,
2612 metadata: BTreeMap::new(),
2613 }
2614 }
2615
2616 fn data_binding(node_id: &NodeId) -> DataBinding {
2617 DataBinding {
2618 node_id: node_id.clone(),
2619 input_name: "x".to_string(),
2620 request_id: "nir-to-tabular".to_string(),
2621 schema_fingerprint: "f97b37872fa22134b508f98fd8e207e5b776b52594fb8f6f5c3e15bee212246b"
2622 .to_string(),
2623 plan_fingerprint: "7c5431d85574b3f337022fa5d25971d5b5cf445b90331b49938f573ff6901e4d"
2624 .to_string(),
2625 relation_fingerprint: Some(
2626 "a3a7e329df35db9f2883a17b8611b7fae6dcaa031875e3ec2c9be1b9e29cbe10".to_string(),
2627 ),
2628 output_representation: "tabular_numeric".to_string(),
2629 feature_set_id: Some("x".to_string()),
2630 source_ids: vec!["nir".to_string()],
2631 require_relations: true,
2632 view_policy: Default::default(),
2633 metadata: BTreeMap::new(),
2634 }
2635 }
2636
2637 fn levels_as_strings(levels: &[Vec<NodeId>]) -> Vec<Vec<String>> {
2638 levels
2639 .iter()
2640 .map(|level| level.iter().map(ToString::to_string).collect())
2641 .collect()
2642 }
2643
2644 #[test]
2645 fn published_campaign_spec_schema_declares_current_contract() {
2646 let schema: serde_json::Value = serde_json::from_str(include_str!(
2647 "../../../docs/contracts/campaign_spec.schema.json"
2648 ))
2649 .unwrap();
2650
2651 assert_eq!(schema["$id"], CAMPAIGN_SPEC_SCHEMA_ID);
2652 assert!(schema["required"]
2653 .as_array()
2654 .unwrap()
2655 .iter()
2656 .any(|field| field.as_str() == Some("id")));
2657 assert!(schema["$defs"]["split_invocation"]["properties"]
2658 .as_object()
2659 .unwrap()
2660 .contains_key("fold_set"));
2661 assert!(schema["$defs"]["aggregation_policy"]["properties"]
2662 .as_object()
2663 .unwrap()
2664 .contains_key("selection_metric_level"));
2665 assert!(schema["$defs"]["aggregation_policy"]["properties"]
2666 .as_object()
2667 .unwrap()
2668 .contains_key("custom_controller"));
2669 assert!(schema["$defs"]["data_binding"]["properties"]
2670 .as_object()
2671 .unwrap()
2672 .contains_key("view_policy"));
2673 assert!(schema["properties"]
2674 .as_object()
2675 .unwrap()
2676 .contains_key("branch_view_plans"));
2677 assert!(schema["$defs"]["branch_view_plan"]["properties"]
2678 .as_object()
2679 .unwrap()
2680 .contains_key("selector"));
2681 }
2682
2683 #[test]
2684 fn published_execution_plan_schema_declares_current_contract() {
2685 let schema: serde_json::Value = serde_json::from_str(include_str!(
2686 "../../../docs/contracts/execution_plan.schema.json"
2687 ))
2688 .unwrap();
2689
2690 assert_eq!(schema["$id"], EXECUTION_PLAN_SCHEMA_ID);
2691 assert!(schema["required"]
2692 .as_array()
2693 .unwrap()
2694 .iter()
2695 .any(|field| field.as_str() == Some("node_plans")));
2696 assert!(schema["properties"]
2697 .as_object()
2698 .unwrap()
2699 .contains_key("controller_fingerprint"));
2700 assert!(schema["$defs"]["node_plan"]["properties"]
2701 .as_object()
2702 .unwrap()
2703 .contains_key("shape_plan"));
2704 assert!(schema["$defs"]["variant_plan"]["properties"]
2705 .as_object()
2706 .unwrap()
2707 .contains_key("choices"));
2708 }
2709
2710 #[test]
2711 fn published_execution_plan_fixture_validates_current_contract() {
2712 let plan: ExecutionPlan = serde_json::from_str(include_str!(
2713 "../../../examples/fixtures/runtime/execution_plan_branch_merge_executable.json"
2714 ))
2715 .unwrap();
2716
2717 plan.validate().unwrap();
2718 assert_eq!(plan.id, "plan:fixture.execution.branch_merge");
2719 assert_eq!(plan.variants.len(), 2);
2720 assert_eq!(plan.node_plans.len(), plan.graph_plan.graph.nodes.len());
2721 }
2722
2723 #[test]
2724 #[ignore = "perf sanity probe; run with --release --ignored --nocapture"]
2725 fn build_execution_plan_large_linear_graph_under_1500ms() {
2726 let started = Instant::now();
2727 let plan = build_execution_plan(
2728 "plan:perf.linear",
2729 large_linear_graph(400),
2730 campaign("campaign:perf.linear"),
2731 ®istry(),
2732 )
2733 .unwrap();
2734 let elapsed = started.elapsed();
2735
2736 assert_eq!(plan.graph_plan.topological_order.len(), 401);
2737 assert_eq!(plan.node_plans.len(), 401);
2738 assert!(
2739 elapsed <= Duration::from_millis(1_500),
2740 "large execution-plan build took {elapsed:?}"
2741 );
2742 }
2743
2744 #[test]
2745 fn builds_execution_plan_with_shape_and_fold_contracts() {
2746 let model_id = NodeId::new("model:pls").unwrap();
2747 let campaign = CampaignSpec {
2748 inner_cv: None,
2749 id: "campaign:oof".to_string(),
2750 root_seed: Some(7),
2751 leakage_policy: LeakageUnitPolicy::default(),
2752 aggregation_policy: AggregationPolicy::default(),
2753 split_invocation: Some(SplitInvocation {
2754 id: "split:outer".to_string(),
2755 controller_id: None,
2756 leakage_policy: LeakageUnitPolicy::default(),
2757 params: BTreeMap::new(),
2758 fold_set: Some(FoldSet {
2759 id: "outer".to_string(),
2760 sample_ids: vec![SampleId::new("s1").unwrap(), SampleId::new("s2").unwrap()],
2761 folds: vec![
2762 crate::fold::FoldAssignment {
2763 fold_id: FoldId::new("fold0").unwrap(),
2764 train_sample_ids: vec![SampleId::new("s2").unwrap()],
2765 validation_sample_ids: vec![SampleId::new("s1").unwrap()],
2766 metadata: BTreeMap::new(),
2767 },
2768 crate::fold::FoldAssignment {
2769 fold_id: FoldId::new("fold1").unwrap(),
2770 train_sample_ids: vec![SampleId::new("s1").unwrap()],
2771 validation_sample_ids: vec![SampleId::new("s2").unwrap()],
2772 metadata: BTreeMap::new(),
2773 },
2774 ],
2775 sample_groups: BTreeMap::new(),
2776 partition_mode: FoldPartitionMode::Partition,
2777 }),
2778 }),
2779 generation: Default::default(),
2780 shape_plans: BTreeMap::from([(
2781 model_id.clone(),
2782 DataModelShapePlan {
2783 node_id: model_id.clone(),
2784 input_granularity: Granularity::Observation,
2785 ..DataModelShapePlan {
2786 node_id: model_id.clone(),
2787 input_granularity: Granularity::Sample,
2788 target_granularity: Granularity::Sample,
2789 fit_rows: crate::policy::FitBoundary::FoldTrain,
2790 predict_rows: crate::policy::FitBoundary::FoldValidation,
2791 feature_namespace: None,
2792 feature_schema_fingerprint: None,
2793 target_space: "raw".to_string(),
2794 aggregation_policy: AggregationPolicy::default(),
2795 augmentation_policy: crate::policy::AugmentationPolicy::default(),
2796 selection_policy: crate::policy::FeatureSelectionPolicy::default(),
2797 }
2798 },
2799 )]),
2800 data_bindings: BTreeMap::from([(model_id.clone(), vec![data_binding(&model_id)])]),
2801 branch_view_plans: Vec::new(),
2802 metadata: BTreeMap::new(),
2803 };
2804
2805 let plan = build_execution_plan("plan:oof", graph(), campaign, ®istry()).unwrap();
2806
2807 assert_eq!(
2808 plan.graph_plan
2809 .topological_order
2810 .iter()
2811 .map(ToString::to_string)
2812 .collect::<Vec<_>>(),
2813 vec!["transform:snv", "model:pls"]
2814 );
2815 assert_eq!(
2816 levels_as_strings(&plan.graph_plan.parallel_levels),
2817 vec![vec!["transform:snv"], vec!["model:pls"]]
2818 );
2819 assert!(plan.node_plans[&model_id]
2820 .controller_capabilities
2821 .contains(&ControllerCapability::EmitsPredictions));
2822 assert!(plan.fold_set.is_some());
2823 let schedule = plan.campaign_phase_schedule(Phase::FitCv).unwrap();
2824 assert_eq!(schedule.scopes.len(), 2);
2825 assert!(schedule.scopes[0].scope_id.starts_with("scope:FIT_CV:"));
2826 assert!(schedule
2827 .scopes
2828 .iter()
2829 .all(|scope| levels_as_strings(&scope.node_levels)
2830 == vec![vec!["transform:snv"], vec!["model:pls"]]));
2831 assert_eq!(
2832 schedule
2833 .scopes
2834 .iter()
2835 .filter_map(|scope| scope.fold_id.as_ref().map(ToString::to_string))
2836 .collect::<Vec<_>>(),
2837 vec!["fold0", "fold1"]
2838 );
2839 assert_eq!(
2840 plan.node_plans
2841 .get(&model_id)
2842 .unwrap()
2843 .controller_id
2844 .as_str(),
2845 "controller:model"
2846 );
2847 assert_eq!(
2848 plan.node_plans.get(&model_id).unwrap().data_bindings.len(),
2849 1
2850 );
2851
2852 let mut bad_plan = plan.clone();
2853 bad_plan.graph_plan.parallel_levels =
2854 vec![vec![model_id], vec![NodeId::new("transform:snv").unwrap()]];
2855 assert!(bad_plan
2856 .validate()
2857 .unwrap_err()
2858 .to_string()
2859 .contains("parallel levels"));
2860
2861 let bad_envelope = ExternalDataPlanEnvelope {
2862 schema_version: crate::data::EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION,
2863 schema_fingerprint: "f97b37872fa22134b508f98fd8e207e5b776b52594fb8f6f5c3e15bee212246b"
2864 .to_string(),
2865 plan_fingerprint: "7c5431d85574b3f337022fa5d25971d5b5cf445b90331b49938f573ff6901e4d"
2866 .to_string(),
2867 relation_fingerprint: None,
2868 data_content_fingerprint: None,
2869 target_content_fingerprint: None,
2870 coordinator_relations: Some(SampleRelationSet {
2871 records: vec![{
2872 let mut relation = SampleRelation::new(
2873 ObservationId::new("obs:outside").unwrap(),
2874 SampleId::new("sample:outside").unwrap(),
2875 );
2876 relation.target_id = Some(TargetId::new("target:outside").unwrap());
2877 relation.source_id = Some("nir".to_string());
2878 relation
2879 }],
2880 }),
2881 };
2882 assert!(plan
2883 .campaign
2884 .validate_data_envelope_relations(&bad_envelope)
2885 .unwrap_err()
2886 .to_string()
2887 .contains("outside fold set"));
2888 }
2889
2890 #[test]
2891 fn planning_refuses_shape_plan_for_unknown_node() {
2892 let campaign = CampaignSpec {
2893 inner_cv: None,
2894 id: "campaign:oof".to_string(),
2895 root_seed: Some(7),
2896 leakage_policy: LeakageUnitPolicy::default(),
2897 aggregation_policy: AggregationPolicy::default(),
2898 split_invocation: None,
2899 generation: Default::default(),
2900 shape_plans: BTreeMap::from([(
2901 NodeId::new("model:missing").unwrap(),
2902 DataModelShapePlan {
2903 node_id: NodeId::new("model:missing").unwrap(),
2904 input_granularity: Granularity::Sample,
2905 target_granularity: Granularity::Sample,
2906 fit_rows: crate::policy::FitBoundary::FoldTrain,
2907 predict_rows: crate::policy::FitBoundary::FoldValidation,
2908 feature_namespace: None,
2909 feature_schema_fingerprint: None,
2910 target_space: "raw".to_string(),
2911 aggregation_policy: AggregationPolicy::default(),
2912 augmentation_policy: crate::policy::AugmentationPolicy::default(),
2913 selection_policy: crate::policy::FeatureSelectionPolicy::default(),
2914 },
2915 )]),
2916 data_bindings: BTreeMap::new(),
2917 branch_view_plans: Vec::new(),
2918 metadata: BTreeMap::new(),
2919 };
2920
2921 assert!(build_execution_plan("plan:oof", graph(), campaign, ®istry()).is_err());
2922 }
2923
2924 #[test]
2925 fn planning_refuses_oof_edge_without_controller_capabilities() {
2926 let mut registry = ControllerRegistry::new();
2927 let mut model_manifest = manifest("controller:model", NodeKind::Model);
2928 model_manifest
2929 .capabilities
2930 .remove(&ControllerCapability::ConsumesOofPredictions);
2931 registry.register(model_manifest).unwrap();
2932
2933 let err = build_execution_plan(
2934 "plan:oof.capability",
2935 oof_graph(),
2936 CampaignSpec {
2937 inner_cv: None,
2938 id: "campaign:oof.capability".to_string(),
2939 root_seed: Some(11),
2940 leakage_policy: Default::default(),
2941 aggregation_policy: Default::default(),
2942 split_invocation: None,
2943 generation: Default::default(),
2944 shape_plans: BTreeMap::new(),
2945 data_bindings: BTreeMap::new(),
2946 branch_view_plans: Vec::new(),
2947 metadata: BTreeMap::new(),
2948 },
2949 ®istry,
2950 )
2951 .unwrap_err();
2952
2953 assert!(err.to_string().contains("consumes_oof_predictions"));
2954 }
2955
2956 #[test]
2957 fn planning_refuses_raw_prediction_sibling_port_into_fitting_node() {
2958 let mut graph = oof_graph();
2959 graph.id = "g:oof.raw-sibling".to_string();
2960 let base = graph
2961 .nodes
2962 .iter_mut()
2963 .find(|node| node.id.as_str() == "model:base")
2964 .unwrap();
2965 base.ports.outputs.push(port("aux", PortKind::Prediction));
2966 let meta = graph
2967 .nodes
2968 .iter_mut()
2969 .find(|node| node.id.as_str() == "model:meta")
2970 .unwrap();
2971 meta.ports.inputs.push(port("aux", PortKind::Prediction));
2972 graph.edges.push(EdgeSpec {
2973 source: PortRef {
2974 node_id: NodeId::new("model:base").unwrap(),
2975 port_name: "aux".to_string(),
2976 },
2977 target: PortRef {
2978 node_id: NodeId::new("model:meta").unwrap(),
2979 port_name: "aux".to_string(),
2980 },
2981 contract: EdgeContract::new(PortKind::Prediction, None),
2982 });
2983
2984 let error = build_execution_plan(
2985 "plan:oof.raw-sibling",
2986 graph,
2987 campaign("campaign:oof.raw-sibling"),
2988 ®istry(),
2989 )
2990 .unwrap_err()
2991 .to_string();
2992
2993 assert!(error.contains("enters fitting controller"));
2994 assert!(error.contains("must require OOF"));
2995 }
2996
2997 #[test]
2998 fn parallel_controller_capability_validation_requires_safe_manifest() {
2999 let mut registry = ControllerRegistry::new();
3000 let mut transform_manifest = manifest("controller:transform", NodeKind::Transform);
3001 transform_manifest
3002 .capabilities
3003 .remove(&ControllerCapability::ThreadSafe);
3004 transform_manifest
3005 .capabilities
3006 .remove(&ControllerCapability::ProcessSafe);
3007 registry.register(transform_manifest).unwrap();
3008 registry
3009 .register(manifest("controller:model", NodeKind::Model))
3010 .unwrap();
3011 let plan = build_execution_plan(
3012 "plan:parallel.capability",
3013 graph(),
3014 CampaignSpec {
3015 inner_cv: None,
3016 id: "campaign:parallel.capability".to_string(),
3017 root_seed: Some(11),
3018 leakage_policy: Default::default(),
3019 aggregation_policy: Default::default(),
3020 split_invocation: None,
3021 generation: Default::default(),
3022 shape_plans: BTreeMap::new(),
3023 data_bindings: BTreeMap::new(),
3024 branch_view_plans: Vec::new(),
3025 metadata: BTreeMap::new(),
3026 },
3027 ®istry,
3028 )
3029 .unwrap();
3030
3031 assert!(plan
3032 .validate_parallel_controller_capabilities(1, Phase::FitCv)
3033 .is_ok());
3034 let err = plan
3035 .validate_parallel_controller_capabilities(2, Phase::FitCv)
3036 .unwrap_err();
3037 assert!(err.to_string().contains("thread_safe or process_safe"));
3038 }
3039
3040 #[test]
3041 fn planning_refuses_generation_override_for_unknown_node() {
3042 let campaign = CampaignSpec {
3043 inner_cv: None,
3044 id: "campaign:oof".to_string(),
3045 root_seed: Some(7),
3046 leakage_policy: LeakageUnitPolicy::default(),
3047 aggregation_policy: AggregationPolicy::default(),
3048 split_invocation: None,
3049 generation: GenerationSpec {
3050 strategy: GenerationStrategy::Cartesian,
3051 dimensions: vec![GenerationDimension {
3052 name: "model_family".to_string(),
3053 choices: vec![GenerationChoice {
3054 label: "pls".to_string(),
3055 value: serde_json::json!("pls"),
3056 param_overrides: vec![GenerationParamOverride {
3057 node_id: NodeId::new("model:missing").unwrap(),
3058 params: BTreeMap::from([(
3059 "n_components".to_string(),
3060 serde_json::json!(8),
3061 )]),
3062 }],
3063 active_subsequence: None,
3064 }],
3065 }],
3066 max_variants: Some(1),
3067 constraints: GenerationConstraints::default(),
3068 },
3069 shape_plans: BTreeMap::new(),
3070 data_bindings: BTreeMap::new(),
3071 branch_view_plans: Vec::new(),
3072 metadata: BTreeMap::new(),
3073 };
3074
3075 let error = build_execution_plan("plan:oof", graph(), campaign, ®istry())
3076 .unwrap_err()
3077 .to_string();
3078
3079 assert!(error.contains("overrides params for unknown node"));
3080 }
3081
3082 #[test]
3083 fn planning_validates_declared_search_space_fingerprint() {
3084 let campaign = CampaignSpec {
3085 inner_cv: None,
3086 id: "campaign:search.fingerprint".to_string(),
3087 root_seed: Some(7),
3088 leakage_policy: LeakageUnitPolicy::default(),
3089 aggregation_policy: AggregationPolicy::default(),
3090 split_invocation: None,
3091 generation: GenerationSpec {
3092 strategy: GenerationStrategy::Cartesian,
3093 dimensions: vec![GenerationDimension {
3094 name: "model_family".to_string(),
3095 choices: vec![GenerationChoice {
3096 label: "pls".to_string(),
3097 value: serde_json::json!("pls"),
3098 param_overrides: vec![GenerationParamOverride {
3099 node_id: NodeId::new("model:pls").unwrap(),
3100 params: BTreeMap::from([(
3101 "n_components".to_string(),
3102 serde_json::json!(8),
3103 )]),
3104 }],
3105 active_subsequence: None,
3106 }],
3107 }],
3108 max_variants: Some(1),
3109 constraints: GenerationConstraints::default(),
3110 },
3111 shape_plans: BTreeMap::new(),
3112 data_bindings: BTreeMap::new(),
3113 branch_view_plans: Vec::new(),
3114 metadata: BTreeMap::new(),
3115 };
3116 let mut graph = graph();
3117 graph.search_space_fingerprint =
3118 Some(generation_spec_fingerprint(&campaign.generation).unwrap());
3119
3120 let plan = build_execution_plan(
3121 "plan:search.fingerprint",
3122 graph.clone(),
3123 campaign.clone(),
3124 ®istry(),
3125 )
3126 .unwrap();
3127 assert_eq!(plan.variants.len(), 1);
3128
3129 graph.search_space_fingerprint = Some("sha256:not-the-generation-spec".to_string());
3130 let error = build_execution_plan("plan:search.fingerprint", graph, campaign, ®istry())
3131 .unwrap_err()
3132 .to_string();
3133 assert!(error.contains("search_space_fingerprint"));
3134 }
3135
3136 #[test]
3137 fn branch_view_lookup_helpers_match_by_branch_id_and_innermost_path() {
3138 use crate::data::{BranchViewMode, DataViewSelector};
3139
3140 let outer = BranchViewPlan {
3141 view_id: "branch_view:outer".to_string(),
3142 branch_id: "branch:outer".to_string(),
3143 mode: BranchViewMode::BySource,
3144 selector: DataViewSelector {
3145 source_ids: vec!["nir".to_string()],
3146 ..Default::default()
3147 },
3148 allow_overlap: false,
3149 metadata: BTreeMap::new(),
3150 };
3151 let inner = BranchViewPlan {
3152 view_id: "branch_view:inner".to_string(),
3153 branch_id: "branch:inner".to_string(),
3154 mode: BranchViewMode::Separation,
3155 selector: DataViewSelector {
3156 source_ids: vec!["chem".to_string()],
3157 ..Default::default()
3158 },
3159 allow_overlap: false,
3160 metadata: BTreeMap::new(),
3161 };
3162 let plans = vec![outer.clone(), inner.clone()];
3163
3164 assert_eq!(
3165 super::branch_view_for_in(&plans, "branch:outer"),
3166 Some(&outer)
3167 );
3168 assert_eq!(
3169 super::branch_view_for_in(&plans, "branch:inner"),
3170 Some(&inner)
3171 );
3172 assert_eq!(super::branch_view_for_in(&plans, "branch:missing"), None);
3173
3174 let path = vec!["branch:outer".to_string(), "branch:inner".to_string()];
3175 assert_eq!(super::branch_view_for_path_in(&plans, &path), Some(&inner));
3177
3178 let path_outer_only = vec!["branch:outer".to_string()];
3179 assert_eq!(
3180 super::branch_view_for_path_in(&plans, &path_outer_only),
3181 Some(&outer)
3182 );
3183
3184 let empty_path: Vec<String> = Vec::new();
3185 assert_eq!(super::branch_view_for_path_in(&plans, &empty_path), None);
3186
3187 let path_no_match = vec!["branch:other".to_string()];
3188 assert_eq!(super::branch_view_for_path_in(&plans, &path_no_match), None);
3189 }
3190}