dag_ml_core/runtime/
mod.rs

1//! Runtime execution: schedulers, controllers, stores, OOF/merge logic.
2//!
3//! Split from the former monolithic `runtime.rs` into cohesive submodules
4//! (pure refactor — code moved verbatim). `mod.rs` owns the run context,
5//! the controller registry, the custom-aggregation dispatch entry points,
6//! native variant selection, and re-exports the full runtime surface so
7//! `pub use runtime::*` in `lib.rs` resolves identically.
8
9pub(crate) use std::cell::RefCell;
10pub(crate) use std::collections::{BTreeMap, BTreeSet};
11pub(crate) use std::fs;
12pub(crate) use std::io::Read;
13pub(crate) use std::path::{Path, PathBuf};
14
15pub(crate) use serde::{Deserialize, Serialize};
16pub(crate) use sha2::{Digest, Sha256};
17
18pub(crate) use crate::aggregation::{
19    aggregate_observation_predictions, aggregate_sample_predictions_by_unit,
20    reduce_predictions_across_branches, reduce_proba_mean_across_branches,
21    AggregatedPredictionBlock, AggregationControllerInput, AggregationControllerOutput,
22    AggregationControllerResult, AggregationControllerTask, ObservationPredictionBlock,
23    PredictionUnitId,
24};
25pub(crate) use crate::bundle::{
26    build_aggregated_prediction_cache_payload, build_prediction_cache_payload,
27    bundle_prediction_requirement_key, validate_prediction_cache_payload_matches_record,
28    BundlePredictionCachePayload, BundlePredictionCachePayloadSet, BundlePredictionCacheRecord,
29    BundlePredictionRequirement, ExecutionBundle, RefitArtifactRecord, ReplayPhaseRequest,
30};
31pub(crate) use crate::campaign::stable_json_fingerprint;
32pub(crate) use crate::controller::{capabilities_support_fit_influence, ControllerCapability};
33pub(crate) use crate::criteria::{EarlyStoppingRecord, LossExecutionAttestation};
34pub(crate) use crate::data::{
35    data_binding_requirement_key, DataBinding, DataRequestPartition, ExternalDataPlanEnvelope,
36    RepresentationCompatibilityReport, RepresentationPlan, RepresentationReplayManifest,
37};
38pub(crate) use crate::error::{DagMlError, Result};
39pub(crate) use crate::fold::{FoldAssignment, FoldPartitionMode, FoldSet};
40pub(crate) use crate::generation::{
41    enumerate_variants, GenerationChoice, OperatorVariantModel, VariantPlan,
42};
43pub(crate) use crate::graph::{EdgeSpec, NodeKind, PortKind};
44pub(crate) use crate::ids::{
45    ArtifactId, BranchId, BundleId, ControllerId, FoldId, LineageId, NodeId, RunId, SampleId,
46    VariantId,
47};
48pub(crate) use crate::metrics::{
49    cross_fold_validation_reports, reassemble_merge_targets, score_regression_aggregated_block,
50    score_regression_prediction_block, OofAverageBlock, RegressionMetricKind,
51    RegressionMetricReport, RegressionTargetBlock, RegressionTargetRecord, ScoreSet,
52    SCORE_SET_SCHEMA_VERSION,
53};
54pub(crate) use crate::oof::{
55    PredictionBlock, PredictionPartition, StackingOofRefitContract, StackingOofRefitDecision,
56    StackingOofRefitPolicy,
57};
58pub(crate) use crate::phase::Phase;
59pub(crate) use crate::plan::{prune_plan_to_active, CampaignSpec, ExecutionPlan, NodePlan};
60pub(crate) use crate::policy::{
61    AggregationPolicy, FitInfluencePolicy, PredictionLevel, ShapeDelta, ShapeDeltaKind,
62};
63pub(crate) use crate::relation::SampleRelationSet;
64pub(crate) use crate::rng::SeedContext;
65pub(crate) use crate::selection::{
66    select_candidate, CandidateScore, SelectionDecision, SelectionMetric, SelectionPolicy,
67};
68
69mod artifact;
70mod dataview;
71mod merge;
72mod methods_replay;
73mod oof;
74mod prediction_store;
75mod scheduler;
76mod scoring;
77mod task;
78
79pub use artifact::*;
80pub use dataview::*;
81pub(crate) use merge::*;
82#[cfg(feature = "methods-optimizer")]
83pub use methods_replay::*;
84pub use oof::*;
85pub use prediction_store::*;
86pub use scheduler::*;
87pub(crate) use scoring::*;
88pub use task::*;
89
90pub struct BundleReplayExecution<'a> {
91    pub plan: &'a ExecutionPlan,
92    pub bundle: &'a ExecutionBundle,
93    pub replay_request: &'a ReplayPhaseRequest,
94    pub prediction_cache_store: Option<&'a dyn RuntimePredictionCacheStore>,
95    pub controllers: &'a RuntimeControllerRegistry,
96    pub data_provider: &'a dyn RuntimeDataProvider,
97    pub artifact_store: &'a dyn RuntimeArtifactStore,
98    pub data_envelopes: &'a BTreeMap<String, ExternalDataPlanEnvelope>,
99}
100
101#[derive(Default)]
102pub struct RuntimeControllerRegistry {
103    controllers: BTreeMap<ControllerId, Box<dyn RuntimeController>>,
104}
105
106impl RuntimeControllerRegistry {
107    pub fn new() -> Self {
108        Self::default()
109    }
110
111    pub fn register(&mut self, controller: Box<dyn RuntimeController>) -> Result<()> {
112        let id = controller.controller_id().clone();
113        if self.controllers.insert(id.clone(), controller).is_some() {
114            return Err(DagMlError::RuntimeValidation(format!(
115                "duplicate runtime controller `{id}`"
116            )));
117        }
118        Ok(())
119    }
120
121    pub fn get(&self, controller_id: &ControllerId) -> Option<&dyn RuntimeController> {
122        self.controllers.get(controller_id).map(Box::as_ref)
123    }
124}
125
126pub fn dispatch_custom_observation_aggregation(
127    plan: &ExecutionPlan,
128    controllers: &RuntimeControllerRegistry,
129    task_id: impl Into<String>,
130    block: ObservationPredictionBlock,
131    relations: SampleRelationSet,
132    policy: AggregationPolicy,
133    requested_sample_order: Vec<SampleId>,
134) -> Result<PredictionBlock> {
135    let controller_id = custom_aggregation_controller_id(&policy)?;
136    ensure_aggregation_controller_capability(plan, controller_id)?;
137    let task = AggregationControllerTask {
138        schema_version: crate::aggregation::AGGREGATION_CONTROLLER_TASK_SCHEMA_VERSION,
139        task_id: task_id.into(),
140        controller_id: controller_id.clone(),
141        policy,
142        reduction_plan: None,
143        input: AggregationControllerInput::ObservationToSample {
144            block,
145            relations,
146            requested_sample_order,
147        },
148    };
149    let result = dispatch_custom_aggregation_task(controllers, &task)?;
150    match result.output {
151        AggregationControllerOutput::Sample { block } => Ok(block),
152        AggregationControllerOutput::Unit { .. } => Err(DagMlError::RuntimeValidation(format!(
153            "aggregation controller task `{}` returned unit output for observation input",
154            task.task_id
155        ))),
156    }
157}
158
159pub fn dispatch_custom_sample_aggregation(
160    plan: &ExecutionPlan,
161    controllers: &RuntimeControllerRegistry,
162    task_id: impl Into<String>,
163    block: PredictionBlock,
164    relations: SampleRelationSet,
165    policy: AggregationPolicy,
166    requested_unit_order: Vec<PredictionUnitId>,
167) -> Result<AggregatedPredictionBlock> {
168    let controller_id = custom_aggregation_controller_id(&policy)?;
169    ensure_aggregation_controller_capability(plan, controller_id)?;
170    let task = AggregationControllerTask {
171        schema_version: crate::aggregation::AGGREGATION_CONTROLLER_TASK_SCHEMA_VERSION,
172        task_id: task_id.into(),
173        controller_id: controller_id.clone(),
174        policy,
175        reduction_plan: None,
176        input: AggregationControllerInput::SampleToUnit {
177            block,
178            relations,
179            requested_unit_order,
180        },
181    };
182    let result = dispatch_custom_aggregation_task(controllers, &task)?;
183    match result.output {
184        AggregationControllerOutput::Unit { block } => Ok(block),
185        AggregationControllerOutput::Sample { .. } => Err(DagMlError::RuntimeValidation(format!(
186            "aggregation controller task `{}` returned sample output for sample input",
187            task.task_id
188        ))),
189    }
190}
191
192pub fn dispatch_custom_aggregation_task(
193    controllers: &RuntimeControllerRegistry,
194    task: &AggregationControllerTask,
195) -> Result<AggregationControllerResult> {
196    task.validate()?;
197    let controller = controllers.get(&task.controller_id).ok_or_else(|| {
198        DagMlError::RuntimeValidation(format!(
199            "aggregation runtime controller `{}` is not registered",
200            task.controller_id
201        ))
202    })?;
203    let result = controller.invoke_aggregation(task)?;
204    result.validate_for_task(task)?;
205    Ok(result)
206}
207
208pub(crate) fn custom_aggregation_controller_id(
209    policy: &AggregationPolicy,
210) -> Result<&ControllerId> {
211    policy.validate()?;
212    policy
213        .custom_controller
214        .as_ref()
215        .map(|controller| &controller.controller_id)
216        .ok_or_else(|| {
217            DagMlError::RuntimeValidation(
218                "custom aggregation dispatch requires a custom_controller policy".to_string(),
219            )
220        })
221}
222
223pub(crate) fn ensure_aggregation_controller_capability(
224    plan: &ExecutionPlan,
225    controller_id: &ControllerId,
226) -> Result<()> {
227    let manifest = plan
228        .controller_manifests
229        .get(controller_id)
230        .ok_or_else(|| {
231            DagMlError::Planning(format!(
232                "missing aggregation controller manifest `{controller_id}`"
233            ))
234        })?;
235    if !manifest
236        .capabilities
237        .contains(&ControllerCapability::AggregatesPredictions)
238    {
239        return Err(DagMlError::Planning(format!(
240            "aggregation controller `{controller_id}` must declare aggregates_predictions"
241        )));
242    }
243    Ok(())
244}
245
246#[derive(Clone, Debug)]
247pub struct RunContext {
248    pub run_id: RunId,
249    pub root_seed: Option<u64>,
250    pub variant_id: Option<VariantId>,
251    pub prediction_store: InMemoryPredictionStore,
252    pub aggregated_prediction_store: InMemoryAggregatedPredictionStore,
253    pub lineage: InMemoryLineageRecorder,
254    /// Native per-fold/per-partition score reports collected during the run (when the host emits
255    /// `regression_targets`).
256    pub score_collector: Vec<RegressionMetricReport>,
257    /// Per-fold `y_true` records, kept so cross-fold ensembles (the OOF average) can be scored.
258    pub regression_target_records: Vec<RegressionTargetRecord>,
259    /// The per-sample cross-fold OOF average blocks (+ `y_true`) collected alongside the scalar OOF
260    /// average reports — one per scored producer. Surfaced so the host can fill the `(validation, avg)`
261    /// row's per-sample y_pred; populated by `collect_cross_fold_validation_scores`, empty otherwise.
262    pub oof_average_blocks: Vec<OofAverageBlock>,
263    /// Declarative per-producer aggregation contracts that are applied only after every
264    /// validation fold has contributed its raw sample-level OOF block.  This is deliberately
265    /// separate from the per-task aggregation path: a semantic unit may span CV folds, so
266    /// reducing it inside one fold would change the experiment being scored.
267    pub(crate) global_oof_aggregation: BTreeMap<NodeId, GlobalOofAggregationSpec>,
268}
269
270#[derive(Clone, Debug)]
271pub(crate) struct GlobalOofAggregationSpec {
272    pub(crate) policy: AggregationPolicy,
273    pub(crate) relations: SampleRelationSet,
274}
275
276impl RunContext {
277    pub fn new(run_id: RunId, root_seed: Option<u64>) -> Self {
278        Self {
279            run_id,
280            root_seed,
281            variant_id: None,
282            prediction_store: InMemoryPredictionStore::new(),
283            aggregated_prediction_store: InMemoryAggregatedPredictionStore::new(),
284            lineage: InMemoryLineageRecorder::new(),
285            score_collector: Vec::new(),
286            regression_target_records: Vec::new(),
287            oof_average_blocks: Vec::new(),
288            global_oof_aggregation: BTreeMap::new(),
289        }
290    }
291
292    /// Bind the relation-attested aggregation policies needed for global OOF scoring before the
293    /// caller executes `FIT_CV`.  The provider remains the authority for relations; this method
294    /// only records a validated, immutable snapshot in the run context so finalisation cannot
295    /// silently use a different grouping from execution.
296    pub(crate) fn configure_global_oof_aggregation(
297        &mut self,
298        plan: &ExecutionPlan,
299        data_provider: &dyn RuntimeDataProvider,
300    ) -> Result<()> {
301        self.global_oof_aggregation = global_oof_aggregation_specs(plan, data_provider)?;
302        Ok(())
303    }
304
305    /// Score the cross-fold OOF average from the collected per-fold validation predictions + targets
306    /// and append the reports (one per producer, `fold_id = "avg"`) to the score collector, plus —
307    /// additively — the per-sample OOF average block + `y_true` each report was computed from to
308    /// [`oof_average_blocks`](Self::oof_average_blocks) (so the host can fill the `(validation, avg)`
309    /// row's per-sample y_pred). Call after FIT_CV; a no-op when nothing was scored or no producer has
310    /// more than one fold.
311    ///
312    /// `partition_mode` is the campaign's [`FoldPartitionMode`]: `Partition` (KFold) requires a unique
313    /// per-producer OOF set, while `Resampled` (ShuffleSplit / repeated CV) permits a sample to be
314    /// validated in multiple folds (averaged when scored). Pass the plan's
315    /// [`fold_set`](ExecutionPlan::fold_set) mode (default `Partition` when there is no fold set).
316    pub fn collect_cross_fold_validation_scores(
317        &mut self,
318        partition_mode: FoldPartitionMode,
319    ) -> Result<()> {
320        let outcome = cross_fold_validation_reports(
321            self.prediction_store.blocks(),
322            &self.regression_target_records,
323            SCORE_METRICS,
324            partition_mode,
325        )?;
326        let outcome = apply_global_oof_aggregation(outcome, &self.global_oof_aggregation)?;
327        self.score_collector.extend(outcome.reports);
328        self.oof_average_blocks.extend(outcome.oof_averages);
329        Ok(())
330    }
331
332    /// Build a [`ScoreSet`] from the collected reports (or `None` if scoring was off / produced
333    /// nothing), e.g. to attach to the [`ExecutionBundle`].
334    pub fn build_score_set(
335        &self,
336        plan_id: impl Into<String>,
337        selection_metric: Option<String>,
338    ) -> Option<ScoreSet> {
339        if self.score_collector.is_empty() {
340            return None;
341        }
342        Some(ScoreSet {
343            schema_version: SCORE_SET_SCHEMA_VERSION,
344            plan_id: plan_id.into(),
345            selection_metric,
346            reports: self.score_collector.clone(),
347        })
348    }
349}
350
351/// Outcome of native variant selection: the winning variant plus EVERY scored variant's
352/// cross-validation reports, each tagged with its own `variant_id`.
353///
354/// The reports are the per-fold + cross-fold-OOF-average VALIDATION (OOF) reports collected while
355/// ranking. They are emitted so a generated sweep can surface every variant's CV score — not only
356/// the winner's — to match the legacy per-variant `num_predictions`. These are REPORT-ONLY
357/// validation scores of non-selected models: they never feed any downstream training/feature path
358/// (no prediction blocks, no `RegressionTargetRecord`s, no handles leave selection — see
359/// [`select_best_variant_by_cv`]), so the OOF/leakage invariants are unaffected.
360#[derive(Clone, Debug)]
361pub struct VariantSelection {
362    /// The winning variant, ranked by `selection_metric`. The SELECT DECISION is identical to the
363    /// pre-existing behavior; `validation_reports` is purely additive context.
364    pub selected_variant_id: VariantId,
365    /// Per-variant VALIDATION (OOF) reports for ALL ranked variants (winner included), each tagged
366    /// with its `variant_id`. The cross-fold OOF average per producer is re-tagged with the variant
367    /// id (its native form has `variant_id = None`); the per-fold reports already carry it.
368    pub validation_reports: Vec<RegressionMetricReport>,
369    /// Per-variant VALIDATION (OOF) PREDICTIONS for ALL ranked variants (winner included), captured
370    /// from each variant's transient FIT_CV [`RunContext`] BEFORE it is dropped, re-tagged with the
371    /// variant's id + content fingerprint. The scalar [`validation_reports`](Self::validation_reports)
372    /// above carry only the score; these carry the per-sample y_pred (+ id-matched y_true) so a host
373    /// can fill a non-selected variant's per-fold prediction rows, not just its CV score.
374    ///
375    /// LEAKAGE: these are each variant's OWN validation (OOF) predictions, re-tagged with that
376    /// variant's id (which prevents cross-variant mixing). They are surfaced for host
377    /// persistence/display only — every transient CV run executes FIT_CV ONLY (no Final/Test/refit),
378    /// so by construction this carries no train/refit predictions, and the captured blocks never feed
379    /// a training/feature path or cross a `requires_oof` edge. This is strictly ADDITIVE — the same
380    /// values the scalar reports were computed from, exposed per sample — analogous to the additive
381    /// OOF-average block surfacing; no leakage validator is relaxed.
382    pub variant_validation_predictions: Vec<VariantValidationPredictions>,
383}
384
385/// Extended result of native variant selection.
386///
387/// The historical [`VariantSelection`] remains source-compatible for callers
388/// that construct or destructure it. Training orchestration uses this additive
389/// result to retain the exact [`SelectionDecision`] produced by the one and
390/// only ranking pass.
391#[derive(Clone, Debug)]
392pub struct VariantSelectionOutcome {
393    pub selection: VariantSelection,
394    pub decision: SelectionDecision,
395}
396
397/// One scored variant's VALIDATION (OOF) predictions, captured from its transient FIT_CV
398/// [`RunContext`] and re-tagged with the variant's id + content fingerprint so a host can fill that
399/// variant's per-sample prediction rows. REPORT-grade output paired with
400/// [`VariantSelection::validation_reports`]: it never feeds a training/feature path (see the field
401/// docs on [`VariantSelection::variant_validation_predictions`]).
402#[derive(Clone, Debug)]
403pub struct VariantValidationPredictions {
404    /// The variant these predictions belong to — the re-tag that keeps them from mixing with another
405    /// variant's predictions.
406    pub variant_id: VariantId,
407    /// The variant's Phase-5 content fingerprint (`variant_label`), `None` for param-variant /
408    /// single-variant SELECT (which carry no operator-variant fingerprint).
409    pub variant_label: Option<String>,
410    /// Per-fold VALIDATION (OOF) prediction blocks (`partition = Validation`), one per `(producer,
411    /// fold)`, paired POSITION-FOR-POSITION with [`regression_targets`](Self::regression_targets) (the
412    /// matching y_true for the same producer/fold/samples).
413    pub predictions: Vec<PredictionBlock>,
414    /// The id-matched y_true blocks for [`predictions`](Self::predictions), one per prediction block in
415    /// the SAME order.
416    pub regression_targets: Vec<RegressionTargetBlock>,
417    /// The per-sample cross-fold OOF AVERAGE block (+ id-matched y_true), if the variant produced one
418    /// (`None` for a single-fold splitter). The same averaged values the variant's scalar `avg` report
419    /// was computed from, exposed per sample.
420    pub oof_average: Option<OofAverageBlock>,
421}
422
423/// Pick the best variant of a multi-variant plan by its cross-validation score, natively.
424///
425/// "Option A": each variant is scored with its OWN single-variant FIT_CV — the plan is cloned with
426/// `variants = vec![variant]` so the existing per-producer cross-fold OOF averaging
427/// ([`RunContext::collect_cross_fold_validation_scores`]) is unambiguous (one variant in scope, so a
428/// validation `PredictionBlock` belongs to exactly one variant). The OOF-average report per variant
429/// becomes a [`CandidateScore`], and [`select_candidate`] ranks them by `selection_metric` (the
430/// metric's [`objective`](RegressionMetricKind::objective) drives the direction — RMSE minimizes,
431/// accuracy maximizes). The winning candidate id maps back to its [`VariantId`].
432///
433/// Beyond ranking, every scored variant's VALIDATION (OOF) reports — the per-fold reports and the
434/// cross-fold OOF average, each tagged with its `variant_id` — are accumulated and returned in
435/// [`VariantSelection::validation_reports`] so the caller can surface ALL variants' CV scores (not
436/// just the winner's) in the final bundle. This is OOF-safe: the per-variant CV runs happen in
437/// transient `RunContext`s whose prediction stores and `RegressionTargetRecord`s are dropped here;
438/// only the scalar score reports (derived from `y_true`) survive, so a non-selected variant's OOF
439/// predictions can NEVER reach any downstream training/feature path.
440///
441/// Native scoring is opt-in: it only happens when the host emits `regression_targets`. So this
442/// returns `Ok(None)` when NO variant produced a cross-fold OOF average (scoring is off, the normal
443/// case today) — the caller should then fall back to its default variant, behaving exactly as before.
444/// When EVERY variant scored, it returns `Ok(Some(best))`. A partially-scored set (some variants
445/// scored, others not) is an inconsistent host and is rejected so variants are never ranked unfairly.
446///
447/// `run_single_variant_fit_cv` runs FIT_CV for the single-variant plan into the supplied context
448/// (the caller supplies the scheduler/data-provider wiring); this keeps the selection logic free of
449/// host runtime details and unit-testable with mock controllers. Cloning a one-variant plan is
450/// valid: `node_plans`/`fold_set` are plan-level (not keyed per variant) and variant params are
451/// applied per-node at task build time, so the per-variant CV is isolated.
452pub fn select_best_variant_by_cv<F>(
453    plan: &ExecutionPlan,
454    run_id: &RunId,
455    root_seed: Option<u64>,
456    selection_metric: RegressionMetricKind,
457    run_single_variant_fit_cv: F,
458) -> Result<Option<VariantSelection>>
459where
460    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
461{
462    Ok(select_best_variant_outcome_by_cv(
463        plan,
464        run_id,
465        root_seed,
466        selection_metric,
467        run_single_variant_fit_cv,
468    )?
469    .map(|outcome| outcome.selection))
470}
471
472/// Select the best plan variant and retain the exact decision produced by the
473/// shared native ranking pass.
474///
475/// This is the training-operation counterpart of
476/// [`select_best_variant_by_cv`]. It does not perform an additional SELECT;
477/// the legacy helper simply projects this result back to its historical type.
478pub fn select_best_variant_outcome_by_cv<F>(
479    plan: &ExecutionPlan,
480    run_id: &RunId,
481    root_seed: Option<u64>,
482    selection_metric: RegressionMetricKind,
483    mut run_single_variant_fit_cv: F,
484) -> Result<Option<VariantSelectionOutcome>>
485where
486    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
487{
488    plan.validate()?;
489    if plan.variants.is_empty() {
490        return Err(DagMlError::RuntimeValidation(
491            "cannot select a variant for a plan with no variants".to_string(),
492        ));
493    }
494    // Mechanism A: each variant is the FULL union plan narrowed to that single variant — params are
495    // applied per-node at task-build time, so cloning a one-variant plan is the per-variant scope.
496    score_and_rank_variants_by_cv(
497        &plan.variants,
498        run_id,
499        root_seed,
500        selection_metric,
501        plan_oof_partition_mode(plan),
502        None,
503        |variant| {
504            Ok(ExecutionPlan {
505                variants: vec![variant.clone()],
506                ..plan.clone()
507            })
508        },
509        // Param-variant SELECT (Mechanism A) has no operator-variant content fingerprint, so reports
510        // carry `variant_id` only (no `variant_label`) — exactly the pre-Phase-5 shape.
511        |_variant| Ok(None),
512        &mut run_single_variant_fit_cv,
513    )
514}
515
516/// Select a plan variant using only the cross-fold OOF average emitted by one
517/// explicitly resolved score-target producer. All producers' validation
518/// reports remain retained in the returned outcome for audit.
519#[allow(clippy::too_many_arguments)]
520pub fn select_best_variant_outcome_by_cv_for_target<F>(
521    plan: &ExecutionPlan,
522    run_id: &RunId,
523    root_seed: Option<u64>,
524    selection_metric: RegressionMetricKind,
525    score_target: &NodeId,
526    score_target_port: Option<&str>,
527    score_target_level: PredictionLevel,
528    mut run_single_variant_fit_cv: F,
529) -> Result<Option<VariantSelectionOutcome>>
530where
531    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
532{
533    plan.validate()?;
534    if !plan.node_plans.contains_key(score_target) {
535        return Err(DagMlError::RuntimeValidation(format!(
536            "native SELECT score target `{score_target}` is absent from plan"
537        )));
538    }
539    score_and_rank_variants_by_cv(
540        &plan.variants,
541        run_id,
542        root_seed,
543        selection_metric,
544        plan_oof_partition_mode(plan),
545        Some((score_target, score_target_port, score_target_level)),
546        |variant| {
547            Ok(ExecutionPlan {
548                variants: vec![variant.clone()],
549                ..plan.clone()
550            })
551        },
552        |_variant| Ok(None),
553        &mut run_single_variant_fit_cv,
554    )
555}
556
557/// Pick the best OPERATOR variant of an operator-generator UNION plan by its cross-validation score.
558///
559/// Where [`select_best_variant_by_cv`] narrows the SAME union plan to one variant (Mechanism A: param
560/// variants), operator-SELECT scores each candidate on its PRUNED plan: the Mechanism-B union
561/// compiles an operator generator as a STACKING graph (`choice -> merge:generator_predictions ->
562/// model:meta`), but operator `_or_` is SELECT, not stacking — so each candidate is the union pruned
563/// down to one choice's sub-sequence + the shared prefix, with the generator merge + meta-model +
564/// every inactive choice ELIDED (see [`prune_plan_to_active`]). The pruned candidate has exactly ONE
565/// terminal producer, so the single-producer guard in the shared ranking loop is satisfied.
566///
567/// `model` is the [`OperatorVariantModel`] lowered from the (single, flat) operator generator;
568/// `union_plan` is the compiled UNION plan; `selection_metric` drives the ranking direction
569/// (`RegressionMetricKind::objective`). MULTIPLE operator generators are REJECTED here (consistent
570/// with the Phase-3 nested-rejection: this phase scopes to a flat single operator generator).
571///
572/// LEAKAGE: each variant runs in a fresh, variant-pinned [`RunContext`] over its PRUNED graph — the
573/// inactive choices' models are physically absent, so they are never fit and no `requires_oof` edge
574/// can pull an inactive variant's OOF. The non-selected variants' OOF predictions never leave their
575/// transient contexts (only their scalar VALIDATION reports survive), exactly as in
576/// [`select_best_variant_by_cv`].
577///
578/// Returns `Ok(None)` when scoring is off (no host targets) — the caller keeps its default — and
579/// `Ok(Some(best))` when every variant scored.
580pub fn select_best_operator_variant_by_cv<F>(
581    union_plan: &ExecutionPlan,
582    model: &OperatorVariantModel,
583    run_id: &RunId,
584    root_seed: Option<u64>,
585    selection_metric: RegressionMetricKind,
586    mut run_single_variant_fit_cv: F,
587) -> Result<Option<VariantSelection>>
588where
589    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
590{
591    union_plan.validate()?;
592    model.validate()?;
593    let variants = enumerate_variants(&model.generation_spec(), root_seed)?;
594    if variants.is_empty() {
595        return Err(DagMlError::RuntimeValidation(format!(
596            "operator variant model `{}` produced no variants",
597            model.generator_id
598        )));
599    }
600    // The union of every choice's active set: subtracted from each candidate's ancestors so a prune
601    // never pulls in a sibling choice (or the elided merge/meta).
602    let all_choice_nodes = model
603        .active_nodes
604        .values()
605        .flatten()
606        .cloned()
607        .collect::<BTreeSet<NodeId>>();
608    // Map each enumerated variant back to its operator choice (the choice's `active_subsequence`
609    // keys `active_nodes`) via `operator_variant_active_subsequence`. The model is a single operator
610    // dimension, so each variant carries exactly one choice.
611    Ok(score_and_rank_variants_by_cv(
612        &variants,
613        run_id,
614        root_seed,
615        selection_metric,
616        plan_oof_partition_mode(union_plan),
617        None,
618        |variant| {
619            let active_subsequence = operator_variant_active_subsequence(model, variant)?;
620            let active_nodes = model.active_nodes.get(active_subsequence).ok_or_else(|| {
621                DagMlError::RuntimeValidation(format!(
622                    "operator variant model `{}` has no active-node set for `{active_subsequence}`",
623                    model.generator_id
624                ))
625            })?;
626            prune_plan_to_active(union_plan, active_nodes, &all_choice_nodes, variant)
627        },
628        // Phase 5: stamp the choice's cross-language content fingerprint on every report. The
629        // operator model's `variant_labels` is the choice-keyed sha256; when a model was hand-built
630        // without labels (the older execution fixtures), the map is empty and reports carry no label.
631        |variant| {
632            let active_subsequence = operator_variant_active_subsequence(model, variant)?;
633            Ok(model.variant_labels.get(active_subsequence).cloned())
634        },
635        &mut run_single_variant_fit_cv,
636    )?
637    .map(|outcome| outcome.selection))
638}
639
640/// Resolve the `active_subsequence` (choice key) of an enumerated operator variant against its
641/// model's single operator dimension. Shared by the prune-plan and the `variant_label` resolvers so
642/// both agree on the choice a variant names.
643fn operator_variant_active_subsequence<'a>(
644    model: &OperatorVariantModel,
645    variant: &'a VariantPlan,
646) -> Result<&'a str> {
647    let dimension_name = &model.dimension.name;
648    let choice = variant.choices.get(dimension_name).ok_or_else(|| {
649        DagMlError::RuntimeValidation(format!(
650            "operator variant `{}` is missing the operator dimension `{dimension_name}`",
651            variant.variant_id
652        ))
653    })?;
654    choice.active_subsequence.as_deref().ok_or_else(|| {
655        DagMlError::RuntimeValidation(format!(
656            "operator variant `{}` choice `{}` has no active_subsequence",
657            variant.variant_id, choice.label
658        ))
659    })
660}
661
662/// Route operator-SELECT from the operator-variant models lowered off a pipeline DSL
663/// ([`compile_operator_variant_models`](crate::compile_operator_variant_models)).
664///
665/// This phase scopes to a FLAT, SINGLE operator generator (consistent with the Phase-3
666/// nested-generator rejection), so MORE THAN ONE operator generator is rejected with a clear error.
667/// An empty slice means the spec has no operator generator at all — there is nothing to operator-SELECT,
668/// so it returns `Ok(None)` (the caller keeps its default variant). Exactly one model delegates to
669/// [`select_best_operator_variant_by_cv`].
670pub fn select_best_operator_variant_from_models<F>(
671    union_plan: &ExecutionPlan,
672    models: &[OperatorVariantModel],
673    run_id: &RunId,
674    root_seed: Option<u64>,
675    selection_metric: RegressionMetricKind,
676    run_single_variant_fit_cv: F,
677) -> Result<Option<VariantSelection>>
678where
679    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
680{
681    match models {
682        [] => Ok(None),
683        [model] => select_best_operator_variant_by_cv(
684            union_plan,
685            model,
686            run_id,
687            root_seed,
688            selection_metric,
689            run_single_variant_fit_cv,
690        ),
691        _ => Err(DagMlError::RuntimeValidation(format!(
692            "operator-SELECT does not support {} operator generators in one pipeline; this phase scopes to a flat single operator generator (generators: {})",
693            models.len(),
694            models
695                .iter()
696                .map(|model| model.generator_id.to_string())
697                .collect::<Vec<_>>()
698                .join(", ")
699        ))),
700    }
701}
702
703/// The shared scoring + ranking loop behind [`select_best_variant_by_cv`] and
704/// [`select_best_operator_variant_by_cv`]: per variant, build its per-variant plan (`make_variant_plan`),
705/// run FIT_CV into a fresh variant-pinned [`RunContext`], collect the cross-fold OOF average, and
706/// rank by `selection_metric`. The two callers differ ONLY in `make_variant_plan` (clone-the-union
707/// vs. prune-to-active); everything below — the single-producer guard, the all-or-nothing scoring
708/// gate, the loser-report retention, and [`select_candidate`] ranking — is identical and lives here.
709/// `resolve_variant_label` resolves each variant's Phase-5 content fingerprint (the two closures
710/// keep the shared loop free of caller-specific plumbing).
711#[allow(clippy::too_many_arguments)]
712fn score_and_rank_variants_by_cv<M, L, F>(
713    variants: &[VariantPlan],
714    run_id: &RunId,
715    root_seed: Option<u64>,
716    selection_metric: RegressionMetricKind,
717    partition_mode: FoldPartitionMode,
718    score_target: Option<(&NodeId, Option<&str>, PredictionLevel)>,
719    mut make_variant_plan: M,
720    mut resolve_variant_label: L,
721    run_single_variant_fit_cv: &mut F,
722) -> Result<Option<VariantSelectionOutcome>>
723where
724    M: FnMut(&VariantPlan) -> Result<ExecutionPlan>,
725    L: FnMut(&VariantPlan) -> Result<Option<String>>,
726    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
727{
728    if variants.is_empty() {
729        return Err(DagMlError::RuntimeValidation(
730            "cannot select a variant for a plan with no variants".to_string(),
731        ));
732    }
733
734    let mut candidates: Vec<CandidateScore> = Vec::with_capacity(variants.len());
735    // Every ranked variant's VALIDATION (OOF) reports, each tagged with its variant_id, accumulated
736    // so the caller can emit ALL variants' CV scores (not just the winner's) in the bundle.
737    let mut variant_validation_reports: Vec<RegressionMetricReport> = Vec::new();
738    // Every ranked variant's VALIDATION (OOF) PREDICTIONS, captured from its transient ctx and
739    // re-tagged with its variant id + content fingerprint, so the caller can fill a non-selected
740    // variant's per-sample prediction rows (not just its scalar CV score). Captured per variant; the
741    // caller filters to the LOSERS (the winner's predictions come fresh from the real FIT_CV pass).
742    let mut variant_validation_predictions: Vec<VariantValidationPredictions> = Vec::new();
743    // Tracks whether ANY variant emitted scores at all (host targets present), so an empty candidate
744    // set can be told apart from "scoring genuinely off" (no targets) — see the post-loop branch.
745    let mut any_scores_seen = false;
746    for variant in variants {
747        let variant_plan = make_variant_plan(variant)?;
748        // Phase 5: the operator-variant content fingerprint for this variant (the choice's
749        // `variant_label`), resolved the SAME way `variant_id` is — `None` for param-variant /
750        // single-variant SELECT, `Some(<sha256>)` for an operator choice.
751        let variant_label = resolve_variant_label(variant)?;
752        let mut ctx = RunContext::new(run_id.clone(), root_seed);
753        ctx.variant_id = Some(variant.variant_id.clone());
754        run_single_variant_fit_cv(&variant_plan, &mut ctx)?;
755        ctx.collect_cross_fold_validation_scores(partition_mode)?;
756        if !ctx.score_collector.is_empty() {
757            any_scores_seen = true;
758        }
759        // ADDITIVE prediction capture (paired with the scalar report retention below). Each per-fold
760        // VALIDATION (OOF) `PredictionBlock` in this variant's transient store is captured together
761        // with its id-matched y_true, plus the cross-fold OOF AVERAGE block — re-tagged with the
762        // variant's id + content fingerprint. Only `Validation` blocks are captured: the transient run
763        // executes FIT_CV ONLY (no Final/Test/refit), so this is OOF-only by construction, and the
764        // re-tag prevents cross-variant mixing. The same values the scalar reports were computed from,
765        // exposed per sample — strictly additive (the captured blocks never feed a training/feature
766        // path or cross a `requires_oof` edge).
767        let captured = capture_variant_validation_predictions(
768            &variant.variant_id,
769            variant_label.clone(),
770            &ctx,
771        );
772        if !captured.predictions.is_empty() || captured.oof_average.is_some() {
773            variant_validation_predictions.push(captured);
774        }
775        // `cross_fold_validation_reports` emits one cross-fold OOF average PER producer. Native SELECT
776        // ranks a variant by a single score, so a multi-producer DAG is ambiguous and refused rather
777        // than silently ranked on whichever producer happened to be first (an explicit score-target
778        // producer is a future extension). For operator-SELECT the pruned candidate has exactly one
779        // terminal producer, so this guard is satisfied by construction.
780        let avg_reports = ctx
781            .score_collector
782            .iter()
783            .filter(|report| {
784                report.partition == PredictionPartition::Validation
785                    && score_target.is_none_or(|(target, target_port, level)| {
786                        &report.producer_node == target
787                            && target_port
788                                .is_none_or(|port| report.producer_port.as_deref() == Some(port))
789                            && report.level == level
790                    })
791                    && report
792                        .fold_id
793                        .as_ref()
794                        .is_some_and(|fold| fold.as_str() == "avg")
795            })
796            .collect::<Vec<_>>();
797        match avg_reports.as_slice() {
798            [] => {}
799            [report] => candidates.push(
800                (*report)
801                    .clone()
802                    .into_candidate_score(variant.variant_id.as_str())?,
803            ),
804            _ => {
805                return Err(DagMlError::RuntimeValidation(format!(
806                    "variant `{}` produced {} cross-fold OOF averages (multiple prediction producers); native SELECT needs a single score target",
807                    variant.variant_id,
808                    avg_reports.len()
809                )));
810            }
811        }
812        // Retain this variant's VALIDATION reports (per-fold + cross-fold avg) tagged with its own
813        // variant_id. The avg report's native form has `variant_id = None`, so stamp it here; the
814        // per-fold reports already carry it from `apply_result_scoring`. Only Validation reports are
815        // kept — the transient CV runs FIT_CV only (no Final/Test), so this is OOF-only by
816        // construction, but the filter makes the report-only guarantee explicit.
817        for mut report in ctx.score_collector {
818            if report.partition != PredictionPartition::Validation {
819                continue;
820            }
821            report.variant_id = Some(variant.variant_id.clone());
822            report.variant_label = variant_label.clone();
823            variant_validation_reports.push(report);
824        }
825    }
826
827    if candidates.is_empty() {
828        if any_scores_seen {
829            // Targets WERE emitted, but no producer yielded a cross-fold average (e.g. a single fold,
830            // where the average is skipped). We cannot rank — surface it instead of falling back.
831            return Err(DagMlError::RuntimeValidation(
832                "variants produced scores but no cross-fold OOF average; cannot rank — need >=2 folds or an explicit score target".to_string(),
833            ));
834        }
835        // Native scoring is genuinely off (no host targets) — let the caller keep its default variant.
836        return Ok(None);
837    }
838    if candidates.len() != variants.len() {
839        return Err(DagMlError::RuntimeValidation(format!(
840            "native variant SELECT scored only {} of {} variants; cannot rank variants fairly",
841            candidates.len(),
842            variants.len()
843        )));
844    }
845
846    let policy = SelectionPolicy {
847        id: format!("select:variant:{}", selection_metric.name()),
848        metric: SelectionMetric {
849            name: selection_metric.name().to_string(),
850            objective: selection_metric.objective(),
851        },
852        required_metric_level: None,
853        require_finite: true,
854        evaluation_scope: None,
855        refit_slot_plan: None,
856        stacking_fit_contract: None,
857        reduction_id: None,
858    };
859    let decision = select_candidate(&policy, &candidates)?;
860    let selected_variant_id =
861        VariantId::new(decision.selected_candidate_id.clone()).map_err(|error| {
862            DagMlError::RuntimeValidation(format!("selected variant id is invalid: {error}"))
863        })?;
864    Ok(Some(VariantSelectionOutcome {
865        selection: VariantSelection {
866            selected_variant_id,
867            validation_reports: variant_validation_reports,
868            variant_validation_predictions,
869        },
870        decision,
871    }))
872}
873
874/// Capture one variant's per-fold VALIDATION (OOF) predictions (paired with id-matched y_true) and
875/// its cross-fold OOF AVERAGE block from a transient FIT_CV [`RunContext`], re-tagged with the
876/// variant's id + content fingerprint. ADDITIVE + leakage-safe: only `Validation` blocks are read (a
877/// transient run is FIT_CV-only, so no Final/Test/refit block exists), and the captured blocks are
878/// copies surfaced for host display — they never feed a training/feature path. The per-fold y_true is
879/// the same record `apply_result_scoring` retained for the score, found by `(producer, fold)`; a
880/// prediction with no matching record is skipped (it could not have been scored either).
881///
882/// The matched target record covers exactly the prediction block's SAMPLE SET (see
883/// `sample_targets_match_block`) but its rows may be in a DIFFERENT ORDER than `block.sample_ids` — a
884/// host controller may validly emit its `regression_targets` in any order. The scoring path realigns
885/// by unit id, but the host surfaces these blocks POSITIONALLY (y_pred from `block.sample_ids`/`values`
886/// paired row-for-row with `regression_targets.values`), so the y_true is REBUILT in `block.sample_ids`
887/// order here — exactly as [`oof_average_block`](crate::metrics) does for the avg — so a host pairs
888/// y_pred ↔ y_true per sample without re-sorting.
889pub(crate) fn capture_variant_validation_predictions(
890    variant_id: &VariantId,
891    variant_label: Option<String>,
892    ctx: &RunContext,
893) -> VariantValidationPredictions {
894    let mut predictions = Vec::new();
895    let mut regression_targets = Vec::new();
896    for block in ctx.prediction_store.blocks() {
897        if block.partition != PredictionPartition::Validation {
898            continue;
899        }
900        let Some(record) = ctx.regression_target_records.iter().find(|record| {
901            record.producer_node == block.producer_node
902                && record.producer_port == block.producer_port
903                && record.partition == PredictionPartition::Validation
904                && record.fold_id == block.fold_id
905        }) else {
906            continue;
907        };
908        predictions.push(block.clone());
909        regression_targets.push(target_block_aligned_to_samples(
910            &block.sample_ids,
911            &record.block,
912        ));
913    }
914    VariantValidationPredictions {
915        variant_id: variant_id.clone(),
916        variant_label,
917        predictions,
918        regression_targets,
919        // A target/group global OOF aggregate is appended after its source sample average.  Prefer
920        // that semantic score surface when present; legacy/sample-only executions retain the
921        // historical first average unchanged.
922        oof_average: ctx
923            .oof_average_blocks
924            .iter()
925            .rev()
926            .find(|block| block.predictions.level != PredictionLevel::Sample)
927            .cloned()
928            .or_else(|| ctx.oof_average_blocks.first().cloned()),
929    }
930}
931
932/// Rebuild a per-fold VALIDATION `y_true` block in `sample_ids` ORDER so a host can pair it
933/// POSITIONALLY with the prediction block's `values` (the host surfaces direct prediction/target pairs
934/// by row position, not by id). `targets` covers exactly the same SAMPLE SET as `sample_ids` (the
935/// `sample_targets_match_block` precondition under which this record was retained), so every sample has
936/// a row; a missing one would indicate a broken invariant, so the original block is returned unchanged
937/// rather than dropping rows. Mirrors the avg realignment in [`oof_average_block`](crate::metrics).
938fn target_block_aligned_to_samples(
939    sample_ids: &[SampleId],
940    targets: &RegressionTargetBlock,
941) -> RegressionTargetBlock {
942    let value_by_sample: BTreeMap<&SampleId, &Vec<f64>> = targets
943        .unit_ids
944        .iter()
945        .zip(&targets.values)
946        .filter_map(|(unit_id, row)| match unit_id {
947            PredictionUnitId::Sample(sample_id) => Some((sample_id, row)),
948            _ => None,
949        })
950        .collect();
951    if sample_ids
952        .iter()
953        .any(|sample_id| !value_by_sample.contains_key(sample_id))
954    {
955        return targets.clone();
956    }
957    RegressionTargetBlock {
958        level: PredictionLevel::Sample,
959        unit_ids: sample_ids
960            .iter()
961            .cloned()
962            .map(PredictionUnitId::Sample)
963            .collect(),
964        values: sample_ids
965            .iter()
966            .map(|sample_id| value_by_sample[sample_id].clone())
967            .collect(),
968        target_names: targets.target_names.clone(),
969    }
970}
971
972#[cfg(test)]
973mod explain_contract_tests {
974    use super::*;
975
976    fn block(method: &str) -> ExplanationBlock {
977        ExplanationBlock {
978            producer_node: NodeId::new("model:base").unwrap(),
979            producer_port: None,
980            method: method.to_string(),
981            target_name: Some("y".to_string()),
982            payload: serde_json::json!({"feature_importance": [0.5, 0.3, 0.2]}),
983        }
984    }
985
986    #[test]
987    fn validates_well_formed_explanation() {
988        assert!(block("shap").validate().is_ok());
989    }
990
991    #[test]
992    fn rejects_empty_method() {
993        assert!(block("  ").validate().is_err());
994    }
995
996    #[test]
997    fn rejects_empty_target_name() {
998        let mut b = block("shap");
999        b.target_name = Some(String::new());
1000        assert!(b.validate().is_err());
1001    }
1002
1003    #[test]
1004    fn round_trips_through_json() {
1005        let b = block("permutation_importance");
1006        let json = serde_json::to_string(&b).expect("serialize");
1007        let parsed: ExplanationBlock = serde_json::from_str(&json).expect("deserialize");
1008        assert_eq!(parsed, b);
1009        // `target_name` is omitted when absent.
1010        let mut without = block("shap");
1011        without.target_name = None;
1012        let json = serde_json::to_string(&without).expect("serialize");
1013        assert!(!json.contains("target_name"));
1014    }
1015}
1016
1017#[cfg(test)]
1018mod tests;
dag_ml_core/runtime/mod.rs

dag_ml_core/runtime/
mod.rs
