eredu_runtime/
speculative.rs

1//! Backend-neutral speculative request lifecycle and fair scheduling.
2
3use eredu_core::{
4    BoundedCompletion, CompletedSpeculativeSchedule, PreparedSpeculativeLane,
5    SpeculativeConstraint, SpeculativeDriverError, SpeculativeExecutor,
6    SpeculativeGenerationBatchOutput, SpeculativeGenerationOutput, SpeculativeGenerationVisitor,
7    SpeculativePublisher, SpeculativeRequestTable, SpeculativeSampling,
8};
9
10/// Neutral owner of speculative request registration and fair scheduling.
11pub struct SpeculativeScheduler<'a, E, S, C, P>
12where
13    E: SpeculativeExecutor,
14    S: SpeculativeSampling<Logits = E::Logits, Error = E::Error>,
15    C: SpeculativeConstraint,
16    P: SpeculativePublisher<C>,
17{
18    executor: &'a mut E,
19    context: E::Context<'a>,
20    optimistic_execution_available: bool,
21    component_timings_collected: bool,
22    requests: SpeculativeRequestTable<'a, E, S, C, P>,
23}
24
25impl<'a, E, S, C, P> SpeculativeScheduler<'a, E, S, C, P>
26where
27    E: SpeculativeExecutor + 'a,
28    S: SpeculativeSampling<Logits = E::Logits, Error = E::Error, Context<'a> = E::Context<'a>> + 'a,
29    C: SpeculativeConstraint,
30    P: SpeculativePublisher<C>,
31{
32    /// Creates a scheduler for one prepared executor and placement.
33    #[allow(clippy::too_many_arguments)]
34    pub fn new(
35        executor: &'a mut E,
36        options: eredu_core::generation::SpeculativeSchedulerOptions,
37        topology: eredu_core::SpeculativeExecutionTopology,
38        optimistic_execution_available: bool,
39        component_timings_collected: bool,
40        context: E::Context<'a>,
41    ) -> Result<Self, SpeculativeDriverError<E::Error>> {
42        let completion_wait = options
43            .completion_wait()
44            .map_err(SpeculativeDriverError::Generation)?;
45        if !E::Completion::supports_cancellation(completion_wait.cancellation()) {
46            return Err(SpeculativeDriverError::UnsupportedCompletionCancellation {
47                cancellation: completion_wait.cancellation(),
48            });
49        }
50        executor.set_telemetry_enabled(component_timings_collected);
51        Ok(Self {
52            executor,
53            context,
54            optimistic_execution_available,
55            component_timings_collected,
56            requests: SpeculativeRequestTable::new(options, topology)
57                .map_err(SpeculativeDriverError::Generation)?,
58        })
59    }
60
61    /// Registers and prefills one independently progressing lane.
62    pub fn submit(
63        &mut self,
64        mut lane: PreparedSpeculativeLane<'a, E, S, C, P>,
65    ) -> Result<eredu_core::generation::SpeculativeRequestId, SpeculativeDriverError<E::Error>>
66    {
67        self.requests.submit(
68            self.executor,
69            lane.take_cache(),
70            lane.take_input(),
71            lane.take_config(),
72            lane.take_runtime(),
73            lane.take_randomness(),
74            self.component_timings_collected,
75            self.context,
76        )
77    }
78
79    /// Performs one fairly selected lifecycle action.
80    pub fn step(&mut self) -> Result<bool, SpeculativeDriverError<E::Error>> {
81        self.requests.step(
82            self.executor,
83            self.optimistic_execution_available,
84            self.context,
85        )
86    }
87
88    /// Drives all registered lanes to terminal states.
89    pub fn run(&mut self) -> Result<(), SpeculativeDriverError<E::Error>> {
90        while self.step()? {}
91        Ok(())
92    }
93
94    /// Returns the current portable status for one lane.
95    pub fn status(
96        &self,
97        id: eredu_core::generation::SpeculativeRequestId,
98    ) -> Option<eredu_core::generation::SpeculativeRequestStatus> {
99        self.requests.status(id)
100    }
101
102    /// Requests cancellation at the next exact safe boundary.
103    pub fn cancel(
104        &mut self,
105        id: eredu_core::generation::SpeculativeRequestId,
106    ) -> Result<(), SpeculativeDriverError<E::Error>> {
107        self.requests.cancel(id)
108    }
109
110    /// Whether all registered lanes are terminal.
111    pub fn is_finished(&self) -> bool {
112        self.requests.is_finished()
113    }
114
115    /// Consumes a terminal scheduler into stable ordered results.
116    pub fn finish(
117        self,
118    ) -> Result<CompletedSpeculativeSchedule<S>, SpeculativeDriverError<E::Error>> {
119        self.requests.finish()
120    }
121}
122
123/// Facade-selected speculative generation driver.
124///
125/// Backends lend prepared native resources through
126/// [`SpeculativeGenerationVisitor`]. This driver alone registers lanes, runs
127/// the fair schedule, observes exact completions, and constructs public
128/// terminal outputs.
129#[derive(Debug, Clone, Copy, Default)]
130pub struct RunSpeculativeGeneration {
131    options: eredu_core::generation::SpeculativeSchedulerOptions,
132}
133
134impl RunSpeculativeGeneration {
135    /// Creates a driver with facade-selected scheduling and lookahead controls.
136    pub const fn new(options: eredu_core::generation::SpeculativeSchedulerOptions) -> Self {
137        Self { options }
138    }
139}
140
141impl SpeculativeGenerationVisitor for RunSpeculativeGeneration {
142    fn run<'a, E, S, C, P>(
143        self,
144        executor: &'a mut E,
145        lanes: Vec<PreparedSpeculativeLane<'a, E, S, C, P>>,
146        topology: eredu_core::SpeculativeExecutionTopology,
147        optimistic_execution_available: bool,
148        component_timings_collected: bool,
149        context: E::Context<'a>,
150    ) -> Result<SpeculativeGenerationBatchOutput, SpeculativeDriverError<E::Error>>
151    where
152        E: SpeculativeExecutor + 'a,
153        S: SpeculativeSampling<Logits = E::Logits, Error = E::Error, Context<'a> = E::Context<'a>>
154            + 'a,
155        C: SpeculativeConstraint,
156        P: SpeculativePublisher<C>,
157    {
158        let mut scheduler = SpeculativeScheduler::new(
159            executor,
160            self.options,
161            topology,
162            optimistic_execution_available,
163            component_timings_collected,
164            context,
165        )?;
166        for lane in lanes {
167            scheduler.submit(lane)?;
168        }
169        scheduler.run()?;
170        let mut completed = scheduler.finish()?;
171        let requests = completed
172            .take_requests()
173            .into_iter()
174            .map(|request| -> Result<_, SpeculativeDriverError<E::Error>> {
175                let finish_reason = request.finish_reason().ok_or_else(|| {
176                    SpeculativeDriverError::Generation(
177                        eredu_core::generation::GenerationError::MissingSpeculativeFinishReason {
178                            index: request.id().index(),
179                        },
180                    )
181                })?;
182                Ok(SpeculativeGenerationOutput::new(
183                    request.token_ids().to_vec(),
184                    finish_reason,
185                    request.stats().clone(),
186                ))
187            })
188            .collect::<Result<Vec<_>, _>>()?;
189        Ok(SpeculativeGenerationBatchOutput::new(
190            requests,
191            completed.take_scheduler(),
192        ))
193    }
194}
eredu_runtime/speculative.rs

eredu_runtime/
speculative.rs
