dag_ml_core/
terminal_prediction.rs

1//! Closed, bundle-backed terminal PREDICT execution.
2//!
3//! This module is intentionally narrow.  It consumes the selected variant and
4//! REFIT artifacts already captured in an [`ExecutionBundle`], runs exactly one
5//! scheduler-owned PREDICT replay, and attests one explicitly selected terminal
6//! prediction port against a V2 [`PredictCohort`].  It never asks a host binding
7//! to recover a model by key, replay a Python-side model store, or substitute a
8//! cohort after REFIT.
9
10use std::collections::BTreeMap;
11
12use serde::{Deserialize, Serialize};
13
14use crate::aggregation::{AggregatedPredictionBlock, ObservationPredictionBlock};
15use crate::bundle::{ExecutionBundle, RefitArtifactRecord, ReplayPhaseRequest};
16use crate::campaign::stable_json_fingerprint;
17use crate::controller::ControllerCapability;
18use crate::data::{
19    ExternalDataPlanEnvelope, PredictCohort, EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION_V2,
20};
21use crate::error::{DagMlError, Result};
22use crate::graph::PortKind;
23use crate::ids::{BundleId, NodeId, VariantId};
24use crate::oof::{PredictionBlock, PredictionPartition};
25use crate::phase::Phase;
26use crate::plan::ExecutionPlan;
27use crate::policy::PredictionLevel;
28use crate::runtime::{
29    BundleReplayExecution, EnvelopeAttestedRuntimeDataProvider, RunContext, RuntimeArtifactStore,
30    RuntimeControllerRegistry, RuntimeDataProvider, SequentialScheduler,
31};
32
33/// First public receipt shape for the V2 terminal-prediction boundary.
34pub const TERMINAL_PREDICTION_RECEIPT_SCHEMA_VERSION: u32 = 1;
35
36/// Explicit terminal output selected by the caller.
37///
38/// A node alone is not enough: graphs may expose more than one prediction
39/// output, so the selected producer port is an integrity boundary as well.
40#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
41#[serde(deny_unknown_fields)]
42pub struct TerminalPredictionSelector {
43    pub node_id: NodeId,
44    pub port: String,
45}
46
47impl TerminalPredictionSelector {
48    pub fn new(node_id: NodeId, port: impl Into<String>) -> Result<Self> {
49        let selector = Self {
50            node_id,
51            port: port.into(),
52        };
53        selector.validate()?;
54        Ok(selector)
55    }
56
57    pub fn validate(&self) -> Result<()> {
58        if self.port.trim().is_empty() {
59            return Err(DagMlError::RuntimeValidation(
60                "terminal prediction selector has an empty port".to_string(),
61            ));
62        }
63        Ok(())
64    }
65}
66
67/// Durable attestation for one terminal PREDICT result.
68///
69/// The receipt contains only logical artifact references from the bundle; it
70/// deliberately never serializes invocation-local `HandleRef` values.
71#[derive(Clone, Debug, PartialEq, Serialize)]
72pub struct TerminalPredictionReceipt {
73    schema_version: u32,
74    bundle_id: BundleId,
75    plan_id: String,
76    graph_fingerprint: String,
77    campaign_fingerprint: String,
78    controller_fingerprint: String,
79    selected_variant_id: VariantId,
80    terminal_node_id: NodeId,
81    terminal_port: String,
82    cohort_fingerprint: String,
83    refit_artifacts: Vec<RefitArtifactRecord>,
84    output_fingerprint: String,
85}
86
87impl TerminalPredictionReceipt {
88    pub fn schema_version(&self) -> u32 {
89        self.schema_version
90    }
91
92    pub fn bundle_id(&self) -> &BundleId {
93        &self.bundle_id
94    }
95
96    pub fn plan_id(&self) -> &str {
97        &self.plan_id
98    }
99
100    pub fn graph_fingerprint(&self) -> &str {
101        &self.graph_fingerprint
102    }
103
104    pub fn campaign_fingerprint(&self) -> &str {
105        &self.campaign_fingerprint
106    }
107
108    pub fn controller_fingerprint(&self) -> &str {
109        &self.controller_fingerprint
110    }
111
112    pub fn selected_variant_id(&self) -> &VariantId {
113        &self.selected_variant_id
114    }
115
116    pub fn terminal_node_id(&self) -> &NodeId {
117        &self.terminal_node_id
118    }
119
120    pub fn terminal_port(&self) -> &str {
121        &self.terminal_port
122    }
123
124    pub fn cohort_fingerprint(&self) -> &str {
125        &self.cohort_fingerprint
126    }
127
128    pub fn refit_artifacts(&self) -> &[RefitArtifactRecord] {
129        &self.refit_artifacts
130    }
131
132    pub fn output_fingerprint(&self) -> &str {
133        &self.output_fingerprint
134    }
135}
136
137/// One fully attested terminal PREDICT outcome.
138#[derive(Clone, Debug, PartialEq, Serialize)]
139pub struct TerminalPredictionExecution {
140    prediction: PredictionBlock,
141    receipt: TerminalPredictionReceipt,
142}
143
144impl TerminalPredictionExecution {
145    pub fn prediction(&self) -> &PredictionBlock {
146        &self.prediction
147    }
148
149    pub fn receipt(&self) -> &TerminalPredictionReceipt {
150        &self.receipt
151    }
152}
153
154/// Runtime resources for one closed terminal PREDICT replay.
155///
156/// This groups the borrowed execution boundary so the public API remains
157/// explicit without leaking a positional list of host/runtime dependencies.
158pub struct TerminalPredictionReplay<'a> {
159    pub plan: &'a ExecutionPlan,
160    pub bundle: &'a ExecutionBundle,
161    pub envelope: &'a ExternalDataPlanEnvelope,
162    pub selector: &'a TerminalPredictionSelector,
163    pub controllers: &'a RuntimeControllerRegistry,
164    pub data_provider: &'a dyn RuntimeDataProvider,
165    pub artifact_store: &'a dyn RuntimeArtifactStore,
166}
167
168/// Require the V2, separately attested PREDICT cohort used by this API.
169///
170/// This is intentionally stricter than generic envelope validation: valid V1
171/// envelopes remain supported by frozen CV/REFIT interfaces, but cannot enter
172/// this terminal PREDICT route.
173pub fn require_terminal_predict_cohort(
174    envelope: &ExternalDataPlanEnvelope,
175) -> Result<&PredictCohort> {
176    if envelope.schema_version != EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION_V2 {
177        return Err(DagMlError::RuntimeValidation(format!(
178            "terminal PREDICT requires external data-plan envelope V2, got V{}",
179            envelope.schema_version
180        )));
181    }
182    envelope.validate()?;
183    let cohort = envelope.predict_cohort.as_ref().ok_or_else(|| {
184        DagMlError::RuntimeValidation(
185            "terminal PREDICT requires a V2 envelope with predict_cohort".to_string(),
186        )
187    })?;
188    if cohort.target_names.is_empty() {
189        return Err(DagMlError::RuntimeValidation(
190            "terminal PREDICT requires V2 predict_cohort target_names to be non-empty".to_string(),
191        ));
192    }
193    Ok(cohort)
194}
195
196/// Validate the V2 terminal route before any training-phase controller is invoked.
197///
198/// This is deliberately separate from bundle validation: Python and other
199/// bindings must reject an invalid terminal selector or unsupported aggregation
200/// before they enter CV/REFIT to construct that bundle.  The bundle-backed
201/// checks remain in [`validate_terminal_prediction_request`].
202pub fn validate_terminal_prediction_preflight(
203    plan: &ExecutionPlan,
204    envelope: &ExternalDataPlanEnvelope,
205    selector: &TerminalPredictionSelector,
206) -> Result<()> {
207    let _cohort = require_terminal_predict_cohort(envelope)?;
208    selector.validate()?;
209    plan.validate()?;
210
211    let node_plan = plan.node_plans.get(&selector.node_id).ok_or_else(|| {
212        DagMlError::RuntimeValidation(format!(
213            "terminal PREDICT selector references unknown node `{}`",
214            selector.node_id
215        ))
216    })?;
217    if !node_plan.supported_phases.contains(&Phase::Predict) {
218        return Err(DagMlError::RuntimeValidation(format!(
219            "terminal PREDICT node `{}` does not support PREDICT",
220            selector.node_id
221        )));
222    }
223    if !node_plan
224        .controller_capabilities
225        .contains(&ControllerCapability::EmitsPredictions)
226    {
227        return Err(DagMlError::RuntimeValidation(format!(
228            "terminal PREDICT node `{}` lacks the emits_predictions capability",
229            selector.node_id
230        )));
231    }
232
233    let graph_node = plan
234        .graph_plan
235        .graph
236        .nodes
237        .iter()
238        .find(|node| node.id == selector.node_id)
239        .ok_or_else(|| {
240            DagMlError::RuntimeValidation(format!(
241                "terminal PREDICT selector node `{}` is absent from the graph",
242                selector.node_id
243            ))
244        })?;
245    let output = graph_node
246        .ports
247        .outputs
248        .iter()
249        .find(|port| port.name == selector.port)
250        .ok_or_else(|| {
251            DagMlError::RuntimeValidation(format!(
252                "terminal PREDICT selector `{}` has no output port `{}`",
253                selector.node_id, selector.port
254            ))
255        })?;
256    if output.kind != PortKind::Prediction {
257        return Err(DagMlError::RuntimeValidation(format!(
258            "terminal PREDICT selector `{}.{}` is not a prediction port",
259            selector.node_id, selector.port
260        )));
261    }
262    if plan.graph_plan.graph.edges.iter().any(|edge| {
263        edge.source.node_id == selector.node_id && edge.source.port_name == selector.port
264    }) {
265        return Err(DagMlError::RuntimeValidation(format!(
266            "terminal PREDICT selector `{}.{}` is consumed by another graph node",
267            selector.node_id, selector.port
268        )));
269    }
270
271    if plan
272        .graph_plan
273        .graph
274        .edges
275        .iter()
276        .any(|edge| edge.contract.requires_oof)
277    {
278        return Err(DagMlError::RuntimeValidation(
279            "terminal PREDICT first slice does not replay requires_oof edges; use the ordinary replay contract until prediction-cache closure is captured"
280                .to_string(),
281        ));
282    }
283
284    if plan.campaign.aggregation_policy.aggregation_level != PredictionLevel::Sample
285        || plan.node_plans.values().any(|node| {
286            node.supported_phases.contains(&Phase::Predict)
287                && node.shape_plan.as_ref().is_some_and(|shape_plan| {
288                    shape_plan.aggregation_policy.aggregation_level != PredictionLevel::Sample
289                })
290        })
291    {
292        return Err(DagMlError::RuntimeValidation(format!(
293            "terminal PREDICT selector `{}.{}` uses unsupported non-sample aggregation",
294            selector.node_id, selector.port
295        )));
296    }
297    Ok(())
298}
299
300/// Validate the closed V2 request before any PREDICT controller is invoked.
301///
302/// The selected output must be a graph-terminal prediction port.  This first
303/// slice intentionally supports direct sample-level predictions only; relation
304/// aggregation is refused rather than accidentally using the CV coordinator
305/// relation universe for an external cohort.
306pub fn validate_terminal_prediction_request(
307    plan: &ExecutionPlan,
308    bundle: &ExecutionBundle,
309    envelope: &ExternalDataPlanEnvelope,
310    selector: &TerminalPredictionSelector,
311) -> Result<()> {
312    validate_terminal_prediction_preflight(plan, envelope, selector)?;
313    bundle.validate_against_plan(plan)?;
314
315    let _selected_variant = bundle.selected_variant_id.as_ref().ok_or_else(|| {
316        DagMlError::RuntimeValidation(format!(
317            "terminal PREDICT requires bundle `{}` to select one variant",
318            bundle.bundle_id
319        ))
320    })?;
321
322    if bundle.data_requirements.is_empty() {
323        return Err(DagMlError::RuntimeValidation(format!(
324            "terminal PREDICT bundle `{}` has no relation-attested data requirement",
325            bundle.bundle_id
326        )));
327    }
328    let envelopes = terminal_prediction_envelopes(bundle, envelope);
329    bundle.validate_replay_envelopes(&envelopes)?;
330
331    for stateful_node in plan.node_plans.values().filter(|node| {
332        node.supported_phases.contains(&Phase::Predict)
333            && node
334                .controller_capabilities
335                .contains(&ControllerCapability::Stateful)
336    }) {
337        if !bundle
338            .refit_artifacts
339            .iter()
340            .any(|artifact| artifact.node_id == stateful_node.node_id)
341        {
342            return Err(DagMlError::RuntimeValidation(format!(
343                "terminal PREDICT stateful node `{}` has no REFIT artifact in bundle `{}`",
344                stateful_node.node_id, bundle.bundle_id
345            )));
346        }
347    }
348    Ok(())
349}
350
351/// Ask the provider to attest every PREDICT-visible binding before replay
352/// materializes an artifact, data handle, or controller task.  The scheduler
353/// repeats this comparison at each request through the envelope wrapper below
354/// so a mutable host provider cannot change cohorts between preflight and
355/// materialization.
356fn validate_terminal_provider_cohorts(
357    plan: &ExecutionPlan,
358    envelope: &ExternalDataPlanEnvelope,
359    data_provider: &dyn RuntimeDataProvider,
360) -> Result<()> {
361    let expected = require_terminal_predict_cohort(envelope)?;
362    for node in plan
363        .node_plans
364        .values()
365        .filter(|node| node.supported_phases.contains(&Phase::Predict))
366    {
367        for binding in &node.data_bindings {
368            binding.validate_envelope(envelope)?;
369            let supplied = data_provider.predict_cohort(binding, Phase::Predict)?;
370            if supplied.as_ref() != Some(expected) {
371                return Err(DagMlError::RuntimeValidation(format!(
372                    "PREDICT cohort supplied by runtime provider for binding `{}` does not exactly match its terminal V2 envelope",
373                    crate::data::data_binding_requirement_key(&binding.node_id, &binding.input_name)
374                )));
375            }
376        }
377    }
378    Ok(())
379}
380
381/// Execute one scheduler-owned PREDICT replay from a captured bundle.
382///
383/// `artifact_store` is the runtime's REFIT artifact store.  It is the only
384/// model source used here: no Python model registry, model refetch, or host
385/// replay fallback is consulted by this API.
386pub fn execute_terminal_prediction(
387    replay: TerminalPredictionReplay<'_>,
388    ctx: &mut RunContext,
389) -> Result<TerminalPredictionExecution> {
390    let TerminalPredictionReplay {
391        plan,
392        bundle,
393        envelope,
394        selector,
395        controllers,
396        data_provider,
397        artifact_store,
398    } = replay;
399    validate_terminal_prediction_request(plan, bundle, envelope, selector)?;
400    validate_terminal_provider_cohorts(plan, envelope, data_provider)?;
401
402    let data_envelopes = terminal_prediction_envelopes(bundle, envelope);
403    // Ask the host provider to attest the same V2 cohort before it can
404    // materialize or view any data.  The wrapper returns the envelope-bound
405    // cohort to the scheduler only after that exact comparison succeeds.
406    let borrowed_data_provider = BorrowedRuntimeDataProvider(data_provider);
407    let envelope_attested_data_provider = EnvelopeAttestedRuntimeDataProvider::new(
408        borrowed_data_provider,
409        terminal_prediction_bindings(plan),
410        data_envelopes.clone(),
411    )?;
412    let attested_data_provider = TerminalCohortAttestedRuntimeDataProvider {
413        inner: envelope_attested_data_provider,
414        source: data_provider,
415        expected_cohort: require_terminal_predict_cohort(envelope)?.clone(),
416    };
417    let replay_request = ReplayPhaseRequest {
418        bundle_id: bundle.bundle_id.clone(),
419        phase: Phase::Predict,
420        data_envelope_keys: data_envelopes.keys().cloned().collect(),
421    };
422    let results = SequentialScheduler.execute_direct_sample_bundle_replay(
423        BundleReplayExecution {
424            plan,
425            bundle,
426            replay_request: &replay_request,
427            prediction_cache_store: None,
428            controllers,
429            data_provider: &attested_data_provider,
430            artifact_store,
431            data_envelopes: &data_envelopes,
432        },
433        ctx,
434    )?;
435
436    let predictions = results
437        .iter()
438        .flat_map(|result| result.predictions.iter().cloned())
439        .collect::<Vec<_>>();
440    let observation_predictions = results
441        .iter()
442        .flat_map(|result| result.observation_predictions.iter().cloned())
443        .collect::<Vec<_>>();
444    let aggregated_predictions = results
445        .iter()
446        .flat_map(|result| result.aggregated_predictions.iter().cloned())
447        .collect::<Vec<_>>();
448    attest_terminal_prediction_output(
449        plan,
450        bundle,
451        envelope,
452        selector,
453        &predictions,
454        &observation_predictions,
455        &aggregated_predictions,
456    )
457}
458
459/// Attest scheduler-produced blocks as one exact, sample-level terminal result.
460///
461/// This is intentionally private to the scheduler-owned replay route.  Public
462/// callers receive its receipt only through [`execute_terminal_prediction`],
463/// so they cannot turn host-supplied prediction blocks into a bundle-backed
464/// terminal claim.
465fn attest_terminal_prediction_output(
466    plan: &ExecutionPlan,
467    bundle: &ExecutionBundle,
468    envelope: &ExternalDataPlanEnvelope,
469    selector: &TerminalPredictionSelector,
470    predictions: &[PredictionBlock],
471    observation_predictions: &[ObservationPredictionBlock],
472    aggregated_predictions: &[AggregatedPredictionBlock],
473) -> Result<TerminalPredictionExecution> {
474    validate_terminal_prediction_request(plan, bundle, envelope, selector)?;
475    let cohort = require_terminal_predict_cohort(envelope)?;
476
477    if observation_predictions.iter().any(|block| {
478        block.producer_node == selector.node_id
479            && block.producer_port.as_deref() == Some(selector.port.as_str())
480    }) || aggregated_predictions.iter().any(|block| {
481        block.producer_node == selector.node_id
482            && block.producer_port.as_deref() == Some(selector.port.as_str())
483    }) {
484        return Err(DagMlError::RuntimeValidation(format!(
485            "terminal PREDICT selector `{}.{}` emitted unsupported aggregated predictions",
486            selector.node_id, selector.port
487        )));
488    }
489
490    let matching = predictions
491        .iter()
492        .filter(|block| {
493            block.producer_node == selector.node_id
494                && block.producer_port.as_deref() == Some(selector.port.as_str())
495        })
496        .collect::<Vec<_>>();
497    let [prediction] = matching.as_slice() else {
498        return Err(DagMlError::RuntimeValidation(format!(
499            "terminal PREDICT selector `{}.{}` must emit exactly one prediction block, got {}",
500            selector.node_id,
501            selector.port,
502            matching.len()
503        )));
504    };
505    prediction.validate_content()?;
506    if prediction.partition != PredictionPartition::Final || prediction.fold_id.is_some() {
507        return Err(DagMlError::RuntimeValidation(format!(
508            "terminal PREDICT selector `{}.{}` must emit a top-level Final prediction block",
509            selector.node_id, selector.port
510        )));
511    }
512    if prediction.sample_ids != cohort.physical_sample_ids {
513        return Err(DagMlError::RuntimeValidation(format!(
514            "terminal PREDICT selector `{}.{}` sample identities do not exactly match the V2 predict cohort",
515            selector.node_id, selector.port
516        )));
517    }
518    if prediction.target_names != cohort.target_names {
519        return Err(DagMlError::RuntimeValidation(format!(
520            "terminal PREDICT selector `{}.{}` target names do not exactly match the V2 predict cohort",
521            selector.node_id, selector.port
522        )));
523    }
524
525    let selected_variant_id = bundle.selected_variant_id.clone().ok_or_else(|| {
526        DagMlError::RuntimeValidation(format!(
527            "terminal PREDICT bundle `{}` lost its selected variant during attestation",
528            bundle.bundle_id
529        ))
530    })?;
531    Ok(TerminalPredictionExecution {
532        prediction: (*prediction).clone(),
533        receipt: TerminalPredictionReceipt {
534            schema_version: TERMINAL_PREDICTION_RECEIPT_SCHEMA_VERSION,
535            bundle_id: bundle.bundle_id.clone(),
536            plan_id: bundle.plan_id.clone(),
537            graph_fingerprint: bundle.graph_fingerprint.clone(),
538            campaign_fingerprint: bundle.campaign_fingerprint.clone(),
539            controller_fingerprint: bundle.controller_fingerprint.clone(),
540            selected_variant_id,
541            terminal_node_id: selector.node_id.clone(),
542            terminal_port: selector.port.clone(),
543            cohort_fingerprint: cohort.cohort_fingerprint.clone(),
544            refit_artifacts: bundle.refit_artifacts.clone(),
545            output_fingerprint: stable_json_fingerprint(prediction)?,
546        },
547    })
548}
549
550fn terminal_prediction_envelopes(
551    bundle: &ExecutionBundle,
552    envelope: &ExternalDataPlanEnvelope,
553) -> BTreeMap<String, ExternalDataPlanEnvelope> {
554    bundle
555        .data_requirements
556        .iter()
557        .map(|requirement| (requirement.key(), envelope.clone()))
558        .collect()
559}
560
561fn terminal_prediction_bindings(plan: &ExecutionPlan) -> Vec<crate::data::DataBinding> {
562    plan.node_plans
563        .values()
564        .flat_map(|node| node.data_bindings.iter().cloned())
565        .collect()
566}
567
568/// Private borrowing adapter used only to compose the public provider trait
569/// with the existing envelope-attestation wrapper.  Keeping this newtype
570/// local avoids adding a blanket trait implementation that downstream crates
571/// could conflict with.
572struct BorrowedRuntimeDataProvider<'a>(&'a dyn RuntimeDataProvider);
573
574impl RuntimeDataProvider for BorrowedRuntimeDataProvider<'_> {
575    fn materialize(
576        &self,
577        request: &crate::runtime::DataMaterializationRequest,
578    ) -> Result<crate::runtime::HandleRef> {
579        self.0.materialize(request)
580    }
581
582    fn make_view(
583        &self,
584        request: &crate::runtime::DataViewRequest,
585    ) -> Result<crate::runtime::HandleRef> {
586        self.0.make_view(request)
587    }
588
589    fn training_data_identity(
590        &self,
591        binding: &crate::data::DataBinding,
592    ) -> Result<Option<crate::training::TrainingDataIdentity>> {
593        self.0.training_data_identity(binding)
594    }
595
596    fn coordinator_relations(
597        &self,
598        binding: &crate::data::DataBinding,
599    ) -> Result<Option<crate::relation::SampleRelationSet>> {
600        self.0.coordinator_relations(binding)
601    }
602
603    fn predict_cohort(
604        &self,
605        binding: &crate::data::DataBinding,
606        phase: Phase,
607    ) -> Result<Option<PredictCohort>> {
608        self.0.predict_cohort(binding, phase)
609    }
610
611    fn methods_pls_capability(&self) -> Result<()> {
612        self.0.methods_pls_capability()
613    }
614
615    fn preflight_methods_pls(&self, request: &crate::runtime::MethodsPlsDataRequest) -> Result<()> {
616        self.0.preflight_methods_pls(request)
617    }
618
619    fn methods_pls_data(
620        &self,
621        request: &crate::runtime::MethodsPlsDataRequest,
622    ) -> Result<crate::runtime::MethodsPlsData> {
623        self.0.methods_pls_data(request)
624    }
625}
626
627/// Terminal-only cohort guard layered outside the generic envelope wrapper.
628///
629/// Generic V2 replay may obtain its cohort solely from an attested envelope:
630/// older C-ABI providers intentionally do not expose the optional
631/// `predict_cohort` callback.  This narrow route has a stronger contract: the
632/// host must independently attest the exact V2 terminal cohort before every
633/// PREDICT materialization.  Keeping that comparison here preserves the
634/// generic C-ABI surface while preventing a mutable terminal provider from
635/// substituting a cohort after the all-binding preflight.
636struct TerminalCohortAttestedRuntimeDataProvider<'a, P> {
637    inner: P,
638    source: &'a dyn RuntimeDataProvider,
639    expected_cohort: PredictCohort,
640}
641
642impl<P: RuntimeDataProvider> RuntimeDataProvider
643    for TerminalCohortAttestedRuntimeDataProvider<'_, P>
644{
645    fn materialize(
646        &self,
647        request: &crate::runtime::DataMaterializationRequest,
648    ) -> Result<crate::runtime::HandleRef> {
649        self.inner.materialize(request)
650    }
651
652    fn make_view(
653        &self,
654        request: &crate::runtime::DataViewRequest,
655    ) -> Result<crate::runtime::HandleRef> {
656        self.inner.make_view(request)
657    }
658
659    fn training_data_identity(
660        &self,
661        binding: &crate::data::DataBinding,
662    ) -> Result<Option<crate::training::TrainingDataIdentity>> {
663        self.inner.training_data_identity(binding)
664    }
665
666    fn coordinator_relations(
667        &self,
668        binding: &crate::data::DataBinding,
669    ) -> Result<Option<crate::relation::SampleRelationSet>> {
670        self.inner.coordinator_relations(binding)
671    }
672
673    fn predict_cohort(
674        &self,
675        binding: &crate::data::DataBinding,
676        phase: Phase,
677    ) -> Result<Option<PredictCohort>> {
678        let supplied = self.source.predict_cohort(binding, phase)?;
679        if supplied.as_ref() != Some(&self.expected_cohort) {
680            return Err(DagMlError::RuntimeValidation(format!(
681                "PREDICT cohort supplied by runtime provider for binding `{}` does not exactly match its terminal V2 envelope",
682                crate::data::data_binding_requirement_key(&binding.node_id, &binding.input_name)
683            )));
684        }
685        let attested = self.inner.predict_cohort(binding, phase)?;
686        if attested.as_ref() != Some(&self.expected_cohort) {
687            return Err(DagMlError::RuntimeValidation(format!(
688                "terminal envelope cohort for binding `{}` changed during replay",
689                crate::data::data_binding_requirement_key(&binding.node_id, &binding.input_name)
690            )));
691        }
692        Ok(attested)
693    }
694
695    fn methods_pls_capability(&self) -> Result<()> {
696        self.inner.methods_pls_capability()
697    }
698
699    fn preflight_methods_pls(&self, request: &crate::runtime::MethodsPlsDataRequest) -> Result<()> {
700        self.inner.preflight_methods_pls(request)
701    }
702
703    fn methods_pls_data(
704        &self,
705        request: &crate::runtime::MethodsPlsDataRequest,
706    ) -> Result<crate::runtime::MethodsPlsData> {
707        self.inner.methods_pls_data(request)
708    }
709}
710
711#[cfg(test)]
712mod tests {
713    use std::cell::Cell;
714    use std::collections::BTreeMap;
715    use std::sync::{
716        atomic::{AtomicUsize, Ordering},
717        Arc,
718    };
719
720    use crate::bundle::build_execution_bundle;
721    use crate::controller::{ControllerManifest, ControllerRegistry};
722    use crate::data::{InMemoryDataProvider, PredictCohortRole};
723    use crate::graph::GraphSpec;
724    use crate::ids::{ArtifactId, ControllerId, LineageId, RunId};
725    use crate::plan::{build_execution_plan, CampaignSpec};
726    use crate::policy::{
727        AggregationControllerSpec, AggregationMethod, AggregationPolicy, DataModelShapePlan,
728        FitBoundary, Granularity,
729    };
730    use crate::relation::SampleRelationSet;
731    use crate::runtime::{
732        AggregationControllerResult, AggregationControllerTask, ArtifactRef,
733        DataMaterializationRequest, DataViewRequest, HandleKind, HandleRef, InMemoryArtifactStore,
734        LineageRecord, NodeResult, NodeTask, RuntimeController,
735    };
736
737    use super::*;
738
739    const SCHEMA_FINGERPRINT: &str =
740        "f97b37872fa22134b508f98fd8e207e5b776b52594fb8f6f5c3e15bee212246b";
741    const PLAN_FINGERPRINT: &str =
742        "7c5431d85574b3f337022fa5d25971d5b5cf445b90331b49938f573ff6901e4d";
743
744    fn terminal_fixture() -> (ExecutionPlan, ExecutionBundle, ExternalDataPlanEnvelope) {
745        terminal_fixture_with_custom_aggregation(false)
746    }
747
748    fn terminal_fixture_with_custom_aggregation(
749        custom_aggregation: bool,
750    ) -> (ExecutionPlan, ExecutionBundle, ExternalDataPlanEnvelope) {
751        let cv_envelope: ExternalDataPlanEnvelope = serde_json::from_str(include_str!(
752            "../tests/fixtures/package/data/coordinator_data_plan_envelope_sample12.json"
753        ))
754        .expect("fixture envelope parses");
755        assert_eq!(cv_envelope.schema_fingerprint, SCHEMA_FINGERPRINT);
756        assert_eq!(cv_envelope.plan_fingerprint, PLAN_FINGERPRINT);
757
758        let graph: GraphSpec = serde_json::from_str(
759            r#"{
760              "id": "graph:terminal.predict",
761              "interface": {"inputs": [], "outputs": []},
762              "nodes": [{
763                "id": "model:terminal",
764                "kind": "model",
765                "operator": null,
766                "params": {},
767                "ports": {
768                  "inputs": [{"name": "x", "kind": "data", "representation": null, "cardinality": "one", "description": ""}],
769                  "outputs": [{"name": "prediction", "kind": "prediction", "representation": null, "cardinality": "one", "description": ""}]
770                },
771                "metadata": {},
772                "seed_label": null
773              }],
774              "edges": [],
775              "search_space_fingerprint": null,
776              "metadata": {}
777            }"#,
778        )
779        .expect("terminal graph parses");
780        let mut campaign: CampaignSpec = serde_json::from_str(&format!(
781            r#"{{
782              "id": "campaign:terminal.predict",
783              "root_seed": 7,
784              "leakage_policy": {{"split_unit": "sample", "forbid_origin_cross_fold": true,
785                "allow_observation_split_with_shared_target": false, "require_group_ids": false, "unsafe_flags": []}},
786              "aggregation_policy": {{"aggregation_level": "sample", "method": "mean", "weights": "none",
787                "emit_parallel_metrics": true, "selection_metric_level": "sample",
788                "store_raw_predictions": true, "store_aggregated_predictions": true}},
789              "split_invocation": {{
790                "id": "split:terminal.predict", "controller_id": null,
791                "leakage_policy": {{"split_unit": "sample", "forbid_origin_cross_fold": true,
792                  "allow_observation_split_with_shared_target": false, "require_group_ids": false, "unsafe_flags": []}},
793                "params": {{}},
794                "fold_set": {{
795                  "id": "folds:terminal.predict", "sample_ids": ["sample:1", "sample:2"],
796                  "folds": [
797                    {{"fold_id": "fold:0", "train_sample_ids": ["sample:2"], "validation_sample_ids": ["sample:1"], "metadata": {{}}}},
798                    {{"fold_id": "fold:1", "train_sample_ids": ["sample:1"], "validation_sample_ids": ["sample:2"], "metadata": {{}}}}
799                  ], "sample_groups": {{}}
800                }}
801              }},
802              "generation": {{"strategy": "none", "dimensions": [], "max_variants": 1}},
803              "shape_plans": {{}},
804              "data_bindings": {{"model:terminal": [{{
805                "node_id": "model:terminal", "input_name": "x", "request_id": "nir-to-tabular",
806                "schema_fingerprint": "{SCHEMA_FINGERPRINT}", "plan_fingerprint": "{PLAN_FINGERPRINT}",
807                "relation_fingerprint": "{}", "output_representation": "tabular_numeric",
808                "feature_set_id": "x", "source_ids": ["nir"], "require_relations": true
809              }}]}},
810              "metadata": {{}}
811            }}"#,
812            cv_envelope
813                .relation_fingerprint
814                .as_deref()
815                .expect("fixture relation fingerprint"),
816        ))
817        .expect("terminal campaign parses");
818        if custom_aggregation {
819            let node_id = NodeId::new("model:terminal").unwrap();
820            campaign.shape_plans.insert(
821                node_id.clone(),
822                DataModelShapePlan {
823                    node_id,
824                    input_granularity: Granularity::Observation,
825                    target_granularity: Granularity::Sample,
826                    fit_rows: FitBoundary::FoldTrain,
827                    predict_rows: FitBoundary::FoldValidation,
828                    feature_namespace: Some("nir".to_string()),
829                    feature_schema_fingerprint: None,
830                    target_space: "regression:protein".to_string(),
831                    aggregation_policy: AggregationPolicy {
832                        aggregation_level: PredictionLevel::Sample,
833                        method: AggregationMethod::CustomController,
834                        custom_controller: Some(AggregationControllerSpec {
835                            controller_id: ControllerId::new("controller:agg.custom").unwrap(),
836                            params: serde_json::json!({"terminal": true}),
837                        }),
838                        ..AggregationPolicy::default()
839                    },
840                    augmentation_policy: Default::default(),
841                    selection_policy: Default::default(),
842                },
843            );
844        }
845        let manifest: ControllerManifest = serde_json::from_str(
846            r#"{
847              "controller_id": "controller:model",
848              "controller_version": "0.1.0",
849              "operator_kind": "model",
850              "priority": 0,
851              "supported_phases": ["FIT_CV", "REFIT", "PREDICT"],
852              "input_ports": [],
853              "output_ports": [],
854              "data_requirements": null,
855              "capabilities": ["deterministic", "thread_safe", "process_safe", "emits_predictions", "emits_artifacts", "stateful"],
856              "fit_scope": "fold_train",
857              "rng_policy": "uses_core_seed",
858              "artifact_policy": "serializable"
859            }"#,
860        )
861        .expect("terminal controller manifest parses");
862        let mut controllers = ControllerRegistry::new();
863        controllers.register(manifest).expect("manifest registers");
864        if custom_aggregation {
865            let aggregation_manifest: ControllerManifest = serde_json::from_str(
866                r#"{
867                  "controller_id": "controller:agg.custom",
868                  "controller_version": "0.1.0",
869                  "operator_kind": "aggregator",
870                  "priority": 0,
871                  "supported_phases": ["PLAN"],
872                  "input_ports": [],
873                  "output_ports": [],
874                  "data_requirements": null,
875                  "capabilities": ["deterministic", "thread_safe", "process_safe", "aggregates_predictions"],
876                  "fit_scope": "inference_only",
877                  "rng_policy": "uses_core_seed",
878                  "artifact_policy": "serializable"
879                }"#,
880            )
881            .expect("aggregation controller manifest parses");
882            controllers
883                .register(aggregation_manifest)
884                .expect("aggregation controller registers");
885        }
886        let plan = build_execution_plan("plan:terminal.predict", graph, campaign, &controllers)
887            .expect("terminal plan builds");
888
889        let heldout_relations: SampleRelationSet = serde_json::from_str(
890            r#"{
891              "records": [
892                {"observation_id": "obs.H001", "sample_id": "sample:heldout:1", "target_id": "target:heldout:1", "group_id": "group:heldout", "origin_sample_id": null, "source_id": "nir", "is_augmented": false},
893                {"observation_id": "obs.H002", "sample_id": "sample:heldout:2", "target_id": "target:heldout:2", "group_id": "group:heldout", "origin_sample_id": null, "source_id": "nir", "is_augmented": false}
894              ]
895            }"#,
896        )
897        .expect("heldout relations parse");
898        let cohort = PredictCohort::from_relations(
899            PredictCohortRole::ExternalTest,
900            heldout_relations,
901            vec!["protein".to_string()],
902            "a".repeat(64),
903            Some("b".repeat(64)),
904        )
905        .expect("predict cohort builds");
906        let mut envelope = cv_envelope;
907        envelope.schema_version = EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION_V2;
908        envelope.predict_cohort = Some(cohort);
909        envelope.validate().expect("V2 envelope validates");
910        plan.campaign
911            .validate_data_envelope_relations(&envelope)
912            .expect("campaign accepts closed external cohort");
913
914        let selected_variant_id = plan
915            .variants
916            .first()
917            .expect("single base variant")
918            .variant_id
919            .clone();
920        let node_plan = plan
921            .node_plans
922            .get(&NodeId::new("model:terminal").unwrap())
923            .expect("terminal node plan");
924        let artifact = RefitArtifactRecord {
925            node_id: node_plan.node_id.clone(),
926            controller_id: node_plan.controller_id.clone(),
927            artifact: ArtifactRef {
928                id: ArtifactId::new("artifact:model:terminal:refit").unwrap(),
929                kind: "mock_model".to_string(),
930                controller_id: ControllerId::new("controller:model").unwrap(),
931                backend: None,
932                uri: None,
933                content_fingerprint: None,
934                size_bytes: Some(1),
935                plugin: None,
936                plugin_version: None,
937            },
938            params_fingerprint: node_plan.params_fingerprint.clone(),
939            training_loss_fingerprint: None,
940            data_requirement_keys: vec!["model:terminal.x".to_string()],
941            prediction_requirement_keys: Vec::new(),
942        };
943        let bundle = build_execution_bundle(
944            BundleId::new("bundle:terminal.predict").unwrap(),
945            &plan,
946            Some(selected_variant_id),
947            BTreeMap::new(),
948            vec![artifact],
949        )
950        .expect("terminal bundle builds");
951        (plan, bundle, envelope)
952    }
953
954    fn terminal_prediction(envelope: &ExternalDataPlanEnvelope) -> PredictionBlock {
955        let cohort = envelope.predict_cohort.as_ref().expect("V2 cohort");
956        PredictionBlock {
957            prediction_id: Some("prediction:terminal".to_string()),
958            producer_node: NodeId::new("model:terminal").unwrap(),
959            producer_port: Some("prediction".to_string()),
960            partition: PredictionPartition::Final,
961            fold_id: None,
962            sample_ids: cohort.physical_sample_ids.clone(),
963            values: vec![vec![0.1], vec![0.2]],
964            target_names: cohort.target_names.clone(),
965        }
966    }
967
968    fn selector() -> TerminalPredictionSelector {
969        TerminalPredictionSelector::new(NodeId::new("model:terminal").unwrap(), "prediction")
970            .unwrap()
971    }
972
973    struct NeverInvokedController {
974        id: ControllerId,
975    }
976
977    impl RuntimeController for NeverInvokedController {
978        fn controller_id(&self) -> &ControllerId {
979            &self.id
980        }
981
982        fn invoke(&self, _task: &NodeTask) -> Result<NodeResult> {
983            Err(DagMlError::RuntimeValidation(
984                "terminal test controller must not be invoked".to_string(),
985            ))
986        }
987    }
988
989    struct SubstitutedCohortProvider {
990        cohort: PredictCohort,
991        materialize_calls: Cell<usize>,
992        view_calls: Cell<usize>,
993    }
994
995    impl RuntimeDataProvider for SubstitutedCohortProvider {
996        fn materialize(&self, _request: &DataMaterializationRequest) -> Result<HandleRef> {
997            self.materialize_calls.set(self.materialize_calls.get() + 1);
998            Err(DagMlError::RuntimeValidation(
999                "substituted provider was allowed to materialize data".to_string(),
1000            ))
1001        }
1002
1003        fn make_view(&self, _request: &DataViewRequest) -> Result<HandleRef> {
1004            self.view_calls.set(self.view_calls.get() + 1);
1005            Err(DagMlError::RuntimeValidation(
1006                "substituted provider was allowed to construct a view".to_string(),
1007            ))
1008        }
1009
1010        fn predict_cohort(
1011            &self,
1012            _binding: &crate::data::DataBinding,
1013            phase: Phase,
1014        ) -> Result<Option<PredictCohort>> {
1015            assert_eq!(phase, Phase::Predict);
1016            Ok(Some(self.cohort.clone()))
1017        }
1018    }
1019
1020    struct ObservationTerminalController {
1021        id: ControllerId,
1022    }
1023
1024    impl RuntimeController for ObservationTerminalController {
1025        fn controller_id(&self) -> &ControllerId {
1026            &self.id
1027        }
1028
1029        fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
1030            let cohort = task
1031                .data_views
1032                .get("data:x")
1033                .and_then(|view| view.sample_ids.as_ref())
1034                .ok_or_else(|| {
1035                    DagMlError::RuntimeValidation(
1036                        "terminal observation test did not receive the V2 cohort view".to_string(),
1037                    )
1038                })?;
1039            let observation_ids = cohort
1040                .iter()
1041                .enumerate()
1042                .map(|(index, _)| crate::ids::ObservationId::new(format!("obs:terminal:{index}")))
1043                .collect::<Result<Vec<_>>>()?;
1044            Ok(NodeResult {
1045                schema_version: None,
1046                node_id: task.node_plan.node_id.clone(),
1047                outputs: BTreeMap::new(),
1048                predictions: Vec::new(),
1049                observation_predictions: vec![ObservationPredictionBlock {
1050                    prediction_id: Some("prediction:terminal:observation".to_string()),
1051                    producer_node: task.node_plan.node_id.clone(),
1052                    producer_port: Some("prediction".to_string()),
1053                    partition: PredictionPartition::Final,
1054                    fold_id: None,
1055                    observation_ids,
1056                    values: vec![vec![0.1]; cohort.len()],
1057                    weights: Vec::new(),
1058                    target_names: vec!["protein".to_string()],
1059                }],
1060                aggregated_predictions: Vec::new(),
1061                explanations: Vec::new(),
1062                shape_deltas: Vec::new(),
1063                artifacts: Vec::new(),
1064                artifact_handles: BTreeMap::new(),
1065                fit_influence_diagnostics: Vec::new(),
1066                regression_targets: Vec::new(),
1067                lineage: LineageRecord {
1068                    record_id: LineageId::new(format!(
1069                        "lineage:terminal:observation:{}",
1070                        task.phase.as_str()
1071                    ))?,
1072                    run_id: task.run_id.clone(),
1073                    node_id: task.node_plan.node_id.clone(),
1074                    phase: task.phase,
1075                    controller_id: self.id.clone(),
1076                    controller_version: task.node_plan.controller_version.clone(),
1077                    variant_id: task.variant_id.clone(),
1078                    fold_id: task.fold_id.clone(),
1079                    branch_path: task.branch_path.clone(),
1080                    input_lineage: Vec::new(),
1081                    artifact_refs: Vec::new(),
1082                    params_fingerprint: task.node_plan.params_fingerprint.clone(),
1083                    data_model_shape_fingerprint: None,
1084                    aggregation_policy_fingerprint: None,
1085                    seed: task.seed,
1086                    unsafe_flags: Default::default(),
1087                    metrics: BTreeMap::new(),
1088                    loss_attestations: Vec::new(),
1089                    early_stopping_records: Vec::new(),
1090                },
1091            })
1092        }
1093    }
1094
1095    struct CountingAggregationController {
1096        id: ControllerId,
1097        calls: Arc<AtomicUsize>,
1098    }
1099
1100    impl RuntimeController for CountingAggregationController {
1101        fn controller_id(&self) -> &ControllerId {
1102            &self.id
1103        }
1104
1105        fn invoke(&self, _task: &NodeTask) -> Result<NodeResult> {
1106            Err(DagMlError::RuntimeValidation(
1107                "aggregation controller must not receive a node task".to_string(),
1108            ))
1109        }
1110
1111        fn invoke_aggregation(
1112            &self,
1113            _task: &AggregationControllerTask,
1114        ) -> Result<AggregationControllerResult> {
1115            self.calls.fetch_add(1, Ordering::SeqCst);
1116            Err(DagMlError::RuntimeValidation(
1117                "terminal direct replay must not invoke aggregation".to_string(),
1118            ))
1119        }
1120    }
1121
1122    fn registered_terminal_artifact_store(bundle: &ExecutionBundle) -> InMemoryArtifactStore {
1123        let artifact = bundle
1124            .refit_artifacts
1125            .first()
1126            .expect("terminal fixture has one REFIT artifact");
1127        let mut store = InMemoryArtifactStore::new();
1128        store
1129            .register(
1130                artifact,
1131                HandleRef {
1132                    handle: 901,
1133                    kind: HandleKind::Model,
1134                    owner_controller: artifact.controller_id.clone(),
1135                },
1136            )
1137            .expect("terminal artifact registers");
1138        store
1139    }
1140
1141    #[test]
1142    fn terminal_receipt_binds_exact_v2_cohort_and_logical_refit_artifact() {
1143        let (plan, bundle, envelope) = terminal_fixture();
1144        let prediction = terminal_prediction(&envelope);
1145        let execution = attest_terminal_prediction_output(
1146            &plan,
1147            &bundle,
1148            &envelope,
1149            &selector(),
1150            std::slice::from_ref(&prediction),
1151            &[],
1152            &[],
1153        )
1154        .expect("exact terminal prediction is accepted");
1155        assert_eq!(execution.prediction(), &prediction);
1156        assert_eq!(
1157            execution.receipt().cohort_fingerprint(),
1158            envelope
1159                .predict_cohort
1160                .as_ref()
1161                .expect("V2 cohort")
1162                .cohort_fingerprint
1163        );
1164        assert_eq!(execution.receipt().refit_artifacts().len(), 1);
1165        assert_ne!(execution.receipt().output_fingerprint(), "");
1166        let receipt_json = serde_json::to_value(execution.receipt()).unwrap();
1167        assert!(receipt_json.get("handle").is_none());
1168    }
1169
1170    #[test]
1171    fn terminal_request_refuses_v1_or_v2_without_cohort() {
1172        let (plan, bundle, envelope) = terminal_fixture();
1173        let mut v1 = envelope.clone();
1174        v1.schema_version = crate::data::EXTERNAL_DATA_PLAN_ENVELOPE_SCHEMA_VERSION_V1;
1175        v1.predict_cohort = None;
1176        let error = validate_terminal_prediction_request(&plan, &bundle, &v1, &selector())
1177            .expect_err("V1 must not enter terminal PREDICT");
1178        assert!(error
1179            .to_string()
1180            .contains("requires external data-plan envelope V2"));
1181
1182        let mut missing = envelope;
1183        missing.predict_cohort = None;
1184        let error = require_terminal_predict_cohort(&missing)
1185            .expect_err("V2 without a cohort must fail closed");
1186        assert!(error.to_string().contains("V2 requires predict_cohort"));
1187    }
1188
1189    #[test]
1190    fn terminal_request_refuses_targetless_v2_inference_cohort() {
1191        let (_plan, _bundle, envelope) = terminal_fixture();
1192        let relations = envelope
1193            .predict_cohort
1194            .as_ref()
1195            .expect("fixture has a V2 cohort")
1196            .relations
1197            .clone();
1198        let error = PredictCohort::from_relations(
1199            PredictCohortRole::Inference,
1200            relations,
1201            Vec::new(),
1202            "c".repeat(64),
1203            None,
1204        )
1205        .expect_err("V2 predict cohorts must bind an output width");
1206        assert!(error
1207            .to_string()
1208            .contains("target_names must be a non-empty list"));
1209    }
1210
1211    #[test]
1212    fn terminal_execution_refuses_provider_cohort_substitution_before_data_access() {
1213        let (plan, bundle, envelope) = terminal_fixture();
1214        let mut substituted = envelope
1215            .predict_cohort
1216            .as_ref()
1217            .expect("fixture V2 cohort")
1218            .clone();
1219        substituted.data_content_fingerprint = "e".repeat(64);
1220        substituted.cohort_fingerprint = substituted.fingerprint().unwrap();
1221        let provider = SubstitutedCohortProvider {
1222            cohort: substituted,
1223            materialize_calls: Cell::new(0),
1224            view_calls: Cell::new(0),
1225        };
1226        let mut controllers = RuntimeControllerRegistry::new();
1227        controllers
1228            .register(Box::new(NeverInvokedController {
1229                id: ControllerId::new("controller:model").unwrap(),
1230            }))
1231            .unwrap();
1232        // Deliberately leave the artifact store empty: the provider mismatch
1233        // must be refused before replay is allowed to hydrate even REFIT
1234        // artifacts, let alone materialize scientific data.
1235        let artifact_store = InMemoryArtifactStore::new();
1236        let mut ctx = RunContext::new(RunId::new("run:terminal.substitution").unwrap(), Some(7));
1237
1238        let error = execute_terminal_prediction(
1239            TerminalPredictionReplay {
1240                plan: &plan,
1241                bundle: &bundle,
1242                envelope: &envelope,
1243                selector: &selector(),
1244                controllers: &controllers,
1245                data_provider: &provider,
1246                artifact_store: &artifact_store,
1247            },
1248            &mut ctx,
1249        )
1250        .expect_err("a provider cohort substitution must fail before data materialization");
1251        assert!(error.to_string().contains("supplied by runtime provider"));
1252        assert_eq!(provider.materialize_calls.get(), 0);
1253        assert_eq!(provider.view_calls.get(), 0);
1254    }
1255
1256    #[test]
1257    fn terminal_execution_rejects_observation_output_without_custom_aggregation() {
1258        let (plan, bundle, envelope) = terminal_fixture_with_custom_aggregation(true);
1259        let provider = InMemoryDataProvider::with_envelope(
1260            ControllerId::new("controller:data.terminal").unwrap(),
1261            envelope.clone(),
1262        )
1263        .unwrap();
1264        let aggregation_calls = Arc::new(AtomicUsize::new(0));
1265        let mut controllers = RuntimeControllerRegistry::new();
1266        controllers
1267            .register(Box::new(ObservationTerminalController {
1268                id: ControllerId::new("controller:model").unwrap(),
1269            }))
1270            .unwrap();
1271        controllers
1272            .register(Box::new(CountingAggregationController {
1273                id: ControllerId::new("controller:agg.custom").unwrap(),
1274                calls: aggregation_calls.clone(),
1275            }))
1276            .unwrap();
1277        let artifact_store = registered_terminal_artifact_store(&bundle);
1278        let mut ctx = RunContext::new(RunId::new("run:terminal.no-aggregation").unwrap(), Some(7));
1279
1280        let error = execute_terminal_prediction(
1281            TerminalPredictionReplay {
1282                plan: &plan,
1283                bundle: &bundle,
1284                envelope: &envelope,
1285                selector: &selector(),
1286                controllers: &controllers,
1287                data_provider: &provider,
1288                artifact_store: &artifact_store,
1289            },
1290            &mut ctx,
1291        )
1292        .expect_err("direct terminal replay must refuse observation output");
1293        assert!(error.to_string().contains("observation-level predictions"));
1294        assert_eq!(aggregation_calls.load(Ordering::SeqCst), 0);
1295    }
1296
1297    #[test]
1298    fn terminal_receipt_refuses_identity_mismatch_and_aggregation() {
1299        let (plan, bundle, envelope) = terminal_fixture();
1300        let mut altered = terminal_prediction(&envelope);
1301        altered.sample_ids.reverse();
1302        let error = attest_terminal_prediction_output(
1303            &plan,
1304            &bundle,
1305            &envelope,
1306            &selector(),
1307            &[altered],
1308            &[],
1309            &[],
1310        )
1311        .expect_err("reordered cohort identities must fail");
1312        assert!(error
1313            .to_string()
1314            .contains("sample identities do not exactly match"));
1315
1316        let observation = ObservationPredictionBlock {
1317            prediction_id: Some("prediction:terminal:observation".to_string()),
1318            producer_node: NodeId::new("model:terminal").unwrap(),
1319            producer_port: Some("prediction".to_string()),
1320            partition: PredictionPartition::Final,
1321            fold_id: None,
1322            observation_ids: Vec::new(),
1323            values: Vec::new(),
1324            weights: Vec::new(),
1325            target_names: vec!["protein".to_string()],
1326        };
1327        let error = attest_terminal_prediction_output(
1328            &plan,
1329            &bundle,
1330            &envelope,
1331            &selector(),
1332            &[terminal_prediction(&envelope)],
1333            &[observation],
1334            &[],
1335        )
1336        .expect_err("aggregation output must not enter this first terminal slice");
1337        assert!(error
1338            .to_string()
1339            .contains("unsupported aggregated predictions"));
1340    }
1341}
dag_ml_core/terminal_prediction.rs

dag_ml_core/
terminal_prediction.rs
