eredu_runtime/
intervention.rs

1//! Prospective intervention scheduling under the ordinary capture/run owner.
2//! There is no additional sampling loop or native-resource owner here.
3
4use crate::capture::{
5    bounded_diagnostic, capture_value, metadata_reservation, CaptureExecutionError, CaptureSession,
6};
7use eredu_core::{capture::*, intervention::*, ObservationPosition};
8
9mod activation;
10mod session;
11pub use activation::apply_activation;
12pub(crate) use session::validate_continuation;
13pub use session::{install_session, validate_session, CaptureObserver};
14
15pub(crate) struct InterventionRun {
16    pub(crate) plan: AdmittedInterventionPlan,
17    pub(crate) records: Option<Vec<InterventionRecord>>,
18    pub(crate) routing_pending: Option<usize>,
19    estimator: std::sync::Arc<dyn InterventionEstimator>,
20}
21
22impl InterventionRun {
23    pub(crate) fn begin_step(
24        &mut self,
25        ledger: &mut CaptureLedger,
26        phase: CapturePhase,
27        prediction: u64,
28    ) -> Result<(), CaptureError> {
29        if self.records.is_some() {
30            return Err(CaptureError::Invalid(
31                "previous intervention step has not been consumed".into(),
32            ));
33        }
34        let mut records = Vec::new();
35        for (operation, point) in self.plan.plan().operations.iter().zip(self.plan.points()) {
36            let charged = intervention_metadata(operation, point, self.plan.identity())?;
37            reserve_envelope(ledger, charged)?;
38            let active = operation.schedule.includes(phase, prediction);
39            let mut evidence = Vec::new();
40            for (selection, geometry) in evidence_selections(operation, point) {
41                let charged = metadata_reservation(&selection, &geometry)?;
42                reserve_envelope(ledger, charged)?;
43                evidence.push(CaptureRecord {
44                    schema_version: CAPTURE_SCHEMA_VERSION,
45                    selection_id: selection.id,
46                    path: geometry.path,
47                    node_id: geometry.node_id,
48                    position: geometry.position,
49                    source_shape: None,
50                    selected_shape: None,
51                    outcome: if active {
52                        CaptureOutcome::Missing
53                    } else {
54                        CaptureOutcome::Skipped {
55                            reason: CaptureSkipReason::Schedule,
56                        }
57                    },
58                    payload: None,
59                    charged,
60                });
61            }
62            records.push(InterventionRecord {
63                schema_version: INTERVENTION_SCHEMA_VERSION,
64                plan_id: self.plan.identity().into(),
65                operation_id: operation.id.clone(),
66                target: operation.target.clone(),
67                node_id: point.node_id.clone(),
68                phase,
69                prediction_index: prediction,
70                outcome: if active {
71                    InterventionOutcome::Missing
72                } else {
73                    InterventionOutcome::Inactive
74                },
75                evidence,
76                charged,
77            });
78        }
79        self.records = Some(records);
80        self.routing_pending = None;
81        Ok(())
82    }
83
84    pub(crate) fn take_records(&mut self) -> Vec<InterventionRecord> {
85        self.records.take().unwrap_or_default()
86    }
87}
88
89impl CaptureSession {
90    /// Installs an immutable intervention plan before the first step. Both plans
91    /// share one ledger; capture-none budgets must still allow outcome metadata.
92    pub fn enable_interventions(
93        &mut self,
94        plan: AdmittedInterventionPlan,
95        estimator: std::sync::Arc<dyn InterventionEstimator>,
96    ) -> Result<(), CaptureError> {
97        if self.records.is_some()
98            || self.interventions.is_some()
99            || self.ledger.total() != CaptureUsage::default()
100        {
101            return Err(CaptureError::Invalid(
102                "interventions must be installed once before generation".into(),
103            ));
104        }
105        if plan.request() != self.plan.request() {
106            return Err(CaptureError::Invalid(
107                "capture and intervention request geometry differs".into(),
108            ));
109        }
110        if !plan.is_empty() {
111            self.interventions = Some(InterventionRun {
112                plan,
113                records: None,
114                routing_pending: None,
115                estimator,
116            });
117        }
118        Ok(())
119    }
120
121    /// Applies scheduled activation operations in plan order. Ordinary observation
122    /// must precede this call at the exact same hook. Returns no replacement when
123    /// no operation is active, avoiding source cloning and native materialization.
124    pub fn intervene<B: InterventionBackend>(
125        &mut self,
126        backend: &mut B,
127        path: &str,
128        tensor: &B::Tensor,
129    ) -> Result<Option<B::Tensor>, CaptureExecutionError<B::Error>> {
130        let Some(run) = &mut self.interventions else {
131            return Ok(None);
132        };
133        let records = run
134            .records
135            .as_mut()
136            .ok_or_else(|| CaptureError::Invalid("intervention step not started".into()))?;
137        let mut effective = None;
138        for (index, ((operation, point), record)) in run
139            .plan
140            .plan()
141            .operations
142            .iter()
143            .zip(run.plan.points())
144            .zip(records)
145            .enumerate()
146        {
147            if operation.target != path
148                || point.routing.is_some()
149                || record.outcome == InterventionOutcome::Inactive
150            {
151                continue;
152            }
153            if record.outcome != InterventionOutcome::Missing {
154                return Err(CaptureError::Invalid(
155                    "intervention point executed more than once in one step".into(),
156                )
157                .into());
158            }
159            let started = std::time::Instant::now();
160            let result = (|| {
161                let input = effective.as_ref().unwrap_or(tensor);
162                let shape = backend
163                    .shape(input)
164                    .map_err(CaptureExecutionError::Backend)?;
165                let dtype = backend
166                    .intervention_dtype(input)
167                    .map_err(CaptureExecutionError::Backend)?;
168                let slice = run.plan.validate_actual(
169                    index,
170                    self.phase,
171                    self.prediction,
172                    &shape,
173                    Some(dtype),
174                )?;
175                run.estimator.validate_geometry(&shape, &slice)?;
176                let evidence = evidence_selections(operation, point);
177                if let Some((selection, geometry)) = evidence.first() {
178                    capture_evidence(
179                        backend,
180                        input,
181                        selection,
182                        geometry,
183                        &mut record.evidence[0],
184                        run.plan.request(),
185                        self.phase,
186                        self.prediction,
187                        &mut self.ledger,
188                    )?;
189                }
190                let output = apply_activation(backend, input, &operation.action, &slice)?;
191                // Backend conformance checks: an implementation cannot replace a
192                // value with a different shape or silently promote its dtype.
193                let output_shape = backend
194                    .shape(&output)
195                    .map_err(CaptureExecutionError::Backend)?;
196                let output_dtype = backend
197                    .intervention_dtype(&output)
198                    .map_err(CaptureExecutionError::Backend)?;
199                run.plan.validate_actual(
200                    index,
201                    self.phase,
202                    self.prediction,
203                    &output_shape,
204                    Some(output_dtype),
205                )?;
206                if let Some((selection, geometry)) = evidence.get(1) {
207                    capture_evidence(
208                        backend,
209                        &output,
210                        selection,
211                        geometry,
212                        &mut record.evidence[1],
213                        run.plan.request(),
214                        self.phase,
215                        self.prediction,
216                        &mut self.ledger,
217                    )?;
218                }
219                Ok(output)
220            })();
221            self.capture_seconds += started.elapsed().as_secs_f64();
222            match result {
223                Ok(output) => {
224                    effective = Some(output);
225                    record.outcome = InterventionOutcome::Applied;
226                }
227                Err(error) => {
228                    record.outcome = InterventionOutcome::Failed {
229                        message: bounded_diagnostic(&error),
230                    };
231                    return Err(error);
232                }
233            }
234        }
235        Ok(effective)
236    }
237
238    /// Checks scheduled targets before committing a prediction. Missing targets are
239    /// preserved in records and fail the attempt; this does not establish rollback.
240    pub fn finish_interventions(&self) -> Result<(), CaptureError> {
241        if let Some(run) = &self.interventions {
242            if run.routing_pending.is_some() {
243                return Err(CaptureError::Invalid(
244                    "routing intervention has not resolved before the boundary".into(),
245                ));
246            }
247            let records = run
248                .records
249                .as_ref()
250                .ok_or_else(|| CaptureError::Invalid("intervention step not started".into()))?;
251            if let Some(record) = records.iter().find(|r| {
252                matches!(
253                    r.outcome,
254                    InterventionOutcome::Missing | InterventionOutcome::Failed { .. }
255                )
256            }) {
257                return Err(CaptureError::Invalid(format!(
258                    "scheduled intervention {} at {} did not complete",
259                    record.operation_id, record.target
260                )));
261            }
262        }
263        Ok(())
264    }
265}
266
267fn reserve_envelope(ledger: &mut CaptureLedger, usage: CaptureUsage) -> Result<(), CaptureError> {
268    match ledger.reserve(usage)? {
269        Some(CaptureSkipReason::Limit { budget, cumulative }) => {
270            Err(CaptureError::Limit { budget, cumulative })
271        }
272        _ => Ok(()),
273    }
274}
275
276fn intervention_metadata(
277    operation: &InterventionOperation,
278    point: &InterventionPoint,
279    identity: &str,
280) -> Result<CaptureUsage, CaptureError> {
281    let strings = add(
282        add(operation.id.len() as u64, point.path.len() as u64)?,
283        add(point.node_id.len() as u64, identity.len() as u64)?,
284    )?;
285    Ok(CaptureUsage {
286        captures: 0,
287        retained_bytes: 0,
288        host_bytes: add(1024, strings)?,
289        encoded_bytes: add(4096, mul(strings, 6)?)?,
290    })
291}
292
293/// Generates only declared evidence fields, with exact shape and attribution.
294pub(crate) fn evidence_selections(
295    operation: &InterventionOperation,
296    point: &InterventionPoint,
297) -> Vec<(CaptureSelection, eredu_core::ObservationPoint)> {
298    let transform = match operation.evidence {
299        InterventionEvidence::None => return vec![],
300        InterventionEvidence::Preview { max_elements } => {
301            CaptureTransform::Preview { max_elements }
302        }
303        InterventionEvidence::Summary => CaptureTransform::Summary,
304    };
305    let mut entries = Vec::new();
306    for (label, position) in [
307        ("before", ObservationPosition::BeforeIntervention),
308        ("after", ObservationPosition::AfterIntervention),
309    ] {
310        let fields: &[Option<eredu_core::RoutingObservationField>] = if point.routing.is_some() {
311            &[
312                Some(eredu_core::RoutingObservationField::SelectedExperts),
313                Some(eredu_core::RoutingObservationField::Coefficients),
314            ]
315        } else {
316            &[None]
317        };
318        for field in fields {
319            let mut geometry = point.observation_geometry();
320            geometry.position = position;
321            if let Some(field) = field {
322                geometry.path = field.path(&point.path);
323                if *field == eredu_core::RoutingObservationField::SelectedExperts {
324                    geometry.dtype = eredu_core::ObservationDtype::Integer;
325                }
326            }
327            entries.push((
328                CaptureSelection {
329                    id: format!("{}:{label}:{:?}", operation.id, field),
330                    path: geometry.path.clone(),
331                    schedule: operation.schedule.clone(),
332                    slices: operation.slices.clone(),
333                    transform: transform.clone(),
334                },
335                geometry,
336            ));
337        }
338    }
339    entries
340}
341
342#[allow(clippy::too_many_arguments)]
343pub(crate) fn capture_evidence<B: CaptureBackend>(
344    backend: &mut B,
345    tensor: &B::Tensor,
346    selection: &CaptureSelection,
347    point: &eredu_core::ObservationPoint,
348    record: &mut CaptureRecord,
349    request: CaptureRequestShape,
350    phase: CapturePhase,
351    prediction: u64,
352    ledger: &mut CaptureLedger,
353) -> Result<(), CaptureExecutionError<B::Error>> {
354    let result = capture_value(
355        backend, tensor, selection, point, record, request, phase, prediction, ledger,
356    );
357    if let Err(error) = &result {
358        record.payload = None;
359        record.outcome = CaptureOutcome::Failed {
360            reason: match error {
361                CaptureExecutionError::Admission(CaptureError::Limit { budget, cumulative }) => {
362                    CaptureFailureReason::Limit {
363                        budget: *budget,
364                        cumulative: *cumulative,
365                    }
366                }
367                CaptureExecutionError::Admission(_) => CaptureFailureReason::Invalid,
368                CaptureExecutionError::Backend(_) => CaptureFailureReason::Native,
369            },
370            message: bounded_diagnostic(error),
371        };
372    }
373    result
374}
375
376/// Joint preflight counts every ordinary capture, intervention diagnostic and
377/// before/after transform against the same per-step and cumulative budgets.
378pub fn preflight(
379    capture: &AdmittedCapturePlan,
380    intervention: &AdmittedInterventionPlan,
381    estimator: &dyn InterventionEstimator,
382) -> Result<(), CaptureError> {
383    preflight_continuation(capture, intervention, estimator, 0, CaptureUsage::default())
384}
385
386pub(crate) fn preflight_continuation(
387    capture: &AdmittedCapturePlan,
388    intervention: &AdmittedInterventionPlan,
389    estimator: &dyn InterventionEstimator,
390    next_prediction: u64,
391    inherited: CaptureUsage,
392) -> Result<(), CaptureError> {
393    if capture.request() != intervention.request() {
394        return Err(CaptureError::Invalid(
395            "capture/intervention geometry mismatch".into(),
396        ));
397    }
398    let mut base = CaptureUsage::default();
399    let mut extra = Vec::new();
400    let mut scheduled_costs = Vec::new();
401    for (index, (operation, point)) in intervention
402        .plan()
403        .operations
404        .iter()
405        .zip(intervention.points())
406        .enumerate()
407    {
408        base = base.checked_add(intervention_metadata(
409            operation,
410            point,
411            intervention.identity(),
412        )?)?;
413        extra.extend(evidence_selections(operation, point));
414        let mut costs = [CaptureUsage::default(); 2];
415        for (phase_index, phase) in [CapturePhase::Prefill, CapturePhase::Decode]
416            .into_iter()
417            .enumerate()
418        {
419            let Some((_, last)) = operation.schedule.count_and_last_from(
420                phase,
421                next_prediction,
422                intervention.request().max_predictions,
423            )?
424            else {
425                continue;
426            };
427            if let Some(shape) =
428                intervention
429                    .request()
430                    .resolve(&point.observation_geometry(), phase, last)?
431            {
432                let slice = intervention.validate_actual(
433                    index,
434                    phase,
435                    last,
436                    &shape,
437                    operation.action.dtype(),
438                )?;
439                estimator.validate_geometry(&shape, &slice)?;
440            }
441            if let Some(policy) = &point.routing {
442                if operation.evidence != InterventionEvidence::None {
443                    let rows = mul(
444                        intervention.request().batch,
445                        if phase == CapturePhase::Prefill {
446                            intervention.request().prompt_tokens
447                        } else {
448                            1
449                        },
450                    )?;
451                    costs[phase_index] = original_route_cost(estimator, policy, rows)?;
452                }
453            }
454        }
455        if point.routing.is_some() && operation.evidence != InterventionEvidence::None {
456            scheduled_costs.push((operation.schedule.clone(), costs));
457        }
458    }
459    crate::capture::preflight_continuation(
460        capture,
461        &extra,
462        base,
463        &scheduled_costs,
464        next_prediction,
465        inherited,
466        |source, selection, slice| estimator.capture_usage(source, selection, slice),
467    )
468}
469
470fn original_route_cost(
471    estimator: &dyn InterventionEstimator,
472    policy: &InterventionRoutingPolicy,
473    rows: u64,
474) -> Result<CaptureUsage, CaptureError> {
475    let cost = estimator.original_route_usage(policy, rows)?;
476    if cost.captures != 0 || cost.encoded_bytes != 0 {
477        return Err(CaptureError::Invalid(
478            "original routing estimate must exclude evidence transforms and diagnostic encoding"
479                .into(),
480        ));
481    }
482    Ok(cost)
483}
484
485#[cfg(test)]
486mod tests;
487
488impl CaptureSession {
489    /// Validates a scheduled routing operation against actual flattened token rows
490    /// and reserves any original-decision work before the selector runs.
491    pub fn routing_control(
492        &mut self,
493        path: &str,
494        token_rows: u64,
495    ) -> Result<Option<eredu_nn::routing_intervention::GroupSelectionControl>, CaptureError> {
496        use eredu_nn::routing_intervention::{
497            GroupScoreStage, GroupSelectionAction, GroupSelectionControl,
498        };
499        let Some(run) = &mut self.interventions else {
500            return Ok(None);
501        };
502        let records = run
503            .records
504            .as_mut()
505            .ok_or_else(|| CaptureError::Invalid("intervention step not started".into()))?;
506        let Some(index) = run
507            .plan
508            .plan()
509            .operations
510            .iter()
511            .zip(records.iter())
512            .position(|(op, record)| {
513                op.target == path && record.outcome != InterventionOutcome::Inactive
514            })
515        else {
516            return Ok(None);
517        };
518        let operation = &run.plan.plan().operations[index];
519        let point = &run.plan.points()[index];
520        let record = &mut records[index];
521        let result = (|| {
522            if run.routing_pending.is_some() || record.outcome != InterventionOutcome::Missing {
523                return Err(CaptureError::Invalid(
524                    "routing target executed more than once or overlaps another pending target"
525                        .into(),
526                ));
527            }
528            let policy = point.routing.as_ref().ok_or_else(|| {
529                CaptureError::Invalid("routing control reached activation target".into())
530            })?;
531            let slice = run.plan.validate_actual(
532                index,
533                self.phase,
534                self.prediction,
535                &[token_rows, policy.top_k as u64],
536                None,
537            )?;
538            run.estimator
539                .validate_geometry(&[token_rows, policy.top_k as u64], &slice)?;
540            let action = match &operation.action {
541                InterventionAction::ExcludeExperts { expert_ids } => {
542                    GroupSelectionAction::Exclude(expert_ids.clone())
543                }
544                InterventionAction::ZeroExpertContribution { expert_ids } => {
545                    GroupSelectionAction::ZeroContribution(expert_ids.clone())
546                }
547                InterventionAction::ForceExperts { expert_ids, .. } => {
548                    GroupSelectionAction::Force(expert_ids.clone())
549                }
550                InterventionAction::BiasRoutingScores {
551                    stage,
552                    expert_ids,
553                    biases,
554                } => GroupSelectionAction::Bias {
555                    stage: match stage {
556                        RoutingScoreStage::RawLogits => GroupScoreStage::RawLogits,
557                        RoutingScoreStage::TransformedScores => GroupScoreStage::TransformedScores,
558                        RoutingScoreStage::RankingScores => GroupScoreStage::RankingScores,
559                    },
560                    ids: expert_ids.clone(),
561                    values: biases.clone(),
562                },
563                _ => {
564                    return Err(CaptureError::Invalid(
565                        "activation operation cannot control routing".into(),
566                    ))
567                }
568            };
569            let expected = eredu_nn::TopKGroupSelectionSpec::new(
570                i32::try_from(policy.expert_count).map_err(|_| CaptureError::Overflow)?,
571                i32::try_from(policy.top_k).map_err(|_| CaptureError::Overflow)?,
572                match policy.scoring {
573                    RoutingScoring::Softmax => eredu_nn::GroupScoring::Softmax,
574                    RoutingScoring::SelectedSoftmax => eredu_nn::GroupScoring::SelectedSoftmax,
575                    RoutingScoring::Sigmoid => eredu_nn::GroupScoring::Sigmoid,
576                    RoutingScoring::SqrtSoftplus => eredu_nn::GroupScoring::SqrtSoftplus,
577                },
578                policy.normalize_selected,
579            )
580            .and_then(|spec| spec.with_groups(policy.groups as i32, policy.selected_groups as i32))
581            .and_then(|spec| {
582                spec.with_weight_policy(policy.normalization_epsilon, policy.coefficient_scale)
583            })
584            .map_err(|error| CaptureError::Invalid(error.to_string()))?;
585            let capture_original = operation.evidence != InterventionEvidence::None;
586            if capture_original {
587                let cost = original_route_cost(run.estimator.as_ref(), policy, token_rows)?;
588                reserve_envelope(&mut self.ledger, cost)?;
589                record.charged = record.charged.checked_add(cost)?;
590            }
591            Ok(GroupSelectionControl {
592                expected,
593                learned_coefficient_scale: policy.learned_coefficient_scale,
594                first_row: slice.starts[0],
595                end_row: slice.ends[0],
596                row_stride: slice.strides[0],
597                action,
598                capture_original,
599            })
600        })();
601        match result {
602            Ok(control) => {
603                run.routing_pending = Some(index);
604                Ok(Some(control))
605            }
606            Err(error) => {
607                record.outcome = InterventionOutcome::Failed {
608                    message: bounded_diagnostic(&error),
609                };
610                Err(error)
611            }
612        }
613    }
614
615    /// Captures attributed IDs/coefficients before the expert provider runs.
616    pub fn routing_applied<B: CaptureBackend>(
617        &mut self,
618        backend: &mut B,
619        path: &str,
620        original: Option<crate::RoutingDecision<'_, B::Tensor>>,
621        effective: crate::RoutingDecision<'_, B::Tensor>,
622    ) -> Result<(), CaptureExecutionError<B::Error>> {
623        let run = self
624            .interventions
625            .as_mut()
626            .ok_or_else(|| CaptureError::Invalid("unsolicited routing result".into()))?;
627        let index = run
628            .routing_pending
629            .take()
630            .ok_or_else(|| CaptureError::Invalid("routing result has no pending control".into()))?;
631        let operation = &run.plan.plan().operations[index];
632        if operation.target != path {
633            return Err(
634                CaptureError::Invalid("routing result target differs from control".into()).into(),
635            );
636        }
637        let record = &mut run
638            .records
639            .as_mut()
640            .ok_or_else(|| CaptureError::Invalid("intervention step not started".into()))?[index];
641        let started = std::time::Instant::now();
642        let result = (|| {
643            let shape = backend
644                .shape(effective.ids)
645                .map_err(CaptureExecutionError::Backend)?;
646            run.plan
647                .validate_actual(index, self.phase, self.prediction, &shape, None)?;
648            if backend
649                .shape(effective.coefficients)
650                .map_err(CaptureExecutionError::Backend)?
651                != shape
652            {
653                return Err(CaptureError::Invalid(
654                    "effective route IDs and coefficient shapes differ".into(),
655                )
656                .into());
657            }
658            let selections = evidence_selections(operation, &run.plan.points()[index]);
659            if !selections.is_empty() {
660                let original = original.ok_or_else(|| {
661                    CaptureError::Invalid("requested original routing decision is missing".into())
662                })?;
663                for ((tensor, (selection, geometry)), evidence) in [
664                    original.ids,
665                    original.coefficients,
666                    effective.ids,
667                    effective.coefficients,
668                ]
669                .into_iter()
670                .zip(&selections)
671                .zip(&mut record.evidence)
672                {
673                    capture_evidence(
674                        backend,
675                        tensor,
676                        selection,
677                        geometry,
678                        evidence,
679                        run.plan.request(),
680                        self.phase,
681                        self.prediction,
682                        &mut self.ledger,
683                    )?;
684                }
685            }
686            Ok(())
687        })();
688        self.capture_seconds += started.elapsed().as_secs_f64();
689        record.outcome = match &result {
690            Ok(()) => InterventionOutcome::Applied,
691            Err(error) => InterventionOutcome::Failed {
692                message: bounded_diagnostic(error),
693            },
694        };
695        result
696    }
697
698    /// Retains a bounded selector failure without treating it as proof of rollback.
699    pub fn routing_failed(&mut self, path: &str, message: &str) {
700        if let Some(run) = &mut self.interventions {
701            if let Some(index) = run.routing_pending.take() {
702                if run.plan.plan().operations[index].target == path {
703                    if let Some(records) = &mut run.records {
704                        records[index].outcome = InterventionOutcome::Failed {
705                            message: bounded_diagnostic(&message),
706                        };
707                    }
708                }
709            }
710        }
711    }
712}
eredu_runtime/intervention.rs

eredu_runtime/
intervention.rs
