dag_ml_core/runtime/
task.rs

1// Auto-split from the former monolithic `runtime.rs` (pure refactor).
2use super::*;
3use crate::TrainingLossRoleReference;
4
5#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
6pub struct PredictionInputSpec {
7    pub producer_node: NodeId,
8    pub source_port: String,
9    pub target_port: String,
10    pub partition: PredictionPartition,
11    #[serde(default = "default_runtime_prediction_level")]
12    pub prediction_level: PredictionLevel,
13    pub fold_id: Option<FoldId>,
14    #[serde(default)]
15    pub fold_ids: Vec<FoldId>,
16    #[serde(default, skip_serializing_if = "Vec::is_empty")]
17    pub unit_ids: Vec<PredictionUnitId>,
18    #[serde(default)]
19    pub sample_ids: Vec<SampleId>,
20    /// Per-sample OOF prediction rows, aligned 1:1 with `sample_ids`
21    /// (width == `prediction_width`). Sourced only from Validation OOF blocks
22    /// so a host can build a stacking meta-feature matrix during FIT_CV/REFIT.
23    #[serde(default)]
24    pub values: Vec<Vec<f64>>,
25    pub prediction_width: usize,
26    #[serde(default)]
27    pub target_names: Vec<String>,
28}
29
30#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
31pub struct ArtifactInputSpec {
32    pub node_id: NodeId,
33    pub controller_id: ControllerId,
34    pub artifact: ArtifactRef,
35    pub params_fingerprint: String,
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub training_loss_fingerprint: Option<String>,
38    #[serde(default)]
39    pub data_requirement_keys: Vec<String>,
40    #[serde(default)]
41    pub prediction_requirement_keys: Vec<String>,
42}
43
44impl ArtifactInputSpec {
45    pub(crate) fn from_refit_record(record: &RefitArtifactRecord) -> Result<Self> {
46        record.validate()?;
47        Ok(Self {
48            node_id: record.node_id.clone(),
49            controller_id: record.controller_id.clone(),
50            artifact: record.artifact.clone(),
51            params_fingerprint: record.params_fingerprint.clone(),
52            training_loss_fingerprint: record.training_loss_fingerprint.clone(),
53            data_requirement_keys: record.data_requirement_keys.clone(),
54            prediction_requirement_keys: record.prediction_requirement_keys.clone(),
55        })
56    }
57}
58
59pub(crate) fn default_runtime_prediction_level() -> PredictionLevel {
60    PredictionLevel::Sample
61}
62
63#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
64pub struct NodeTask {
65    pub run_id: RunId,
66    pub node_plan: NodePlan,
67    pub phase: Phase,
68    pub variant_id: Option<VariantId>,
69    #[serde(default)]
70    pub variant: Option<VariantExecutionSpec>,
71    pub fold_id: Option<FoldId>,
72    #[serde(default)]
73    pub branch_path: Vec<BranchId>,
74    #[serde(default)]
75    pub input_handles: BTreeMap<String, HandleRef>,
76    #[serde(default)]
77    pub data_views: BTreeMap<String, DataProviderViewSpec>,
78    #[serde(default)]
79    pub prediction_inputs: BTreeMap<String, PredictionInputSpec>,
80    #[serde(default)]
81    pub artifact_inputs: BTreeMap<String, ArtifactInputSpec>,
82    /// Native-produced attestation templates for the training losses that must
83    /// execute in this task. The order matches `node_plan.training_losses`
84    /// after filtering for `phase`. A controller may copy an entry into its
85    /// lineage only after the corresponding local or built-in loss succeeds.
86    #[serde(default, skip_serializing_if = "Vec::is_empty")]
87    pub required_loss_attestations: Vec<LossExecutionAttestation>,
88    /// Nested (inner) CV fold set for this node in the current outer fold, built
89    /// by the runtime from the outer fold's training samples when an effective
90    /// `inner_cv` policy applies (FIT_CV only). `None` otherwise. Leakage-safe by
91    /// construction (inner ⊆ outer-train); see [`crate::fold::NestedCvSpec`].
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub inner_fold_set: Option<FoldSet>,
94    #[serde(default, skip_serializing_if = "FitInfluenceTask::is_default")]
95    pub fit_influence: FitInfluenceTask,
96    pub seed: Option<u64>,
97}
98
99impl NodeTask {
100    pub fn required_loss_attestations_for(
101        node_plan: &NodePlan,
102        phase: Phase,
103    ) -> Result<Vec<LossExecutionAttestation>> {
104        node_plan
105            .training_losses_for_phase(phase)
106            .map(|role| LossExecutionAttestation::for_role(role, phase))
107            .collect()
108    }
109
110    pub fn validate_required_loss_attestations(&self) -> Result<()> {
111        let expected = Self::required_loss_attestations_for(&self.node_plan, self.phase)?;
112        if self.required_loss_attestations != expected {
113            return Err(DagMlError::RuntimeValidation(format!(
114                "task for node `{}` has loss execution requirements that do not match its ordered training losses for phase {:?}",
115                self.node_plan.node_id, self.phase
116            )));
117        }
118        Ok(())
119    }
120
121    /// Return one active training-loss role and its exact native attestation.
122    ///
123    /// The index addresses losses after filtering the node plan for the task's
124    /// phase. This keeps host bindings from reconstructing role ordering or
125    /// trusting controller-supplied attestation fields.
126    pub fn training_loss_binding(
127        &self,
128        role_index: usize,
129    ) -> Result<(&TrainingLossRoleReference, &LossExecutionAttestation)> {
130        if !matches!(self.phase, Phase::FitCv | Phase::Refit) {
131            return Err(DagMlError::RuntimeValidation(
132                "training loss phase must be FIT_CV or REFIT".to_string(),
133            ));
134        }
135        self.validate_required_loss_attestations()?;
136        let role = self
137            .node_plan
138            .training_losses_for_phase(self.phase)
139            .nth(role_index)
140            .ok_or_else(|| {
141                DagMlError::RuntimeValidation(format!(
142                    "role_index {role_index} is outside the active training loss range"
143                ))
144            })?;
145        let attestation = self
146            .required_loss_attestations
147            .get(role_index)
148            .ok_or_else(|| {
149                DagMlError::RuntimeValidation(
150                    "validated training loss role has no matching attestation".to_string(),
151                )
152            })?;
153        Ok((role, attestation))
154    }
155}
156
157#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
158#[serde(rename_all = "snake_case")]
159pub enum FitInfluenceMechanism {
160    UniformRows,
161    SampleWeights,
162    RowResampling,
163    BackendLossWeights,
164    ScorerOnly,
165}
166
167#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
168pub struct FitInfluenceTask {
169    pub requested_policy: FitInfluencePolicy,
170    pub effective_policy: FitInfluencePolicy,
171    pub mechanism: FitInfluenceMechanism,
172    #[serde(default, skip_serializing_if = "Vec::is_empty")]
173    pub row_weights: Vec<f64>,
174    #[serde(default, skip_serializing_if = "Vec::is_empty")]
175    pub warnings: Vec<String>,
176}
177
178impl Default for FitInfluenceTask {
179    fn default() -> Self {
180        Self {
181            requested_policy: FitInfluencePolicy::UniformRows,
182            effective_policy: FitInfluencePolicy::UniformRows,
183            mechanism: FitInfluenceMechanism::UniformRows,
184            row_weights: Vec::new(),
185            warnings: Vec::new(),
186        }
187    }
188}
189
190impl FitInfluenceTask {
191    fn is_default(&self) -> bool {
192        self == &Self::default()
193    }
194
195    pub fn diagnostic(&self) -> FitInfluenceDiagnostic {
196        FitInfluenceDiagnostic {
197            requested_policy: self.requested_policy,
198            effective_policy: self.effective_policy,
199            mechanism: self.mechanism,
200            fallback_used: !self.warnings.is_empty(),
201            row_weight_count: self.row_weights.len(),
202            warnings: self.warnings.clone(),
203        }
204    }
205
206    pub fn validate(&self) -> Result<()> {
207        if !self
208            .row_weights
209            .iter()
210            .all(|weight| weight.is_finite() && *weight > 0.0)
211        {
212            return Err(DagMlError::RuntimeValidation(
213                "fit influence row_weights must be finite and > 0".to_string(),
214            ));
215        }
216        if self
217            .warnings
218            .iter()
219            .any(|warning| warning.trim().is_empty())
220        {
221            return Err(DagMlError::RuntimeValidation(
222                "fit influence warnings must not be empty".to_string(),
223            ));
224        }
225        match self.effective_policy {
226            FitInfluencePolicy::EqualSampleInfluence | FitInfluencePolicy::BackendLossWeight
227                if self.row_weights.is_empty() =>
228            {
229                return Err(DagMlError::RuntimeValidation(format!(
230                    "fit influence {:?} requires row_weights",
231                    self.effective_policy
232                )));
233            }
234            _ => {}
235        }
236        if self.requested_policy == FitInfluencePolicy::StrictWeightSupport
237            && self.effective_policy == FitInfluencePolicy::UniformRows
238        {
239            return Err(DagMlError::RuntimeValidation(
240                "strict fit influence cannot fall back to uniform_rows".to_string(),
241            ));
242        }
243        Ok(())
244    }
245}
246
247#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
248#[serde(deny_unknown_fields)]
249pub struct FitInfluenceDiagnostic {
250    pub requested_policy: FitInfluencePolicy,
251    pub effective_policy: FitInfluencePolicy,
252    pub mechanism: FitInfluenceMechanism,
253    #[serde(default)]
254    pub fallback_used: bool,
255    #[serde(default)]
256    pub row_weight_count: usize,
257    #[serde(default, skip_serializing_if = "Vec::is_empty")]
258    pub warnings: Vec<String>,
259}
260
261impl FitInfluenceDiagnostic {
262    pub fn validate(&self, task: &NodeTask) -> Result<()> {
263        if self.requested_policy != task.fit_influence.requested_policy {
264            return Err(DagMlError::RuntimeValidation(format!(
265                "fit influence diagnostic requested_policy {:?} does not match task {:?}",
266                self.requested_policy, task.fit_influence.requested_policy
267            )));
268        }
269        if self.effective_policy != task.fit_influence.effective_policy {
270            return Err(DagMlError::RuntimeValidation(format!(
271                "fit influence diagnostic effective_policy {:?} does not match task {:?}",
272                self.effective_policy, task.fit_influence.effective_policy
273            )));
274        }
275        if self.mechanism != task.fit_influence.mechanism {
276            return Err(DagMlError::RuntimeValidation(format!(
277                "fit influence diagnostic mechanism {:?} does not match task {:?}",
278                self.mechanism, task.fit_influence.mechanism
279            )));
280        }
281        if self.row_weight_count != task.fit_influence.row_weights.len() {
282            return Err(DagMlError::RuntimeValidation(format!(
283                "fit influence diagnostic row_weight_count {} does not match task {}",
284                self.row_weight_count,
285                task.fit_influence.row_weights.len()
286            )));
287        }
288        if self.fallback_used == task.fit_influence.warnings.is_empty() {
289            return Err(DagMlError::RuntimeValidation(
290                "fit influence diagnostic fallback_used does not match task warnings".to_string(),
291            ));
292        }
293        if self.warnings != task.fit_influence.warnings {
294            return Err(DagMlError::RuntimeValidation(
295                "fit influence diagnostic warnings do not match task warnings".to_string(),
296            ));
297        }
298        if self
299            .warnings
300            .iter()
301            .any(|warning| warning.trim().is_empty())
302        {
303            return Err(DagMlError::RuntimeValidation(
304                "fit influence diagnostic warnings must not be empty".to_string(),
305            ));
306        }
307        Ok(())
308    }
309}
310
311#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
312pub struct VariantExecutionSpec {
313    pub variant_id: VariantId,
314    #[serde(default)]
315    pub choices: BTreeMap<String, GenerationChoice>,
316    pub fingerprint: String,
317    pub seed: Option<u64>,
318}
319
320impl VariantExecutionSpec {
321    pub fn from_plan(variant: &VariantPlan) -> Self {
322        Self {
323            variant_id: variant.variant_id.clone(),
324            choices: variant.choices.clone(),
325            fingerprint: variant.fingerprint.clone(),
326            seed: variant.seed,
327        }
328    }
329
330    pub fn validate(&self) -> Result<()> {
331        if self.fingerprint.trim().is_empty() {
332            return Err(DagMlError::RuntimeValidation(format!(
333                "variant `{}` has an empty fingerprint in task context",
334                self.variant_id
335            )));
336        }
337        for (dimension_name, choice) in &self.choices {
338            if dimension_name.trim().is_empty() {
339                return Err(DagMlError::RuntimeValidation(format!(
340                    "variant `{}` has an empty generation dimension name",
341                    self.variant_id
342                )));
343            }
344            if choice.label.trim().is_empty() {
345                return Err(DagMlError::RuntimeValidation(format!(
346                    "variant `{}` has an empty choice label for dimension `{dimension_name}`",
347                    self.variant_id
348                )));
349            }
350            for override_spec in &choice.param_overrides {
351                if override_spec.params.is_empty() {
352                    return Err(DagMlError::RuntimeValidation(format!(
353                        "variant `{}` has an empty param override for node `{}`",
354                        self.variant_id, override_spec.node_id
355                    )));
356                }
357                for param_key in override_spec.params.keys() {
358                    if param_key.trim().is_empty() {
359                        return Err(DagMlError::RuntimeValidation(format!(
360                            "variant `{}` has an empty param override key for node `{}`",
361                            self.variant_id, override_spec.node_id
362                        )));
363                    }
364                }
365            }
366        }
367        self.param_overrides_by_node()?;
368        Ok(())
369    }
370
371    pub fn effective_params_for_node(
372        &self,
373        node_id: &NodeId,
374        base_params: &BTreeMap<String, serde_json::Value>,
375    ) -> Result<BTreeMap<String, serde_json::Value>> {
376        let overrides_by_node = self.param_overrides_by_node()?;
377        let Some(overrides) = overrides_by_node.get(node_id) else {
378            return Ok(base_params.clone());
379        };
380        let mut params = base_params.clone();
381        params.extend(overrides.clone());
382        Ok(params)
383    }
384
385    fn param_overrides_by_node(
386        &self,
387    ) -> Result<BTreeMap<NodeId, BTreeMap<String, serde_json::Value>>> {
388        let mut overrides = BTreeMap::<NodeId, BTreeMap<String, serde_json::Value>>::new();
389        let mut owners = BTreeMap::<(NodeId, String), String>::new();
390        for (dimension_name, choice) in &self.choices {
391            for override_spec in &choice.param_overrides {
392                for (param_key, value) in &override_spec.params {
393                    let owner_key = (override_spec.node_id.clone(), param_key.clone());
394                    if let Some(previous) =
395                        owners.insert(owner_key, format!("{dimension_name}:{}", choice.label))
396                    {
397                        return Err(DagMlError::RuntimeValidation(format!(
398                            "variant `{}` has conflicting generation overrides for `{}.{}` from `{previous}` and `{}:{}`",
399                            self.variant_id,
400                            override_spec.node_id,
401                            param_key,
402                            dimension_name,
403                            choice.label
404                        )));
405                    }
406                    overrides
407                        .entry(override_spec.node_id.clone())
408                        .or_default()
409                        .insert(param_key.clone(), value.clone());
410                }
411            }
412        }
413        Ok(overrides)
414    }
415}
416
417/// An EXPLAIN-phase output block (ADR-12 explain contract). Explanations are a
418/// node *output* returned in the [`NodeResult`] — like predictions, they cross as
419/// data, not as an opaque host handle. The `payload` shape is controller-defined
420/// (e.g. per-feature importances); the core does not interpret it. Explanations
421/// are only valid in the `EXPLAIN` phase.
422#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
423#[serde(deny_unknown_fields)]
424pub struct ExplanationBlock {
425    /// Node whose model the explanation describes (must equal the producing node).
426    pub producer_node: NodeId,
427    #[serde(default, skip_serializing_if = "Option::is_none")]
428    pub producer_port: Option<String>,
429    /// Stable explanation method identifier, e.g. `shap`, `permutation_importance`.
430    pub method: String,
431    /// Optional target/output name the explanation pertains to.
432    #[serde(default, skip_serializing_if = "Option::is_none")]
433    pub target_name: Option<String>,
434    /// Controller-defined explanation payload as canonical JSON.
435    pub payload: serde_json::Value,
436}
437
438impl ExplanationBlock {
439    /// Validate the intrinsic shape of the explanation block (method/target_name
440    /// non-empty). Producer identity is checked against the node in
441    /// [`NodeResult::validate_for_task`].
442    pub fn validate(&self) -> Result<()> {
443        if self.method.trim().is_empty() {
444            return Err(DagMlError::RuntimeValidation(
445                "explanation method must be a non-empty identifier".to_string(),
446            ));
447        }
448        if let Some(name) = &self.target_name {
449            if name.trim().is_empty() {
450                return Err(DagMlError::RuntimeValidation(
451                    "explanation target_name must be non-empty when present".to_string(),
452                ));
453            }
454        }
455        Ok(())
456    }
457}
458
459#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
460#[serde(deny_unknown_fields)]
461pub struct NodeResult {
462    #[serde(default, skip_serializing_if = "Option::is_none")]
463    pub schema_version: Option<u32>,
464    pub node_id: NodeId,
465    #[serde(default)]
466    pub outputs: BTreeMap<String, HandleRef>,
467    #[serde(default)]
468    pub predictions: Vec<PredictionBlock>,
469    #[serde(default)]
470    pub observation_predictions: Vec<ObservationPredictionBlock>,
471    #[serde(default)]
472    pub aggregated_predictions: Vec<AggregatedPredictionBlock>,
473    #[serde(default)]
474    pub explanations: Vec<ExplanationBlock>,
475    #[serde(default)]
476    pub shape_deltas: Vec<ShapeDelta>,
477    #[serde(default)]
478    pub artifacts: Vec<ArtifactRef>,
479    #[serde(default)]
480    pub artifact_handles: BTreeMap<ArtifactId, HandleRef>,
481    #[serde(default, skip_serializing_if = "Vec::is_empty")]
482    pub fit_influence_diagnostics: Vec<FitInfluenceDiagnostic>,
483    /// Optional ground-truth targets the host controller emits alongside predictions so the core
484    /// can score natively (the runtime never sees feature matrices; `y_true` is data-tier and may
485    /// cross the ABI per the ownership table). Each block is identity-keyed by `unit_ids`.
486    #[serde(default, skip_serializing_if = "Vec::is_empty")]
487    pub regression_targets: Vec<RegressionTargetBlock>,
488    pub lineage: LineageRecord,
489}
490
491impl NodeResult {
492    pub fn validate_for_task(&self, task: &NodeTask) -> Result<()> {
493        if self.node_id != task.node_plan.node_id {
494            return Err(DagMlError::RuntimeValidation(format!(
495                "task for `{}` returned result for `{}`",
496                task.node_plan.node_id, self.node_id
497            )));
498        }
499        if self.lineage.node_id != task.node_plan.node_id {
500            return Err(DagMlError::RuntimeValidation(format!(
501                "lineage for task `{}` references node `{}`",
502                task.node_plan.node_id, self.lineage.node_id
503            )));
504        }
505        if self.lineage.phase != task.phase {
506            return Err(DagMlError::RuntimeValidation(format!(
507                "lineage for node `{}` has phase {:?}, expected {:?}",
508                task.node_plan.node_id, self.lineage.phase, task.phase
509            )));
510        }
511        if self.lineage.run_id != task.run_id {
512            return Err(DagMlError::RuntimeValidation(format!(
513                "lineage for node `{}` has run `{}`, expected `{}`",
514                task.node_plan.node_id, self.lineage.run_id, task.run_id
515            )));
516        }
517        if self.lineage.controller_id != task.node_plan.controller_id {
518            return Err(DagMlError::RuntimeValidation(format!(
519                "lineage for node `{}` has controller `{}`, expected `{}`",
520                task.node_plan.node_id, self.lineage.controller_id, task.node_plan.controller_id
521            )));
522        }
523        if self.lineage.controller_version != task.node_plan.controller_version {
524            return Err(DagMlError::RuntimeValidation(format!(
525                "lineage for node `{}` has controller version `{}`, expected `{}`",
526                task.node_plan.node_id,
527                self.lineage.controller_version,
528                task.node_plan.controller_version
529            )));
530        }
531        if self.lineage.variant_id != task.variant_id {
532            return Err(DagMlError::RuntimeValidation(format!(
533                "lineage for node `{}` has variant {:?}, expected {:?}",
534                task.node_plan.node_id, self.lineage.variant_id, task.variant_id
535            )));
536        }
537        if let Some(variant) = &task.variant {
538            variant.validate()?;
539            if Some(&variant.variant_id) != task.variant_id.as_ref() {
540                return Err(DagMlError::RuntimeValidation(format!(
541                    "task for node `{}` has variant context `{}` but variant_id {:?}",
542                    task.node_plan.node_id, variant.variant_id, task.variant_id
543                )));
544            }
545        }
546        if self.lineage.fold_id != task.fold_id {
547            return Err(DagMlError::RuntimeValidation(format!(
548                "lineage for node `{}` has fold {:?}, expected {:?}",
549                task.node_plan.node_id, self.lineage.fold_id, task.fold_id
550            )));
551        }
552        if self.lineage.branch_path != task.branch_path {
553            return Err(DagMlError::RuntimeValidation(format!(
554                "lineage for node `{}` has branch path {:?}, expected {:?}",
555                task.node_plan.node_id, self.lineage.branch_path, task.branch_path
556            )));
557        }
558        if self.lineage.seed != task.seed {
559            return Err(DagMlError::RuntimeValidation(format!(
560                "lineage for node `{}` has seed {:?}, expected {:?}",
561                task.node_plan.node_id, self.lineage.seed, task.seed
562            )));
563        }
564        if self.lineage.params_fingerprint != task.node_plan.params_fingerprint {
565            return Err(DagMlError::RuntimeValidation(format!(
566                "lineage for node `{}` has params fingerprint `{}`, expected `{}`",
567                task.node_plan.node_id,
568                self.lineage.params_fingerprint,
569                task.node_plan.params_fingerprint
570            )));
571        }
572        task.validate_required_loss_attestations()?;
573        let expected_losses = task
574            .node_plan
575            .training_losses_for_phase(task.phase)
576            .collect::<Vec<_>>();
577        if self.lineage.loss_attestations.len() != expected_losses.len() {
578            return Err(DagMlError::RuntimeValidation(format!(
579                "node `{}` returned {} loss attestations for {} resolved losses in phase {:?}",
580                task.node_plan.node_id,
581                self.lineage.loss_attestations.len(),
582                expected_losses.len(),
583                task.phase
584            )));
585        }
586        for (attestation, role) in self.lineage.loss_attestations.iter().zip(expected_losses) {
587            attestation.validate_against(role, &task.node_plan.node_id, task.phase)?;
588        }
589        if self.lineage.loss_attestations != task.required_loss_attestations {
590            return Err(DagMlError::RuntimeValidation(format!(
591                "node `{}` returned loss attestations that do not match the task requirements",
592                task.node_plan.node_id
593            )));
594        }
595        task.fit_influence.validate()?;
596        for diagnostic in &self.fit_influence_diagnostics {
597            diagnostic.validate(task)?;
598        }
599        validate_lineage_shape_fingerprints(&self.lineage, task)?;
600        if !self.explanations.is_empty() && task.phase != Phase::Explain {
601            return Err(DagMlError::RuntimeValidation(format!(
602                "node `{}` returned explanations outside the EXPLAIN phase",
603                task.node_plan.node_id
604            )));
605        }
606        for explanation in &self.explanations {
607            explanation.validate()?;
608            if explanation.producer_node != self.node_id {
609                return Err(DagMlError::RuntimeValidation(format!(
610                    "node `{}` returned an explanation produced by `{}`",
611                    self.node_id, explanation.producer_node
612                )));
613            }
614        }
615        for (port, handle) in &self.outputs {
616            if handle.owner_controller != task.node_plan.controller_id {
617                return Err(DagMlError::RuntimeValidation(format!(
618                    "node `{}` output `{port}` is owned by `{}`, expected `{}`",
619                    task.node_plan.node_id, handle.owner_controller, task.node_plan.controller_id
620                )));
621            }
622        }
623        let mut artifact_ids = BTreeSet::new();
624        for artifact in &self.artifacts {
625            artifact.validate()?;
626            if !artifact_ids.insert(artifact.id.clone()) {
627                return Err(DagMlError::RuntimeValidation(format!(
628                    "node `{}` emitted duplicate artifact `{}`",
629                    task.node_plan.node_id, artifact.id
630                )));
631            }
632            if artifact.controller_id != task.node_plan.controller_id {
633                return Err(DagMlError::RuntimeValidation(format!(
634                    "node `{}` emitted artifact `{}` for controller `{}`, expected `{}`",
635                    task.node_plan.node_id,
636                    artifact.id,
637                    artifact.controller_id,
638                    task.node_plan.controller_id
639                )));
640            }
641            let handle = self.artifact_handles.get(&artifact.id).ok_or_else(|| {
642                DagMlError::RuntimeValidation(format!(
643                    "node `{}` emitted artifact `{}` without artifact handle",
644                    task.node_plan.node_id, artifact.id
645                ))
646            })?;
647            if !matches!(handle.kind, HandleKind::Model | HandleKind::Artifact) {
648                return Err(DagMlError::RuntimeValidation(format!(
649                    "node `{}` emitted artifact `{}` with non-artifact/model handle kind {:?}",
650                    task.node_plan.node_id, artifact.id, handle.kind
651                )));
652            }
653            if handle.owner_controller != task.node_plan.controller_id {
654                return Err(DagMlError::RuntimeValidation(format!(
655                    "node `{}` emitted artifact `{}` owned by `{}`, expected `{}`",
656                    task.node_plan.node_id,
657                    artifact.id,
658                    handle.owner_controller,
659                    task.node_plan.controller_id
660                )));
661            }
662        }
663        for artifact_id in self.artifact_handles.keys() {
664            if !self
665                .artifacts
666                .iter()
667                .any(|artifact| &artifact.id == artifact_id)
668            {
669                return Err(DagMlError::RuntimeValidation(format!(
670                    "node `{}` emitted artifact handle for undeclared artifact `{artifact_id}`",
671                    task.node_plan.node_id
672                )));
673            }
674        }
675        for artifact in &self.artifacts {
676            if !self
677                .lineage
678                .artifact_refs
679                .iter()
680                .any(|lineage_artifact| lineage_artifact == artifact)
681            {
682                return Err(DagMlError::RuntimeValidation(format!(
683                    "node `{}` emitted artifact `{}` without matching lineage artifact ref",
684                    task.node_plan.node_id, artifact.id
685                )));
686            }
687        }
688        for artifact in &self.lineage.artifact_refs {
689            if !self
690                .artifacts
691                .iter()
692                .any(|emitted_artifact| emitted_artifact == artifact)
693            {
694                return Err(DagMlError::RuntimeValidation(format!(
695                    "node `{}` lineage references undeclared artifact `{}`",
696                    task.node_plan.node_id, artifact.id
697                )));
698            }
699        }
700        for prediction in &self.predictions {
701            prediction.validate_shape()?;
702            if prediction.producer_node != task.node_plan.node_id {
703                return Err(DagMlError::RuntimeValidation(format!(
704                    "node `{}` emitted prediction for producer `{}`",
705                    task.node_plan.node_id, prediction.producer_node
706                )));
707            }
708            validate_prediction_scope(prediction, task)?;
709        }
710        for prediction in &self.observation_predictions {
711            prediction.validate_shape()?;
712            if prediction.producer_node != task.node_plan.node_id {
713                return Err(DagMlError::RuntimeValidation(format!(
714                    "node `{}` emitted observation prediction for producer `{}`",
715                    task.node_plan.node_id, prediction.producer_node
716                )));
717            }
718            validate_observation_prediction_scope(prediction, task)?;
719        }
720        for prediction in &self.aggregated_predictions {
721            prediction.validate_shape()?;
722            if prediction.producer_node != task.node_plan.node_id {
723                return Err(DagMlError::RuntimeValidation(format!(
724                    "node `{}` emitted aggregated prediction for producer `{}`",
725                    task.node_plan.node_id, prediction.producer_node
726                )));
727            }
728            validate_aggregated_prediction_scope(prediction, task)?;
729        }
730        for delta in &self.shape_deltas {
731            delta.validate()?;
732            if delta.node_id != task.node_plan.node_id {
733                return Err(DagMlError::RuntimeValidation(format!(
734                    "node `{}` emitted shape delta for `{}`",
735                    task.node_plan.node_id, delta.node_id
736                )));
737            }
738            validate_shape_delta_for_task(delta, task)?;
739        }
740        for target in &self.regression_targets {
741            target.validate_shape()?;
742        }
743        self.lineage.validate()
744    }
745}
746
747pub(crate) fn validate_lineage_shape_fingerprints(
748    lineage: &LineageRecord,
749    task: &NodeTask,
750) -> Result<()> {
751    let Some(shape_plan) = &task.node_plan.shape_plan else {
752        if lineage.data_model_shape_fingerprint.is_some()
753            || lineage.aggregation_policy_fingerprint.is_some()
754        {
755            return Err(DagMlError::RuntimeValidation(format!(
756                "lineage for node `{}` carries shape fingerprints but the node has no shape plan",
757                task.node_plan.node_id
758            )));
759        }
760        return Ok(());
761    };
762
763    if let Some(actual) = &lineage.data_model_shape_fingerprint {
764        let expected = stable_json_fingerprint(shape_plan)?;
765        if actual != &expected {
766            return Err(DagMlError::RuntimeValidation(format!(
767                "lineage for node `{}` has data/model shape fingerprint `{actual}`, expected `{expected}`",
768                task.node_plan.node_id
769            )));
770        }
771    }
772    if let Some(actual) = &lineage.aggregation_policy_fingerprint {
773        let expected = stable_json_fingerprint(&shape_plan.aggregation_policy)?;
774        if actual != &expected {
775            return Err(DagMlError::RuntimeValidation(format!(
776                "lineage for node `{}` has aggregation policy fingerprint `{actual}`, expected `{expected}`",
777                task.node_plan.node_id
778            )));
779        }
780    }
781    Ok(())
782}
783
784pub(crate) fn validate_shape_delta_for_task(delta: &ShapeDelta, task: &NodeTask) -> Result<()> {
785    let Some(shape_plan) = &task.node_plan.shape_plan else {
786        return Ok(());
787    };
788    if delta.kind == ShapeDeltaKind::Feature {
789        if let Some(expected) = &shape_plan.feature_schema_fingerprint {
790            if &delta.before_fingerprint != expected {
791                return Err(DagMlError::RuntimeValidation(format!(
792                    "node `{}` emitted feature shape delta from `{}`, expected current schema `{expected}`",
793                    task.node_plan.node_id, delta.before_fingerprint
794                )));
795            }
796        }
797    }
798    Ok(())
799}
800
801pub(crate) fn validate_prediction_scope(
802    prediction: &PredictionBlock,
803    task: &NodeTask,
804) -> Result<()> {
805    if prediction.partition != PredictionPartition::Validation {
806        return Ok(());
807    }
808    if prediction.fold_id != task.fold_id {
809        return Err(DagMlError::RuntimeValidation(format!(
810            "node `{}` emitted validation predictions for fold {:?}, expected {:?}",
811            task.node_plan.node_id, prediction.fold_id, task.fold_id
812        )));
813    }
814    if task.phase == Phase::FitCv
815        && task.fold_id.is_some()
816        && (!task.node_plan.data_bindings.is_empty() || !task.data_views.is_empty())
817    {
818        let validation_sample_ids = validation_view_sample_ids(task).ok_or_else(|| {
819            DagMlError::RuntimeValidation(format!(
820                "node `{}` emitted validation predictions without a fold-validation data view",
821                task.node_plan.node_id
822            ))
823        })?;
824        for sample_id in &prediction.sample_ids {
825            if !validation_sample_ids.contains(sample_id) {
826                return Err(DagMlError::RuntimeValidation(format!(
827                    "node `{}` emitted validation prediction for sample `{}` outside its validation view",
828                    task.node_plan.node_id, sample_id
829                )));
830            }
831        }
832    }
833    Ok(())
834}
835
836pub(crate) fn validate_observation_prediction_scope(
837    prediction: &ObservationPredictionBlock,
838    task: &NodeTask,
839) -> Result<()> {
840    if prediction.partition != PredictionPartition::Validation {
841        return Ok(());
842    }
843    if prediction.fold_id != task.fold_id {
844        return Err(DagMlError::RuntimeValidation(format!(
845            "node `{}` emitted observation validation predictions for fold {:?}, expected {:?}",
846            task.node_plan.node_id, prediction.fold_id, task.fold_id
847        )));
848    }
849    Ok(())
850}
851
852pub(crate) fn validate_aggregated_prediction_scope(
853    prediction: &AggregatedPredictionBlock,
854    task: &NodeTask,
855) -> Result<()> {
856    if prediction.partition != PredictionPartition::Validation {
857        return Ok(());
858    }
859    if prediction.fold_id != task.fold_id {
860        return Err(DagMlError::RuntimeValidation(format!(
861            "node `{}` emitted aggregated validation predictions for fold {:?}, expected {:?}",
862            task.node_plan.node_id, prediction.fold_id, task.fold_id
863        )));
864    }
865    // Sample-level aggregated validation units must stay inside this fold's
866    // validation view, mirroring `validate_prediction_scope`. Target / group
867    // units are checked against their relation set in the aggregation path.
868    if prediction.level == PredictionLevel::Sample
869        && task.phase == Phase::FitCv
870        && task.fold_id.is_some()
871        && (!task.node_plan.data_bindings.is_empty() || !task.data_views.is_empty())
872    {
873        if let Some(validation_sample_ids) = validation_view_sample_ids(task) {
874            for unit_id in &prediction.unit_ids {
875                if let PredictionUnitId::Sample(sample_id) = unit_id {
876                    if !validation_sample_ids.contains(sample_id) {
877                        return Err(DagMlError::RuntimeValidation(format!(
878                            "node `{}` emitted aggregated validation prediction for sample `{}` outside its validation view",
879                            task.node_plan.node_id, sample_id
880                        )));
881                    }
882                }
883            }
884        }
885    }
886    Ok(())
887}
888
889pub(crate) fn validation_view_sample_ids(task: &NodeTask) -> Option<BTreeSet<SampleId>> {
890    let mut sample_ids = BTreeSet::new();
891    for view in task
892        .data_views
893        .values()
894        .filter(|view| view.partition == DataRequestPartition::FoldValidation)
895    {
896        if let Some(view_sample_ids) = &view.sample_ids {
897            sample_ids.extend(view_sample_ids.iter().cloned());
898        }
899    }
900    (!sample_ids.is_empty()).then_some(sample_ids)
901}
902
903pub(crate) fn fit_influence_task_for_node(
904    plan: &ExecutionPlan,
905    node_plan: &NodePlan,
906    data_views: &BTreeMap<String, DataProviderViewSpec>,
907) -> Result<FitInfluenceTask> {
908    let manifest = plan
909        .controller_manifests
910        .get(&node_plan.controller_id)
911        .ok_or_else(|| {
912            DagMlError::RuntimeValidation(format!(
913                "node `{}` references missing controller manifest `{}`",
914                node_plan.node_id, node_plan.controller_id
915            ))
916        })?;
917    let Some(model_input_spec) = manifest.model_input_spec()? else {
918        return Ok(FitInfluenceTask::default());
919    };
920    let Some(requested_policy) = model_input_spec.fit_influence_policy else {
921        return Ok(FitInfluenceTask::default());
922    };
923    resolve_fit_influence_task(
924        requested_policy,
925        &node_plan.controller_capabilities,
926        data_views,
927    )
928}
929
930pub(crate) fn resolve_fit_influence_task(
931    requested_policy: FitInfluencePolicy,
932    capabilities: &BTreeSet<ControllerCapability>,
933    data_views: &BTreeMap<String, DataProviderViewSpec>,
934) -> Result<FitInfluenceTask> {
935    let row_weights = equal_sample_influence_weights(data_views);
936    match requested_policy {
937        FitInfluencePolicy::UniformRows => Ok(FitInfluenceTask {
938            requested_policy,
939            effective_policy: FitInfluencePolicy::UniformRows,
940            mechanism: FitInfluenceMechanism::UniformRows,
941            row_weights: Vec::new(),
942            warnings: Vec::new(),
943        }),
944        FitInfluencePolicy::ScorerOnly => Ok(FitInfluenceTask {
945            requested_policy,
946            effective_policy: FitInfluencePolicy::ScorerOnly,
947            mechanism: FitInfluenceMechanism::ScorerOnly,
948            row_weights: Vec::new(),
949            warnings: Vec::new(),
950        }),
951        FitInfluencePolicy::EqualSampleInfluence => {
952            require_fit_influence_support(capabilities, requested_policy)?;
953            let weights = row_weights.ok_or_else(|| {
954                DagMlError::RuntimeValidation(
955                    "equal_sample_influence requires task row sample ids".to_string(),
956                )
957            })?;
958            Ok(FitInfluenceTask {
959                requested_policy,
960                effective_policy: FitInfluencePolicy::EqualSampleInfluence,
961                mechanism: FitInfluenceMechanism::SampleWeights,
962                row_weights: weights,
963                warnings: Vec::new(),
964            })
965        }
966        FitInfluencePolicy::ResampleEqualized => {
967            require_fit_influence_support(capabilities, requested_policy)?;
968            Ok(FitInfluenceTask {
969                requested_policy,
970                effective_policy: FitInfluencePolicy::ResampleEqualized,
971                mechanism: FitInfluenceMechanism::RowResampling,
972                row_weights: Vec::new(),
973                warnings: Vec::new(),
974            })
975        }
976        FitInfluencePolicy::BackendLossWeight => {
977            require_fit_influence_support(capabilities, requested_policy)?;
978            let weights = row_weights.ok_or_else(|| {
979                DagMlError::RuntimeValidation(
980                    "backend_loss_weight requires task row sample ids".to_string(),
981                )
982            })?;
983            Ok(FitInfluenceTask {
984                requested_policy,
985                effective_policy: FitInfluencePolicy::BackendLossWeight,
986                mechanism: FitInfluenceMechanism::BackendLossWeights,
987                row_weights: weights,
988                warnings: Vec::new(),
989            })
990        }
991        FitInfluencePolicy::StrictWeightSupport => {
992            require_fit_influence_support(capabilities, requested_policy)?;
993            strict_fit_influence_task(capabilities, row_weights, requested_policy)
994        }
995        FitInfluencePolicy::Auto => Ok(auto_fit_influence_task(capabilities, row_weights)),
996    }
997}
998
999pub(crate) fn require_fit_influence_support(
1000    capabilities: &BTreeSet<ControllerCapability>,
1001    policy: FitInfluencePolicy,
1002) -> Result<()> {
1003    if capabilities_support_fit_influence(capabilities, policy) {
1004        return Ok(());
1005    }
1006    Err(DagMlError::RuntimeValidation(format!(
1007        "controller capabilities do not support requested fit influence policy {:?}",
1008        policy
1009    )))
1010}
1011
1012pub(crate) fn strict_fit_influence_task(
1013    capabilities: &BTreeSet<ControllerCapability>,
1014    row_weights: Option<Vec<f64>>,
1015    requested_policy: FitInfluencePolicy,
1016) -> Result<FitInfluenceTask> {
1017    if capabilities.contains(&ControllerCapability::SupportsBackendLossWeights) {
1018        let weights = row_weights.ok_or_else(|| {
1019            DagMlError::RuntimeValidation(
1020                "strict_weight_support with backend loss weights requires task row sample ids"
1021                    .to_string(),
1022            )
1023        })?;
1024        return Ok(FitInfluenceTask {
1025            requested_policy,
1026            effective_policy: FitInfluencePolicy::BackendLossWeight,
1027            mechanism: FitInfluenceMechanism::BackendLossWeights,
1028            row_weights: weights,
1029            warnings: Vec::new(),
1030        });
1031    }
1032    if capabilities.contains(&ControllerCapability::SupportsSampleWeights) {
1033        let weights = row_weights.ok_or_else(|| {
1034            DagMlError::RuntimeValidation(
1035                "strict_weight_support with sample weights requires task row sample ids"
1036                    .to_string(),
1037            )
1038        })?;
1039        return Ok(FitInfluenceTask {
1040            requested_policy,
1041            effective_policy: FitInfluencePolicy::EqualSampleInfluence,
1042            mechanism: FitInfluenceMechanism::SampleWeights,
1043            row_weights: weights,
1044            warnings: Vec::new(),
1045        });
1046    }
1047    Ok(FitInfluenceTask {
1048        requested_policy,
1049        effective_policy: FitInfluencePolicy::ResampleEqualized,
1050        mechanism: FitInfluenceMechanism::RowResampling,
1051        row_weights: Vec::new(),
1052        warnings: Vec::new(),
1053    })
1054}
1055
1056pub(crate) fn auto_fit_influence_task(
1057    capabilities: &BTreeSet<ControllerCapability>,
1058    row_weights: Option<Vec<f64>>,
1059) -> FitInfluenceTask {
1060    if capabilities.contains(&ControllerCapability::SupportsSampleWeights) {
1061        if let Some(weights) = row_weights.clone() {
1062            return FitInfluenceTask {
1063                requested_policy: FitInfluencePolicy::Auto,
1064                effective_policy: FitInfluencePolicy::EqualSampleInfluence,
1065                mechanism: FitInfluenceMechanism::SampleWeights,
1066                row_weights: weights,
1067                warnings: Vec::new(),
1068            };
1069        }
1070    }
1071    if capabilities.contains(&ControllerCapability::SupportsRowResampling) {
1072        return FitInfluenceTask {
1073            requested_policy: FitInfluencePolicy::Auto,
1074            effective_policy: FitInfluencePolicy::ResampleEqualized,
1075            mechanism: FitInfluenceMechanism::RowResampling,
1076            row_weights: Vec::new(),
1077            warnings: Vec::new(),
1078        };
1079    }
1080    if capabilities.contains(&ControllerCapability::SupportsBackendLossWeights) {
1081        if let Some(weights) = row_weights {
1082            return FitInfluenceTask {
1083                requested_policy: FitInfluencePolicy::Auto,
1084                effective_policy: FitInfluencePolicy::BackendLossWeight,
1085                mechanism: FitInfluenceMechanism::BackendLossWeights,
1086                row_weights: weights,
1087                warnings: Vec::new(),
1088            };
1089        }
1090    }
1091    FitInfluenceTask {
1092        requested_policy: FitInfluencePolicy::Auto,
1093        effective_policy: FitInfluencePolicy::UniformRows,
1094        mechanism: FitInfluenceMechanism::UniformRows,
1095        row_weights: Vec::new(),
1096        warnings: vec![
1097            "auto fit influence fell back to uniform_rows because no supported weighting capability was usable".to_string(),
1098        ],
1099    }
1100}
1101
1102pub(crate) fn equal_sample_influence_weights(
1103    data_views: &BTreeMap<String, DataProviderViewSpec>,
1104) -> Option<Vec<f64>> {
1105    let row_sample_ids = data_views
1106        .values()
1107        .filter(|view| {
1108            matches!(
1109                view.partition,
1110                DataRequestPartition::FoldTrain | DataRequestPartition::FullTrain
1111            )
1112        })
1113        .filter_map(|view| view.sample_ids.as_ref())
1114        .find(|sample_ids| !sample_ids.is_empty())
1115        .or_else(|| {
1116            data_views
1117                .values()
1118                .filter_map(|view| view.sample_ids.as_ref())
1119                .find(|sample_ids| !sample_ids.is_empty())
1120        })?;
1121    let mut counts = BTreeMap::<&SampleId, usize>::new();
1122    for sample_id in row_sample_ids {
1123        *counts.entry(sample_id).or_default() += 1;
1124    }
1125    Some(
1126        row_sample_ids
1127            .iter()
1128            .map(|sample_id| 1.0 / *counts.get(sample_id).expect("counted sample id") as f64)
1129            .collect(),
1130    )
1131}
1132
1133pub(crate) fn record_fit_influence_diagnostic(task: &NodeTask, result: &mut NodeResult) {
1134    if task.fit_influence.is_default() || !result.fit_influence_diagnostics.is_empty() {
1135        return;
1136    }
1137    result
1138        .fit_influence_diagnostics
1139        .push(task.fit_influence.diagnostic());
1140}
dag_ml_core/runtime/task.rs

dag_ml_core/runtime/
task.rs
