1use std::{cell::RefCell, sync::Arc};
8
9use eredu_core::{
10 scheduler::{DistributedTransitionOutput, TransitionOutput, WorkDescriptor},
11 Completion, RealtimeInputFrame,
12};
13
14use crate::{
15 complete_realtime_frame, prepare_realtime_frame, CompletedRealtimeFrame,
16 MaterializedRealtimeInput, RealtimeCompletionAttachmentError, RealtimeFrameInterpretationError,
17 RealtimeFrameTensorMechanisms, RealtimeGenerationBranch, RealtimeHostTokenMaterializer,
18 RealtimeIngressContract, RealtimeIngressError, RealtimePayloadBranch, RealtimePayloadContract,
19 RealtimePayloadContractError, RealtimePayloadHistory, Sampler, SamplingBackend,
20 SequentialDecisionDriver, SequentialDecisionError, SequentialDecisionPlan,
21 SequentialDecisionPlanError,
22};
23
24pub trait PreparedRealtimeFrameExecutor<B, S, M>
30where
31 B: SamplingBackend,
32 S: Sampler<B>,
33{
34 type Error;
36 type Retained;
38
39 fn execute(
42 &mut self,
43 model_state: &mut M,
44 temporal: &[B::Token],
45 driver: &mut SequentialDecisionDriver<B, S>,
46 context: &B::Context,
47 ) -> Result<Self::Retained, Self::Error>;
48}
49
50pub trait RealtimeFrameCompletionMechanism<T, M, R> {
55 type Completion;
57 type Error;
59
60 fn complete(
62 &mut self,
63 input: MaterializedRealtimeInput<T>,
64 output: &CompletedRealtimeFrame<T, T>,
65 model_state: &M,
66 payload_history: &RealtimePayloadHistory<T>,
67 execution: Option<R>,
68 ) -> Result<Self::Completion, RealtimeCompletionCreationError<Self::Completion, Self::Error>>;
69
70 fn retained_resources(&self, _completion: &Self::Completion) -> usize {
72 0
73 }
74}
75
76pub enum RealtimeCompletionCreationError<C, E> {
78 BeforeSubmission(E),
80 AfterSubmission {
82 error: E,
84 completion: C,
86 },
87}
88
89impl<C, E> RealtimeCompletionCreationError<C, E> {
90 pub const fn before_submission(error: E) -> Self {
92 Self::BeforeSubmission(error)
93 }
94
95 pub const fn after_submission(error: E, completion: C) -> Self {
97 Self::AfterSubmission { error, completion }
98 }
99}
100
101pub struct SubmittedRealtimeFrame<T, C> {
103 frame: CompletedRealtimeFrame<T, T>,
104 completion: C,
105 retained_resources: usize,
106}
107
108impl<T, C> SubmittedRealtimeFrame<T, C> {
109 pub const fn frame(&self) -> &CompletedRealtimeFrame<T, T> {
111 &self.frame
112 }
113
114 pub const fn completion(&self) -> &C {
116 &self.completion
117 }
118
119 pub fn into_parts(self) -> (CompletedRealtimeFrame<T, T>, C) {
121 (self.frame, self.completion)
122 }
123}
124
125impl<T, C> TransitionOutput for SubmittedRealtimeFrame<T, C>
126where
127 C: Completion,
128{
129 type Error = C::Error;
130
131 fn is_complete(&self) -> Result<bool, Self::Error> {
132 self.completion.is_complete()
133 }
134
135 fn retained_resources(&self) -> usize {
136 self.retained_resources
137 }
138}
139
140pub trait RealtimeFrameHostObserver<T> {
146 type Output;
148 type Error: std::error::Error + 'static;
150
151 fn observe(
153 &mut self,
154 frame: &CompletedRealtimeFrame<T, T>,
155 ) -> Result<Self::Output, Self::Error>;
156}
157
158enum HostObservationState<H: RealtimeFrameHostObserver<T>, T> {
159 Pending(H),
160 Ready(H::Output),
161 Failed(Arc<H::Error>),
162}
163
164pub struct PrepublicationRealtimeFrame<T, C, H>
172where
173 H: RealtimeFrameHostObserver<T>,
174{
175 submitted: SubmittedRealtimeFrame<T, C>,
176 observation: RefCell<HostObservationState<H, T>>,
177}
178
179impl<T, C, H> PrepublicationRealtimeFrame<T, C, H>
180where
181 H: RealtimeFrameHostObserver<T>,
182{
183 pub fn new(submitted: SubmittedRealtimeFrame<T, C>, observer: H) -> Self {
185 Self {
186 submitted,
187 observation: RefCell::new(HostObservationState::Pending(observer)),
188 }
189 }
190
191 pub fn into_host_output(self) -> Result<H::Output, RealtimeHostOutputUnavailable<H::Error>> {
196 match self.observation.into_inner() {
197 HostObservationState::Pending(_) => Err(RealtimeHostOutputUnavailable::Pending),
198 HostObservationState::Ready(output) => Ok(output),
199 HostObservationState::Failed(error) => {
200 Err(RealtimeHostOutputUnavailable::Observation(error))
201 }
202 }
203 }
204}
205
206impl<T, C, H> TransitionOutput for PrepublicationRealtimeFrame<T, C, H>
207where
208 C: Completion,
209 C::Error: 'static,
210 H: RealtimeFrameHostObserver<T>,
211{
212 type Error = RealtimePrepublicationError<C::Error, H::Error>;
213
214 fn is_complete(&self) -> Result<bool, Self::Error> {
215 {
216 let observation = self.observation.borrow();
217 match &*observation {
218 HostObservationState::Ready(_) => return Ok(true),
219 HostObservationState::Failed(error) => {
220 return Err(RealtimePrepublicationError::Observation(Arc::clone(error)))
221 }
222 HostObservationState::Pending(_) => {}
223 }
224 }
225
226 if !self
227 .submitted
228 .completion()
229 .is_complete()
230 .map_err(RealtimePrepublicationError::Completion)?
231 {
232 return Ok(false);
233 }
234 self.submitted
235 .completion()
236 .wait()
237 .map_err(RealtimePrepublicationError::Completion)?;
238
239 let mut observation = self.observation.borrow_mut();
240 let HostObservationState::Pending(observer) = &mut *observation else {
241 unreachable!("prepublication observation state was checked before completion wait")
242 };
243 match observer.observe(self.submitted.frame()) {
244 Ok(output) => {
245 *observation = HostObservationState::Ready(output);
246 Ok(true)
247 }
248 Err(error) => {
249 let error = Arc::new(error);
250 *observation = HostObservationState::Failed(Arc::clone(&error));
251 Err(RealtimePrepublicationError::Observation(error))
252 }
253 }
254 }
255
256 fn retained_resources(&self) -> usize {
257 self.submitted.retained_resources()
258 }
259}
260
261impl<T, C, H> DistributedTransitionOutput for PrepublicationRealtimeFrame<T, C, H>
262where
263 C: Completion,
264 C::Error: 'static,
265 H: RealtimeFrameHostObserver<T>,
266 H::Output: WorkDescriptor,
267{
268 fn encode_distributed_output(&self, output: &mut Vec<u32>) -> Result<(), String> {
269 match &*self.observation.borrow() {
270 HostObservationState::Ready(observed) => observed
271 .encode_descriptor(output)
272 .map_err(|error| error.to_string()),
273 HostObservationState::Pending(_) => {
274 Err("realtime output observation is still pending".into())
275 }
276 HostObservationState::Failed(error) => {
277 Err(format!("realtime output observation failed: {error}"))
278 }
279 }
280 }
281}
282
283#[derive(Debug)]
285pub enum RealtimePrepublicationError<C, O> {
286 Completion(C),
288 Observation(Arc<O>),
290}
291
292impl<C, O> std::fmt::Display for RealtimePrepublicationError<C, O>
293where
294 C: std::error::Error,
295 O: std::error::Error,
296{
297 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
298 match self {
299 Self::Completion(error) => {
300 write!(formatter, "realtime exact completion failed: {error}")
301 }
302 Self::Observation(error) => {
303 write!(formatter, "realtime host observation failed: {error}")
304 }
305 }
306 }
307}
308
309impl<C, O> std::error::Error for RealtimePrepublicationError<C, O>
310where
311 C: std::error::Error + 'static,
312 O: std::error::Error + 'static,
313{
314 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
315 match self {
316 Self::Completion(error) => Some(error),
317 Self::Observation(error) => Some(error.as_ref()),
318 }
319 }
320}
321
322#[derive(Debug)]
324pub enum RealtimeHostOutputUnavailable<O> {
325 Pending,
327 Observation(Arc<O>),
329}
330
331impl<O> std::fmt::Display for RealtimeHostOutputUnavailable<O>
332where
333 O: std::error::Error,
334{
335 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
336 match self {
337 Self::Pending => formatter.write_str("realtime host output is not observed yet"),
338 Self::Observation(error) => {
339 write!(formatter, "realtime host observation failed: {error}")
340 }
341 }
342 }
343}
344
345impl<O> std::error::Error for RealtimeHostOutputUnavailable<O>
346where
347 O: std::error::Error + 'static,
348{
349 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
350 match self {
351 Self::Pending => None,
352 Self::Observation(error) => Some(error.as_ref()),
353 }
354 }
355}
356
357#[derive(Debug, Clone, PartialEq)]
359pub struct RealtimeDecisionExecution {
360 allow_fully_forced_tail_skip: bool,
361}
362
363impl RealtimeDecisionExecution {
364 pub const fn new(allow_fully_forced_tail_skip: bool) -> Self {
366 Self {
367 allow_fully_forced_tail_skip,
368 }
369 }
370
371 pub const fn allows_fully_forced_tail_skip(&self) -> bool {
373 self.allow_fully_forced_tail_skip
374 }
375}
376
377#[allow(clippy::too_many_arguments, clippy::type_complexity)]
386pub fn execute_realtime_frame<B, S, M, T, C, H, F, E, K>(
387 contract: &RealtimeIngressContract,
388 payload_contract: &RealtimePayloadContract,
389 frame: &RealtimeInputFrame,
390 branch: &mut RealtimeGenerationBranch<RealtimePayloadBranch<M, T>, S, B::RandomState, C>,
391 decisions: &RealtimeDecisionExecution,
392 host_materializer: &mut H,
393 tensor_mechanisms: &mut F,
394 model_executor: &mut E,
395 completion_mechanism: &mut K,
396 context: &B::Context,
397) -> Result<
398 SubmittedRealtimeFrame<T, C>,
399 RealtimeFrameCoordinatorError<H::Error, F::Error, E::Error, B::Error, K::Error>,
400>
401where
402 B: SamplingBackend<Token = T, Logits = T>,
403 B::RandomState: Clone,
404 C: Completion + Clone,
405 S: Sampler<B> + Clone,
406 T: Clone,
407 H: RealtimeHostTokenMaterializer<Tensor = T>,
408 F: RealtimeFrameTensorMechanisms<Tensor = T>,
409 E: PreparedRealtimeFrameExecutor<B, S, M>,
410 K: RealtimeFrameCompletionMechanism<T, M, E::Retained, Completion = C>,
411{
412 if branch.has_submission_completion() {
413 return Err(RealtimeFrameCoordinatorError::CompletionAttachment(
414 RealtimeCompletionAttachmentError::AlreadyAttached,
415 ));
416 }
417 branch
418 .schedule_state()
419 .validate_schedule(contract.schedule())
420 .map_err(|error| {
421 RealtimeFrameCoordinatorError::Interpretation(
422 RealtimeFrameInterpretationError::Schedule(error),
423 )
424 })?;
425 if payload_contract.schedule() != contract.schedule() {
426 return Err(RealtimeFrameCoordinatorError::PayloadContract(
427 RealtimePayloadContractError::ScheduleMismatch,
428 ));
429 }
430 if payload_contract.batch().get() != frame.batch() {
431 return Err(RealtimeFrameCoordinatorError::PayloadContract(
432 RealtimePayloadContractError::BatchMismatch,
433 ));
434 }
435 if payload_contract.text_domain() != contract.text_domain() {
436 return Err(RealtimeFrameCoordinatorError::PayloadContract(
437 RealtimePayloadContractError::TextDomainMismatch,
438 ));
439 }
440 if payload_contract.audio_domain() != contract.audio_domain() {
441 return Err(RealtimeFrameCoordinatorError::PayloadContract(
442 RealtimePayloadContractError::AudioDomainMismatch,
443 ));
444 }
445 let mut payload_history = branch.model_state().payload_history().clone();
446 payload_history
447 .bind_or_validate_contract(payload_contract)
448 .map_err(|error| {
449 RealtimeFrameCoordinatorError::Interpretation(
450 RealtimeFrameInterpretationError::History(error),
451 )
452 })?;
453 let validated = contract
454 .validate(frame)
455 .map_err(RealtimeFrameCoordinatorError::Ingress)?;
456 let input = validated
457 .materialize(host_materializer)
458 .map_err(RealtimeFrameCoordinatorError::Materialization)?;
459
460 let schedule = contract.schedule();
461 let mut schedule_state = branch.schedule_state().clone();
462 let prepared = prepare_realtime_frame(
463 schedule,
464 &mut schedule_state,
465 &mut payload_history,
466 &input,
467 tensor_mechanisms,
468 )
469 .map_err(RealtimeFrameCoordinatorError::Interpretation)?;
470
471 let (completed, execution_retained) = if prepared.transition().model_call_required() {
472 let plan = SequentialDecisionPlan::new(
473 prepared.directives().iter().cloned(),
474 prepared.retains_diagnostics(),
475 decisions.allow_fully_forced_tail_skip,
476 )
477 .map_err(RealtimeFrameCoordinatorError::DecisionPlan)?;
478 let sampling = branch.sampling();
479 let temperatures = std::iter::once(sampling.text_temperature())
480 .chain(std::iter::repeat_n(
481 sampling.audio_temperature(),
482 schedule.depth_audio_codebooks(),
483 ))
484 .collect();
485 let mut driver = branch
486 .decision_driver::<B>(plan, temperatures)
487 .map_err(RealtimeFrameCoordinatorError::DecisionPlan)?;
488 let execution_retained = model_executor
489 .execute(
490 branch.model_state_mut().model_state_mut(),
491 prepared.temporal(),
492 &mut driver,
493 context,
494 )
495 .map_err(RealtimeFrameCoordinatorError::Model)?;
496 driver
497 .finish()
498 .map_err(RealtimeFrameCoordinatorError::Decision)?;
499 let resolved = driver
500 .decisions()
501 .iter()
502 .map(|decision| decision.token().clone())
503 .collect::<Vec<_>>();
504 let diagnostics = driver
505 .diagnostics()
506 .iter()
507 .map(|diagnostic| diagnostic.logits().clone())
508 .collect::<Vec<_>>();
509 let completed = complete_realtime_frame(
510 schedule,
511 &mut payload_history,
512 prepared,
513 resolved,
514 diagnostics,
515 tensor_mechanisms,
516 )
517 .map_err(RealtimeFrameCoordinatorError::Interpretation)?;
518 branch
519 .adopt_decision_driver(driver)
520 .map_err(RealtimeFrameCoordinatorError::Decision)?;
521 (completed, Some(execution_retained))
522 } else {
523 (
524 complete_realtime_frame(
525 schedule,
526 &mut payload_history,
527 prepared,
528 Vec::new(),
529 Vec::new(),
530 tensor_mechanisms,
531 )
532 .map_err(RealtimeFrameCoordinatorError::Interpretation)?,
533 None,
534 )
535 };
536
537 let completion = match completion_mechanism.complete(
538 input,
539 &completed,
540 branch.model_state_mut().model_state(),
541 &payload_history,
542 execution_retained,
543 ) {
544 Ok(completion) => completion,
545 Err(RealtimeCompletionCreationError::BeforeSubmission(error)) => {
546 return Err(RealtimeFrameCoordinatorError::Completion(error));
547 }
548 Err(RealtimeCompletionCreationError::AfterSubmission { error, completion }) => {
549 branch
550 .attach_submission_completion(completion)
551 .map_err(RealtimeFrameCoordinatorError::CompletionAttachment)?;
552 return Err(RealtimeFrameCoordinatorError::CompletionAfterSubmission(
553 error,
554 ));
555 }
556 };
557 let retained_resources = completion_mechanism.retained_resources(&completion);
558
559 *branch.schedule_state_mut() = schedule_state;
560 *branch.model_state_mut().payload_history_mut() = payload_history;
561 branch
562 .attach_submission_completion(completion.clone())
563 .map_err(RealtimeFrameCoordinatorError::CompletionAttachment)?;
564 Ok(SubmittedRealtimeFrame {
565 frame: completed,
566 completion,
567 retained_resources,
568 })
569}
570
571#[derive(Debug, thiserror::Error)]
573#[non_exhaustive]
574pub enum RealtimeFrameCoordinatorError<H, F, E, B, K> {
575 #[error(transparent)]
577 PayloadContract(RealtimePayloadContractError),
578 #[error(transparent)]
580 Ingress(RealtimeIngressError),
581 #[error("realtime host payload materialization failed")]
583 Materialization(H),
584 #[error("realtime frame interpretation failed")]
586 Interpretation(RealtimeFrameInterpretationError<F>),
587 #[error(transparent)]
589 DecisionPlan(SequentialDecisionPlanError),
590 #[error("realtime ordered decision failed")]
592 Decision(SequentialDecisionError<B>),
593 #[error("realtime prepared model execution failed")]
595 Model(E),
596 #[error("realtime exact completion creation failed")]
598 Completion(K),
599 #[error("realtime exact completion creation failed after native submission")]
601 CompletionAfterSubmission(K),
602 #[error(transparent)]
604 CompletionAttachment(RealtimeCompletionAttachmentError),
605}
606
607#[cfg(test)]
608mod tests {
609 use std::{cell::Cell, convert::Infallible, rc::Rc, time::Instant};
610
611 use eredu_core::{
612 scheduler::{RequestId, Scheduler, SchedulerLimits, SemanticStateTransaction},
613 Completion, RealtimeFrameConvention, RealtimeSpeechConfig, TokenFilter,
614 };
615
616 use super::*;
617 use crate::{
618 PenaltyConfig, RealtimeGenerationState, RealtimePayloadGeneration,
619 RealtimePayloadOwnerIdentity, RealtimePayloadState, TokenDomain,
620 };
621
622 #[derive(Debug, Clone, Eq, PartialEq)]
623 struct Matrix {
624 values: Vec<i32>,
625 shape: [usize; 2],
626 }
627
628 #[derive(Debug, Clone, Copy, thiserror::Error)]
629 #[error("test mechanism failed")]
630 struct MechanismError;
631
632 #[derive(Default)]
633 struct Mechanisms {
634 calls: usize,
635 }
636
637 impl RealtimeHostTokenMaterializer for Mechanisms {
638 type Tensor = Matrix;
639 type Error = MechanismError;
640
641 fn materialize_i32(
642 &mut self,
643 values: &[i32],
644 shape: [usize; 2],
645 ) -> Result<Self::Tensor, Self::Error> {
646 self.calls += 1;
647 Ok(Matrix {
648 values: values.to_vec(),
649 shape,
650 })
651 }
652 }
653
654 impl RealtimeFrameTensorMechanisms for Mechanisms {
655 type Tensor = Matrix;
656 type Error = MechanismError;
657
658 fn column(
659 &mut self,
660 matrix: &Self::Tensor,
661 column: usize,
662 ) -> Result<Self::Tensor, Self::Error> {
663 self.calls += 1;
664 let columns = matrix.shape[1];
665 Ok(Matrix {
666 values: (0..matrix.shape[0])
667 .map(|row| matrix.values[row * columns + column])
668 .collect(),
669 shape: [matrix.shape[0], 1],
670 })
671 }
672
673 fn filled_column(&mut self, token: i32, batch: usize) -> Result<Self::Tensor, Self::Error> {
674 self.calls += 1;
675 Ok(Matrix {
676 values: vec![token; batch],
677 shape: [batch, 1],
678 })
679 }
680
681 fn stack_columns(
682 &mut self,
683 columns: &[Self::Tensor],
684 batch: usize,
685 ) -> Result<Self::Tensor, Self::Error> {
686 self.calls += 1;
687 let values = (0..batch)
688 .flat_map(|row| columns.iter().map(move |column| column.values[row]))
689 .collect();
690 Ok(Matrix {
691 values,
692 shape: [batch, columns.len()],
693 })
694 }
695 }
696
697 struct TestBackend;
698
699 impl SamplingBackend for TestBackend {
700 type Logits = Matrix;
701 type Token = Matrix;
702 type RandomState = usize;
703 type Context = ();
704 type Error = MechanismError;
705
706 fn error(_message: String) -> Self::Error {
707 MechanismError
708 }
709
710 fn validate_token(
711 token: &Self::Token,
712 domain: TokenDomain,
713 _context: &Self::Context,
714 ) -> Result<Self::Token, Self::Error> {
715 token
716 .values
717 .iter()
718 .all(|value| {
719 usize::try_from(*value).is_ok_and(|value| value < domain.cardinality())
720 })
721 .then(|| token.clone())
722 .ok_or(MechanismError)
723 }
724
725 fn scale_temperature(
726 logits: &Self::Logits,
727 _temperature: f32,
728 _context: &Self::Context,
729 ) -> Result<Self::Logits, Self::Error> {
730 Ok(logits.clone())
731 }
732
733 fn apply_penalties(
734 logits: &Self::Logits,
735 _history: &[u32],
736 _penalties: PenaltyConfig,
737 _context: &Self::Context,
738 ) -> Result<Self::Logits, Self::Error> {
739 Ok(logits.clone())
740 }
741
742 fn apply_top_k(
743 logits: Self::Logits,
744 _top_k: i32,
745 _context: &Self::Context,
746 ) -> Result<Self::Logits, Self::Error> {
747 Ok(logits)
748 }
749
750 fn apply_top_p(
751 logits: Self::Logits,
752 _top_p: f32,
753 _context: &Self::Context,
754 ) -> Result<Self::Logits, Self::Error> {
755 Ok(logits)
756 }
757
758 fn apply_min_p(
759 logits: Self::Logits,
760 _min_p: f32,
761 _context: &Self::Context,
762 ) -> Result<Self::Logits, Self::Error> {
763 Ok(logits)
764 }
765
766 fn apply_token_filter(
767 logits: &Self::Logits,
768 _filter: &TokenFilter,
769 _context: &Self::Context,
770 ) -> Result<Self::Logits, Self::Error> {
771 Ok(logits.clone())
772 }
773
774 fn apply_mirostat(
775 logits: &Self::Logits,
776 _history: &[u32],
777 _penalties: PenaltyConfig,
778 _tau: f32,
779 _eta: f32,
780 _context: &Self::Context,
781 ) -> Result<Self::Logits, Self::Error> {
782 Ok(logits.clone())
783 }
784
785 fn sample_raw(
786 logits: &Self::Logits,
787 _temperature: f32,
788 random: Option<&mut Self::RandomState>,
789 _context: &Self::Context,
790 ) -> Result<Self::Token, Self::Error> {
791 if let Some(random) = random {
792 *random += 1;
793 }
794 Ok(logits.clone())
795 }
796
797 fn sample_processed(
798 logits: &Self::Logits,
799 temperature: f32,
800 random: Option<&mut Self::RandomState>,
801 context: &Self::Context,
802 ) -> Result<Self::Token, Self::Error> {
803 Self::sample_raw(logits, temperature, random, context)
804 }
805
806 fn token_id(token: &Self::Token, _context: &Self::Context) -> Result<u32, Self::Error> {
807 token
808 .values
809 .first()
810 .copied()
811 .and_then(|value| u32::try_from(value).ok())
812 .ok_or(MechanismError)
813 }
814
815 fn token_probability(
816 _logits: &Self::Logits,
817 _token: u32,
818 _context: &Self::Context,
819 ) -> Result<f32, Self::Error> {
820 Ok(1.0)
821 }
822 }
823
824 #[derive(Debug, Clone, Eq, PartialEq)]
825 struct TestSampler;
826
827 impl Sampler<TestBackend> for TestSampler {
828 fn sample(
829 &mut self,
830 logits: &Matrix,
831 temperature: f32,
832 random: Option<&mut usize>,
833 context: &(),
834 ) -> Result<Matrix, MechanismError> {
835 TestBackend::sample_raw(logits, temperature, random, context)
836 }
837 }
838
839 #[derive(Debug, Clone)]
840 struct ModelState {
841 executions: usize,
842 discards: Rc<Cell<usize>>,
843 }
844
845 #[derive(Debug, Clone, Copy, thiserror::Error)]
846 #[error("test model-state transaction failed")]
847 struct ModelStateError;
848
849 impl SemanticStateTransaction for ModelState {
850 type Branch = Self;
851 type Error = ModelStateError;
852
853 fn branch(&self) -> Result<Self::Branch, Self::Error> {
854 Ok(self.clone())
855 }
856
857 fn commit_branch(&mut self, branch: Self::Branch) -> Result<(), Self::Error> {
858 *self = branch;
859 Ok(())
860 }
861
862 fn discard_branch(branch: Self::Branch) -> Result<(), Self::Error> {
863 branch.discards.set(branch.discards.get() + 1);
864 Ok(())
865 }
866 }
867
868 #[derive(Debug, Clone, Default)]
869 struct TestCompletion {
870 waits: Rc<Cell<usize>>,
871 }
872
873 #[derive(Debug, Clone, Copy, thiserror::Error)]
874 #[error("test completion failed")]
875 struct CompletionError;
876
877 impl Completion for TestCompletion {
878 type Error = CompletionError;
879
880 fn is_complete(&self) -> Result<bool, Self::Error> {
881 Ok(true)
882 }
883
884 fn wait(&self) -> Result<(), Self::Error> {
885 self.waits.set(self.waits.get() + 1);
886 Ok(())
887 }
888 }
889
890 struct Executor {
891 fail: bool,
892 calls: usize,
893 }
894
895 #[derive(Debug, Clone, Copy, thiserror::Error)]
896 #[error("test model execution failed")]
897 struct ExecutionError;
898
899 impl PreparedRealtimeFrameExecutor<TestBackend, TestSampler, ModelState> for Executor {
900 type Error = ExecutionError;
901 type Retained = Matrix;
902
903 fn execute(
904 &mut self,
905 model_state: &mut ModelState,
906 _temporal: &[Matrix],
907 driver: &mut SequentialDecisionDriver<TestBackend, TestSampler>,
908 _context: &(),
909 ) -> Result<Self::Retained, Self::Error> {
910 self.calls += 1;
911 model_state.executions += 1;
912 if self.fail {
913 return Err(ExecutionError);
914 }
915 for prediction in 0..driver.plan().len() {
916 let value = if prediction == 0 { 2 } else { 6 };
917 let domain = TokenDomain::new(if prediction == 0 { 10 } else { 9 });
918 driver
919 .resolve(
920 prediction,
921 &Matrix {
922 values: vec![value],
923 shape: [1, 1],
924 },
925 domain,
926 &(),
927 )
928 .map_err(|_| ExecutionError)?;
929 }
930 Ok(Matrix {
931 values: vec![7],
932 shape: [1, 1],
933 })
934 }
935 }
936
937 #[derive(Default)]
938 struct CompletionMechanism {
939 calls: usize,
940 retained_calls: usize,
941 fail_after_submission: bool,
942 waits: Rc<Cell<usize>>,
943 }
944
945 impl RealtimeFrameCompletionMechanism<Matrix, ModelState, Matrix> for CompletionMechanism {
946 type Completion = TestCompletion;
947 type Error = CompletionError;
948
949 fn complete(
950 &mut self,
951 _input: MaterializedRealtimeInput<Matrix>,
952 _output: &CompletedRealtimeFrame<Matrix, Matrix>,
953 _model_state: &ModelState,
954 _payload_history: &RealtimePayloadHistory<Matrix>,
955 execution: Option<Matrix>,
956 ) -> Result<Self::Completion, RealtimeCompletionCreationError<Self::Completion, Self::Error>>
957 {
958 self.calls += 1;
959 if let Some(execution) = execution {
960 self.retained_calls += 1;
961 assert_eq!(execution.values, vec![7]);
962 }
963 let completion = TestCompletion {
964 waits: self.waits.clone(),
965 };
966 if self.fail_after_submission {
967 return Err(RealtimeCompletionCreationError::after_submission(
968 CompletionError,
969 completion,
970 ));
971 }
972 Ok(completion)
973 }
974 }
975
976 #[derive(Clone)]
977 struct ControlledCompletion {
978 ready: Rc<Cell<bool>>,
979 fail_wait: bool,
980 waits: Rc<Cell<usize>>,
981 }
982
983 impl Completion for ControlledCompletion {
984 type Error = CompletionError;
985
986 fn is_complete(&self) -> Result<bool, Self::Error> {
987 Ok(self.ready.get())
988 }
989
990 fn wait(&self) -> Result<(), Self::Error> {
991 self.waits.set(self.waits.get() + 1);
992 if self.fail_wait {
993 Err(CompletionError)
994 } else {
995 Ok(())
996 }
997 }
998 }
999
1000 #[derive(Debug, Clone, Copy, thiserror::Error)]
1001 #[error("test host observation failed")]
1002 struct ObservationError;
1003
1004 struct Observer {
1005 calls: Rc<Cell<usize>>,
1006 fail: bool,
1007 }
1008
1009 impl RealtimeFrameHostObserver<Matrix> for Observer {
1010 type Output = Vec<i32>;
1011 type Error = ObservationError;
1012
1013 fn observe(
1014 &mut self,
1015 frame: &CompletedRealtimeFrame<Matrix, Matrix>,
1016 ) -> Result<Self::Output, Self::Error> {
1017 self.calls.set(self.calls.get() + 1);
1018 if self.fail {
1019 Err(ObservationError)
1020 } else {
1021 Ok(frame.text().values.clone())
1022 }
1023 }
1024 }
1025
1026 #[derive(Clone)]
1027 struct PublicationState {
1028 value: usize,
1029 published: Rc<Cell<usize>>,
1030 discards: Rc<Cell<usize>>,
1031 }
1032
1033 impl SemanticStateTransaction for PublicationState {
1034 type Branch = Self;
1035 type Error = Infallible;
1036
1037 fn branch(&self) -> Result<Self::Branch, Self::Error> {
1038 Ok(self.clone())
1039 }
1040
1041 fn commit_branch(&mut self, branch: Self::Branch) -> Result<(), Self::Error> {
1042 branch.published.set(branch.value);
1043 *self = branch;
1044 Ok(())
1045 }
1046
1047 fn discard_branch(branch: Self::Branch) -> Result<(), Self::Error> {
1048 branch.discards.set(branch.discards.get() + 1);
1049 Ok(())
1050 }
1051 }
1052
1053 type Generation = RealtimeGenerationState<
1054 RealtimePayloadState<ModelState, Matrix>,
1055 TestSampler,
1056 usize,
1057 TestCompletion,
1058 >;
1059
1060 fn schedule(convention: RealtimeFrameConvention) -> RealtimeSpeechConfig {
1061 RealtimeSpeechConfig::new(2, 1, 1, 1, 9, 8, convention, vec![0, 0, 1]).unwrap()
1062 }
1063
1064 fn state(schedule: &RealtimeSpeechConfig, discards: Rc<Cell<usize>>) -> Generation {
1065 let payload = RealtimePayloadState::new(
1066 ModelState {
1067 executions: 0,
1068 discards,
1069 },
1070 RealtimePayloadHistory::new(schedule.clone()),
1071 schedule,
1072 )
1073 .unwrap();
1074 Generation::new(
1075 payload,
1076 schedule.clone(),
1077 eredu_core::RealtimeSampling::greedy(),
1078 vec![TestSampler, TestSampler],
1079 Some(0),
1080 )
1081 .unwrap()
1082 }
1083
1084 fn contract(schedule: &RealtimeSpeechConfig) -> RealtimeIngressContract {
1085 RealtimeIngressContract::new(schedule.clone(), TokenDomain::new(10), TokenDomain::new(9))
1086 .unwrap()
1087 }
1088
1089 fn payload_contract(schedule: &RealtimeSpeechConfig) -> RealtimePayloadContract {
1090 RealtimePayloadContract::new(
1091 schedule.clone(),
1092 1,
1093 TokenDomain::new(10),
1094 TokenDomain::new(9),
1095 RealtimePayloadGeneration::new(1).unwrap(),
1096 RealtimePayloadOwnerIdentity::new(1).unwrap(),
1097 )
1098 .unwrap()
1099 }
1100
1101 fn frame() -> RealtimeInputFrame {
1102 RealtimeInputFrame::new(1, vec![4])
1103 .with_forced_text(vec![3])
1104 .with_partially_forced_generated_audio(vec![5], vec![false])
1105 }
1106
1107 fn submitted_test_frame<C>(completion: C) -> SubmittedRealtimeFrame<Matrix, C> {
1108 let schedule = schedule(RealtimeFrameConvention::FeedbackAlignedHistory);
1109 let state = state(&schedule, Rc::new(Cell::new(0)));
1110 let mut branch = state.branch().unwrap();
1111 let mut host = Mechanisms::default();
1112 let mut tensors = Mechanisms::default();
1113 let mut executor = Executor {
1114 fail: false,
1115 calls: 0,
1116 };
1117 let mut completion_mechanism = CompletionMechanism::default();
1118 let submitted = execute_realtime_frame::<TestBackend, _, _, _, _, _, _, _, _>(
1119 &contract(&schedule),
1120 &payload_contract(&schedule),
1121 &frame(),
1122 &mut branch,
1123 &RealtimeDecisionExecution::new(true),
1124 &mut host,
1125 &mut tensors,
1126 &mut executor,
1127 &mut completion_mechanism,
1128 &(),
1129 )
1130 .unwrap();
1131 let SubmittedRealtimeFrame {
1132 frame,
1133 retained_resources,
1134 ..
1135 } = submitted;
1136 SubmittedRealtimeFrame {
1137 frame,
1138 completion,
1139 retained_resources,
1140 }
1141 }
1142
1143 #[test]
1144 fn complete_frame_changes_only_the_unpublished_branch_until_commit() {
1145 let schedule = schedule(RealtimeFrameConvention::FeedbackAlignedHistory);
1146 let discards = Rc::new(Cell::new(0));
1147 let mut state = state(&schedule, discards);
1148 let mut branch = state.branch().unwrap();
1149 let mut host = Mechanisms::default();
1150 let mut tensors = Mechanisms::default();
1151 let mut executor = Executor {
1152 fail: false,
1153 calls: 0,
1154 };
1155 let mut completion = CompletionMechanism::default();
1156
1157 let output = execute_realtime_frame::<TestBackend, _, _, _, _, _, _, _, _>(
1158 &contract(&schedule),
1159 &payload_contract(&schedule),
1160 &frame(),
1161 &mut branch,
1162 &RealtimeDecisionExecution::new(true),
1163 &mut host,
1164 &mut tensors,
1165 &mut executor,
1166 &mut completion,
1167 &(),
1168 )
1169 .unwrap();
1170
1171 assert_eq!(output.frame().text().values, vec![3]);
1172 assert_eq!(output.frame().sampled_audio().values, vec![6]);
1173 assert_eq!(executor.calls, 1);
1174 assert_eq!(completion.calls, 1);
1175 assert_eq!(completion.retained_calls, 1);
1176 assert_eq!(state.schedule_state().frontier(), 0);
1177 assert_eq!(state.model_state().model_state().executions, 0);
1178 assert!(state.model_state().payload_history().is_empty());
1179
1180 state.commit_branch(branch).unwrap();
1181 assert_eq!(state.schedule_state().frontier(), 1);
1182 assert_eq!(state.model_state().model_state().executions, 1);
1183 assert!(!state.model_state().payload_history().is_empty());
1184 assert_eq!(state.random_state(), Some(&1));
1185 }
1186
1187 #[test]
1188 fn model_failure_requires_discard_and_never_changes_canonical_state() {
1189 let schedule = schedule(RealtimeFrameConvention::FeedbackAlignedHistory);
1190 let discards = Rc::new(Cell::new(0));
1191 let state = state(&schedule, discards.clone());
1192 let mut branch = state.branch().unwrap();
1193 let mut host = Mechanisms::default();
1194 let mut tensors = Mechanisms::default();
1195 let mut executor = Executor {
1196 fail: true,
1197 calls: 0,
1198 };
1199 let mut completion = CompletionMechanism::default();
1200
1201 assert!(matches!(
1202 execute_realtime_frame::<TestBackend, _, _, _, _, _, _, _, _>(
1203 &contract(&schedule),
1204 &payload_contract(&schedule),
1205 &frame(),
1206 &mut branch,
1207 &RealtimeDecisionExecution::new(true),
1208 &mut host,
1209 &mut tensors,
1210 &mut executor,
1211 &mut completion,
1212 &(),
1213 ),
1214 Err(RealtimeFrameCoordinatorError::Model(ExecutionError))
1215 ));
1216 assert_eq!(state.schedule_state().frontier(), 0);
1217 assert_eq!(state.model_state().model_state().executions, 0);
1218 assert!(state.model_state().payload_history().is_empty());
1219 assert_eq!(completion.calls, 0);
1220
1221 Generation::discard_branch(branch).unwrap();
1222 assert_eq!(discards.get(), 1);
1223 }
1224
1225 #[test]
1226 fn post_submission_completion_failure_is_quarantined_until_discard_waits() {
1227 let schedule = schedule(RealtimeFrameConvention::FeedbackAlignedHistory);
1228 let discards = Rc::new(Cell::new(0));
1229 let state = state(&schedule, discards.clone());
1230 let mut branch = state.branch().unwrap();
1231 let mut host = Mechanisms::default();
1232 let mut tensors = Mechanisms::default();
1233 let mut executor = Executor {
1234 fail: false,
1235 calls: 0,
1236 };
1237 let waits = Rc::new(Cell::new(0));
1238 let mut completion = CompletionMechanism {
1239 fail_after_submission: true,
1240 waits: waits.clone(),
1241 ..CompletionMechanism::default()
1242 };
1243
1244 assert!(matches!(
1245 execute_realtime_frame::<TestBackend, _, _, _, _, _, _, _, _>(
1246 &contract(&schedule),
1247 &payload_contract(&schedule),
1248 &frame(),
1249 &mut branch,
1250 &RealtimeDecisionExecution::new(true),
1251 &mut host,
1252 &mut tensors,
1253 &mut executor,
1254 &mut completion,
1255 &(),
1256 ),
1257 Err(RealtimeFrameCoordinatorError::CompletionAfterSubmission(
1258 CompletionError
1259 ))
1260 ));
1261 assert!(branch.has_submission_completion());
1262 assert_eq!(waits.get(), 0);
1263 assert_eq!(state.schedule_state().frontier(), 0);
1264 assert_eq!(state.model_state().model_state().executions, 0);
1265
1266 Generation::discard_branch(branch).unwrap();
1267 assert_eq!(waits.get(), 1);
1268 assert_eq!(discards.get(), 1);
1269 }
1270
1271 #[test]
1272 fn initialization_bypasses_the_nonempty_decision_plan_and_model() {
1273 let schedule = schedule(RealtimeFrameConvention::AbsoluteDelayedSlots);
1274 let discards = Rc::new(Cell::new(0));
1275 let state = state(&schedule, discards);
1276 let mut branch = state.branch().unwrap();
1277 let mut host = Mechanisms::default();
1278 let mut tensors = Mechanisms::default();
1279 let mut executor = Executor {
1280 fail: true,
1281 calls: 0,
1282 };
1283 let mut completion = CompletionMechanism::default();
1284
1285 let output = execute_realtime_frame::<TestBackend, _, _, _, _, _, _, _, _>(
1286 &contract(&schedule),
1287 &payload_contract(&schedule),
1288 &frame(),
1289 &mut branch,
1290 &RealtimeDecisionExecution::new(true),
1291 &mut host,
1292 &mut tensors,
1293 &mut executor,
1294 &mut completion,
1295 &(),
1296 )
1297 .unwrap();
1298
1299 assert_eq!(executor.calls, 0);
1300 assert_eq!(completion.calls, 1);
1301 assert_eq!(completion.retained_calls, 0);
1302 assert_eq!(output.frame().text().values, vec![9]);
1303 assert_eq!(branch.schedule_state().frontier(), 1);
1304 }
1305
1306 #[test]
1307 fn branch_and_schedule_preflight_precede_every_native_mechanism() {
1308 let active_schedule = schedule(RealtimeFrameConvention::FeedbackAlignedHistory);
1309 let state = state(&active_schedule, Rc::new(Cell::new(0)));
1310 let mut branch = state.branch().unwrap();
1311 branch
1312 .attach_submission_completion(TestCompletion::default())
1313 .unwrap();
1314 let mut host = Mechanisms::default();
1315 let mut tensors = Mechanisms::default();
1316 let mut executor = Executor {
1317 fail: false,
1318 calls: 0,
1319 };
1320 let mut completion = CompletionMechanism::default();
1321
1322 assert!(matches!(
1323 execute_realtime_frame::<TestBackend, _, _, _, _, _, _, _, _>(
1324 &contract(&active_schedule),
1325 &payload_contract(&active_schedule),
1326 &frame(),
1327 &mut branch,
1328 &RealtimeDecisionExecution::new(true),
1329 &mut host,
1330 &mut tensors,
1331 &mut executor,
1332 &mut completion,
1333 &(),
1334 ),
1335 Err(RealtimeFrameCoordinatorError::CompletionAttachment(
1336 RealtimeCompletionAttachmentError::AlreadyAttached
1337 ))
1338 ));
1339 assert_eq!(host.calls, 0);
1340 assert_eq!(tensors.calls, 0);
1341 assert_eq!(executor.calls, 0);
1342 assert_eq!(completion.calls, 0);
1343
1344 let mut branch = state.branch().unwrap();
1345 let mismatched = schedule(RealtimeFrameConvention::AbsoluteDelayedSlots);
1346 assert!(matches!(
1347 execute_realtime_frame::<TestBackend, _, _, _, _, _, _, _, _>(
1348 &contract(&mismatched),
1349 &payload_contract(&mismatched),
1350 &frame(),
1351 &mut branch,
1352 &RealtimeDecisionExecution::new(true),
1353 &mut host,
1354 &mut tensors,
1355 &mut executor,
1356 &mut completion,
1357 &(),
1358 ),
1359 Err(RealtimeFrameCoordinatorError::Interpretation(
1360 RealtimeFrameInterpretationError::Schedule(_)
1361 ))
1362 ));
1363 assert_eq!(host.calls, 0);
1364 assert_eq!(tensors.calls, 0);
1365 assert_eq!(executor.calls, 0);
1366 assert_eq!(completion.calls, 0);
1367 }
1368
1369 #[test]
1370 fn prepublication_observation_waits_for_exact_completion_and_is_cached_once() {
1371 let ready = Rc::new(Cell::new(false));
1372 let waits = Rc::new(Cell::new(0));
1373 let observer_calls = Rc::new(Cell::new(0));
1374 let transition = PrepublicationRealtimeFrame::new(
1375 submitted_test_frame(ControlledCompletion {
1376 ready: ready.clone(),
1377 fail_wait: false,
1378 waits: waits.clone(),
1379 }),
1380 Observer {
1381 calls: observer_calls.clone(),
1382 fail: false,
1383 },
1384 );
1385
1386 assert!(!transition.is_complete().unwrap());
1387 assert_eq!(waits.get(), 0);
1388 assert_eq!(observer_calls.get(), 0);
1389
1390 ready.set(true);
1391 assert!(transition.is_complete().unwrap());
1392 assert!(transition.is_complete().unwrap());
1393 assert_eq!(waits.get(), 1);
1394 assert_eq!(observer_calls.get(), 1);
1395 assert_eq!(transition.into_host_output().unwrap(), vec![3]);
1396 }
1397
1398 #[test]
1399 fn completion_and_observation_failures_never_expose_host_output() {
1400 let observer_calls = Rc::new(Cell::new(0));
1401 let completion_failure = PrepublicationRealtimeFrame::new(
1402 submitted_test_frame(ControlledCompletion {
1403 ready: Rc::new(Cell::new(true)),
1404 fail_wait: true,
1405 waits: Rc::new(Cell::new(0)),
1406 }),
1407 Observer {
1408 calls: observer_calls.clone(),
1409 fail: false,
1410 },
1411 );
1412 assert!(matches!(
1413 completion_failure.is_complete(),
1414 Err(RealtimePrepublicationError::Completion(CompletionError))
1415 ));
1416 assert_eq!(observer_calls.get(), 0);
1417 assert!(matches!(
1418 completion_failure.into_host_output(),
1419 Err(RealtimeHostOutputUnavailable::Pending)
1420 ));
1421
1422 let observer_calls = Rc::new(Cell::new(0));
1423 let observation_failure = PrepublicationRealtimeFrame::new(
1424 submitted_test_frame(ControlledCompletion {
1425 ready: Rc::new(Cell::new(true)),
1426 fail_wait: false,
1427 waits: Rc::new(Cell::new(0)),
1428 }),
1429 Observer {
1430 calls: observer_calls.clone(),
1431 fail: true,
1432 },
1433 );
1434 assert!(matches!(
1435 observation_failure.is_complete(),
1436 Err(RealtimePrepublicationError::Observation(_))
1437 ));
1438 assert!(matches!(
1439 observation_failure.is_complete(),
1440 Err(RealtimePrepublicationError::Observation(_))
1441 ));
1442 assert_eq!(observer_calls.get(), 1);
1443 assert!(matches!(
1444 observation_failure.into_host_output(),
1445 Err(RealtimeHostOutputUnavailable::Observation(_))
1446 ));
1447 }
1448
1449 #[test]
1450 fn scheduler_commits_only_after_host_observation_succeeds() {
1451 let limits = SchedulerLimits::new(1, 1).unwrap();
1452 let request = RequestId::new(7);
1453 let published = Rc::new(Cell::new(0));
1454 let discards = Rc::new(Cell::new(0));
1455 let mut scheduler = Scheduler::new(limits).unwrap();
1456 scheduler
1457 .register(
1458 request,
1459 PublicationState {
1460 value: 0,
1461 published: published.clone(),
1462 discards: discards.clone(),
1463 },
1464 )
1465 .unwrap();
1466 scheduler.enqueue(request, frame()).unwrap();
1467 let observer_calls = Rc::new(Cell::new(0));
1468 let progress = scheduler
1469 .run_local_turn(Instant::now(), |_, _, branch| {
1470 branch.value += 1;
1471 Ok::<_, Infallible>(PrepublicationRealtimeFrame::new(
1472 submitted_test_frame(ControlledCompletion {
1473 ready: Rc::new(Cell::new(true)),
1474 fail_wait: false,
1475 waits: Rc::new(Cell::new(0)),
1476 }),
1477 Observer {
1478 calls: observer_calls.clone(),
1479 fail: true,
1480 },
1481 ))
1482 })
1483 .unwrap();
1484 assert!(progress.committed.is_empty());
1485 assert_eq!(progress.failed.len(), 1);
1486 assert_eq!(published.get(), 0);
1487 assert_eq!(discards.get(), 1);
1488 assert_eq!(observer_calls.get(), 1);
1489
1490 let request = RequestId::new(8);
1491 let published = Rc::new(Cell::new(0));
1492 let discards = Rc::new(Cell::new(0));
1493 let mut scheduler = Scheduler::new(limits).unwrap();
1494 scheduler
1495 .register(
1496 request,
1497 PublicationState {
1498 value: 0,
1499 published: published.clone(),
1500 discards: discards.clone(),
1501 },
1502 )
1503 .unwrap();
1504 scheduler.enqueue(request, frame()).unwrap();
1505 let observer_calls = Rc::new(Cell::new(0));
1506 let mut progress = scheduler
1507 .run_local_turn(Instant::now(), |_, _, branch| {
1508 branch.value += 1;
1509 Ok::<_, Infallible>(PrepublicationRealtimeFrame::new(
1510 submitted_test_frame(ControlledCompletion {
1511 ready: Rc::new(Cell::new(true)),
1512 fail_wait: false,
1513 waits: Rc::new(Cell::new(0)),
1514 }),
1515 Observer {
1516 calls: observer_calls.clone(),
1517 fail: false,
1518 },
1519 ))
1520 })
1521 .unwrap();
1522 assert!(progress.failed.is_empty());
1523 assert_eq!(progress.committed.len(), 1);
1524 assert_eq!(published.get(), 1);
1525 assert_eq!(discards.get(), 0);
1526 assert_eq!(observer_calls.get(), 1);
1527 let (_, _, transition) = progress.committed.pop().unwrap();
1528 assert_eq!(transition.into_host_output().unwrap(), vec![3]);
1529 }
1530}