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

dag_ml_core/runtime/
mod.rs
