dag_ml_core/runtime/
task.rs

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

dag_ml_core/runtime/
task.rs
