1use super::{SnapshotBudget, SnapshotReservation};
8use crate::capture::{
9 CaptureCheckpoint, CaptureForkRequest, CaptureSession, InterventionForkRequest,
10};
11use eredu_core::{
12 capture::CaptureError,
13 execution_control::{
14 ExecutionControlError, NativeTextStateBackend, SnapshotEstimate, SnapshotResourceKind,
15 },
16 ModelRuntime, PendingTextInput, TextContinuationBoundary, TextContinuationIdentity,
17 TokenFilterController,
18};
19
20pub trait TextSnapshotBackend: NativeTextStateBackend {
24 type SamplingState;
27
28 fn sampling_state(state: &Self::TextGenerationState) -> &Self::SamplingState;
30 fn install_sampling_state(state: &mut Self::TextGenerationState, sampling: Self::SamplingState);
33 fn assemble_generation_state(
36 sampling: Self::SamplingState,
37 capture: Option<CaptureSession>,
38 ) -> Self::TextGenerationState;
39 fn sampling_prediction(sampling: &Self::SamplingState) -> u64;
41 fn estimate_sampling_state(
43 runtime: &ModelRuntime<Self>,
44 sampling: &Self::SamplingState,
45 ) -> Result<Option<SnapshotEstimate>, Self::Error>;
46 fn copy_sampling_state(
48 runtime: &mut ModelRuntime<Self>,
49 sampling: &Self::SamplingState,
50 ) -> Result<Self::SamplingState, Self::Error>;
51 fn estimate_pending_input(
54 runtime: &ModelRuntime<Self>,
55 input: Option<PendingTextInput<&Self::Prompt, &Self::Token>>,
56 ) -> Result<Option<SnapshotEstimate>, Self::Error>;
57 fn continuation_input_tokens(
60 _input: Option<PendingTextInput<&Self::Prompt, &Self::Token>>,
61 _predictions: u64,
62 ) -> Option<u64> {
63 None
64 }
65 fn estimate_sampling_growth(
68 _runtime: &ModelRuntime<Self>,
69 _sampling: &Self::SamplingState,
70 _predictions: u64,
71 ) -> Result<Option<u64>, Self::Error> {
72 Ok(None)
73 }
74 #[allow(clippy::type_complexity)]
76 fn copy_pending_input(
77 runtime: &mut ModelRuntime<Self>,
78 input: Option<PendingTextInput<&Self::Prompt, &Self::Token>>,
79 ) -> Result<Option<PendingTextInput<Self::Prompt, Self::Token>>, Self::Error>;
80 fn capture_run(state: &Self::TextGenerationState) -> Option<&CaptureSession>;
82 fn capture_run_mut(state: &mut Self::TextGenerationState) -> Option<&mut CaptureSession>;
84
85 fn estimate_child_capture(
87 runtime: &ModelRuntime<Self>,
88 shape: &[u64],
89 selection: &eredu_core::capture::CaptureSelection,
90 slice: &eredu_core::capture::ResolvedCaptureSlice,
91 ) -> Result<eredu_core::capture::CaptureUsage, CaptureError>;
92 fn child_intervention_estimator(
94 runtime: &ModelRuntime<Self>,
95 ) -> Result<std::sync::Arc<dyn eredu_core::intervention::InterventionEstimator>, CaptureError>;
96}
97
98pub trait SnapshotTokenController: TokenFilterController + Sized {
101 fn snapshot_storage_bytes(&self) -> Option<u64>;
104 fn fork_snapshot(&self) -> Result<Self, String>;
107}
108
109#[derive(Debug, thiserror::Error)]
111pub enum TextSnapshotError<E: std::error::Error + 'static> {
112 #[error("host continuation snapshot failed: {0}")]
114 Host(String),
115 #[error("constraint snapshot failed: {0}")]
117 Controller(String),
118 #[error("snapshot backend operation failed: {0}")]
120 Backend(#[source] E),
121 #[error(transparent)]
123 Capture(#[from] CaptureError),
124 #[error(transparent)]
126 Control(#[from] ExecutionControlError),
127 #[error("snapshot belongs to another continuation")]
129 IncompatibleRun,
130 #[error("snapshot prediction or capture ownership differs")]
132 InconsistentState,
133 #[error("unsupported continuation: {0}")]
135 Unsupported(&'static str),
136}
137
138pub struct TextBranchRequest<'a> {
141 pub session_id: &'a str,
143 pub max_predictions: u64,
145 pub capture_limits: Option<eredu_core::capture::CaptureLimits>,
147 pub intervention: Option<eredu_core::intervention::InterventionPlan>,
149 pub host_bytes: Option<u64>,
152 pub continuation_growth_bytes: Option<u64>,
155}
156
157pub struct TextContinuationBranch<B: TextSnapshotBackend, C: TokenFilterController> {
160 continuation: ManagedTextContinuation<B, C>,
161 native: B::NativeTextState,
162}
163
164pub struct ManagedTextContinuation<B: eredu_core::TextGenerationBackend, C: TokenFilterController> {
168 state: eredu_core::TextGenerationContinuation<B, C>,
169 reservation: Option<SnapshotReservation>,
170}
171
172impl<B: eredu_core::TextGenerationBackend, C: TokenFilterController> ManagedTextContinuation<B, C> {
173 pub fn root(state: eredu_core::TextGenerationContinuation<B, C>) -> Self {
176 Self {
177 state,
178 reservation: None,
179 }
180 }
181 #[allow(clippy::type_complexity)]
183 pub fn advance(
184 &mut self,
185 driver: &mut eredu_core::TextGenerationDriver<'_, B>,
186 ) -> Result<
187 Option<eredu_core::ControlledToken<B::Token>>,
188 eredu_core::TextContinuationError<B::Error, C::Error>,
189 > {
190 driver.advance(&mut self.state)
191 }
192 pub fn take_completed_step(
194 &mut self,
195 driver: &mut eredu_core::TextGenerationDriver<'_, B>,
196 ) -> Result<
197 Option<eredu_core::capture::CapturedStep>,
198 eredu_core::TextContinuationError<B::Error, C::Error>,
199 > {
200 driver.take_completed_step(&mut self.state)
201 }
202 pub fn boundary<'d, 's>(
204 &'s mut self,
205 driver: &'d mut eredu_core::TextGenerationDriver<'_, B>,
206 ) -> Result<
207 TextContinuationBoundary<'d, 's, B, C>,
208 eredu_core::TextContinuationError<B::Error, C::Error>,
209 > {
210 driver.quiescent(&mut self.state)
211 }
212 pub fn controller(&self) -> &C {
214 self.state.controller()
215 }
216 pub fn controller_mut(&mut self) -> &mut C {
218 self.state.controller_mut()
219 }
220 pub fn retained_branch_bytes(&self) -> Option<u64> {
222 self.reservation
223 .as_ref()
224 .map(SnapshotReservation::retained_bytes)
225 }
226}
227impl<B: TextSnapshotBackend, C: TokenFilterController> TextContinuationBranch<B, C> {
228 pub fn exchange(
231 &mut self,
232 driver: &mut eredu_core::TextGenerationDriver<'_, B>,
233 active: &mut ManagedTextContinuation<B, C>,
234 ) -> Result<(), eredu_core::TextContinuationError<B::Error, C::Error>> {
235 driver
236 .quiescent(&mut active.state)?
237 .exchange_branch(&mut self.continuation.state, &mut self.native)?;
238 std::mem::swap(&mut active.reservation, &mut self.continuation.reservation);
239 Ok(())
240 }
241}
242
243pub struct TextContinuationSnapshot<B: TextSnapshotBackend, C: TokenFilterController> {
247 driver: eredu_core::TextDriverIdentity,
248 identity: TextContinuationIdentity,
249 native: B::NativeTextState,
250 sampling: B::SamplingState,
251 pending: Option<PendingTextInput<B::Prompt, B::Token>>,
252 controller: C,
253 remaining_tokens: Option<usize>,
254 capture: Option<CaptureCheckpoint>,
255 host_bytes: u64,
256 _reservation: SnapshotReservation,
257}
258
259impl<B: TextSnapshotBackend, C: SnapshotTokenController> TextContinuationSnapshot<B, C> {
260 pub fn retained_bytes(&self) -> u64 {
262 self._reservation.retained_bytes()
263 }
264 pub fn next_prediction(&self) -> u64 {
266 B::sampling_prediction(&self.sampling)
267 }
268
269 pub fn controller(&self) -> &C {
271 &self.controller
272 }
273
274 pub fn capture_checkpoint(&self) -> Option<&CaptureCheckpoint> {
276 self.capture.as_ref()
277 }
278
279 pub fn native_continuation_growth(
283 &self,
284 runtime: &ModelRuntime<B>,
285 max_predictions: u64,
286 ) -> Result<u64, TextSnapshotError<B::Error>> {
287 let predictions = max_predictions
288 .checked_sub(self.next_prediction())
289 .ok_or(TextSnapshotError::InconsistentState)?;
290 let input = B::continuation_input_tokens(
291 self.pending.as_ref().map(PendingTextInput::as_ref),
292 predictions,
293 )
294 .ok_or(ExecutionControlError::UnknownEstimate)?;
295 let native = B::estimate_native_text_growth(runtime, &self.native, input)
296 .map_err(TextSnapshotError::Backend)?
297 .ok_or(ExecutionControlError::UnknownEstimate)?;
298 let sampling = B::estimate_sampling_growth(runtime, &self.sampling, predictions)
299 .map_err(TextSnapshotError::Backend)?
300 .ok_or(ExecutionControlError::UnknownEstimate)?;
301 native
302 .checked_add(sampling)
303 .ok_or_else(|| ExecutionControlError::Overflow.into())
304 }
305
306 pub fn capture(
310 boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
311 budget: &SnapshotBudget,
312 host_bytes: Option<u64>,
313 ) -> Result<Self, TextSnapshotError<B::Error>> {
314 let host_bytes = host_bytes.ok_or(ExecutionControlError::UnknownEstimate)?;
315 let identity = boundary.identity();
316 let driver = boundary.driver_identity();
317 let remaining_tokens = boundary.remaining_tokens();
318 let (runtime, state, pending) = boundary.parts();
319 let discovery = B::capture_run(state)
320 .map(|_| B::capture_discovery(runtime))
321 .transpose()?;
322 let capture_bytes = match (B::capture_run(state), discovery.as_ref()) {
323 (Some(run), Some(discovery)) => run
324 .checkpoint_storage_bytes(discovery)
325 .ok_or(ExecutionControlError::UnknownEstimate)?,
326 _ => 0,
327 };
328 let estimate = combine_estimates(
329 [
330 B::estimate_native_text_state(runtime, None).map_err(TextSnapshotError::Backend)?,
331 B::estimate_sampling_state(runtime, B::sampling_state(state))
332 .map_err(TextSnapshotError::Backend)?,
333 B::estimate_pending_input(runtime, pending).map_err(TextSnapshotError::Backend)?,
334 ],
335 host_bytes
336 .checked_add(
337 boundary
338 .controller()
339 .snapshot_storage_bytes()
340 .ok_or(ExecutionControlError::UnknownEstimate)?,
341 )
342 .ok_or(ExecutionControlError::Overflow)?,
343 capture_bytes,
344 std::mem::size_of::<Self>(),
345 )?;
346 let reservation = budget.reserve(SnapshotResourceKind::Snapshot, Some(estimate))?;
347 let controller = boundary
350 .controller()
351 .fork_snapshot()
352 .map_err(TextSnapshotError::Controller)?;
353 let (runtime, state, pending) = boundary.mechanism_parts();
354 let capture = match (B::capture_run(state), discovery.as_ref()) {
355 (Some(run), Some(discovery)) => Some(run.checkpoint(discovery)?),
356 _ => None,
357 };
358 if capture.as_ref().is_some_and(|capture| {
359 capture.next_prediction() != B::sampling_prediction(B::sampling_state(state))
360 }) {
361 return Err(TextSnapshotError::InconsistentState);
362 }
363 let pending =
364 B::copy_pending_input(runtime, pending).map_err(TextSnapshotError::Backend)?;
365 let sampling = B::copy_sampling_state(runtime, B::sampling_state(state))
366 .map_err(TextSnapshotError::Backend)?;
367 let native = B::capture_native_text_state(runtime).map_err(TextSnapshotError::Backend)?;
368 Ok(Self {
369 driver,
370 identity,
371 native,
372 sampling,
373 pending,
374 controller,
375 remaining_tokens,
376 capture,
377 host_bytes,
378 _reservation: reservation,
379 })
380 }
381
382 fn copy_estimate(
383 &self,
384 runtime: &ModelRuntime<B>,
385 ) -> Result<SnapshotEstimate, TextSnapshotError<B::Error>> {
386 combine_estimates(
387 [
388 B::estimate_native_text_state(runtime, Some(&self.native))
389 .map_err(TextSnapshotError::Backend)?,
390 B::estimate_sampling_state(runtime, &self.sampling)
391 .map_err(TextSnapshotError::Backend)?,
392 B::estimate_pending_input(
393 runtime,
394 self.pending.as_ref().map(PendingTextInput::as_ref),
395 )
396 .map_err(TextSnapshotError::Backend)?,
397 ],
398 self.host_bytes
399 .checked_add(
400 self.controller
401 .snapshot_storage_bytes()
402 .ok_or(ExecutionControlError::UnknownEstimate)?,
403 )
404 .ok_or(ExecutionControlError::Overflow)?,
405 match &self.capture {
406 Some(capture) => capture
407 .logical_storage_bytes()
408 .ok_or(ExecutionControlError::UnknownEstimate)?,
409 None => 0,
410 },
411 std::mem::size_of::<Self>(),
412 )
413 .map_err(Into::into)
414 }
415
416 pub fn restore(
421 &self,
422 boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
423 budget: &SnapshotBudget,
424 ) -> Result<(), TextSnapshotError<B::Error>> {
425 self.restore_with(boundary, budget, || Ok(()))
426 }
427
428 pub fn restore_with<H>(
433 &self,
434 boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
435 budget: &SnapshotBudget,
436 prepare_host: impl FnOnce() -> Result<H, String>,
437 ) -> Result<H, TextSnapshotError<B::Error>> {
438 if boundary.identity() != self.identity {
439 return Err(TextSnapshotError::IncompatibleRun);
440 }
441 let (runtime, state, _) = boundary.parts();
442 B::validate_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
443 match (B::capture_run(state), &self.capture) {
444 (Some(run), Some(saved)) => run.validate_restore(saved)?,
445 (None, None) => {}
446 _ => return Err(TextSnapshotError::InconsistentState),
447 }
448 let _reservation = budget.reserve(
449 SnapshotResourceKind::Restore,
450 Some(self.copy_estimate(runtime)?),
451 )?;
452 let host = prepare_host().map_err(TextSnapshotError::Host)?;
453 let controller = self
454 .controller
455 .fork_snapshot()
456 .map_err(TextSnapshotError::Controller)?;
457 let (runtime, state, _) = boundary.mechanism_parts();
458 let pending =
459 B::copy_pending_input(runtime, self.pending.as_ref().map(PendingTextInput::as_ref))
460 .map_err(TextSnapshotError::Backend)?;
461 let sampling =
462 B::copy_sampling_state(runtime, &self.sampling).map_err(TextSnapshotError::Backend)?;
463 let mut native =
464 B::copy_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
465 let capture_restore = match (B::capture_run_mut(state), &self.capture) {
466 (Some(run), Some(saved)) => Some(run.prepare_restore(saved)?),
467 (None, None) => None,
468 _ => return Err(TextSnapshotError::InconsistentState),
469 };
470 B::exchange_native_text_state(runtime, &mut native).map_err(TextSnapshotError::Backend)?;
471 if let Some(restore) = capture_restore {
472 restore.commit();
473 }
474 B::install_sampling_state(state, sampling);
475 boundary.install_host_state(controller, pending, self.remaining_tokens);
476 Ok(host)
477 }
478
479 pub fn fork(
483 &self,
484 boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
485 budget: &SnapshotBudget,
486 request: TextBranchRequest<'_>,
487 ) -> Result<TextContinuationBranch<B, C>, TextSnapshotError<B::Error>> {
488 self.fork_with(boundary, budget, request, |_, _| Ok(()))
489 .map(|(branch, ())| branch)
490 }
491
492 #[allow(clippy::type_complexity)]
497 pub fn fork_with<H>(
498 &self,
499 boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
500 budget: &SnapshotBudget,
501 request: TextBranchRequest<'_>,
502 prepare: impl FnOnce(
503 &mut ModelRuntime<B>,
504 &mut B::TextGenerationState,
505 ) -> Result<H, TextSnapshotError<B::Error>>,
506 ) -> Result<(TextContinuationBranch<B, C>, H), TextSnapshotError<B::Error>> {
507 if self.driver != boundary.driver_identity() {
508 return Err(TextSnapshotError::IncompatibleRun);
509 }
510 if request.session_id.is_empty() || request.max_predictions < self.next_prediction() {
511 return Err(TextSnapshotError::InconsistentState);
512 }
513 let host_bytes = request
514 .host_bytes
515 .ok_or(ExecutionControlError::UnknownEstimate)?;
516 let growth_bytes = request
517 .continuation_growth_bytes
518 .ok_or(ExecutionControlError::UnknownEstimate)?;
519 let remaining = usize::try_from(request.max_predictions - self.next_prediction())
520 .map_err(|_| ExecutionControlError::Overflow)?;
521 let (runtime, _, _) = boundary.parts();
522 B::validate_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
523 if self.capture.is_none() && request.intervention.is_some() {
524 return Err(TextSnapshotError::Unsupported(
525 "adding interventions requires retained request admission geometry",
526 ));
527 }
528 let discovery = self
529 .capture
530 .as_ref()
531 .map(|_| B::capture_discovery(runtime))
532 .transpose()?;
533 let needs_intervention = self
534 .capture
535 .as_ref()
536 .is_some_and(|saved| saved.intervention_plan().is_some())
537 || request.intervention.is_some();
538 let intervention_discovery = needs_intervention
539 .then(|| B::intervention_discovery(runtime))
540 .transpose()?;
541 let child = match discovery.as_ref() {
542 Some(discovery) => Some(CaptureForkRequest {
543 discovery,
544 max_predictions: request.max_predictions,
545 limits: request
546 .capture_limits
547 .ok_or(TextSnapshotError::Unsupported(
548 "captured branches require explicit child limits",
549 ))?,
550 intervention: match intervention_discovery.as_ref() {
551 Some(discovery) => Some(InterventionForkRequest {
552 discovery,
553 session_id: request.session_id,
554 replacement: request.intervention,
555 estimator: B::child_intervention_estimator(runtime)?,
556 }),
557 None => None,
558 },
559 }),
560 None => None,
561 };
562 let child_bytes = match (&self.capture, &child) {
563 (Some(saved), Some(child)) => saved
564 .fork_storage_bytes(child)
565 .ok_or(ExecutionControlError::UnknownEstimate)?,
566 _ => 0,
567 };
568 let mut estimate = self.copy_estimate(runtime)?;
569 let extra = host_bytes
570 .checked_add(child_bytes)
571 .ok_or(ExecutionControlError::Overflow)?;
572 estimate.retained_bytes = estimate
573 .retained_bytes
574 .checked_add(extra)
575 .and_then(|n| n.checked_add(growth_bytes))
576 .ok_or(ExecutionControlError::Overflow)?;
577 estimate.copy_bytes = estimate
578 .copy_bytes
579 .checked_add(extra)
580 .ok_or(ExecutionControlError::Overflow)?;
581 let reservation = budget.reserve(SnapshotResourceKind::Branch, Some(estimate))?;
582 let capture = match (&self.capture, child) {
583 (Some(saved), Some(child)) => Some(saved.fork(child, |shape, selection, slice| {
584 B::estimate_child_capture(runtime, shape, selection, slice)
585 })?),
586 _ => None,
587 };
588 let controller = self
589 .controller
590 .fork_snapshot()
591 .map_err(TextSnapshotError::Controller)?;
592 let (runtime, _, _) = boundary.mechanism_parts();
593 let pending =
594 B::copy_pending_input(runtime, self.pending.as_ref().map(PendingTextInput::as_ref))
595 .map_err(TextSnapshotError::Backend)?;
596 let sampling =
597 B::copy_sampling_state(runtime, &self.sampling).map_err(TextSnapshotError::Backend)?;
598 let native =
599 B::copy_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
600 let mut generation = B::assemble_generation_state(sampling, capture);
601 let host = prepare(runtime, &mut generation)?;
602 let state = boundary.fork_host_state(generation, controller, pending, Some(remaining));
603 Ok((
604 TextContinuationBranch {
605 continuation: ManagedTextContinuation {
606 state,
607 reservation: Some(reservation),
608 },
609 native,
610 },
611 host,
612 ))
613 }
614}
615
616fn combine_estimates<const N: usize>(
617 estimates: [Option<SnapshotEstimate>; N],
618 host: u64,
619 capture: u64,
620 inline: usize,
621) -> Result<SnapshotEstimate, ExecutionControlError> {
622 let base = host
623 .checked_add(capture)
624 .and_then(|n| n.checked_add(u64::try_from(inline).ok()?))
625 .ok_or(ExecutionControlError::Overflow)?;
626 let mut total = SnapshotEstimate {
627 retained_bytes: base,
628 copy_bytes: base,
629 };
630 for estimate in estimates {
631 let estimate = estimate.ok_or(ExecutionControlError::UnknownEstimate)?;
632 total.retained_bytes = total
633 .retained_bytes
634 .checked_add(estimate.retained_bytes)
635 .ok_or(ExecutionControlError::Overflow)?;
636 total.copy_bytes = total
637 .copy_bytes
638 .checked_add(estimate.copy_bytes)
639 .ok_or(ExecutionControlError::Overflow)?;
640 }
641 Ok(total)
642}