dag_ml_core/runtime/
dataview.rs

1// Auto-split from the former monolithic `runtime.rs` (pure refactor).
2use super::*;
3
4#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
5pub struct DataMaterializationRequest {
6    pub run_id: RunId,
7    pub node_id: NodeId,
8    pub input_name: String,
9    pub phase: Phase,
10    pub variant_id: Option<VariantId>,
11    pub fold_id: Option<FoldId>,
12    pub binding: crate::data::DataBinding,
13}
14
15#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
16pub struct DataProviderViewSpec {
17    #[serde(default)]
18    pub sample_ids: Option<Vec<SampleId>>,
19    pub partition: DataRequestPartition,
20    #[serde(default)]
21    pub fold_id: Option<FoldId>,
22    #[serde(default)]
23    pub source_ids: Option<Vec<String>>,
24    #[serde(default)]
25    pub columns: Option<Vec<String>>,
26    pub include_augmented: bool,
27    pub include_excluded: bool,
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub branch_view: Option<crate::data::BranchViewPlan>,
30    #[serde(default)]
31    pub extra: BTreeMap<String, serde_json::Value>,
32}
33
34pub const DATA_OUTPUT_PROVENANCE_KEY: &str = "dag_ml_output";
35pub const DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION: u32 = 1;
36pub const DATA_OUTPUT_PROVENANCE_SCHEMA_ID: &str =
37    "https://github.com/GBeurier/dag-ml/schemas/data_output_provenance.v1.schema.json";
38pub const NODE_TASK_SCHEMA_VERSION: u32 = 1;
39pub const NODE_TASK_SCHEMA_ID: &str =
40    "https://github.com/GBeurier/dag-ml/schemas/node_task.v1.schema.json";
41pub const NODE_RESULT_SCHEMA_VERSION: u32 = 1;
42pub const NODE_RESULT_SCHEMA_ID: &str =
43    "https://github.com/GBeurier/dag-ml/schemas/node_result.v1.schema.json";
44
45pub(crate) fn default_data_output_provenance_schema_version() -> u32 {
46    DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION
47}
48
49impl DataProviderViewSpec {
50    pub fn validate(&self) -> Result<()> {
51        validate_optional_ids("sample id", &self.sample_ids)?;
52        validate_optional_strings("source id", &self.source_ids)?;
53        validate_optional_strings("column", &self.columns)?;
54        match self.partition {
55            DataRequestPartition::FoldTrain | DataRequestPartition::FoldValidation => {
56                if self.sample_ids.is_some() && self.fold_id.is_none() {
57                    return Err(DagMlError::RuntimeValidation(format!(
58                        "data provider view {:?} with explicit sample ids requires a fold id",
59                        self.partition
60                    )));
61                }
62            }
63            DataRequestPartition::FullTrain | DataRequestPartition::Predict => {
64                if self.fold_id.is_some() {
65                    return Err(DagMlError::RuntimeValidation(format!(
66                        "data provider view {:?} must not carry a fold id",
67                        self.partition
68                    )));
69                }
70            }
71        }
72        for key in self.extra.keys() {
73            if key.trim().is_empty() {
74                return Err(DagMlError::RuntimeValidation(
75                    "data provider view extra contains an empty key".to_string(),
76                ));
77            }
78        }
79        if let Some(branch_view) = &self.branch_view {
80            branch_view.validate()?;
81        }
82        self.output_provenance()?;
83        Ok(())
84    }
85
86    pub fn output_provenance(&self) -> Result<Option<DataOutputProvenance>> {
87        let Some(value) = self.extra.get(DATA_OUTPUT_PROVENANCE_KEY) else {
88            return Ok(None);
89        };
90        let provenance: DataOutputProvenance = serde_json::from_value(value.clone())?;
91        provenance.validate()?;
92        Ok(Some(provenance))
93    }
94}
95
96#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
97pub struct DataOutputProvenance {
98    #[serde(default = "default_data_output_provenance_schema_version")]
99    pub schema_version: u32,
100    pub producer_node: NodeId,
101    pub producer_port: String,
102    pub producer_phase: Phase,
103    #[serde(default)]
104    pub variant_id: Option<VariantId>,
105    #[serde(default)]
106    pub fold_id: Option<FoldId>,
107    #[serde(default)]
108    pub shape_plan_fingerprint: Option<String>,
109    #[serde(default)]
110    pub aggregation_policy_fingerprint: Option<String>,
111    #[serde(default)]
112    pub feature_namespace: Option<String>,
113    #[serde(default)]
114    pub feature_schema_fingerprint: Option<String>,
115    #[serde(default, skip_serializing_if = "Option::is_none")]
116    pub representation_plan: Option<RepresentationPlan>,
117    #[serde(default, skip_serializing_if = "Option::is_none")]
118    pub representation_replay_manifest: Option<RepresentationReplayManifest>,
119    #[serde(default, skip_serializing_if = "Option::is_none")]
120    pub representation_compatibility: Option<RepresentationCompatibilityReport>,
121    #[serde(default, skip_serializing_if = "Option::is_none")]
122    pub relation_delta_fingerprint: Option<String>,
123    #[serde(default)]
124    pub shape_deltas: Vec<ShapeDelta>,
125}
126
127impl DataOutputProvenance {
128    pub fn validate(&self) -> Result<()> {
129        if self.schema_version != DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION {
130            return Err(DagMlError::RuntimeValidation(format!(
131                "data output provenance for `{}` uses unsupported schema_version {}, expected {}",
132                self.producer_node, self.schema_version, DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION
133            )));
134        }
135        if self.producer_port.trim().is_empty() {
136            return Err(DagMlError::RuntimeValidation(format!(
137                "data output provenance for `{}` has empty producer_port",
138                self.producer_node
139            )));
140        }
141        validate_optional_fingerprint(
142            "shape_plan_fingerprint",
143            &self.shape_plan_fingerprint,
144            &self.producer_node,
145        )?;
146        validate_optional_fingerprint(
147            "aggregation_policy_fingerprint",
148            &self.aggregation_policy_fingerprint,
149            &self.producer_node,
150        )?;
151        validate_optional_fingerprint(
152            "feature_schema_fingerprint",
153            &self.feature_schema_fingerprint,
154            &self.producer_node,
155        )?;
156        validate_optional_fingerprint(
157            "relation_delta_fingerprint",
158            &self.relation_delta_fingerprint,
159            &self.producer_node,
160        )?;
161        if let Some(representation_plan) = &self.representation_plan {
162            representation_plan.validate().map_err(|error| {
163                DagMlError::RuntimeValidation(format!(
164                    "data output provenance for `{}` has invalid representation_plan: {error}",
165                    self.producer_node
166                ))
167            })?;
168        }
169        if let Some(replay_manifest) = &self.representation_replay_manifest {
170            replay_manifest.validate().map_err(|error| {
171                DagMlError::RuntimeValidation(format!(
172                    "data output provenance for `{}` has invalid representation_replay_manifest: {error}",
173                    self.producer_node
174                ))
175            })?;
176        }
177        if let Some(report) = &self.representation_compatibility {
178            report.validate().map_err(|error| {
179                DagMlError::RuntimeValidation(format!(
180                    "data output provenance for `{}` has invalid representation_compatibility: {error}",
181                    self.producer_node
182                ))
183            })?;
184        }
185        if self
186            .feature_namespace
187            .as_ref()
188            .is_some_and(|namespace| namespace.trim().is_empty())
189        {
190            return Err(DagMlError::RuntimeValidation(format!(
191                "data output provenance for `{}` has empty feature_namespace",
192                self.producer_node
193            )));
194        }
195        for delta in &self.shape_deltas {
196            delta.validate()?;
197            if delta.node_id != self.producer_node {
198                return Err(DagMlError::RuntimeValidation(format!(
199                    "data output provenance for `{}` contains shape delta for `{}`",
200                    self.producer_node, delta.node_id
201                )));
202            }
203        }
204        if let Some(feature_schema_fingerprint) = &self.feature_schema_fingerprint {
205            if let Some(last_feature_delta) = self
206                .shape_deltas
207                .iter()
208                .rev()
209                .find(|delta| delta.kind == ShapeDeltaKind::Feature)
210            {
211                if &last_feature_delta.after_fingerprint != feature_schema_fingerprint {
212                    return Err(DagMlError::RuntimeValidation(format!(
213                        "data output provenance for `{}` has feature_schema_fingerprint `{feature_schema_fingerprint}` but last feature delta ends at `{}`",
214                        self.producer_node, last_feature_delta.after_fingerprint
215                    )));
216                }
217            }
218        }
219        Ok(())
220    }
221}
222
223pub(crate) fn validate_optional_fingerprint(
224    label: &str,
225    fingerprint: &Option<String>,
226    producer_node: &NodeId,
227) -> Result<()> {
228    let Some(fingerprint) = fingerprint else {
229        return Ok(());
230    };
231    if fingerprint.len() != 64 || !fingerprint.bytes().all(|byte| byte.is_ascii_hexdigit()) {
232        return Err(DagMlError::RuntimeValidation(format!(
233            "data output provenance for `{producer_node}` has invalid {label}"
234        )));
235    }
236    Ok(())
237}
238
239pub(crate) fn validate_optional_ids<T>(label: &str, values: &Option<Vec<T>>) -> Result<()>
240where
241    T: Ord + ToString,
242{
243    let Some(values) = values else {
244        return Ok(());
245    };
246    if values.is_empty() {
247        return Err(DagMlError::RuntimeValidation(format!(
248            "data provider view {label} list is empty"
249        )));
250    }
251    let mut seen = BTreeSet::new();
252    for value in values {
253        if !seen.insert(value) {
254            return Err(DagMlError::RuntimeValidation(format!(
255                "data provider view has duplicate {label} `{}`",
256                value.to_string()
257            )));
258        }
259    }
260    Ok(())
261}
262
263pub(crate) fn validate_optional_strings(label: &str, values: &Option<Vec<String>>) -> Result<()> {
264    let Some(values) = values else {
265        return Ok(());
266    };
267    if values.is_empty() {
268        return Err(DagMlError::RuntimeValidation(format!(
269            "data provider view {label} list is empty"
270        )));
271    }
272    let mut seen = BTreeSet::new();
273    for value in values {
274        if value.trim().is_empty() {
275            return Err(DagMlError::RuntimeValidation(format!(
276                "data provider view contains an empty {label}"
277            )));
278        }
279        if !seen.insert(value.as_str()) {
280            return Err(DagMlError::RuntimeValidation(format!(
281                "data provider view has duplicate {label} `{value}`"
282            )));
283        }
284    }
285    Ok(())
286}
287
288#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
289pub struct DataViewRequest {
290    pub run_id: RunId,
291    pub node_id: NodeId,
292    pub input_name: String,
293    pub phase: Phase,
294    pub variant_id: Option<VariantId>,
295    pub fold_id: Option<FoldId>,
296    pub binding: crate::data::DataBinding,
297    pub data_handle: HandleRef,
298    pub view: DataProviderViewSpec,
299}
300
301pub trait RuntimeDataProvider {
302    fn materialize(&self, request: &DataMaterializationRequest) -> Result<HandleRef>;
303    fn make_view(&self, request: &DataViewRequest) -> Result<HandleRef>;
304    fn coordinator_relations(&self, _binding: &DataBinding) -> Result<Option<SampleRelationSet>> {
305        Ok(None)
306    }
307}
308
309pub trait RuntimeController: Send + Sync {
310    fn controller_id(&self) -> &ControllerId;
311    fn invoke(&self, task: &NodeTask) -> Result<NodeResult>;
312
313    fn invoke_aggregation(
314        &self,
315        task: &AggregationControllerTask,
316    ) -> Result<AggregationControllerResult> {
317        Err(DagMlError::RuntimeValidation(format!(
318            "runtime controller `{}` does not implement aggregation task `{}`",
319            self.controller_id(),
320            task.task_id
321        )))
322    }
323}
324pub(crate) struct CollectedInputs {
325    pub(crate) handles: BTreeMap<String, HandleRef>,
326    pub(crate) data_views: BTreeMap<String, DataProviderViewSpec>,
327    pub(crate) prediction_inputs: BTreeMap<String, PredictionInputSpec>,
328    pub(crate) skip_node: bool,
329}
330
331pub(crate) fn data_view_key(input_name: &str) -> String {
332    format!("data:{input_name}")
333}
334
335pub(crate) fn validation_data_view_key(input_name: &str) -> String {
336    format!("{input_name}:validation")
337}
338
339pub(crate) fn derive_output_data_views(
340    plan: &ExecutionPlan,
341    task: &NodeTask,
342    result: &NodeResult,
343) -> Result<BTreeMap<String, DataProviderViewSpec>> {
344    let node = plan
345        .graph_plan
346        .graph
347        .nodes
348        .iter()
349        .find(|node| node.id == task.node_plan.node_id)
350        .expect("execution plan was validated");
351    let mut views = BTreeMap::new();
352    for port in node
353        .ports
354        .outputs
355        .iter()
356        .filter(|port| port.kind == PortKind::Data)
357    {
358        let Some(handle) = result.outputs.get(&port.name) else {
359            continue;
360        };
361        if !matches!(handle.kind, HandleKind::Data | HandleKind::DataView) {
362            return Err(DagMlError::RuntimeValidation(format!(
363                "node `{}` emitted data output `{}` with non-data/data-view handle kind {:?}",
364                task.node_plan.node_id, port.name, handle.kind
365            )));
366        }
367        if let Some(view) = primary_output_data_view(task) {
368            views.insert(
369                port.name.clone(),
370                output_data_view_for_port(task, result, &port.name, view)?,
371            );
372        }
373        if let Some(validation_view) = validation_output_data_view(task) {
374            views.insert(
375                validation_data_view_key(&port.name),
376                output_data_view_for_port(task, result, &port.name, validation_view)?,
377            );
378        }
379    }
380    Ok(views)
381}
382
383pub(crate) fn output_data_view_for_port(
384    task: &NodeTask,
385    result: &NodeResult,
386    port_name: &str,
387    base_view: &DataProviderViewSpec,
388) -> Result<DataProviderViewSpec> {
389    let mut view = base_view.clone();
390    if let Some(upstream_provenance) = view.extra.remove(DATA_OUTPUT_PROVENANCE_KEY) {
391        let provenance: DataOutputProvenance =
392            serde_json::from_value(upstream_provenance).map_err(|error| {
393                DagMlError::RuntimeValidation(format!(
394                    "node `{}` cannot propagate data output `{port_name}` because upstream data output provenance is invalid JSON: {error}",
395                    task.node_plan.node_id
396                ))
397            })?;
398        provenance.validate().map_err(|error| {
399            DagMlError::RuntimeValidation(format!(
400                "node `{}` cannot propagate data output `{port_name}` because upstream data output provenance is invalid: {error}",
401                task.node_plan.node_id
402            ))
403        })?;
404    }
405    let shape_deltas = result
406        .shape_deltas
407        .iter()
408        .filter(|delta| delta.node_id == task.node_plan.node_id)
409        .cloned()
410        .collect::<Vec<_>>();
411    let mut provenance = DataOutputProvenance {
412        schema_version: DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION,
413        producer_node: task.node_plan.node_id.clone(),
414        producer_port: port_name.to_string(),
415        producer_phase: task.phase,
416        variant_id: task.variant_id.clone(),
417        fold_id: task.fold_id.clone(),
418        shape_plan_fingerprint: None,
419        aggregation_policy_fingerprint: None,
420        feature_namespace: None,
421        feature_schema_fingerprint: None,
422        representation_plan: None,
423        representation_replay_manifest: None,
424        representation_compatibility: None,
425        relation_delta_fingerprint: None,
426        shape_deltas,
427    };
428    if let Some(shape_plan) = &task.node_plan.shape_plan {
429        provenance.shape_plan_fingerprint = Some(stable_json_fingerprint(shape_plan)?);
430        provenance.aggregation_policy_fingerprint =
431            Some(stable_json_fingerprint(&shape_plan.aggregation_policy)?);
432        provenance.feature_namespace = shape_plan.feature_namespace.clone();
433        provenance.feature_schema_fingerprint =
434            output_feature_schema_fingerprint(shape_plan, result);
435    }
436    provenance.validate()?;
437
438    view.extra.insert(
439        DATA_OUTPUT_PROVENANCE_KEY.to_string(),
440        serde_json::to_value(provenance)?,
441    );
442    view.validate()?;
443    Ok(view)
444}
445
446pub(crate) fn output_feature_schema_fingerprint(
447    shape_plan: &crate::policy::DataModelShapePlan,
448    result: &NodeResult,
449) -> Option<String> {
450    result
451        .shape_deltas
452        .iter()
453        .rev()
454        .find(|delta| delta.kind == ShapeDeltaKind::Feature)
455        .map(|delta| delta.after_fingerprint.clone())
456        .or_else(|| shape_plan.feature_schema_fingerprint.clone())
457}
458
459pub(crate) fn primary_output_data_view(task: &NodeTask) -> Option<&DataProviderViewSpec> {
460    task.data_views
461        .values()
462        .find(|view| view.partition != DataRequestPartition::FoldValidation)
463        .or_else(|| task.data_views.values().next())
464}
465
466pub(crate) fn validation_output_data_view(task: &NodeTask) -> Option<&DataProviderViewSpec> {
467    task.data_views
468        .values()
469        .find(|view| view.partition == DataRequestPartition::FoldValidation)
470}
471
472pub(crate) fn make_data_view_handle(
473    data_provider: &dyn RuntimeDataProvider,
474    ctx: &RunContext,
475    node_plan: &NodePlan,
476    scope: &PhaseScope,
477    binding: &DataBinding,
478    data_handle: &HandleRef,
479    view: &DataProviderViewSpec,
480) -> Result<HandleRef> {
481    view.validate()?;
482    let view_handle = data_provider.make_view(&DataViewRequest {
483        run_id: ctx.run_id.clone(),
484        node_id: node_plan.node_id.clone(),
485        input_name: binding.input_name.clone(),
486        phase: scope.phase,
487        variant_id: scope.variant_id.clone(),
488        fold_id: scope.fold_id.clone(),
489        binding: binding.clone(),
490        data_handle: data_handle.clone(),
491        view: view.clone(),
492    })?;
493    // A data view is delivered to the controller as a data input, so the
494    // provider must return a data-bearing handle. Refuse a model / artifact /
495    // prediction / relation handle masquerading as a view across the ABI.
496    if !matches!(view_handle.kind, HandleKind::Data | HandleKind::DataView) {
497        return Err(DagMlError::RuntimeValidation(format!(
498            "node `{}` data view `{}` resolved to a non-data/data-view handle kind {:?}",
499            node_plan.node_id, binding.input_name, view_handle.kind
500        )));
501    }
502    Ok(view_handle)
503}
504
505pub(crate) fn data_view_for_scope(
506    binding: &DataBinding,
507    fold_set: Option<&FoldSet>,
508    scope: &PhaseScope,
509    branch_view: Option<&crate::data::BranchViewPlan>,
510    excluded_samples: &BTreeSet<SampleId>,
511) -> Result<DataProviderViewSpec> {
512    let partition = data_partition_for_scope(binding, scope);
513    // During FIT_CV and REFIT this primary view IS the training input; during
514    // PREDICT/EXPLAIN (and the planning phases) it is a non-fit read.
515    let role = match scope.phase {
516        Phase::FitCv | Phase::Refit => DataViewRole::Fit,
517        _ => DataViewRole::NonFit,
518    };
519    data_view_for_partition(
520        binding,
521        fold_set,
522        scope,
523        partition,
524        branch_view,
525        role,
526        excluded_samples,
527    )
528}
529
530pub(crate) fn validation_data_view_for_scope(
531    binding: &DataBinding,
532    fold_set: Option<&FoldSet>,
533    scope: &PhaseScope,
534    branch_view: Option<&crate::data::BranchViewPlan>,
535    excluded_samples: &BTreeSet<SampleId>,
536) -> Result<Option<DataProviderViewSpec>> {
537    if scope.phase != Phase::FitCv || scope.fold_id.is_none() {
538        return Ok(None);
539    }
540    let partition = binding.view_policy.predict_partition;
541    if partition == data_partition_for_scope(binding, scope) {
542        return Ok(None);
543    }
544    // This is the validation companion read, never the training input.
545    data_view_for_partition(
546        binding,
547        fold_set,
548        scope,
549        partition,
550        branch_view,
551        DataViewRole::NonFit,
552        excluded_samples,
553    )
554    .map(Some)
555}
dag_ml_core/runtime/dataview.rs

dag_ml_core/runtime/
dataview.rs
