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, 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.
490pub fn select_best_variant_outcome_by_cv_for_target<F>(
491    plan: &ExecutionPlan,
492    run_id: &RunId,
493    root_seed: Option<u64>,
494    selection_metric: RegressionMetricKind,
495    score_target: (&NodeId, Option<&str>, PredictionLevel),
496    mut run_single_variant_fit_cv: F,
497) -> Result<Option<VariantSelectionOutcome>>
498where
499    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
500{
501    let (score_target, score_target_port, score_target_level) = score_target;
502    plan.validate()?;
503    if !plan.node_plans.contains_key(score_target) {
504        return Err(DagMlError::RuntimeValidation(format!(
505            "native SELECT score target `{score_target}` is absent from plan"
506        )));
507    }
508    score_and_rank_variants_by_cv(
509        &plan.variants,
510        run_id,
511        root_seed,
512        selection_metric,
513        plan_oof_partition_mode(plan),
514        Some((score_target, score_target_port, score_target_level)),
515        |variant| {
516            Ok(ExecutionPlan {
517                variants: vec![variant.clone()],
518                ..plan.clone()
519            })
520        },
521        |_variant| Ok(None),
522        &mut run_single_variant_fit_cv,
523    )
524}
525
526/// Pick the best OPERATOR variant of an operator-generator UNION plan by its cross-validation score.
527///
528/// Where [`select_best_variant_by_cv`] narrows the SAME union plan to one variant (Mechanism A: param
529/// variants), operator-SELECT scores each candidate on its PRUNED plan: the Mechanism-B union
530/// compiles an operator generator as a STACKING graph (`choice -> merge:generator_predictions ->
531/// model:meta`), but operator `_or_` is SELECT, not stacking — so each candidate is the union pruned
532/// down to one choice's sub-sequence + the shared prefix, with the generator merge + meta-model +
533/// every inactive choice ELIDED (see [`prune_plan_to_active`]). The pruned candidate has exactly ONE
534/// terminal producer, so the single-producer guard in the shared ranking loop is satisfied.
535///
536/// `model` is the [`OperatorVariantModel`] lowered from the (single, flat) operator generator;
537/// `union_plan` is the compiled UNION plan; `selection_metric` drives the ranking direction
538/// (`RegressionMetricKind::objective`). MULTIPLE operator generators are REJECTED here (consistent
539/// with the Phase-3 nested-rejection: this phase scopes to a flat single operator generator).
540///
541/// LEAKAGE: each variant runs in a fresh, variant-pinned [`RunContext`] over its PRUNED graph — the
542/// inactive choices' models are physically absent, so they are never fit and no `requires_oof` edge
543/// can pull an inactive variant's OOF. The non-selected variants' OOF predictions never leave their
544/// transient contexts (only their scalar VALIDATION reports survive), exactly as in
545/// [`select_best_variant_by_cv`].
546///
547/// Returns `Ok(None)` when scoring is off (no host targets) — the caller keeps its default — and
548/// `Ok(Some(best))` when every variant scored.
549pub fn select_best_operator_variant_by_cv<F>(
550    union_plan: &ExecutionPlan,
551    model: &OperatorVariantModel,
552    run_id: &RunId,
553    root_seed: Option<u64>,
554    selection_metric: RegressionMetricKind,
555    mut run_single_variant_fit_cv: F,
556) -> Result<Option<VariantSelection>>
557where
558    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
559{
560    union_plan.validate()?;
561    model.validate()?;
562    let variants = enumerate_variants(&model.generation_spec(), root_seed)?;
563    if variants.is_empty() {
564        return Err(DagMlError::RuntimeValidation(format!(
565            "operator variant model `{}` produced no variants",
566            model.generator_id
567        )));
568    }
569    // The union of every choice's active set: subtracted from each candidate's ancestors so a prune
570    // never pulls in a sibling choice (or the elided merge/meta).
571    let all_choice_nodes = model
572        .active_nodes
573        .values()
574        .flatten()
575        .cloned()
576        .collect::<BTreeSet<NodeId>>();
577    // Map each enumerated variant back to its operator choice (the choice's `active_subsequence`
578    // keys `active_nodes`) via `operator_variant_active_subsequence`. The model is a single operator
579    // dimension, so each variant carries exactly one choice.
580    Ok(score_and_rank_variants_by_cv(
581        &variants,
582        run_id,
583        root_seed,
584        selection_metric,
585        plan_oof_partition_mode(union_plan),
586        None,
587        |variant| {
588            let active_subsequence = operator_variant_active_subsequence(model, variant)?;
589            let active_nodes = model.active_nodes.get(active_subsequence).ok_or_else(|| {
590                DagMlError::RuntimeValidation(format!(
591                    "operator variant model `{}` has no active-node set for `{active_subsequence}`",
592                    model.generator_id
593                ))
594            })?;
595            prune_plan_to_active(union_plan, active_nodes, &all_choice_nodes, variant)
596        },
597        // Phase 5: stamp the choice's cross-language content fingerprint on every report. The
598        // operator model's `variant_labels` is the choice-keyed sha256; when a model was hand-built
599        // without labels (the older execution fixtures), the map is empty and reports carry no label.
600        |variant| {
601            let active_subsequence = operator_variant_active_subsequence(model, variant)?;
602            Ok(model.variant_labels.get(active_subsequence).cloned())
603        },
604        &mut run_single_variant_fit_cv,
605    )?
606    .map(|outcome| outcome.selection))
607}
608
609/// Resolve the `active_subsequence` (choice key) of an enumerated operator variant against its
610/// model's single operator dimension. Shared by the prune-plan and the `variant_label` resolvers so
611/// both agree on the choice a variant names.
612fn operator_variant_active_subsequence<'a>(
613    model: &OperatorVariantModel,
614    variant: &'a VariantPlan,
615) -> Result<&'a str> {
616    let dimension_name = &model.dimension.name;
617    let choice = variant.choices.get(dimension_name).ok_or_else(|| {
618        DagMlError::RuntimeValidation(format!(
619            "operator variant `{}` is missing the operator dimension `{dimension_name}`",
620            variant.variant_id
621        ))
622    })?;
623    choice.active_subsequence.as_deref().ok_or_else(|| {
624        DagMlError::RuntimeValidation(format!(
625            "operator variant `{}` choice `{}` has no active_subsequence",
626            variant.variant_id, choice.label
627        ))
628    })
629}
630
631/// Route operator-SELECT from the operator-variant models lowered off a pipeline DSL
632/// ([`compile_operator_variant_models`](crate::compile_operator_variant_models)).
633///
634/// This phase scopes to a FLAT, SINGLE operator generator (consistent with the Phase-3
635/// nested-generator rejection), so MORE THAN ONE operator generator is rejected with a clear error.
636/// An empty slice means the spec has no operator generator at all — there is nothing to operator-SELECT,
637/// so it returns `Ok(None)` (the caller keeps its default variant). Exactly one model delegates to
638/// [`select_best_operator_variant_by_cv`].
639pub fn select_best_operator_variant_from_models<F>(
640    union_plan: &ExecutionPlan,
641    models: &[OperatorVariantModel],
642    run_id: &RunId,
643    root_seed: Option<u64>,
644    selection_metric: RegressionMetricKind,
645    run_single_variant_fit_cv: F,
646) -> Result<Option<VariantSelection>>
647where
648    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
649{
650    match models {
651        [] => Ok(None),
652        [model] => select_best_operator_variant_by_cv(
653            union_plan,
654            model,
655            run_id,
656            root_seed,
657            selection_metric,
658            run_single_variant_fit_cv,
659        ),
660        _ => Err(DagMlError::RuntimeValidation(format!(
661            "operator-SELECT does not support {} operator generators in one pipeline; this phase scopes to a flat single operator generator (generators: {})",
662            models.len(),
663            models
664                .iter()
665                .map(|model| model.generator_id.to_string())
666                .collect::<Vec<_>>()
667                .join(", ")
668        ))),
669    }
670}
671
672/// The shared scoring + ranking loop behind [`select_best_variant_by_cv`] and
673/// [`select_best_operator_variant_by_cv`]: per variant, build its per-variant plan (`make_variant_plan`),
674/// run FIT_CV into a fresh variant-pinned [`RunContext`], collect the cross-fold OOF average, and
675/// rank by `selection_metric`. The two callers differ ONLY in `make_variant_plan` (clone-the-union
676/// vs. prune-to-active); everything below — the single-producer guard, the all-or-nothing scoring
677/// gate, the loser-report retention, and [`select_candidate`] ranking — is identical and lives here.
678/// `resolve_variant_label` resolves each variant's Phase-5 content fingerprint (the two closures
679/// keep the shared loop free of caller-specific plumbing).
680#[allow(clippy::too_many_arguments)]
681fn score_and_rank_variants_by_cv<M, L, F>(
682    variants: &[VariantPlan],
683    run_id: &RunId,
684    root_seed: Option<u64>,
685    selection_metric: RegressionMetricKind,
686    partition_mode: FoldPartitionMode,
687    score_target: Option<(&NodeId, Option<&str>, PredictionLevel)>,
688    mut make_variant_plan: M,
689    mut resolve_variant_label: L,
690    run_single_variant_fit_cv: &mut F,
691) -> Result<Option<VariantSelectionOutcome>>
692where
693    M: FnMut(&VariantPlan) -> Result<ExecutionPlan>,
694    L: FnMut(&VariantPlan) -> Result<Option<String>>,
695    F: FnMut(&ExecutionPlan, &mut RunContext) -> Result<()>,
696{
697    if variants.is_empty() {
698        return Err(DagMlError::RuntimeValidation(
699            "cannot select a variant for a plan with no variants".to_string(),
700        ));
701    }
702
703    let mut candidates: Vec<CandidateScore> = Vec::with_capacity(variants.len());
704    // Every ranked variant's VALIDATION (OOF) reports, each tagged with its variant_id, accumulated
705    // so the caller can emit ALL variants' CV scores (not just the winner's) in the bundle.
706    let mut variant_validation_reports: Vec<RegressionMetricReport> = Vec::new();
707    // Every ranked variant's VALIDATION (OOF) PREDICTIONS, captured from its transient ctx and
708    // re-tagged with its variant id + content fingerprint, so the caller can fill a non-selected
709    // variant's per-sample prediction rows (not just its scalar CV score). Captured per variant; the
710    // caller filters to the LOSERS (the winner's predictions come fresh from the real FIT_CV pass).
711    let mut variant_validation_predictions: Vec<VariantValidationPredictions> = Vec::new();
712    // Tracks whether ANY variant emitted scores at all (host targets present), so an empty candidate
713    // set can be told apart from "scoring genuinely off" (no targets) — see the post-loop branch.
714    let mut any_scores_seen = false;
715    for variant in variants {
716        let variant_plan = make_variant_plan(variant)?;
717        // Phase 5: the operator-variant content fingerprint for this variant (the choice's
718        // `variant_label`), resolved the SAME way `variant_id` is — `None` for param-variant /
719        // single-variant SELECT, `Some(<sha256>)` for an operator choice.
720        let variant_label = resolve_variant_label(variant)?;
721        let mut ctx = RunContext::new(run_id.clone(), root_seed);
722        ctx.variant_id = Some(variant.variant_id.clone());
723        run_single_variant_fit_cv(&variant_plan, &mut ctx)?;
724        ctx.collect_cross_fold_validation_scores(partition_mode)?;
725        if !ctx.score_collector.is_empty() {
726            any_scores_seen = true;
727        }
728        // ADDITIVE prediction capture (paired with the scalar report retention below). Each per-fold
729        // VALIDATION (OOF) `PredictionBlock` in this variant's transient store is captured together
730        // with its id-matched y_true, plus the cross-fold OOF AVERAGE block — re-tagged with the
731        // variant's id + content fingerprint. Only `Validation` blocks are captured: the transient run
732        // executes FIT_CV ONLY (no Final/Test/refit), so this is OOF-only by construction, and the
733        // re-tag prevents cross-variant mixing. The same values the scalar reports were computed from,
734        // exposed per sample — strictly additive (the captured blocks never feed a training/feature
735        // path or cross a `requires_oof` edge).
736        let captured = capture_variant_validation_predictions(
737            &variant.variant_id,
738            variant_label.clone(),
739            &ctx,
740        );
741        if !captured.predictions.is_empty() || captured.oof_average.is_some() {
742            variant_validation_predictions.push(captured);
743        }
744        // `cross_fold_validation_reports` emits one cross-fold OOF average PER producer. Native SELECT
745        // ranks a variant by a single score, so a multi-producer DAG is ambiguous and refused rather
746        // than silently ranked on whichever producer happened to be first (an explicit score-target
747        // producer is a future extension). For operator-SELECT the pruned candidate has exactly one
748        // terminal producer, so this guard is satisfied by construction.
749        let avg_reports = ctx
750            .score_collector
751            .iter()
752            .filter(|report| {
753                report.partition == PredictionPartition::Validation
754                    && score_target.is_none_or(|(target, target_port, level)| {
755                        &report.producer_node == target
756                            && target_port
757                                .is_none_or(|port| report.producer_port.as_deref() == Some(port))
758                            && report.level == level
759                    })
760                    && report
761                        .fold_id
762                        .as_ref()
763                        .is_some_and(|fold| fold.as_str() == "avg")
764            })
765            .collect::<Vec<_>>();
766        match avg_reports.as_slice() {
767            [] => {}
768            [report] => candidates.push(
769                (*report)
770                    .clone()
771                    .into_candidate_score(variant.variant_id.as_str())?,
772            ),
773            _ => {
774                return Err(DagMlError::RuntimeValidation(format!(
775                    "variant `{}` produced {} cross-fold OOF averages (multiple prediction producers); native SELECT needs a single score target",
776                    variant.variant_id,
777                    avg_reports.len()
778                )));
779            }
780        }
781        // Retain this variant's VALIDATION reports (per-fold + cross-fold avg) tagged with its own
782        // variant_id. The avg report's native form has `variant_id = None`, so stamp it here; the
783        // per-fold reports already carry it from `apply_result_scoring`. Only Validation reports are
784        // kept — the transient CV runs FIT_CV only (no Final/Test), so this is OOF-only by
785        // construction, but the filter makes the report-only guarantee explicit.
786        for mut report in ctx.score_collector {
787            if report.partition != PredictionPartition::Validation {
788                continue;
789            }
790            report.variant_id = Some(variant.variant_id.clone());
791            report.variant_label = variant_label.clone();
792            variant_validation_reports.push(report);
793        }
794    }
795
796    if candidates.is_empty() {
797        if any_scores_seen {
798            // Targets WERE emitted, but no producer yielded a cross-fold average (e.g. a single fold,
799            // where the average is skipped). We cannot rank — surface it instead of falling back.
800            return Err(DagMlError::RuntimeValidation(
801                "variants produced scores but no cross-fold OOF average; cannot rank — need >=2 folds or an explicit score target".to_string(),
802            ));
803        }
804        // Native scoring is genuinely off (no host targets) — let the caller keep its default variant.
805        return Ok(None);
806    }
807    if candidates.len() != variants.len() {
808        return Err(DagMlError::RuntimeValidation(format!(
809            "native variant SELECT scored only {} of {} variants; cannot rank variants fairly",
810            candidates.len(),
811            variants.len()
812        )));
813    }
814
815    let policy = SelectionPolicy {
816        id: format!("select:variant:{}", selection_metric.name()),
817        metric: SelectionMetric {
818            name: selection_metric.name().to_string(),
819            objective: selection_metric.objective(),
820        },
821        required_metric_level: None,
822        require_finite: true,
823        evaluation_scope: None,
824        refit_slot_plan: None,
825        stacking_fit_contract: None,
826        reduction_id: None,
827    };
828    let decision = select_candidate(&policy, &candidates)?;
829    let selected_variant_id =
830        VariantId::new(decision.selected_candidate_id.clone()).map_err(|error| {
831            DagMlError::RuntimeValidation(format!("selected variant id is invalid: {error}"))
832        })?;
833    Ok(Some(VariantSelectionOutcome {
834        selection: VariantSelection {
835            selected_variant_id,
836            validation_reports: variant_validation_reports,
837            variant_validation_predictions,
838        },
839        decision,
840    }))
841}
842
843/// Capture one variant's per-fold VALIDATION (OOF) predictions (paired with id-matched y_true) and
844/// its cross-fold OOF AVERAGE block from a transient FIT_CV [`RunContext`], re-tagged with the
845/// variant's id + content fingerprint. ADDITIVE + leakage-safe: only `Validation` blocks are read (a
846/// transient run is FIT_CV-only, so no Final/Test/refit block exists), and the captured blocks are
847/// copies surfaced for host display — they never feed a training/feature path. The per-fold y_true is
848/// the same record `apply_result_scoring` retained for the score, found by `(producer, fold)`; a
849/// prediction with no matching record is skipped (it could not have been scored either).
850///
851/// The matched target record covers exactly the prediction block's SAMPLE SET (see
852/// `sample_targets_match_block`) but its rows may be in a DIFFERENT ORDER than `block.sample_ids` — a
853/// host controller may validly emit its `regression_targets` in any order. The scoring path realigns
854/// by unit id, but the host surfaces these blocks POSITIONALLY (y_pred from `block.sample_ids`/`values`
855/// paired row-for-row with `regression_targets.values`), so the y_true is REBUILT in `block.sample_ids`
856/// order here — exactly as [`oof_average_block`](crate::metrics) does for the avg — so a host pairs
857/// y_pred ↔ y_true per sample without re-sorting.
858fn capture_variant_validation_predictions(
859    variant_id: &VariantId,
860    variant_label: Option<String>,
861    ctx: &RunContext,
862) -> VariantValidationPredictions {
863    let mut predictions = Vec::new();
864    let mut regression_targets = Vec::new();
865    for block in ctx.prediction_store.blocks() {
866        if block.partition != PredictionPartition::Validation {
867            continue;
868        }
869        let Some(record) = ctx.regression_target_records.iter().find(|record| {
870            record.producer_node == block.producer_node
871                && record.producer_port == block.producer_port
872                && record.partition == PredictionPartition::Validation
873                && record.fold_id == block.fold_id
874        }) else {
875            continue;
876        };
877        predictions.push(block.clone());
878        regression_targets.push(target_block_aligned_to_samples(
879            &block.sample_ids,
880            &record.block,
881        ));
882    }
883    VariantValidationPredictions {
884        variant_id: variant_id.clone(),
885        variant_label,
886        predictions,
887        regression_targets,
888        oof_average: ctx.oof_average_blocks.first().cloned(),
889    }
890}
891
892/// Rebuild a per-fold VALIDATION `y_true` block in `sample_ids` ORDER so a host can pair it
893/// POSITIONALLY with the prediction block's `values` (the host surfaces direct prediction/target pairs
894/// by row position, not by id). `targets` covers exactly the same SAMPLE SET as `sample_ids` (the
895/// `sample_targets_match_block` precondition under which this record was retained), so every sample has
896/// a row; a missing one would indicate a broken invariant, so the original block is returned unchanged
897/// rather than dropping rows. Mirrors the avg realignment in [`oof_average_block`](crate::metrics).
898fn target_block_aligned_to_samples(
899    sample_ids: &[SampleId],
900    targets: &RegressionTargetBlock,
901) -> RegressionTargetBlock {
902    let value_by_sample: BTreeMap<&SampleId, &Vec<f64>> = targets
903        .unit_ids
904        .iter()
905        .zip(&targets.values)
906        .filter_map(|(unit_id, row)| match unit_id {
907            PredictionUnitId::Sample(sample_id) => Some((sample_id, row)),
908            _ => None,
909        })
910        .collect();
911    if sample_ids
912        .iter()
913        .any(|sample_id| !value_by_sample.contains_key(sample_id))
914    {
915        return targets.clone();
916    }
917    RegressionTargetBlock {
918        level: PredictionLevel::Sample,
919        unit_ids: sample_ids
920            .iter()
921            .cloned()
922            .map(PredictionUnitId::Sample)
923            .collect(),
924        values: sample_ids
925            .iter()
926            .map(|sample_id| value_by_sample[sample_id].clone())
927            .collect(),
928        target_names: targets.target_names.clone(),
929    }
930}
931
932#[cfg(test)]
933mod explain_contract_tests {
934    use super::*;
935
936    fn block(method: &str) -> ExplanationBlock {
937        ExplanationBlock {
938            producer_node: NodeId::new("model:base").unwrap(),
939            producer_port: None,
940            method: method.to_string(),
941            target_name: Some("y".to_string()),
942            payload: serde_json::json!({"feature_importance": [0.5, 0.3, 0.2]}),
943        }
944    }
945
946    #[test]
947    fn validates_well_formed_explanation() {
948        assert!(block("shap").validate().is_ok());
949    }
950
951    #[test]
952    fn rejects_empty_method() {
953        assert!(block("  ").validate().is_err());
954    }
955
956    #[test]
957    fn rejects_empty_target_name() {
958        let mut b = block("shap");
959        b.target_name = Some(String::new());
960        assert!(b.validate().is_err());
961    }
962
963    #[test]
964    fn round_trips_through_json() {
965        let b = block("permutation_importance");
966        let json = serde_json::to_string(&b).expect("serialize");
967        let parsed: ExplanationBlock = serde_json::from_str(&json).expect("deserialize");
968        assert_eq!(parsed, b);
969        // `target_name` is omitted when absent.
970        let mut without = block("shap");
971        without.target_name = None;
972        let json = serde_json::to_string(&without).expect("serialize");
973        assert!(!json.contains("target_name"));
974    }
975}
976
977#[cfg(test)]
978mod tests;
dag_ml_core/runtime/mod.rs

dag_ml_core/runtime/
mod.rs
