eredu_runtime/intervention/
session.rs1use super::*;
3use std::sync::Arc;
4
5pub fn validate_session(
8 capture: &AdmittedCapturePlan,
9 plan: &AdmittedInterventionPlan,
10 discovery: &InterventionDiscovery,
11 estimator: &dyn InterventionEstimator,
12) -> Result<(), CaptureError> {
13 validate_continuation(
14 capture,
15 plan,
16 discovery,
17 estimator,
18 0,
19 CaptureUsage::default(),
20 )
21}
22
23pub(crate) fn validate_continuation(
24 capture: &AdmittedCapturePlan,
25 plan: &AdmittedInterventionPlan,
26 discovery: &InterventionDiscovery,
27 estimator: &dyn InterventionEstimator,
28 next_prediction: u64,
29 inherited: CaptureUsage,
30) -> Result<(), CaptureError> {
31 if plan.request().batch != 1 {
32 return Err(CaptureError::Unsupported(
33 "interventions require single-sequence text generation".into(),
34 ));
35 }
36 let checked = plan
37 .plan()
38 .clone()
39 .admit(discovery, plan.request(), plan.session_id())?;
40 if checked.identity() != plan.identity() {
41 return Err(CaptureError::Invalid(
42 "intervention admission differs from loaded source/session capabilities".into(),
43 ));
44 }
45 preflight_continuation(capture, &checked, estimator, next_prediction, inherited)
46}
47
48pub fn install_session(
51 slot: &mut Option<CaptureSession>,
52 capture: AdmittedCapturePlan,
53 intervention: Option<(AdmittedInterventionPlan, Arc<dyn InterventionEstimator>)>,
54) -> Result<(), CaptureError> {
55 if slot.is_some() {
56 return Err(CaptureError::Invalid(
57 "capture/intervention plan already installed".into(),
58 ));
59 }
60 if capture.is_empty()
61 && intervention
62 .as_ref()
63 .is_none_or(|(plan, _)| plan.is_empty())
64 {
65 return Ok(());
66 }
67 let mut session = CaptureSession::new(capture);
68 if let Some((plan, estimator)) = intervention {
69 session.enable_interventions(plan, estimator)?;
70 }
71 *slot = Some(session);
72 Ok(())
73}
74
75pub struct CaptureObserver<'a, B, F> {
78 session: &'a mut CaptureSession,
79 backend: B,
80 map_error: F,
81}
82impl<'a, B, F> CaptureObserver<'a, B, F> {
83 pub fn new(session: &'a mut CaptureSession, backend: B, map_error: F) -> Self {
85 Self {
86 session,
87 backend,
88 map_error,
89 }
90 }
91}
92impl<B, E, F> crate::ActivationObserver<B::Tensor, E> for CaptureObserver<'_, B, F>
93where
94 B: InterventionBackend,
95 F: Fn(CaptureExecutionError<B::Error>) -> E,
96{
97 fn observe(&mut self, path: &str, value: &B::Tensor) -> Result<(), E> {
98 self.session
99 .observe(&mut self.backend, path, value)
100 .map_err(&self.map_error)
101 }
102 fn intervene(&mut self, path: &str, value: &B::Tensor) -> Result<Option<B::Tensor>, E> {
103 self.session
104 .intervene(&mut self.backend, path, value)
105 .map_err(&self.map_error)
106 }
107 fn routing_control(
108 &mut self,
109 path: &str,
110 rows: u64,
111 ) -> Result<Option<eredu_nn::routing_intervention::GroupSelectionControl>, E> {
112 self.session
113 .routing_control(path, rows)
114 .map_err(|error| (self.map_error)(error.into()))
115 }
116 fn routing_applied(
117 &mut self,
118 path: &str,
119 original: Option<crate::RoutingDecision<'_, B::Tensor>>,
120 effective: crate::RoutingDecision<'_, B::Tensor>,
121 ) -> Result<(), E> {
122 self.session
123 .routing_applied(&mut self.backend, path, original, effective)
124 .map_err(&self.map_error)
125 }
126 fn routing_failed(&mut self, path: &str, message: &str) {
127 self.session.routing_failed(path, message);
128 }
129 fn finish(&mut self) -> Result<(), E> {
130 self.session
131 .finish_interventions()
132 .map_err(|error| (self.map_error)(error.into()))
133 }
134 fn observe_routing(
135 &mut self,
136 routing: crate::RoutingObservation<'_, B::Tensor>,
137 ) -> Result<(), E> {
138 let mut result = Ok(());
139 routing.for_each_tensor(|path, value| {
140 if result.is_ok() {
141 result = self.observe(&path, value);
142 }
143 });
144 result
145 }
146}