1use crate::{KvCacheHandle, RecurrentStateHandle, RecurrentStateSpec, TensorRef};
7use async_trait::async_trait;
8use ferrum_types::{ExecutorAdmissionLimits, FerrumError, ModelInfo, RequestId, Result, TokenId};
9use serde::{Deserialize, Serialize};
10use std::{
11 collections::{hash_map::DefaultHasher, HashMap, HashSet},
12 future::Future,
13 hash::{Hash, Hasher},
14 num::NonZeroU64,
15 ops::Range,
16 pin::Pin,
17 sync::Arc,
18};
19
20mod prefix_capture;
21mod prefix_restore;
22pub use prefix_capture::{
23 PrefixCaptureBoundary, PrefixCaptureLease, PrefixCapturePlan, PrefixCaptureRequest,
24 PrefixCaptureStatus,
25};
26pub use prefix_restore::{
27 PlanRuntimePrefixRestoreDeferral, PlanRuntimePrefixRestoreInput,
28 PlanRuntimePrefixRestoreOutcome, PlanRuntimePrefixRestoreOutput, PrefixRestoreDecision,
29 PrefixRestoreObservation, PrefixRestoreSource,
30};
31
32#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40pub struct KvSlotRequest {
41 pub cache_id: String,
42 pub target_len: usize,
43 pub admission_target_len: Option<usize>,
44}
45
46#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
48pub struct KvSlotAllocation {
49 pub cache_id: String,
50 pub blocks_before: usize,
51 pub blocks_after: usize,
52 pub new_blocks: usize,
53}
54
55#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61pub struct KvSlotReservation {
62 pub block_size: usize,
63 pub total_blocks: usize,
64 pub free_blocks_before: usize,
65 pub free_blocks_after: usize,
66 pub allocations: Vec<KvSlotAllocation>,
67}
68
69#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct KvSlotCapacitySnapshot {
77 pub block_size: usize,
78 pub total_blocks: usize,
79 pub free_blocks: usize,
80}
81
82#[derive(Clone)]
89pub struct TokenSelectionMask {
90 pub fingerprint: u64,
91 pub valid_token_mask: Arc<[i8]>,
92}
93
94impl TokenSelectionMask {
95 pub fn new(valid_token_mask: Vec<i8>) -> Self {
96 let fingerprint = Self::fingerprint(&valid_token_mask);
97 Self {
98 fingerprint,
99 valid_token_mask: Arc::from(valid_token_mask),
100 }
101 }
102
103 fn fingerprint(valid_token_mask: &[i8]) -> u64 {
104 let mut hasher = DefaultHasher::new();
105 valid_token_mask.hash(&mut hasher);
106 hasher.finish()
107 }
108
109 pub fn set_tokens_validity(&mut self, token_ids: &[u32], valid: bool) -> bool {
114 let value = i8::from(valid);
115 let slots = Arc::make_mut(&mut self.valid_token_mask);
116 let mut changed = false;
117 for &token_id in token_ids {
118 if let Some(slot) = slots.get_mut(token_id as usize) {
119 if *slot != value {
120 *slot = value;
121 changed = true;
122 }
123 }
124 }
125 if changed {
126 self.fingerprint = Self::fingerprint(slots);
127 }
128 changed
129 }
130
131 pub fn len(&self) -> usize {
132 self.valid_token_mask.len()
133 }
134
135 pub fn is_empty(&self) -> bool {
136 self.valid_token_mask.is_empty()
137 }
138}
139
140impl std::fmt::Debug for TokenSelectionMask {
141 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
142 let valid_count = self.valid_token_mask.iter().filter(|&&v| v != 0).count();
143 f.debug_struct("TokenSelectionMask")
144 .field("fingerprint", &self.fingerprint)
145 .field("len", &self.valid_token_mask.len())
146 .field("valid_count", &valid_count)
147 .finish()
148 }
149}
150
151#[cfg(test)]
152mod token_selection_mask_tests {
153 use super::TokenSelectionMask;
154
155 #[test]
156 fn response_completion_mask_is_copy_on_write_and_restores_fingerprint() {
157 let mut mask = TokenSelectionMask::new(vec![1, 1, 1]);
158 let original = mask.clone();
159 let original_fingerprint = mask.fingerprint;
160
161 assert!(mask.set_tokens_validity(&[1], false));
162 assert_eq!(mask.valid_token_mask.as_ref(), &[1, 0, 1]);
163 assert_eq!(original.valid_token_mask.as_ref(), &[1, 1, 1]);
164 assert_ne!(mask.fingerprint, original_fingerprint);
165
166 let masked_fingerprint = mask.fingerprint;
167 assert!(!mask.set_tokens_validity(&[1], false));
168 assert_eq!(mask.fingerprint, masked_fingerprint);
169
170 assert!(mask.set_tokens_validity(&[1], true));
171 assert_eq!(mask.fingerprint, original_fingerprint);
172 }
173}
174
175#[derive(Clone, Debug)]
176pub enum LogitsReturnPolicy {
177 FullLogits,
178 GreedyArgmax {
179 token_mask: Option<TokenSelectionMask>,
180 repetition_penalty: Option<GreedyRepetitionPenalty>,
181 },
182}
183
184impl Default for LogitsReturnPolicy {
185 fn default() -> Self {
186 Self::FullLogits
187 }
188}
189
190impl LogitsReturnPolicy {
191 pub fn requires_full_logits(&self) -> bool {
192 matches!(self, Self::FullLogits)
193 }
194}
195
196#[derive(Debug, Clone, PartialEq)]
203pub enum ExecutorSamplingOutput {
204 FullLogits(Vec<f32>),
205 GreedyToken(TokenId),
206}
207
208impl ExecutorSamplingOutput {
209 pub fn full_logits(logits: Vec<f32>) -> Result<Self> {
210 if logits.is_empty() {
211 return Err(FerrumError::backend(
212 "plan-runtime sampling output requires non-empty logits",
213 ));
214 }
215 Ok(Self::FullLogits(logits))
216 }
217
218 pub const fn greedy_token(token: TokenId) -> Self {
219 Self::GreedyToken(token)
220 }
221
222 pub fn validate_for_policy(
228 &self,
229 policy: &LogitsReturnPolicy,
230 vocabulary_size: usize,
231 ) -> Result<()> {
232 match self {
233 Self::FullLogits(logits) if logits.len() != vocabulary_size => {
234 return Err(FerrumError::backend(format!(
235 "plan runtime returned {} logits for vocabulary {vocabulary_size}",
236 logits.len()
237 )));
238 }
239 Self::GreedyToken(_) if policy.requires_full_logits() => {
240 return Err(FerrumError::backend(
241 "plan runtime returned a greedy token for a full-logits request",
242 ));
243 }
244 Self::GreedyToken(token)
245 if usize::try_from(token.get())
246 .ok()
247 .is_none_or(|token| token >= vocabulary_size) =>
248 {
249 return Err(FerrumError::backend(format!(
250 "plan runtime returned token {} outside vocabulary {vocabulary_size}",
251 token.get()
252 )));
253 }
254 _ => {}
255 }
256 Ok(())
257 }
258
259 pub fn into_full_logits(self) -> Result<Vec<f32>> {
260 match self {
261 Self::FullLogits(logits) => Ok(logits),
262 Self::GreedyToken(_) => Err(FerrumError::backend(
263 "plan-runtime prefill unexpectedly returned a selected token",
264 )),
265 }
266 }
267}
268
269#[cfg(test)]
270mod executor_sampling_output_tests {
271 use super::{ExecutorSamplingOutput, LogitsReturnPolicy};
272 use ferrum_types::TokenId;
273
274 #[test]
275 fn full_logits_require_exact_vocabulary_width() {
276 let output = ExecutorSamplingOutput::full_logits(vec![0.0; 4]).unwrap();
277 assert!(output
278 .validate_for_policy(&LogitsReturnPolicy::FullLogits, 4)
279 .is_ok());
280 assert!(output
281 .validate_for_policy(&LogitsReturnPolicy::FullLogits, 5)
282 .is_err());
283 }
284
285 #[test]
286 fn greedy_token_requires_greedy_policy_and_in_vocabulary_token() {
287 let allowed = LogitsReturnPolicy::GreedyArgmax {
288 token_mask: None,
289 repetition_penalty: None,
290 };
291 let output = ExecutorSamplingOutput::greedy_token(TokenId::new(3));
292 assert!(output.validate_for_policy(&allowed, 4).is_ok());
293 assert!(output.validate_for_policy(&allowed, 3).is_err());
294 assert!(output
295 .validate_for_policy(&LogitsReturnPolicy::FullLogits, 4)
296 .is_err());
297 }
298
299 #[test]
300 fn full_logits_are_a_legal_greedy_batch_fallback() {
301 let policy = LogitsReturnPolicy::GreedyArgmax {
302 token_mask: None,
303 repetition_penalty: None,
304 };
305 let output = ExecutorSamplingOutput::full_logits(vec![0.0; 4]).unwrap();
306 assert!(output.validate_for_policy(&policy, 4).is_ok());
307 }
308}
309
310#[derive(Clone, Debug)]
316pub struct GreedyRepetitionPenalty {
317 penalty: f32,
318 token_ids: Arc<[u32]>,
319}
320
321impl GreedyRepetitionPenalty {
322 pub fn new(penalty: f32, mut token_ids: Vec<u32>) -> Self {
323 let mut seen = HashSet::with_capacity(token_ids.len());
324 token_ids.retain(|token| seen.insert(*token));
325 Self {
326 penalty,
327 token_ids: Arc::from(token_ids),
328 }
329 }
330
331 pub const fn penalty(&self) -> f32 {
332 self.penalty
333 }
334
335 pub fn token_ids(&self) -> &[u32] {
336 &self.token_ids
337 }
338
339 pub fn is_empty(&self) -> bool {
340 self.token_ids.is_empty() || self.penalty == 1.0
341 }
342}
343
344#[cfg(test)]
345mod greedy_repetition_penalty_tests {
346 use super::GreedyRepetitionPenalty;
347
348 #[test]
349 fn constructor_preserves_first_seen_order_and_removes_duplicates() {
350 let repetition = GreedyRepetitionPenalty::new(1.1, vec![7, 3, 7, 9, 3]);
351 assert_eq!(repetition.penalty(), 1.1);
352 assert_eq!(repetition.token_ids(), [7, 3, 9]);
353 }
354}
355
356#[derive(Debug, Clone)]
358pub struct PrefillInput {
359 pub request_id: Option<RequestId>,
361 pub maximum_sequence_tokens: Option<usize>,
364 pub chunk: Option<PrefillChunk>,
369 pub input_ids: TensorRef,
371 pub attention_mask: Option<TensorRef>,
373 pub position_ids: Option<TensorRef>,
375 pub kv_cache: Option<Arc<dyn KvCacheHandle>>,
377 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
379 pub metadata: HashMap<String, serde_json::Value>,
381}
382
383impl PrefillInput {
384 pub fn new(input_ids: TensorRef) -> Self {
386 Self {
387 request_id: None,
388 maximum_sequence_tokens: None,
389 chunk: None,
390 input_ids,
391 attention_mask: None,
392 position_ids: None,
393 kv_cache: None,
394 recurrent_state: None,
395 metadata: HashMap::new(),
396 }
397 }
398
399 pub fn with_request_context(
401 mut self,
402 request_id: RequestId,
403 maximum_sequence_tokens: usize,
404 ) -> Self {
405 self.request_id = Some(request_id);
406 self.maximum_sequence_tokens = Some(maximum_sequence_tokens);
407 self
408 }
409
410 pub fn with_chunk(mut self, chunk: PrefillChunk) -> Self {
412 self.chunk = Some(chunk);
413 self
414 }
415
416 pub fn with_kv_cache(mut self, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
418 self.kv_cache = Some(kv_cache);
419 self
420 }
421
422 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
424 self.recurrent_state = Some(recurrent_state);
425 self
426 }
427
428 pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
430 self.metadata = metadata;
431 self
432 }
433
434 pub fn with_attention_mask(mut self, mask: TensorRef) -> Self {
436 self.attention_mask = Some(mask);
437 self
438 }
439
440 pub fn with_position_ids(mut self, positions: TensorRef) -> Self {
442 self.position_ids = Some(positions);
443 self
444 }
445
446 pub fn batch_size(&self) -> usize {
448 self.input_ids.shape()[0]
449 }
450
451 pub fn sequence_length(&self) -> usize {
453 if self.input_ids.shape().len() >= 2 {
454 self.input_ids.shape()[1]
455 } else {
456 1
457 }
458 }
459}
460
461#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
463pub struct PrefillChunk {
464 tokens_processed: usize,
465 tokens_to_process: usize,
466 total_prompt_tokens: usize,
467}
468
469#[cfg(test)]
470mod prefill_chunk_tests {
471 use super::PrefillChunk;
472
473 #[test]
474 fn validates_exact_progress_and_finality() {
475 let first = PrefillChunk::new(0, 3, 8).unwrap();
476 assert_eq!(first.range(), 0..3);
477 assert_eq!(first.end(), 3);
478 assert!(!first.is_final());
479
480 let final_chunk = PrefillChunk::new(3, 5, 8).unwrap();
481 assert_eq!(final_chunk.range(), 3..8);
482 assert!(final_chunk.is_final());
483 }
484
485 #[test]
486 fn rejects_empty_out_of_bounds_and_overflowing_progress() {
487 assert!(PrefillChunk::new(0, 0, 8).is_err());
488 assert!(PrefillChunk::new(0, 1, 0).is_err());
489 assert!(PrefillChunk::new(7, 2, 8).is_err());
490 assert!(PrefillChunk::new(usize::MAX, 1, usize::MAX).is_err());
491 }
492}
493
494impl PrefillChunk {
495 pub fn new(
496 tokens_processed: usize,
497 tokens_to_process: usize,
498 total_prompt_tokens: usize,
499 ) -> Result<Self> {
500 let end = tokens_processed
501 .checked_add(tokens_to_process)
502 .ok_or_else(|| {
503 ferrum_types::FerrumError::request_validation("prefill chunk overflows")
504 })?;
505 if tokens_to_process == 0 || total_prompt_tokens == 0 || end > total_prompt_tokens {
506 return Err(ferrum_types::FerrumError::request_validation(
507 "prefill chunk must be non-empty and within the full prompt",
508 ));
509 }
510 Ok(Self {
511 tokens_processed,
512 tokens_to_process,
513 total_prompt_tokens,
514 })
515 }
516
517 pub const fn tokens_processed(self) -> usize {
518 self.tokens_processed
519 }
520
521 pub const fn tokens_to_process(self) -> usize {
522 self.tokens_to_process
523 }
524
525 pub const fn total_prompt_tokens(self) -> usize {
526 self.total_prompt_tokens
527 }
528
529 pub fn range(self) -> Range<usize> {
530 self.tokens_processed..self.tokens_processed + self.tokens_to_process
531 }
532
533 pub const fn end(self) -> usize {
534 self.tokens_processed + self.tokens_to_process
535 }
536
537 pub const fn is_final(self) -> bool {
538 self.end() == self.total_prompt_tokens
539 }
540}
541
542#[derive(Debug, Clone)]
544pub struct PrefillOutput {
545 pub logits: TensorRef,
547 pub kv_cache: Arc<dyn KvCacheHandle>,
549 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
551 pub hidden_states: Option<Vec<TensorRef>>,
553 pub attention_weights: Option<Vec<TensorRef>>,
555}
556
557impl PrefillOutput {
558 pub fn new(logits: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
560 Self {
561 logits,
562 kv_cache,
563 recurrent_state: None,
564 hidden_states: None,
565 attention_weights: None,
566 }
567 }
568
569 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
571 self.recurrent_state = Some(recurrent_state);
572 self
573 }
574
575 pub fn last_token_logits(&self) -> Result<TensorRef> {
577 let shape = self.logits.shape();
578 if shape.len() != 3 {
579 return Err(ferrum_types::FerrumError::backend(
580 "Expected 3D logits tensor [batch, seq, vocab]",
581 ));
582 }
583
584 let seq_len = shape[1];
585 if seq_len == 0 {
586 return Err(ferrum_types::FerrumError::backend("Empty sequence"));
587 }
588
589 self.logits
591 .view(&[0, seq_len - 1, 0], &[shape[0], seq_len, shape[2]])
592 }
593}
594
595#[derive(Debug, Clone)]
597pub struct DecodeInput {
598 pub request_id: Option<RequestId>,
600 pub input_ids: TensorRef,
602 pub kv_cache: Arc<dyn KvCacheHandle>,
604 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
606 pub position_ids: Option<TensorRef>,
608 pub metadata: HashMap<String, serde_json::Value>,
610 pub logits_policy: LogitsReturnPolicy,
612}
613
614impl DecodeInput {
615 pub fn new(input_ids: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
617 Self {
618 request_id: None,
619 input_ids,
620 kv_cache,
621 recurrent_state: None,
622 position_ids: None,
623 metadata: HashMap::new(),
624 logits_policy: LogitsReturnPolicy::FullLogits,
625 }
626 }
627
628 pub fn with_request_id(mut self, request_id: RequestId) -> Self {
630 self.request_id = Some(request_id);
631 self
632 }
633
634 pub fn with_position_ids(mut self, positions: TensorRef) -> Self {
636 self.position_ids = Some(positions);
637 self
638 }
639
640 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
642 self.recurrent_state = Some(recurrent_state);
643 self
644 }
645
646 pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
648 self.metadata = metadata;
649 self
650 }
651
652 pub fn with_logits_policy(mut self, policy: LogitsReturnPolicy) -> Self {
653 self.logits_policy = policy;
654 self
655 }
656
657 pub fn batch_size(&self) -> usize {
659 self.input_ids.shape()[0]
660 }
661}
662
663#[derive(Clone)]
677pub struct UnifiedBatchItem {
678 pub seq_id: String,
680 pub q_tokens: Vec<u32>,
683 pub kv_cache: Arc<dyn KvCacheHandle>,
685 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
687 pub pos_offset: usize,
691 pub is_final_chunk: bool,
696 pub metadata: HashMap<String, serde_json::Value>,
698 pub logits_policy: LogitsReturnPolicy,
700}
701
702impl std::fmt::Debug for UnifiedBatchItem {
703 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
704 f.debug_struct("UnifiedBatchItem")
705 .field("seq_id", &self.seq_id)
706 .field("q_len", &self.q_tokens.len())
707 .field("has_recurrent_state", &self.recurrent_state.is_some())
708 .field("pos_offset", &self.pos_offset)
709 .field("is_final_chunk", &self.is_final_chunk)
710 .finish()
711 }
712}
713
714#[derive(Debug, Clone, Default)]
721pub struct UnifiedBatch {
722 pub items: Vec<UnifiedBatchItem>,
723}
724
725impl UnifiedBatch {
726 pub fn new() -> Self {
727 Self::default()
728 }
729
730 pub fn total_q_tokens(&self) -> usize {
733 self.items.iter().map(|it| it.q_tokens.len()).sum()
734 }
735
736 pub fn num_sampled_items(&self) -> usize {
739 self.items.iter().filter(|it| it.is_final_chunk).count()
740 }
741}
742
743#[derive(Debug, Clone)]
750pub struct PlanRuntimeDecodeInput {
751 pub request_id: RequestId,
752 pub input_token: TokenId,
753 pub kv_cache: Arc<dyn KvCacheHandle>,
754 pub logits_policy: LogitsReturnPolicy,
755}
756
757impl PlanRuntimeDecodeInput {
758 pub fn new(
759 request_id: RequestId,
760 input_token: TokenId,
761 kv_cache: Arc<dyn KvCacheHandle>,
762 ) -> Self {
763 Self {
764 request_id,
765 input_token,
766 kv_cache,
767 logits_policy: LogitsReturnPolicy::FullLogits,
768 }
769 }
770
771 pub fn with_logits_policy(mut self, logits_policy: LogitsReturnPolicy) -> Self {
772 self.logits_policy = logits_policy;
773 self
774 }
775}
776
777#[derive(Debug, Clone)]
784pub struct PlanRuntimePrefillInput {
785 pub request_id: RequestId,
786 pub input_tokens: Arc<[TokenId]>,
787 pub maximum_sequence_tokens: usize,
788 pub chunk: PrefillChunk,
789}
790
791impl PlanRuntimePrefillInput {
792 pub fn new(
793 request_id: RequestId,
794 input_tokens: impl Into<Arc<[TokenId]>>,
795 maximum_sequence_tokens: usize,
796 chunk: PrefillChunk,
797 ) -> Result<Self> {
798 let input_tokens = input_tokens.into();
799 if input_tokens.is_empty() {
800 return Err(FerrumError::request_validation(
801 "plan-runtime prefill requires at least one input token",
802 ));
803 }
804 if chunk.total_prompt_tokens() != input_tokens.len() {
805 return Err(FerrumError::request_validation(format!(
806 "plan-runtime prefill chunk declares {} prompt tokens for input length {}",
807 chunk.total_prompt_tokens(),
808 input_tokens.len()
809 )));
810 }
811 if maximum_sequence_tokens < input_tokens.len() {
812 return Err(FerrumError::request_validation(format!(
813 "plan-runtime sequence ceiling {maximum_sequence_tokens} does not cover prompt length {}",
814 input_tokens.len()
815 )));
816 }
817 Ok(Self {
818 request_id,
819 input_tokens,
820 maximum_sequence_tokens,
821 chunk,
822 })
823 }
824}
825
826#[derive(Debug, Clone)]
828pub struct DecodeOutput {
829 pub logits: TensorRef,
831 pub kv_cache: Arc<dyn KvCacheHandle>,
833 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
835 pub hidden_state: Option<TensorRef>,
837 pub attention_weights: Option<Vec<TensorRef>>,
839}
840
841impl DecodeOutput {
842 pub fn new(logits: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
844 Self {
845 logits,
846 kv_cache,
847 recurrent_state: None,
848 hidden_state: None,
849 attention_weights: None,
850 }
851 }
852
853 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
855 self.recurrent_state = Some(recurrent_state);
856 self
857 }
858}
859
860#[derive(Debug, Clone)]
862pub struct PlanRuntimeDecodeOutput {
863 pub sampling_output: ExecutorSamplingOutput,
864 pub kv_cache: Arc<dyn KvCacheHandle>,
865}
866
867impl PlanRuntimeDecodeOutput {
868 pub fn new(sampling_output: ExecutorSamplingOutput, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
869 Self {
870 sampling_output,
871 kv_cache,
872 }
873 }
874}
875
876#[derive(Debug)]
878pub enum PlanRuntimePrefillProduct {
879 Intermediate,
881 FinalLogits(Vec<f32>),
883}
884
885#[derive(Debug)]
887pub struct PlanRuntimePrefillAuthority {
888 request_id: RequestId,
889 committed_tokens: usize,
890 kv_cache: Arc<dyn KvCacheHandle>,
891}
892
893impl PlanRuntimePrefillAuthority {
894 pub fn request_id(&self) -> &RequestId {
895 &self.request_id
896 }
897
898 pub const fn committed_tokens(&self) -> usize {
899 self.committed_tokens
900 }
901
902 pub fn kv_cache(&self) -> &Arc<dyn KvCacheHandle> {
903 &self.kv_cache
904 }
905
906 pub fn into_cache(self) -> Arc<dyn KvCacheHandle> {
907 self.kv_cache
908 }
909}
910
911#[derive(Debug)]
913pub struct PlanRuntimePrefillOutput {
914 authority: PlanRuntimePrefillAuthority,
915 product: PlanRuntimePrefillProduct,
916}
917
918impl PlanRuntimePrefillOutput {
919 pub fn intermediate(
920 request_id: RequestId,
921 committed_tokens: usize,
922 kv_cache: Arc<dyn KvCacheHandle>,
923 ) -> Self {
924 Self {
925 authority: PlanRuntimePrefillAuthority {
926 request_id,
927 committed_tokens,
928 kv_cache,
929 },
930 product: PlanRuntimePrefillProduct::Intermediate,
931 }
932 }
933
934 pub fn final_logits(
935 request_id: RequestId,
936 committed_tokens: usize,
937 logits: Vec<f32>,
938 kv_cache: Arc<dyn KvCacheHandle>,
939 ) -> Result<Self> {
940 if logits.is_empty() {
941 return Err(FerrumError::backend(
942 "plan-runtime final prefill returned empty logits",
943 ));
944 }
945 Ok(Self {
946 authority: PlanRuntimePrefillAuthority {
947 request_id,
948 committed_tokens,
949 kv_cache,
950 },
951 product: PlanRuntimePrefillProduct::FinalLogits(logits),
952 })
953 }
954
955 pub fn request_id(&self) -> &RequestId {
956 self.authority.request_id()
957 }
958
959 pub const fn committed_tokens(&self) -> usize {
960 self.authority.committed_tokens()
961 }
962
963 pub fn product(&self) -> &PlanRuntimePrefillProduct {
964 &self.product
965 }
966
967 pub fn kv_cache(&self) -> &Arc<dyn KvCacheHandle> {
968 self.authority.kv_cache()
969 }
970
971 pub fn validate_for_completion(
972 &self,
973 expected_request_id: &RequestId,
974 completed_chunk: PrefillChunk,
975 vocabulary_size: usize,
976 ) -> Result<()> {
977 if self.request_id() != expected_request_id {
978 return Err(FerrumError::backend(format!(
979 "plan runtime returned prefill output for request {}, expected {expected_request_id}",
980 self.request_id()
981 )));
982 }
983 if self.committed_tokens() != completed_chunk.end() {
984 return Err(FerrumError::backend(format!(
985 "plan runtime returned prefill extent {}, expected {}",
986 self.committed_tokens(),
987 completed_chunk.end()
988 )));
989 }
990 if self.kv_cache().num_tokens() != self.committed_tokens() {
991 return Err(FerrumError::backend(format!(
992 "plan runtime prefill cache `{}` reports {} tokens for committed extent {}",
993 self.kv_cache().cache_id(),
994 self.kv_cache().num_tokens(),
995 self.committed_tokens()
996 )));
997 }
998 match (&self.product, completed_chunk.is_final()) {
999 (PlanRuntimePrefillProduct::Intermediate, false) => Ok(()),
1000 (PlanRuntimePrefillProduct::FinalLogits(logits), true)
1001 if logits.len() == vocabulary_size =>
1002 {
1003 Ok(())
1004 }
1005 (PlanRuntimePrefillProduct::FinalLogits(logits), true) => {
1006 Err(FerrumError::backend(format!(
1007 "plan runtime returned {} final prefill logits for vocabulary {vocabulary_size}",
1008 logits.len()
1009 )))
1010 }
1011 (PlanRuntimePrefillProduct::Intermediate, true) => Err(FerrumError::backend(
1012 "plan runtime returned an intermediate product for a final prefill chunk",
1013 )),
1014 (PlanRuntimePrefillProduct::FinalLogits(_), false) => Err(FerrumError::backend(
1015 "plan runtime returned final logits for an intermediate prefill chunk",
1016 )),
1017 }
1018 }
1019
1020 pub fn into_parts(self) -> (PlanRuntimePrefillAuthority, PlanRuntimePrefillProduct) {
1021 (self.authority, self.product)
1022 }
1023}
1024
1025#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1033pub struct ExecutorSequenceCompletion {
1034 request_id: RequestId,
1035 cache_id: String,
1036 input_tokens: u64,
1037 output_tokens: u64,
1038}
1039
1040impl ExecutorSequenceCompletion {
1041 pub fn new(
1042 request_id: RequestId,
1043 cache_id: String,
1044 input_tokens: usize,
1045 output_tokens: usize,
1046 ) -> Result<Self> {
1047 if cache_id.is_empty() {
1048 return Err(FerrumError::request_validation(
1049 "executor sequence completion requires a cache identity",
1050 ));
1051 }
1052 let input_tokens = u64::try_from(input_tokens).map_err(|_| {
1053 FerrumError::request_validation("executor completion input token count exceeds u64")
1054 })?;
1055 let output_tokens = u64::try_from(output_tokens).map_err(|_| {
1056 FerrumError::request_validation("executor completion output token count exceeds u64")
1057 })?;
1058 Ok(Self {
1059 request_id,
1060 cache_id,
1061 input_tokens,
1062 output_tokens,
1063 })
1064 }
1065
1066 pub fn request_id(&self) -> &RequestId {
1067 &self.request_id
1068 }
1069
1070 pub fn cache_id(&self) -> &str {
1071 &self.cache_id
1072 }
1073
1074 pub const fn input_tokens(&self) -> u64 {
1075 self.input_tokens
1076 }
1077
1078 pub const fn output_tokens(&self) -> u64 {
1079 self.output_tokens
1080 }
1081}
1082
1083pub use ferrum_types::ExecutionResourceAuthority;
1084
1085#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1091pub struct ExecutorExecutionCapacityPreemption {
1092 request_id: RequestId,
1093 cache_id: String,
1094}
1095
1096impl ExecutorExecutionCapacityPreemption {
1097 pub fn new(request_id: RequestId, cache_id: String) -> Result<Self> {
1098 if cache_id.is_empty() {
1099 return Err(FerrumError::request_validation(
1100 "execution-capacity preemption requires a cache identity",
1101 ));
1102 }
1103 Ok(Self {
1104 request_id,
1105 cache_id,
1106 })
1107 }
1108
1109 pub fn request_id(&self) -> &RequestId {
1110 &self.request_id
1111 }
1112
1113 pub fn cache_id(&self) -> &str {
1114 &self.cache_id
1115 }
1116}
1117
1118#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
1119#[serde(rename_all = "snake_case")]
1120pub enum ExecutorExecutionCapacityPreemptionAuthority {
1121 RetainedPrefill,
1122 ActiveSequence,
1123}
1124
1125#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1129pub struct ExecutorExecutionCapacityPreemptionReceipt {
1130 request_id: RequestId,
1131 cache_id: String,
1132 authority: ExecutorExecutionCapacityPreemptionAuthority,
1133}
1134
1135impl ExecutorExecutionCapacityPreemptionReceipt {
1136 pub fn new(
1137 request_id: RequestId,
1138 cache_id: String,
1139 authority: ExecutorExecutionCapacityPreemptionAuthority,
1140 ) -> Self {
1141 Self {
1142 request_id,
1143 cache_id,
1144 authority,
1145 }
1146 }
1147
1148 pub fn request_id(&self) -> &RequestId {
1149 &self.request_id
1150 }
1151
1152 pub fn cache_id(&self) -> &str {
1153 &self.cache_id
1154 }
1155
1156 pub const fn authority(&self) -> ExecutorExecutionCapacityPreemptionAuthority {
1157 self.authority
1158 }
1159}
1160
1161#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1169pub struct PlanRuntimeResourceSnapshot {
1170 device_capacity_bytes: u64,
1171 usable_capacity_bytes: u64,
1172 process_claimed_bytes: u64,
1173 plan_claimed_bytes: u64,
1174 static_bytes: u64,
1175 dynamic_resident_bytes: u64,
1176 dynamic_free_bytes: u64,
1177 pending_growth_bytes: u64,
1178 quarantined_bytes: u64,
1179}
1180
1181impl PlanRuntimeResourceSnapshot {
1182 #[allow(clippy::too_many_arguments)]
1183 pub fn new(
1184 device_capacity_bytes: u64,
1185 usable_capacity_bytes: u64,
1186 process_claimed_bytes: u64,
1187 plan_claimed_bytes: u64,
1188 static_bytes: u64,
1189 dynamic_resident_bytes: u64,
1190 dynamic_free_bytes: u64,
1191 pending_growth_bytes: u64,
1192 quarantined_bytes: u64,
1193 ) -> Result<Self> {
1194 let snapshot = Self {
1195 device_capacity_bytes,
1196 usable_capacity_bytes,
1197 process_claimed_bytes,
1198 plan_claimed_bytes,
1199 static_bytes,
1200 dynamic_resident_bytes,
1201 dynamic_free_bytes,
1202 pending_growth_bytes,
1203 quarantined_bytes,
1204 };
1205 snapshot.validate()?;
1206 Ok(snapshot)
1207 }
1208
1209 pub fn validate(&self) -> Result<()> {
1212 if self.usable_capacity_bytes > self.device_capacity_bytes {
1213 return Err(ferrum_types::FerrumError::internal(format!(
1214 "plan runtime usable capacity {} exceeds device capacity {}",
1215 self.usable_capacity_bytes, self.device_capacity_bytes
1216 )));
1217 }
1218 if self.process_claimed_bytes > self.usable_capacity_bytes {
1219 return Err(ferrum_types::FerrumError::internal(format!(
1220 "plan runtime process claims {} exceed usable capacity {}",
1221 self.process_claimed_bytes, self.usable_capacity_bytes
1222 )));
1223 }
1224 if self.plan_claimed_bytes > self.process_claimed_bytes {
1225 return Err(ferrum_types::FerrumError::internal(format!(
1226 "plan runtime plan claims {} exceed process claims {}",
1227 self.plan_claimed_bytes, self.process_claimed_bytes
1228 )));
1229 }
1230 if self.dynamic_free_bytes > self.dynamic_resident_bytes {
1231 return Err(ferrum_types::FerrumError::internal(format!(
1232 "plan runtime dynamic free bytes {} exceed resident bytes {}",
1233 self.dynamic_free_bytes, self.dynamic_resident_bytes
1234 )));
1235 }
1236 let minimum_plan_claim = self
1237 .static_bytes
1238 .checked_add(self.dynamic_resident_bytes)
1239 .and_then(|bytes| bytes.checked_add(self.quarantined_bytes))
1240 .ok_or_else(|| {
1241 ferrum_types::FerrumError::internal(
1242 "plan runtime static, resident, and quarantined bytes overflow u64",
1243 )
1244 })?;
1245 if minimum_plan_claim > self.plan_claimed_bytes {
1246 return Err(ferrum_types::FerrumError::internal(format!(
1247 "plan runtime accounted plan bytes {minimum_plan_claim} exceed plan claims {}",
1248 self.plan_claimed_bytes
1249 )));
1250 }
1251 Ok(())
1252 }
1253
1254 pub const fn device_capacity_bytes(&self) -> u64 {
1255 self.device_capacity_bytes
1256 }
1257
1258 pub const fn usable_capacity_bytes(&self) -> u64 {
1259 self.usable_capacity_bytes
1260 }
1261
1262 pub const fn process_claimed_bytes(&self) -> u64 {
1263 self.process_claimed_bytes
1264 }
1265
1266 pub const fn plan_claimed_bytes(&self) -> u64 {
1267 self.plan_claimed_bytes
1268 }
1269
1270 pub const fn static_bytes(&self) -> u64 {
1271 self.static_bytes
1272 }
1273
1274 pub const fn dynamic_resident_bytes(&self) -> u64 {
1275 self.dynamic_resident_bytes
1276 }
1277
1278 pub const fn dynamic_free_bytes(&self) -> u64 {
1279 self.dynamic_free_bytes
1280 }
1281
1282 pub const fn dynamic_used_bytes(&self) -> u64 {
1283 self.dynamic_resident_bytes - self.dynamic_free_bytes
1284 }
1285
1286 pub const fn pending_growth_bytes(&self) -> u64 {
1287 self.pending_growth_bytes
1288 }
1289
1290 pub const fn quarantined_bytes(&self) -> u64 {
1291 self.quarantined_bytes
1292 }
1293
1294 pub fn available_bytes(&self) -> Result<u64> {
1298 self.usable_capacity_bytes
1299 .checked_sub(self.process_claimed_bytes)
1300 .and_then(|bytes| bytes.checked_add(self.dynamic_free_bytes))
1301 .ok_or_else(|| {
1302 ferrum_types::FerrumError::internal(
1303 "plan runtime available capacity calculation overflowed",
1304 )
1305 })
1306 }
1307
1308 pub fn used_bytes(&self) -> Result<u64> {
1309 self.available_bytes().and_then(|available| {
1310 self.usable_capacity_bytes
1311 .checked_sub(available)
1312 .ok_or_else(|| {
1313 ferrum_types::FerrumError::internal(
1314 "plan runtime available bytes exceed usable capacity",
1315 )
1316 })
1317 })
1318 }
1319}
1320
1321#[cfg(test)]
1322mod plan_runtime_resource_snapshot_tests {
1323 use super::PlanRuntimeResourceSnapshot;
1324
1325 #[test]
1326 fn separates_static_and_dynamic_usage() {
1327 let snapshot =
1328 PlanRuntimeResourceSnapshot::new(1_000, 900, 710, 710, 400, 300, 200, 20, 10).unwrap();
1329
1330 assert_eq!(snapshot.available_bytes().unwrap(), 390);
1331 assert_eq!(snapshot.used_bytes().unwrap(), 510);
1332 assert_eq!(snapshot.dynamic_resident_bytes(), 300);
1333 assert_eq!(snapshot.dynamic_used_bytes(), 100);
1334 assert_eq!(snapshot.dynamic_free_bytes(), 200);
1335 assert_eq!(snapshot.pending_growth_bytes(), 20);
1336 assert_eq!(snapshot.quarantined_bytes(), 10);
1337 }
1338
1339 #[test]
1340 fn rejects_incoherent_capacity_evidence() {
1341 assert!(PlanRuntimeResourceSnapshot::new(1_000, 1_001, 0, 0, 0, 0, 0, 0, 0).is_err());
1342 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 901, 0, 0, 0, 0, 0, 0).is_err());
1343 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 500, 501, 0, 0, 0, 0, 0).is_err());
1344 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 100, 100, 0, 100, 101, 0, 0).is_err());
1345 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 500, 500, 400, 100, 0, 0, 1).is_err());
1346 }
1347}
1348
1349#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1351pub enum ExecutorRequestOrigin {
1352 Product,
1353 Startup,
1354 Diagnostic,
1355}
1356
1357impl ExecutorRequestOrigin {
1358 pub const fn namespace(self) -> &'static str {
1359 match self {
1360 Self::Product => "product",
1361 Self::Startup => "startup",
1362 Self::Diagnostic => "diagnostic",
1363 }
1364 }
1365
1366 pub fn from_namespaced_request_identity(identity: &str) -> Option<Self> {
1367 let suffix = identity.strip_prefix("request.")?;
1368 let (namespace, request_id) = suffix.split_once('.')?;
1369 if request_id.is_empty() {
1370 return None;
1371 }
1372 match namespace {
1373 "product" => Some(Self::Product),
1374 "startup" => Some(Self::Startup),
1375 "diagnostic" => Some(Self::Diagnostic),
1376 _ => None,
1377 }
1378 }
1379}
1380
1381#[derive(Debug, Clone, Copy)]
1390pub struct ExecutorPrefillAdmission<'a> {
1391 pub request_id: &'a RequestId,
1392 pub input_tokens: &'a [TokenId],
1393 pub maximum_sequence_tokens: usize,
1394 pub product_prompt_tokens: usize,
1396 pub replayed_output_tokens: usize,
1398 pub request_origin: ExecutorRequestOrigin,
1399}
1400
1401impl<'a> ExecutorPrefillAdmission<'a> {
1402 pub const fn for_startup(
1403 request_id: &'a RequestId,
1404 input_tokens: &'a [TokenId],
1405 maximum_sequence_tokens: usize,
1406 ) -> Self {
1407 Self {
1408 request_id,
1409 input_tokens,
1410 maximum_sequence_tokens,
1411 product_prompt_tokens: input_tokens.len(),
1412 replayed_output_tokens: 0,
1413 request_origin: ExecutorRequestOrigin::Startup,
1414 }
1415 }
1416
1417 pub const fn for_diagnostic(
1418 request_id: &'a RequestId,
1419 input_tokens: &'a [TokenId],
1420 maximum_sequence_tokens: usize,
1421 ) -> Self {
1422 Self {
1423 request_id,
1424 input_tokens,
1425 maximum_sequence_tokens,
1426 product_prompt_tokens: input_tokens.len(),
1427 replayed_output_tokens: 0,
1428 request_origin: ExecutorRequestOrigin::Diagnostic,
1429 }
1430 }
1431
1432 pub fn for_product_request(
1435 request_id: &'a RequestId,
1436 input_tokens: &'a [TokenId],
1437 maximum_sequence_tokens: usize,
1438 product_prompt_tokens: usize,
1439 replayed_output_tokens: usize,
1440 ) -> Result<Self> {
1441 let admission = Self {
1442 request_id,
1443 input_tokens,
1444 maximum_sequence_tokens,
1445 product_prompt_tokens,
1446 replayed_output_tokens,
1447 request_origin: ExecutorRequestOrigin::Product,
1448 };
1449 admission.validate()?;
1450 Ok(admission)
1451 }
1452
1453 pub fn validate(&self) -> Result<()> {
1454 if self.input_tokens.is_empty() {
1455 return Err(FerrumError::request_validation(
1456 "executor prefill admission requires at least one execution-context token",
1457 ));
1458 }
1459 if self.product_prompt_tokens == 0 {
1460 return Err(FerrumError::request_validation(
1461 "executor prefill admission requires at least one product prompt token",
1462 ));
1463 }
1464 let execution_context_tokens = self
1465 .product_prompt_tokens
1466 .checked_add(self.replayed_output_tokens)
1467 .ok_or_else(|| {
1468 FerrumError::request_validation(
1469 "executor prefill product token accounting exceeds usize",
1470 )
1471 })?;
1472 if execution_context_tokens != self.input_tokens.len() {
1473 return Err(FerrumError::request_validation(format!(
1474 "executor prefill execution context has {} tokens but product accounting declares {} prompt + {} replayed output",
1475 self.input_tokens.len(),
1476 self.product_prompt_tokens,
1477 self.replayed_output_tokens
1478 )));
1479 }
1480 if self.maximum_sequence_tokens < execution_context_tokens {
1481 return Err(FerrumError::request_validation(format!(
1482 "executor prefill sequence ceiling {} does not cover execution context {execution_context_tokens}",
1483 self.maximum_sequence_tokens
1484 )));
1485 }
1486 Ok(())
1487 }
1488}
1489
1490#[cfg(test)]
1491mod executor_prefill_admission_tests {
1492 use super::{ExecutorPrefillAdmission, ExecutorRequestOrigin};
1493 use ferrum_types::{RequestId, TokenId};
1494
1495 #[test]
1496 fn product_accounting_distinguishes_replayed_output_from_prompt() {
1497 let request_id = RequestId::new();
1498 let tokens = [1, 2, 3, 4, 5]
1499 .into_iter()
1500 .map(TokenId::new)
1501 .collect::<Vec<_>>();
1502
1503 let admission =
1504 ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 8, 3, 2)
1505 .expect("recompute accounting must be accepted");
1506
1507 assert_eq!(admission.product_prompt_tokens, 3);
1508 assert_eq!(admission.replayed_output_tokens, 2);
1509 assert_eq!(admission.request_origin, ExecutorRequestOrigin::Product);
1510 assert_eq!(
1511 ExecutorPrefillAdmission::for_startup(&request_id, &tokens, 8).request_origin,
1512 ExecutorRequestOrigin::Startup
1513 );
1514 assert_eq!(
1515 ExecutorPrefillAdmission::for_diagnostic(&request_id, &tokens, 8).request_origin,
1516 ExecutorRequestOrigin::Diagnostic
1517 );
1518 assert_eq!(ExecutorRequestOrigin::Product.namespace(), "product");
1519 assert_eq!(ExecutorRequestOrigin::Startup.namespace(), "startup");
1520 assert_eq!(ExecutorRequestOrigin::Diagnostic.namespace(), "diagnostic");
1521 assert_eq!(
1522 ExecutorRequestOrigin::from_namespaced_request_identity("request.product.123"),
1523 Some(ExecutorRequestOrigin::Product)
1524 );
1525 assert_eq!(
1526 ExecutorRequestOrigin::from_namespaced_request_identity("request.startup.123"),
1527 Some(ExecutorRequestOrigin::Startup)
1528 );
1529 assert_eq!(
1530 ExecutorRequestOrigin::from_namespaced_request_identity("request.diagnostic.123"),
1531 Some(ExecutorRequestOrigin::Diagnostic)
1532 );
1533 assert_eq!(
1534 ExecutorRequestOrigin::from_namespaced_request_identity("request.product."),
1535 None
1536 );
1537 assert_eq!(
1538 ExecutorRequestOrigin::from_namespaced_request_identity("request/external"),
1539 None
1540 );
1541 }
1542
1543 #[test]
1544 fn product_accounting_rejects_context_drift_and_short_ceiling() {
1545 let request_id = RequestId::new();
1546 let tokens = [1, 2, 3].into_iter().map(TokenId::new).collect::<Vec<_>>();
1547
1548 assert!(
1549 ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 3, 2, 0).is_err()
1550 );
1551 assert!(
1552 ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 2, 2, 1).is_err()
1553 );
1554 }
1555}
1556
1557#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1560pub struct ExecutorPrefillAdmissionReceipt {
1561 pub request_id: RequestId,
1562}
1563
1564#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1566pub struct ExecutorAdmissionEpochs {
1567 pub coordinator_id: NonZeroU64,
1568 pub release_epoch: u64,
1569 pub capacity_epoch: u64,
1570}
1571
1572impl ExecutorAdmissionEpochs {
1573 pub const fn new(coordinator_id: NonZeroU64, release_epoch: u64, capacity_epoch: u64) -> Self {
1574 Self {
1575 coordinator_id,
1576 release_epoch,
1577 capacity_epoch,
1578 }
1579 }
1580
1581 pub fn from_capacity(epochs: crate::vnext::CapacityEpochs) -> Self {
1582 Self::new(
1583 NonZeroU64::new(epochs.coordinator_id().get())
1584 .expect("core-issued admission coordinator ids are non-zero"),
1585 epochs.release_epoch(),
1586 epochs.capacity_epoch(),
1587 )
1588 }
1589}
1590
1591type ExecutorCapacityWaitFuture =
1592 Pin<Box<dyn Future<Output = Result<ExecutorAdmissionEpochs>> + Send + 'static>>;
1593
1594#[must_use = "capacity wait registrations must be awaited or explicitly dropped"]
1601pub struct ExecutorCapacityWaitRegistration {
1602 future: ExecutorCapacityWaitFuture,
1603}
1604
1605impl ExecutorCapacityWaitRegistration {
1606 pub fn new<F>(future: F) -> Self
1607 where
1608 F: Future<Output = Result<ExecutorAdmissionEpochs>> + Send + 'static,
1609 {
1610 Self {
1611 future: Box::pin(future),
1612 }
1613 }
1614
1615 pub async fn wait_for_change(self) -> Result<ExecutorAdmissionEpochs> {
1616 self.future.await
1617 }
1618}
1619
1620#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1622#[serde(rename_all = "snake_case")]
1623pub enum ExecutorExecutionCapacityStage {
1624 SequenceExtension,
1625 StepAdmission,
1626 SubmissionWave,
1627}
1628
1629#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1632pub struct ExecutorExecutionMaintenanceMutation {
1633 pool_id: crate::vnext::DynamicBackingPoolId,
1634 domain_id: crate::vnext::CapacityDomainId,
1635 chunk: crate::vnext::BackingChunkIdentity,
1636 chunk_bytes: u64,
1637 published_capacity_bytes: u64,
1638 capacity_epoch: u64,
1639}
1640
1641impl ExecutorExecutionMaintenanceMutation {
1642 pub fn pool_id(&self) -> &crate::vnext::DynamicBackingPoolId {
1643 &self.pool_id
1644 }
1645
1646 pub const fn domain_id(&self) -> crate::vnext::CapacityDomainId {
1647 self.domain_id
1648 }
1649
1650 pub fn chunk(&self) -> &crate::vnext::BackingChunkIdentity {
1651 &self.chunk
1652 }
1653
1654 pub const fn chunk_bytes(&self) -> u64 {
1655 self.chunk_bytes
1656 }
1657
1658 pub const fn published_capacity_bytes(&self) -> u64 {
1659 self.published_capacity_bytes
1660 }
1661
1662 pub const fn capacity_epoch(&self) -> u64 {
1663 self.capacity_epoch
1664 }
1665}
1666
1667#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1675pub struct ExecutorExecutionMaintenanceProgress {
1676 attempts: u32,
1677 coordinator_id: NonZeroU64,
1678 mutations: Vec<ExecutorExecutionMaintenanceMutation>,
1679 latest_capacity_epoch: u64,
1680}
1681
1682impl ExecutorExecutionMaintenanceProgress {
1683 pub fn from_growth_receipts(
1684 attempts: u32,
1685 observed: ExecutorAdmissionEpochs,
1686 receipts: &[crate::vnext::DynamicPoolGrowthBatchReceipt],
1687 pools: &[crate::vnext::DynamicPoolStatus],
1688 ) -> Result<Self> {
1689 if attempts == 0 || receipts.is_empty() || receipts.len() > attempts as usize {
1690 return Err(FerrumError::internal(
1691 "execution maintenance retry requires bounded, non-empty growth receipts",
1692 ));
1693 }
1694
1695 let mut mutations = Vec::new();
1696 let mut previous_capacity_epoch = None;
1697 for receipt in receipts {
1698 if receipt.coordinator_id().get() != observed.coordinator_id.get() {
1699 return Err(FerrumError::internal(
1700 "execution maintenance receipt belongs to another capacity coordinator",
1701 ));
1702 }
1703 if receipt.growths().is_empty()
1704 || previous_capacity_epoch
1705 .is_some_and(|previous| receipt.capacity_epoch() <= previous)
1706 {
1707 return Err(FerrumError::internal(
1708 "execution maintenance receipts contain no new ordered capacity mutation",
1709 ));
1710 }
1711 previous_capacity_epoch = Some(receipt.capacity_epoch());
1712
1713 for growth in receipt.growths() {
1714 let pool = pools
1715 .iter()
1716 .find(|pool| pool.pool_id() == growth.pool_id())
1717 .ok_or_else(|| {
1718 FerrumError::internal(
1719 "execution maintenance receipt references an unknown dynamic pool",
1720 )
1721 })?;
1722 if growth.chunk().pool_id() != growth.pool_id()
1723 || growth.chunk_bytes() == 0
1724 || growth.published_capacity_bytes() == 0
1725 || growth.capacity_epoch() != receipt.capacity_epoch()
1726 {
1727 return Err(FerrumError::internal(
1728 "execution maintenance receipt contains an invalid pool mutation",
1729 ));
1730 }
1731 if mutations
1732 .iter()
1733 .any(|mutation: &ExecutorExecutionMaintenanceMutation| {
1734 mutation.pool_id() == growth.pool_id() && mutation.chunk() == growth.chunk()
1735 })
1736 {
1737 return Err(FerrumError::internal(
1738 "execution maintenance receipts repeat one physical pool mutation",
1739 ));
1740 }
1741 mutations.push(ExecutorExecutionMaintenanceMutation {
1742 pool_id: growth.pool_id().clone(),
1743 domain_id: pool.domain_id(),
1744 chunk: growth.chunk().clone(),
1745 chunk_bytes: growth.chunk_bytes(),
1746 published_capacity_bytes: growth.published_capacity_bytes(),
1747 capacity_epoch: growth.capacity_epoch(),
1748 });
1749 }
1750 }
1751
1752 let latest_capacity_epoch = previous_capacity_epoch.expect("receipts are non-empty");
1753 if latest_capacity_epoch > observed.capacity_epoch {
1754 return Err(FerrumError::internal(
1755 "execution maintenance receipt is newer than the exported capacity observation",
1756 ));
1757 }
1758 mutations.sort_by(|left, right| {
1759 (
1760 left.capacity_epoch,
1761 left.pool_id.as_str(),
1762 left.chunk.ordinal(),
1763 left.chunk.generation(),
1764 )
1765 .cmp(&(
1766 right.capacity_epoch,
1767 right.pool_id.as_str(),
1768 right.chunk.ordinal(),
1769 right.chunk.generation(),
1770 ))
1771 });
1772 Ok(Self {
1773 attempts,
1774 coordinator_id: observed.coordinator_id,
1775 mutations,
1776 latest_capacity_epoch,
1777 })
1778 }
1779
1780 pub const fn attempts(&self) -> u32 {
1781 self.attempts
1782 }
1783
1784 pub const fn coordinator_id(&self) -> NonZeroU64 {
1785 self.coordinator_id
1786 }
1787
1788 pub fn mutations(&self) -> &[ExecutorExecutionMaintenanceMutation] {
1789 &self.mutations
1790 }
1791
1792 pub const fn latest_capacity_epoch(&self) -> u64 {
1793 self.latest_capacity_epoch
1794 }
1795}
1796
1797#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1803pub struct ExecutorExecutionMaintenanceRetry {
1804 affected_request_ids: Vec<RequestId>,
1805 progress: ExecutorExecutionMaintenanceProgress,
1806}
1807
1808impl ExecutorExecutionMaintenanceRetry {
1809 fn new(
1810 affected_request_ids: Vec<RequestId>,
1811 progress: ExecutorExecutionMaintenanceProgress,
1812 ) -> Result<Self> {
1813 let unique = affected_request_ids.iter().collect::<HashSet<_>>();
1814 if affected_request_ids.is_empty() || unique.len() != affected_request_ids.len() {
1815 return Err(FerrumError::internal(
1816 "execution maintenance retry requires unique affected requests",
1817 ));
1818 }
1819 if progress.mutations().is_empty() {
1820 return Err(FerrumError::internal(
1821 "execution maintenance retry requires physical mutations",
1822 ));
1823 }
1824 Ok(Self {
1825 affected_request_ids,
1826 progress,
1827 })
1828 }
1829
1830 pub fn affected_request_ids(&self) -> &[RequestId] {
1831 &self.affected_request_ids
1832 }
1833
1834 pub const fn progress(&self) -> &ExecutorExecutionMaintenanceProgress {
1835 &self.progress
1836 }
1837}
1838
1839#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1846#[serde(rename_all = "snake_case")]
1847pub enum ExecutorExecutionCapacityEvidenceOwner {
1848 Logical,
1849 Backing,
1850}
1851
1852#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1853#[serde(tag = "kind", rename_all = "snake_case")]
1854enum ExecutorExecutionCapacityEvidenceKind {
1855 Logical {
1856 shortfalls: Vec<crate::vnext::CapacityShortfall>,
1857 #[serde(skip_serializing_if = "Option::is_none")]
1858 pressure: Option<crate::vnext::DynamicBackingPressure>,
1859 },
1860 BackingDeferred {
1861 blockers: Vec<crate::vnext::DynamicBackingBlocker>,
1862 #[serde(skip_serializing_if = "Option::is_none")]
1863 pressure: Option<crate::vnext::DynamicBackingPressure>,
1864 },
1865 BackingPressure {
1866 pressure: crate::vnext::DynamicBackingPressure,
1867 },
1868}
1869
1870#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1875pub struct ExecutorExecutionCapacityEvidence {
1876 owner: ExecutorExecutionCapacityEvidenceOwner,
1877 #[serde(flatten)]
1878 kind: ExecutorExecutionCapacityEvidenceKind,
1879 #[serde(skip_serializing_if = "Option::is_none")]
1880 maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
1881}
1882
1883impl ExecutorExecutionCapacityEvidence {
1884 fn logical(shortfalls: Vec<crate::vnext::CapacityShortfall>) -> Result<Self> {
1885 Self::logical_with_pressure(shortfalls, None)
1886 }
1887
1888 fn logical_with_pressure(
1889 shortfalls: Vec<crate::vnext::CapacityShortfall>,
1890 pressure: Option<crate::vnext::DynamicBackingPressure>,
1891 ) -> Result<Self> {
1892 if shortfalls.is_empty() {
1893 return Err(FerrumError::internal(
1894 "logical execution deferral requires at least one shortfall",
1895 ));
1896 }
1897 Ok(Self {
1898 owner: ExecutorExecutionCapacityEvidenceOwner::Logical,
1899 kind: ExecutorExecutionCapacityEvidenceKind::Logical {
1900 shortfalls,
1901 pressure,
1902 },
1903 maintenance_boundary: None,
1904 })
1905 }
1906
1907 fn backing_deferred(blockers: Vec<crate::vnext::DynamicBackingBlocker>) -> Result<Self> {
1908 Self::backing_deferred_with_pressure(blockers, None)
1909 }
1910
1911 fn backing_deferred_with_pressure(
1912 blockers: Vec<crate::vnext::DynamicBackingBlocker>,
1913 pressure: Option<crate::vnext::DynamicBackingPressure>,
1914 ) -> Result<Self> {
1915 if blockers.is_empty() {
1916 return Err(FerrumError::internal(
1917 "physical execution deferral requires at least one backing blocker",
1918 ));
1919 }
1920 Ok(Self {
1921 owner: ExecutorExecutionCapacityEvidenceOwner::Backing,
1922 kind: ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, pressure },
1923 maintenance_boundary: None,
1924 })
1925 }
1926
1927 fn direct_backing_pressure(pressure: crate::vnext::DynamicBackingPressure) -> Self {
1928 Self {
1929 owner: ExecutorExecutionCapacityEvidenceOwner::Backing,
1930 kind: ExecutorExecutionCapacityEvidenceKind::BackingPressure { pressure },
1931 maintenance_boundary: None,
1932 }
1933 }
1934
1935 pub const fn owner(&self) -> ExecutorExecutionCapacityEvidenceOwner {
1936 self.owner
1937 }
1938
1939 pub fn shortfalls(&self) -> &[crate::vnext::CapacityShortfall] {
1940 match &self.kind {
1941 ExecutorExecutionCapacityEvidenceKind::Logical { shortfalls, .. } => shortfalls,
1942 ExecutorExecutionCapacityEvidenceKind::BackingDeferred { .. }
1943 | ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => &[],
1944 }
1945 }
1946
1947 pub fn backing_blockers(&self) -> &[crate::vnext::DynamicBackingBlocker] {
1948 match &self.kind {
1949 ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, .. } => blockers,
1950 ExecutorExecutionCapacityEvidenceKind::Logical { .. }
1951 | ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => &[],
1952 }
1953 }
1954
1955 pub const fn backing_pressure(&self) -> Option<&crate::vnext::DynamicBackingPressure> {
1956 match &self.kind {
1957 ExecutorExecutionCapacityEvidenceKind::Logical { pressure, .. }
1958 | ExecutorExecutionCapacityEvidenceKind::BackingDeferred { pressure, .. } => {
1959 pressure.as_ref()
1960 }
1961 ExecutorExecutionCapacityEvidenceKind::BackingPressure { pressure } => Some(pressure),
1962 }
1963 }
1964
1965 pub const fn maintenance_boundary(
1966 &self,
1967 ) -> Option<&crate::vnext::DynamicPoolMaintenanceBoundaryReceipt> {
1968 self.maintenance_boundary.as_ref()
1969 }
1970
1971 fn with_maintenance_boundary(
1972 mut self,
1973 boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
1974 ) -> Result<Self> {
1975 match (self.backing_pressure(), boundary.as_ref()) {
1976 (
1977 Some(crate::vnext::DynamicBackingPressure::DeviceCapacity(pressure)),
1978 Some(boundary),
1979 ) if pressure == boundary.pressure() && !boundary.reclaim_sufficient() => {}
1980 (Some(crate::vnext::DynamicBackingPressure::PoolResident(_)), None) | (None, None) => {}
1981 (Some(crate::vnext::DynamicBackingPressure::DeviceCapacity(_)), None) => {
1982 return Err(FerrumError::internal(
1983 "device-capacity execution maintenance lost its boundary receipt",
1984 ));
1985 }
1986 _ => {
1987 return Err(FerrumError::internal(
1988 "execution maintenance boundary differs from its blocked pressure",
1989 ));
1990 }
1991 }
1992 self.maintenance_boundary = boundary;
1993 Ok(self)
1994 }
1995
1996 fn has_relevant_mutation(&self, mutation: &ExecutorExecutionMaintenanceMutation) -> bool {
1997 let logical_matches = |shortfalls: &[crate::vnext::CapacityShortfall]| {
1998 shortfalls.iter().any(|shortfall| {
1999 shortfall.kind() == crate::vnext::CapacityShortfallKind::BackingGrowthRequired
2000 && shortfall.domain() == Some(mutation.domain_id())
2001 })
2002 };
2003 let backing_matches = |blockers: &[crate::vnext::DynamicBackingBlocker]| {
2004 blockers.iter().any(|blocker| {
2005 blocker.pool_id() == mutation.pool_id()
2006 && blocker.domain_id() == mutation.domain_id()
2007 })
2008 };
2009 match &self.kind {
2010 ExecutorExecutionCapacityEvidenceKind::Logical { shortfalls, .. } => {
2011 logical_matches(shortfalls)
2012 }
2013 ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, .. } => {
2014 backing_matches(blockers)
2015 }
2016 ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => false,
2017 }
2018 }
2019}
2020
2021#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2022pub struct ExecutorExecutionCapacityDeferral {
2023 observed: ExecutorAdmissionEpochs,
2024 wait_condition: crate::vnext::CapacityWaitCondition,
2025 stage: ExecutorExecutionCapacityStage,
2026 evidence: ExecutorExecutionCapacityEvidence,
2027 #[serde(skip_serializing_if = "Option::is_none")]
2028 maintenance_retry: Option<ExecutorExecutionMaintenanceRetry>,
2029}
2030
2031impl ExecutorExecutionCapacityDeferral {
2032 fn with_evidence(
2033 observed: ExecutorAdmissionEpochs,
2034 wait_condition: crate::vnext::CapacityWaitCondition,
2035 stage: ExecutorExecutionCapacityStage,
2036 evidence: ExecutorExecutionCapacityEvidence,
2037 ) -> Result<Self> {
2038 if wait_condition.coordinator_id().get() != observed.coordinator_id.get() {
2039 return Err(ferrum_types::FerrumError::request_validation(
2040 "executor execution deferral belongs to a different capacity coordinator",
2041 ));
2042 }
2043 Ok(Self {
2044 observed,
2045 wait_condition,
2046 stage,
2047 evidence,
2048 maintenance_retry: None,
2049 })
2050 }
2051
2052 pub fn from_backing_pressure(
2055 observed: ExecutorAdmissionEpochs,
2056 wait_condition: crate::vnext::CapacityWaitCondition,
2057 pressure: crate::vnext::DynamicBackingPressure,
2058 stage: ExecutorExecutionCapacityStage,
2059 ) -> Result<Self> {
2060 Self::with_evidence(
2061 observed,
2062 wait_condition,
2063 stage,
2064 ExecutorExecutionCapacityEvidence::direct_backing_pressure(pressure),
2065 )
2066 }
2067
2068 pub fn from_admission(
2069 deferred: &crate::vnext::AdmissionDeferred,
2070 stage: ExecutorExecutionCapacityStage,
2071 ) -> Result<Self> {
2072 if deferred.action() != crate::vnext::DeferredAction::WaitForRelease {
2073 return Err(ferrum_types::FerrumError::internal(
2074 "execution capacity deferral must be reduced to WaitForRelease before export",
2075 ));
2076 }
2077 let evidence = ExecutorExecutionCapacityEvidence::logical(deferred.blockers().to_vec())?;
2078 Self::with_evidence(
2079 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2080 deferred.wait_condition().clone(),
2081 stage,
2082 evidence,
2083 )
2084 }
2085
2086 pub fn from_pending_maintenance(
2093 deferred: &crate::vnext::AdmissionDeferred,
2094 stage: ExecutorExecutionCapacityStage,
2095 ) -> Result<Self> {
2096 if deferred.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2097 return Err(ferrum_types::FerrumError::internal(
2098 "pending execution maintenance must await backing growth",
2099 ));
2100 }
2101 let evidence = ExecutorExecutionCapacityEvidence::logical(deferred.blockers().to_vec())?;
2102 Self::with_evidence(
2103 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2104 deferred.wait_condition().clone(),
2105 stage,
2106 evidence,
2107 )
2108 }
2109
2110 pub fn from_backing(
2113 deferred: &crate::vnext::DynamicBackingDeferred,
2114 stage: ExecutorExecutionCapacityStage,
2115 ) -> Result<Self> {
2116 let evidence =
2117 ExecutorExecutionCapacityEvidence::backing_deferred(deferred.blockers().to_vec())?;
2118 Self::with_evidence(
2119 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2120 deferred.wait_condition().clone(),
2121 stage,
2122 evidence,
2123 )
2124 }
2125
2126 pub fn with_relevant_maintenance_retry(
2130 mut self,
2131 attempts: u32,
2132 receipts: &[crate::vnext::DynamicPoolGrowthBatchReceipt],
2133 pools: &[crate::vnext::DynamicPoolStatus],
2134 affected_request_ids: Vec<RequestId>,
2135 ) -> Result<Self> {
2136 if receipts.is_empty() {
2137 return Ok(self);
2138 }
2139 let progress = ExecutorExecutionMaintenanceProgress::from_growth_receipts(
2140 attempts,
2141 self.observed,
2142 receipts,
2143 pools,
2144 )?;
2145 if progress.coordinator_id() != self.observed.coordinator_id
2146 || progress.latest_capacity_epoch() > self.observed.capacity_epoch
2147 || progress.mutations().is_empty()
2148 {
2149 return Err(FerrumError::internal(
2150 "execution maintenance progress does not match the exported deferral",
2151 ));
2152 }
2153 let relevant_mutation = progress
2154 .mutations()
2155 .iter()
2156 .any(|mutation| self.evidence.has_relevant_mutation(mutation));
2157 if !relevant_mutation {
2158 return Ok(self);
2159 }
2160 let retry = ExecutorExecutionMaintenanceRetry::new(affected_request_ids, progress)?;
2161 if self.stage == ExecutorExecutionCapacityStage::SequenceExtension
2162 && retry.affected_request_ids().len() != 1
2163 {
2164 return Err(FerrumError::internal(
2165 "sequence-extension maintenance retry must affect exactly one request",
2166 ));
2167 }
2168 self.maintenance_retry = Some(retry);
2169 Ok(self)
2170 }
2171
2172 pub fn from_admission_maintenance(
2173 source: &crate::vnext::AdmissionDeferred,
2174 observed: ExecutorAdmissionEpochs,
2175 wait_condition: crate::vnext::CapacityWaitCondition,
2176 pressure: crate::vnext::DynamicBackingPressure,
2177 maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
2178 stage: ExecutorExecutionCapacityStage,
2179 ) -> Result<Self> {
2180 if source.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2181 return Err(ferrum_types::FerrumError::internal(
2182 "execution maintenance source must await backing growth",
2183 ));
2184 }
2185 let evidence = ExecutorExecutionCapacityEvidence::logical_with_pressure(
2186 source.blockers().to_vec(),
2187 Some(pressure),
2188 )?
2189 .with_maintenance_boundary(maintenance_boundary)?;
2190 if evidence.maintenance_boundary().is_some_and(|boundary| {
2191 boundary.coordinator_id().get() != observed.coordinator_id.get()
2192 }) {
2193 return Err(FerrumError::internal(
2194 "execution maintenance boundary belongs to another coordinator",
2195 ));
2196 }
2197 Self::with_evidence(observed, wait_condition, stage, evidence)
2198 }
2199
2200 pub fn from_backing_maintenance(
2201 source: &crate::vnext::DynamicBackingDeferred,
2202 observed: ExecutorAdmissionEpochs,
2203 wait_condition: crate::vnext::CapacityWaitCondition,
2204 pressure: crate::vnext::DynamicBackingPressure,
2205 maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
2206 stage: ExecutorExecutionCapacityStage,
2207 ) -> Result<Self> {
2208 let evidence = ExecutorExecutionCapacityEvidence::backing_deferred_with_pressure(
2209 source.blockers().to_vec(),
2210 Some(pressure),
2211 )?
2212 .with_maintenance_boundary(maintenance_boundary)?;
2213 if evidence.maintenance_boundary().is_some_and(|boundary| {
2214 boundary.coordinator_id().get() != observed.coordinator_id.get()
2215 }) {
2216 return Err(FerrumError::internal(
2217 "execution maintenance boundary belongs to another coordinator",
2218 ));
2219 }
2220 Self::with_evidence(observed, wait_condition, stage, evidence)
2221 }
2222
2223 pub const fn observed(&self) -> ExecutorAdmissionEpochs {
2224 self.observed
2225 }
2226
2227 pub fn wait_condition(&self) -> &crate::vnext::CapacityWaitCondition {
2228 &self.wait_condition
2229 }
2230
2231 pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2232 self.stage
2233 }
2234
2235 pub const fn evidence(&self) -> &ExecutorExecutionCapacityEvidence {
2236 &self.evidence
2237 }
2238
2239 pub fn shortfalls(&self) -> &[crate::vnext::CapacityShortfall] {
2240 self.evidence.shortfalls()
2241 }
2242
2243 pub fn backing_blockers(&self) -> &[crate::vnext::DynamicBackingBlocker] {
2244 self.evidence.backing_blockers()
2245 }
2246
2247 pub const fn backing_pressure(&self) -> Option<&crate::vnext::DynamicBackingPressure> {
2248 self.evidence.backing_pressure()
2249 }
2250
2251 pub const fn maintenance_boundary(
2252 &self,
2253 ) -> Option<&crate::vnext::DynamicPoolMaintenanceBoundaryReceipt> {
2254 self.evidence.maintenance_boundary()
2255 }
2256
2257 pub fn maintenance_retry(&self) -> Option<&ExecutorExecutionMaintenanceRetry> {
2258 self.maintenance_retry.as_ref()
2259 }
2260
2261 pub fn validated_maintenance_retry_scope(
2264 &self,
2265 current_request_ids: &[RequestId],
2266 ) -> Result<Option<&ExecutorExecutionMaintenanceRetry>> {
2267 let Some(retry) = self.maintenance_retry.as_ref() else {
2268 return Ok(None);
2269 };
2270 let current = current_request_ids.iter().collect::<HashSet<_>>();
2271 if current_request_ids.is_empty() || current.len() != current_request_ids.len() {
2272 return Err(FerrumError::internal(
2273 "execution maintenance retry received an invalid current request cohort",
2274 ));
2275 }
2276 let affected = retry.affected_request_ids().iter().collect::<HashSet<_>>();
2277 if !affected.is_subset(¤t) {
2278 return Err(FerrumError::internal(
2279 "execution maintenance retry affects a request outside the current cohort",
2280 ));
2281 }
2282 match self.stage {
2283 ExecutorExecutionCapacityStage::SequenceExtension => {
2284 if affected.len() != 1 {
2285 return Err(FerrumError::internal(
2286 "sequence-extension maintenance retry must affect one current request",
2287 ));
2288 }
2289 }
2290 ExecutorExecutionCapacityStage::StepAdmission
2291 | ExecutorExecutionCapacityStage::SubmissionWave => {
2292 if affected != current {
2293 return Err(FerrumError::internal(
2294 "cohort maintenance retry must cover the complete current cohort",
2295 ));
2296 }
2297 }
2298 }
2299 Ok(Some(retry))
2300 }
2301
2302 pub fn narrower_prefill_tokens(&self, attempted_tokens: usize) -> Option<usize> {
2310 if attempted_tokens <= 1 {
2311 return None;
2312 }
2313 let maximum_next = attempted_tokens
2314 .saturating_sub(attempted_tokens.div_ceil(4))
2315 .max(1);
2316 let proportional = self
2317 .shortfalls()
2318 .iter()
2319 .filter_map(|shortfall| {
2320 let requested = shortfall.requested().get();
2321 let available = shortfall.available().get();
2322 (requested > available).then(|| {
2323 let scaled = (attempted_tokens as u128).saturating_mul(available as u128)
2324 / requested as u128;
2325 usize::try_from(scaled)
2326 .unwrap_or(usize::MAX)
2327 .clamp(1, attempted_tokens - 1)
2328 })
2329 })
2330 .min();
2331 Some(
2332 proportional
2333 .unwrap_or_else(|| attempted_tokens.div_ceil(2))
2334 .min(maximum_next)
2335 .max(1),
2336 )
2337 }
2338}
2339
2340#[derive(Debug, Clone, Serialize)]
2348pub struct ExecutorRequestStateDeferral {
2349 stage: ExecutorExecutionCapacityStage,
2350 request_ids: Vec<RequestId>,
2351 hazard: crate::vnext::RequestStateHazardDeferral,
2352}
2353
2354impl ExecutorRequestStateDeferral {
2355 pub fn new(
2356 stage: ExecutorExecutionCapacityStage,
2357 request_ids: Vec<RequestId>,
2358 hazard: crate::vnext::RequestStateHazardDeferral,
2359 ) -> Result<Self> {
2360 let unique = request_ids.iter().collect::<HashSet<_>>();
2361 if request_ids.is_empty() || unique.len() != request_ids.len() {
2362 return Err(FerrumError::internal(
2363 "request-state execution deferral requires a non-empty unique product cohort",
2364 ));
2365 }
2366 if hazard.blockers().is_empty() {
2367 return Err(FerrumError::internal(
2368 "request-state execution deferral requires exact blockers",
2369 ));
2370 }
2371 Ok(Self {
2372 stage,
2373 request_ids,
2374 hazard,
2375 })
2376 }
2377
2378 pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2379 self.stage
2380 }
2381
2382 pub fn request_ids(&self) -> &[RequestId] {
2383 &self.request_ids
2384 }
2385
2386 pub const fn hazard(&self) -> &crate::vnext::RequestStateHazardDeferral {
2387 &self.hazard
2388 }
2389
2390 pub fn register_waiter(&self) -> Result<crate::vnext::RequestStateHazardWaitRegistration> {
2391 self.hazard
2392 .register_waiter()
2393 .map_err(|error| FerrumError::backend(error.to_string()))
2394 }
2395}
2396
2397#[derive(Debug, Clone, Serialize)]
2402#[serde(tag = "reason", content = "evidence", rename_all = "snake_case")]
2403pub enum ExecutorExecutionDeferral {
2404 Capacity(ExecutorExecutionCapacityDeferral),
2405 RequestState(ExecutorRequestStateDeferral),
2406}
2407
2408impl ExecutorExecutionDeferral {
2409 pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2410 match self {
2411 Self::Capacity(deferral) => deferral.stage(),
2412 Self::RequestState(deferral) => deferral.stage(),
2413 }
2414 }
2415
2416 pub const fn as_capacity(&self) -> Option<&ExecutorExecutionCapacityDeferral> {
2417 match self {
2418 Self::Capacity(deferral) => Some(deferral),
2419 Self::RequestState(_) => None,
2420 }
2421 }
2422
2423 pub const fn as_request_state(&self) -> Option<&ExecutorRequestStateDeferral> {
2424 match self {
2425 Self::Capacity(_) => None,
2426 Self::RequestState(deferral) => Some(deferral),
2427 }
2428 }
2429}
2430
2431impl From<ExecutorExecutionCapacityDeferral> for ExecutorExecutionDeferral {
2432 fn from(deferral: ExecutorExecutionCapacityDeferral) -> Self {
2433 Self::Capacity(deferral)
2434 }
2435}
2436
2437impl From<ExecutorRequestStateDeferral> for ExecutorExecutionDeferral {
2438 fn from(deferral: ExecutorRequestStateDeferral) -> Self {
2439 Self::RequestState(deferral)
2440 }
2441}
2442
2443#[cfg(test)]
2444mod execution_capacity_deferral_tests {
2445 use super::{
2446 ExecutorAdmissionEpochs, ExecutorExecutionCapacityDeferral,
2447 ExecutorExecutionCapacityEvidenceOwner, ExecutorExecutionCapacityStage,
2448 ExecutorExecutionMaintenanceProgress, ExecutorExecutionMaintenanceRetry,
2449 };
2450 use crate::vnext::{
2451 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
2452 DeviceCapacityPressure, DeviceCapacityPressureScope, DynamicBackingPressure,
2453 };
2454 use ferrum_types::RequestId;
2455 use std::num::NonZeroU64;
2456
2457 fn test_progress() -> ExecutorExecutionMaintenanceProgress {
2458 ExecutorExecutionMaintenanceProgress {
2459 attempts: 1,
2460 coordinator_id: NonZeroU64::new(19).unwrap(),
2461 mutations: Vec::new(),
2462 latest_capacity_epoch: 5,
2463 }
2464 }
2465
2466 fn test_deferral(stage: ExecutorExecutionCapacityStage) -> ExecutorExecutionCapacityDeferral {
2467 let observed =
2468 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 7)
2469 .unwrap();
2470 let condition = CapacityWaitCondition::from_observation(19, vec![observed]).unwrap();
2471 ExecutorExecutionCapacityDeferral::from_backing_pressure(
2472 ExecutorAdmissionEpochs::new(NonZeroU64::new(19).unwrap(), 3, 5),
2473 condition,
2474 test_pressure(),
2475 stage,
2476 )
2477 .unwrap()
2478 }
2479
2480 fn test_pressure() -> DynamicBackingPressure {
2481 DeviceCapacityPressure::new(
2482 DeviceCapacityPressureScope::PlanBudget,
2483 "device.execution-capacity-test".to_owned(),
2484 1,
2485 1,
2486 1,
2487 1,
2488 1,
2489 )
2490 .unwrap()
2491 .into()
2492 }
2493
2494 #[test]
2495 fn prefill_narrowing_is_strict_bounded_and_stops_at_one_token() {
2496 let observed =
2497 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 7)
2498 .unwrap();
2499 let condition = CapacityWaitCondition::from_observation(19, vec![observed]).unwrap();
2500 let deferred = ExecutorExecutionCapacityDeferral::from_backing_pressure(
2501 ExecutorAdmissionEpochs::new(NonZeroU64::new(19).unwrap(), 3, 5),
2502 condition,
2503 test_pressure(),
2504 ExecutorExecutionCapacityStage::StepAdmission,
2505 )
2506 .unwrap();
2507
2508 assert_eq!(deferred.narrower_prefill_tokens(342), Some(171));
2509 assert_eq!(deferred.narrower_prefill_tokens(2), Some(1));
2510 assert_eq!(deferred.narrower_prefill_tokens(1), None);
2511 }
2512
2513 #[test]
2514 fn backing_pressure_serializes_one_typed_evidence_owner() {
2515 let deferred = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension);
2516 let serialized = serde_json::to_value(&deferred).unwrap();
2517
2518 assert_eq!(
2519 deferred.evidence().owner(),
2520 ExecutorExecutionCapacityEvidenceOwner::Backing
2521 );
2522 assert!(deferred.shortfalls().is_empty());
2523 assert!(deferred.backing_blockers().is_empty());
2524 assert!(deferred.backing_pressure().is_some());
2525 assert_eq!(serialized["evidence"]["owner"], "backing");
2526 assert_eq!(serialized["evidence"]["kind"], "backing_pressure");
2527 assert!(serialized["evidence"]["pressure"].is_object());
2528 }
2529
2530 #[test]
2531 fn empty_maintenance_receipts_remain_an_ordinary_typed_deferral() {
2532 let request_id = RequestId::new();
2533 let deferred = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension)
2534 .with_relevant_maintenance_retry(2, &[], &[], vec![request_id])
2535 .unwrap();
2536
2537 assert!(deferred.maintenance_retry().is_none());
2538 }
2539
2540 #[test]
2541 fn maintenance_retry_rejects_empty_duplicate_or_unproven_scope() {
2542 let request_id = RequestId::new();
2543 assert!(ExecutorExecutionMaintenanceRetry::new(Vec::new(), test_progress()).is_err());
2544 assert!(ExecutorExecutionMaintenanceRetry::new(
2545 vec![request_id.clone(), request_id.clone()],
2546 test_progress(),
2547 )
2548 .is_err());
2549 assert!(ExecutorExecutionMaintenanceRetry::new(vec![request_id], test_progress()).is_err());
2550 }
2551
2552 #[test]
2553 fn maintenance_retry_scope_is_fail_closed_for_sequence_and_cohort_stages() {
2554 let first = RequestId::new();
2555 let second = RequestId::new();
2556 let retry = |affected_request_ids| ExecutorExecutionMaintenanceRetry {
2557 affected_request_ids,
2558 progress: test_progress(),
2559 };
2560
2561 let mut sequence = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension);
2562 sequence.maintenance_retry = Some(retry(vec![second.clone()]));
2563 assert_eq!(
2564 sequence
2565 .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2566 .unwrap()
2567 .unwrap()
2568 .affected_request_ids(),
2569 [second.clone()]
2570 );
2571 sequence.maintenance_retry = Some(retry(vec![first.clone(), second.clone()]));
2572 assert!(sequence
2573 .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2574 .is_err());
2575
2576 let mut cohort = test_deferral(ExecutorExecutionCapacityStage::SubmissionWave);
2577 cohort.maintenance_retry = Some(retry(vec![second.clone()]));
2578 assert!(cohort
2579 .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2580 .is_err());
2581 cohort.maintenance_retry = Some(retry(vec![first.clone(), second.clone()]));
2582 assert!(cohort
2583 .validated_maintenance_retry_scope(&[first, second])
2584 .unwrap()
2585 .is_some());
2586 }
2587}
2588
2589pub enum ExecutorBatchDecodeOutcome {
2595 Completed(Vec<DecodeOutput>),
2596 Deferred(ExecutorExecutionDeferral),
2597}
2598
2599pub enum PlanRuntimeBatchDecodeOutcome {
2601 Completed(Vec<PlanRuntimeDecodeOutput>),
2602 Deferred(ExecutorExecutionDeferral),
2603}
2604
2605pub struct PlanRuntimePrefillCompletion {
2607 output: PlanRuntimePrefillOutput,
2608 planned_chunk: PrefillChunk,
2609 completed_chunk: PrefillChunk,
2610 capacity_probe_count: u32,
2611}
2612
2613impl PlanRuntimePrefillCompletion {
2614 pub fn new(
2615 output: PlanRuntimePrefillOutput,
2616 planned_chunk: PrefillChunk,
2617 completed_chunk: PrefillChunk,
2618 capacity_probe_count: u32,
2619 ) -> Result<Self> {
2620 validate_prefill_completion_shape(planned_chunk, completed_chunk, capacity_probe_count)?;
2621 Ok(Self {
2622 output,
2623 planned_chunk,
2624 completed_chunk,
2625 capacity_probe_count,
2626 })
2627 }
2628
2629 pub fn exact(output: PlanRuntimePrefillOutput, chunk: PrefillChunk) -> Self {
2630 Self {
2631 output,
2632 planned_chunk: chunk,
2633 completed_chunk: chunk,
2634 capacity_probe_count: 0,
2635 }
2636 }
2637
2638 pub const fn planned_chunk(&self) -> PrefillChunk {
2639 self.planned_chunk
2640 }
2641
2642 pub const fn completed_chunk(&self) -> PrefillChunk {
2643 self.completed_chunk
2644 }
2645
2646 pub const fn capacity_probe_count(&self) -> u32 {
2647 self.capacity_probe_count
2648 }
2649
2650 pub fn output(&self) -> &PlanRuntimePrefillOutput {
2651 &self.output
2652 }
2653
2654 pub fn validate_for(
2655 &self,
2656 expected_request_id: &RequestId,
2657 expected_planned_chunk: PrefillChunk,
2658 vocabulary_size: usize,
2659 ) -> Result<()> {
2660 if self.planned_chunk != expected_planned_chunk {
2661 return Err(FerrumError::backend(format!(
2662 "plan runtime completed prefill frontier {:?}, expected {:?}",
2663 self.planned_chunk.range(),
2664 expected_planned_chunk.range()
2665 )));
2666 }
2667 validate_prefill_completion_shape(
2668 self.planned_chunk,
2669 self.completed_chunk,
2670 self.capacity_probe_count,
2671 )?;
2672 self.output.validate_for_completion(
2673 expected_request_id,
2674 self.completed_chunk,
2675 vocabulary_size,
2676 )
2677 }
2678
2679 pub fn into_parts(self) -> (PlanRuntimePrefillOutput, PrefillChunk, PrefillChunk, u32) {
2680 (
2681 self.output,
2682 self.planned_chunk,
2683 self.completed_chunk,
2684 self.capacity_probe_count,
2685 )
2686 }
2687}
2688
2689pub enum PlanRuntimePrefillOutcome {
2690 Completed(PlanRuntimePrefillCompletion),
2691 Deferred(ExecutorExecutionDeferral),
2692}
2693
2694pub enum PlanRuntimeBatchPrefillOutcome {
2695 Completed(Vec<PlanRuntimePrefillCompletion>),
2696 NotSubmitted(ExecutorExecutionDeferral),
2697 Unsupported,
2698}
2699
2700pub struct ExecutorPrefillCompletion {
2706 output: PrefillOutput,
2707 planned_chunk: PrefillChunk,
2708 completed_chunk: PrefillChunk,
2709 capacity_probe_count: u32,
2710}
2711
2712impl ExecutorPrefillCompletion {
2713 pub fn new(
2714 output: PrefillOutput,
2715 planned_chunk: PrefillChunk,
2716 completed_chunk: PrefillChunk,
2717 capacity_probe_count: u32,
2718 ) -> Result<Self> {
2719 validate_prefill_completion_shape(planned_chunk, completed_chunk, capacity_probe_count)?;
2720 Ok(Self {
2721 output,
2722 planned_chunk,
2723 completed_chunk,
2724 capacity_probe_count,
2725 })
2726 }
2727
2728 pub fn exact(output: PrefillOutput, chunk: PrefillChunk) -> Self {
2729 Self {
2730 output,
2731 planned_chunk: chunk,
2732 completed_chunk: chunk,
2733 capacity_probe_count: 0,
2734 }
2735 }
2736
2737 pub const fn planned_chunk(&self) -> PrefillChunk {
2738 self.planned_chunk
2739 }
2740
2741 pub const fn completed_chunk(&self) -> PrefillChunk {
2742 self.completed_chunk
2743 }
2744
2745 pub const fn capacity_probe_count(&self) -> u32 {
2746 self.capacity_probe_count
2747 }
2748
2749 pub fn into_parts(self) -> (PrefillOutput, PrefillChunk, PrefillChunk, u32) {
2750 (
2751 self.output,
2752 self.planned_chunk,
2753 self.completed_chunk,
2754 self.capacity_probe_count,
2755 )
2756 }
2757}
2758
2759fn validate_prefill_completion_shape(
2760 planned_chunk: PrefillChunk,
2761 completed_chunk: PrefillChunk,
2762 capacity_probe_count: u32,
2763) -> Result<()> {
2764 if completed_chunk.tokens_processed() != planned_chunk.tokens_processed()
2765 || completed_chunk.total_prompt_tokens() != planned_chunk.total_prompt_tokens()
2766 || completed_chunk.tokens_to_process() > planned_chunk.tokens_to_process()
2767 {
2768 return Err(ferrum_types::FerrumError::internal(
2769 "completed prefill chunk is not a non-empty prefix of its planned chunk",
2770 ));
2771 }
2772 if completed_chunk != planned_chunk && capacity_probe_count == 0 {
2773 return Err(ferrum_types::FerrumError::internal(
2774 "partial prefill completion requires a failed capacity probe",
2775 ));
2776 }
2777 Ok(())
2778}
2779
2780pub enum ExecutorPrefillOutcome {
2781 Completed(ExecutorPrefillCompletion),
2782 Deferred(ExecutorExecutionDeferral),
2783}
2784
2785pub enum ExecutorBatchPrefillOutcome {
2793 Completed(Vec<ExecutorPrefillCompletion>),
2794 NotSubmitted(ExecutorExecutionDeferral),
2795 Unsupported,
2796}
2797
2798#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
2803#[serde(rename_all = "snake_case")]
2804pub enum ExecutorPrefillMaintenanceStage {
2805 LogicalCapacity,
2806 PhysicalBacking,
2807}
2808
2809#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2812#[serde(tag = "source", rename_all = "snake_case")]
2813pub enum ExecutorPrefillMaintenanceBlocker {
2814 Capacity {
2815 domain_id: Option<u32>,
2816 kind: crate::vnext::CapacityShortfallKind,
2817 requested: u64,
2818 available: u64,
2819 current_total: u64,
2820 maximum_total: u64,
2821 },
2822 Backing {
2823 pool_id: String,
2824 domain_id: u32,
2825 lifetime: crate::vnext::DynamicBackingClaimScope,
2826 reason: crate::vnext::DynamicBackingDeferralReason,
2827 requested_bytes: u64,
2828 free_bytes: u64,
2829 largest_contiguous_bytes: u64,
2830 },
2831}
2832
2833#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2839pub struct ExecutorPrefillMaintenanceDeferral {
2840 request_id: RequestId,
2841 observed: ExecutorAdmissionEpochs,
2842 wait_condition: crate::vnext::CapacityWaitCondition,
2843 stage: ExecutorPrefillMaintenanceStage,
2844 blockers: Vec<ExecutorPrefillMaintenanceBlocker>,
2845}
2846
2847impl ExecutorPrefillMaintenanceDeferral {
2848 pub fn new(
2849 request_id: RequestId,
2850 observed: ExecutorAdmissionEpochs,
2851 wait_condition: crate::vnext::CapacityWaitCondition,
2852 stage: ExecutorPrefillMaintenanceStage,
2853 blockers: Vec<ExecutorPrefillMaintenanceBlocker>,
2854 ) -> Result<Self> {
2855 if blockers.is_empty() {
2856 return Err(ferrum_types::FerrumError::request_validation(
2857 "executor prefill maintenance deferral requires at least one blocker",
2858 ));
2859 }
2860 if wait_condition.coordinator_id().get() != observed.coordinator_id.get() {
2861 return Err(ferrum_types::FerrumError::request_validation(
2862 "executor prefill maintenance wait condition belongs to a different coordinator",
2863 ));
2864 }
2865 Ok(Self {
2866 request_id,
2867 observed,
2868 wait_condition,
2869 stage,
2870 blockers,
2871 })
2872 }
2873
2874 pub fn from_admission(
2875 request_id: &RequestId,
2876 deferred: &crate::vnext::AdmissionDeferred,
2877 ) -> Result<Self> {
2878 if deferred.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2879 return Err(ferrum_types::FerrumError::internal(
2880 "logical prefill maintenance projection requires AwaitBackingGrowth",
2881 ));
2882 }
2883 let blockers = deferred
2884 .blockers()
2885 .iter()
2886 .map(|blocker| ExecutorPrefillMaintenanceBlocker::Capacity {
2887 domain_id: blocker.domain().map(|domain| domain.get()),
2888 kind: blocker.kind(),
2889 requested: blocker.requested().get(),
2890 available: blocker.available().get(),
2891 current_total: blocker.current_total().get(),
2892 maximum_total: blocker.maximum_total().get(),
2893 })
2894 .collect();
2895 Self::new(
2896 request_id.clone(),
2897 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2898 deferred.wait_condition().clone(),
2899 ExecutorPrefillMaintenanceStage::LogicalCapacity,
2900 blockers,
2901 )
2902 }
2903
2904 pub fn from_backing(
2905 request_id: &RequestId,
2906 deferred: &crate::vnext::DynamicBackingDeferred,
2907 ) -> Result<Self> {
2908 let blockers = deferred
2909 .blockers()
2910 .iter()
2911 .map(|blocker| ExecutorPrefillMaintenanceBlocker::Backing {
2912 pool_id: blocker.pool_id().as_str().to_string(),
2913 domain_id: blocker.domain_id().get(),
2914 lifetime: deferred.scope(),
2915 reason: blocker.reason(),
2916 requested_bytes: blocker.requested_bytes(),
2917 free_bytes: blocker.free_bytes(),
2918 largest_contiguous_bytes: blocker.largest_contiguous_bytes(),
2919 })
2920 .collect();
2921 Self::new(
2922 request_id.clone(),
2923 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2924 deferred.wait_condition().clone(),
2925 ExecutorPrefillMaintenanceStage::PhysicalBacking,
2926 blockers,
2927 )
2928 }
2929
2930 pub fn request_id(&self) -> &RequestId {
2931 &self.request_id
2932 }
2933
2934 pub const fn observed(&self) -> ExecutorAdmissionEpochs {
2935 self.observed
2936 }
2937
2938 pub fn wait_condition(&self) -> &crate::vnext::CapacityWaitCondition {
2939 &self.wait_condition
2940 }
2941
2942 pub const fn stage(&self) -> ExecutorPrefillMaintenanceStage {
2943 self.stage
2944 }
2945
2946 pub fn blockers(&self) -> &[ExecutorPrefillMaintenanceBlocker] {
2947 &self.blockers
2948 }
2949}
2950
2951#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2953#[serde(tag = "outcome", rename_all = "snake_case")]
2954pub enum ExecutorPrefillMaintenanceOutcome {
2955 NoLongerPending,
2958 RetryAdmission { current: ExecutorAdmissionEpochs },
2963 WaitForRelease {
2967 current: ExecutorAdmissionEpochs,
2968 wait_condition: crate::vnext::CapacityWaitCondition,
2969 pressure: crate::vnext::DynamicBackingPressure,
2970 },
2971 Maintained {
2974 current: ExecutorAdmissionEpochs,
2975 pools_grown: usize,
2976 allocated_bytes: u64,
2977 pools_reclaimed: usize,
2978 chunks_reclaimed: usize,
2979 reclaimed_bytes: u64,
2980 rebalance: Option<crate::vnext::DynamicPoolRebalanceReceipt>,
2983 },
2984}
2985
2986#[derive(Debug, Clone)]
2992pub enum ExecutorPrefillAdmissionDecision {
2993 Admitted(ExecutorPrefillAdmissionReceipt),
2994 Deferred(crate::vnext::AdmissionDeferred),
2995 MaintenanceDeferred(ExecutorPrefillMaintenanceDeferral),
2996 PermanentRejected(crate::vnext::AdmissionRejected),
2997}
2998
2999#[async_trait]
3001pub trait ModelExecutor: Send + Sync {
3002 fn plan_prompt_tail_capture_boundary(&self, _chunk: PrefillChunk) -> Option<PrefixCapturePlan> {
3008 None
3009 }
3010
3011 fn plan_prefix_capture_boundary(
3014 &self,
3015 _input: PrefixCaptureBoundary<'_>,
3016 ) -> Option<PrefixCapturePlan> {
3017 None
3018 }
3019
3020 fn retain_prefix_capture_interest(
3024 &self,
3025 _input: PrefixCaptureRequest<'_>,
3026 ) -> Result<Option<Arc<dyn PrefixCaptureLease>>> {
3027 Ok(None)
3028 }
3029
3030 fn info(&self) -> &ModelInfo;
3032
3033 fn execution_resource_authority(&self) -> ExecutionResourceAuthority {
3039 ExecutionResourceAuthority::LegacyEngine
3040 }
3041
3042 fn admission_limits(&self) -> Result<Option<ExecutorAdmissionLimits>> {
3046 Ok(None)
3047 }
3048
3049 fn resolved_model_plan(&self) -> Option<&crate::vnext::ResolvedModelPlan> {
3054 None
3055 }
3056
3057 fn plan_runtime_resource_snapshot(&self) -> Result<Option<PlanRuntimeResourceSnapshot>> {
3061 Ok(None)
3062 }
3063
3064 fn supports_native_unified_decode(&self) -> bool {
3073 false
3074 }
3075
3076 fn kv_capacity(&self) -> Option<usize> {
3079 None
3080 }
3081
3082 fn startup_memory_plan(&self) -> Option<&ferrum_types::StartupMemoryPlan> {
3085 None
3086 }
3087
3088 fn attach_execution_event_sink(&self, _sink: Arc<dyn crate::vnext::ExecutionEventSink>) {}
3094
3095 fn execution_capacity_epochs(&self) -> Result<Option<ExecutorAdmissionEpochs>> {
3099 Ok(None)
3100 }
3101
3102 fn write_execution_capacity_snapshot(
3107 &self,
3108 availability: &mut Vec<crate::vnext::CapacityAvailabilityEpoch>,
3109 ) -> Result<Option<ExecutorAdmissionEpochs>> {
3110 availability.clear();
3111 self.execution_capacity_epochs()
3112 }
3113
3114 fn register_execution_capacity_waiter(
3122 &self,
3123 _observed: &crate::vnext::CapacityWaitCondition,
3124 ) -> Result<Option<ExecutorCapacityWaitRegistration>> {
3125 Ok(None)
3126 }
3127
3128 fn try_admit_prefill(
3132 &self,
3133 _input: ExecutorPrefillAdmission<'_>,
3134 ) -> Result<ExecutorPrefillAdmissionDecision> {
3135 Err(ferrum_types::FerrumError::unsupported(
3136 "plan-runtime prefill admission is not implemented",
3137 ))
3138 }
3139
3140 fn cancel_prefill_admission(&self, _request_id: &RequestId) -> bool {
3144 false
3145 }
3146
3147 fn supports_plan_runtime_prefix_restore(&self) -> bool {
3150 false
3151 }
3152
3153 async fn try_restore_plan_runtime_prefix(
3160 &self,
3161 _input: PlanRuntimePrefixRestoreInput<'_>,
3162 ) -> Result<PlanRuntimePrefixRestoreOutcome> {
3163 Ok(PlanRuntimePrefixRestoreOutcome::Unavailable)
3164 }
3165
3166 fn write_execution_capacity_release_sources(
3175 &self,
3176 _preemption: &ExecutorExecutionCapacityPreemption,
3177 sources: &mut Vec<crate::vnext::CapacityAvailabilitySource>,
3178 ) -> Result<bool> {
3179 sources.clear();
3180 Ok(false)
3181 }
3182
3183 async fn preempt_execution_capacity(
3190 &self,
3191 _preemption: ExecutorExecutionCapacityPreemption,
3192 ) -> Result<ExecutorExecutionCapacityPreemptionReceipt> {
3193 Err(FerrumError::unsupported(
3194 "request-scoped execution-capacity preemption is not implemented",
3195 ))
3196 }
3197
3198 fn maintain_prefill_backing(
3203 &self,
3204 _request_id: &RequestId,
3205 ) -> Result<ExecutorPrefillMaintenanceOutcome> {
3206 Err(ferrum_types::FerrumError::unsupported(
3207 "plan-runtime prefill backing maintenance is not implemented",
3208 ))
3209 }
3210
3211 fn reserve_kv_slots(&self, _requests: &[KvSlotRequest]) -> Result<Option<KvSlotReservation>> {
3218 Ok(None)
3219 }
3220
3221 fn kv_slot_capacity_snapshot(&self) -> Option<KvSlotCapacitySnapshot> {
3225 None
3226 }
3227
3228 fn recurrent_state_spec(
3235 &self,
3236 _request_id: &RequestId,
3237 _input_tokens: &[TokenId],
3238 ) -> Result<Option<RecurrentStateSpec>> {
3239 Ok(None)
3240 }
3241
3242 async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput>;
3244
3245 async fn prefill_with_capacity(&self, input: &PrefillInput) -> Result<ExecutorPrefillOutcome> {
3248 let output = self.prefill(input).await?;
3249 let chunk = match input.chunk {
3250 Some(chunk) => chunk,
3251 None => PrefillChunk::new(0, input.sequence_length(), input.sequence_length())?,
3252 };
3253 Ok(ExecutorPrefillOutcome::Completed(
3254 ExecutorPrefillCompletion::exact(output, chunk),
3255 ))
3256 }
3257
3258 async fn batch_prefill(&self, inputs: &[PrefillInput]) -> Result<Vec<PrefillOutput>> {
3271 let mut outputs = Vec::with_capacity(inputs.len());
3272 for input in inputs {
3273 outputs.push(self.prefill(input).await?);
3274 }
3275 Ok(outputs)
3276 }
3277
3278 async fn batch_prefill_with_capacity(
3285 &self,
3286 _inputs: &[PrefillInput],
3287 ) -> Result<ExecutorBatchPrefillOutcome> {
3288 Ok(ExecutorBatchPrefillOutcome::Unsupported)
3289 }
3290
3291 async fn plan_runtime_prefill_with_capacity(
3297 &self,
3298 _input: &PlanRuntimePrefillInput,
3299 ) -> Result<PlanRuntimePrefillOutcome> {
3300 Err(FerrumError::unsupported(
3301 "tensor-free plan-runtime prefill is not implemented",
3302 ))
3303 }
3304
3305 async fn plan_runtime_batch_prefill_with_capacity(
3311 &self,
3312 _inputs: &[PlanRuntimePrefillInput],
3313 ) -> Result<PlanRuntimeBatchPrefillOutcome> {
3314 Ok(PlanRuntimeBatchPrefillOutcome::Unsupported)
3315 }
3316
3317 fn discard_plan_runtime_prefill(&self, authority: PlanRuntimePrefillAuthority) -> Result<()> {
3323 self.release_cache(&authority.kv_cache().cache_id());
3324 Ok(())
3325 }
3326
3327 async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput>;
3329
3330 async fn batch_decode(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>> {
3341 let mut outputs = Vec::with_capacity(inputs.len());
3342 for input in inputs {
3343 outputs.push(self.decode(input).await?);
3344 }
3345 Ok(outputs)
3346 }
3347
3348 async fn batch_decode_with_capacity(
3354 &self,
3355 inputs: &[DecodeInput],
3356 ) -> Result<ExecutorBatchDecodeOutcome> {
3357 self.batch_decode(inputs)
3358 .await
3359 .map(ExecutorBatchDecodeOutcome::Completed)
3360 }
3361
3362 async fn plan_runtime_batch_decode_with_capacity(
3370 &self,
3371 _inputs: &[PlanRuntimeDecodeInput],
3372 ) -> Result<PlanRuntimeBatchDecodeOutcome> {
3373 Err(FerrumError::unsupported(
3374 "tensor-free plan-runtime batch decode is not implemented",
3375 ))
3376 }
3377
3378 async fn unified_decode(&self, _batch: &UnifiedBatch) -> Result<Vec<Option<Vec<f32>>>> {
3399 Err(ferrum_types::FerrumError::unsupported(
3400 "unified_decode not implemented for this executor",
3401 ))
3402 }
3403
3404 async fn forward(&self, _input: &TensorRef) -> Result<TensorRef> {
3406 Err(ferrum_types::FerrumError::unsupported(
3408 "Full forward pass not supported by this executor",
3409 ))
3410 }
3411
3412 async fn truncate_kv(
3418 &self,
3419 _kv_cache: &std::sync::Arc<dyn crate::KvCacheHandle>,
3420 _new_len: usize,
3421 ) -> Result<()> {
3422 Ok(())
3423 }
3424
3425 async fn forward_verify(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>> {
3435 let mut out = Vec::with_capacity(inputs.len());
3436 for input in inputs {
3437 out.push(self.decode(input).await?);
3438 }
3439 Ok(out)
3440 }
3441
3442 fn capabilities(&self) -> ExecutorCapabilities;
3444
3445 fn status(&self) -> ExecutorStatus;
3447
3448 fn cache_metrics_snapshot(&self) -> Option<serde_json::Value> {
3454 None
3455 }
3456
3457 fn execution_attribution_snapshot(&self) -> Option<serde_json::Value> {
3461 None
3462 }
3463
3464 fn lora_metrics_snapshot(&self) -> Option<serde_json::Value> {
3466 None
3467 }
3468
3469 async fn prepare_startup(&self) -> Result<()> {
3477 Ok(())
3478 }
3479
3480 async fn warmup(&mut self) -> Result<()> {
3482 Ok(())
3484 }
3485
3486 async fn shutdown(&mut self) -> Result<()> {
3488 Ok(())
3490 }
3491
3492 async fn complete_cache(&self, completion: ExecutorSequenceCompletion) -> Result<()> {
3501 self.release_cache(completion.cache_id());
3502 Ok(())
3503 }
3504
3505 fn release_cache(&self, _cache_id: &str) {
3511 }
3513}
3514
3515#[derive(Debug, Clone, Serialize, Deserialize)]
3517pub struct ExecutorCapabilities {
3518 pub max_batch_size: usize,
3520 pub max_sequence_length: usize,
3522 pub attention_mechanisms: Vec<AttentionType>,
3524 pub supports_dynamic_batching: bool,
3526 pub supports_continuous_batching: bool,
3528 pub supports_speculative_decoding: bool,
3530 pub supports_tensor_parallelism: bool,
3532 pub supports_pipeline_parallelism: bool,
3534 pub supported_dtypes: Vec<ferrum_types::DataType>,
3536 pub supported_devices: Vec<ferrum_types::Device>,
3538 pub memory_requirements: MemoryRequirements,
3540}
3541
3542#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
3544pub enum AttentionType {
3545 MultiHead,
3547 MultiQuery,
3549 GroupedQuery,
3551 Flash,
3553 Paged,
3555 SlidingWindow,
3557}
3558
3559#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
3562pub struct TypedSequenceStateMemory {
3563 pub kv_bytes_per_token: u64,
3565 pub other_token_scaled_bytes_per_token: u64,
3567 pub fixed_bytes_per_sequence: u64,
3569}
3570
3571#[derive(Debug, Clone, Deserialize)]
3574pub struct MemoryRequirements {
3575 pub parameter_memory: u64,
3577 pub activation_memory_per_token: usize,
3579 #[serde(default)]
3582 pub kv_cache_memory_per_token: usize,
3583 #[serde(default)]
3586 pub typed_sequence_state: Option<TypedSequenceStateMemory>,
3587 pub overhead_memory: u64,
3589}
3590
3591impl Serialize for MemoryRequirements {
3592 fn serialize<S: serde::Serializer>(
3593 &self,
3594 serializer: S,
3595 ) -> std::result::Result<S::Ok, S::Error> {
3596 use serde::ser::SerializeStruct;
3597 let mut value = serializer.serialize_struct("MemoryRequirements", 4)?;
3598 value.serialize_field("parameter_memory", &self.parameter_memory)?;
3599 value.serialize_field(
3600 "activation_memory_per_token",
3601 &self.activation_memory_per_token,
3602 )?;
3603 match &self.typed_sequence_state {
3604 Some(state) => value.serialize_field("typed_sequence_state", state)?,
3605 None => value
3606 .serialize_field("kv_cache_memory_per_token", &self.kv_cache_memory_per_token)?,
3607 }
3608 value.serialize_field("overhead_memory", &self.overhead_memory)?;
3609 value.end()
3610 }
3611}
3612
3613impl MemoryRequirements {
3614 pub fn checked_calculate_total_memory(
3617 &self,
3618 batch_size: usize,
3619 sequence_length: usize,
3620 num_layers: usize,
3621 ) -> Option<u64> {
3622 let batch = u64::try_from(batch_size).ok()?;
3623 let tokens = u64::try_from(sequence_length).ok()?;
3624 let token_count = batch.checked_mul(tokens)?;
3625 let activation_mem = u64::try_from(self.activation_memory_per_token)
3626 .ok()?
3627 .checked_mul(token_count)?;
3628 let state_mem = match self.typed_sequence_state {
3629 Some(state) => state
3630 .kv_bytes_per_token
3631 .checked_add(state.other_token_scaled_bytes_per_token)?
3632 .checked_mul(token_count)?
3633 .checked_add(state.fixed_bytes_per_sequence.checked_mul(batch)?)?,
3634 None => u64::try_from(self.kv_cache_memory_per_token)
3635 .ok()?
3636 .checked_mul(token_count)?
3637 .checked_mul(u64::try_from(num_layers).ok()?)?,
3638 };
3639 self.parameter_memory
3640 .checked_add(activation_mem)?
3641 .checked_add(state_mem)?
3642 .checked_add(self.overhead_memory)
3643 }
3644
3645 pub fn calculate_total_memory(
3648 &self,
3649 batch_size: usize,
3650 sequence_length: usize,
3651 num_layers: usize,
3652 ) -> u64 {
3653 self.checked_calculate_total_memory(batch_size, sequence_length, num_layers)
3654 .unwrap_or(u64::MAX)
3655 }
3656}
3657
3658#[cfg(test)]
3659mod memory_requirements_tests {
3660 use super::{MemoryRequirements, TypedSequenceStateMemory};
3661
3662 #[test]
3663 fn typed_sequence_memory_counts_scales_and_fixed_state_once_per_sequence() {
3664 let memory = MemoryRequirements {
3665 parameter_memory: 100,
3666 activation_memory_per_token: 4,
3667 kv_cache_memory_per_token: 999,
3668 typed_sequence_state: Some(TypedSequenceStateMemory {
3669 kv_bytes_per_token: 528,
3670 other_token_scaled_bytes_per_token: 8,
3671 fixed_bytes_per_sequence: 64,
3672 }),
3673 overhead_memory: 20,
3674 };
3675 let expected = 100 + 4 * 2 * 3 + (528 + 8) * 2 * 3 + 64 * 2 + 20;
3676 for layers in [1, 3, 32] {
3677 assert_eq!(
3678 memory.checked_calculate_total_memory(2, 3, layers),
3679 Some(expected)
3680 );
3681 }
3682 let wire = serde_json::to_value(&memory).unwrap();
3683 assert!(wire.get("kv_cache_memory_per_token").is_none());
3684 assert_eq!(wire["typed_sequence_state"]["kv_bytes_per_token"], 528);
3685 let decoded: MemoryRequirements = serde_json::from_value(wire).unwrap();
3686 assert_eq!(decoded.calculate_total_memory(2, 3, 32), expected);
3687 }
3688
3689 #[test]
3690 fn legacy_memory_wire_keeps_per_layer_calculation_and_rejects_overflow() {
3691 let wire = serde_json::json!({
3692 "parameter_memory": 100,
3693 "activation_memory_per_token": 4,
3694 "kv_cache_memory_per_token": 16,
3695 "overhead_memory": 20,
3696 });
3697 let mut memory: MemoryRequirements = serde_json::from_value(wire.clone()).unwrap();
3698 assert!(memory.typed_sequence_state.is_none());
3699 assert_eq!(
3700 memory.calculate_total_memory(2, 3, 5),
3701 100 + 4 * 2 * 3 + 16 * 2 * 3 * 5 + 20
3702 );
3703 assert_eq!(serde_json::to_value(&memory).unwrap(), wire);
3704 memory.parameter_memory = u64::MAX;
3705 assert_eq!(memory.checked_calculate_total_memory(1, 1, 1), None);
3706 assert_eq!(memory.calculate_total_memory(1, 1, 1), u64::MAX);
3707 memory.parameter_memory = 0;
3708 memory.typed_sequence_state = Some(TypedSequenceStateMemory {
3709 kv_bytes_per_token: u64::MAX,
3710 other_token_scaled_bytes_per_token: 1,
3711 fixed_bytes_per_sequence: 0,
3712 });
3713 assert_eq!(memory.checked_calculate_total_memory(1, 1, 1), None);
3714 }
3715}
3716
3717#[derive(Debug, Clone, Serialize, Deserialize)]
3719pub struct ExecutorStatus {
3720 pub state: ExecutorState,
3722 pub is_ready: bool,
3724 pub current_batch_size: usize,
3726 pub prefill_operations: u64,
3728 pub decode_operations: u64,
3730 pub avg_prefill_time_ms: f64,
3732 pub avg_decode_time_ms: f64,
3734 pub memory_usage: ExecutorMemoryUsage,
3736 #[serde(skip)]
3738 pub last_operation: Option<std::time::Instant>,
3739}
3740
3741#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
3743pub enum ExecutorState {
3744 Initializing,
3746 Ready,
3748 Busy,
3750 Error,
3752 Shutdown,
3754}
3755
3756#[derive(Debug, Clone, Serialize, Deserialize)]
3758pub struct ExecutorMemoryUsage {
3759 pub allocated_bytes: usize,
3761 pub used_bytes: usize,
3763 pub peak_bytes: usize,
3765 pub utilization_percent: f32,
3767}
3768
3769#[async_trait]
3771pub trait BatchModelExecutor: ModelExecutor {
3772 async fn batch_prefill(&self, inputs: &[PrefillInput]) -> Result<Vec<PrefillOutput>>;
3774
3775 async fn batch_decode(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>>;
3777
3778 fn optimal_batch_size(&self) -> usize;
3780
3781 fn supports_batch_size(&self, batch_size: usize) -> bool;
3783}
3784
3785#[async_trait]
3787pub trait SpeculativeExecutor: ModelExecutor {
3788 async fn speculative_decode(
3790 &self,
3791 input: &DecodeInput,
3792 draft_tokens: &[ferrum_types::TokenId],
3793 acceptance_threshold: f32,
3794 ) -> Result<SpeculativeDecodeOutput>;
3795}
3796
3797#[derive(Debug, Clone)]
3799pub struct SpeculativeDecodeOutput {
3800 pub accepted_tokens: Vec<ferrum_types::TokenId>,
3802 pub next_logits: TensorRef,
3804 pub kv_cache: Arc<dyn KvCacheHandle>,
3806 pub acceptance_count: usize,
3808}
3809
3810#[async_trait]
3812pub trait ModelExecutorFactory: Send + Sync {
3813 async fn create_executor(&self, config: &ExecutorConfig) -> Result<Box<dyn ModelExecutor>>;
3815
3816 async fn create_batch_executor(
3818 &self,
3819 config: &ExecutorConfig,
3820 ) -> Result<Box<dyn BatchModelExecutor>>;
3821
3822 fn supported_types(&self) -> Vec<ExecutorType>;
3824
3825 fn validate_config(&self, config: &ExecutorConfig) -> Result<()>;
3827}
3828
3829#[derive(Debug, Clone, Serialize, Deserialize)]
3831pub struct ExecutorConfig {
3832 pub model_info: ModelInfo,
3834 pub device: ferrum_types::Device,
3836 pub dtype: ferrum_types::DataType,
3838 pub max_batch_size: usize,
3840 pub max_sequence_length: usize,
3842 pub attention_config: ExecutorAttentionConfig,
3844 pub memory_config: ExecutorMemoryConfig,
3846 pub optimization_config: OptimizationConfig,
3848 pub executor_options: HashMap<String, serde_json::Value>,
3850}
3851
3852#[derive(Debug, Clone, Serialize, Deserialize)]
3858pub struct ExecutorAttentionConfig {
3859 pub attention_type: AttentionType,
3861 pub enable_flash_attention: bool,
3863 pub enable_paged_attention: bool,
3865 pub block_size: Option<usize>,
3867 pub sliding_window_size: Option<usize>,
3869}
3870
3871#[derive(Debug, Clone, Serialize, Deserialize)]
3873pub struct ExecutorMemoryConfig {
3874 pub enable_memory_pooling: bool,
3876 pub memory_pool_size: Option<usize>,
3878 pub enable_kv_cache_sharing: bool,
3880 pub max_memory_usage: f32,
3882}
3883
3884#[derive(Debug, Clone, Serialize, Deserialize)]
3886pub struct OptimizationConfig {
3887 pub enable_cuda_graphs: bool,
3889 pub enable_kernel_fusion: bool,
3891 pub enable_mixed_precision: bool,
3893 pub optimization_level: u8,
3895 pub custom_flags: HashMap<String, bool>,
3897}
3898
3899#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
3901pub enum ExecutorType {
3902 Sequential,
3904 Batch,
3906 ContinuousBatch,
3908 Speculative,
3910 PipelineParallel,
3912 TensorParallel,
3914}
3915
3916#[derive(Debug, Clone, Serialize, Deserialize)]
3918pub struct ExecutorMetrics {
3919 pub total_operations: u64,
3921 pub prefill_operations: u64,
3923 pub decode_operations: u64,
3925 pub avg_prefill_latency: f64,
3927 pub avg_decode_latency: f64,
3929 pub p95_prefill_latency: f64,
3931 pub p95_decode_latency: f64,
3933 pub throughput_tps: f64,
3935 pub memory_efficiency: f32,
3937 pub batch_utilization: f32,
3939}
3940
3941pub trait ExecutorRegistry: Send + Sync {
3943 fn register(&mut self, name: &str, executor: Box<dyn ModelExecutor>) -> Result<()>;
3945
3946 fn get(&self, name: &str) -> Option<&dyn ModelExecutor>;
3948
3949 fn remove(&mut self, name: &str) -> Option<Box<dyn ModelExecutor>>;
3951
3952 fn list_names(&self) -> Vec<String>;
3954
3955 fn get_metrics(&self, name: &str) -> Option<ExecutorMetrics>;
3957}