1use 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 #[serde(default, skip_serializing_if = "Option::is_none")]
17 pub predict_cohort: Option<crate::data::PredictCohort>,
18}
19
20#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
21pub struct DataProviderViewSpec {
22 #[serde(default)]
23 pub sample_ids: Option<Vec<SampleId>>,
24 pub partition: DataRequestPartition,
25 #[serde(default)]
26 pub fold_id: Option<FoldId>,
27 #[serde(default)]
28 pub source_ids: Option<Vec<String>>,
29 #[serde(default)]
30 pub columns: Option<Vec<String>>,
31 pub include_augmented: bool,
32 pub include_excluded: bool,
33 #[serde(default, skip_serializing_if = "Option::is_none")]
34 pub branch_view: Option<crate::data::BranchViewPlan>,
35 #[serde(default)]
36 pub extra: BTreeMap<String, serde_json::Value>,
37}
38
39pub const DATA_OUTPUT_PROVENANCE_KEY: &str = "dag_ml_output";
40pub const DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION: u32 = 1;
41pub const DATA_OUTPUT_PROVENANCE_SCHEMA_ID: &str =
42 "https://github.com/GBeurier/dag-ml/schemas/data_output_provenance.v1.schema.json";
43pub const NODE_TASK_SCHEMA_VERSION: u32 = 1;
44pub const NODE_TASK_SCHEMA_ID: &str =
45 "https://github.com/GBeurier/dag-ml/schemas/node_task.v1.schema.json";
46pub const NODE_RESULT_SCHEMA_VERSION: u32 = 1;
47pub const NODE_RESULT_SCHEMA_ID: &str =
48 "https://github.com/GBeurier/dag-ml/schemas/node_result.v1.schema.json";
49
50pub(crate) fn default_data_output_provenance_schema_version() -> u32 {
51 DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION
52}
53
54impl DataProviderViewSpec {
55 pub fn validate(&self) -> Result<()> {
56 validate_optional_ids("sample id", &self.sample_ids)?;
57 validate_optional_strings("source id", &self.source_ids)?;
58 validate_optional_strings("column", &self.columns)?;
59 match self.partition {
60 DataRequestPartition::FoldTrain | DataRequestPartition::FoldValidation => {
61 if self.sample_ids.is_some() && self.fold_id.is_none() {
62 return Err(DagMlError::RuntimeValidation(format!(
63 "data provider view {:?} with explicit sample ids requires a fold id",
64 self.partition
65 )));
66 }
67 }
68 DataRequestPartition::FullTrain | DataRequestPartition::Predict => {
69 if self.fold_id.is_some() {
70 return Err(DagMlError::RuntimeValidation(format!(
71 "data provider view {:?} must not carry a fold id",
72 self.partition
73 )));
74 }
75 }
76 }
77 for key in self.extra.keys() {
78 if key.trim().is_empty() {
79 return Err(DagMlError::RuntimeValidation(
80 "data provider view extra contains an empty key".to_string(),
81 ));
82 }
83 }
84 if let Some(branch_view) = &self.branch_view {
85 branch_view.validate()?;
86 }
87 self.output_provenance()?;
88 Ok(())
89 }
90
91 pub fn output_provenance(&self) -> Result<Option<DataOutputProvenance>> {
92 let Some(value) = self.extra.get(DATA_OUTPUT_PROVENANCE_KEY) else {
93 return Ok(None);
94 };
95 let provenance: DataOutputProvenance = serde_json::from_value(value.clone())?;
96 provenance.validate()?;
97 Ok(Some(provenance))
98 }
99}
100
101#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
102pub struct DataOutputProvenance {
103 #[serde(default = "default_data_output_provenance_schema_version")]
104 pub schema_version: u32,
105 pub producer_node: NodeId,
106 pub producer_port: String,
107 pub producer_phase: Phase,
108 #[serde(default)]
109 pub variant_id: Option<VariantId>,
110 #[serde(default)]
111 pub fold_id: Option<FoldId>,
112 #[serde(default)]
113 pub shape_plan_fingerprint: Option<String>,
114 #[serde(default)]
115 pub aggregation_policy_fingerprint: Option<String>,
116 #[serde(default)]
117 pub feature_namespace: Option<String>,
118 #[serde(default)]
119 pub feature_schema_fingerprint: Option<String>,
120 #[serde(default, skip_serializing_if = "Option::is_none")]
121 pub representation_plan: Option<RepresentationPlan>,
122 #[serde(default, skip_serializing_if = "Option::is_none")]
123 pub representation_replay_manifest: Option<RepresentationReplayManifest>,
124 #[serde(default, skip_serializing_if = "Option::is_none")]
125 pub representation_compatibility: Option<RepresentationCompatibilityReport>,
126 #[serde(default, skip_serializing_if = "Option::is_none")]
127 pub relation_delta_fingerprint: Option<String>,
128 #[serde(default)]
129 pub shape_deltas: Vec<ShapeDelta>,
130}
131
132impl DataOutputProvenance {
133 pub fn validate(&self) -> Result<()> {
134 if self.schema_version != DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION {
135 return Err(DagMlError::RuntimeValidation(format!(
136 "data output provenance for `{}` uses unsupported schema_version {}, expected {}",
137 self.producer_node, self.schema_version, DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION
138 )));
139 }
140 if self.producer_port.trim().is_empty() {
141 return Err(DagMlError::RuntimeValidation(format!(
142 "data output provenance for `{}` has empty producer_port",
143 self.producer_node
144 )));
145 }
146 validate_optional_fingerprint(
147 "shape_plan_fingerprint",
148 &self.shape_plan_fingerprint,
149 &self.producer_node,
150 )?;
151 validate_optional_fingerprint(
152 "aggregation_policy_fingerprint",
153 &self.aggregation_policy_fingerprint,
154 &self.producer_node,
155 )?;
156 validate_optional_fingerprint(
157 "feature_schema_fingerprint",
158 &self.feature_schema_fingerprint,
159 &self.producer_node,
160 )?;
161 validate_optional_fingerprint(
162 "relation_delta_fingerprint",
163 &self.relation_delta_fingerprint,
164 &self.producer_node,
165 )?;
166 if let Some(representation_plan) = &self.representation_plan {
167 representation_plan.validate().map_err(|error| {
168 DagMlError::RuntimeValidation(format!(
169 "data output provenance for `{}` has invalid representation_plan: {error}",
170 self.producer_node
171 ))
172 })?;
173 }
174 if let Some(replay_manifest) = &self.representation_replay_manifest {
175 replay_manifest.validate().map_err(|error| {
176 DagMlError::RuntimeValidation(format!(
177 "data output provenance for `{}` has invalid representation_replay_manifest: {error}",
178 self.producer_node
179 ))
180 })?;
181 }
182 if let Some(report) = &self.representation_compatibility {
183 report.validate().map_err(|error| {
184 DagMlError::RuntimeValidation(format!(
185 "data output provenance for `{}` has invalid representation_compatibility: {error}",
186 self.producer_node
187 ))
188 })?;
189 }
190 if self
191 .feature_namespace
192 .as_ref()
193 .is_some_and(|namespace| namespace.trim().is_empty())
194 {
195 return Err(DagMlError::RuntimeValidation(format!(
196 "data output provenance for `{}` has empty feature_namespace",
197 self.producer_node
198 )));
199 }
200 for delta in &self.shape_deltas {
201 delta.validate()?;
202 if delta.node_id != self.producer_node {
203 return Err(DagMlError::RuntimeValidation(format!(
204 "data output provenance for `{}` contains shape delta for `{}`",
205 self.producer_node, delta.node_id
206 )));
207 }
208 }
209 if let Some(feature_schema_fingerprint) = &self.feature_schema_fingerprint {
210 if let Some(last_feature_delta) = self
211 .shape_deltas
212 .iter()
213 .rev()
214 .find(|delta| delta.kind == ShapeDeltaKind::Feature)
215 {
216 if &last_feature_delta.after_fingerprint != feature_schema_fingerprint {
217 return Err(DagMlError::RuntimeValidation(format!(
218 "data output provenance for `{}` has feature_schema_fingerprint `{feature_schema_fingerprint}` but last feature delta ends at `{}`",
219 self.producer_node, last_feature_delta.after_fingerprint
220 )));
221 }
222 }
223 }
224 Ok(())
225 }
226}
227
228pub(crate) fn validate_optional_fingerprint(
229 label: &str,
230 fingerprint: &Option<String>,
231 producer_node: &NodeId,
232) -> Result<()> {
233 let Some(fingerprint) = fingerprint else {
234 return Ok(());
235 };
236 if fingerprint.len() != 64 || !fingerprint.bytes().all(|byte| byte.is_ascii_hexdigit()) {
237 return Err(DagMlError::RuntimeValidation(format!(
238 "data output provenance for `{producer_node}` has invalid {label}"
239 )));
240 }
241 Ok(())
242}
243
244pub(crate) fn validate_optional_ids<T>(label: &str, values: &Option<Vec<T>>) -> Result<()>
245where
246 T: Ord + ToString,
247{
248 let Some(values) = values else {
249 return Ok(());
250 };
251 if values.is_empty() {
252 return Err(DagMlError::RuntimeValidation(format!(
253 "data provider view {label} list is empty"
254 )));
255 }
256 let mut seen = BTreeSet::new();
257 for value in values {
258 if !seen.insert(value) {
259 return Err(DagMlError::RuntimeValidation(format!(
260 "data provider view has duplicate {label} `{}`",
261 value.to_string()
262 )));
263 }
264 }
265 Ok(())
266}
267
268pub(crate) fn validate_optional_strings(label: &str, values: &Option<Vec<String>>) -> Result<()> {
269 let Some(values) = values else {
270 return Ok(());
271 };
272 if values.is_empty() {
273 return Err(DagMlError::RuntimeValidation(format!(
274 "data provider view {label} list is empty"
275 )));
276 }
277 let mut seen = BTreeSet::new();
278 for value in values {
279 if value.trim().is_empty() {
280 return Err(DagMlError::RuntimeValidation(format!(
281 "data provider view contains an empty {label}"
282 )));
283 }
284 if !seen.insert(value.as_str()) {
285 return Err(DagMlError::RuntimeValidation(format!(
286 "data provider view has duplicate {label} `{value}`"
287 )));
288 }
289 }
290 Ok(())
291}
292
293#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
294pub struct DataViewRequest {
295 pub run_id: RunId,
296 pub node_id: NodeId,
297 pub input_name: String,
298 pub phase: Phase,
299 pub variant_id: Option<VariantId>,
300 pub fold_id: Option<FoldId>,
301 pub binding: crate::data::DataBinding,
302 pub data_handle: HandleRef,
303 pub view: DataProviderViewSpec,
304 #[serde(default, skip_serializing_if = "Option::is_none")]
308 pub predict_cohort: Option<crate::data::PredictCohort>,
309}
310
311pub trait RuntimeDataProvider {
312 fn materialize(&self, request: &DataMaterializationRequest) -> Result<HandleRef>;
313 fn make_view(&self, request: &DataViewRequest) -> Result<HandleRef>;
314 fn training_data_identity(
320 &self,
321 _binding: &DataBinding,
322 ) -> Result<Option<crate::training::TrainingDataIdentity>> {
323 Ok(None)
324 }
325 fn coordinator_relations(&self, _binding: &DataBinding) -> Result<Option<SampleRelationSet>> {
326 Ok(None)
327 }
328
329 fn refit_sample_ids(&self, _binding: &DataBinding) -> Result<Option<Vec<SampleId>>> {
334 Ok(None)
335 }
336
337 fn predict_cohort(
341 &self,
342 _binding: &DataBinding,
343 phase: Phase,
344 ) -> Result<Option<crate::data::PredictCohort>> {
345 validate_predict_cohort_phase(phase)?;
346 Ok(None)
347 }
348
349 fn methods_pls_capability(&self) -> Result<()> {
354 Err(DagMlError::RuntimeValidation(
355 "runtime data provider does not implement the portable Methods PLS numeric view"
356 .to_string(),
357 ))
358 }
359
360 fn preflight_methods_pls(&self, request: &MethodsPlsDataRequest) -> Result<()> {
361 request.validate()?;
362 self.methods_pls_capability()
363 }
364
365 fn methods_pls_data(&self, _request: &MethodsPlsDataRequest) -> Result<MethodsPlsData> {
370 Err(DagMlError::RuntimeValidation(
371 "runtime data provider does not implement the portable Methods PLS numeric view"
372 .to_string(),
373 ))
374 }
375}
376
377fn validate_predict_cohort_phase(phase: Phase) -> Result<()> {
378 if phase != Phase::Predict {
379 return Err(DagMlError::RuntimeValidation(format!(
380 "predict cohort may be requested only during PREDICT, got {phase:?}"
381 )));
382 }
383 Ok(())
384}
385
386#[derive(Clone, Debug, PartialEq)]
389pub struct MethodsPlsMatrix {
390 pub values: Vec<f64>,
391 pub rows: usize,
392 pub cols: usize,
393}
394
395impl MethodsPlsMatrix {
396 pub fn validate(&self, label: &str) -> Result<()> {
397 if self.rows == 0
398 || self.cols == 0
399 || self.rows.checked_mul(self.cols) != Some(self.values.len())
400 {
401 return Err(DagMlError::RuntimeValidation(format!(
402 "portable Methods PLS {label} matrix has invalid row-major dimensions"
403 )));
404 }
405 if self.values.iter().any(|value| !value.is_finite()) {
406 return Err(DagMlError::RuntimeValidation(format!(
407 "portable Methods PLS {label} matrix contains a non-finite value"
408 )));
409 }
410 Ok(())
411 }
412}
413
414#[derive(Clone, Debug, PartialEq)]
416pub struct MethodsPlsDataset {
417 pub sample_ids: Vec<SampleId>,
418 pub x: MethodsPlsMatrix,
419 pub y: Option<MethodsPlsMatrix>,
423 pub target_names: Vec<String>,
424}
425
426impl MethodsPlsDataset {
427 pub fn validate(&self, label: &str, require_targets: bool) -> Result<()> {
428 self.x.validate(&format!("{label}.x"))?;
429 if self.sample_ids.len() != self.x.rows {
430 return Err(DagMlError::RuntimeValidation(format!(
431 "portable Methods PLS {label} rows do not match sample identities"
432 )));
433 }
434 if self.target_names.is_empty()
435 || self.target_names.iter().any(|name| name.trim().is_empty())
436 {
437 return Err(DagMlError::RuntimeValidation(format!(
438 "portable Methods PLS {label} has invalid target names"
439 )));
440 }
441 match &self.y {
442 Some(y) => {
443 y.validate(&format!("{label}.y"))?;
444 if self.sample_ids.len() != y.rows || self.target_names.len() != y.cols {
445 return Err(DagMlError::RuntimeValidation(format!(
446 "portable Methods PLS {label} targets do not match sample identities or target names"
447 )));
448 }
449 }
450 None if require_targets => {
451 return Err(DagMlError::RuntimeValidation(format!(
452 "portable Methods PLS {label} requires targets for fitting or CV scoring"
453 )))
454 }
455 None => {}
456 }
457 let unique = self.sample_ids.iter().collect::<BTreeSet<_>>();
458 if unique.len() != self.sample_ids.len() {
459 return Err(DagMlError::RuntimeValidation(format!(
460 "portable Methods PLS {label} contains duplicate sample identities"
461 )));
462 }
463 Ok(())
464 }
465}
466
467#[derive(Clone, Debug, PartialEq)]
469pub struct MethodsPlsDataRequest {
470 pub node_id: NodeId,
471 pub phase: Phase,
472 pub variant_id: Option<VariantId>,
473 pub fold_id: Option<FoldId>,
474 pub binding: DataBinding,
477 pub identity: Option<crate::training::TrainingDataIdentity>,
485 pub fit_view: DataProviderViewSpec,
486 pub prediction_view: Option<DataProviderViewSpec>,
487}
488
489impl MethodsPlsDataRequest {
490 pub fn validate(&self) -> Result<()> {
491 self.binding.validate()?;
492 match &self.identity {
493 Some(identity) => {
494 identity.validate()?;
495 if identity.requirement_key
496 != crate::data::data_binding_requirement_key(
497 &self.binding.node_id,
498 &self.binding.input_name,
499 )
500 {
501 return Err(DagMlError::RuntimeValidation(
502 "portable Methods PLS identity is not bound to its data binding"
503 .to_string(),
504 ));
505 }
506 }
507 None if self.phase != Phase::Predict => {
508 return Err(DagMlError::RuntimeValidation(
509 "portable Methods PLS FIT_CV/REFIT requires a target-bound training data identity"
510 .to_string(),
511 ));
512 }
513 None => {}
514 }
515 self.fit_view.validate()?;
516 if let Some(view) = &self.prediction_view {
517 view.validate()?;
518 }
519 Ok(())
520 }
521}
522
523#[derive(Clone, Debug, PartialEq)]
525pub struct MethodsPlsData {
526 pub fit: MethodsPlsDataset,
527 pub prediction: Option<MethodsPlsDataset>,
528}
529
530impl MethodsPlsData {
531 pub fn validate_for(&self, request: &MethodsPlsDataRequest) -> Result<()> {
532 request.validate()?;
533 self.fit.validate("fit", request.phase != Phase::Predict)?;
534 if let Some(expected_sample_ids) = &request.fit_view.sample_ids {
535 if self.fit.sample_ids != *expected_sample_ids {
536 return Err(DagMlError::RuntimeValidation(
537 "portable Methods PLS fit rows do not exactly match the scheduler-selected identity view".to_string(),
538 ));
539 }
540 } else if request.phase != Phase::Predict {
541 return Err(DagMlError::RuntimeValidation(
542 "portable Methods PLS fit view must carry scheduler-selected sample identities"
543 .to_string(),
544 ));
545 }
546 if let Some(prediction) = &self.prediction {
547 prediction.validate("prediction", request.phase == Phase::FitCv)?;
548 if prediction.sample_ids != request.prediction_view_sample_ids()? {
549 return Err(DagMlError::RuntimeValidation(
550 "portable Methods PLS prediction rows do not exactly match the scheduler-selected identity view".to_string(),
551 ));
552 }
553 if prediction.x.cols != self.fit.x.cols
554 || prediction.target_names != self.fit.target_names
555 || matches!((&prediction.y, &self.fit.y), (Some(left), Some(right)) if left.cols != right.cols)
556 {
557 return Err(DagMlError::RuntimeValidation(
558 "portable Methods PLS prediction schema differs from fit schema".to_string(),
559 ));
560 }
561 }
562 if request.prediction_view.is_some() != self.prediction.is_some() {
563 return Err(DagMlError::RuntimeValidation(
564 "portable Methods PLS provider did not return exactly the requested prediction view".to_string(),
565 ));
566 }
567 Ok(())
568 }
569}
570
571pub const METHODS_PLS_PREDICT_CONTENT_PROFILE: &str = "n4a-matrix-f64-le.v1";
580
581pub fn methods_pls_predict_feature_content_fingerprint(
588 matrix: &MethodsPlsMatrix,
589) -> Result<String> {
590 matrix.validate("PREDICT feature fingerprint")?;
591 let rows = u64::try_from(matrix.rows).map_err(|_| {
592 DagMlError::RuntimeValidation(
593 "portable Methods PLS PREDICT matrix row count does not fit the content identity profile"
594 .to_string(),
595 )
596 })?;
597 let cols = u64::try_from(matrix.cols).map_err(|_| {
598 DagMlError::RuntimeValidation(
599 "portable Methods PLS PREDICT matrix column count does not fit the content identity profile"
600 .to_string(),
601 )
602 })?;
603 let mut hasher = Sha256::new();
604 hasher.update(METHODS_PLS_PREDICT_CONTENT_PROFILE.as_bytes());
605 hasher.update([0]);
606 hasher.update(rows.to_le_bytes());
607 hasher.update(cols.to_le_bytes());
608 for value in &matrix.values {
609 hasher.update(value.to_bits().to_le_bytes());
610 }
611 Ok(format!("{:x}", hasher.finalize()))
612}
613
614#[derive(Clone, Debug, PartialEq)]
615pub struct MethodsPlsPredictInput {
616 pub data_content_profile: String,
618 pub data_content_fingerprint: String,
619 pub dataset: MethodsPlsDataset,
620}
621
622#[derive(Debug)]
630pub struct MethodsPlsPredictDataProvider {
631 inner: EnvelopeAttestedRuntimeDataProvider<crate::data::InMemoryDataProvider>,
632 inputs: BTreeMap<String, MethodsPlsPredictInput>,
633}
634
635impl MethodsPlsPredictDataProvider {
636 pub fn new<I>(
637 owner_controller: ControllerId,
638 bindings: I,
639 envelopes: BTreeMap<String, ExternalDataPlanEnvelope>,
640 inputs: BTreeMap<String, MethodsPlsPredictInput>,
641 ) -> Result<Self>
642 where
643 I: IntoIterator<Item = DataBinding>,
644 {
645 let bindings = bindings.into_iter().collect::<Vec<_>>();
646 let expected_keys = bindings
647 .iter()
648 .map(|binding| data_binding_requirement_key(&binding.node_id, &binding.input_name))
649 .collect::<BTreeSet<_>>();
650 let input_keys = inputs.keys().cloned().collect::<BTreeSet<_>>();
651 if input_keys.is_empty() || !input_keys.is_subset(&expected_keys) {
652 return Err(DagMlError::RuntimeValidation(format!(
653 "portable Methods PLS PREDICT inputs must name registered runtime bindings (unexpected: [{}])",
654 input_keys
655 .difference(&expected_keys)
656 .cloned()
657 .collect::<Vec<_>>()
658 .join(", "),
659 )));
660 }
661 for (key, input) in &inputs {
662 input.dataset.validate("predict input", false)?;
663 if input.data_content_profile != METHODS_PLS_PREDICT_CONTENT_PROFILE {
664 return Err(DagMlError::RuntimeValidation(format!(
665 "portable Methods PLS PREDICT input `{key}` has unsupported feature content profile `{}`",
666 input.data_content_profile,
667 )));
668 }
669 let actual_fingerprint =
670 methods_pls_predict_feature_content_fingerprint(&input.dataset.x)?;
671 if input.data_content_fingerprint != actual_fingerprint {
672 return Err(DagMlError::RuntimeValidation(format!(
673 "portable Methods PLS PREDICT input `{key}` feature content fingerprint does not match its row-major f64 values"
674 )));
675 }
676 if input.dataset.y.is_some() {
677 return Err(DagMlError::RuntimeValidation(format!(
678 "portable Methods PLS PREDICT input `{key}` must not carry targets"
679 )));
680 }
681 let envelope = envelopes.get(key).ok_or_else(|| {
682 DagMlError::RuntimeValidation(format!(
683 "portable Methods PLS PREDICT input `{key}` has no external envelope"
684 ))
685 })?;
686 envelope.validate()?;
687 let expected_fingerprint = envelope.data_content_fingerprint.as_deref().ok_or_else(|| {
688 DagMlError::RuntimeValidation(format!(
689 "portable Methods PLS PREDICT envelope `{key}` has no feature content fingerprint"
690 ))
691 })?;
692 if input.data_content_fingerprint != expected_fingerprint {
693 return Err(DagMlError::RuntimeValidation(format!(
694 "portable Methods PLS PREDICT input `{key}` feature content fingerprint does not match its envelope"
695 )));
696 }
697 if envelope.target_content_fingerprint.is_some() {
698 return Err(DagMlError::RuntimeValidation(format!(
699 "portable Methods PLS PREDICT input `{key}` requires a target-free envelope"
700 )));
701 }
702 }
703
704 let mut raw = crate::data::InMemoryDataProvider::new(owner_controller);
705 for envelope in envelopes.values().cloned() {
706 raw.register_envelope(envelope)?;
707 }
708 let inner = EnvelopeAttestedRuntimeDataProvider::new(raw, bindings, envelopes)?;
709 Ok(Self { inner, inputs })
710 }
711
712 fn input_for(&self, request: &MethodsPlsDataRequest) -> Result<&MethodsPlsPredictInput> {
713 request.validate()?;
714 if request.phase != Phase::Predict || request.identity.is_some() {
715 return Err(DagMlError::RuntimeValidation(
716 "portable Methods PLS target-free provider supports only PREDICT without a training identity"
717 .to_string(),
718 ));
719 }
720 if request.prediction_view.is_some() {
721 return Err(DagMlError::RuntimeValidation(
722 "portable Methods PLS target-free provider does not support FIT_CV validation views"
723 .to_string(),
724 ));
725 }
726 let key =
727 data_binding_requirement_key(&request.binding.node_id, &request.binding.input_name);
728 let input = self.inputs.get(&key).ok_or_else(|| {
729 DagMlError::RuntimeValidation(format!(
730 "portable Methods PLS target-free provider has no input for `{key}`"
731 ))
732 })?;
733 if let Some(expected_sample_ids) = &request.fit_view.sample_ids {
734 if input.dataset.sample_ids != *expected_sample_ids {
735 return Err(DagMlError::RuntimeValidation(
736 "portable Methods PLS target-free input rows do not match the scheduler-selected identity view"
737 .to_string(),
738 ));
739 }
740 }
741 Ok(input)
742 }
743}
744
745impl RuntimeDataProvider for MethodsPlsPredictDataProvider {
746 fn materialize(&self, request: &DataMaterializationRequest) -> Result<HandleRef> {
747 self.inner.materialize(request)
748 }
749
750 fn make_view(&self, request: &DataViewRequest) -> Result<HandleRef> {
751 self.inner.make_view(request)
752 }
753
754 fn training_data_identity(
755 &self,
756 binding: &DataBinding,
757 ) -> Result<Option<crate::training::TrainingDataIdentity>> {
758 self.inner.training_data_identity(binding)
759 }
760
761 fn coordinator_relations(&self, binding: &DataBinding) -> Result<Option<SampleRelationSet>> {
762 self.inner.coordinator_relations(binding)
763 }
764
765 fn predict_cohort(
766 &self,
767 binding: &DataBinding,
768 phase: Phase,
769 ) -> Result<Option<crate::data::PredictCohort>> {
770 self.inner.predict_cohort(binding, phase)
771 }
772
773 fn methods_pls_capability(&self) -> Result<()> {
774 Ok(())
775 }
776
777 fn preflight_methods_pls(&self, request: &MethodsPlsDataRequest) -> Result<()> {
778 self.input_for(request)?;
779 Ok(())
780 }
781
782 fn methods_pls_data(&self, request: &MethodsPlsDataRequest) -> Result<MethodsPlsData> {
783 let input = self.input_for(request)?;
784 Ok(MethodsPlsData {
785 fit: input.dataset.clone(),
786 prediction: None,
787 })
788 }
789}
790
791impl MethodsPlsDataRequest {
792 fn prediction_view_sample_ids(&self) -> Result<Vec<SampleId>> {
793 self.prediction_view
794 .as_ref()
795 .and_then(|view| view.sample_ids.clone())
796 .ok_or_else(|| {
797 DagMlError::RuntimeValidation(
798 "portable Methods PLS prediction view must carry scheduler-selected sample identities".to_string(),
799 )
800 })
801 }
802}
803
804#[derive(Debug)]
805struct EnvelopeAttestation {
806 binding: DataBinding,
807 envelope: ExternalDataPlanEnvelope,
808 identity: Option<crate::training::TrainingDataIdentity>,
812}
813
814#[derive(Debug)]
823pub struct EnvelopeAttestedRuntimeDataProvider<P> {
824 inner: P,
825 attestations: BTreeMap<String, EnvelopeAttestation>,
826}
827
828impl<P> EnvelopeAttestedRuntimeDataProvider<P> {
829 pub fn new<I>(
830 inner: P,
831 bindings: I,
832 mut envelopes: BTreeMap<String, ExternalDataPlanEnvelope>,
833 ) -> Result<Self>
834 where
835 I: IntoIterator<Item = DataBinding>,
836 {
837 let mut bindings_by_key: BTreeMap<String, DataBinding> = BTreeMap::new();
838 for binding in bindings {
839 binding.validate()?;
840 let key = data_binding_requirement_key(&binding.node_id, &binding.input_name);
841 if let Some(previous) = bindings_by_key.get(&key) {
842 let detail = if previous.node_id == binding.node_id
843 && previous.input_name == binding.input_name
844 {
845 "duplicates the same coordinates"
846 } else {
847 "uses distinct coordinates that collide under the V1 node.input spelling"
848 };
849 return Err(DagMlError::RuntimeValidation(format!(
850 "data binding requirement key `{key}` {detail}"
851 )));
852 }
853 bindings_by_key.insert(key, binding);
854 }
855
856 let expected_keys = bindings_by_key.keys().cloned().collect::<BTreeSet<_>>();
857 let actual_keys = envelopes.keys().cloned().collect::<BTreeSet<_>>();
858 if expected_keys != actual_keys {
859 let missing = expected_keys
860 .difference(&actual_keys)
861 .cloned()
862 .collect::<Vec<_>>();
863 let unexpected = actual_keys
864 .difference(&expected_keys)
865 .cloned()
866 .collect::<Vec<_>>();
867 return Err(DagMlError::RuntimeValidation(format!(
868 "attested data envelopes must exactly cover runtime bindings (missing: [{}]; unexpected: [{}])",
869 missing.join(", "),
870 unexpected.join(", ")
871 )));
872 }
873
874 let mut attestations = BTreeMap::new();
875 for (key, binding) in bindings_by_key {
876 let envelope = envelopes
877 .remove(&key)
878 .expect("exact key coverage was checked above");
879 let identity = if envelope.relation_fingerprint.is_some()
880 && envelope.data_content_fingerprint.is_some()
881 && envelope.target_content_fingerprint.is_some()
882 {
883 Some(
884 crate::training::TrainingDataIdentity::from_binding_envelope(
885 &binding, &envelope,
886 )?,
887 )
888 } else {
889 None
890 };
891 attestations.insert(
892 key,
893 EnvelopeAttestation {
894 binding,
895 envelope,
896 identity,
897 },
898 );
899 }
900
901 Ok(Self {
902 inner,
903 attestations,
904 })
905 }
906
907 pub fn inner(&self) -> &P {
908 &self.inner
909 }
910
911 pub fn into_inner(self) -> P {
912 self.inner
913 }
914
915 fn attestation_for_binding(&self, binding: &DataBinding) -> Result<&EnvelopeAttestation> {
916 binding.validate()?;
917 let key = data_binding_requirement_key(&binding.node_id, &binding.input_name);
918 let attestation = self.attestations.get(&key).ok_or_else(|| {
919 DagMlError::RuntimeValidation(format!(
920 "runtime data binding `{key}` has no registered envelope attestation"
921 ))
922 })?;
923 if attestation.binding != *binding {
924 return Err(DagMlError::RuntimeValidation(format!(
925 "runtime data binding `{key}` does not exactly match its attested binding"
926 )));
927 }
928 Ok(attestation)
929 }
930
931 fn validate_request_binding(
932 &self,
933 node_id: &NodeId,
934 input_name: &str,
935 binding: &DataBinding,
936 ) -> Result<()> {
937 if node_id != &binding.node_id || input_name != binding.input_name {
938 return Err(DagMlError::RuntimeValidation(format!(
939 "runtime data request coordinates `{node_id}.{input_name}` do not match binding `{}`",
940 data_binding_requirement_key(&binding.node_id, &binding.input_name)
941 )));
942 }
943 self.attestation_for_binding(binding)?;
944 Ok(())
945 }
946
947 fn validate_predict_cohort_request(
948 &self,
949 binding: &DataBinding,
950 phase: Phase,
951 supplied: &Option<crate::data::PredictCohort>,
952 ) -> Result<()> {
953 let attestation = self.attestation_for_binding(binding)?;
954 match phase {
955 Phase::Predict => {
956 if let Some(cohort) = supplied {
957 cohort.validate()?;
958 }
959 if supplied != &attestation.envelope.predict_cohort {
960 return Err(DagMlError::RuntimeValidation(format!(
961 "PREDICT cohort for runtime binding `{}` does not exactly match its envelope attestation",
962 data_binding_requirement_key(&binding.node_id, &binding.input_name)
963 )));
964 }
965 }
966 _ if supplied.is_some() => {
967 return Err(DagMlError::RuntimeValidation(format!(
968 "runtime binding `{}` carries a PREDICT cohort during non-PREDICT phase {phase:?}",
969 data_binding_requirement_key(&binding.node_id, &binding.input_name)
970 )));
971 }
972 _ => {}
973 }
974 Ok(())
975 }
976}
977
978impl<P: RuntimeDataProvider> RuntimeDataProvider for EnvelopeAttestedRuntimeDataProvider<P> {
979 fn materialize(&self, request: &DataMaterializationRequest) -> Result<HandleRef> {
980 self.validate_request_binding(&request.node_id, &request.input_name, &request.binding)?;
981 self.validate_predict_cohort_request(
982 &request.binding,
983 request.phase,
984 &request.predict_cohort,
985 )?;
986 self.inner.materialize(request)
987 }
988
989 fn make_view(&self, request: &DataViewRequest) -> Result<HandleRef> {
990 request.view.validate()?;
991 self.validate_request_binding(&request.node_id, &request.input_name, &request.binding)?;
992 self.validate_predict_cohort_request(
993 &request.binding,
994 request.phase,
995 &request.predict_cohort,
996 )?;
997 self.inner.make_view(request)
998 }
999
1000 fn training_data_identity(
1001 &self,
1002 binding: &DataBinding,
1003 ) -> Result<Option<crate::training::TrainingDataIdentity>> {
1004 Ok(self.attestation_for_binding(binding)?.identity.clone())
1005 }
1006
1007 fn coordinator_relations(&self, binding: &DataBinding) -> Result<Option<SampleRelationSet>> {
1008 Ok(self
1009 .attestation_for_binding(binding)?
1010 .envelope
1011 .coordinator_relations
1012 .clone())
1013 }
1014
1015 fn predict_cohort(
1016 &self,
1017 binding: &DataBinding,
1018 phase: Phase,
1019 ) -> Result<Option<crate::data::PredictCohort>> {
1020 validate_predict_cohort_phase(phase)?;
1021 Ok(self
1022 .attestation_for_binding(binding)?
1023 .envelope
1024 .predict_cohort
1025 .clone())
1026 }
1027
1028 fn methods_pls_capability(&self) -> Result<()> {
1029 self.inner.methods_pls_capability()
1030 }
1031
1032 fn preflight_methods_pls(&self, request: &MethodsPlsDataRequest) -> Result<()> {
1033 request.validate()?;
1034 self.inner.preflight_methods_pls(request)
1035 }
1036
1037 fn methods_pls_data(&self, request: &MethodsPlsDataRequest) -> Result<MethodsPlsData> {
1038 request.validate()?;
1039 self.inner.methods_pls_data(request)
1040 }
1041}
1042
1043pub trait RuntimeController: Send + Sync {
1044 fn controller_id(&self) -> &ControllerId;
1045 fn invoke(&self, task: &NodeTask) -> Result<NodeResult>;
1046
1047 fn export_artifact_payload(&self, _artifact_id: &ArtifactId) -> Result<Option<Vec<u8>>> {
1051 Ok(None)
1052 }
1053
1054 fn hydrate_artifact_payload(
1060 &self,
1061 _request: &ArtifactMaterializationRequest,
1062 _payload: &[u8],
1063 ) -> Result<HandleRef> {
1064 Err(DagMlError::RuntimeValidation(format!(
1065 "runtime controller `{}` cannot hydrate a raw portable artifact payload",
1066 self.controller_id()
1067 )))
1068 }
1069
1070 fn release_hydrated_artifact_payload(&self, _handle: &HandleRef) -> Result<()> {
1075 Err(DagMlError::RuntimeValidation(format!(
1076 "runtime controller `{}` cannot release a hydrated raw portable artifact payload",
1077 self.controller_id()
1078 )))
1079 }
1080
1081 fn invoke_with_data_provider(
1085 &self,
1086 task: &NodeTask,
1087 _data_provider: &dyn RuntimeDataProvider,
1088 ) -> Result<NodeResult> {
1089 self.invoke(task)
1090 }
1091
1092 fn create_tuner_session(
1099 &self,
1100 task: &RuntimeHpoCampaignTask,
1101 _context: &RuntimeHpoExecutionContext,
1102 ) -> Result<Box<dyn RuntimeTunerSession>> {
1103 Err(DagMlError::RuntimeValidation(format!(
1104 "runtime controller `{}` does not implement an execution-local tuner session for HPO campaign `{}`",
1105 self.controller_id(), task.operation_id
1106 )))
1107 }
1108
1109 fn invoke_aggregation(
1110 &self,
1111 task: &AggregationControllerTask,
1112 ) -> Result<AggregationControllerResult> {
1113 Err(DagMlError::RuntimeValidation(format!(
1114 "runtime controller `{}` does not implement aggregation task `{}`",
1115 self.controller_id(),
1116 task.task_id
1117 )))
1118 }
1119}
1120
1121pub trait RuntimeTunerSession {
1128 fn trial_history_len(&self) -> Result<u32>;
1132
1133 fn ask(&mut self) -> Result<Option<RuntimeHpoProposal>>;
1134
1135 fn report_intermediate(
1136 &mut self,
1137 intermediate: RuntimeHpoIntermediate,
1138 ) -> Result<RuntimeHpoIntermediateOutcome>;
1139
1140 fn tell(&mut self, trial_id: i64, terminal: RuntimeHpoTerminal) -> Result<()>;
1141
1142 fn incumbent(&self, variants: &BTreeMap<i64, VariantId>)
1146 -> Result<Option<RuntimeHpoIncumbent>>;
1147
1148 fn terminal_trial_snapshots(
1152 &self,
1153 variants: &BTreeMap<i64, VariantId>,
1154 ) -> Result<Vec<RuntimeHpoTerminalSnapshot>>;
1155
1156 fn checkpoint(&self) -> Result<crate::hpo::N4moptCheckpointArtifact>;
1160}
1161pub(crate) struct CollectedInputs {
1162 pub(crate) handles: BTreeMap<String, HandleRef>,
1163 pub(crate) data_views: BTreeMap<String, DataProviderViewSpec>,
1164 pub(crate) prediction_inputs: BTreeMap<String, PredictionInputSpec>,
1165 pub(crate) skip_node: bool,
1166}
1167
1168pub(crate) fn data_view_key(input_name: &str) -> String {
1169 format!("data:{input_name}")
1170}
1171
1172pub(crate) fn validation_data_view_key(input_name: &str) -> String {
1173 format!("{input_name}:validation")
1174}
1175
1176pub(crate) fn derive_output_data_views(
1177 plan: &ExecutionPlan,
1178 task: &NodeTask,
1179 result: &NodeResult,
1180) -> Result<BTreeMap<String, DataProviderViewSpec>> {
1181 let node = plan
1182 .graph_plan
1183 .graph
1184 .nodes
1185 .iter()
1186 .find(|node| node.id == task.node_plan.node_id)
1187 .expect("execution plan was validated");
1188 let mut views = BTreeMap::new();
1189 for port in node
1190 .ports
1191 .outputs
1192 .iter()
1193 .filter(|port| port.kind == PortKind::Data)
1194 {
1195 let Some(handle) = result.outputs.get(&port.name) else {
1196 continue;
1197 };
1198 if !matches!(handle.kind, HandleKind::Data | HandleKind::DataView) {
1199 return Err(DagMlError::RuntimeValidation(format!(
1200 "node `{}` emitted data output `{}` with non-data/data-view handle kind {:?}",
1201 task.node_plan.node_id, port.name, handle.kind
1202 )));
1203 }
1204 if let Some(view) = primary_output_data_view(task) {
1205 views.insert(
1206 port.name.clone(),
1207 output_data_view_for_port(task, result, &port.name, view)?,
1208 );
1209 }
1210 if let Some(validation_view) = validation_output_data_view(task) {
1211 views.insert(
1212 validation_data_view_key(&port.name),
1213 output_data_view_for_port(task, result, &port.name, validation_view)?,
1214 );
1215 }
1216 }
1217 Ok(views)
1218}
1219
1220pub(crate) fn output_data_view_for_port(
1221 task: &NodeTask,
1222 result: &NodeResult,
1223 port_name: &str,
1224 base_view: &DataProviderViewSpec,
1225) -> Result<DataProviderViewSpec> {
1226 let mut view = base_view.clone();
1227 if let Some(upstream_provenance) = view.extra.remove(DATA_OUTPUT_PROVENANCE_KEY) {
1228 let provenance: DataOutputProvenance =
1229 serde_json::from_value(upstream_provenance).map_err(|error| {
1230 DagMlError::RuntimeValidation(format!(
1231 "node `{}` cannot propagate data output `{port_name}` because upstream data output provenance is invalid JSON: {error}",
1232 task.node_plan.node_id
1233 ))
1234 })?;
1235 provenance.validate().map_err(|error| {
1236 DagMlError::RuntimeValidation(format!(
1237 "node `{}` cannot propagate data output `{port_name}` because upstream data output provenance is invalid: {error}",
1238 task.node_plan.node_id
1239 ))
1240 })?;
1241 }
1242 let shape_deltas = result
1243 .shape_deltas
1244 .iter()
1245 .filter(|delta| delta.node_id == task.node_plan.node_id)
1246 .cloned()
1247 .collect::<Vec<_>>();
1248 let mut provenance = DataOutputProvenance {
1249 schema_version: DATA_OUTPUT_PROVENANCE_SCHEMA_VERSION,
1250 producer_node: task.node_plan.node_id.clone(),
1251 producer_port: port_name.to_string(),
1252 producer_phase: task.phase,
1253 variant_id: task.variant_id.clone(),
1254 fold_id: task.fold_id.clone(),
1255 shape_plan_fingerprint: None,
1256 aggregation_policy_fingerprint: None,
1257 feature_namespace: None,
1258 feature_schema_fingerprint: None,
1259 representation_plan: None,
1260 representation_replay_manifest: None,
1261 representation_compatibility: None,
1262 relation_delta_fingerprint: None,
1263 shape_deltas,
1264 };
1265 if let Some(shape_plan) = &task.node_plan.shape_plan {
1266 provenance.shape_plan_fingerprint = Some(stable_json_fingerprint(shape_plan)?);
1267 provenance.aggregation_policy_fingerprint =
1268 Some(stable_json_fingerprint(&shape_plan.aggregation_policy)?);
1269 provenance.feature_namespace = shape_plan.feature_namespace.clone();
1270 provenance.feature_schema_fingerprint =
1271 output_feature_schema_fingerprint(shape_plan, result);
1272 }
1273 provenance.validate()?;
1274
1275 view.extra.insert(
1276 DATA_OUTPUT_PROVENANCE_KEY.to_string(),
1277 serde_json::to_value(provenance)?,
1278 );
1279 view.validate()?;
1280 Ok(view)
1281}
1282
1283pub(crate) fn output_feature_schema_fingerprint(
1284 shape_plan: &crate::policy::DataModelShapePlan,
1285 result: &NodeResult,
1286) -> Option<String> {
1287 result
1288 .shape_deltas
1289 .iter()
1290 .rev()
1291 .find(|delta| delta.kind == ShapeDeltaKind::Feature)
1292 .map(|delta| delta.after_fingerprint.clone())
1293 .or_else(|| shape_plan.feature_schema_fingerprint.clone())
1294}
1295
1296pub(crate) fn primary_output_data_view(task: &NodeTask) -> Option<&DataProviderViewSpec> {
1297 task.data_views
1298 .values()
1299 .find(|view| view.partition != DataRequestPartition::FoldValidation)
1300 .or_else(|| task.data_views.values().next())
1301}
1302
1303pub(crate) fn validation_output_data_view(task: &NodeTask) -> Option<&DataProviderViewSpec> {
1304 task.data_views
1305 .values()
1306 .find(|view| view.partition == DataRequestPartition::FoldValidation)
1307}
1308
1309pub(crate) struct DataViewHandleInput<'a> {
1311 pub(crate) data_handle: &'a HandleRef,
1312 pub(crate) view: &'a DataProviderViewSpec,
1313 pub(crate) predict_cohort: Option<&'a crate::data::PredictCohort>,
1314}
1315
1316pub(crate) fn make_data_view_handle(
1317 data_provider: &dyn RuntimeDataProvider,
1318 ctx: &RunContext,
1319 node_plan: &NodePlan,
1320 scope: &PhaseScope,
1321 binding: &DataBinding,
1322 input: DataViewHandleInput<'_>,
1323) -> Result<HandleRef> {
1324 input.view.validate()?;
1325 let view_handle = data_provider.make_view(&DataViewRequest {
1326 run_id: ctx.run_id.clone(),
1327 node_id: node_plan.node_id.clone(),
1328 input_name: binding.input_name.clone(),
1329 phase: scope.phase,
1330 variant_id: scope.variant_id.clone(),
1331 fold_id: scope.fold_id.clone(),
1332 binding: binding.clone(),
1333 data_handle: input.data_handle.clone(),
1334 view: input.view.clone(),
1335 predict_cohort: input.predict_cohort.cloned(),
1336 })?;
1337 if !matches!(view_handle.kind, HandleKind::Data | HandleKind::DataView) {
1341 return Err(DagMlError::RuntimeValidation(format!(
1342 "node `{}` data view `{}` resolved to a non-data/data-view handle kind {:?}",
1343 node_plan.node_id, binding.input_name, view_handle.kind
1344 )));
1345 }
1346 Ok(view_handle)
1347}
1348
1349pub(crate) fn data_view_for_scope(
1350 binding: &DataBinding,
1351 fold_set: Option<&FoldSet>,
1352 scope: &PhaseScope,
1353 branch_view: Option<&crate::data::BranchViewPlan>,
1354 excluded_samples: &BTreeSet<SampleId>,
1355) -> Result<DataProviderViewSpec> {
1356 let partition = data_partition_for_scope(binding, scope);
1357 let role = match scope.phase {
1360 Phase::FitCv | Phase::Refit => DataViewRole::Fit,
1361 _ => DataViewRole::NonFit,
1362 };
1363 data_view_for_partition(
1364 binding,
1365 fold_set,
1366 scope,
1367 partition,
1368 branch_view,
1369 role,
1370 excluded_samples,
1371 )
1372}
1373
1374pub(crate) fn bind_predict_cohort_to_view(
1382 view: &mut DataProviderViewSpec,
1383 cohort: &crate::data::PredictCohort,
1384) -> Result<()> {
1385 cohort.validate()?;
1386 if view.partition != DataRequestPartition::Predict || view.fold_id.is_some() {
1387 return Err(DagMlError::RuntimeValidation(
1388 "PREDICT cohort may only bind a top-level Predict data view".to_string(),
1389 ));
1390 }
1391 view.sample_ids = Some(cohort.physical_sample_ids.clone());
1392 view.validate()
1393}
1394
1395pub(crate) fn validation_data_view_for_scope(
1396 binding: &DataBinding,
1397 fold_set: Option<&FoldSet>,
1398 scope: &PhaseScope,
1399 branch_view: Option<&crate::data::BranchViewPlan>,
1400 excluded_samples: &BTreeSet<SampleId>,
1401) -> Result<Option<DataProviderViewSpec>> {
1402 if scope.phase != Phase::FitCv || scope.fold_id.is_none() {
1403 return Ok(None);
1404 }
1405 let partition = binding.view_policy.predict_partition;
1406 if partition == data_partition_for_scope(binding, scope) {
1407 return Ok(None);
1408 }
1409 data_view_for_partition(
1411 binding,
1412 fold_set,
1413 scope,
1414 partition,
1415 branch_view,
1416 DataViewRole::NonFit,
1417 excluded_samples,
1418 )
1419 .map(Some)
1420}
1421
1422#[cfg(test)]
1423mod envelope_attested_provider_tests {
1424 use std::cell::Cell;
1425
1426 use super::*;
1427
1428 #[derive(Debug, Default)]
1429 struct ProbeProvider {
1430 materialize_calls: Cell<usize>,
1431 make_view_calls: Cell<usize>,
1432 }
1433
1434 impl RuntimeDataProvider for ProbeProvider {
1435 fn materialize(&self, _request: &DataMaterializationRequest) -> Result<HandleRef> {
1436 self.materialize_calls.set(self.materialize_calls.get() + 1);
1437 Ok(HandleRef {
1438 handle: 41,
1439 kind: HandleKind::Data,
1440 owner_controller: ControllerId::new("controller:data.probe").unwrap(),
1441 })
1442 }
1443
1444 fn make_view(&self, _request: &DataViewRequest) -> Result<HandleRef> {
1445 self.make_view_calls.set(self.make_view_calls.get() + 1);
1446 Ok(HandleRef {
1447 handle: 42,
1448 kind: HandleKind::DataView,
1449 owner_controller: ControllerId::new("controller:data.probe").unwrap(),
1450 })
1451 }
1452 }
1453
1454 fn complete_envelope() -> ExternalDataPlanEnvelope {
1455 let mut envelope: ExternalDataPlanEnvelope = serde_json::from_str(include_str!(
1456 "../../tests/fixtures/package/data/coordinator_data_plan_envelope_sample12.json"
1457 ))
1458 .unwrap();
1459 envelope.data_content_fingerprint = Some("a".repeat(64));
1460 envelope.target_content_fingerprint = Some("b".repeat(64));
1461 envelope
1462 }
1463
1464 fn inference_predict_cohort(envelope: &ExternalDataPlanEnvelope) -> crate::data::PredictCohort {
1465 let relations = envelope
1466 .coordinator_relations
1467 .clone()
1468 .expect("complete test envelope carries coordinator relations");
1469 let physical_sample_ids = relations
1470 .records
1471 .iter()
1472 .map(|record| record.sample_id.clone())
1473 .collect::<BTreeSet<_>>()
1474 .into_iter()
1475 .collect::<Vec<_>>();
1476 let origin_sample_ids = relations
1477 .records
1478 .iter()
1479 .map(|record| {
1480 record
1481 .origin_sample_id
1482 .clone()
1483 .unwrap_or_else(|| record.sample_id.clone())
1484 })
1485 .collect::<BTreeSet<_>>()
1486 .into_iter()
1487 .collect::<Vec<_>>();
1488 let mut cohort = crate::data::PredictCohort {
1489 role: crate::data::PredictCohortRole::Inference,
1490 physical_sample_ids,
1491 origin_sample_ids,
1492 target_names: vec!["y".to_string()],
1493 relation_fingerprint: relations.fingerprint().unwrap(),
1494 relations,
1495 data_content_fingerprint: "c".repeat(64),
1496 target_content_fingerprint: None,
1497 cohort_fingerprint: String::new(),
1498 };
1499 cohort.cohort_fingerprint = cohort.fingerprint().unwrap();
1500 cohort
1501 }
1502
1503 fn binding_for(
1504 node_id: &str,
1505 input_name: &str,
1506 envelope: &ExternalDataPlanEnvelope,
1507 ) -> DataBinding {
1508 DataBinding {
1509 node_id: NodeId::new(node_id).unwrap(),
1510 input_name: input_name.to_string(),
1511 request_id: "request:data.probe".to_string(),
1512 schema_fingerprint: envelope.schema_fingerprint.clone(),
1513 plan_fingerprint: envelope.plan_fingerprint.clone(),
1514 relation_fingerprint: envelope.relation_fingerprint.clone(),
1515 output_representation: "tabular_numeric".to_string(),
1516 feature_set_id: Some(input_name.to_string()),
1517 source_ids: vec!["source:probe".to_string()],
1518 require_relations: true,
1519 view_policy: Default::default(),
1520 metadata: BTreeMap::new(),
1521 }
1522 }
1523
1524 fn envelopes_for(
1525 binding: &DataBinding,
1526 envelope: ExternalDataPlanEnvelope,
1527 ) -> BTreeMap<String, ExternalDataPlanEnvelope> {
1528 BTreeMap::from([(
1529 data_binding_requirement_key(&binding.node_id, &binding.input_name),
1530 envelope,
1531 )])
1532 }
1533
1534 fn materialization_request(binding: &DataBinding) -> DataMaterializationRequest {
1535 DataMaterializationRequest {
1536 run_id: RunId::new("run:attested.provider").unwrap(),
1537 node_id: binding.node_id.clone(),
1538 input_name: binding.input_name.clone(),
1539 phase: Phase::Refit,
1540 variant_id: None,
1541 fold_id: None,
1542 binding: binding.clone(),
1543 predict_cohort: None,
1544 }
1545 }
1546
1547 #[test]
1548 fn envelope_attested_provider_delegates_and_returns_exact_attestations() {
1549 let envelope = complete_envelope();
1550 let binding = binding_for("model:base", "x", &envelope);
1551 let expected_identity =
1552 crate::training::TrainingDataIdentity::from_binding_envelope(&binding, &envelope)
1553 .unwrap();
1554 let expected_relations = envelope.coordinator_relations.clone();
1555 let provider = EnvelopeAttestedRuntimeDataProvider::new(
1556 ProbeProvider::default(),
1557 vec![binding.clone()],
1558 envelopes_for(&binding, envelope),
1559 )
1560 .unwrap();
1561
1562 assert_eq!(
1563 provider.training_data_identity(&binding).unwrap(),
1564 Some(expected_identity)
1565 );
1566 assert_eq!(
1567 provider.coordinator_relations(&binding).unwrap(),
1568 expected_relations
1569 );
1570
1571 let materialization = materialization_request(&binding);
1572 let data_handle = provider.materialize(&materialization).unwrap();
1573 assert_eq!(data_handle.handle, 41);
1574 let view_handle = provider
1575 .make_view(&DataViewRequest {
1576 run_id: materialization.run_id,
1577 node_id: binding.node_id.clone(),
1578 input_name: binding.input_name.clone(),
1579 phase: Phase::Refit,
1580 variant_id: None,
1581 fold_id: None,
1582 binding: binding.clone(),
1583 data_handle,
1584 view: DataProviderViewSpec {
1585 sample_ids: None,
1586 partition: DataRequestPartition::FullTrain,
1587 fold_id: None,
1588 source_ids: None,
1589 columns: None,
1590 include_augmented: true,
1591 include_excluded: false,
1592 branch_view: None,
1593 extra: BTreeMap::new(),
1594 },
1595 predict_cohort: None,
1596 })
1597 .unwrap();
1598 assert_eq!(view_handle.handle, 42);
1599 assert_eq!(provider.inner().materialize_calls.get(), 1);
1600 assert_eq!(provider.inner().make_view_calls.get(), 1);
1601
1602 let inner = provider.into_inner();
1603 assert_eq!(inner.materialize_calls.get(), 1);
1604 assert_eq!(inner.make_view_calls.get(), 1);
1605 }
1606
1607 #[test]
1608 fn envelope_attested_provider_refuses_substituted_or_non_predict_cohorts() {
1609 let mut envelope = complete_envelope();
1610 envelope.schema_version = crate::data::EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION_V2;
1611 let expected = inference_predict_cohort(&envelope);
1612 envelope.predict_cohort = Some(expected.clone());
1613 envelope.validate().unwrap();
1614 let binding = binding_for("model:base", "x", &envelope);
1615 let provider = EnvelopeAttestedRuntimeDataProvider::new(
1616 ProbeProvider::default(),
1617 vec![binding.clone()],
1618 envelopes_for(&binding, envelope),
1619 )
1620 .unwrap();
1621
1622 let mut request = materialization_request(&binding);
1623 request.phase = Phase::Predict;
1624 request.predict_cohort = Some(expected.clone());
1625 provider.materialize(&request).unwrap();
1626 assert_eq!(provider.inner().materialize_calls.get(), 1);
1627
1628 let mut substituted = expected.clone();
1629 substituted.data_content_fingerprint = "d".repeat(64);
1630 substituted.cohort_fingerprint = substituted.fingerprint().unwrap();
1631 request.predict_cohort = Some(substituted);
1632 let error = provider.materialize(&request).unwrap_err().to_string();
1633 assert!(error.contains("does not exactly match its envelope attestation"));
1634 assert_eq!(provider.inner().materialize_calls.get(), 1);
1635
1636 request.phase = Phase::Refit;
1637 request.predict_cohort = Some(expected);
1638 let error = provider.materialize(&request).unwrap_err().to_string();
1639 assert!(error.contains("during non-PREDICT phase"));
1640 assert_eq!(provider.inner().materialize_calls.get(), 1);
1641 }
1642
1643 #[test]
1644 fn envelope_attested_provider_preserves_target_free_predict_envelopes() {
1645 let mut envelope = complete_envelope();
1646 envelope.target_content_fingerprint = None;
1647 let binding = binding_for("model:base", "x", &envelope);
1648 let provider = EnvelopeAttestedRuntimeDataProvider::new(
1649 ProbeProvider::default(),
1650 vec![binding.clone()],
1651 envelopes_for(&binding, envelope.clone()),
1652 )
1653 .unwrap();
1654
1655 assert_eq!(provider.training_data_identity(&binding).unwrap(), None);
1659 assert_eq!(
1660 provider.coordinator_relations(&binding).unwrap(),
1661 envelope.coordinator_relations
1662 );
1663 let mut request = materialization_request(&binding);
1664 request.phase = Phase::Predict;
1665 assert_eq!(provider.materialize(&request).unwrap().handle, 41);
1666 }
1667
1668 #[test]
1669 fn envelope_attested_provider_requires_exact_envelope_coverage() {
1670 let envelope = complete_envelope();
1671 let binding = binding_for("model:base", "x", &envelope);
1672
1673 let missing = EnvelopeAttestedRuntimeDataProvider::new(
1674 ProbeProvider::default(),
1675 vec![binding.clone()],
1676 BTreeMap::new(),
1677 )
1678 .unwrap_err();
1679 assert!(missing.to_string().contains("exactly cover"));
1680 assert!(missing.to_string().contains("model:base.x"));
1681
1682 let mut unexpected = envelopes_for(&binding, envelope.clone());
1683 unexpected.insert("model:other.x".to_string(), envelope);
1684 let extra = EnvelopeAttestedRuntimeDataProvider::new(
1685 ProbeProvider::default(),
1686 vec![binding],
1687 unexpected,
1688 )
1689 .unwrap_err();
1690 assert!(extra.to_string().contains("exactly cover"));
1691 assert!(extra.to_string().contains("model:other.x"));
1692 }
1693
1694 #[test]
1695 fn envelope_attested_provider_rejects_rendered_key_collisions() {
1696 let envelope = complete_envelope();
1697 let left = binding_for("a.b", "c", &envelope);
1698 let right = binding_for("a", "b.c", &envelope);
1699 assert_eq!(
1700 data_binding_requirement_key(&left.node_id, &left.input_name),
1701 data_binding_requirement_key(&right.node_id, &right.input_name)
1702 );
1703
1704 let error = EnvelopeAttestedRuntimeDataProvider::new(
1705 ProbeProvider::default(),
1706 vec![left.clone(), right],
1707 envelopes_for(&left, envelope),
1708 )
1709 .unwrap_err();
1710 assert!(error.to_string().contains("distinct coordinates"));
1711 assert!(error.to_string().contains("a.b.c"));
1712 }
1713
1714 #[test]
1715 fn envelope_attested_provider_refuses_unattested_binding_before_delegation() {
1716 let envelope = complete_envelope();
1717 let binding = binding_for("model:base", "x", &envelope);
1718 let provider = EnvelopeAttestedRuntimeDataProvider::new(
1719 ProbeProvider::default(),
1720 vec![binding.clone()],
1721 envelopes_for(&binding, envelope),
1722 )
1723 .unwrap();
1724 let mut changed = binding;
1725 changed.request_id = "request:data.changed".to_string();
1726
1727 let error = provider
1728 .materialize(&materialization_request(&changed))
1729 .unwrap_err();
1730 assert!(error.to_string().contains("does not exactly match"));
1731 assert_eq!(provider.inner().materialize_calls.get(), 0);
1732 }
1733
1734 #[test]
1735 fn envelope_attested_provider_marks_incomplete_envelope_as_non_training() {
1736 let mut envelope = complete_envelope();
1737 envelope.data_content_fingerprint = None;
1738 let binding = binding_for("model:base", "x", &envelope);
1739 let provider = EnvelopeAttestedRuntimeDataProvider::new(
1740 ProbeProvider::default(),
1741 vec![binding.clone()],
1742 envelopes_for(&binding, envelope),
1743 )
1744 .unwrap();
1745 assert_eq!(provider.training_data_identity(&binding).unwrap(), None);
1746 }
1747
1748 #[test]
1749 fn methods_pls_request_allows_target_free_predict_but_not_training() {
1750 let envelope = complete_envelope();
1751 let binding = binding_for("model:base", "x", &envelope);
1752 let predict_view = DataProviderViewSpec {
1753 sample_ids: Some(vec![SampleId::new("sample:1").unwrap()]),
1754 partition: DataRequestPartition::Predict,
1755 fold_id: None,
1756 source_ids: None,
1757 columns: None,
1758 include_augmented: false,
1759 include_excluded: false,
1760 branch_view: None,
1761 extra: BTreeMap::new(),
1762 };
1763 let request = MethodsPlsDataRequest {
1764 node_id: binding.node_id.clone(),
1765 phase: Phase::Predict,
1766 variant_id: None,
1767 fold_id: None,
1768 binding: binding.clone(),
1769 identity: None,
1770 fit_view: predict_view.clone(),
1771 prediction_view: None,
1772 };
1773 request.validate().unwrap();
1774
1775 let mut refit = request;
1776 refit.phase = Phase::Refit;
1777 refit.fit_view.partition = DataRequestPartition::FullTrain;
1778 let error = refit.validate().unwrap_err();
1779 assert!(error
1780 .to_string()
1781 .contains("FIT_CV/REFIT requires a target-bound training data identity"));
1782 }
1783
1784 #[test]
1785 fn methods_pls_predict_provider_binds_x_only_rows_to_the_envelope() {
1786 let mut envelope = complete_envelope();
1787 envelope.target_content_fingerprint = None;
1788 let binding = binding_for("model:base", "x", &envelope);
1789 let key = data_binding_requirement_key(&binding.node_id, &binding.input_name);
1790 let input = MethodsPlsPredictInput {
1791 data_content_profile: METHODS_PLS_PREDICT_CONTENT_PROFILE.to_string(),
1792 data_content_fingerprint: methods_pls_predict_feature_content_fingerprint(
1793 &MethodsPlsMatrix {
1794 values: vec![1.0, 2.0],
1795 rows: 1,
1796 cols: 2,
1797 },
1798 )
1799 .unwrap(),
1800 dataset: MethodsPlsDataset {
1801 sample_ids: vec![SampleId::new("sample:1").unwrap()],
1802 x: MethodsPlsMatrix {
1803 values: vec![1.0, 2.0],
1804 rows: 1,
1805 cols: 2,
1806 },
1807 y: None,
1808 target_names: vec!["protein".to_string()],
1809 },
1810 };
1811 let provider = MethodsPlsPredictDataProvider::new(
1812 ControllerId::new("controller:data.methods.predict").unwrap(),
1813 vec![binding.clone()],
1814 envelopes_for(
1815 &binding,
1816 complete_envelope_with_target_free_fingerprint(
1817 input.data_content_fingerprint.clone(),
1818 ),
1819 ),
1820 BTreeMap::from([(key, input.clone())]),
1821 )
1822 .unwrap();
1823 let request = MethodsPlsDataRequest {
1824 node_id: binding.node_id.clone(),
1825 phase: Phase::Predict,
1826 variant_id: None,
1827 fold_id: None,
1828 binding,
1829 identity: None,
1830 fit_view: DataProviderViewSpec {
1831 sample_ids: Some(input.dataset.sample_ids.clone()),
1832 partition: DataRequestPartition::Predict,
1833 fold_id: None,
1834 source_ids: None,
1835 columns: None,
1836 include_augmented: false,
1837 include_excluded: false,
1838 branch_view: None,
1839 extra: BTreeMap::new(),
1840 },
1841 prediction_view: None,
1842 };
1843 assert_eq!(
1844 provider.methods_pls_data(&request).unwrap().fit,
1845 input.dataset
1846 );
1847
1848 let mut wrong_fingerprint = input;
1849 wrong_fingerprint.data_content_fingerprint = "f".repeat(64);
1850 let error = MethodsPlsPredictDataProvider::new(
1851 ControllerId::new("controller:data.methods.predict").unwrap(),
1852 vec![request.binding.clone()],
1853 envelopes_for(
1854 &request.binding,
1855 complete_envelope_with_target_free_fingerprint(
1856 methods_pls_predict_feature_content_fingerprint(&wrong_fingerprint.dataset.x)
1857 .unwrap(),
1858 ),
1859 ),
1860 BTreeMap::from([(
1861 data_binding_requirement_key(&request.binding.node_id, &request.binding.input_name),
1862 wrong_fingerprint,
1863 )]),
1864 )
1865 .unwrap_err();
1866 assert!(error.to_string().contains("feature content fingerprint"));
1867 }
1868
1869 #[test]
1870 fn methods_pls_predict_content_profile_matches_the_python_reference_vector() {
1871 let fingerprint = methods_pls_predict_feature_content_fingerprint(&MethodsPlsMatrix {
1872 values: vec![1.0, 2.0, 3.0, 4.0],
1873 rows: 2,
1874 cols: 2,
1875 })
1876 .unwrap();
1877 assert_eq!(METHODS_PLS_PREDICT_CONTENT_PROFILE, "n4a-matrix-f64-le.v1");
1878 assert_eq!(
1879 fingerprint,
1880 "ca93722602866b81462d63044d1857ea9acb31ee9532e1a891dcb69a2fd41981"
1881 );
1882 }
1883
1884 fn complete_envelope_with_target_free_fingerprint(
1885 data_content_fingerprint: String,
1886 ) -> ExternalDataPlanEnvelope {
1887 let mut envelope = complete_envelope();
1888 envelope.data_content_fingerprint = Some(data_content_fingerprint);
1889 envelope.target_content_fingerprint = None;
1890 envelope
1891 }
1892}