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

dag_ml_core/runtime/
mod.rs
