use eredu_core::{capture::*, execution_control::*, TextGenerationDriver, TokenFilterController};
use eredu_runtime::execution_control::{
apply_sampling_override, ManagedTextContinuation, SamplingOverride, SnapshotBudget,
SnapshotTokenController, TextBranchRequest, TextContinuationSnapshot,
TextSamplingControlBackend, TextSnapshotBackend, TextSnapshotError, TokenChoiceController,
};
use std::fmt::Debug;
#[cfg(test)]
mod tests;
pub struct ContinuationFixtureLimits {
pub host_bytes: u64,
pub growth_bytes: u64,
pub max_predictions: u64,
pub capture: Option<CaptureLimits>,
}
fn step<B: TextSnapshotBackend, C: TokenFilterController>(
driver: &mut TextGenerationDriver<'_, B>,
state: &mut ManagedTextContinuation<B, C>,
) -> (u32, Option<CapturedStep>) {
let token = state.advance(driver).unwrap().expect("fixture ended early");
assert!(matches!(
state.boundary(driver),
Err(eredu_core::TextContinuationError::NotQuiescent)
));
let records = state.take_completed_step(driver).unwrap().map(|mut step| {
step.capture_seconds = 0.0;
step.cumulative_usage = Default::default();
step
});
(token.token_id(), records)
}
fn branch_values(mut value: (u32, Option<CapturedStep>)) -> (u32, Option<CapturedStep>) {
if let Some(step) = &mut value.1 {
for record in &mut step.interventions {
record.plan_id.clear();
}
}
value
}
pub fn sampling_override_conformance<B, C, P>(
driver: &mut TextGenerationDriver<'_, B>,
state: &mut ManagedTextContinuation<B, C>,
limits: &ContinuationFixtureLimits,
probe: impl Fn() -> P,
) where
B: TextSnapshotBackend + TextSamplingControlBackend,
C: SnapshotTokenController + Clone + PartialEq + Debug,
P: PartialEq + Debug,
{
let budget = SnapshotBudget::new(SnapshotLimits {
max_snapshots: 2,
max_branches: 1,
retained_bytes: 64_000_000,
cumulative_copy_bytes: 512_000_000,
});
let initial = TextContinuationSnapshot::capture(
&mut state.boundary(driver).unwrap(),
&budget,
Some(limits.host_bytes),
)
.unwrap();
let facts = B::sampling_control_facts(state.boundary(driver).unwrap().parts().1);
let baseline: Vec<_> = (0..3).map(|_| step(driver, state)).collect();
initial
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
let before = probe();
assert!(apply_sampling_override(
&mut state.boundary(driver).unwrap(),
SamplingOverride {
temperature: Some(f32::NAN),
reseed: Some(123),
}
)
.is_err());
let greedy = apply_sampling_override(
&mut state.boundary(driver).unwrap(),
SamplingOverride {
temperature: Some(0.0),
reseed: None,
},
);
if facts.requires_positive_temperature {
assert!(greedy.is_err());
} else {
assert!(
greedy.unwrap().has_rng,
"temporary greedy selection discarded inherited RNG"
);
}
apply_sampling_override(
&mut state.boundary(driver).unwrap(),
SamplingOverride {
temperature: Some(facts.temperature),
reseed: None,
},
)
.unwrap();
assert_eq!(probe(), before);
let actual: Vec<_> = (0..3).map(|_| step(driver, state)).collect();
assert_eq!(
actual, baseline,
"invalid/no-draw changes lost RNG or sampler state"
);
initial
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
let mut child = initial
.fork(
&mut state.boundary(driver).unwrap(),
&budget,
TextBranchRequest {
session_id: "conformance-sampling-child",
max_predictions: limits.max_predictions,
capture_limits: limits.capture.clone(),
intervention: None,
host_bytes: Some(limits.host_bytes),
continuation_growth_bytes: Some(limits.growth_bytes),
},
)
.unwrap();
child.exchange(driver, state).unwrap();
let before = probe();
let updated = apply_sampling_override(
&mut state.boundary(driver).unwrap(),
SamplingOverride {
temperature: Some(0.35),
reseed: Some(711),
},
)
.unwrap();
assert_eq!(updated.temperature, 0.35);
assert!(updated.has_rng);
let modified = TextContinuationSnapshot::capture(
&mut state.boundary(driver).unwrap(),
&budget,
Some(limits.host_bytes),
)
.unwrap();
assert_eq!(probe(), before, "sampling change executed model work");
let changed: Vec<_> = (0..3).map(|_| step(driver, state)).collect();
modified
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
assert_eq!(
B::sampling_control_facts(state.boundary(driver).unwrap().parts().1),
updated
);
let again: Vec<_> = (0..3).map(|_| step(driver, state)).collect();
assert_eq!(
again, changed,
"snapshot lost explicit reseed or temperature change"
);
child.exchange(driver, state).unwrap();
assert_eq!(
B::sampling_control_facts(state.boundary(driver).unwrap().parts().1),
facts
);
let parent: Vec<_> = (0..3).map(|_| step(driver, state)).collect();
assert_eq!(parent, baseline, "sampling override changed the parent");
initial
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
}
pub fn forced_choice_conformance<B, C, P>(
driver: &mut TextGenerationDriver<'_, B>,
state: &mut ManagedTextContinuation<B, TokenChoiceController<C>>,
limits: &ContinuationFixtureLimits,
vocabulary: usize,
probe: impl Fn() -> P,
) where
B: TextSnapshotBackend,
C: SnapshotTokenController + Clone + PartialEq + Debug,
P: PartialEq + Debug,
{
let budget = SnapshotBudget::new(SnapshotLimits {
max_snapshots: 2,
max_branches: 1,
retained_bytes: 64_000_000,
cumulative_copy_bytes: 512_000_000,
});
let initial = TextContinuationSnapshot::capture(
&mut state.boundary(driver).unwrap(),
&budget,
Some(limits.host_bytes),
)
.unwrap();
let baseline = step(driver, state).0;
initial
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
let filter = state.controller_mut().current_filter().unwrap();
let alternative = (0..vocabulary as u32)
.find(|&token| {
token != baseline
&& filter
.allowed_mask()
.is_none_or(|mask| mask[token as usize])
})
.expect("fixture must admit an alternative canonical token");
let mut expected = state.controller().inner().clone();
expected.commit_token(alternative).unwrap();
let before = probe();
let mut child = initial
.fork(
&mut state.boundary(driver).unwrap(),
&budget,
TextBranchRequest {
session_id: "conformance-forced-child",
max_predictions: limits.max_predictions,
capture_limits: limits.capture.clone(),
intervention: None,
host_bytes: Some(limits.host_bytes),
continuation_growth_bytes: Some(limits.growth_bytes),
},
)
.unwrap();
child.exchange(driver, state).unwrap();
state.controller_mut().force_next(alternative).unwrap();
let pending = TextContinuationSnapshot::capture(
&mut state.boundary(driver).unwrap(),
&budget,
Some(limits.host_bytes),
)
.unwrap();
assert_eq!(probe(), before, "choice/fork/capture executed the model");
assert_eq!(step(driver, state).0, alternative);
assert!(state.controller().last_committed_was_forced());
assert_eq!(state.controller().inner(), &expected);
let continuation: Vec<_> = (0..3).map(|_| step(driver, state)).collect();
assert!(!state.controller().last_committed_was_forced());
for _ in 0..2 {
let before = probe();
pending
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
assert_eq!(probe(), before, "restore replayed the forced prefix");
assert_eq!(state.controller().pending_forced(), Some(alternative));
assert_eq!(step(driver, state).0, alternative);
assert_eq!(state.controller().inner(), &expected);
let actual: Vec<_> = (0..3).map(|_| step(driver, state)).collect();
assert_eq!(actual, continuation);
}
child.exchange(driver, state).unwrap();
assert_eq!(state.controller().pending_forced(), None);
assert_eq!(step(driver, state).0, baseline, "child changed the parent");
initial
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
}
pub fn continuation_conformance<B, C, P>(
driver: &mut TextGenerationDriver<'_, B>,
mut state: ManagedTextContinuation<B, C>,
limits: ContinuationFixtureLimits,
probe: impl Fn() -> P,
) where
B: TextSnapshotBackend,
C: SnapshotTokenController + Clone + PartialEq + Debug,
P: PartialEq + Debug,
{
assert!(limits.max_predictions >= 8);
let budget = SnapshotBudget::new(SnapshotLimits {
max_snapshots: 2,
max_branches: 2,
retained_bytes: 64_000_000,
cumulative_copy_bytes: 512_000_000,
});
let before = probe();
assert!(matches!(
TextContinuationSnapshot::capture(&mut state.boundary(driver).unwrap(), &budget, None),
Err(TextSnapshotError::Control(
ExecutionControlError::UnknownEstimate
))
));
assert_eq!(budget.usage(), SnapshotUsage::default());
assert_eq!(probe(), before);
let initial = TextContinuationSnapshot::capture(
&mut state.boundary(driver).unwrap(),
&budget,
Some(limits.host_bytes),
)
.unwrap();
assert_eq!(initial.next_prediction(), 0);
let initial_native = B::estimate_native_text_state(driver.runtime(), None)
.unwrap()
.unwrap()
.retained_bytes;
let initial_growth = initial
.native_continuation_growth(driver.runtime(), limits.max_predictions)
.unwrap();
assert_eq!(
probe(),
before,
"growth estimation executed or rebuilt the model"
);
let prefix: Vec<_> = (0..3).map(|_| step(driver, &mut state)).collect();
let prefix_controller = state.controller().clone();
let saved = TextContinuationSnapshot::capture(
&mut state.boundary(driver).unwrap(),
&budget,
Some(limits.host_bytes),
)
.unwrap();
assert_eq!(saved.next_prediction(), 3);
let usage = budget.usage();
let before = probe();
assert!(matches!(
TextContinuationSnapshot::capture(
&mut state.boundary(driver).unwrap(),
&budget,
Some(limits.host_bytes),
),
Err(TextSnapshotError::Control(ExecutionControlError::Limit(
"snapshot count"
)))
));
assert_eq!(budget.usage(), usage);
let native_growth = saved
.native_continuation_growth(driver.runtime(), limits.max_predictions)
.unwrap();
let child = |session_id| TextBranchRequest {
session_id,
max_predictions: limits.max_predictions,
capture_limits: limits.capture.clone(),
intervention: None,
host_bytes: Some(limits.host_bytes),
continuation_growth_bytes: Some(limits.growth_bytes.checked_add(native_growth).unwrap()),
};
let mut left = saved
.fork(
&mut state.boundary(driver).unwrap(),
&budget,
child("conformance-left"),
)
.unwrap();
let mut right = saved
.fork(
&mut state.boundary(driver).unwrap(),
&budget,
child("conformance-right"),
)
.unwrap();
assert_eq!(probe(), before, "copy/fork executed or rebuilt the model");
let baseline: Vec<_> = (0..5).map(|_| step(driver, &mut state)).collect();
let native_after = B::estimate_native_text_state(driver.runtime(), None)
.unwrap()
.unwrap()
.retained_bytes;
assert!(
native_after <= initial_native.checked_add(initial_growth).unwrap(),
"retained native state exceeded the pre-generation growth allowance"
);
left.exchange(driver, &mut state).unwrap();
assert!(matches!(
saved.restore(&mut state.boundary(driver).unwrap(), &budget),
Err(TextSnapshotError::IncompatibleRun)
));
if let Some(checkpoint) = saved.capture_checkpoint() {
let boundary = state.boundary(driver).unwrap();
assert_eq!(
B::capture_run(boundary.parts().1)
.unwrap()
.cumulative_usage(),
checkpoint.inherited_usage()
);
}
let first_left = step(driver, &mut state);
if let (Some(child), Some(parent)) = (&first_left.1, &baseline[0].1) {
for (child, parent) in child.interventions.iter().zip(&parent.interventions) {
assert_ne!(
child.plan_id, parent.plan_id,
"child reused a session-bound admission"
);
}
}
assert_eq!(
branch_values(first_left),
branch_values(baseline[0].clone())
);
left.exchange(driver, &mut state).unwrap();
right.exchange(driver, &mut state).unwrap();
let other: Vec<_> = (0..5)
.map(|_| branch_values(step(driver, &mut state)))
.collect();
assert_eq!(
other,
baseline
.iter()
.cloned()
.map(branch_values)
.collect::<Vec<_>>()
);
right.exchange(driver, &mut state).unwrap();
left.exchange(driver, &mut state).unwrap();
let rest: Vec<_> = (0..4)
.map(|_| branch_values(step(driver, &mut state)))
.collect();
assert_eq!(
rest,
baseline[1..]
.iter()
.cloned()
.map(branch_values)
.collect::<Vec<_>>()
);
left.exchange(driver, &mut state).unwrap();
let capture_usage = {
let boundary = state.boundary(driver).unwrap();
B::capture_run(boundary.parts().1).map(|run| run.cumulative_usage())
};
for _ in 0..2 {
let before = probe();
saved
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
assert_eq!(probe(), before, "restore executed or rebuilt the model");
assert_eq!(state.controller(), &prefix_controller);
let actual: Vec<_> = (0..5).map(|_| step(driver, &mut state)).collect();
assert_eq!(actual, baseline);
}
if let Some(before) = capture_usage {
let boundary = state.boundary(driver).unwrap();
let after = B::capture_run(boundary.parts().1)
.unwrap()
.cumulative_usage();
assert!(after.captures >= before.captures);
assert!(after.encoded_bytes >= before.encoded_bytes);
if before.encoded_bytes != 0 {
assert!(after.encoded_bytes > before.encoded_bytes);
}
}
initial
.restore(&mut state.boundary(driver).unwrap(), &budget)
.unwrap();
let actual: Vec<_> = (0..3).map(|_| step(driver, &mut state)).collect();
assert_eq!(actual, prefix);
drop(right);
assert_eq!(budget.usage().branches, 1);
left.exchange(driver, &mut state).unwrap();
assert!(state.retained_branch_bytes().is_some());
let usage = budget.usage();
drop(left); assert_eq!(
budget.usage(),
usage,
"dropping parent released active child retention"
);
drop(state);
assert_eq!(budget.usage().branches, 0);
drop((initial, saved));
assert_eq!(budget.usage().retained_bytes, 0);
assert_eq!(budget.usage().snapshots, 0);
assert_eq!(
budget.usage().cumulative_copy_bytes,
usage.cumulative_copy_bytes
);
}