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> {
294 self.validate()?;
295 let mut roles_by_node = BTreeMap::<NodeId, Vec<TrainingLossRoleReference>>::new();
296 for role in roles {
297 role.validate()?;
298 if !self.node_plans.contains_key(&role.node_id) {
299 return Err(DagMlError::Planning(format!(
300 "training loss references unknown plan node `{}`",
301 role.node_id
302 )));
303 }
304 roles_by_node
305 .entry(role.node_id.clone())
306 .or_default()
307 .push(role);
308 }
309 for node_roles in roles_by_node.values_mut() {
310 node_roles.sort_by(|left, right| {
311 (&left.output_id, &left.phases).cmp(&(&right.output_id, &right.phases))
312 });
313 }
314 for (node_id, node_plan) in &mut self.node_plans {
315 node_plan.training_losses = roles_by_node.remove(node_id).unwrap_or_default();
316 }
317 self.validate()?;
318 Ok(self)
319 }
320
321 pub fn validate(&self) -> Result<()> {
322 self.graph_plan.graph.validate()?;
323 self.campaign.validate()?;
324 if !self.graph_plan.parallel_levels.is_empty()
327 && self.graph_plan.parallel_levels != self.graph_plan.graph.parallel_levels()?
328 {
329 return Err(DagMlError::Planning(
330 "graph plan parallel levels do not match graph".to_string(),
331 ));
332 }
333
334 for (controller_id, manifest) in &self.controller_manifests {
338 manifest.validate()?;
339 if controller_id != &manifest.controller_id {
340 return Err(DagMlError::Planning(format!(
341 "controller manifest keyed `{controller_id}` declares id `{}`",
342 manifest.controller_id
343 )));
344 }
345 }
346
347 let recomputed_graph_fingerprint = stable_json_fingerprint(&self.graph_plan.graph)?;
358 if recomputed_graph_fingerprint != self.graph_fingerprint {
359 return Err(DagMlError::Planning(
360 "execution plan graph_fingerprint does not match the embedded graph".to_string(),
361 ));
362 }
363 let recomputed_campaign_fingerprint = stable_json_fingerprint(&self.campaign)?;
364 if recomputed_campaign_fingerprint != self.campaign_fingerprint {
365 return Err(DagMlError::Planning(
366 "execution plan campaign_fingerprint does not match the embedded campaign"
367 .to_string(),
368 ));
369 }
370 let recomputed_controller_fingerprint =
371 stable_json_fingerprint(&self.controller_manifests)?;
372 if recomputed_controller_fingerprint != self.controller_fingerprint {
373 return Err(DagMlError::Planning(
374 "execution plan controller_fingerprint does not match the embedded controller manifests"
375 .to_string(),
376 ));
377 }
378
379 let graph_node_ids = self
384 .graph_plan
385 .graph
386 .nodes
387 .iter()
388 .map(|node| node.id.clone())
389 .collect::<BTreeSet<_>>();
390 for (node_id, plan) in &self.node_plans {
391 if node_id != &plan.node_id {
392 return Err(DagMlError::Planning(format!(
393 "node plan keyed `{node_id}` declares node_id `{}`",
394 plan.node_id
395 )));
396 }
397 }
398 let plan_node_ids = self.node_plans.keys().cloned().collect::<BTreeSet<_>>();
399 if plan_node_ids != graph_node_ids {
400 return Err(DagMlError::Planning(
401 "execution plan node_plans do not exactly cover the graph node-id set".to_string(),
402 ));
403 }
404
405 if self.graph_plan.topological_order != self.graph_plan.graph.topological_order()? {
410 return Err(DagMlError::Planning(
411 "execution plan topological_order does not match the graph".to_string(),
412 ));
413 }
414
415 let graph_nodes_by_id = self
416 .graph_plan
417 .graph
418 .nodes
419 .iter()
420 .map(|node| (node.id.clone(), node))
421 .collect::<BTreeMap<_, _>>();
422 for (node_id, plan) in &self.node_plans {
423 let graph_node = graph_nodes_by_id
424 .get(node_id)
425 .expect("node_plans keys equal graph node ids");
426 if plan.kind != graph_node.kind {
431 return Err(DagMlError::Planning(format!(
432 "node plan `{node_id}` kind does not match graph node kind"
433 )));
434 }
435 if plan.input_nodes != self.graph_plan.graph.upstream_nodes(node_id)
436 || plan.output_nodes != self.graph_plan.graph.downstream_nodes(node_id)
437 {
438 return Err(DagMlError::Planning(format!(
439 "node plan `{node_id}` input/output adjacency does not match the graph"
440 )));
441 }
442 let manifest = self
443 .controller_manifests
444 .get(&plan.controller_id)
445 .ok_or_else(|| {
446 DagMlError::Planning(format!(
447 "missing controller manifest `{}` for node `{node_id}`",
448 plan.controller_id
449 ))
450 })?;
451 if manifest.operator_kind != plan.kind
456 || manifest.controller_version != plan.controller_version
457 || manifest.supported_phases != plan.supported_phases
458 || manifest.capabilities != plan.controller_capabilities
459 || manifest.fit_scope != plan.fit_scope
460 || manifest.rng_policy != plan.rng_policy
461 || manifest.artifact_policy != plan.artifact_policy
462 {
463 return Err(DagMlError::Planning(format!(
464 "node `{node_id}` node plan does not match controller manifest `{}`",
465 manifest.controller_id
466 )));
467 }
468 for binding in &plan.data_bindings {
469 if binding.node_id != *node_id {
470 return Err(DagMlError::Planning(format!(
471 "node plan `{node_id}` contains data binding for `{}`",
472 binding.node_id
473 )));
474 }
475 binding.validate()?;
476 }
477 validate_data_binding_requirements(node_id, plan, manifest, graph_node)?;
478 validate_node_training_losses(plan)?;
479 let actual_params_fingerprint = stable_json_fingerprint(&plan.params)?;
480 if actual_params_fingerprint != plan.params_fingerprint {
481 return Err(DagMlError::Planning(format!(
482 "node plan `{node_id}` params fingerprint does not match params"
483 )));
484 }
485 if let Some(inner_cv) = &plan.inner_cv {
489 inner_cv.validate().map_err(|error| {
490 DagMlError::Planning(format!(
491 "node plan `{node_id}` has invalid inner_cv: {error}"
492 ))
493 })?;
494 }
495 }
496 self.validate_oof_controller_capabilities()?;
497 if let Some(fold_set) = &self.fold_set {
498 fold_set.validate()?;
499 }
500 if self.variants.is_empty() {
501 return Err(DagMlError::Planning(
502 "execution plan has no variants".to_string(),
503 ));
504 }
505 for variant in &self.variants {
506 variant.validate()?;
507 }
508 Ok(())
509 }
510
511 pub fn validate_parallel_controller_capabilities(
512 &self,
513 max_workers: usize,
514 phase: Phase,
515 ) -> Result<()> {
516 if max_workers <= 1 {
517 return Ok(());
518 }
519 let node_ids = self
520 .node_parallel_levels_for_phase(phase)?
521 .into_iter()
522 .flatten()
523 .collect::<Vec<_>>();
524 for node_id in node_ids {
525 let node_plan = self.node_plans.get(&node_id).ok_or_else(|| {
526 DagMlError::Planning(format!("missing node plan for `{node_id}`"))
527 })?;
528 let manifest = self
529 .controller_manifests
530 .get(&node_plan.controller_id)
531 .ok_or_else(|| {
532 DagMlError::Planning(format!(
533 "missing controller manifest `{}` for node `{}`",
534 node_plan.controller_id, node_plan.node_id
535 ))
536 })?;
537 if !manifest.supports_parallel_invocation() {
538 return Err(DagMlError::Planning(format!(
539 "parallel scheduler with {max_workers} workers requires controller `{}` for node `{}` to declare thread_safe or process_safe",
540 manifest.controller_id, node_plan.node_id
541 )));
542 }
543 }
544 Ok(())
545 }
546
547 fn validate_oof_controller_capabilities(&self) -> Result<()> {
548 for edge in &self.graph_plan.graph.edges {
549 if edge.contract.kind != PortKind::Prediction {
550 continue;
551 }
552 let target_plan = self.node_plans.get(&edge.target.node_id).ok_or_else(|| {
553 DagMlError::Planning(format!(
554 "prediction edge target node `{}` has no node plan",
555 edge.target.node_id
556 ))
557 })?;
558 let target_fits = matches!(
559 target_plan.fit_scope,
560 ControllerFitScope::FoldTrain | ControllerFitScope::FullTrain
561 ) && target_plan
562 .supported_phases
563 .iter()
564 .any(|phase| matches!(phase, Phase::FitCv | Phase::Refit));
565 if target_fits && !edge.contract.requires_oof {
566 return Err(DagMlError::Planning(format!(
567 "prediction edge `{}.{}` -> `{}.{}` enters fitting controller `{}` and must require OOF",
568 edge.source.node_id,
569 edge.source.port_name,
570 edge.target.node_id,
571 edge.target.port_name,
572 target_plan.controller_id
573 )));
574 }
575 if !edge.contract.requires_oof {
576 continue;
577 }
578 let source_plan = self.node_plans.get(&edge.source.node_id).ok_or_else(|| {
579 DagMlError::Planning(format!(
580 "OOF edge source node `{}` has no node plan",
581 edge.source.node_id
582 ))
583 })?;
584 if !source_plan
585 .controller_capabilities
586 .contains(&ControllerCapability::EmitsPredictions)
587 {
588 return Err(DagMlError::Planning(format!(
589 "OOF edge `{}.{}` -> `{}.{}` requires source controller `{}` to declare emits_predictions",
590 edge.source.node_id,
591 edge.source.port_name,
592 edge.target.node_id,
593 edge.target.port_name,
594 source_plan.controller_id
595 )));
596 }
597 if !target_plan
598 .controller_capabilities
599 .contains(&ControllerCapability::ConsumesOofPredictions)
600 {
601 return Err(DagMlError::Planning(format!(
602 "OOF edge `{}.{}` -> `{}.{}` requires target controller `{}` to declare consumes_oof_predictions",
603 edge.source.node_id,
604 edge.source.port_name,
605 edge.target.node_id,
606 edge.target.port_name,
607 target_plan.controller_id
608 )));
609 }
610 }
611 Ok(())
612 }
613
614 pub fn node_parallel_levels_for_phase(&self, phase: Phase) -> Result<Vec<Vec<NodeId>>> {
615 let levels = self
616 .graph_plan
617 .parallel_levels()?
618 .into_iter()
619 .map(|level| {
620 level
621 .into_iter()
622 .filter(|node_id| {
623 self.node_plans
624 .get(node_id)
625 .is_some_and(|node_plan| node_plan.supported_phases.contains(&phase))
626 })
627 .collect::<Vec<_>>()
628 })
629 .filter(|level| !level.is_empty())
630 .collect::<Vec<_>>();
631 Ok(levels)
632 }
633
634 pub fn campaign_phase_schedule(&self, phase: Phase) -> Result<PhaseExecutionSchedule> {
635 self.validate()?;
636 let node_levels = self.node_parallel_levels_for_phase(phase)?;
637 let fold_ids = if phase == Phase::FitCv {
638 self.fold_set
639 .as_ref()
640 .map(|fold_set| {
641 fold_set
642 .folds
643 .iter()
644 .map(|fold| Some(fold.fold_id.clone()))
645 .collect::<Vec<_>>()
646 })
647 .unwrap_or_else(|| vec![None])
648 } else {
649 vec![None]
650 };
651 let mut scopes = Vec::new();
652 for variant in &self.variants {
653 for fold_id in &fold_ids {
654 scopes.push(ExecutionScopePlan {
655 scope_id: execution_scope_id(
656 phase,
657 Some(&variant.variant_id),
658 fold_id.as_ref(),
659 ),
660 phase,
661 variant_id: Some(variant.variant_id.clone()),
662 variant_seed: variant.seed,
663 fold_id: fold_id.clone(),
664 node_levels: node_levels.clone(),
665 });
666 }
667 }
668 Ok(PhaseExecutionSchedule {
669 plan_id: self.id.clone(),
670 phase,
671 scopes,
672 })
673 }
674
675 pub fn branch_view_for(&self, branch_id: &str) -> Option<&BranchViewPlan> {
679 branch_view_for_in(&self.campaign.branch_view_plans, branch_id)
680 }
681
682 pub fn branch_view_for_path(&self, branch_path: &[String]) -> Option<&BranchViewPlan> {
689 branch_view_for_path_in(&self.campaign.branch_view_plans, branch_path)
690 }
691}
692
693fn validate_data_binding_requirements(
694 node_id: &NodeId,
695 plan: &NodePlan,
696 manifest: &ControllerManifest,
697 node: &NodeSpec,
698) -> Result<()> {
699 let branch_view = branch_view_plan_from_node_metadata(node)?;
700 let Some(model_input) = manifest.model_input_spec()? else {
701 for binding in &plan.data_bindings {
702 let effective_source_ids = effective_binding_source_ids(binding, branch_view.as_ref())?;
703 if effective_source_ids.len() > 1 {
704 return Err(data_requirement_refusal(
705 "dagml.data_requirement.missing_data_requirements",
706 node_id,
707 binding,
708 manifest,
709 "multisource",
710 &effective_source_ids,
711 "multi-source data binding requires controller data_requirements".to_string(),
712 ));
713 }
714 }
715 return Ok(());
716 };
717 for binding in &plan.data_bindings {
718 let Some(port) = model_input
719 .ports
720 .iter()
721 .find(|port| port.name == binding.input_name)
722 else {
723 return Err(DagMlError::Planning(format!(
724 "node `{node_id}` data binding `{}` is not declared by controller `{}` data_requirements",
725 binding.input_name, manifest.controller_id
726 )));
727 };
728 if !port
729 .accepted_representations
730 .iter()
731 .any(|representation| representation == &binding.output_representation)
732 {
733 return Err(DagMlError::Planning(format!(
734 "node `{node_id}` data binding `{}` output representation `{}` is not accepted by controller `{}` data_requirements port `{}`",
735 binding.input_name,
736 binding.output_representation,
737 manifest.controller_id,
738 port.name
739 )));
740 }
741 if let Some(type_id) = representation_type_id(&binding.output_representation) {
742 if !port
743 .accepted_types
744 .iter()
745 .any(|accepted_type| accepted_type.as_str() == type_id)
746 {
747 return Err(DagMlError::Planning(format!(
748 "node `{node_id}` data binding `{}` output representation `{}` has registered type `{type_id}` but controller `{}` data_requirements port `{}` accepts types {:?}",
749 binding.input_name,
750 binding.output_representation,
751 manifest.controller_id,
752 port.name,
753 port.accepted_types
754 )));
755 }
756 }
757 validate_data_binding_source_shape(
758 node_id,
759 binding,
760 port,
761 &model_input,
762 manifest,
763 branch_view.as_ref(),
764 )?;
765 }
766 Ok(())
767}
768
769fn validate_node_training_losses(plan: &NodePlan) -> Result<()> {
770 let mut previous_key: Option<(Option<String>, BTreeSet<Phase>)> = None;
771 let mut occupied_phases = BTreeSet::new();
772 for role in &plan.training_losses {
773 role.validate()?;
774 if role.node_id != plan.node_id {
775 return Err(DagMlError::Planning(format!(
776 "node plan `{}` contains training loss for `{}`",
777 plan.node_id, role.node_id
778 )));
779 }
780 let key = (role.output_id.clone(), role.phases.clone());
781 if previous_key
782 .as_ref()
783 .is_some_and(|previous| previous >= &key)
784 {
785 return Err(DagMlError::Planning(format!(
786 "node plan `{}` training losses must be strictly sorted by output_id and phases",
787 plan.node_id
788 )));
789 }
790 previous_key = Some(key);
791 for phase in &role.phases {
792 if !plan.supported_phases.contains(phase) {
793 return Err(DagMlError::Planning(format!(
794 "node `{}` has a training loss for unsupported phase {phase:?}",
795 plan.node_id
796 )));
797 }
798 if !occupied_phases.insert((role.output_id.clone(), *phase)) {
799 return Err(DagMlError::Planning(format!(
800 "node `{}` has overlapping training losses for output {:?} in phase {phase:?}",
801 plan.node_id, role.output_id
802 )));
803 }
804 }
805 if !plan
806 .controller_capabilities
807 .contains(&ControllerCapability::SupportsConfigurableLoss)
808 {
809 return Err(DagMlError::Planning(format!(
810 "node `{}` configures a training loss but its controller does not support configurable loss",
811 plan.node_id
812 )));
813 }
814 if role.loss.spec.kind == SemanticSpecKind::Custom
815 && !plan
816 .controller_capabilities
817 .contains(&ControllerCapability::SupportsCustomLoss)
818 {
819 return Err(DagMlError::Planning(format!(
820 "node `{}` configures a custom loss but its controller does not support custom loss",
821 plan.node_id
822 )));
823 }
824 if role
825 .loss
826 .spec
827 .capabilities
828 .contains(&LossCapability::Differentiable)
829 && !plan
830 .controller_capabilities
831 .contains(&ControllerCapability::SupportsDifferentiableLoss)
832 {
833 return Err(DagMlError::Planning(format!(
834 "node `{}` configures a differentiable loss but its controller does not support differentiable loss",
835 plan.node_id
836 )));
837 }
838 if role
839 .loss
840 .spec
841 .required_inputs
842 .contains(&CriterionInput::SampleWeight)
843 && !plan
844 .controller_capabilities
845 .contains(&ControllerCapability::SupportsSampleWeights)
846 {
847 return Err(DagMlError::Planning(format!(
848 "node `{}` loss requires sample weights but its controller does not support them",
849 plan.node_id
850 )));
851 }
852 if role
853 .loss
854 .spec
855 .required_inputs
856 .contains(&CriterionInput::MissingMask)
857 && !plan
858 .controller_capabilities
859 .contains(&ControllerCapability::SupportsMissingMasks)
860 {
861 return Err(DagMlError::Planning(format!(
862 "node `{}` loss requires missing masks but its controller does not support them",
863 plan.node_id
864 )));
865 }
866 if role
867 .loss
868 .implementation
869 .capabilities
870 .contains(&ImplementationCapability::NeedsGil)
871 && !plan
872 .controller_capabilities
873 .contains(&ControllerCapability::NeedsPythonGil)
874 {
875 return Err(DagMlError::Planning(format!(
876 "node `{}` loss implementation needs the Python GIL but its controller does not declare it",
877 plan.node_id
878 )));
879 }
880 }
881 Ok(())
882}
883
884fn branch_view_plan_from_node_metadata(node: &NodeSpec) -> Result<Option<BranchViewPlan>> {
885 let Some(value) = node.metadata.get("dsl_branch_view_plan") else {
886 return Ok(None);
887 };
888 let plan: BranchViewPlan = serde_json::from_value(value.clone()).map_err(|error| {
889 DagMlError::Planning(format!(
890 "node `{}` carries malformed `dsl_branch_view_plan` metadata: {error}",
891 node.id
892 ))
893 })?;
894 plan.validate()
895 .map_err(|error| DagMlError::Planning(error.to_string()))?;
896 Ok(Some(plan))
897}
898
899fn validate_data_binding_source_shape(
900 node_id: &NodeId,
901 binding: &DataBinding,
902 port: &ModelInputPortSpec,
903 model_input: &ModelInputSpec,
904 manifest: &ControllerManifest,
905 branch_view: Option<&BranchViewPlan>,
906) -> Result<()> {
907 let effective_source_ids = effective_binding_source_ids(binding, branch_view)?;
908 if effective_source_ids.len() < 2 {
909 return Ok(());
910 }
911 if !port.multi_source {
912 return Err(data_requirement_refusal(
913 "dagml.data_requirement.multi_source_port_not_supported",
914 node_id,
915 binding,
916 manifest,
917 "multisource",
918 &effective_source_ids,
919 format!(
920 "controller `{}` data_requirements port `{}` does not declare multi_source=true",
921 manifest.controller_id, port.name
922 ),
923 ));
924 }
925 let Some(fusion) = &model_input.default_fusion else {
926 return Err(data_requirement_refusal(
927 "dagml.data_requirement.missing_multisource_fusion",
928 node_id,
929 binding,
930 manifest,
931 "multisource",
932 &effective_source_ids,
933 "multi-source data binding requires an explicit default_fusion policy".to_string(),
934 ));
935 };
936 validate_fusion_sources_match_binding(
937 node_id,
938 binding,
939 manifest,
940 fusion.representation_plan.as_ref(),
941 &effective_source_ids,
942 )?;
943 match fusion.mode {
944 ModelInputFusionMode::ConcatenateFeatures => {
945 if binding
946 .metadata
947 .get(SOURCE_INDEX_METADATA_KEY)
948 .and_then(serde_json::Value::as_object)
949 .is_none()
950 {
951 return Err(data_requirement_refusal(
952 "dagml.data_requirement.source_concat_requires_source_index",
953 node_id,
954 binding,
955 manifest,
956 "source_concat",
957 &effective_source_ids,
958 "source-concat feature fusion requires data binding metadata.source_index so feature-axis blocks are explicit".to_string(),
959 ));
960 }
961 }
962 ModelInputFusionMode::DictBySource | ModelInputFusionMode::Custom => {}
963 ModelInputFusionMode::SingleSource | ModelInputFusionMode::StackSamples => {
964 let fusion_mode = fusion_mode_label(fusion.mode);
965 return Err(data_requirement_refusal(
966 "dagml.data_requirement.unsupported_multisource_fusion_mode",
967 node_id,
968 binding,
969 manifest,
970 fusion_mode,
971 &effective_source_ids,
972 format!(
973 "multi-source data binding cannot be planned with default_fusion.mode={fusion_mode}"
974 ),
975 ));
976 }
977 }
978 Ok(())
979}
980
981fn fusion_mode_label(mode: ModelInputFusionMode) -> &'static str {
982 match mode {
983 ModelInputFusionMode::SingleSource => "single_source",
984 ModelInputFusionMode::ConcatenateFeatures => "concatenate_features",
985 ModelInputFusionMode::StackSamples => "stack_samples",
986 ModelInputFusionMode::DictBySource => "dict_by_source",
987 ModelInputFusionMode::Custom => "custom",
988 }
989}
990
991fn effective_binding_source_ids(
992 binding: &DataBinding,
993 branch_view: Option<&BranchViewPlan>,
994) -> Result<Vec<String>> {
995 let Some(branch_view) = branch_view else {
996 return Ok(binding.source_ids.clone());
997 };
998 if branch_view.mode != BranchViewMode::BySource {
999 return Ok(binding.source_ids.clone());
1000 }
1001 if branch_view.selector.source_ids.len() != 1 {
1002 return Err(data_requirement_refusal_for_branch(
1003 "dagml.data_requirement.unsupported_by_source_shape",
1004 binding,
1005 branch_view,
1006 "by_source branch views must select exactly one source_id for per-source X-chain fit semantics".to_string(),
1007 ));
1008 }
1009 if !binding.source_ids.is_empty() {
1010 let declared = binding.source_ids.iter().collect::<BTreeSet<_>>();
1011 for source_id in &branch_view.selector.source_ids {
1012 if !declared.contains(source_id) {
1013 return Err(data_requirement_refusal_for_branch(
1014 "dagml.data_requirement.by_source_selector_outside_binding",
1015 binding,
1016 branch_view,
1017 format!(
1018 "by_source branch selector source `{source_id}` is not declared by data binding source_ids"
1019 ),
1020 ));
1021 }
1022 }
1023 }
1024 Ok(branch_view.selector.source_ids.clone())
1025}
1026
1027fn validate_fusion_sources_match_binding(
1028 node_id: &NodeId,
1029 binding: &DataBinding,
1030 manifest: &ControllerManifest,
1031 representation_plan: Option<&RepresentationPlan>,
1032 effective_source_ids: &[String],
1033) -> Result<()> {
1034 let Some(representation_plan) = representation_plan else {
1035 return Ok(());
1036 };
1037 let component_sources = representation_plan_component_sources(representation_plan);
1038 if component_sources.is_empty() {
1039 return Ok(());
1040 }
1041 let declared = component_sources.iter().cloned().collect::<BTreeSet<_>>();
1042 let effective = effective_source_ids.iter().collect::<BTreeSet<_>>();
1043 if declared != effective {
1044 return Err(data_requirement_refusal(
1045 "dagml.data_requirement.representation_sources_mismatch",
1046 node_id,
1047 binding,
1048 manifest,
1049 "multisource",
1050 effective_source_ids,
1051 format!(
1052 "default_fusion.representation_plan component_source_ids {:?} do not match binding source_ids {:?}",
1053 component_sources, effective_source_ids
1054 ),
1055 ));
1056 }
1057 Ok(())
1058}
1059
1060fn representation_plan_component_sources(plan: &RepresentationPlan) -> Vec<&String> {
1061 match plan {
1062 RepresentationPlan::Aggregate(_) => Vec::new(),
1063 RepresentationPlan::CartesianProduct(plan) => {
1064 plan.combination_plan.component_source_ids.iter().collect()
1065 }
1066 RepresentationPlan::MonteCarloCartesian(plan) => {
1067 plan.combination_plan.component_source_ids.iter().collect()
1068 }
1069 RepresentationPlan::StackFixed(plan) => plan.component_source_ids.iter().collect(),
1070 RepresentationPlan::StackPaddedMasked(plan) => plan.component_source_ids.iter().collect(),
1071 }
1072}
1073
1074fn data_requirement_refusal(
1075 code: &'static str,
1076 node_id: &NodeId,
1077 binding: &DataBinding,
1078 manifest: &ControllerManifest,
1079 shape: &str,
1080 source_ids: &[String],
1081 message: String,
1082) -> DagMlError {
1083 DagMlError::Planning(format!(
1084 "data requirement refusal: {}",
1085 serde_json::json!({
1086 "schema_version": 1,
1087 "code": code,
1088 "node_id": node_id.to_string(),
1089 "input_name": binding.input_name.as_str(),
1090 "controller_id": manifest.controller_id.to_string(),
1091 "shape": shape,
1092 "source_ids": source_ids,
1093 "message": message
1094 })
1095 ))
1096}
1097
1098fn data_requirement_refusal_for_branch(
1099 code: &'static str,
1100 binding: &DataBinding,
1101 branch_view: &BranchViewPlan,
1102 message: String,
1103) -> DagMlError {
1104 DagMlError::Planning(format!(
1105 "data requirement refusal: {}",
1106 serde_json::json!({
1107 "schema_version": 1,
1108 "code": code,
1109 "node_id": binding.node_id.to_string(),
1110 "input_name": binding.input_name.as_str(),
1111 "branch_view_id": branch_view.view_id.as_str(),
1112 "branch_id": branch_view.branch_id.as_str(),
1113 "shape": "by_source",
1114 "source_ids": &branch_view.selector.source_ids,
1115 "message": message
1116 })
1117 ))
1118}
1119
1120fn branch_view_for_in<'a>(
1121 plans: &'a [BranchViewPlan],
1122 branch_id: &str,
1123) -> Option<&'a BranchViewPlan> {
1124 plans.iter().find(|plan| plan.branch_id == branch_id)
1125}
1126
1127fn branch_view_for_path_in<'a>(
1128 plans: &'a [BranchViewPlan],
1129 branch_path: &[String],
1130) -> Option<&'a BranchViewPlan> {
1131 for branch_id in branch_path.iter().rev() {
1132 if let Some(plan) = branch_view_for_in(plans, branch_id) {
1133 return Some(plan);
1134 }
1135 }
1136 None
1137}
1138
1139fn execution_scope_id(
1140 phase: Phase,
1141 variant_id: Option<&VariantId>,
1142 fold_id: Option<&FoldId>,
1143) -> String {
1144 format!(
1145 "scope:{}:{}:{}",
1146 phase_scope_label(phase),
1147 variant_id
1148 .map(ToString::to_string)
1149 .unwrap_or_else(|| "base".to_string()),
1150 fold_id
1151 .map(ToString::to_string)
1152 .unwrap_or_else(|| "nofold".to_string())
1153 )
1154}
1155
1156fn phase_scope_label(phase: Phase) -> &'static str {
1157 match phase {
1158 Phase::Compile => "COMPILE",
1159 Phase::Plan => "PLAN",
1160 Phase::FitCv => "FIT_CV",
1161 Phase::Select => "SELECT",
1162 Phase::Refit => "REFIT",
1163 Phase::Predict => "PREDICT",
1164 Phase::Explain => "EXPLAIN",
1165 }
1166}
1167
1168pub fn build_execution_plan(
1169 id: impl Into<String>,
1170 graph: GraphSpec,
1171 campaign: CampaignSpec,
1172 registry: &ControllerRegistry,
1173) -> Result<ExecutionPlan> {
1174 let id = id.into();
1175 if id.trim().is_empty() {
1176 return Err(DagMlError::Planning(
1177 "execution plan id is empty".to_string(),
1178 ));
1179 }
1180 campaign.validate()?;
1181 let graph_plan = GraphPlan::from_graph(graph)?;
1182 validate_campaign_node_targets(&graph_plan.graph, &campaign)?;
1183
1184 let mut node_plans = BTreeMap::new();
1185 let mut controller_manifests = BTreeMap::new();
1186 for node_id in &graph_plan.topological_order {
1187 let node = graph_plan
1188 .graph
1189 .nodes
1190 .iter()
1191 .find(|node| &node.id == node_id)
1192 .expect("topological node exists");
1193 let manifest = registry.resolve_for_node(node)?;
1194 let params = node.params.clone();
1195 let params_fingerprint = stable_json_fingerprint(¶ms)?;
1196 let inner_cv = match node.metadata.get("dsl_inner_cv") {
1200 Some(value) => {
1201 let spec =
1202 serde_json::from_value::<NestedCvSpec>(value.clone()).map_err(|error| {
1203 DagMlError::Planning(format!(
1204 "node `{}` has invalid dsl_inner_cv metadata: {error}",
1205 node.id
1206 ))
1207 })?;
1208 spec.validate().map_err(|error| {
1211 DagMlError::Planning(format!(
1212 "node `{}` has invalid dsl_inner_cv metadata: {error}",
1213 node.id
1214 ))
1215 })?;
1216 Some(spec)
1217 }
1218 None => None,
1219 };
1220 let shape_plan = campaign.shape_plans.get(&node.id).cloned();
1221 let data_bindings = campaign
1222 .data_bindings
1223 .get(&node.id)
1224 .cloned()
1225 .unwrap_or_default();
1226 node_plans.insert(
1227 node.id.clone(),
1228 NodePlan {
1229 inner_cv,
1230 node_id: node.id.clone(),
1231 kind: node.kind.clone(),
1232 controller_id: manifest.controller_id.clone(),
1233 controller_version: manifest.controller_version.clone(),
1234 supported_phases: manifest.supported_phases.clone(),
1235 controller_capabilities: manifest.capabilities.clone(),
1236 training_losses: Vec::new(),
1237 fit_scope: manifest.fit_scope,
1238 rng_policy: manifest.rng_policy,
1239 artifact_policy: manifest.artifact_policy,
1240 input_nodes: graph_plan.graph.upstream_nodes(&node.id),
1241 output_nodes: graph_plan.graph.downstream_nodes(&node.id),
1242 shape_plan,
1243 data_bindings,
1244 params,
1245 params_fingerprint,
1246 },
1247 );
1248 controller_manifests.insert(manifest.controller_id.clone(), manifest);
1249 }
1250
1251 let fold_set = campaign
1252 .split_invocation
1253 .as_ref()
1254 .and_then(|split| split.fold_set.clone());
1255 validate_search_space_fingerprint(&graph_plan.graph, &campaign)?;
1256 let variants = enumerate_variants(&campaign.generation, campaign.root_seed)?;
1257 validate_generation_override_targets(&graph_plan.graph, &variants)?;
1258 let graph_fingerprint = stable_json_fingerprint(&graph_plan.graph)?;
1259 let campaign_fingerprint = stable_json_fingerprint(&campaign)?;
1260 let controller_fingerprint = stable_json_fingerprint(&controller_manifests)?;
1261 let plan = ExecutionPlan {
1262 id,
1263 graph_plan,
1264 campaign,
1265 node_plans,
1266 controller_manifests,
1267 variants,
1268 fold_set,
1269 graph_fingerprint,
1270 campaign_fingerprint,
1271 controller_fingerprint,
1272 };
1273 plan.validate()?;
1274 Ok(plan)
1275}
1276
1277fn validate_search_space_fingerprint(graph: &GraphSpec, campaign: &CampaignSpec) -> Result<()> {
1278 let Some(expected_fingerprint) = &graph.search_space_fingerprint else {
1279 return Ok(());
1280 };
1281 if expected_fingerprint.trim().is_empty() {
1282 return Err(DagMlError::Planning(format!(
1283 "graph `{}` has empty search_space_fingerprint",
1284 graph.id
1285 )));
1286 }
1287 let actual_fingerprint = generation_spec_fingerprint(&campaign.generation)?;
1288 if expected_fingerprint != &actual_fingerprint {
1289 return Err(DagMlError::Planning(format!(
1290 "graph `{}` search_space_fingerprint does not match campaign generation spec",
1291 graph.id
1292 )));
1293 }
1294 Ok(())
1295}
1296
1297fn validate_generation_override_targets(graph: &GraphSpec, variants: &[VariantPlan]) -> Result<()> {
1298 let node_ids = graph
1299 .nodes
1300 .iter()
1301 .map(|node| node.id.clone())
1302 .collect::<BTreeSet<_>>();
1303 for variant in variants {
1304 for node_id in variant.param_override_targets()? {
1305 if !node_ids.contains(&node_id) {
1306 return Err(DagMlError::Planning(format!(
1307 "variant `{}` overrides params for unknown node `{node_id}`",
1308 variant.variant_id
1309 )));
1310 }
1311 }
1312 }
1313 Ok(())
1314}
1315
1316fn validate_campaign_node_targets(graph: &GraphSpec, campaign: &CampaignSpec) -> Result<()> {
1317 let node_ids = graph
1318 .nodes
1319 .iter()
1320 .map(|node| &node.id)
1321 .collect::<BTreeSet<_>>();
1322 for node_id in campaign.shape_plans.keys() {
1323 if !node_ids.contains(node_id) {
1324 return Err(DagMlError::Planning(format!(
1325 "shape plan references unknown node `{node_id}`"
1326 )));
1327 }
1328 }
1329 for node_id in campaign.data_bindings.keys() {
1330 if !node_ids.contains(node_id) {
1331 return Err(DagMlError::Planning(format!(
1332 "data binding references unknown node `{node_id}`"
1333 )));
1334 }
1335 }
1336 Ok(())
1337}
1338
1339pub fn prune_plan_to_active(
1371 plan: &ExecutionPlan,
1372 active_nodes: &BTreeSet<NodeId>,
1373 all_choice_nodes: &BTreeSet<NodeId>,
1374 variant: &VariantPlan,
1375) -> Result<ExecutionPlan> {
1376 plan.validate()?;
1377 variant.validate()?;
1378 for node_id in active_nodes {
1379 if !plan.node_plans.contains_key(node_id) {
1380 return Err(DagMlError::Planning(format!(
1381 "operator-SELECT prune: active node `{node_id}` is not in the union plan"
1382 )));
1383 }
1384 }
1385 if active_nodes.is_empty() {
1386 return Err(DagMlError::Planning(
1387 "operator-SELECT prune: active node set is empty".to_string(),
1388 ));
1389 }
1390
1391 let graph = &plan.graph_plan.graph;
1394 let mut ancestors = BTreeSet::<NodeId>::new();
1395 let mut stack: Vec<NodeId> = active_nodes.iter().cloned().collect();
1396 while let Some(node_id) = stack.pop() {
1397 for upstream in graph.upstream_nodes(&node_id) {
1398 if ancestors.insert(upstream.clone()) {
1399 stack.push(upstream);
1400 }
1401 }
1402 }
1403 let shared_prefix = ancestors
1404 .into_iter()
1405 .filter(|node_id| !all_choice_nodes.contains(node_id))
1406 .collect::<BTreeSet<_>>();
1407
1408 let keep = shared_prefix
1409 .iter()
1410 .chain(active_nodes.iter())
1411 .cloned()
1412 .collect::<BTreeSet<_>>();
1413
1414 let mut pruned_graph = graph.clone();
1417 pruned_graph.nodes.retain(|node| keep.contains(&node.id));
1418 pruned_graph
1419 .edges
1420 .retain(|edge| keep.contains(&edge.source.node_id) && keep.contains(&edge.target.node_id));
1421
1422 let graph_plan = GraphPlan::from_graph(pruned_graph)?;
1423
1424 let mut node_plans = BTreeMap::new();
1427 for (node_id, node_plan) in &plan.node_plans {
1428 if !keep.contains(node_id) {
1429 continue;
1430 }
1431 let mut pruned_node_plan = node_plan.clone();
1432 pruned_node_plan.input_nodes = graph_plan.graph.upstream_nodes(node_id);
1433 pruned_node_plan.output_nodes = graph_plan.graph.downstream_nodes(node_id);
1434 node_plans.insert(node_id.clone(), pruned_node_plan);
1435 }
1436
1437 let graph_fingerprint = stable_json_fingerprint(&graph_plan.graph)?;
1438 let pruned = ExecutionPlan {
1439 id: plan.id.clone(),
1440 graph_plan,
1441 campaign: plan.campaign.clone(),
1442 node_plans,
1443 controller_manifests: plan.controller_manifests.clone(),
1444 variants: vec![variant.clone()],
1445 fold_set: plan.fold_set.clone(),
1446 graph_fingerprint,
1447 campaign_fingerprint: plan.campaign_fingerprint.clone(),
1448 controller_fingerprint: plan.controller_fingerprint.clone(),
1449 };
1450 pruned.validate()?;
1451 validate_active_inputs(&pruned, graph)?;
1452 Ok(pruned)
1453}
1454
1455fn validate_active_inputs(plan: &ExecutionPlan, union_graph: &GraphSpec) -> Result<()> {
1469 let pruned_graph = &plan.graph_plan.graph;
1470 let kept: BTreeSet<&NodeId> = pruned_graph.nodes.iter().map(|node| &node.id).collect();
1471
1472 let mut union_port_sources = BTreeMap::<(NodeId, String), usize>::new();
1473 for edge in &union_graph.edges {
1474 if !kept.contains(&edge.target.node_id) {
1475 continue;
1476 }
1477 *union_port_sources
1478 .entry((edge.target.node_id.clone(), edge.target.port_name.clone()))
1479 .or_insert(0) += 1;
1480 }
1481 let mut pruned_port_sources = BTreeMap::<(NodeId, String), usize>::new();
1482 for edge in &pruned_graph.edges {
1483 *pruned_port_sources
1484 .entry((edge.target.node_id.clone(), edge.target.port_name.clone()))
1485 .or_insert(0) += 1;
1486 }
1487
1488 for (key, union_count) in &union_port_sources {
1489 let pruned_count = pruned_port_sources.get(key).copied().unwrap_or(0);
1490 let (node_id, port_name) = key;
1491 if *union_count >= 1 && pruned_count == 0 {
1492 return Err(DagMlError::Planning(format!(
1493 "operator-SELECT prune left node `{node_id}` required input port `{port_name}` with zero surviving sources (dangling): its producer was pruned away"
1494 )));
1495 }
1496 if pruned_count > 1 {
1497 return Err(DagMlError::Planning(format!(
1498 "operator-SELECT prune left node `{node_id}` input port `{port_name}` fed by {pruned_count} surviving sources (ambiguous)"
1499 )));
1500 }
1501 }
1502 Ok(())
1503}
1504
1505#[cfg(test)]
1506mod tests {
1507 use std::collections::{BTreeMap, BTreeSet};
1508 use std::time::{Duration, Instant};
1509
1510 use super::*;
1511 use crate::controller::{
1512 ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest, RngPolicy,
1513 };
1514 use crate::fold::FoldPartitionMode;
1515
1516 #[test]
1517 fn params_fingerprint_pins_serde_json_binary64_spelling() {
1518 let params = BTreeMap::from([
1519 (
1520 "scope".to_string(),
1521 serde_json::Value::String("train_only".to_string()),
1522 ),
1523 ("std".to_string(), serde_json::Value::from(1e-7_f64)),
1524 ]);
1525 assert_eq!(
1526 stable_json_fingerprint(¶ms).unwrap(),
1527 "3f417903752f65005bc9b69bcd23dfcf3ede2cda010e4f1ead6090a1a407b851"
1528 );
1529
1530 for (value, expected) in [
1533 (1e-7_f64, "1e-7"),
1534 (1e-6_f64, "1e-6"),
1535 (1e-5_f64, "0.00001"),
1536 (1e20_f64, "1e+20"),
1537 (1e21_f64, "1e+21"),
1538 (-0.0_f64, "-0.0"),
1539 (0.1_f64, "0.1"),
1540 (2.0_f64, "2.0"),
1541 (f64::from_bits(1), "5e-324"),
1542 ] {
1543 assert_eq!(serde_json::to_string(&value).unwrap(), expected);
1544 }
1545 }
1546
1547 #[test]
1548 fn inner_cv_is_declarable_at_campaign_and_node_level() {
1549 let campaign_json = r#"{"id":"c","root_seed":null,"inner_cv":{"kind":"kfold","n_splits":3,"shuffle":false,"seed":5}}"#;
1551 let campaign: CampaignSpec = serde_json::from_str(campaign_json).unwrap();
1552 campaign.validate().unwrap();
1553 assert!(campaign.inner_cv.is_some());
1554
1555 let node_inner = crate::fold::NestedCvSpec::KFold(crate::fold::KFoldSpec {
1557 n_splits: 4,
1558 shuffle: false,
1559 seed: Some(6),
1560 });
1561 let resolved = crate::fold::resolve_inner_cv(Some(&node_inner), campaign.inner_cv.as_ref());
1562 assert_eq!(resolved, Some(&node_inner));
1563
1564 let bare = r#"{"id":"c","root_seed":null}"#;
1566 let bare_campaign: CampaignSpec = serde_json::from_str(bare).unwrap();
1567 assert!(bare_campaign.inner_cv.is_none());
1568 let reserialized = serde_json::to_string(&bare_campaign).unwrap();
1569 assert!(!reserialized.contains("inner_cv"));
1570
1571 let bad: CampaignSpec = serde_json::from_str(
1574 r#"{"id":"c","root_seed":null,"inner_cv":{"kind":"kfold","n_splits":1,"shuffle":false,"seed":null}}"#,
1575 )
1576 .unwrap();
1577 let error = bad.validate().unwrap_err();
1578 assert!(error.to_string().contains("at least two splits"));
1579 }
1580
1581 #[test]
1582 fn execution_plan_validate_rejects_invalid_node_local_inner_cv() {
1583 let campaign = CampaignSpec {
1587 inner_cv: None,
1588 id: "campaign:plan-validate".to_string(),
1589 root_seed: Some(7),
1590 leakage_policy: LeakageUnitPolicy::default(),
1591 aggregation_policy: AggregationPolicy::default(),
1592 split_invocation: None,
1593 generation: Default::default(),
1594 shape_plans: BTreeMap::new(),
1595 data_bindings: BTreeMap::new(),
1596 branch_view_plans: Vec::new(),
1597 metadata: BTreeMap::new(),
1598 };
1599 let mut plan =
1600 build_execution_plan("plan:validate", graph(), campaign, ®istry()).unwrap();
1601 plan.validate().unwrap();
1602 plan.node_plans
1603 .get_mut(&NodeId::new("model:pls").unwrap())
1604 .unwrap()
1605 .inner_cv = Some(crate::fold::NestedCvSpec::KFold(crate::fold::KFoldSpec {
1606 n_splits: 1,
1607 shuffle: false,
1608 seed: None,
1609 }));
1610 let error = plan.validate().unwrap_err();
1611 assert!(matches!(error, DagMlError::Planning(_)));
1612 assert!(error.to_string().contains("invalid inner_cv"));
1613 assert!(error.to_string().contains("at least two splits"));
1614 }
1615
1616 fn hardening_plan() -> ExecutionPlan {
1617 let campaign = CampaignSpec {
1618 inner_cv: None,
1619 id: "campaign:harden".to_string(),
1620 root_seed: Some(7),
1621 leakage_policy: LeakageUnitPolicy::default(),
1622 aggregation_policy: AggregationPolicy::default(),
1623 split_invocation: None,
1624 generation: Default::default(),
1625 shape_plans: BTreeMap::new(),
1626 data_bindings: BTreeMap::new(),
1627 branch_view_plans: Vec::new(),
1628 metadata: BTreeMap::new(),
1629 };
1630 build_execution_plan("plan:harden", graph(), campaign, ®istry()).unwrap()
1631 }
1632
1633 #[test]
1634 fn execution_plan_external_reader_rejects_positional_struct_sequences() {
1635 let plan = hardening_plan();
1636 let mut wire = serde_json::to_value(&plan).unwrap();
1637 wire["campaign"]["leakage_policy"] = serde_json::json!([]);
1638
1639 let permissive: ExecutionPlan = serde_json::from_value(wire.clone()).unwrap();
1642 permissive.validate().unwrap();
1643
1644 let error = ExecutionPlan::from_json(&serde_json::to_string(&wire).unwrap()).unwrap_err();
1646 assert!(error.to_string().contains("must use a JSON object"));
1647 assert!(error.to_string().contains("campaign.leakage_policy"));
1648 }
1649
1650 #[test]
1651 fn execution_plan_external_reader_preserves_typed_serde_compatibility() {
1652 let plan = hardening_plan();
1653 let mut wire = serde_json::to_value(&plan).unwrap();
1654 wire["graph_plan"]
1655 .as_object_mut()
1656 .unwrap()
1657 .remove("parallel_levels");
1658 wire["graph_plan"]["graph"]
1659 .as_object_mut()
1660 .unwrap()
1661 .insert("forward_compatible".to_string(), serde_json::json!(true));
1662
1663 let parsed = ExecutionPlan::from_json(&serde_json::to_string(&wire).unwrap()).unwrap();
1664 assert!(parsed.graph_plan.parallel_levels.is_empty());
1665 assert_eq!(parsed.graph_plan.graph, plan.graph_plan.graph);
1666 }
1667
1668 #[test]
1669 fn validate_rejects_manifest_keyed_under_foreign_controller_id() {
1670 let mut plan = hardening_plan();
1671 plan.validate().unwrap();
1672 plan.controller_manifests
1673 .get_mut(&ControllerId::new("controller:model").unwrap())
1674 .unwrap()
1675 .controller_id = ControllerId::new("controller:imposter").unwrap();
1676 let error = plan.validate().unwrap_err();
1677 assert!(matches!(error, DagMlError::Planning(_)));
1678 assert!(error.to_string().contains("declares id"));
1679 }
1680
1681 #[test]
1682 fn validate_rejects_node_plan_keyed_under_foreign_node_id() {
1683 let mut plan = hardening_plan();
1684 plan.validate().unwrap();
1685 plan.node_plans
1686 .get_mut(&NodeId::new("model:pls").unwrap())
1687 .unwrap()
1688 .node_id = NodeId::new("transform:snv").unwrap();
1689 let error = plan.validate().unwrap_err();
1690 assert!(matches!(error, DagMlError::Planning(_)));
1691 assert!(error.to_string().contains("declares node_id"));
1692 }
1693
1694 #[test]
1695 fn validate_rejects_node_plans_not_covering_graph_node_set() {
1696 let mut plan = hardening_plan();
1697 plan.validate().unwrap();
1698 let mut ghost = plan
1701 .node_plans
1702 .remove(&NodeId::new("model:pls").unwrap())
1703 .unwrap();
1704 let bogus = NodeId::new("model:ghost").unwrap();
1705 ghost.node_id = bogus.clone();
1706 plan.node_plans.insert(bogus, ghost);
1707 let error = plan.validate().unwrap_err();
1708 assert!(matches!(error, DagMlError::Planning(_)));
1709 assert!(error
1710 .to_string()
1711 .contains("do not exactly cover the graph node-id set"));
1712 }
1713
1714 #[test]
1715 fn validate_rejects_topological_order_that_omits_a_node() {
1716 let mut plan = hardening_plan();
1717 plan.validate().unwrap();
1718 plan.graph_plan.topological_order = vec![NodeId::new("transform:snv").unwrap()];
1721 let error = plan.validate().unwrap_err();
1722 assert!(matches!(error, DagMlError::Planning(_)));
1723 assert!(error
1724 .to_string()
1725 .contains("topological_order does not match"));
1726 }
1727
1728 #[test]
1729 fn validate_rejects_node_plan_kind_that_differs_from_graph_node() {
1730 let mut plan = hardening_plan();
1731 plan.validate().unwrap();
1732 plan.node_plans
1733 .get_mut(&NodeId::new("model:pls").unwrap())
1734 .unwrap()
1735 .kind = NodeKind::Transform;
1736 let error = plan.validate().unwrap_err();
1737 assert!(matches!(error, DagMlError::Planning(_)));
1738 assert!(error
1739 .to_string()
1740 .contains("kind does not match graph node kind"));
1741 }
1742
1743 #[test]
1744 fn validate_rejects_forged_input_adjacency() {
1745 let mut plan = hardening_plan();
1746 plan.validate().unwrap();
1747 plan.node_plans
1750 .get_mut(&NodeId::new("model:pls").unwrap())
1751 .unwrap()
1752 .input_nodes = vec![];
1753 let error = plan.validate().unwrap_err();
1754 assert!(matches!(error, DagMlError::Planning(_)));
1755 assert!(error
1756 .to_string()
1757 .contains("adjacency does not match the graph"));
1758 }
1759
1760 #[test]
1761 fn validate_rejects_node_plan_supported_phases_that_diverge_from_manifest() {
1762 let mut plan = hardening_plan();
1763 plan.validate().unwrap();
1764 plan.node_plans
1767 .get_mut(&NodeId::new("model:pls").unwrap())
1768 .unwrap()
1769 .supported_phases
1770 .insert(Phase::Explain);
1771 let error = plan.validate().unwrap_err();
1772 assert!(matches!(error, DagMlError::Planning(_)));
1773 assert!(error
1774 .to_string()
1775 .contains("does not match controller manifest"));
1776 }
1777
1778 #[test]
1779 fn validate_rejects_node_plan_controller_version_that_diverges_from_manifest() {
1780 let mut plan = hardening_plan();
1781 plan.validate().unwrap();
1782 plan.node_plans
1783 .get_mut(&NodeId::new("model:pls").unwrap())
1784 .unwrap()
1785 .controller_version = "9.9.9".to_string();
1786 let error = plan.validate().unwrap_err();
1787 assert!(matches!(error, DagMlError::Planning(_)));
1788 assert!(error
1789 .to_string()
1790 .contains("does not match controller manifest"));
1791 }
1792
1793 #[test]
1794 fn validate_rejects_stale_graph_fingerprint_after_graph_mutation() {
1795 let mut plan = hardening_plan();
1796 plan.validate().unwrap();
1797 plan.graph_plan
1802 .graph
1803 .metadata
1804 .insert("tampered".to_string(), serde_json::json!(true));
1805 let error = plan.validate().unwrap_err();
1806 assert!(matches!(error, DagMlError::Planning(_)));
1807 assert!(error
1808 .to_string()
1809 .contains("graph_fingerprint does not match"));
1810 }
1811
1812 #[test]
1813 fn validate_rejects_forged_graph_fingerprint_field() {
1814 let mut plan = hardening_plan();
1815 plan.validate().unwrap();
1816 plan.graph_fingerprint = "sha256:forged".to_string();
1819 let error = plan.validate().unwrap_err();
1820 assert!(matches!(error, DagMlError::Planning(_)));
1821 assert!(error
1822 .to_string()
1823 .contains("graph_fingerprint does not match"));
1824 }
1825
1826 #[test]
1827 fn validate_rejects_stale_campaign_fingerprint_after_campaign_mutation() {
1828 let mut plan = hardening_plan();
1829 plan.validate().unwrap();
1830 plan.campaign
1832 .metadata
1833 .insert("tampered".to_string(), serde_json::json!("x"));
1834 let error = plan.validate().unwrap_err();
1835 assert!(matches!(error, DagMlError::Planning(_)));
1836 assert!(error
1837 .to_string()
1838 .contains("campaign_fingerprint does not match"));
1839 }
1840
1841 #[test]
1842 fn validate_rejects_forged_campaign_fingerprint_field() {
1843 let mut plan = hardening_plan();
1844 plan.validate().unwrap();
1845 plan.campaign_fingerprint = "sha256:forged".to_string();
1846 let error = plan.validate().unwrap_err();
1847 assert!(matches!(error, DagMlError::Planning(_)));
1848 assert!(error
1849 .to_string()
1850 .contains("campaign_fingerprint does not match"));
1851 }
1852
1853 #[test]
1854 fn validate_rejects_stale_controller_fingerprint_after_manifest_mutation() {
1855 let mut plan = hardening_plan();
1856 plan.validate().unwrap();
1857 plan.controller_manifests
1863 .get_mut(&ControllerId::new("controller:model").unwrap())
1864 .unwrap()
1865 .priority = 7;
1866 let error = plan.validate().unwrap_err();
1867 assert!(matches!(error, DagMlError::Planning(_)));
1868 assert!(error
1869 .to_string()
1870 .contains("controller_fingerprint does not match"));
1871 }
1872
1873 #[test]
1874 fn validate_rejects_forged_controller_fingerprint_field() {
1875 let mut plan = hardening_plan();
1876 plan.validate().unwrap();
1877 plan.controller_fingerprint = "sha256:forged".to_string();
1878 let error = plan.validate().unwrap_err();
1879 assert!(matches!(error, DagMlError::Planning(_)));
1880 assert!(error
1881 .to_string()
1882 .contains("controller_fingerprint does not match"));
1883 }
1884
1885 #[test]
1886 fn build_execution_plan_lowers_dsl_inner_cv_metadata_into_node_plan() {
1887 let mut graph = graph();
1888 graph
1889 .nodes
1890 .iter_mut()
1891 .find(|node| node.id.as_str() == "model:pls")
1892 .unwrap()
1893 .metadata
1894 .insert(
1895 "dsl_inner_cv".to_string(),
1896 serde_json::json!({"kind": "kfold", "n_splits": 3, "shuffle": false, "seed": 9}),
1897 );
1898
1899 let campaign = CampaignSpec {
1900 inner_cv: None,
1901 id: "campaign:inner-cv".to_string(),
1902 root_seed: Some(7),
1903 leakage_policy: LeakageUnitPolicy::default(),
1904 aggregation_policy: AggregationPolicy::default(),
1905 split_invocation: None,
1906 generation: Default::default(),
1907 shape_plans: BTreeMap::new(),
1908 data_bindings: BTreeMap::new(),
1909 branch_view_plans: Vec::new(),
1910 metadata: BTreeMap::new(),
1911 };
1912
1913 let plan = build_execution_plan("plan:inner-cv", graph, campaign, ®istry()).unwrap();
1914 match &plan.node_plans[&NodeId::new("model:pls").unwrap()].inner_cv {
1915 Some(crate::fold::NestedCvSpec::KFold(k)) => {
1916 assert_eq!(k.n_splits, 3);
1917 assert_eq!(k.seed, Some(9));
1918 }
1919 other => panic!("expected lowered KFold inner_cv, got {other:?}"),
1920 }
1921 assert!(plan.node_plans[&NodeId::new("transform:snv").unwrap()]
1922 .inner_cv
1923 .is_none());
1924 }
1925
1926 #[test]
1927 fn build_execution_plan_rejects_malformed_dsl_inner_cv_metadata() {
1928 let mut graph = graph();
1929 graph
1930 .nodes
1931 .iter_mut()
1932 .find(|node| node.id.as_str() == "model:pls")
1933 .unwrap()
1934 .metadata
1935 .insert(
1936 "dsl_inner_cv".to_string(),
1937 serde_json::json!({"kind": "not_a_real_kind"}),
1938 );
1939
1940 let campaign = CampaignSpec {
1941 inner_cv: None,
1942 id: "campaign:inner-cv.bad".to_string(),
1943 root_seed: Some(7),
1944 leakage_policy: LeakageUnitPolicy::default(),
1945 aggregation_policy: AggregationPolicy::default(),
1946 split_invocation: None,
1947 generation: Default::default(),
1948 shape_plans: BTreeMap::new(),
1949 data_bindings: BTreeMap::new(),
1950 branch_view_plans: Vec::new(),
1951 metadata: BTreeMap::new(),
1952 };
1953
1954 let error =
1955 build_execution_plan("plan:inner-cv.bad", graph, campaign, ®istry()).unwrap_err();
1956 assert!(matches!(error, DagMlError::Planning(_)));
1957 assert!(error.to_string().contains("invalid dsl_inner_cv metadata"));
1958 }
1959
1960 #[test]
1961 fn build_execution_plan_rejects_semantically_invalid_dsl_inner_cv() {
1962 let mut graph = graph();
1965 graph
1966 .nodes
1967 .iter_mut()
1968 .find(|node| node.id.as_str() == "model:pls")
1969 .unwrap()
1970 .metadata
1971 .insert(
1972 "dsl_inner_cv".to_string(),
1973 serde_json::json!({"kind": "kfold", "n_splits": 1, "shuffle": false, "seed": null}),
1974 );
1975
1976 let campaign = CampaignSpec {
1977 inner_cv: None,
1978 id: "campaign:inner-cv.nsplits".to_string(),
1979 root_seed: Some(7),
1980 leakage_policy: LeakageUnitPolicy::default(),
1981 aggregation_policy: AggregationPolicy::default(),
1982 split_invocation: None,
1983 generation: Default::default(),
1984 shape_plans: BTreeMap::new(),
1985 data_bindings: BTreeMap::new(),
1986 branch_view_plans: Vec::new(),
1987 metadata: BTreeMap::new(),
1988 };
1989
1990 let error = build_execution_plan("plan:inner-cv.nsplits", graph, campaign, ®istry())
1991 .unwrap_err();
1992 assert!(matches!(error, DagMlError::Planning(_)));
1993 assert!(error.to_string().contains("at least two splits"));
1994 }
1995 use crate::data::{
1996 BranchViewMode, BranchViewPlan, DataBinding, DataViewSelector, SOURCE_INDEX_METADATA_KEY,
1997 };
1998 use crate::generation::{
1999 GenerationChoice, GenerationConstraints, GenerationDimension, GenerationParamOverride,
2000 GenerationStrategy,
2001 };
2002 use crate::graph::{
2003 EdgeContract, EdgeSpec, GraphInterface, NodeSpec, PortCardinality, PortKind, PortRef,
2004 PortSchema, PortSpec,
2005 };
2006 use crate::ids::{ControllerId, FoldId, ObservationId, SampleId, TargetId};
2007 use crate::phase::Phase;
2008 use crate::policy::{DataModelShapePlan, Granularity};
2009 use crate::relation::{SampleRelation, SampleRelationSet};
2010
2011 fn port(name: &str, kind: PortKind) -> PortSpec {
2012 PortSpec {
2013 name: name.to_string(),
2014 kind,
2015 representation: None,
2016 cardinality: PortCardinality::One,
2017 unit_level: None,
2018 alignment_key: None,
2019 target_level: None,
2020 description: String::new(),
2021 }
2022 }
2023
2024 fn node(id: &str, kind: NodeKind, inputs: Vec<PortSpec>, outputs: Vec<PortSpec>) -> NodeSpec {
2025 NodeSpec {
2026 id: NodeId::new(id).unwrap(),
2027 kind,
2028 operator: None,
2029 params: BTreeMap::new(),
2030 ports: PortSchema { inputs, outputs },
2031 metadata: BTreeMap::new(),
2032 seed_label: None,
2033 }
2034 }
2035
2036 fn graph() -> GraphSpec {
2037 GraphSpec {
2038 id: "g".to_string(),
2039 interface: GraphInterface::default(),
2040 nodes: vec![
2041 node(
2042 "transform:snv",
2043 NodeKind::Transform,
2044 vec![],
2045 vec![port("x", PortKind::Data)],
2046 ),
2047 node(
2048 "model:pls",
2049 NodeKind::Model,
2050 vec![port("x", PortKind::Data)],
2051 vec![port("pred", PortKind::Prediction)],
2052 ),
2053 ],
2054 edges: vec![EdgeSpec {
2055 source: PortRef {
2056 node_id: NodeId::new("transform:snv").unwrap(),
2057 port_name: "x".to_string(),
2058 },
2059 target: PortRef {
2060 node_id: NodeId::new("model:pls").unwrap(),
2061 port_name: "x".to_string(),
2062 },
2063 contract: EdgeContract {
2064 requires_oof: false,
2065 requires_fold_alignment: false,
2066 ..EdgeContract::new(PortKind::Data, None)
2067 },
2068 }],
2069 search_space_fingerprint: None,
2070 metadata: BTreeMap::new(),
2071 }
2072 }
2073
2074 fn manifest(id: &str, kind: NodeKind) -> ControllerManifest {
2075 let mut capabilities = BTreeSet::from([
2076 ControllerCapability::Deterministic,
2077 ControllerCapability::ThreadSafe,
2078 ControllerCapability::ProcessSafe,
2079 ]);
2080 if kind == NodeKind::Model {
2081 capabilities.insert(ControllerCapability::EmitsPredictions);
2082 capabilities.insert(ControllerCapability::ConsumesOofPredictions);
2083 }
2084 ControllerManifest {
2085 controller_id: ControllerId::new(id).unwrap(),
2086 controller_version: "0.1.0".to_string(),
2087 operator_kind: kind,
2088 priority: 0,
2089 supported_phases: BTreeSet::from([Phase::FitCv, Phase::Refit, Phase::Predict]),
2090 input_ports: Vec::new(),
2091 output_ports: Vec::new(),
2092 data_requirements: None,
2093 capabilities,
2094 operator_selectors: Vec::new(),
2095 fit_scope: ControllerFitScope::FoldTrain,
2096 rng_policy: RngPolicy::UsesCoreSeed,
2097 artifact_policy: ArtifactPolicy::Serializable,
2098 }
2099 }
2100
2101 fn registry() -> ControllerRegistry {
2102 let mut registry = ControllerRegistry::new();
2103 registry
2104 .register(manifest("controller:transform", NodeKind::Transform))
2105 .unwrap();
2106 registry
2107 .register(manifest("controller:model", NodeKind::Model))
2108 .unwrap();
2109 registry
2110 }
2111
2112 fn registry_with_model_data_requirements(
2113 accepted_representations: &[&str],
2114 accepted_types: &[&str],
2115 ) -> ControllerRegistry {
2116 let mut registry = ControllerRegistry::new();
2117 registry
2118 .register(manifest("controller:transform", NodeKind::Transform))
2119 .unwrap();
2120 let mut model = manifest("controller:model", NodeKind::Model);
2121 model.data_requirements = Some(serde_json::json!({
2122 "schema_version": 1,
2123 "ports": [
2124 {
2125 "name": "x",
2126 "accepted_representations": accepted_representations,
2127 "accepted_types": accepted_types,
2128 "rank": 2,
2129 "multi_source": true,
2130 "optional": false
2131 }
2132 ],
2133 "metadata": {
2134 "source": "plan-test"
2135 }
2136 }));
2137 registry.register(model).unwrap();
2138 registry
2139 }
2140
2141 fn registry_with_model_data_requirements_json(
2142 data_requirements: serde_json::Value,
2143 ) -> ControllerRegistry {
2144 let mut registry = ControllerRegistry::new();
2145 registry
2146 .register(manifest("controller:transform", NodeKind::Transform))
2147 .unwrap();
2148 let mut model = manifest("controller:model", NodeKind::Model);
2149 model.data_requirements = Some(data_requirements);
2150 registry.register(model).unwrap();
2151 registry
2152 }
2153
2154 fn model_data_requirements(
2155 multi_source: bool,
2156 default_fusion: Option<serde_json::Value>,
2157 ) -> serde_json::Value {
2158 let mut spec = serde_json::json!({
2159 "schema_version": 1,
2160 "ports": [
2161 {
2162 "name": "x",
2163 "accepted_representations": ["tabular_numeric"],
2164 "accepted_types": ["table"],
2165 "rank": 2,
2166 "multi_source": multi_source,
2167 "optional": false
2168 }
2169 ],
2170 "metadata": {
2171 "source": "plan-test"
2172 }
2173 });
2174 if let Some(default_fusion) = default_fusion {
2175 spec.as_object_mut()
2176 .unwrap()
2177 .insert("default_fusion".to_string(), default_fusion);
2178 }
2179 spec
2180 }
2181
2182 fn multisource_binding(node_id: &NodeId) -> DataBinding {
2183 let mut binding = data_binding(node_id);
2184 binding.request_id = "nir-chem-source-concat".to_string();
2185 binding.feature_set_id = Some("x_fused".to_string());
2186 binding.source_ids = vec!["nir".to_string(), "chem".to_string()];
2187 binding
2188 }
2189
2190 fn add_source_index(binding: &mut DataBinding) {
2191 binding.metadata.insert(
2192 SOURCE_INDEX_METADATA_KEY.to_string(),
2193 serde_json::json!({
2194 "nir": 0,
2195 "chem": 1
2196 }),
2197 );
2198 }
2199
2200 fn source_concat_fusion() -> serde_json::Value {
2201 serde_json::json!({
2202 "mode": "concatenate_features",
2203 "alignment": "sample_id",
2204 "adapter_id": null,
2205 "params": {
2206 "namespace_columns": true
2207 }
2208 })
2209 }
2210
2211 fn by_source_graph(source_ids: Vec<&str>) -> GraphSpec {
2212 let mut graph = graph();
2213 let branch_view = BranchViewPlan {
2214 view_id: "branch_view:source".to_string(),
2215 branch_id: "branch:source".to_string(),
2216 mode: BranchViewMode::BySource,
2217 selector: DataViewSelector {
2218 source_ids: source_ids.into_iter().map(str::to_string).collect(),
2219 ..Default::default()
2220 },
2221 allow_overlap: false,
2222 metadata: BTreeMap::new(),
2223 };
2224 graph
2225 .nodes
2226 .iter_mut()
2227 .find(|node| node.id.as_str() == "model:pls")
2228 .unwrap()
2229 .metadata
2230 .insert(
2231 "dsl_branch_view_plan".to_string(),
2232 serde_json::to_value(branch_view).unwrap(),
2233 );
2234 graph
2235 }
2236
2237 fn refusal_payload(error: DagMlError) -> serde_json::Value {
2238 let message = error.to_string();
2239 let payload = message
2240 .split_once("data requirement refusal: ")
2241 .unwrap_or_else(|| panic!("missing structured refusal payload in: {message}"))
2242 .1;
2243 serde_json::from_str(payload).unwrap()
2244 }
2245
2246 fn campaign(id: &str) -> CampaignSpec {
2247 CampaignSpec {
2248 id: id.to_string(),
2249 root_seed: Some(7),
2250 leakage_policy: LeakageUnitPolicy::default(),
2251 aggregation_policy: AggregationPolicy::default(),
2252 split_invocation: None,
2253 generation: Default::default(),
2254 shape_plans: BTreeMap::new(),
2255 data_bindings: BTreeMap::new(),
2256 branch_view_plans: Vec::new(),
2257 inner_cv: None,
2258 metadata: BTreeMap::new(),
2259 }
2260 }
2261
2262 #[cfg(dag_ml_workspace_contract_fixtures)]
2263 fn custom_loss_role(node_id: &str, output_id: &str) -> TrainingLossRoleReference {
2264 let fixture: serde_json::Value = serde_json::from_str(include_str!(
2265 "../../../examples/fixtures/criteria/javascript_local_implementations.v1.json"
2266 ))
2267 .unwrap();
2268 let mut role: TrainingLossRoleReference =
2269 serde_json::from_value(fixture["training_loss_role"].clone()).unwrap();
2270 role.node_id = NodeId::new(node_id).unwrap();
2271 role.output_id = Some(output_id.to_string());
2272 role
2273 }
2274
2275 #[cfg(dag_ml_workspace_contract_fixtures)]
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 #[cfg(dag_ml_workspace_contract_fixtures)]
2645 #[test]
2646 fn published_campaign_spec_schema_declares_current_contract() {
2647 let schema: serde_json::Value = serde_json::from_str(include_str!(
2648 "../../../docs/contracts/campaign_spec.schema.json"
2649 ))
2650 .unwrap();
2651
2652 assert_eq!(schema["$id"], CAMPAIGN_SPEC_SCHEMA_ID);
2653 assert!(schema["required"]
2654 .as_array()
2655 .unwrap()
2656 .iter()
2657 .any(|field| field.as_str() == Some("id")));
2658 assert!(schema["$defs"]["split_invocation"]["properties"]
2659 .as_object()
2660 .unwrap()
2661 .contains_key("fold_set"));
2662 assert!(schema["$defs"]["aggregation_policy"]["properties"]
2663 .as_object()
2664 .unwrap()
2665 .contains_key("selection_metric_level"));
2666 assert!(schema["$defs"]["aggregation_policy"]["properties"]
2667 .as_object()
2668 .unwrap()
2669 .contains_key("custom_controller"));
2670 assert!(schema["$defs"]["data_binding"]["properties"]
2671 .as_object()
2672 .unwrap()
2673 .contains_key("view_policy"));
2674 assert!(schema["properties"]
2675 .as_object()
2676 .unwrap()
2677 .contains_key("branch_view_plans"));
2678 assert!(schema["$defs"]["branch_view_plan"]["properties"]
2679 .as_object()
2680 .unwrap()
2681 .contains_key("selector"));
2682 }
2683
2684 #[cfg(dag_ml_workspace_contract_fixtures)]
2685 #[test]
2686 fn published_execution_plan_schema_declares_current_contract() {
2687 let schema: serde_json::Value = serde_json::from_str(include_str!(
2688 "../../../docs/contracts/execution_plan.schema.json"
2689 ))
2690 .unwrap();
2691
2692 assert_eq!(schema["$id"], EXECUTION_PLAN_SCHEMA_ID);
2693 assert!(schema["required"]
2694 .as_array()
2695 .unwrap()
2696 .iter()
2697 .any(|field| field.as_str() == Some("node_plans")));
2698 assert!(schema["properties"]
2699 .as_object()
2700 .unwrap()
2701 .contains_key("controller_fingerprint"));
2702 assert!(schema["$defs"]["node_plan"]["properties"]
2703 .as_object()
2704 .unwrap()
2705 .contains_key("shape_plan"));
2706 assert!(schema["$defs"]["variant_plan"]["properties"]
2707 .as_object()
2708 .unwrap()
2709 .contains_key("choices"));
2710 }
2711
2712 #[cfg(dag_ml_workspace_contract_fixtures)]
2713 #[test]
2714 fn published_execution_plan_fixture_validates_current_contract() {
2715 let plan: ExecutionPlan = serde_json::from_str(include_str!(
2716 "../../../examples/fixtures/runtime/execution_plan_branch_merge_executable.json"
2717 ))
2718 .unwrap();
2719
2720 plan.validate().unwrap();
2721 assert_eq!(plan.id, "plan:fixture.execution.branch_merge");
2722 assert_eq!(plan.variants.len(), 2);
2723 assert_eq!(plan.node_plans.len(), plan.graph_plan.graph.nodes.len());
2724 }
2725
2726 #[test]
2727 #[ignore = "perf sanity probe; run with --release --ignored --nocapture"]
2728 fn build_execution_plan_large_linear_graph_under_1500ms() {
2729 let started = Instant::now();
2730 let plan = build_execution_plan(
2731 "plan:perf.linear",
2732 large_linear_graph(400),
2733 campaign("campaign:perf.linear"),
2734 ®istry(),
2735 )
2736 .unwrap();
2737 let elapsed = started.elapsed();
2738
2739 assert_eq!(plan.graph_plan.topological_order.len(), 401);
2740 assert_eq!(plan.node_plans.len(), 401);
2741 assert!(
2742 elapsed <= Duration::from_millis(1_500),
2743 "large execution-plan build took {elapsed:?}"
2744 );
2745 }
2746
2747 #[test]
2748 fn builds_execution_plan_with_shape_and_fold_contracts() {
2749 let model_id = NodeId::new("model:pls").unwrap();
2750 let campaign = CampaignSpec {
2751 inner_cv: None,
2752 id: "campaign:oof".to_string(),
2753 root_seed: Some(7),
2754 leakage_policy: LeakageUnitPolicy::default(),
2755 aggregation_policy: AggregationPolicy::default(),
2756 split_invocation: Some(SplitInvocation {
2757 id: "split:outer".to_string(),
2758 controller_id: None,
2759 leakage_policy: LeakageUnitPolicy::default(),
2760 params: BTreeMap::new(),
2761 fold_set: Some(FoldSet {
2762 id: "outer".to_string(),
2763 sample_ids: vec![SampleId::new("s1").unwrap(), SampleId::new("s2").unwrap()],
2764 folds: vec![
2765 crate::fold::FoldAssignment {
2766 fold_id: FoldId::new("fold0").unwrap(),
2767 train_sample_ids: vec![SampleId::new("s2").unwrap()],
2768 validation_sample_ids: vec![SampleId::new("s1").unwrap()],
2769 metadata: BTreeMap::new(),
2770 },
2771 crate::fold::FoldAssignment {
2772 fold_id: FoldId::new("fold1").unwrap(),
2773 train_sample_ids: vec![SampleId::new("s1").unwrap()],
2774 validation_sample_ids: vec![SampleId::new("s2").unwrap()],
2775 metadata: BTreeMap::new(),
2776 },
2777 ],
2778 sample_groups: BTreeMap::new(),
2779 partition_mode: FoldPartitionMode::Partition,
2780 }),
2781 }),
2782 generation: Default::default(),
2783 shape_plans: BTreeMap::from([(
2784 model_id.clone(),
2785 DataModelShapePlan {
2786 node_id: model_id.clone(),
2787 input_granularity: Granularity::Observation,
2788 ..DataModelShapePlan {
2789 node_id: model_id.clone(),
2790 input_granularity: Granularity::Sample,
2791 target_granularity: Granularity::Sample,
2792 fit_rows: crate::policy::FitBoundary::FoldTrain,
2793 predict_rows: crate::policy::FitBoundary::FoldValidation,
2794 feature_namespace: None,
2795 feature_schema_fingerprint: None,
2796 target_space: "raw".to_string(),
2797 aggregation_policy: AggregationPolicy::default(),
2798 augmentation_policy: crate::policy::AugmentationPolicy::default(),
2799 selection_policy: crate::policy::FeatureSelectionPolicy::default(),
2800 }
2801 },
2802 )]),
2803 data_bindings: BTreeMap::from([(model_id.clone(), vec![data_binding(&model_id)])]),
2804 branch_view_plans: Vec::new(),
2805 metadata: BTreeMap::new(),
2806 };
2807
2808 let plan = build_execution_plan("plan:oof", graph(), campaign, ®istry()).unwrap();
2809
2810 assert_eq!(
2811 plan.graph_plan
2812 .topological_order
2813 .iter()
2814 .map(ToString::to_string)
2815 .collect::<Vec<_>>(),
2816 vec!["transform:snv", "model:pls"]
2817 );
2818 assert_eq!(
2819 levels_as_strings(&plan.graph_plan.parallel_levels),
2820 vec![vec!["transform:snv"], vec!["model:pls"]]
2821 );
2822 assert!(plan.node_plans[&model_id]
2823 .controller_capabilities
2824 .contains(&ControllerCapability::EmitsPredictions));
2825 assert!(plan.fold_set.is_some());
2826 let schedule = plan.campaign_phase_schedule(Phase::FitCv).unwrap();
2827 assert_eq!(schedule.scopes.len(), 2);
2828 assert!(schedule.scopes[0].scope_id.starts_with("scope:FIT_CV:"));
2829 assert!(schedule
2830 .scopes
2831 .iter()
2832 .all(|scope| levels_as_strings(&scope.node_levels)
2833 == vec![vec!["transform:snv"], vec!["model:pls"]]));
2834 assert_eq!(
2835 schedule
2836 .scopes
2837 .iter()
2838 .filter_map(|scope| scope.fold_id.as_ref().map(ToString::to_string))
2839 .collect::<Vec<_>>(),
2840 vec!["fold0", "fold1"]
2841 );
2842 assert_eq!(
2843 plan.node_plans
2844 .get(&model_id)
2845 .unwrap()
2846 .controller_id
2847 .as_str(),
2848 "controller:model"
2849 );
2850 assert_eq!(
2851 plan.node_plans.get(&model_id).unwrap().data_bindings.len(),
2852 1
2853 );
2854
2855 let mut bad_plan = plan.clone();
2856 bad_plan.graph_plan.parallel_levels =
2857 vec![vec![model_id], vec![NodeId::new("transform:snv").unwrap()]];
2858 assert!(bad_plan
2859 .validate()
2860 .unwrap_err()
2861 .to_string()
2862 .contains("parallel levels"));
2863
2864 let bad_envelope = ExternalDataPlanEnvelope {
2865 schema_version: crate::data::EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION,
2866 schema_fingerprint: "f97b37872fa22134b508f98fd8e207e5b776b52594fb8f6f5c3e15bee212246b"
2867 .to_string(),
2868 plan_fingerprint: "7c5431d85574b3f337022fa5d25971d5b5cf445b90331b49938f573ff6901e4d"
2869 .to_string(),
2870 relation_fingerprint: None,
2871 data_content_fingerprint: None,
2872 target_content_fingerprint: None,
2873 coordinator_relations: Some(SampleRelationSet {
2874 records: vec![{
2875 let mut relation = SampleRelation::new(
2876 ObservationId::new("obs:outside").unwrap(),
2877 SampleId::new("sample:outside").unwrap(),
2878 );
2879 relation.target_id = Some(TargetId::new("target:outside").unwrap());
2880 relation.source_id = Some("nir".to_string());
2881 relation
2882 }],
2883 }),
2884 };
2885 assert!(plan
2886 .campaign
2887 .validate_data_envelope_relations(&bad_envelope)
2888 .unwrap_err()
2889 .to_string()
2890 .contains("outside fold set"));
2891 }
2892
2893 #[test]
2894 fn planning_refuses_shape_plan_for_unknown_node() {
2895 let campaign = CampaignSpec {
2896 inner_cv: None,
2897 id: "campaign:oof".to_string(),
2898 root_seed: Some(7),
2899 leakage_policy: LeakageUnitPolicy::default(),
2900 aggregation_policy: AggregationPolicy::default(),
2901 split_invocation: None,
2902 generation: Default::default(),
2903 shape_plans: BTreeMap::from([(
2904 NodeId::new("model:missing").unwrap(),
2905 DataModelShapePlan {
2906 node_id: NodeId::new("model:missing").unwrap(),
2907 input_granularity: Granularity::Sample,
2908 target_granularity: Granularity::Sample,
2909 fit_rows: crate::policy::FitBoundary::FoldTrain,
2910 predict_rows: crate::policy::FitBoundary::FoldValidation,
2911 feature_namespace: None,
2912 feature_schema_fingerprint: None,
2913 target_space: "raw".to_string(),
2914 aggregation_policy: AggregationPolicy::default(),
2915 augmentation_policy: crate::policy::AugmentationPolicy::default(),
2916 selection_policy: crate::policy::FeatureSelectionPolicy::default(),
2917 },
2918 )]),
2919 data_bindings: BTreeMap::new(),
2920 branch_view_plans: Vec::new(),
2921 metadata: BTreeMap::new(),
2922 };
2923
2924 assert!(build_execution_plan("plan:oof", graph(), campaign, ®istry()).is_err());
2925 }
2926
2927 #[test]
2928 fn planning_refuses_oof_edge_without_controller_capabilities() {
2929 let mut registry = ControllerRegistry::new();
2930 let mut model_manifest = manifest("controller:model", NodeKind::Model);
2931 model_manifest
2932 .capabilities
2933 .remove(&ControllerCapability::ConsumesOofPredictions);
2934 registry.register(model_manifest).unwrap();
2935
2936 let err = build_execution_plan(
2937 "plan:oof.capability",
2938 oof_graph(),
2939 CampaignSpec {
2940 inner_cv: None,
2941 id: "campaign:oof.capability".to_string(),
2942 root_seed: Some(11),
2943 leakage_policy: Default::default(),
2944 aggregation_policy: Default::default(),
2945 split_invocation: None,
2946 generation: Default::default(),
2947 shape_plans: BTreeMap::new(),
2948 data_bindings: BTreeMap::new(),
2949 branch_view_plans: Vec::new(),
2950 metadata: BTreeMap::new(),
2951 },
2952 ®istry,
2953 )
2954 .unwrap_err();
2955
2956 assert!(err.to_string().contains("consumes_oof_predictions"));
2957 }
2958
2959 #[test]
2960 fn planning_refuses_raw_prediction_sibling_port_into_fitting_node() {
2961 let mut graph = oof_graph();
2962 graph.id = "g:oof.raw-sibling".to_string();
2963 let base = graph
2964 .nodes
2965 .iter_mut()
2966 .find(|node| node.id.as_str() == "model:base")
2967 .unwrap();
2968 base.ports.outputs.push(port("aux", PortKind::Prediction));
2969 let meta = graph
2970 .nodes
2971 .iter_mut()
2972 .find(|node| node.id.as_str() == "model:meta")
2973 .unwrap();
2974 meta.ports.inputs.push(port("aux", PortKind::Prediction));
2975 graph.edges.push(EdgeSpec {
2976 source: PortRef {
2977 node_id: NodeId::new("model:base").unwrap(),
2978 port_name: "aux".to_string(),
2979 },
2980 target: PortRef {
2981 node_id: NodeId::new("model:meta").unwrap(),
2982 port_name: "aux".to_string(),
2983 },
2984 contract: EdgeContract::new(PortKind::Prediction, None),
2985 });
2986
2987 let error = build_execution_plan(
2988 "plan:oof.raw-sibling",
2989 graph,
2990 campaign("campaign:oof.raw-sibling"),
2991 ®istry(),
2992 )
2993 .unwrap_err()
2994 .to_string();
2995
2996 assert!(error.contains("enters fitting controller"));
2997 assert!(error.contains("must require OOF"));
2998 }
2999
3000 #[test]
3001 fn parallel_controller_capability_validation_requires_safe_manifest() {
3002 let mut registry = ControllerRegistry::new();
3003 let mut transform_manifest = manifest("controller:transform", NodeKind::Transform);
3004 transform_manifest
3005 .capabilities
3006 .remove(&ControllerCapability::ThreadSafe);
3007 transform_manifest
3008 .capabilities
3009 .remove(&ControllerCapability::ProcessSafe);
3010 registry.register(transform_manifest).unwrap();
3011 registry
3012 .register(manifest("controller:model", NodeKind::Model))
3013 .unwrap();
3014 let plan = build_execution_plan(
3015 "plan:parallel.capability",
3016 graph(),
3017 CampaignSpec {
3018 inner_cv: None,
3019 id: "campaign:parallel.capability".to_string(),
3020 root_seed: Some(11),
3021 leakage_policy: Default::default(),
3022 aggregation_policy: Default::default(),
3023 split_invocation: None,
3024 generation: Default::default(),
3025 shape_plans: BTreeMap::new(),
3026 data_bindings: BTreeMap::new(),
3027 branch_view_plans: Vec::new(),
3028 metadata: BTreeMap::new(),
3029 },
3030 ®istry,
3031 )
3032 .unwrap();
3033
3034 assert!(plan
3035 .validate_parallel_controller_capabilities(1, Phase::FitCv)
3036 .is_ok());
3037 let err = plan
3038 .validate_parallel_controller_capabilities(2, Phase::FitCv)
3039 .unwrap_err();
3040 assert!(err.to_string().contains("thread_safe or process_safe"));
3041 }
3042
3043 #[test]
3044 fn planning_refuses_generation_override_for_unknown_node() {
3045 let campaign = CampaignSpec {
3046 inner_cv: None,
3047 id: "campaign:oof".to_string(),
3048 root_seed: Some(7),
3049 leakage_policy: LeakageUnitPolicy::default(),
3050 aggregation_policy: AggregationPolicy::default(),
3051 split_invocation: None,
3052 generation: GenerationSpec {
3053 strategy: GenerationStrategy::Cartesian,
3054 dimensions: vec![GenerationDimension {
3055 name: "model_family".to_string(),
3056 choices: vec![GenerationChoice {
3057 label: "pls".to_string(),
3058 value: serde_json::json!("pls"),
3059 param_overrides: vec![GenerationParamOverride {
3060 node_id: NodeId::new("model:missing").unwrap(),
3061 params: BTreeMap::from([(
3062 "n_components".to_string(),
3063 serde_json::json!(8),
3064 )]),
3065 }],
3066 active_subsequence: None,
3067 }],
3068 }],
3069 max_variants: Some(1),
3070 constraints: GenerationConstraints::default(),
3071 },
3072 shape_plans: BTreeMap::new(),
3073 data_bindings: BTreeMap::new(),
3074 branch_view_plans: Vec::new(),
3075 metadata: BTreeMap::new(),
3076 };
3077
3078 let error = build_execution_plan("plan:oof", graph(), campaign, ®istry())
3079 .unwrap_err()
3080 .to_string();
3081
3082 assert!(error.contains("overrides params for unknown node"));
3083 }
3084
3085 #[test]
3086 fn planning_validates_declared_search_space_fingerprint() {
3087 let campaign = CampaignSpec {
3088 inner_cv: None,
3089 id: "campaign:search.fingerprint".to_string(),
3090 root_seed: Some(7),
3091 leakage_policy: LeakageUnitPolicy::default(),
3092 aggregation_policy: AggregationPolicy::default(),
3093 split_invocation: None,
3094 generation: GenerationSpec {
3095 strategy: GenerationStrategy::Cartesian,
3096 dimensions: vec![GenerationDimension {
3097 name: "model_family".to_string(),
3098 choices: vec![GenerationChoice {
3099 label: "pls".to_string(),
3100 value: serde_json::json!("pls"),
3101 param_overrides: vec![GenerationParamOverride {
3102 node_id: NodeId::new("model:pls").unwrap(),
3103 params: BTreeMap::from([(
3104 "n_components".to_string(),
3105 serde_json::json!(8),
3106 )]),
3107 }],
3108 active_subsequence: None,
3109 }],
3110 }],
3111 max_variants: Some(1),
3112 constraints: GenerationConstraints::default(),
3113 },
3114 shape_plans: BTreeMap::new(),
3115 data_bindings: BTreeMap::new(),
3116 branch_view_plans: Vec::new(),
3117 metadata: BTreeMap::new(),
3118 };
3119 let mut graph = graph();
3120 graph.search_space_fingerprint =
3121 Some(generation_spec_fingerprint(&campaign.generation).unwrap());
3122
3123 let plan = build_execution_plan(
3124 "plan:search.fingerprint",
3125 graph.clone(),
3126 campaign.clone(),
3127 ®istry(),
3128 )
3129 .unwrap();
3130 assert_eq!(plan.variants.len(), 1);
3131
3132 graph.search_space_fingerprint = Some("sha256:not-the-generation-spec".to_string());
3133 let error = build_execution_plan("plan:search.fingerprint", graph, campaign, ®istry())
3134 .unwrap_err()
3135 .to_string();
3136 assert!(error.contains("search_space_fingerprint"));
3137 }
3138
3139 #[test]
3140 fn branch_view_lookup_helpers_match_by_branch_id_and_innermost_path() {
3141 use crate::data::{BranchViewMode, DataViewSelector};
3142
3143 let outer = BranchViewPlan {
3144 view_id: "branch_view:outer".to_string(),
3145 branch_id: "branch:outer".to_string(),
3146 mode: BranchViewMode::BySource,
3147 selector: DataViewSelector {
3148 source_ids: vec!["nir".to_string()],
3149 ..Default::default()
3150 },
3151 allow_overlap: false,
3152 metadata: BTreeMap::new(),
3153 };
3154 let inner = BranchViewPlan {
3155 view_id: "branch_view:inner".to_string(),
3156 branch_id: "branch:inner".to_string(),
3157 mode: BranchViewMode::Separation,
3158 selector: DataViewSelector {
3159 source_ids: vec!["chem".to_string()],
3160 ..Default::default()
3161 },
3162 allow_overlap: false,
3163 metadata: BTreeMap::new(),
3164 };
3165 let plans = vec![outer.clone(), inner.clone()];
3166
3167 assert_eq!(
3168 super::branch_view_for_in(&plans, "branch:outer"),
3169 Some(&outer)
3170 );
3171 assert_eq!(
3172 super::branch_view_for_in(&plans, "branch:inner"),
3173 Some(&inner)
3174 );
3175 assert_eq!(super::branch_view_for_in(&plans, "branch:missing"), None);
3176
3177 let path = vec!["branch:outer".to_string(), "branch:inner".to_string()];
3178 assert_eq!(super::branch_view_for_path_in(&plans, &path), Some(&inner));
3180
3181 let path_outer_only = vec!["branch:outer".to_string()];
3182 assert_eq!(
3183 super::branch_view_for_path_in(&plans, &path_outer_only),
3184 Some(&outer)
3185 );
3186
3187 let empty_path: Vec<String> = Vec::new();
3188 assert_eq!(super::branch_view_for_path_in(&plans, &empty_path), None);
3189
3190 let path_no_match = vec!["branch:other".to_string()];
3191 assert_eq!(super::branch_view_for_path_in(&plans, &path_no_match), None);
3192 }
3193}