use super::{SnapshotBudget, SnapshotReservation};
use crate::capture::{
CaptureCheckpoint, CaptureForkRequest, CaptureSession, InterventionForkRequest,
};
use eredu_core::{
capture::CaptureError,
execution_control::{
ExecutionControlError, NativeTextStateBackend, SnapshotEstimate, SnapshotResourceKind,
},
ModelRuntime, PendingTextInput, TextContinuationBoundary, TextContinuationIdentity,
TokenFilterController,
};
pub trait TextSnapshotBackend: NativeTextStateBackend {
type SamplingState;
fn sampling_state(state: &Self::TextGenerationState) -> &Self::SamplingState;
fn install_sampling_state(state: &mut Self::TextGenerationState, sampling: Self::SamplingState);
fn assemble_generation_state(
sampling: Self::SamplingState,
capture: Option<CaptureSession>,
) -> Self::TextGenerationState;
fn sampling_prediction(sampling: &Self::SamplingState) -> u64;
fn estimate_sampling_state(
runtime: &ModelRuntime<Self>,
sampling: &Self::SamplingState,
) -> Result<Option<SnapshotEstimate>, Self::Error>;
fn copy_sampling_state(
runtime: &mut ModelRuntime<Self>,
sampling: &Self::SamplingState,
) -> Result<Self::SamplingState, Self::Error>;
fn estimate_pending_input(
runtime: &ModelRuntime<Self>,
input: Option<PendingTextInput<&Self::Prompt, &Self::Token>>,
) -> Result<Option<SnapshotEstimate>, Self::Error>;
fn continuation_input_tokens(
_input: Option<PendingTextInput<&Self::Prompt, &Self::Token>>,
_predictions: u64,
) -> Option<u64> {
None
}
fn estimate_sampling_growth(
_runtime: &ModelRuntime<Self>,
_sampling: &Self::SamplingState,
_predictions: u64,
) -> Result<Option<u64>, Self::Error> {
Ok(None)
}
#[allow(clippy::type_complexity)]
fn copy_pending_input(
runtime: &mut ModelRuntime<Self>,
input: Option<PendingTextInput<&Self::Prompt, &Self::Token>>,
) -> Result<Option<PendingTextInput<Self::Prompt, Self::Token>>, Self::Error>;
fn capture_run(state: &Self::TextGenerationState) -> Option<&CaptureSession>;
fn capture_run_mut(state: &mut Self::TextGenerationState) -> Option<&mut CaptureSession>;
fn estimate_child_capture(
runtime: &ModelRuntime<Self>,
shape: &[u64],
selection: &eredu_core::capture::CaptureSelection,
slice: &eredu_core::capture::ResolvedCaptureSlice,
) -> Result<eredu_core::capture::CaptureUsage, CaptureError>;
fn child_intervention_estimator(
runtime: &ModelRuntime<Self>,
) -> Result<std::sync::Arc<dyn eredu_core::intervention::InterventionEstimator>, CaptureError>;
}
pub trait SnapshotTokenController: TokenFilterController + Sized {
fn snapshot_storage_bytes(&self) -> Option<u64>;
fn fork_snapshot(&self) -> Result<Self, String>;
}
#[derive(Debug, thiserror::Error)]
pub enum TextSnapshotError<E: std::error::Error + 'static> {
#[error("host continuation snapshot failed: {0}")]
Host(String),
#[error("constraint snapshot failed: {0}")]
Controller(String),
#[error("snapshot backend operation failed: {0}")]
Backend(#[source] E),
#[error(transparent)]
Capture(#[from] CaptureError),
#[error(transparent)]
Control(#[from] ExecutionControlError),
#[error("snapshot belongs to another continuation")]
IncompatibleRun,
#[error("snapshot prediction or capture ownership differs")]
InconsistentState,
#[error("unsupported continuation: {0}")]
Unsupported(&'static str),
}
pub struct TextBranchRequest<'a> {
pub session_id: &'a str,
pub max_predictions: u64,
pub capture_limits: Option<eredu_core::capture::CaptureLimits>,
pub intervention: Option<eredu_core::intervention::InterventionPlan>,
pub host_bytes: Option<u64>,
pub continuation_growth_bytes: Option<u64>,
}
pub struct TextContinuationBranch<B: TextSnapshotBackend, C: TokenFilterController> {
continuation: ManagedTextContinuation<B, C>,
native: B::NativeTextState,
}
pub struct ManagedTextContinuation<B: eredu_core::TextGenerationBackend, C: TokenFilterController> {
state: eredu_core::TextGenerationContinuation<B, C>,
reservation: Option<SnapshotReservation>,
}
impl<B: eredu_core::TextGenerationBackend, C: TokenFilterController> ManagedTextContinuation<B, C> {
pub fn root(state: eredu_core::TextGenerationContinuation<B, C>) -> Self {
Self {
state,
reservation: None,
}
}
#[allow(clippy::type_complexity)]
pub fn advance(
&mut self,
driver: &mut eredu_core::TextGenerationDriver<'_, B>,
) -> Result<
Option<eredu_core::ControlledToken<B::Token>>,
eredu_core::TextContinuationError<B::Error, C::Error>,
> {
driver.advance(&mut self.state)
}
pub fn take_completed_step(
&mut self,
driver: &mut eredu_core::TextGenerationDriver<'_, B>,
) -> Result<
Option<eredu_core::capture::CapturedStep>,
eredu_core::TextContinuationError<B::Error, C::Error>,
> {
driver.take_completed_step(&mut self.state)
}
pub fn boundary<'d, 's>(
&'s mut self,
driver: &'d mut eredu_core::TextGenerationDriver<'_, B>,
) -> Result<
TextContinuationBoundary<'d, 's, B, C>,
eredu_core::TextContinuationError<B::Error, C::Error>,
> {
driver.quiescent(&mut self.state)
}
pub fn controller(&self) -> &C {
self.state.controller()
}
pub fn controller_mut(&mut self) -> &mut C {
self.state.controller_mut()
}
pub fn retained_branch_bytes(&self) -> Option<u64> {
self.reservation
.as_ref()
.map(SnapshotReservation::retained_bytes)
}
}
impl<B: TextSnapshotBackend, C: TokenFilterController> TextContinuationBranch<B, C> {
pub fn exchange(
&mut self,
driver: &mut eredu_core::TextGenerationDriver<'_, B>,
active: &mut ManagedTextContinuation<B, C>,
) -> Result<(), eredu_core::TextContinuationError<B::Error, C::Error>> {
driver
.quiescent(&mut active.state)?
.exchange_branch(&mut self.continuation.state, &mut self.native)?;
std::mem::swap(&mut active.reservation, &mut self.continuation.reservation);
Ok(())
}
}
pub struct TextContinuationSnapshot<B: TextSnapshotBackend, C: TokenFilterController> {
driver: eredu_core::TextDriverIdentity,
identity: TextContinuationIdentity,
native: B::NativeTextState,
sampling: B::SamplingState,
pending: Option<PendingTextInput<B::Prompt, B::Token>>,
controller: C,
remaining_tokens: Option<usize>,
capture: Option<CaptureCheckpoint>,
host_bytes: u64,
_reservation: SnapshotReservation,
}
impl<B: TextSnapshotBackend, C: SnapshotTokenController> TextContinuationSnapshot<B, C> {
pub fn retained_bytes(&self) -> u64 {
self._reservation.retained_bytes()
}
pub fn next_prediction(&self) -> u64 {
B::sampling_prediction(&self.sampling)
}
pub fn controller(&self) -> &C {
&self.controller
}
pub fn capture_checkpoint(&self) -> Option<&CaptureCheckpoint> {
self.capture.as_ref()
}
pub fn native_continuation_growth(
&self,
runtime: &ModelRuntime<B>,
max_predictions: u64,
) -> Result<u64, TextSnapshotError<B::Error>> {
let predictions = max_predictions
.checked_sub(self.next_prediction())
.ok_or(TextSnapshotError::InconsistentState)?;
let input = B::continuation_input_tokens(
self.pending.as_ref().map(PendingTextInput::as_ref),
predictions,
)
.ok_or(ExecutionControlError::UnknownEstimate)?;
let native = B::estimate_native_text_growth(runtime, &self.native, input)
.map_err(TextSnapshotError::Backend)?
.ok_or(ExecutionControlError::UnknownEstimate)?;
let sampling = B::estimate_sampling_growth(runtime, &self.sampling, predictions)
.map_err(TextSnapshotError::Backend)?
.ok_or(ExecutionControlError::UnknownEstimate)?;
native
.checked_add(sampling)
.ok_or_else(|| ExecutionControlError::Overflow.into())
}
pub fn capture(
boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
budget: &SnapshotBudget,
host_bytes: Option<u64>,
) -> Result<Self, TextSnapshotError<B::Error>> {
let host_bytes = host_bytes.ok_or(ExecutionControlError::UnknownEstimate)?;
let identity = boundary.identity();
let driver = boundary.driver_identity();
let remaining_tokens = boundary.remaining_tokens();
let (runtime, state, pending) = boundary.parts();
let discovery = B::capture_run(state)
.map(|_| B::capture_discovery(runtime))
.transpose()?;
let capture_bytes = match (B::capture_run(state), discovery.as_ref()) {
(Some(run), Some(discovery)) => run
.checkpoint_storage_bytes(discovery)
.ok_or(ExecutionControlError::UnknownEstimate)?,
_ => 0,
};
let estimate = combine_estimates(
[
B::estimate_native_text_state(runtime, None).map_err(TextSnapshotError::Backend)?,
B::estimate_sampling_state(runtime, B::sampling_state(state))
.map_err(TextSnapshotError::Backend)?,
B::estimate_pending_input(runtime, pending).map_err(TextSnapshotError::Backend)?,
],
host_bytes
.checked_add(
boundary
.controller()
.snapshot_storage_bytes()
.ok_or(ExecutionControlError::UnknownEstimate)?,
)
.ok_or(ExecutionControlError::Overflow)?,
capture_bytes,
std::mem::size_of::<Self>(),
)?;
let reservation = budget.reserve(SnapshotResourceKind::Snapshot, Some(estimate))?;
let controller = boundary
.controller()
.fork_snapshot()
.map_err(TextSnapshotError::Controller)?;
let (runtime, state, pending) = boundary.mechanism_parts();
let capture = match (B::capture_run(state), discovery.as_ref()) {
(Some(run), Some(discovery)) => Some(run.checkpoint(discovery)?),
_ => None,
};
if capture.as_ref().is_some_and(|capture| {
capture.next_prediction() != B::sampling_prediction(B::sampling_state(state))
}) {
return Err(TextSnapshotError::InconsistentState);
}
let pending =
B::copy_pending_input(runtime, pending).map_err(TextSnapshotError::Backend)?;
let sampling = B::copy_sampling_state(runtime, B::sampling_state(state))
.map_err(TextSnapshotError::Backend)?;
let native = B::capture_native_text_state(runtime).map_err(TextSnapshotError::Backend)?;
Ok(Self {
driver,
identity,
native,
sampling,
pending,
controller,
remaining_tokens,
capture,
host_bytes,
_reservation: reservation,
})
}
fn copy_estimate(
&self,
runtime: &ModelRuntime<B>,
) -> Result<SnapshotEstimate, TextSnapshotError<B::Error>> {
combine_estimates(
[
B::estimate_native_text_state(runtime, Some(&self.native))
.map_err(TextSnapshotError::Backend)?,
B::estimate_sampling_state(runtime, &self.sampling)
.map_err(TextSnapshotError::Backend)?,
B::estimate_pending_input(
runtime,
self.pending.as_ref().map(PendingTextInput::as_ref),
)
.map_err(TextSnapshotError::Backend)?,
],
self.host_bytes
.checked_add(
self.controller
.snapshot_storage_bytes()
.ok_or(ExecutionControlError::UnknownEstimate)?,
)
.ok_or(ExecutionControlError::Overflow)?,
match &self.capture {
Some(capture) => capture
.logical_storage_bytes()
.ok_or(ExecutionControlError::UnknownEstimate)?,
None => 0,
},
std::mem::size_of::<Self>(),
)
.map_err(Into::into)
}
pub fn restore(
&self,
boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
budget: &SnapshotBudget,
) -> Result<(), TextSnapshotError<B::Error>> {
self.restore_with(boundary, budget, || Ok(()))
}
pub fn restore_with<H>(
&self,
boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
budget: &SnapshotBudget,
prepare_host: impl FnOnce() -> Result<H, String>,
) -> Result<H, TextSnapshotError<B::Error>> {
if boundary.identity() != self.identity {
return Err(TextSnapshotError::IncompatibleRun);
}
let (runtime, state, _) = boundary.parts();
B::validate_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
match (B::capture_run(state), &self.capture) {
(Some(run), Some(saved)) => run.validate_restore(saved)?,
(None, None) => {}
_ => return Err(TextSnapshotError::InconsistentState),
}
let _reservation = budget.reserve(
SnapshotResourceKind::Restore,
Some(self.copy_estimate(runtime)?),
)?;
let host = prepare_host().map_err(TextSnapshotError::Host)?;
let controller = self
.controller
.fork_snapshot()
.map_err(TextSnapshotError::Controller)?;
let (runtime, state, _) = boundary.mechanism_parts();
let pending =
B::copy_pending_input(runtime, self.pending.as_ref().map(PendingTextInput::as_ref))
.map_err(TextSnapshotError::Backend)?;
let sampling =
B::copy_sampling_state(runtime, &self.sampling).map_err(TextSnapshotError::Backend)?;
let mut native =
B::copy_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
let capture_restore = match (B::capture_run_mut(state), &self.capture) {
(Some(run), Some(saved)) => Some(run.prepare_restore(saved)?),
(None, None) => None,
_ => return Err(TextSnapshotError::InconsistentState),
};
B::exchange_native_text_state(runtime, &mut native).map_err(TextSnapshotError::Backend)?;
if let Some(restore) = capture_restore {
restore.commit();
}
B::install_sampling_state(state, sampling);
boundary.install_host_state(controller, pending, self.remaining_tokens);
Ok(host)
}
pub fn fork(
&self,
boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
budget: &SnapshotBudget,
request: TextBranchRequest<'_>,
) -> Result<TextContinuationBranch<B, C>, TextSnapshotError<B::Error>> {
self.fork_with(boundary, budget, request, |_, _| Ok(()))
.map(|(branch, ())| branch)
}
#[allow(clippy::type_complexity)]
pub fn fork_with<H>(
&self,
boundary: &mut TextContinuationBoundary<'_, '_, B, C>,
budget: &SnapshotBudget,
request: TextBranchRequest<'_>,
prepare: impl FnOnce(
&mut ModelRuntime<B>,
&mut B::TextGenerationState,
) -> Result<H, TextSnapshotError<B::Error>>,
) -> Result<(TextContinuationBranch<B, C>, H), TextSnapshotError<B::Error>> {
if self.driver != boundary.driver_identity() {
return Err(TextSnapshotError::IncompatibleRun);
}
if request.session_id.is_empty() || request.max_predictions < self.next_prediction() {
return Err(TextSnapshotError::InconsistentState);
}
let host_bytes = request
.host_bytes
.ok_or(ExecutionControlError::UnknownEstimate)?;
let growth_bytes = request
.continuation_growth_bytes
.ok_or(ExecutionControlError::UnknownEstimate)?;
let remaining = usize::try_from(request.max_predictions - self.next_prediction())
.map_err(|_| ExecutionControlError::Overflow)?;
let (runtime, _, _) = boundary.parts();
B::validate_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
if self.capture.is_none() && request.intervention.is_some() {
return Err(TextSnapshotError::Unsupported(
"adding interventions requires retained request admission geometry",
));
}
let discovery = self
.capture
.as_ref()
.map(|_| B::capture_discovery(runtime))
.transpose()?;
let needs_intervention = self
.capture
.as_ref()
.is_some_and(|saved| saved.intervention_plan().is_some())
|| request.intervention.is_some();
let intervention_discovery = needs_intervention
.then(|| B::intervention_discovery(runtime))
.transpose()?;
let child = match discovery.as_ref() {
Some(discovery) => Some(CaptureForkRequest {
discovery,
max_predictions: request.max_predictions,
limits: request
.capture_limits
.ok_or(TextSnapshotError::Unsupported(
"captured branches require explicit child limits",
))?,
intervention: match intervention_discovery.as_ref() {
Some(discovery) => Some(InterventionForkRequest {
discovery,
session_id: request.session_id,
replacement: request.intervention,
estimator: B::child_intervention_estimator(runtime)?,
}),
None => None,
},
}),
None => None,
};
let child_bytes = match (&self.capture, &child) {
(Some(saved), Some(child)) => saved
.fork_storage_bytes(child)
.ok_or(ExecutionControlError::UnknownEstimate)?,
_ => 0,
};
let mut estimate = self.copy_estimate(runtime)?;
let extra = host_bytes
.checked_add(child_bytes)
.ok_or(ExecutionControlError::Overflow)?;
estimate.retained_bytes = estimate
.retained_bytes
.checked_add(extra)
.and_then(|n| n.checked_add(growth_bytes))
.ok_or(ExecutionControlError::Overflow)?;
estimate.copy_bytes = estimate
.copy_bytes
.checked_add(extra)
.ok_or(ExecutionControlError::Overflow)?;
let reservation = budget.reserve(SnapshotResourceKind::Branch, Some(estimate))?;
let capture = match (&self.capture, child) {
(Some(saved), Some(child)) => Some(saved.fork(child, |shape, selection, slice| {
B::estimate_child_capture(runtime, shape, selection, slice)
})?),
_ => None,
};
let controller = self
.controller
.fork_snapshot()
.map_err(TextSnapshotError::Controller)?;
let (runtime, _, _) = boundary.mechanism_parts();
let pending =
B::copy_pending_input(runtime, self.pending.as_ref().map(PendingTextInput::as_ref))
.map_err(TextSnapshotError::Backend)?;
let sampling =
B::copy_sampling_state(runtime, &self.sampling).map_err(TextSnapshotError::Backend)?;
let native =
B::copy_native_text_state(runtime, &self.native).map_err(TextSnapshotError::Backend)?;
let mut generation = B::assemble_generation_state(sampling, capture);
let host = prepare(runtime, &mut generation)?;
let state = boundary.fork_host_state(generation, controller, pending, Some(remaining));
Ok((
TextContinuationBranch {
continuation: ManagedTextContinuation {
state,
reservation: Some(reservation),
},
native,
},
host,
))
}
}
fn combine_estimates<const N: usize>(
estimates: [Option<SnapshotEstimate>; N],
host: u64,
capture: u64,
inline: usize,
) -> Result<SnapshotEstimate, ExecutionControlError> {
let base = host
.checked_add(capture)
.and_then(|n| n.checked_add(u64::try_from(inline).ok()?))
.ok_or(ExecutionControlError::Overflow)?;
let mut total = SnapshotEstimate {
retained_bytes: base,
copy_bytes: base,
};
for estimate in estimates {
let estimate = estimate.ok_or(ExecutionControlError::UnknownEstimate)?;
total.retained_bytes = total
.retained_bytes
.checked_add(estimate.retained_bytes)
.ok_or(ExecutionControlError::Overflow)?;
total.copy_bytes = total
.copy_bytes
.checked_add(estimate.copy_bytes)
.ok_or(ExecutionControlError::Overflow)?;
}
Ok(total)
}