use std::{
num::NonZeroUsize,
sync::atomic::{AtomicU64, Ordering},
time::Instant,
};
use eredu_core::{
consensus::{validate_ranked_identity_bounded, BoundedConsensusTransport},
scheduler::{
RequestId, RequestStatus, Scheduler, SchedulerCapabilities, SchedulerError,
SchedulerLimits, SchedulerProgress, SchedulerReport, SemanticStateTransaction,
TransitionOutput, WorkId,
},
BoundedCompletionWait, Completion, ParallelTopology, RealtimeInputFrame, RealtimeSampling,
RealtimeSpeechConfig,
};
use sha2::{Digest, Sha256};
use crate::{
CommunicationCompletionPolicy, LayerWeightResidency, RealtimeGenerationBranch,
RealtimeGenerationState, RealtimeGenerationTransactionError, RealtimeIdentity,
RealtimeIngressContract, RealtimePayloadContract, RealtimePayloadGeneration,
RealtimePayloadOwnerIdentity, SelectedRealtimeRealization,
};
static NEXT_REALTIME_INCARNATION: AtomicU64 = AtomicU64::new(1);
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct RealtimeModelSessionIdentity {
selected: Option<Box<SelectedRealtimeRealization>>,
architecture: RealtimeIdentity,
source: RealtimeIdentity,
execution: RealtimeIdentity,
schedule_identity: RealtimeIdentity,
schedule: RealtimeSpeechConfig,
state_layout: RealtimeIdentity,
topology_policy: RealtimeIdentity,
topology: ParallelTopology,
rank: usize,
residency: LayerWeightResidency,
completion: CommunicationCompletionPolicy,
}
impl RealtimeModelSessionIdentity {
pub fn from_selected(selected: &SelectedRealtimeRealization) -> Self {
Self {
selected: Some(Box::new(selected.clone())),
architecture: selected.requirements().architecture().clone(),
source: selected.source().clone(),
execution: selected.execution().clone(),
schedule_identity: selected.requirements().speech_schedule_identity().clone(),
schedule: selected.requirements().speech_schedule().clone(),
state_layout: selected.requirements().state_layout_identity().clone(),
topology_policy: selected.requirements().topology().identity().clone(),
topology: selected.topology(),
rank: selected.rank(),
residency: selected.residency(),
completion: selected.completion(),
}
}
#[cfg(test)]
#[allow(clippy::too_many_arguments)]
fn from_parts(
architecture: RealtimeIdentity,
source: RealtimeIdentity,
execution: RealtimeIdentity,
schedule_identity: RealtimeIdentity,
schedule: RealtimeSpeechConfig,
state_layout: RealtimeIdentity,
topology_policy: RealtimeIdentity,
topology: ParallelTopology,
rank: usize,
residency: LayerWeightResidency,
completion: CommunicationCompletionPolicy,
) -> Self {
Self {
selected: None,
architecture,
source,
execution,
schedule_identity,
schedule,
state_layout,
topology_policy,
topology,
rank,
residency,
completion,
}
}
pub const fn schedule(&self) -> &RealtimeSpeechConfig {
&self.schedule
}
pub const fn topology(&self) -> ParallelTopology {
self.topology
}
pub const fn completion(&self) -> CommunicationCompletionPolicy {
self.completion
}
pub fn model_owner(&self) -> RealtimeModelOwnerIdentity {
RealtimeModelOwnerIdentity(self.clone())
}
fn distributed_consensus_identity(&self) -> [u32; 8] {
let mut digest = Sha256::new();
for identity in [
&self.architecture,
&self.source,
&self.execution,
&self.schedule_identity,
&self.topology_policy,
] {
let bytes = identity.as_str().as_bytes();
digest.update(
u64::try_from(bytes.len())
.expect("identity byte length fits u64")
.to_le_bytes(),
);
digest.update(bytes);
}
for value in [
format!("{:?}", self.schedule),
format!("{:?}", self.topology),
format!("{:?}", self.residency),
format!("{:?}", self.completion),
] {
digest.update(
u64::try_from(value.len())
.expect("identity component length fits u64")
.to_le_bytes(),
);
digest.update(value.as_bytes());
}
if let Some(selected) = &self.selected {
for value in [
format!("{:?}", selected.state()),
format!("{:?}", selected.observations()),
] {
digest.update(
u64::try_from(value.len())
.expect("selected identity component length fits u64")
.to_le_bytes(),
);
digest.update(value.as_bytes());
}
}
let digest = digest.finalize();
std::array::from_fn(|index| {
let start = index * 4;
u32::from_le_bytes(
digest[start..start + 4]
.try_into()
.expect("SHA-256 has eight complete u32 words"),
)
})
}
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct RealtimeModelOwnerIdentity(RealtimeModelSessionIdentity);
impl RealtimeModelOwnerIdentity {
pub const fn session_identity(&self) -> &RealtimeModelSessionIdentity {
&self.0
}
}
#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct RealtimeSessionIncarnation(u64);
impl RealtimeSessionIncarnation {
pub const fn value(self) -> u64 {
self.0
}
}
#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct RealtimeHistoryGeneration(u64);
impl RealtimeHistoryGeneration {
pub const fn value(self) -> u64 {
self.0
}
}
pub struct RealtimeSessionState<M, S, R, C> {
model: RealtimeModelSessionIdentity,
owner: RealtimeModelOwnerIdentity,
incarnation: RealtimeSessionIncarnation,
history_generation: RealtimeHistoryGeneration,
committed_batch: Option<NonZeroUsize>,
generation: RealtimeGenerationState<M, S, R, C>,
}
impl<M, S, R, C> RealtimeSessionState<M, S, R, C> {
pub const fn model_identity(&self) -> &RealtimeModelSessionIdentity {
&self.model
}
pub const fn model_owner(&self) -> &RealtimeModelOwnerIdentity {
&self.owner
}
pub const fn incarnation(&self) -> RealtimeSessionIncarnation {
self.incarnation
}
pub const fn history_generation(&self) -> RealtimeHistoryGeneration {
self.history_generation
}
pub const fn committed_batch(&self) -> Option<NonZeroUsize> {
self.committed_batch
}
pub fn payload_contract(
&self,
ingress: &RealtimeIngressContract,
) -> Result<RealtimePayloadContract, RealtimeSessionExecutionError> {
if ingress.schedule() != self.model.schedule() {
return Err(RealtimeSessionExecutionError::IngressScheduleMismatch);
}
let batch = self
.committed_batch
.ok_or(RealtimeSessionExecutionError::BatchNotAdmitted)?;
Ok(RealtimePayloadContract::new(
ingress.schedule().clone(),
batch.get(),
ingress.text_domain(),
ingress.audio_domain(),
RealtimePayloadGeneration::new(self.history_generation.value())
.expect("scheduler history generations are nonzero"),
RealtimePayloadOwnerIdentity::new(self.incarnation.value())
.expect("scheduler session incarnations are nonzero"),
)
.expect("scheduler-admitted payload contract has a positive batch"))
}
pub const fn generation(&self) -> &RealtimeGenerationState<M, S, R, C> {
&self.generation
}
pub fn generation_mut(&mut self) -> &mut RealtimeGenerationState<M, S, R, C> {
&mut self.generation
}
}
pub struct RealtimeSessionBranch<MB, S, R, C> {
model: RealtimeModelSessionIdentity,
owner: RealtimeModelOwnerIdentity,
incarnation: RealtimeSessionIncarnation,
history_generation: RealtimeHistoryGeneration,
committed_batch: Option<NonZeroUsize>,
generation: RealtimeGenerationBranch<MB, S, R, C>,
}
impl<MB, S, R, C> RealtimeSessionBranch<MB, S, R, C> {
pub const fn model_identity(&self) -> &RealtimeModelSessionIdentity {
&self.model
}
pub const fn model_owner(&self) -> &RealtimeModelOwnerIdentity {
&self.owner
}
pub const fn incarnation(&self) -> RealtimeSessionIncarnation {
self.incarnation
}
pub const fn history_generation(&self) -> RealtimeHistoryGeneration {
self.history_generation
}
pub const fn committed_batch(&self) -> Option<NonZeroUsize> {
self.committed_batch
}
pub fn payload_contract(
&self,
ingress: &RealtimeIngressContract,
) -> Result<RealtimePayloadContract, RealtimeSessionExecutionError> {
if ingress.schedule() != self.model.schedule() {
return Err(RealtimeSessionExecutionError::IngressScheduleMismatch);
}
let batch = self
.committed_batch
.ok_or(RealtimeSessionExecutionError::BatchNotAdmitted)?;
Ok(RealtimePayloadContract::new(
ingress.schedule().clone(),
batch.get(),
ingress.text_domain(),
ingress.audio_domain(),
RealtimePayloadGeneration::new(self.history_generation.value())
.expect("scheduler history generations are nonzero"),
RealtimePayloadOwnerIdentity::new(self.incarnation.value())
.expect("scheduler session incarnations are nonzero"),
)
.expect("scheduler-admitted payload contract has a positive batch"))
}
pub fn generation_mut(&mut self) -> &mut RealtimeGenerationBranch<MB, S, R, C> {
&mut self.generation
}
fn admit_batch(&mut self, batch: usize) -> Result<(), RealtimeSessionExecutionError> {
let batch = NonZeroUsize::new(batch).ok_or(RealtimeSessionExecutionError::EmptyBatch)?;
match self.committed_batch {
Some(committed) if committed != batch => Err(RealtimeSessionExecutionError::Batch {
committed: committed.get(),
submitted: batch.get(),
}),
Some(_) => Ok(()),
None => {
self.committed_batch = Some(batch);
Ok(())
}
}
}
}
impl<M, S, R, C> SemanticStateTransaction for RealtimeSessionState<M, S, R, C>
where
M: SemanticStateTransaction,
M::Error: 'static,
S: Clone,
R: Clone,
C: Completion,
{
type Branch = RealtimeSessionBranch<M::Branch, S, R, C>;
type Error = RealtimeSessionTransactionError<M::Error, C::Error>;
fn branch(&self) -> Result<Self::Branch, Self::Error> {
Ok(RealtimeSessionBranch {
model: self.model.clone(),
owner: self.owner.clone(),
incarnation: self.incarnation,
history_generation: self.history_generation,
committed_batch: self.committed_batch,
generation: self
.generation
.branch()
.map_err(RealtimeSessionTransactionError::Generation)?,
})
}
fn commit_branch(&mut self, branch: Self::Branch) -> Result<(), Self::Error> {
if self.model != branch.model
|| self.owner != branch.owner
|| self.incarnation != branch.incarnation
|| self.history_generation != branch.history_generation
{
Self::discard_branch(branch)?;
return Err(RealtimeSessionTransactionError::IdentityMismatch);
}
if let (Some(committed), Some(submitted)) = (self.committed_batch, branch.committed_batch) {
if committed != submitted {
Self::discard_branch(branch)?;
return Err(RealtimeSessionTransactionError::BatchMismatch);
}
}
self.generation
.commit_branch(branch.generation)
.map_err(RealtimeSessionTransactionError::Generation)?;
self.committed_batch = branch.committed_batch;
Ok(())
}
fn discard_branch(branch: Self::Branch) -> Result<(), Self::Error> {
RealtimeGenerationState::<M, S, R, C>::discard_branch(branch.generation)
.map_err(RealtimeSessionTransactionError::Generation)
}
fn permits_parallel_branches(&self) -> bool {
false
}
}
pub struct ReleasedRealtimeSession<M, S, R, C> {
state: RealtimeSessionState<M, S, R, C>,
}
impl<M, S, R, C> ReleasedRealtimeSession<M, S, R, C> {
pub const fn model_identity(&self) -> &RealtimeModelSessionIdentity {
self.state.model_identity()
}
pub const fn incarnation(&self) -> RealtimeSessionIncarnation {
self.state.incarnation()
}
pub const fn committed_batch(&self) -> Option<NonZeroUsize> {
self.state.committed_batch()
}
}
pub struct RealtimeSessionScheduler<M, S, R, C, O>
where
M: SemanticStateTransaction,
M::Error: 'static,
S: Clone,
R: Clone,
C: Completion,
O: TransitionOutput,
{
model: RealtimeModelSessionIdentity,
scheduler: Scheduler<RealtimeInputFrame, RealtimeSessionState<M, S, R, C>, O>,
}
impl<M, S, R, C, O> RealtimeSessionScheduler<M, S, R, C, O>
where
M: SemanticStateTransaction,
M::Error: 'static,
S: Clone,
R: Clone,
C: Completion,
O: TransitionOutput,
{
pub const fn model_identity(&self) -> &RealtimeModelSessionIdentity {
&self.model
}
pub fn new(
model: RealtimeModelSessionIdentity,
limits: SchedulerLimits,
) -> Result<Self, SchedulerError> {
Ok(Self {
model,
scheduler: Scheduler::new(limits)?,
})
}
pub fn register(
&mut self,
request: RequestId,
generation: RealtimeGenerationState<M, S, R, C>,
) -> Result<RealtimeSessionIncarnation, RealtimeSessionError> {
if generation.schedule_state().schedule() != self.model.schedule() {
return Err(RealtimeSessionError::ScheduleMismatch);
}
self.scheduler.validate_registration(request)?;
let incarnation = allocate_incarnation()?;
let state = RealtimeSessionState {
model: self.model.clone(),
owner: self.model.model_owner(),
incarnation,
history_generation: RealtimeHistoryGeneration(incarnation.0),
committed_batch: None,
generation,
};
self.scheduler.register(request, state)?;
Ok(incarnation)
}
pub fn resume(
&mut self,
request: RequestId,
released: ReleasedRealtimeSession<M, S, R, C>,
) -> Result<(), RealtimeSessionResumeError<M, S, R, C>> {
if released.state.model != self.model {
return Err(RealtimeSessionResumeError {
reason: RealtimeSessionError::ModelIdentityMismatch,
released: Box::new(released),
});
}
if let Err(error) = self.scheduler.validate_registration(request) {
return Err(RealtimeSessionResumeError {
reason: RealtimeSessionError::Scheduler(error),
released: Box::new(released),
});
}
self.scheduler
.register(request, released.state)
.expect("prevalidated realtime resumption cannot fail registration");
Ok(())
}
pub fn enqueue(
&mut self,
request: RequestId,
frame: RealtimeInputFrame,
) -> Result<WorkId, SchedulerError> {
self.scheduler.enqueue(request, frame)
}
pub fn enqueue_with_deadline(
&mut self,
request: RequestId,
frame: RealtimeInputFrame,
deadline: Option<Instant>,
) -> Result<WorkId, SchedulerError> {
self.scheduler
.enqueue_with_deadline(request, frame, deadline)
}
pub fn enqueue_batch(
&mut self,
request: RequestId,
frames: Vec<RealtimeInputFrame>,
) -> Result<Vec<WorkId>, SchedulerError> {
self.scheduler.enqueue_batch(request, frames)
}
pub fn run_local_turn<E>(
&mut self,
now: Instant,
mut execute: impl FnMut(
WorkId,
&RealtimeInputFrame,
&mut RealtimeSessionBranch<M::Branch, S, R, C>,
) -> Result<O, E>,
) -> Result<SchedulerProgress<RealtimeInputFrame, O>, SchedulerError>
where
E: std::error::Error,
{
self.ensure_local_topology()?;
self.scheduler.run_local_turn(now, |id, frame, branch| {
branch
.admit_batch(frame.batch())
.map_err(RealtimeSessionSubmissionError::<E>::Session)?;
execute(id, frame, branch).map_err(RealtimeSessionSubmissionError::Execution)
})
}
pub fn run_local_bounded<E>(
&mut self,
now: Instant,
maximum_frames: usize,
mut execute: impl FnMut(
WorkId,
&RealtimeInputFrame,
&mut RealtimeSessionBranch<M::Branch, S, R, C>,
) -> Result<O, E>,
) -> Result<SchedulerProgress<RealtimeInputFrame, O>, SchedulerError>
where
E: std::error::Error,
{
self.ensure_local_topology()?;
let mut progress = self.scheduler.poll_completions(now);
self.scheduler.prepare_bounded(maximum_frames, now)?;
progress.newly_submitted = self.scheduler.submit_prepared(now, |id, frame, branch| {
branch
.admit_batch(frame.batch())
.map_err(RealtimeSessionSubmissionError::<E>::Session)?;
execute(id, frame, branch).map_err(RealtimeSessionSubmissionError::Execution)
})?;
let completed = self.scheduler.poll_completions(now);
progress.committed.extend(completed.committed);
progress.failed.extend(completed.failed);
Ok(progress)
}
pub fn run_distributed_turn<T, E>(
&mut self,
protocol: u64,
transport: &T,
now: Instant,
mut execute: impl FnMut(
WorkId,
&RealtimeInputFrame,
&mut RealtimeSessionBranch<M::Branch, S, R, C>,
) -> Result<O, E>,
) -> Result<SchedulerProgress<RealtimeInputFrame, O>, SchedulerError>
where
T: BoundedConsensusTransport,
<T::Completion as Completion>::Error: std::fmt::Display,
E: std::error::Error,
O: eredu_core::scheduler::DistributedTransitionOutput,
{
if self.model.topology().is_replicated() {
return Err(SchedulerError::Consensus(
"distributed realtime turns require a non-replicated selected topology".into(),
));
}
let participants = self.model.topology().world_size();
if transport.participant_count() != participants {
return Err(SchedulerError::Consensus(format!(
"distributed realtime transport has {} participants; selected topology requires {participants}",
transport.participant_count(),
)));
}
let wait: BoundedCompletionWait = self.model.completion().bounded_wait();
validate_ranked_identity_bounded(
transport,
protocol,
&self.model.distributed_consensus_identity(),
self.model.rank,
wait,
)
.map_err(|error| SchedulerError::Consensus(error.to_string()))?;
self.scheduler
.run_distributed_turn(protocol, transport, wait, now, |id, frame, branch| {
branch
.admit_batch(frame.batch())
.map_err(RealtimeSessionSubmissionError::<E>::Session)?;
execute(id, frame, branch).map_err(RealtimeSessionSubmissionError::Execution)
})
}
fn ensure_local_topology(&self) -> Result<(), SchedulerError> {
if self.model.topology().is_replicated() {
Ok(())
} else {
Err(SchedulerError::Consensus(
"rank-local realtime turns are forbidden for a non-replicated selected topology"
.into(),
))
}
}
pub fn replace_sampling<E>(
&mut self,
request: RequestId,
sampling: RealtimeSampling,
realize: impl FnOnce(RealtimeSampling) -> Result<(Vec<S>, Option<R>), E>,
) -> Result<(), RealtimeSamplingUpdateError<E, M::Error, C::Error>> {
if self.scheduler.queued_for_request(request) != 0 {
return Err(RealtimeSamplingReplacementError::QueuedWork);
}
let state = self
.scheduler
.request_state_mut(request)
.map_err(RealtimeSamplingReplacementError::Scheduler)?;
let (samplers, random) =
realize(sampling).map_err(RealtimeSamplingReplacementError::Realization)?;
state
.generation_mut()
.replace_sampling(sampling, samplers, random)
.map_err(RealtimeSamplingReplacementError::Generation)
}
pub fn cancel(&mut self, request: RequestId) -> Result<(), SchedulerError> {
self.scheduler.cancel(request)
}
pub fn finish(&mut self, request: RequestId) -> Result<(), SchedulerError> {
self.scheduler.finish(request)
}
pub fn release(
&mut self,
request: RequestId,
) -> Result<ReleasedRealtimeSession<M, S, R, C>, SchedulerError> {
self.scheduler
.release(request)
.map(|state| ReleasedRealtimeSession { state })
}
pub fn request_status(&self, request: RequestId) -> Option<RequestStatus> {
self.scheduler.request_status(request)
}
pub fn queued_for_request(&self, request: RequestId) -> usize {
self.scheduler.queued_for_request(request)
}
pub fn forget_terminal(&mut self, request: RequestId) -> Result<RequestStatus, SchedulerError> {
self.scheduler.forget_terminal(request)
}
pub fn request_state(&self, request: RequestId) -> Option<&RealtimeSessionState<M, S, R, C>> {
self.scheduler.request_state(request)
}
pub fn report(&self) -> SchedulerReport {
self.scheduler.report()
}
pub fn capabilities(&self) -> SchedulerCapabilities {
self.scheduler.capabilities()
}
}
fn allocate_incarnation() -> Result<RealtimeSessionIncarnation, RealtimeSessionError> {
NEXT_REALTIME_INCARNATION
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| {
value.checked_add(1)
})
.map(RealtimeSessionIncarnation)
.map_err(|_| RealtimeSessionError::IncarnationExhausted)
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum RealtimeSessionError {
#[error(transparent)]
Scheduler(#[from] SchedulerError),
#[error("realtime generation schedule differs from selected model identity")]
ScheduleMismatch,
#[error("released realtime session model identity does not match")]
ModelIdentityMismatch,
#[error("realtime session incarnation identity space is exhausted")]
IncarnationExhausted,
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, thiserror::Error)]
#[non_exhaustive]
pub enum RealtimeSessionExecutionError {
#[error("realtime session frame batch has not been admitted")]
BatchNotAdmitted,
#[error("realtime session frame batch must be positive")]
EmptyBatch,
#[error("realtime session batch is {submitted}, committed batch is {committed}")]
Batch {
committed: usize,
submitted: usize,
},
#[error("realtime ingress schedule differs from selected session schedule")]
IngressScheduleMismatch,
}
pub struct RealtimeSessionResumeError<M, S, R, C> {
reason: RealtimeSessionError,
released: Box<ReleasedRealtimeSession<M, S, R, C>>,
}
impl<M, S, R, C> RealtimeSessionResumeError<M, S, R, C> {
pub const fn reason(&self) -> &RealtimeSessionError {
&self.reason
}
pub fn into_released(self) -> ReleasedRealtimeSession<M, S, R, C> {
*self.released
}
}
impl<M, S, R, C> std::fmt::Debug for RealtimeSessionResumeError<M, S, R, C> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("RealtimeSessionResumeError")
.field("reason", &self.reason)
.finish_non_exhaustive()
}
}
impl<M, S, R, C> std::fmt::Display for RealtimeSessionResumeError<M, S, R, C> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.reason.fmt(formatter)
}
}
impl<M, S, R, C> std::error::Error for RealtimeSessionResumeError<M, S, R, C> {}
#[derive(Debug, thiserror::Error)]
enum RealtimeSessionSubmissionError<E: std::error::Error> {
#[error(transparent)]
Session(RealtimeSessionExecutionError),
#[error(transparent)]
Execution(E),
}
#[derive(Debug, thiserror::Error)]
pub enum RealtimeSessionTransactionError<M, C>
where
M: std::error::Error,
C: std::error::Error,
{
#[error(transparent)]
Generation(RealtimeGenerationTransactionError<M, C>),
#[error("realtime session branch identity does not match canonical state")]
IdentityMismatch,
#[error("realtime session branch batch does not match canonical state")]
BatchMismatch,
}
#[derive(Debug, thiserror::Error)]
pub enum RealtimeSamplingReplacementError<E, G> {
#[error("cannot replace realtime sampling while frames are queued")]
QueuedWork,
#[error(transparent)]
Scheduler(SchedulerError),
#[error("realtime sampling realization failed")]
Realization(E),
#[error("realtime generation rejected replacement sampling")]
Generation(G),
}
pub type RealtimeSamplingUpdateError<E, M, C> =
RealtimeSamplingReplacementError<E, RealtimeGenerationTransactionError<M, C>>;
#[cfg(test)]
mod tests {
use std::{cell::Cell, convert::Infallible, rc::Rc, time::Duration};
use eredu_core::{
consensus::{BoundedConsensusTransport, ConsensusTransport},
scheduler::{RequestStatus, TransitionOutput},
BoundedCompletion, BoundedCompletionOutcome, BoundedCompletionWait,
CompletionCancellationMode, RealtimeFrameConvention, Submission,
};
use super::*;
use crate::TokenDomain;
#[derive(Debug, Clone, Eq, PartialEq)]
struct ModelState(usize);
#[derive(Debug, Clone, Copy, thiserror::Error)]
#[error("model transaction failed")]
struct ModelError;
impl SemanticStateTransaction for ModelState {
type Branch = Self;
type Error = ModelError;
fn branch(&self) -> Result<Self::Branch, Self::Error> {
Ok(self.clone())
}
fn commit_branch(&mut self, branch: Self::Branch) -> Result<(), Self::Error> {
*self = branch;
Ok(())
}
}
#[derive(Debug, Clone)]
struct MockCompletion(Rc<Cell<bool>>);
#[derive(Debug, Clone, Copy, thiserror::Error)]
#[error("completion failed")]
struct CompletionError;
impl Completion for MockCompletion {
type Error = CompletionError;
fn is_complete(&self) -> Result<bool, Self::Error> {
Ok(self.0.get())
}
fn wait(&self) -> Result<(), Self::Error> {
self.0.get().then_some(()).ok_or(CompletionError)
}
}
#[derive(Debug)]
struct MockOutput {
completion: MockCompletion,
}
impl TransitionOutput for MockOutput {
type Error = CompletionError;
fn is_complete(&self) -> Result<bool, Self::Error> {
self.completion.is_complete()
}
fn retained_resources(&self) -> usize {
1
}
}
impl eredu_core::scheduler::DistributedTransitionOutput for MockOutput {
fn encode_distributed_output(&self, output: &mut Vec<u32>) -> Result<(), String> {
output.push(0);
Ok(())
}
}
type Sessions = RealtimeSessionScheduler<ModelState, (), (), MockCompletion, MockOutput>;
fn schedule() -> RealtimeSpeechConfig {
RealtimeSpeechConfig::new(
2,
1,
1,
1,
9,
8,
RealtimeFrameConvention::FeedbackAlignedHistory,
vec![0, 0, 1],
)
.unwrap()
}
fn identity(suffix: &str) -> RealtimeModelSessionIdentity {
identity_with_topology(suffix, ParallelTopology::new(1, 1, 1, 1).unwrap())
}
fn identity_with_topology(
suffix: &str,
topology: ParallelTopology,
) -> RealtimeModelSessionIdentity {
let id = |prefix: &str| RealtimeIdentity::new(format!("{prefix}-{suffix}")).unwrap();
RealtimeModelSessionIdentity::from_parts(
id("architecture"),
id("source"),
id("execution"),
id("schedule"),
schedule(),
id("state"),
id("topology"),
topology,
0,
LayerWeightResidency::FullyResident,
CommunicationCompletionPolicy::new(
Duration::from_secs(1),
CompletionCancellationMode::QuarantineUntilComplete,
)
.unwrap(),
)
}
#[derive(Debug, Clone, Copy)]
struct ReadyConsensusCompletion;
impl Completion for ReadyConsensusCompletion {
type Error = Infallible;
fn is_complete(&self) -> Result<bool, Self::Error> {
Ok(true)
}
fn wait(&self) -> Result<(), Self::Error> {
Ok(())
}
}
impl BoundedCompletion for ReadyConsensusCompletion {
fn wait_bounded(
self,
_: BoundedCompletionWait,
) -> Result<BoundedCompletionOutcome, Self::Error> {
Ok(BoundedCompletionOutcome::Completed)
}
}
struct DisagreeingTransport {
calls: Cell<usize>,
disagree_on: usize,
}
impl DisagreeingTransport {
fn model_identity() -> Self {
Self {
calls: Cell::new(0),
disagree_on: 1,
}
}
fn schedule() -> Self {
Self {
calls: Cell::new(0),
disagree_on: 4,
}
}
}
impl ConsensusTransport for DisagreeingTransport {
type Error = Infallible;
fn participant_count(&self) -> usize {
2
}
fn all_gather_words(&self, local: &[u32]) -> Result<Vec<u32>, Self::Error> {
let call = self.calls.get() + 1;
self.calls.set(call);
let mut gathered = local.repeat(2);
if call == 1 {
*gathered.last_mut().expect("identity frame has a rank word") = 1;
}
if call == self.disagree_on {
gathered[local.len()] ^= 1;
}
Ok(gathered)
}
}
impl BoundedConsensusTransport for DisagreeingTransport {
type Completion = ReadyConsensusCompletion;
type GatherOutput = Vec<u32>;
fn submit_all_gather_words(
&self,
local: &[u32],
) -> Result<Submission<Self::GatherOutput, Self::Completion>, Self::Error> {
Ok(Submission {
output: self.all_gather_words(local)?,
completion: ReadyConsensusCompletion,
})
}
fn resolve_all_gather_words(
&self,
output: Self::GatherOutput,
) -> Result<Vec<u32>, Self::Error> {
Ok(output)
}
}
fn generation() -> RealtimeGenerationState<ModelState, (), (), MockCompletion> {
RealtimeGenerationState::new(
ModelState(0),
schedule(),
RealtimeSampling::greedy(),
vec![(), ()],
None,
)
.unwrap()
}
fn ingress() -> RealtimeIngressContract {
RealtimeIngressContract::new(schedule(), TokenDomain::new(32), TokenDomain::new(16))
.unwrap()
}
fn frame(batch: usize) -> RealtimeInputFrame {
RealtimeInputFrame::new(batch, vec![1; batch])
}
fn limits(submissions: usize) -> SchedulerLimits {
SchedulerLimits::with_execution_bounds(8, 32, submissions, 8, 1, usize::MAX).unwrap()
}
fn execute_immediately(
_id: WorkId,
_frame: &RealtimeInputFrame,
branch: &mut RealtimeSessionBranch<ModelState, (), (), MockCompletion>,
) -> Result<MockOutput, Infallible> {
branch.generation_mut().model_state_mut().0 += 1;
let completion = MockCompletion(Rc::new(Cell::new(true)));
branch
.generation_mut()
.attach_submission_completion(completion.clone())
.unwrap();
Ok(MockOutput { completion })
}
#[test]
fn one_scheduler_round_robins_single_lane_sessions() {
let mut sessions = Sessions::new(identity("a"), limits(2)).unwrap();
let first = RequestId::new(1);
let second = RequestId::new(2);
let first_incarnation = sessions.register(first, generation()).unwrap();
let second_incarnation = sessions.register(second, generation()).unwrap();
assert_ne!(first_incarnation, second_incarnation);
sessions
.enqueue_batch(first, vec![frame(1), frame(1)])
.unwrap();
sessions
.enqueue_batch(second, vec![frame(1), frame(1)])
.unwrap();
let progress = sessions
.run_local_turn(Instant::now(), execute_immediately)
.unwrap();
assert_eq!(progress.committed.len(), 2);
assert_eq!(progress.committed[0].0.request(), first);
assert_eq!(progress.committed[1].0.request(), second);
assert_eq!(
sessions
.request_state(first)
.unwrap()
.generation()
.model_state()
.0,
1
);
assert_eq!(
sessions
.request_state(second)
.unwrap()
.generation()
.model_state()
.0,
1
);
}
#[test]
fn non_replicated_identity_rejects_local_turn_before_submission() {
let topology = ParallelTopology::new(2, 1, 1, 1).unwrap();
let mut sessions =
Sessions::new(identity_with_topology("tp", topology), limits(1)).unwrap();
let request = RequestId::new(90);
sessions.register(request, generation()).unwrap();
sessions.enqueue(request, frame(1)).unwrap();
let calls = Rc::new(Cell::new(0));
let observed = calls.clone();
let error = sessions
.run_local_turn(Instant::now(), move |id, frame, branch| {
observed.set(observed.get() + 1);
execute_immediately(id, frame, branch)
})
.unwrap_err();
assert!(matches!(error, SchedulerError::Consensus(_)));
assert_eq!(calls.get(), 0);
assert_eq!(sessions.queued_for_request(request), 1);
let error = sessions
.run_local_bounded(Instant::now(), 1, |id, frame, branch| {
calls.set(calls.get() + 1);
execute_immediately(id, frame, branch)
})
.unwrap_err();
assert!(matches!(error, SchedulerError::Consensus(_)));
assert_eq!(calls.get(), 0);
assert_eq!(sessions.queued_for_request(request), 1);
}
#[test]
fn distributed_model_identity_disagreement_submits_and_publishes_nothing() {
let topology = ParallelTopology::new(2, 1, 1, 1).unwrap();
let mut sessions = Sessions::new(
identity_with_topology("tp-disagreement", topology),
limits(1),
)
.unwrap();
let request = RequestId::new(91);
sessions.register(request, generation()).unwrap();
sessions.enqueue(request, frame(1)).unwrap();
let calls = Rc::new(Cell::new(0));
let observed = calls.clone();
let error = sessions
.run_distributed_turn(
17,
&DisagreeingTransport::model_identity(),
Instant::now(),
move |id, frame, branch| {
observed.set(observed.get() + 1);
execute_immediately(id, frame, branch)
},
)
.unwrap_err();
assert!(matches!(error, SchedulerError::Consensus(_)));
assert_eq!(calls.get(), 0);
assert_eq!(
sessions
.request_state(request)
.unwrap()
.generation()
.model_state()
.0,
0
);
assert_eq!(sessions.report().completed_work, 0);
}
#[test]
fn distributed_schedule_disagreement_submits_and_publishes_nothing() {
let topology = ParallelTopology::new(2, 1, 1, 1).unwrap();
let mut sessions =
Sessions::new(identity_with_topology("tp-schedule", topology), limits(1)).unwrap();
let request = RequestId::new(92);
sessions.register(request, generation()).unwrap();
sessions.enqueue(request, frame(1)).unwrap();
let calls = Rc::new(Cell::new(0));
let observed = calls.clone();
let error = sessions
.run_distributed_turn(
18,
&DisagreeingTransport::schedule(),
Instant::now(),
move |id, frame, branch| {
observed.set(observed.get() + 1);
execute_immediately(id, frame, branch)
},
)
.unwrap_err();
assert!(matches!(error, SchedulerError::Consensus(_)));
assert_eq!(calls.get(), 0);
assert_eq!(sessions.report().completed_work, 0);
}
#[test]
fn scheduler_branch_derives_payload_identity_before_native_execution() {
let mut sessions = Sessions::new(identity("payload"), limits(1)).unwrap();
let request = RequestId::new(20);
let incarnation = sessions.register(request, generation()).unwrap();
sessions.enqueue(request, frame(2)).unwrap();
sessions
.run_local_turn(Instant::now(), |id, submitted, branch| {
let contract = branch.payload_contract(&ingress()).unwrap();
assert_eq!(contract.batch().get(), 2);
assert_eq!(contract.owner().value(), incarnation.value());
assert_eq!(
contract.generation().value(),
branch.history_generation().value()
);
execute_immediately(id, submitted, branch)
})
.unwrap();
}
#[test]
fn committed_batch_rejects_a_different_batch_before_execution() {
let mut sessions = Sessions::new(identity("a"), limits(1)).unwrap();
let request = RequestId::new(3);
sessions.register(request, generation()).unwrap();
sessions.enqueue(request, frame(1)).unwrap();
sessions
.run_local_turn(Instant::now(), execute_immediately)
.unwrap();
assert_eq!(
sessions
.request_state(request)
.unwrap()
.committed_batch()
.unwrap()
.get(),
1
);
sessions.enqueue(request, frame(2)).unwrap();
let calls = Rc::new(Cell::new(0));
let observed = calls.clone();
assert!(sessions
.run_local_turn(Instant::now(), move |id, frame, branch| {
observed.set(observed.get() + 1);
execute_immediately(id, frame, branch)
})
.is_err());
assert_eq!(calls.get(), 0);
assert_eq!(
sessions.request_status(request),
Some(RequestStatus::Failed)
);
}
#[test]
fn release_resume_preserves_incarnation_and_wrong_model_keeps_state_recoverable() {
let request = RequestId::new(4);
let mut first = Sessions::new(identity("a"), limits(1)).unwrap();
let incarnation = first.register(request, generation()).unwrap();
let released = first.release(request).unwrap();
assert_eq!(released.incarnation(), incarnation);
let mut wrong = Sessions::new(identity("b"), limits(1)).unwrap();
let error = wrong.resume(request, released).unwrap_err();
assert!(matches!(
error.reason(),
RealtimeSessionError::ModelIdentityMismatch
));
let released = error.into_released();
first.resume(request, released).unwrap();
assert_eq!(
first.request_state(request).unwrap().incarnation(),
incarnation
);
}
#[test]
fn queued_work_blocks_sampling_replacement_before_realization() {
let request = RequestId::new(5);
let mut sessions = Sessions::new(identity("a"), limits(1)).unwrap();
sessions.register(request, generation()).unwrap();
sessions.enqueue(request, frame(1)).unwrap();
let realized = Rc::new(Cell::new(false));
let observed = realized.clone();
let result = sessions.replace_sampling(
request,
RealtimeSampling::greedy(),
move |_| -> Result<_, Infallible> {
observed.set(true);
Ok((vec![(), ()], None))
},
);
assert!(matches!(
result,
Err(RealtimeSamplingReplacementError::QueuedWork)
));
assert!(!realized.get());
}
#[test]
fn cancellation_and_deadline_use_core_terminal_semantics() {
let mut sessions = Sessions::new(identity("a"), limits(1)).unwrap();
let cancelled = RequestId::new(6);
sessions.register(cancelled, generation()).unwrap();
sessions.enqueue(cancelled, frame(1)).unwrap();
sessions.cancel(cancelled).unwrap();
assert_eq!(
sessions.request_status(cancelled),
Some(RequestStatus::Cancelled)
);
let expired = RequestId::new(7);
sessions.register(expired, generation()).unwrap();
sessions
.enqueue_with_deadline(
expired,
frame(1),
Some(Instant::now() - Duration::from_secs(1)),
)
.unwrap();
let calls = Rc::new(Cell::new(0));
let observed = calls.clone();
sessions
.run_local_turn(Instant::now(), move |id, frame, branch| {
observed.set(observed.get() + 1);
execute_immediately(id, frame, branch)
})
.unwrap();
assert_eq!(calls.get(), 0);
assert_eq!(
sessions.request_status(expired),
Some(RequestStatus::DeadlineExceeded)
);
}
#[test]
fn cancelling_submitted_work_retains_it_until_exact_completion() {
let request = RequestId::new(8);
let mut sessions = Sessions::new(identity("a"), limits(1)).unwrap();
sessions.register(request, generation()).unwrap();
sessions.enqueue(request, frame(1)).unwrap();
let completion = Rc::new(Cell::new(false));
let submitted = completion.clone();
sessions
.run_local_turn(Instant::now(), move |_, _, branch| {
let completion = MockCompletion(submitted.clone());
branch
.generation_mut()
.attach_submission_completion(completion.clone())
.unwrap();
Ok::<_, Infallible>(MockOutput { completion })
})
.unwrap();
sessions.cancel(request).unwrap();
assert_eq!(sessions.report().abandoned_in_flight_work, 1);
completion.set(true);
sessions
.run_local_turn(Instant::now(), execute_immediately)
.unwrap();
assert_eq!(sessions.report().abandoned_in_flight_work, 0);
}
}