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
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28pub struct KvSlotRequest {
29 pub cache_id: String,
30 pub target_len: usize,
31 pub admission_target_len: Option<usize>,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
36pub struct KvSlotAllocation {
37 pub cache_id: String,
38 pub blocks_before: usize,
39 pub blocks_after: usize,
40 pub new_blocks: usize,
41}
42
43#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49pub struct KvSlotReservation {
50 pub block_size: usize,
51 pub total_blocks: usize,
52 pub free_blocks_before: usize,
53 pub free_blocks_after: usize,
54 pub allocations: Vec<KvSlotAllocation>,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
64pub struct KvSlotCapacitySnapshot {
65 pub block_size: usize,
66 pub total_blocks: usize,
67 pub free_blocks: usize,
68}
69
70#[derive(Clone)]
77pub struct TokenSelectionMask {
78 pub fingerprint: u64,
79 pub valid_token_mask: Arc<[i8]>,
80}
81
82impl TokenSelectionMask {
83 pub fn new(valid_token_mask: Vec<i8>) -> Self {
84 let fingerprint = Self::fingerprint(&valid_token_mask);
85 Self {
86 fingerprint,
87 valid_token_mask: Arc::from(valid_token_mask),
88 }
89 }
90
91 fn fingerprint(valid_token_mask: &[i8]) -> u64 {
92 let mut hasher = DefaultHasher::new();
93 valid_token_mask.hash(&mut hasher);
94 hasher.finish()
95 }
96
97 pub fn set_tokens_validity(&mut self, token_ids: &[u32], valid: bool) -> bool {
102 let value = i8::from(valid);
103 let slots = Arc::make_mut(&mut self.valid_token_mask);
104 let mut changed = false;
105 for &token_id in token_ids {
106 if let Some(slot) = slots.get_mut(token_id as usize) {
107 if *slot != value {
108 *slot = value;
109 changed = true;
110 }
111 }
112 }
113 if changed {
114 self.fingerprint = Self::fingerprint(slots);
115 }
116 changed
117 }
118
119 pub fn len(&self) -> usize {
120 self.valid_token_mask.len()
121 }
122
123 pub fn is_empty(&self) -> bool {
124 self.valid_token_mask.is_empty()
125 }
126}
127
128impl std::fmt::Debug for TokenSelectionMask {
129 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130 let valid_count = self.valid_token_mask.iter().filter(|&&v| v != 0).count();
131 f.debug_struct("TokenSelectionMask")
132 .field("fingerprint", &self.fingerprint)
133 .field("len", &self.valid_token_mask.len())
134 .field("valid_count", &valid_count)
135 .finish()
136 }
137}
138
139#[cfg(test)]
140mod token_selection_mask_tests {
141 use super::TokenSelectionMask;
142
143 #[test]
144 fn response_completion_mask_is_copy_on_write_and_restores_fingerprint() {
145 let mut mask = TokenSelectionMask::new(vec![1, 1, 1]);
146 let original = mask.clone();
147 let original_fingerprint = mask.fingerprint;
148
149 assert!(mask.set_tokens_validity(&[1], false));
150 assert_eq!(mask.valid_token_mask.as_ref(), &[1, 0, 1]);
151 assert_eq!(original.valid_token_mask.as_ref(), &[1, 1, 1]);
152 assert_ne!(mask.fingerprint, original_fingerprint);
153
154 let masked_fingerprint = mask.fingerprint;
155 assert!(!mask.set_tokens_validity(&[1], false));
156 assert_eq!(mask.fingerprint, masked_fingerprint);
157
158 assert!(mask.set_tokens_validity(&[1], true));
159 assert_eq!(mask.fingerprint, original_fingerprint);
160 }
161}
162
163#[derive(Clone, Debug)]
164pub enum LogitsReturnPolicy {
165 FullLogits,
166 GreedyArgmax {
167 token_mask: Option<TokenSelectionMask>,
168 repetition_penalty: Option<GreedyRepetitionPenalty>,
169 },
170}
171
172impl Default for LogitsReturnPolicy {
173 fn default() -> Self {
174 Self::FullLogits
175 }
176}
177
178impl LogitsReturnPolicy {
179 pub fn requires_full_logits(&self) -> bool {
180 matches!(self, Self::FullLogits)
181 }
182}
183
184#[derive(Debug, Clone, PartialEq)]
191pub enum ExecutorSamplingOutput {
192 FullLogits(Vec<f32>),
193 GreedyToken(TokenId),
194}
195
196impl ExecutorSamplingOutput {
197 pub fn full_logits(logits: Vec<f32>) -> Result<Self> {
198 if logits.is_empty() {
199 return Err(FerrumError::backend(
200 "plan-runtime sampling output requires non-empty logits",
201 ));
202 }
203 Ok(Self::FullLogits(logits))
204 }
205
206 pub const fn greedy_token(token: TokenId) -> Self {
207 Self::GreedyToken(token)
208 }
209
210 pub fn validate_for_policy(
216 &self,
217 policy: &LogitsReturnPolicy,
218 vocabulary_size: usize,
219 ) -> Result<()> {
220 match self {
221 Self::FullLogits(logits) if logits.len() != vocabulary_size => {
222 return Err(FerrumError::backend(format!(
223 "plan runtime returned {} logits for vocabulary {vocabulary_size}",
224 logits.len()
225 )));
226 }
227 Self::GreedyToken(_) if policy.requires_full_logits() => {
228 return Err(FerrumError::backend(
229 "plan runtime returned a greedy token for a full-logits request",
230 ));
231 }
232 Self::GreedyToken(token)
233 if usize::try_from(token.get())
234 .ok()
235 .is_none_or(|token| token >= vocabulary_size) =>
236 {
237 return Err(FerrumError::backend(format!(
238 "plan runtime returned token {} outside vocabulary {vocabulary_size}",
239 token.get()
240 )));
241 }
242 _ => {}
243 }
244 Ok(())
245 }
246
247 pub fn into_full_logits(self) -> Result<Vec<f32>> {
248 match self {
249 Self::FullLogits(logits) => Ok(logits),
250 Self::GreedyToken(_) => Err(FerrumError::backend(
251 "plan-runtime prefill unexpectedly returned a selected token",
252 )),
253 }
254 }
255}
256
257#[cfg(test)]
258mod executor_sampling_output_tests {
259 use super::{ExecutorSamplingOutput, LogitsReturnPolicy};
260 use ferrum_types::TokenId;
261
262 #[test]
263 fn full_logits_require_exact_vocabulary_width() {
264 let output = ExecutorSamplingOutput::full_logits(vec![0.0; 4]).unwrap();
265 assert!(output
266 .validate_for_policy(&LogitsReturnPolicy::FullLogits, 4)
267 .is_ok());
268 assert!(output
269 .validate_for_policy(&LogitsReturnPolicy::FullLogits, 5)
270 .is_err());
271 }
272
273 #[test]
274 fn greedy_token_requires_greedy_policy_and_in_vocabulary_token() {
275 let allowed = LogitsReturnPolicy::GreedyArgmax {
276 token_mask: None,
277 repetition_penalty: None,
278 };
279 let output = ExecutorSamplingOutput::greedy_token(TokenId::new(3));
280 assert!(output.validate_for_policy(&allowed, 4).is_ok());
281 assert!(output.validate_for_policy(&allowed, 3).is_err());
282 assert!(output
283 .validate_for_policy(&LogitsReturnPolicy::FullLogits, 4)
284 .is_err());
285 }
286
287 #[test]
288 fn full_logits_are_a_legal_greedy_batch_fallback() {
289 let policy = LogitsReturnPolicy::GreedyArgmax {
290 token_mask: None,
291 repetition_penalty: None,
292 };
293 let output = ExecutorSamplingOutput::full_logits(vec![0.0; 4]).unwrap();
294 assert!(output.validate_for_policy(&policy, 4).is_ok());
295 }
296}
297
298#[derive(Clone, Debug)]
304pub struct GreedyRepetitionPenalty {
305 penalty: f32,
306 token_ids: Arc<[u32]>,
307}
308
309impl GreedyRepetitionPenalty {
310 pub fn new(penalty: f32, mut token_ids: Vec<u32>) -> Self {
311 let mut seen = HashSet::with_capacity(token_ids.len());
312 token_ids.retain(|token| seen.insert(*token));
313 Self {
314 penalty,
315 token_ids: Arc::from(token_ids),
316 }
317 }
318
319 pub const fn penalty(&self) -> f32 {
320 self.penalty
321 }
322
323 pub fn token_ids(&self) -> &[u32] {
324 &self.token_ids
325 }
326
327 pub fn is_empty(&self) -> bool {
328 self.token_ids.is_empty() || self.penalty == 1.0
329 }
330}
331
332#[cfg(test)]
333mod greedy_repetition_penalty_tests {
334 use super::GreedyRepetitionPenalty;
335
336 #[test]
337 fn constructor_preserves_first_seen_order_and_removes_duplicates() {
338 let repetition = GreedyRepetitionPenalty::new(1.1, vec![7, 3, 7, 9, 3]);
339 assert_eq!(repetition.penalty(), 1.1);
340 assert_eq!(repetition.token_ids(), [7, 3, 9]);
341 }
342}
343
344#[derive(Debug, Clone)]
346pub struct PrefillInput {
347 pub request_id: Option<RequestId>,
349 pub maximum_sequence_tokens: Option<usize>,
352 pub chunk: Option<PrefillChunk>,
357 pub input_ids: TensorRef,
359 pub attention_mask: Option<TensorRef>,
361 pub position_ids: Option<TensorRef>,
363 pub kv_cache: Option<Arc<dyn KvCacheHandle>>,
365 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
367 pub metadata: HashMap<String, serde_json::Value>,
369}
370
371impl PrefillInput {
372 pub fn new(input_ids: TensorRef) -> Self {
374 Self {
375 request_id: None,
376 maximum_sequence_tokens: None,
377 chunk: None,
378 input_ids,
379 attention_mask: None,
380 position_ids: None,
381 kv_cache: None,
382 recurrent_state: None,
383 metadata: HashMap::new(),
384 }
385 }
386
387 pub fn with_request_context(
389 mut self,
390 request_id: RequestId,
391 maximum_sequence_tokens: usize,
392 ) -> Self {
393 self.request_id = Some(request_id);
394 self.maximum_sequence_tokens = Some(maximum_sequence_tokens);
395 self
396 }
397
398 pub fn with_chunk(mut self, chunk: PrefillChunk) -> Self {
400 self.chunk = Some(chunk);
401 self
402 }
403
404 pub fn with_kv_cache(mut self, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
406 self.kv_cache = Some(kv_cache);
407 self
408 }
409
410 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
412 self.recurrent_state = Some(recurrent_state);
413 self
414 }
415
416 pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
418 self.metadata = metadata;
419 self
420 }
421
422 pub fn with_attention_mask(mut self, mask: TensorRef) -> Self {
424 self.attention_mask = Some(mask);
425 self
426 }
427
428 pub fn with_position_ids(mut self, positions: TensorRef) -> Self {
430 self.position_ids = Some(positions);
431 self
432 }
433
434 pub fn batch_size(&self) -> usize {
436 self.input_ids.shape()[0]
437 }
438
439 pub fn sequence_length(&self) -> usize {
441 if self.input_ids.shape().len() >= 2 {
442 self.input_ids.shape()[1]
443 } else {
444 1
445 }
446 }
447}
448
449#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
451pub struct PrefillChunk {
452 tokens_processed: usize,
453 tokens_to_process: usize,
454 total_prompt_tokens: usize,
455}
456
457#[cfg(test)]
458mod prefill_chunk_tests {
459 use super::PrefillChunk;
460
461 #[test]
462 fn validates_exact_progress_and_finality() {
463 let first = PrefillChunk::new(0, 3, 8).unwrap();
464 assert_eq!(first.range(), 0..3);
465 assert_eq!(first.end(), 3);
466 assert!(!first.is_final());
467
468 let final_chunk = PrefillChunk::new(3, 5, 8).unwrap();
469 assert_eq!(final_chunk.range(), 3..8);
470 assert!(final_chunk.is_final());
471 }
472
473 #[test]
474 fn rejects_empty_out_of_bounds_and_overflowing_progress() {
475 assert!(PrefillChunk::new(0, 0, 8).is_err());
476 assert!(PrefillChunk::new(0, 1, 0).is_err());
477 assert!(PrefillChunk::new(7, 2, 8).is_err());
478 assert!(PrefillChunk::new(usize::MAX, 1, usize::MAX).is_err());
479 }
480}
481
482impl PrefillChunk {
483 pub fn new(
484 tokens_processed: usize,
485 tokens_to_process: usize,
486 total_prompt_tokens: usize,
487 ) -> Result<Self> {
488 let end = tokens_processed
489 .checked_add(tokens_to_process)
490 .ok_or_else(|| {
491 ferrum_types::FerrumError::request_validation("prefill chunk overflows")
492 })?;
493 if tokens_to_process == 0 || total_prompt_tokens == 0 || end > total_prompt_tokens {
494 return Err(ferrum_types::FerrumError::request_validation(
495 "prefill chunk must be non-empty and within the full prompt",
496 ));
497 }
498 Ok(Self {
499 tokens_processed,
500 tokens_to_process,
501 total_prompt_tokens,
502 })
503 }
504
505 pub const fn tokens_processed(self) -> usize {
506 self.tokens_processed
507 }
508
509 pub const fn tokens_to_process(self) -> usize {
510 self.tokens_to_process
511 }
512
513 pub const fn total_prompt_tokens(self) -> usize {
514 self.total_prompt_tokens
515 }
516
517 pub fn range(self) -> Range<usize> {
518 self.tokens_processed..self.tokens_processed + self.tokens_to_process
519 }
520
521 pub const fn end(self) -> usize {
522 self.tokens_processed + self.tokens_to_process
523 }
524
525 pub const fn is_final(self) -> bool {
526 self.end() == self.total_prompt_tokens
527 }
528}
529
530#[derive(Debug, Clone)]
532pub struct PrefillOutput {
533 pub logits: TensorRef,
535 pub kv_cache: Arc<dyn KvCacheHandle>,
537 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
539 pub hidden_states: Option<Vec<TensorRef>>,
541 pub attention_weights: Option<Vec<TensorRef>>,
543}
544
545impl PrefillOutput {
546 pub fn new(logits: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
548 Self {
549 logits,
550 kv_cache,
551 recurrent_state: None,
552 hidden_states: None,
553 attention_weights: None,
554 }
555 }
556
557 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
559 self.recurrent_state = Some(recurrent_state);
560 self
561 }
562
563 pub fn last_token_logits(&self) -> Result<TensorRef> {
565 let shape = self.logits.shape();
566 if shape.len() != 3 {
567 return Err(ferrum_types::FerrumError::backend(
568 "Expected 3D logits tensor [batch, seq, vocab]",
569 ));
570 }
571
572 let seq_len = shape[1];
573 if seq_len == 0 {
574 return Err(ferrum_types::FerrumError::backend("Empty sequence"));
575 }
576
577 self.logits
579 .view(&[0, seq_len - 1, 0], &[shape[0], seq_len, shape[2]])
580 }
581}
582
583#[derive(Debug, Clone)]
585pub struct DecodeInput {
586 pub request_id: Option<RequestId>,
588 pub input_ids: TensorRef,
590 pub kv_cache: Arc<dyn KvCacheHandle>,
592 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
594 pub position_ids: Option<TensorRef>,
596 pub metadata: HashMap<String, serde_json::Value>,
598 pub logits_policy: LogitsReturnPolicy,
600}
601
602impl DecodeInput {
603 pub fn new(input_ids: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
605 Self {
606 request_id: None,
607 input_ids,
608 kv_cache,
609 recurrent_state: None,
610 position_ids: None,
611 metadata: HashMap::new(),
612 logits_policy: LogitsReturnPolicy::FullLogits,
613 }
614 }
615
616 pub fn with_request_id(mut self, request_id: RequestId) -> Self {
618 self.request_id = Some(request_id);
619 self
620 }
621
622 pub fn with_position_ids(mut self, positions: TensorRef) -> Self {
624 self.position_ids = Some(positions);
625 self
626 }
627
628 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
630 self.recurrent_state = Some(recurrent_state);
631 self
632 }
633
634 pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
636 self.metadata = metadata;
637 self
638 }
639
640 pub fn with_logits_policy(mut self, policy: LogitsReturnPolicy) -> Self {
641 self.logits_policy = policy;
642 self
643 }
644
645 pub fn batch_size(&self) -> usize {
647 self.input_ids.shape()[0]
648 }
649}
650
651#[derive(Clone)]
665pub struct UnifiedBatchItem {
666 pub seq_id: String,
668 pub q_tokens: Vec<u32>,
671 pub kv_cache: Arc<dyn KvCacheHandle>,
673 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
675 pub pos_offset: usize,
679 pub is_final_chunk: bool,
684 pub metadata: HashMap<String, serde_json::Value>,
686 pub logits_policy: LogitsReturnPolicy,
688}
689
690impl std::fmt::Debug for UnifiedBatchItem {
691 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
692 f.debug_struct("UnifiedBatchItem")
693 .field("seq_id", &self.seq_id)
694 .field("q_len", &self.q_tokens.len())
695 .field("has_recurrent_state", &self.recurrent_state.is_some())
696 .field("pos_offset", &self.pos_offset)
697 .field("is_final_chunk", &self.is_final_chunk)
698 .finish()
699 }
700}
701
702#[derive(Debug, Clone, Default)]
709pub struct UnifiedBatch {
710 pub items: Vec<UnifiedBatchItem>,
711}
712
713impl UnifiedBatch {
714 pub fn new() -> Self {
715 Self::default()
716 }
717
718 pub fn total_q_tokens(&self) -> usize {
721 self.items.iter().map(|it| it.q_tokens.len()).sum()
722 }
723
724 pub fn num_sampled_items(&self) -> usize {
727 self.items.iter().filter(|it| it.is_final_chunk).count()
728 }
729}
730
731#[derive(Debug, Clone)]
738pub struct PlanRuntimeDecodeInput {
739 pub request_id: RequestId,
740 pub input_token: TokenId,
741 pub kv_cache: Arc<dyn KvCacheHandle>,
742 pub logits_policy: LogitsReturnPolicy,
743}
744
745impl PlanRuntimeDecodeInput {
746 pub fn new(
747 request_id: RequestId,
748 input_token: TokenId,
749 kv_cache: Arc<dyn KvCacheHandle>,
750 ) -> Self {
751 Self {
752 request_id,
753 input_token,
754 kv_cache,
755 logits_policy: LogitsReturnPolicy::FullLogits,
756 }
757 }
758
759 pub fn with_logits_policy(mut self, logits_policy: LogitsReturnPolicy) -> Self {
760 self.logits_policy = logits_policy;
761 self
762 }
763}
764
765#[derive(Debug, Clone)]
772pub struct PlanRuntimePrefillInput {
773 pub request_id: RequestId,
774 pub input_tokens: Arc<[TokenId]>,
775 pub maximum_sequence_tokens: usize,
776 pub chunk: PrefillChunk,
777}
778
779impl PlanRuntimePrefillInput {
780 pub fn new(
781 request_id: RequestId,
782 input_tokens: impl Into<Arc<[TokenId]>>,
783 maximum_sequence_tokens: usize,
784 chunk: PrefillChunk,
785 ) -> Result<Self> {
786 let input_tokens = input_tokens.into();
787 if input_tokens.is_empty() {
788 return Err(FerrumError::request_validation(
789 "plan-runtime prefill requires at least one input token",
790 ));
791 }
792 if chunk.total_prompt_tokens() != input_tokens.len() {
793 return Err(FerrumError::request_validation(format!(
794 "plan-runtime prefill chunk declares {} prompt tokens for input length {}",
795 chunk.total_prompt_tokens(),
796 input_tokens.len()
797 )));
798 }
799 if maximum_sequence_tokens < input_tokens.len() {
800 return Err(FerrumError::request_validation(format!(
801 "plan-runtime sequence ceiling {maximum_sequence_tokens} does not cover prompt length {}",
802 input_tokens.len()
803 )));
804 }
805 Ok(Self {
806 request_id,
807 input_tokens,
808 maximum_sequence_tokens,
809 chunk,
810 })
811 }
812}
813
814#[derive(Debug, Clone)]
816pub struct DecodeOutput {
817 pub logits: TensorRef,
819 pub kv_cache: Arc<dyn KvCacheHandle>,
821 pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
823 pub hidden_state: Option<TensorRef>,
825 pub attention_weights: Option<Vec<TensorRef>>,
827}
828
829impl DecodeOutput {
830 pub fn new(logits: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
832 Self {
833 logits,
834 kv_cache,
835 recurrent_state: None,
836 hidden_state: None,
837 attention_weights: None,
838 }
839 }
840
841 pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
843 self.recurrent_state = Some(recurrent_state);
844 self
845 }
846}
847
848#[derive(Debug, Clone)]
850pub struct PlanRuntimeDecodeOutput {
851 pub sampling_output: ExecutorSamplingOutput,
852 pub kv_cache: Arc<dyn KvCacheHandle>,
853}
854
855impl PlanRuntimeDecodeOutput {
856 pub fn new(sampling_output: ExecutorSamplingOutput, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
857 Self {
858 sampling_output,
859 kv_cache,
860 }
861 }
862}
863
864#[derive(Debug)]
866pub enum PlanRuntimePrefillProduct {
867 Intermediate,
869 FinalLogits(Vec<f32>),
871}
872
873#[derive(Debug)]
875pub struct PlanRuntimePrefillAuthority {
876 request_id: RequestId,
877 committed_tokens: usize,
878 kv_cache: Arc<dyn KvCacheHandle>,
879}
880
881impl PlanRuntimePrefillAuthority {
882 pub fn request_id(&self) -> &RequestId {
883 &self.request_id
884 }
885
886 pub const fn committed_tokens(&self) -> usize {
887 self.committed_tokens
888 }
889
890 pub fn kv_cache(&self) -> &Arc<dyn KvCacheHandle> {
891 &self.kv_cache
892 }
893
894 pub fn into_cache(self) -> Arc<dyn KvCacheHandle> {
895 self.kv_cache
896 }
897}
898
899#[derive(Debug)]
901pub struct PlanRuntimePrefillOutput {
902 authority: PlanRuntimePrefillAuthority,
903 product: PlanRuntimePrefillProduct,
904}
905
906impl PlanRuntimePrefillOutput {
907 pub fn intermediate(
908 request_id: RequestId,
909 committed_tokens: usize,
910 kv_cache: Arc<dyn KvCacheHandle>,
911 ) -> Self {
912 Self {
913 authority: PlanRuntimePrefillAuthority {
914 request_id,
915 committed_tokens,
916 kv_cache,
917 },
918 product: PlanRuntimePrefillProduct::Intermediate,
919 }
920 }
921
922 pub fn final_logits(
923 request_id: RequestId,
924 committed_tokens: usize,
925 logits: Vec<f32>,
926 kv_cache: Arc<dyn KvCacheHandle>,
927 ) -> Result<Self> {
928 if logits.is_empty() {
929 return Err(FerrumError::backend(
930 "plan-runtime final prefill returned empty logits",
931 ));
932 }
933 Ok(Self {
934 authority: PlanRuntimePrefillAuthority {
935 request_id,
936 committed_tokens,
937 kv_cache,
938 },
939 product: PlanRuntimePrefillProduct::FinalLogits(logits),
940 })
941 }
942
943 pub fn request_id(&self) -> &RequestId {
944 self.authority.request_id()
945 }
946
947 pub const fn committed_tokens(&self) -> usize {
948 self.authority.committed_tokens()
949 }
950
951 pub fn product(&self) -> &PlanRuntimePrefillProduct {
952 &self.product
953 }
954
955 pub fn kv_cache(&self) -> &Arc<dyn KvCacheHandle> {
956 self.authority.kv_cache()
957 }
958
959 pub fn validate_for_completion(
960 &self,
961 expected_request_id: &RequestId,
962 completed_chunk: PrefillChunk,
963 vocabulary_size: usize,
964 ) -> Result<()> {
965 if self.request_id() != expected_request_id {
966 return Err(FerrumError::backend(format!(
967 "plan runtime returned prefill output for request {}, expected {expected_request_id}",
968 self.request_id()
969 )));
970 }
971 if self.committed_tokens() != completed_chunk.end() {
972 return Err(FerrumError::backend(format!(
973 "plan runtime returned prefill extent {}, expected {}",
974 self.committed_tokens(),
975 completed_chunk.end()
976 )));
977 }
978 if self.kv_cache().num_tokens() != self.committed_tokens() {
979 return Err(FerrumError::backend(format!(
980 "plan runtime prefill cache `{}` reports {} tokens for committed extent {}",
981 self.kv_cache().cache_id(),
982 self.kv_cache().num_tokens(),
983 self.committed_tokens()
984 )));
985 }
986 match (&self.product, completed_chunk.is_final()) {
987 (PlanRuntimePrefillProduct::Intermediate, false) => Ok(()),
988 (PlanRuntimePrefillProduct::FinalLogits(logits), true)
989 if logits.len() == vocabulary_size =>
990 {
991 Ok(())
992 }
993 (PlanRuntimePrefillProduct::FinalLogits(logits), true) => {
994 Err(FerrumError::backend(format!(
995 "plan runtime returned {} final prefill logits for vocabulary {vocabulary_size}",
996 logits.len()
997 )))
998 }
999 (PlanRuntimePrefillProduct::Intermediate, true) => Err(FerrumError::backend(
1000 "plan runtime returned an intermediate product for a final prefill chunk",
1001 )),
1002 (PlanRuntimePrefillProduct::FinalLogits(_), false) => Err(FerrumError::backend(
1003 "plan runtime returned final logits for an intermediate prefill chunk",
1004 )),
1005 }
1006 }
1007
1008 pub fn into_parts(self) -> (PlanRuntimePrefillAuthority, PlanRuntimePrefillProduct) {
1009 (self.authority, self.product)
1010 }
1011}
1012
1013#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1021pub struct ExecutorSequenceCompletion {
1022 request_id: RequestId,
1023 cache_id: String,
1024 input_tokens: u64,
1025 output_tokens: u64,
1026}
1027
1028impl ExecutorSequenceCompletion {
1029 pub fn new(
1030 request_id: RequestId,
1031 cache_id: String,
1032 input_tokens: usize,
1033 output_tokens: usize,
1034 ) -> Result<Self> {
1035 if cache_id.is_empty() {
1036 return Err(FerrumError::request_validation(
1037 "executor sequence completion requires a cache identity",
1038 ));
1039 }
1040 let input_tokens = u64::try_from(input_tokens).map_err(|_| {
1041 FerrumError::request_validation("executor completion input token count exceeds u64")
1042 })?;
1043 let output_tokens = u64::try_from(output_tokens).map_err(|_| {
1044 FerrumError::request_validation("executor completion output token count exceeds u64")
1045 })?;
1046 Ok(Self {
1047 request_id,
1048 cache_id,
1049 input_tokens,
1050 output_tokens,
1051 })
1052 }
1053
1054 pub fn request_id(&self) -> &RequestId {
1055 &self.request_id
1056 }
1057
1058 pub fn cache_id(&self) -> &str {
1059 &self.cache_id
1060 }
1061
1062 pub const fn input_tokens(&self) -> u64 {
1063 self.input_tokens
1064 }
1065
1066 pub const fn output_tokens(&self) -> u64 {
1067 self.output_tokens
1068 }
1069}
1070
1071pub use ferrum_types::ExecutionResourceAuthority;
1072
1073#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1079pub struct ExecutorExecutionCapacityPreemption {
1080 request_id: RequestId,
1081 cache_id: String,
1082}
1083
1084impl ExecutorExecutionCapacityPreemption {
1085 pub fn new(request_id: RequestId, cache_id: String) -> Result<Self> {
1086 if cache_id.is_empty() {
1087 return Err(FerrumError::request_validation(
1088 "execution-capacity preemption requires a cache identity",
1089 ));
1090 }
1091 Ok(Self {
1092 request_id,
1093 cache_id,
1094 })
1095 }
1096
1097 pub fn request_id(&self) -> &RequestId {
1098 &self.request_id
1099 }
1100
1101 pub fn cache_id(&self) -> &str {
1102 &self.cache_id
1103 }
1104}
1105
1106#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
1107#[serde(rename_all = "snake_case")]
1108pub enum ExecutorExecutionCapacityPreemptionAuthority {
1109 RetainedPrefill,
1110 ActiveSequence,
1111}
1112
1113#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1117pub struct ExecutorExecutionCapacityPreemptionReceipt {
1118 request_id: RequestId,
1119 cache_id: String,
1120 authority: ExecutorExecutionCapacityPreemptionAuthority,
1121}
1122
1123impl ExecutorExecutionCapacityPreemptionReceipt {
1124 pub fn new(
1125 request_id: RequestId,
1126 cache_id: String,
1127 authority: ExecutorExecutionCapacityPreemptionAuthority,
1128 ) -> Self {
1129 Self {
1130 request_id,
1131 cache_id,
1132 authority,
1133 }
1134 }
1135
1136 pub fn request_id(&self) -> &RequestId {
1137 &self.request_id
1138 }
1139
1140 pub fn cache_id(&self) -> &str {
1141 &self.cache_id
1142 }
1143
1144 pub const fn authority(&self) -> ExecutorExecutionCapacityPreemptionAuthority {
1145 self.authority
1146 }
1147}
1148
1149#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1157pub struct PlanRuntimeResourceSnapshot {
1158 device_capacity_bytes: u64,
1159 usable_capacity_bytes: u64,
1160 process_claimed_bytes: u64,
1161 plan_claimed_bytes: u64,
1162 static_bytes: u64,
1163 dynamic_resident_bytes: u64,
1164 dynamic_free_bytes: u64,
1165 pending_growth_bytes: u64,
1166 quarantined_bytes: u64,
1167}
1168
1169impl PlanRuntimeResourceSnapshot {
1170 #[allow(clippy::too_many_arguments)]
1171 pub fn new(
1172 device_capacity_bytes: u64,
1173 usable_capacity_bytes: u64,
1174 process_claimed_bytes: u64,
1175 plan_claimed_bytes: u64,
1176 static_bytes: u64,
1177 dynamic_resident_bytes: u64,
1178 dynamic_free_bytes: u64,
1179 pending_growth_bytes: u64,
1180 quarantined_bytes: u64,
1181 ) -> Result<Self> {
1182 let snapshot = Self {
1183 device_capacity_bytes,
1184 usable_capacity_bytes,
1185 process_claimed_bytes,
1186 plan_claimed_bytes,
1187 static_bytes,
1188 dynamic_resident_bytes,
1189 dynamic_free_bytes,
1190 pending_growth_bytes,
1191 quarantined_bytes,
1192 };
1193 snapshot.validate()?;
1194 Ok(snapshot)
1195 }
1196
1197 pub fn validate(&self) -> Result<()> {
1200 if self.usable_capacity_bytes > self.device_capacity_bytes {
1201 return Err(ferrum_types::FerrumError::internal(format!(
1202 "plan runtime usable capacity {} exceeds device capacity {}",
1203 self.usable_capacity_bytes, self.device_capacity_bytes
1204 )));
1205 }
1206 if self.process_claimed_bytes > self.usable_capacity_bytes {
1207 return Err(ferrum_types::FerrumError::internal(format!(
1208 "plan runtime process claims {} exceed usable capacity {}",
1209 self.process_claimed_bytes, self.usable_capacity_bytes
1210 )));
1211 }
1212 if self.plan_claimed_bytes > self.process_claimed_bytes {
1213 return Err(ferrum_types::FerrumError::internal(format!(
1214 "plan runtime plan claims {} exceed process claims {}",
1215 self.plan_claimed_bytes, self.process_claimed_bytes
1216 )));
1217 }
1218 if self.dynamic_free_bytes > self.dynamic_resident_bytes {
1219 return Err(ferrum_types::FerrumError::internal(format!(
1220 "plan runtime dynamic free bytes {} exceed resident bytes {}",
1221 self.dynamic_free_bytes, self.dynamic_resident_bytes
1222 )));
1223 }
1224 let minimum_plan_claim = self
1225 .static_bytes
1226 .checked_add(self.dynamic_resident_bytes)
1227 .and_then(|bytes| bytes.checked_add(self.quarantined_bytes))
1228 .ok_or_else(|| {
1229 ferrum_types::FerrumError::internal(
1230 "plan runtime static, resident, and quarantined bytes overflow u64",
1231 )
1232 })?;
1233 if minimum_plan_claim > self.plan_claimed_bytes {
1234 return Err(ferrum_types::FerrumError::internal(format!(
1235 "plan runtime accounted plan bytes {minimum_plan_claim} exceed plan claims {}",
1236 self.plan_claimed_bytes
1237 )));
1238 }
1239 Ok(())
1240 }
1241
1242 pub const fn device_capacity_bytes(&self) -> u64 {
1243 self.device_capacity_bytes
1244 }
1245
1246 pub const fn usable_capacity_bytes(&self) -> u64 {
1247 self.usable_capacity_bytes
1248 }
1249
1250 pub const fn process_claimed_bytes(&self) -> u64 {
1251 self.process_claimed_bytes
1252 }
1253
1254 pub const fn plan_claimed_bytes(&self) -> u64 {
1255 self.plan_claimed_bytes
1256 }
1257
1258 pub const fn static_bytes(&self) -> u64 {
1259 self.static_bytes
1260 }
1261
1262 pub const fn dynamic_resident_bytes(&self) -> u64 {
1263 self.dynamic_resident_bytes
1264 }
1265
1266 pub const fn dynamic_free_bytes(&self) -> u64 {
1267 self.dynamic_free_bytes
1268 }
1269
1270 pub const fn dynamic_used_bytes(&self) -> u64 {
1271 self.dynamic_resident_bytes - self.dynamic_free_bytes
1272 }
1273
1274 pub const fn pending_growth_bytes(&self) -> u64 {
1275 self.pending_growth_bytes
1276 }
1277
1278 pub const fn quarantined_bytes(&self) -> u64 {
1279 self.quarantined_bytes
1280 }
1281
1282 pub fn available_bytes(&self) -> Result<u64> {
1286 self.usable_capacity_bytes
1287 .checked_sub(self.process_claimed_bytes)
1288 .and_then(|bytes| bytes.checked_add(self.dynamic_free_bytes))
1289 .ok_or_else(|| {
1290 ferrum_types::FerrumError::internal(
1291 "plan runtime available capacity calculation overflowed",
1292 )
1293 })
1294 }
1295
1296 pub fn used_bytes(&self) -> Result<u64> {
1297 self.available_bytes().and_then(|available| {
1298 self.usable_capacity_bytes
1299 .checked_sub(available)
1300 .ok_or_else(|| {
1301 ferrum_types::FerrumError::internal(
1302 "plan runtime available bytes exceed usable capacity",
1303 )
1304 })
1305 })
1306 }
1307}
1308
1309#[cfg(test)]
1310mod plan_runtime_resource_snapshot_tests {
1311 use super::PlanRuntimeResourceSnapshot;
1312
1313 #[test]
1314 fn separates_static_and_dynamic_usage() {
1315 let snapshot =
1316 PlanRuntimeResourceSnapshot::new(1_000, 900, 710, 710, 400, 300, 200, 20, 10).unwrap();
1317
1318 assert_eq!(snapshot.available_bytes().unwrap(), 390);
1319 assert_eq!(snapshot.used_bytes().unwrap(), 510);
1320 assert_eq!(snapshot.dynamic_resident_bytes(), 300);
1321 assert_eq!(snapshot.dynamic_used_bytes(), 100);
1322 assert_eq!(snapshot.dynamic_free_bytes(), 200);
1323 assert_eq!(snapshot.pending_growth_bytes(), 20);
1324 assert_eq!(snapshot.quarantined_bytes(), 10);
1325 }
1326
1327 #[test]
1328 fn rejects_incoherent_capacity_evidence() {
1329 assert!(PlanRuntimeResourceSnapshot::new(1_000, 1_001, 0, 0, 0, 0, 0, 0, 0).is_err());
1330 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 901, 0, 0, 0, 0, 0, 0).is_err());
1331 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 500, 501, 0, 0, 0, 0, 0).is_err());
1332 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 100, 100, 0, 100, 101, 0, 0).is_err());
1333 assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 500, 500, 400, 100, 0, 0, 1).is_err());
1334 }
1335}
1336
1337#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1339pub enum ExecutorRequestOrigin {
1340 Product,
1341 Startup,
1342 Diagnostic,
1343}
1344
1345impl ExecutorRequestOrigin {
1346 pub const fn namespace(self) -> &'static str {
1347 match self {
1348 Self::Product => "product",
1349 Self::Startup => "startup",
1350 Self::Diagnostic => "diagnostic",
1351 }
1352 }
1353
1354 pub fn from_namespaced_request_identity(identity: &str) -> Option<Self> {
1355 let suffix = identity.strip_prefix("request.")?;
1356 let (namespace, request_id) = suffix.split_once('.')?;
1357 if request_id.is_empty() {
1358 return None;
1359 }
1360 match namespace {
1361 "product" => Some(Self::Product),
1362 "startup" => Some(Self::Startup),
1363 "diagnostic" => Some(Self::Diagnostic),
1364 _ => None,
1365 }
1366 }
1367}
1368
1369#[derive(Debug, Clone, Copy)]
1378pub struct ExecutorPrefillAdmission<'a> {
1379 pub request_id: &'a RequestId,
1380 pub input_tokens: &'a [TokenId],
1381 pub maximum_sequence_tokens: usize,
1382 pub product_prompt_tokens: usize,
1384 pub replayed_output_tokens: usize,
1386 pub request_origin: ExecutorRequestOrigin,
1387}
1388
1389impl<'a> ExecutorPrefillAdmission<'a> {
1390 pub const fn for_startup(
1391 request_id: &'a RequestId,
1392 input_tokens: &'a [TokenId],
1393 maximum_sequence_tokens: usize,
1394 ) -> Self {
1395 Self {
1396 request_id,
1397 input_tokens,
1398 maximum_sequence_tokens,
1399 product_prompt_tokens: input_tokens.len(),
1400 replayed_output_tokens: 0,
1401 request_origin: ExecutorRequestOrigin::Startup,
1402 }
1403 }
1404
1405 pub const fn for_diagnostic(
1406 request_id: &'a RequestId,
1407 input_tokens: &'a [TokenId],
1408 maximum_sequence_tokens: usize,
1409 ) -> Self {
1410 Self {
1411 request_id,
1412 input_tokens,
1413 maximum_sequence_tokens,
1414 product_prompt_tokens: input_tokens.len(),
1415 replayed_output_tokens: 0,
1416 request_origin: ExecutorRequestOrigin::Diagnostic,
1417 }
1418 }
1419
1420 pub fn for_product_request(
1423 request_id: &'a RequestId,
1424 input_tokens: &'a [TokenId],
1425 maximum_sequence_tokens: usize,
1426 product_prompt_tokens: usize,
1427 replayed_output_tokens: usize,
1428 ) -> Result<Self> {
1429 let admission = Self {
1430 request_id,
1431 input_tokens,
1432 maximum_sequence_tokens,
1433 product_prompt_tokens,
1434 replayed_output_tokens,
1435 request_origin: ExecutorRequestOrigin::Product,
1436 };
1437 admission.validate()?;
1438 Ok(admission)
1439 }
1440
1441 pub fn validate(&self) -> Result<()> {
1442 if self.input_tokens.is_empty() {
1443 return Err(FerrumError::request_validation(
1444 "executor prefill admission requires at least one execution-context token",
1445 ));
1446 }
1447 if self.product_prompt_tokens == 0 {
1448 return Err(FerrumError::request_validation(
1449 "executor prefill admission requires at least one product prompt token",
1450 ));
1451 }
1452 let execution_context_tokens = self
1453 .product_prompt_tokens
1454 .checked_add(self.replayed_output_tokens)
1455 .ok_or_else(|| {
1456 FerrumError::request_validation(
1457 "executor prefill product token accounting exceeds usize",
1458 )
1459 })?;
1460 if execution_context_tokens != self.input_tokens.len() {
1461 return Err(FerrumError::request_validation(format!(
1462 "executor prefill execution context has {} tokens but product accounting declares {} prompt + {} replayed output",
1463 self.input_tokens.len(),
1464 self.product_prompt_tokens,
1465 self.replayed_output_tokens
1466 )));
1467 }
1468 if self.maximum_sequence_tokens < execution_context_tokens {
1469 return Err(FerrumError::request_validation(format!(
1470 "executor prefill sequence ceiling {} does not cover execution context {execution_context_tokens}",
1471 self.maximum_sequence_tokens
1472 )));
1473 }
1474 Ok(())
1475 }
1476}
1477
1478#[cfg(test)]
1479mod executor_prefill_admission_tests {
1480 use super::{ExecutorPrefillAdmission, ExecutorRequestOrigin};
1481 use ferrum_types::{RequestId, TokenId};
1482
1483 #[test]
1484 fn product_accounting_distinguishes_replayed_output_from_prompt() {
1485 let request_id = RequestId::new();
1486 let tokens = [1, 2, 3, 4, 5]
1487 .into_iter()
1488 .map(TokenId::new)
1489 .collect::<Vec<_>>();
1490
1491 let admission =
1492 ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 8, 3, 2)
1493 .expect("recompute accounting must be accepted");
1494
1495 assert_eq!(admission.product_prompt_tokens, 3);
1496 assert_eq!(admission.replayed_output_tokens, 2);
1497 assert_eq!(admission.request_origin, ExecutorRequestOrigin::Product);
1498 assert_eq!(
1499 ExecutorPrefillAdmission::for_startup(&request_id, &tokens, 8).request_origin,
1500 ExecutorRequestOrigin::Startup
1501 );
1502 assert_eq!(
1503 ExecutorPrefillAdmission::for_diagnostic(&request_id, &tokens, 8).request_origin,
1504 ExecutorRequestOrigin::Diagnostic
1505 );
1506 assert_eq!(ExecutorRequestOrigin::Product.namespace(), "product");
1507 assert_eq!(ExecutorRequestOrigin::Startup.namespace(), "startup");
1508 assert_eq!(ExecutorRequestOrigin::Diagnostic.namespace(), "diagnostic");
1509 assert_eq!(
1510 ExecutorRequestOrigin::from_namespaced_request_identity("request.product.123"),
1511 Some(ExecutorRequestOrigin::Product)
1512 );
1513 assert_eq!(
1514 ExecutorRequestOrigin::from_namespaced_request_identity("request.startup.123"),
1515 Some(ExecutorRequestOrigin::Startup)
1516 );
1517 assert_eq!(
1518 ExecutorRequestOrigin::from_namespaced_request_identity("request.diagnostic.123"),
1519 Some(ExecutorRequestOrigin::Diagnostic)
1520 );
1521 assert_eq!(
1522 ExecutorRequestOrigin::from_namespaced_request_identity("request.product."),
1523 None
1524 );
1525 assert_eq!(
1526 ExecutorRequestOrigin::from_namespaced_request_identity("request/external"),
1527 None
1528 );
1529 }
1530
1531 #[test]
1532 fn product_accounting_rejects_context_drift_and_short_ceiling() {
1533 let request_id = RequestId::new();
1534 let tokens = [1, 2, 3].into_iter().map(TokenId::new).collect::<Vec<_>>();
1535
1536 assert!(
1537 ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 3, 2, 0).is_err()
1538 );
1539 assert!(
1540 ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 2, 2, 1).is_err()
1541 );
1542 }
1543}
1544
1545#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1548pub struct ExecutorPrefillAdmissionReceipt {
1549 pub request_id: RequestId,
1550}
1551
1552#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1554pub struct ExecutorAdmissionEpochs {
1555 pub coordinator_id: NonZeroU64,
1556 pub release_epoch: u64,
1557 pub capacity_epoch: u64,
1558}
1559
1560impl ExecutorAdmissionEpochs {
1561 pub const fn new(coordinator_id: NonZeroU64, release_epoch: u64, capacity_epoch: u64) -> Self {
1562 Self {
1563 coordinator_id,
1564 release_epoch,
1565 capacity_epoch,
1566 }
1567 }
1568
1569 pub fn from_capacity(epochs: crate::vnext::CapacityEpochs) -> Self {
1570 Self::new(
1571 NonZeroU64::new(epochs.coordinator_id().get())
1572 .expect("core-issued admission coordinator ids are non-zero"),
1573 epochs.release_epoch(),
1574 epochs.capacity_epoch(),
1575 )
1576 }
1577}
1578
1579type ExecutorCapacityWaitFuture =
1580 Pin<Box<dyn Future<Output = Result<ExecutorAdmissionEpochs>> + Send + 'static>>;
1581
1582#[must_use = "capacity wait registrations must be awaited or explicitly dropped"]
1589pub struct ExecutorCapacityWaitRegistration {
1590 future: ExecutorCapacityWaitFuture,
1591}
1592
1593impl ExecutorCapacityWaitRegistration {
1594 pub fn new<F>(future: F) -> Self
1595 where
1596 F: Future<Output = Result<ExecutorAdmissionEpochs>> + Send + 'static,
1597 {
1598 Self {
1599 future: Box::pin(future),
1600 }
1601 }
1602
1603 pub async fn wait_for_change(self) -> Result<ExecutorAdmissionEpochs> {
1604 self.future.await
1605 }
1606}
1607
1608#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1610#[serde(rename_all = "snake_case")]
1611pub enum ExecutorExecutionCapacityStage {
1612 SequenceExtension,
1613 StepAdmission,
1614 SubmissionWave,
1615}
1616
1617#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1620pub struct ExecutorExecutionMaintenanceMutation {
1621 pool_id: crate::vnext::DynamicBackingPoolId,
1622 domain_id: crate::vnext::CapacityDomainId,
1623 chunk: crate::vnext::BackingChunkIdentity,
1624 chunk_bytes: u64,
1625 published_capacity_bytes: u64,
1626 capacity_epoch: u64,
1627}
1628
1629impl ExecutorExecutionMaintenanceMutation {
1630 pub fn pool_id(&self) -> &crate::vnext::DynamicBackingPoolId {
1631 &self.pool_id
1632 }
1633
1634 pub const fn domain_id(&self) -> crate::vnext::CapacityDomainId {
1635 self.domain_id
1636 }
1637
1638 pub fn chunk(&self) -> &crate::vnext::BackingChunkIdentity {
1639 &self.chunk
1640 }
1641
1642 pub const fn chunk_bytes(&self) -> u64 {
1643 self.chunk_bytes
1644 }
1645
1646 pub const fn published_capacity_bytes(&self) -> u64 {
1647 self.published_capacity_bytes
1648 }
1649
1650 pub const fn capacity_epoch(&self) -> u64 {
1651 self.capacity_epoch
1652 }
1653}
1654
1655#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1663pub struct ExecutorExecutionMaintenanceProgress {
1664 attempts: u32,
1665 coordinator_id: NonZeroU64,
1666 mutations: Vec<ExecutorExecutionMaintenanceMutation>,
1667 latest_capacity_epoch: u64,
1668}
1669
1670impl ExecutorExecutionMaintenanceProgress {
1671 pub fn from_growth_receipts(
1672 attempts: u32,
1673 observed: ExecutorAdmissionEpochs,
1674 receipts: &[crate::vnext::DynamicPoolGrowthBatchReceipt],
1675 pools: &[crate::vnext::DynamicPoolStatus],
1676 ) -> Result<Self> {
1677 if attempts == 0 || receipts.is_empty() || receipts.len() > attempts as usize {
1678 return Err(FerrumError::internal(
1679 "execution maintenance retry requires bounded, non-empty growth receipts",
1680 ));
1681 }
1682
1683 let mut mutations = Vec::new();
1684 let mut previous_capacity_epoch = None;
1685 for receipt in receipts {
1686 if receipt.coordinator_id().get() != observed.coordinator_id.get() {
1687 return Err(FerrumError::internal(
1688 "execution maintenance receipt belongs to another capacity coordinator",
1689 ));
1690 }
1691 if receipt.growths().is_empty()
1692 || previous_capacity_epoch
1693 .is_some_and(|previous| receipt.capacity_epoch() <= previous)
1694 {
1695 return Err(FerrumError::internal(
1696 "execution maintenance receipts contain no new ordered capacity mutation",
1697 ));
1698 }
1699 previous_capacity_epoch = Some(receipt.capacity_epoch());
1700
1701 for growth in receipt.growths() {
1702 let pool = pools
1703 .iter()
1704 .find(|pool| pool.pool_id() == growth.pool_id())
1705 .ok_or_else(|| {
1706 FerrumError::internal(
1707 "execution maintenance receipt references an unknown dynamic pool",
1708 )
1709 })?;
1710 if growth.chunk().pool_id() != growth.pool_id()
1711 || growth.chunk_bytes() == 0
1712 || growth.published_capacity_bytes() == 0
1713 || growth.capacity_epoch() != receipt.capacity_epoch()
1714 {
1715 return Err(FerrumError::internal(
1716 "execution maintenance receipt contains an invalid pool mutation",
1717 ));
1718 }
1719 if mutations
1720 .iter()
1721 .any(|mutation: &ExecutorExecutionMaintenanceMutation| {
1722 mutation.pool_id() == growth.pool_id() && mutation.chunk() == growth.chunk()
1723 })
1724 {
1725 return Err(FerrumError::internal(
1726 "execution maintenance receipts repeat one physical pool mutation",
1727 ));
1728 }
1729 mutations.push(ExecutorExecutionMaintenanceMutation {
1730 pool_id: growth.pool_id().clone(),
1731 domain_id: pool.domain_id(),
1732 chunk: growth.chunk().clone(),
1733 chunk_bytes: growth.chunk_bytes(),
1734 published_capacity_bytes: growth.published_capacity_bytes(),
1735 capacity_epoch: growth.capacity_epoch(),
1736 });
1737 }
1738 }
1739
1740 let latest_capacity_epoch = previous_capacity_epoch.expect("receipts are non-empty");
1741 if latest_capacity_epoch > observed.capacity_epoch {
1742 return Err(FerrumError::internal(
1743 "execution maintenance receipt is newer than the exported capacity observation",
1744 ));
1745 }
1746 mutations.sort_by(|left, right| {
1747 (
1748 left.capacity_epoch,
1749 left.pool_id.as_str(),
1750 left.chunk.ordinal(),
1751 left.chunk.generation(),
1752 )
1753 .cmp(&(
1754 right.capacity_epoch,
1755 right.pool_id.as_str(),
1756 right.chunk.ordinal(),
1757 right.chunk.generation(),
1758 ))
1759 });
1760 Ok(Self {
1761 attempts,
1762 coordinator_id: observed.coordinator_id,
1763 mutations,
1764 latest_capacity_epoch,
1765 })
1766 }
1767
1768 pub const fn attempts(&self) -> u32 {
1769 self.attempts
1770 }
1771
1772 pub const fn coordinator_id(&self) -> NonZeroU64 {
1773 self.coordinator_id
1774 }
1775
1776 pub fn mutations(&self) -> &[ExecutorExecutionMaintenanceMutation] {
1777 &self.mutations
1778 }
1779
1780 pub const fn latest_capacity_epoch(&self) -> u64 {
1781 self.latest_capacity_epoch
1782 }
1783}
1784
1785#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1791pub struct ExecutorExecutionMaintenanceRetry {
1792 affected_request_ids: Vec<RequestId>,
1793 progress: ExecutorExecutionMaintenanceProgress,
1794}
1795
1796impl ExecutorExecutionMaintenanceRetry {
1797 fn new(
1798 affected_request_ids: Vec<RequestId>,
1799 progress: ExecutorExecutionMaintenanceProgress,
1800 ) -> Result<Self> {
1801 let unique = affected_request_ids.iter().collect::<HashSet<_>>();
1802 if affected_request_ids.is_empty() || unique.len() != affected_request_ids.len() {
1803 return Err(FerrumError::internal(
1804 "execution maintenance retry requires unique affected requests",
1805 ));
1806 }
1807 if progress.mutations().is_empty() {
1808 return Err(FerrumError::internal(
1809 "execution maintenance retry requires physical mutations",
1810 ));
1811 }
1812 Ok(Self {
1813 affected_request_ids,
1814 progress,
1815 })
1816 }
1817
1818 pub fn affected_request_ids(&self) -> &[RequestId] {
1819 &self.affected_request_ids
1820 }
1821
1822 pub const fn progress(&self) -> &ExecutorExecutionMaintenanceProgress {
1823 &self.progress
1824 }
1825}
1826
1827#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1834#[serde(rename_all = "snake_case")]
1835pub enum ExecutorExecutionCapacityEvidenceOwner {
1836 Logical,
1837 Backing,
1838}
1839
1840#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1841#[serde(tag = "kind", rename_all = "snake_case")]
1842enum ExecutorExecutionCapacityEvidenceKind {
1843 Logical {
1844 shortfalls: Vec<crate::vnext::CapacityShortfall>,
1845 #[serde(skip_serializing_if = "Option::is_none")]
1846 pressure: Option<crate::vnext::DynamicBackingPressure>,
1847 },
1848 BackingDeferred {
1849 blockers: Vec<crate::vnext::DynamicBackingBlocker>,
1850 #[serde(skip_serializing_if = "Option::is_none")]
1851 pressure: Option<crate::vnext::DynamicBackingPressure>,
1852 },
1853 BackingPressure {
1854 pressure: crate::vnext::DynamicBackingPressure,
1855 },
1856}
1857
1858#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1863pub struct ExecutorExecutionCapacityEvidence {
1864 owner: ExecutorExecutionCapacityEvidenceOwner,
1865 #[serde(flatten)]
1866 kind: ExecutorExecutionCapacityEvidenceKind,
1867 #[serde(skip_serializing_if = "Option::is_none")]
1868 maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
1869}
1870
1871impl ExecutorExecutionCapacityEvidence {
1872 fn logical(shortfalls: Vec<crate::vnext::CapacityShortfall>) -> Result<Self> {
1873 Self::logical_with_pressure(shortfalls, None)
1874 }
1875
1876 fn logical_with_pressure(
1877 shortfalls: Vec<crate::vnext::CapacityShortfall>,
1878 pressure: Option<crate::vnext::DynamicBackingPressure>,
1879 ) -> Result<Self> {
1880 if shortfalls.is_empty() {
1881 return Err(FerrumError::internal(
1882 "logical execution deferral requires at least one shortfall",
1883 ));
1884 }
1885 Ok(Self {
1886 owner: ExecutorExecutionCapacityEvidenceOwner::Logical,
1887 kind: ExecutorExecutionCapacityEvidenceKind::Logical {
1888 shortfalls,
1889 pressure,
1890 },
1891 maintenance_boundary: None,
1892 })
1893 }
1894
1895 fn backing_deferred(blockers: Vec<crate::vnext::DynamicBackingBlocker>) -> Result<Self> {
1896 Self::backing_deferred_with_pressure(blockers, None)
1897 }
1898
1899 fn backing_deferred_with_pressure(
1900 blockers: Vec<crate::vnext::DynamicBackingBlocker>,
1901 pressure: Option<crate::vnext::DynamicBackingPressure>,
1902 ) -> Result<Self> {
1903 if blockers.is_empty() {
1904 return Err(FerrumError::internal(
1905 "physical execution deferral requires at least one backing blocker",
1906 ));
1907 }
1908 Ok(Self {
1909 owner: ExecutorExecutionCapacityEvidenceOwner::Backing,
1910 kind: ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, pressure },
1911 maintenance_boundary: None,
1912 })
1913 }
1914
1915 fn direct_backing_pressure(pressure: crate::vnext::DynamicBackingPressure) -> Self {
1916 Self {
1917 owner: ExecutorExecutionCapacityEvidenceOwner::Backing,
1918 kind: ExecutorExecutionCapacityEvidenceKind::BackingPressure { pressure },
1919 maintenance_boundary: None,
1920 }
1921 }
1922
1923 pub const fn owner(&self) -> ExecutorExecutionCapacityEvidenceOwner {
1924 self.owner
1925 }
1926
1927 pub fn shortfalls(&self) -> &[crate::vnext::CapacityShortfall] {
1928 match &self.kind {
1929 ExecutorExecutionCapacityEvidenceKind::Logical { shortfalls, .. } => shortfalls,
1930 ExecutorExecutionCapacityEvidenceKind::BackingDeferred { .. }
1931 | ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => &[],
1932 }
1933 }
1934
1935 pub fn backing_blockers(&self) -> &[crate::vnext::DynamicBackingBlocker] {
1936 match &self.kind {
1937 ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, .. } => blockers,
1938 ExecutorExecutionCapacityEvidenceKind::Logical { .. }
1939 | ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => &[],
1940 }
1941 }
1942
1943 pub const fn backing_pressure(&self) -> Option<&crate::vnext::DynamicBackingPressure> {
1944 match &self.kind {
1945 ExecutorExecutionCapacityEvidenceKind::Logical { pressure, .. }
1946 | ExecutorExecutionCapacityEvidenceKind::BackingDeferred { pressure, .. } => {
1947 pressure.as_ref()
1948 }
1949 ExecutorExecutionCapacityEvidenceKind::BackingPressure { pressure } => Some(pressure),
1950 }
1951 }
1952
1953 pub const fn maintenance_boundary(
1954 &self,
1955 ) -> Option<&crate::vnext::DynamicPoolMaintenanceBoundaryReceipt> {
1956 self.maintenance_boundary.as_ref()
1957 }
1958
1959 fn with_maintenance_boundary(
1960 mut self,
1961 boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
1962 ) -> Result<Self> {
1963 match (self.backing_pressure(), boundary.as_ref()) {
1964 (
1965 Some(crate::vnext::DynamicBackingPressure::DeviceCapacity(pressure)),
1966 Some(boundary),
1967 ) if pressure == boundary.pressure() && !boundary.reclaim_sufficient() => {}
1968 (Some(crate::vnext::DynamicBackingPressure::PoolResident(_)), None) | (None, None) => {}
1969 (Some(crate::vnext::DynamicBackingPressure::DeviceCapacity(_)), None) => {
1970 return Err(FerrumError::internal(
1971 "device-capacity execution maintenance lost its boundary receipt",
1972 ));
1973 }
1974 _ => {
1975 return Err(FerrumError::internal(
1976 "execution maintenance boundary differs from its blocked pressure",
1977 ));
1978 }
1979 }
1980 self.maintenance_boundary = boundary;
1981 Ok(self)
1982 }
1983
1984 fn has_relevant_mutation(&self, mutation: &ExecutorExecutionMaintenanceMutation) -> bool {
1985 let logical_matches = |shortfalls: &[crate::vnext::CapacityShortfall]| {
1986 shortfalls.iter().any(|shortfall| {
1987 shortfall.kind() == crate::vnext::CapacityShortfallKind::BackingGrowthRequired
1988 && shortfall.domain() == Some(mutation.domain_id())
1989 })
1990 };
1991 let backing_matches = |blockers: &[crate::vnext::DynamicBackingBlocker]| {
1992 blockers.iter().any(|blocker| {
1993 blocker.pool_id() == mutation.pool_id()
1994 && blocker.domain_id() == mutation.domain_id()
1995 })
1996 };
1997 match &self.kind {
1998 ExecutorExecutionCapacityEvidenceKind::Logical { shortfalls, .. } => {
1999 logical_matches(shortfalls)
2000 }
2001 ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, .. } => {
2002 backing_matches(blockers)
2003 }
2004 ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => false,
2005 }
2006 }
2007}
2008
2009#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2010pub struct ExecutorExecutionCapacityDeferral {
2011 observed: ExecutorAdmissionEpochs,
2012 wait_condition: crate::vnext::CapacityWaitCondition,
2013 stage: ExecutorExecutionCapacityStage,
2014 evidence: ExecutorExecutionCapacityEvidence,
2015 #[serde(skip_serializing_if = "Option::is_none")]
2016 maintenance_retry: Option<ExecutorExecutionMaintenanceRetry>,
2017}
2018
2019impl ExecutorExecutionCapacityDeferral {
2020 fn with_evidence(
2021 observed: ExecutorAdmissionEpochs,
2022 wait_condition: crate::vnext::CapacityWaitCondition,
2023 stage: ExecutorExecutionCapacityStage,
2024 evidence: ExecutorExecutionCapacityEvidence,
2025 ) -> Result<Self> {
2026 if wait_condition.coordinator_id().get() != observed.coordinator_id.get() {
2027 return Err(ferrum_types::FerrumError::request_validation(
2028 "executor execution deferral belongs to a different capacity coordinator",
2029 ));
2030 }
2031 Ok(Self {
2032 observed,
2033 wait_condition,
2034 stage,
2035 evidence,
2036 maintenance_retry: None,
2037 })
2038 }
2039
2040 pub fn from_backing_pressure(
2043 observed: ExecutorAdmissionEpochs,
2044 wait_condition: crate::vnext::CapacityWaitCondition,
2045 pressure: crate::vnext::DynamicBackingPressure,
2046 stage: ExecutorExecutionCapacityStage,
2047 ) -> Result<Self> {
2048 Self::with_evidence(
2049 observed,
2050 wait_condition,
2051 stage,
2052 ExecutorExecutionCapacityEvidence::direct_backing_pressure(pressure),
2053 )
2054 }
2055
2056 pub fn from_admission(
2057 deferred: &crate::vnext::AdmissionDeferred,
2058 stage: ExecutorExecutionCapacityStage,
2059 ) -> Result<Self> {
2060 if deferred.action() != crate::vnext::DeferredAction::WaitForRelease {
2061 return Err(ferrum_types::FerrumError::internal(
2062 "execution capacity deferral must be reduced to WaitForRelease before export",
2063 ));
2064 }
2065 let evidence = ExecutorExecutionCapacityEvidence::logical(deferred.blockers().to_vec())?;
2066 Self::with_evidence(
2067 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2068 deferred.wait_condition().clone(),
2069 stage,
2070 evidence,
2071 )
2072 }
2073
2074 pub fn from_pending_maintenance(
2081 deferred: &crate::vnext::AdmissionDeferred,
2082 stage: ExecutorExecutionCapacityStage,
2083 ) -> Result<Self> {
2084 if deferred.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2085 return Err(ferrum_types::FerrumError::internal(
2086 "pending execution maintenance must await backing growth",
2087 ));
2088 }
2089 let evidence = ExecutorExecutionCapacityEvidence::logical(deferred.blockers().to_vec())?;
2090 Self::with_evidence(
2091 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2092 deferred.wait_condition().clone(),
2093 stage,
2094 evidence,
2095 )
2096 }
2097
2098 pub fn from_backing(
2101 deferred: &crate::vnext::DynamicBackingDeferred,
2102 stage: ExecutorExecutionCapacityStage,
2103 ) -> Result<Self> {
2104 let evidence =
2105 ExecutorExecutionCapacityEvidence::backing_deferred(deferred.blockers().to_vec())?;
2106 Self::with_evidence(
2107 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2108 deferred.wait_condition().clone(),
2109 stage,
2110 evidence,
2111 )
2112 }
2113
2114 pub fn with_relevant_maintenance_retry(
2118 mut self,
2119 attempts: u32,
2120 receipts: &[crate::vnext::DynamicPoolGrowthBatchReceipt],
2121 pools: &[crate::vnext::DynamicPoolStatus],
2122 affected_request_ids: Vec<RequestId>,
2123 ) -> Result<Self> {
2124 if receipts.is_empty() {
2125 return Ok(self);
2126 }
2127 let progress = ExecutorExecutionMaintenanceProgress::from_growth_receipts(
2128 attempts,
2129 self.observed,
2130 receipts,
2131 pools,
2132 )?;
2133 if progress.coordinator_id() != self.observed.coordinator_id
2134 || progress.latest_capacity_epoch() > self.observed.capacity_epoch
2135 || progress.mutations().is_empty()
2136 {
2137 return Err(FerrumError::internal(
2138 "execution maintenance progress does not match the exported deferral",
2139 ));
2140 }
2141 let relevant_mutation = progress
2142 .mutations()
2143 .iter()
2144 .any(|mutation| self.evidence.has_relevant_mutation(mutation));
2145 if !relevant_mutation {
2146 return Ok(self);
2147 }
2148 let retry = ExecutorExecutionMaintenanceRetry::new(affected_request_ids, progress)?;
2149 if self.stage == ExecutorExecutionCapacityStage::SequenceExtension
2150 && retry.affected_request_ids().len() != 1
2151 {
2152 return Err(FerrumError::internal(
2153 "sequence-extension maintenance retry must affect exactly one request",
2154 ));
2155 }
2156 self.maintenance_retry = Some(retry);
2157 Ok(self)
2158 }
2159
2160 pub fn from_admission_maintenance(
2161 source: &crate::vnext::AdmissionDeferred,
2162 observed: ExecutorAdmissionEpochs,
2163 wait_condition: crate::vnext::CapacityWaitCondition,
2164 pressure: crate::vnext::DynamicBackingPressure,
2165 maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
2166 stage: ExecutorExecutionCapacityStage,
2167 ) -> Result<Self> {
2168 if source.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2169 return Err(ferrum_types::FerrumError::internal(
2170 "execution maintenance source must await backing growth",
2171 ));
2172 }
2173 let evidence = ExecutorExecutionCapacityEvidence::logical_with_pressure(
2174 source.blockers().to_vec(),
2175 Some(pressure),
2176 )?
2177 .with_maintenance_boundary(maintenance_boundary)?;
2178 if evidence.maintenance_boundary().is_some_and(|boundary| {
2179 boundary.coordinator_id().get() != observed.coordinator_id.get()
2180 }) {
2181 return Err(FerrumError::internal(
2182 "execution maintenance boundary belongs to another coordinator",
2183 ));
2184 }
2185 Self::with_evidence(observed, wait_condition, stage, evidence)
2186 }
2187
2188 pub fn from_backing_maintenance(
2189 source: &crate::vnext::DynamicBackingDeferred,
2190 observed: ExecutorAdmissionEpochs,
2191 wait_condition: crate::vnext::CapacityWaitCondition,
2192 pressure: crate::vnext::DynamicBackingPressure,
2193 maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
2194 stage: ExecutorExecutionCapacityStage,
2195 ) -> Result<Self> {
2196 let evidence = ExecutorExecutionCapacityEvidence::backing_deferred_with_pressure(
2197 source.blockers().to_vec(),
2198 Some(pressure),
2199 )?
2200 .with_maintenance_boundary(maintenance_boundary)?;
2201 if evidence.maintenance_boundary().is_some_and(|boundary| {
2202 boundary.coordinator_id().get() != observed.coordinator_id.get()
2203 }) {
2204 return Err(FerrumError::internal(
2205 "execution maintenance boundary belongs to another coordinator",
2206 ));
2207 }
2208 Self::with_evidence(observed, wait_condition, stage, evidence)
2209 }
2210
2211 pub const fn observed(&self) -> ExecutorAdmissionEpochs {
2212 self.observed
2213 }
2214
2215 pub fn wait_condition(&self) -> &crate::vnext::CapacityWaitCondition {
2216 &self.wait_condition
2217 }
2218
2219 pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2220 self.stage
2221 }
2222
2223 pub const fn evidence(&self) -> &ExecutorExecutionCapacityEvidence {
2224 &self.evidence
2225 }
2226
2227 pub fn shortfalls(&self) -> &[crate::vnext::CapacityShortfall] {
2228 self.evidence.shortfalls()
2229 }
2230
2231 pub fn backing_blockers(&self) -> &[crate::vnext::DynamicBackingBlocker] {
2232 self.evidence.backing_blockers()
2233 }
2234
2235 pub const fn backing_pressure(&self) -> Option<&crate::vnext::DynamicBackingPressure> {
2236 self.evidence.backing_pressure()
2237 }
2238
2239 pub const fn maintenance_boundary(
2240 &self,
2241 ) -> Option<&crate::vnext::DynamicPoolMaintenanceBoundaryReceipt> {
2242 self.evidence.maintenance_boundary()
2243 }
2244
2245 pub fn maintenance_retry(&self) -> Option<&ExecutorExecutionMaintenanceRetry> {
2246 self.maintenance_retry.as_ref()
2247 }
2248
2249 pub fn validated_maintenance_retry_scope(
2252 &self,
2253 current_request_ids: &[RequestId],
2254 ) -> Result<Option<&ExecutorExecutionMaintenanceRetry>> {
2255 let Some(retry) = self.maintenance_retry.as_ref() else {
2256 return Ok(None);
2257 };
2258 let current = current_request_ids.iter().collect::<HashSet<_>>();
2259 if current_request_ids.is_empty() || current.len() != current_request_ids.len() {
2260 return Err(FerrumError::internal(
2261 "execution maintenance retry received an invalid current request cohort",
2262 ));
2263 }
2264 let affected = retry.affected_request_ids().iter().collect::<HashSet<_>>();
2265 if !affected.is_subset(¤t) {
2266 return Err(FerrumError::internal(
2267 "execution maintenance retry affects a request outside the current cohort",
2268 ));
2269 }
2270 match self.stage {
2271 ExecutorExecutionCapacityStage::SequenceExtension => {
2272 if affected.len() != 1 {
2273 return Err(FerrumError::internal(
2274 "sequence-extension maintenance retry must affect one current request",
2275 ));
2276 }
2277 }
2278 ExecutorExecutionCapacityStage::StepAdmission
2279 | ExecutorExecutionCapacityStage::SubmissionWave => {
2280 if affected != current {
2281 return Err(FerrumError::internal(
2282 "cohort maintenance retry must cover the complete current cohort",
2283 ));
2284 }
2285 }
2286 }
2287 Ok(Some(retry))
2288 }
2289
2290 pub fn narrower_prefill_tokens(&self, attempted_tokens: usize) -> Option<usize> {
2298 if attempted_tokens <= 1 {
2299 return None;
2300 }
2301 let maximum_next = attempted_tokens
2302 .saturating_sub(attempted_tokens.div_ceil(4))
2303 .max(1);
2304 let proportional = self
2305 .shortfalls()
2306 .iter()
2307 .filter_map(|shortfall| {
2308 let requested = shortfall.requested().get();
2309 let available = shortfall.available().get();
2310 (requested > available).then(|| {
2311 let scaled = (attempted_tokens as u128).saturating_mul(available as u128)
2312 / requested as u128;
2313 usize::try_from(scaled)
2314 .unwrap_or(usize::MAX)
2315 .clamp(1, attempted_tokens - 1)
2316 })
2317 })
2318 .min();
2319 Some(
2320 proportional
2321 .unwrap_or_else(|| attempted_tokens.div_ceil(2))
2322 .min(maximum_next)
2323 .max(1),
2324 )
2325 }
2326}
2327
2328#[derive(Debug, Clone, Serialize)]
2336pub struct ExecutorRequestStateDeferral {
2337 stage: ExecutorExecutionCapacityStage,
2338 request_ids: Vec<RequestId>,
2339 hazard: crate::vnext::RequestStateHazardDeferral,
2340}
2341
2342impl ExecutorRequestStateDeferral {
2343 pub fn new(
2344 stage: ExecutorExecutionCapacityStage,
2345 request_ids: Vec<RequestId>,
2346 hazard: crate::vnext::RequestStateHazardDeferral,
2347 ) -> Result<Self> {
2348 let unique = request_ids.iter().collect::<HashSet<_>>();
2349 if request_ids.is_empty() || unique.len() != request_ids.len() {
2350 return Err(FerrumError::internal(
2351 "request-state execution deferral requires a non-empty unique product cohort",
2352 ));
2353 }
2354 if hazard.blockers().is_empty() {
2355 return Err(FerrumError::internal(
2356 "request-state execution deferral requires exact blockers",
2357 ));
2358 }
2359 Ok(Self {
2360 stage,
2361 request_ids,
2362 hazard,
2363 })
2364 }
2365
2366 pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2367 self.stage
2368 }
2369
2370 pub fn request_ids(&self) -> &[RequestId] {
2371 &self.request_ids
2372 }
2373
2374 pub const fn hazard(&self) -> &crate::vnext::RequestStateHazardDeferral {
2375 &self.hazard
2376 }
2377
2378 pub fn register_waiter(&self) -> Result<crate::vnext::RequestStateHazardWaitRegistration> {
2379 self.hazard
2380 .register_waiter()
2381 .map_err(|error| FerrumError::backend(error.to_string()))
2382 }
2383}
2384
2385#[derive(Debug, Clone, Serialize)]
2390#[serde(tag = "reason", content = "evidence", rename_all = "snake_case")]
2391pub enum ExecutorExecutionDeferral {
2392 Capacity(ExecutorExecutionCapacityDeferral),
2393 RequestState(ExecutorRequestStateDeferral),
2394}
2395
2396impl ExecutorExecutionDeferral {
2397 pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2398 match self {
2399 Self::Capacity(deferral) => deferral.stage(),
2400 Self::RequestState(deferral) => deferral.stage(),
2401 }
2402 }
2403
2404 pub const fn as_capacity(&self) -> Option<&ExecutorExecutionCapacityDeferral> {
2405 match self {
2406 Self::Capacity(deferral) => Some(deferral),
2407 Self::RequestState(_) => None,
2408 }
2409 }
2410
2411 pub const fn as_request_state(&self) -> Option<&ExecutorRequestStateDeferral> {
2412 match self {
2413 Self::Capacity(_) => None,
2414 Self::RequestState(deferral) => Some(deferral),
2415 }
2416 }
2417}
2418
2419impl From<ExecutorExecutionCapacityDeferral> for ExecutorExecutionDeferral {
2420 fn from(deferral: ExecutorExecutionCapacityDeferral) -> Self {
2421 Self::Capacity(deferral)
2422 }
2423}
2424
2425impl From<ExecutorRequestStateDeferral> for ExecutorExecutionDeferral {
2426 fn from(deferral: ExecutorRequestStateDeferral) -> Self {
2427 Self::RequestState(deferral)
2428 }
2429}
2430
2431#[cfg(test)]
2432mod execution_capacity_deferral_tests {
2433 use super::{
2434 ExecutorAdmissionEpochs, ExecutorExecutionCapacityDeferral,
2435 ExecutorExecutionCapacityEvidenceOwner, ExecutorExecutionCapacityStage,
2436 ExecutorExecutionMaintenanceProgress, ExecutorExecutionMaintenanceRetry,
2437 };
2438 use crate::vnext::{
2439 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
2440 DeviceCapacityPressure, DeviceCapacityPressureScope, DynamicBackingPressure,
2441 };
2442 use ferrum_types::RequestId;
2443 use std::num::NonZeroU64;
2444
2445 fn test_progress() -> ExecutorExecutionMaintenanceProgress {
2446 ExecutorExecutionMaintenanceProgress {
2447 attempts: 1,
2448 coordinator_id: NonZeroU64::new(19).unwrap(),
2449 mutations: Vec::new(),
2450 latest_capacity_epoch: 5,
2451 }
2452 }
2453
2454 fn test_deferral(stage: ExecutorExecutionCapacityStage) -> ExecutorExecutionCapacityDeferral {
2455 let observed =
2456 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 7)
2457 .unwrap();
2458 let condition = CapacityWaitCondition::from_observation(19, vec![observed]).unwrap();
2459 ExecutorExecutionCapacityDeferral::from_backing_pressure(
2460 ExecutorAdmissionEpochs::new(NonZeroU64::new(19).unwrap(), 3, 5),
2461 condition,
2462 test_pressure(),
2463 stage,
2464 )
2465 .unwrap()
2466 }
2467
2468 fn test_pressure() -> DynamicBackingPressure {
2469 DeviceCapacityPressure::new(
2470 DeviceCapacityPressureScope::PlanBudget,
2471 "device.execution-capacity-test".to_owned(),
2472 1,
2473 1,
2474 1,
2475 1,
2476 1,
2477 )
2478 .unwrap()
2479 .into()
2480 }
2481
2482 #[test]
2483 fn prefill_narrowing_is_strict_bounded_and_stops_at_one_token() {
2484 let observed =
2485 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 7)
2486 .unwrap();
2487 let condition = CapacityWaitCondition::from_observation(19, vec![observed]).unwrap();
2488 let deferred = ExecutorExecutionCapacityDeferral::from_backing_pressure(
2489 ExecutorAdmissionEpochs::new(NonZeroU64::new(19).unwrap(), 3, 5),
2490 condition,
2491 test_pressure(),
2492 ExecutorExecutionCapacityStage::StepAdmission,
2493 )
2494 .unwrap();
2495
2496 assert_eq!(deferred.narrower_prefill_tokens(342), Some(171));
2497 assert_eq!(deferred.narrower_prefill_tokens(2), Some(1));
2498 assert_eq!(deferred.narrower_prefill_tokens(1), None);
2499 }
2500
2501 #[test]
2502 fn backing_pressure_serializes_one_typed_evidence_owner() {
2503 let deferred = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension);
2504 let serialized = serde_json::to_value(&deferred).unwrap();
2505
2506 assert_eq!(
2507 deferred.evidence().owner(),
2508 ExecutorExecutionCapacityEvidenceOwner::Backing
2509 );
2510 assert!(deferred.shortfalls().is_empty());
2511 assert!(deferred.backing_blockers().is_empty());
2512 assert!(deferred.backing_pressure().is_some());
2513 assert_eq!(serialized["evidence"]["owner"], "backing");
2514 assert_eq!(serialized["evidence"]["kind"], "backing_pressure");
2515 assert!(serialized["evidence"]["pressure"].is_object());
2516 }
2517
2518 #[test]
2519 fn empty_maintenance_receipts_remain_an_ordinary_typed_deferral() {
2520 let request_id = RequestId::new();
2521 let deferred = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension)
2522 .with_relevant_maintenance_retry(2, &[], &[], vec![request_id])
2523 .unwrap();
2524
2525 assert!(deferred.maintenance_retry().is_none());
2526 }
2527
2528 #[test]
2529 fn maintenance_retry_rejects_empty_duplicate_or_unproven_scope() {
2530 let request_id = RequestId::new();
2531 assert!(ExecutorExecutionMaintenanceRetry::new(Vec::new(), test_progress()).is_err());
2532 assert!(ExecutorExecutionMaintenanceRetry::new(
2533 vec![request_id.clone(), request_id.clone()],
2534 test_progress(),
2535 )
2536 .is_err());
2537 assert!(ExecutorExecutionMaintenanceRetry::new(vec![request_id], test_progress()).is_err());
2538 }
2539
2540 #[test]
2541 fn maintenance_retry_scope_is_fail_closed_for_sequence_and_cohort_stages() {
2542 let first = RequestId::new();
2543 let second = RequestId::new();
2544 let retry = |affected_request_ids| ExecutorExecutionMaintenanceRetry {
2545 affected_request_ids,
2546 progress: test_progress(),
2547 };
2548
2549 let mut sequence = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension);
2550 sequence.maintenance_retry = Some(retry(vec![second.clone()]));
2551 assert_eq!(
2552 sequence
2553 .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2554 .unwrap()
2555 .unwrap()
2556 .affected_request_ids(),
2557 [second.clone()]
2558 );
2559 sequence.maintenance_retry = Some(retry(vec![first.clone(), second.clone()]));
2560 assert!(sequence
2561 .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2562 .is_err());
2563
2564 let mut cohort = test_deferral(ExecutorExecutionCapacityStage::SubmissionWave);
2565 cohort.maintenance_retry = Some(retry(vec![second.clone()]));
2566 assert!(cohort
2567 .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2568 .is_err());
2569 cohort.maintenance_retry = Some(retry(vec![first.clone(), second.clone()]));
2570 assert!(cohort
2571 .validated_maintenance_retry_scope(&[first, second])
2572 .unwrap()
2573 .is_some());
2574 }
2575}
2576
2577pub enum ExecutorBatchDecodeOutcome {
2583 Completed(Vec<DecodeOutput>),
2584 Deferred(ExecutorExecutionDeferral),
2585}
2586
2587pub enum PlanRuntimeBatchDecodeOutcome {
2589 Completed(Vec<PlanRuntimeDecodeOutput>),
2590 Deferred(ExecutorExecutionDeferral),
2591}
2592
2593pub struct PlanRuntimePrefillCompletion {
2595 output: PlanRuntimePrefillOutput,
2596 planned_chunk: PrefillChunk,
2597 completed_chunk: PrefillChunk,
2598 capacity_probe_count: u32,
2599}
2600
2601impl PlanRuntimePrefillCompletion {
2602 pub fn new(
2603 output: PlanRuntimePrefillOutput,
2604 planned_chunk: PrefillChunk,
2605 completed_chunk: PrefillChunk,
2606 capacity_probe_count: u32,
2607 ) -> Result<Self> {
2608 validate_prefill_completion_shape(planned_chunk, completed_chunk, capacity_probe_count)?;
2609 Ok(Self {
2610 output,
2611 planned_chunk,
2612 completed_chunk,
2613 capacity_probe_count,
2614 })
2615 }
2616
2617 pub fn exact(output: PlanRuntimePrefillOutput, chunk: PrefillChunk) -> Self {
2618 Self {
2619 output,
2620 planned_chunk: chunk,
2621 completed_chunk: chunk,
2622 capacity_probe_count: 0,
2623 }
2624 }
2625
2626 pub const fn planned_chunk(&self) -> PrefillChunk {
2627 self.planned_chunk
2628 }
2629
2630 pub const fn completed_chunk(&self) -> PrefillChunk {
2631 self.completed_chunk
2632 }
2633
2634 pub const fn capacity_probe_count(&self) -> u32 {
2635 self.capacity_probe_count
2636 }
2637
2638 pub fn output(&self) -> &PlanRuntimePrefillOutput {
2639 &self.output
2640 }
2641
2642 pub fn validate_for(
2643 &self,
2644 expected_request_id: &RequestId,
2645 expected_planned_chunk: PrefillChunk,
2646 vocabulary_size: usize,
2647 ) -> Result<()> {
2648 if self.planned_chunk != expected_planned_chunk {
2649 return Err(FerrumError::backend(format!(
2650 "plan runtime completed prefill frontier {:?}, expected {:?}",
2651 self.planned_chunk.range(),
2652 expected_planned_chunk.range()
2653 )));
2654 }
2655 validate_prefill_completion_shape(
2656 self.planned_chunk,
2657 self.completed_chunk,
2658 self.capacity_probe_count,
2659 )?;
2660 self.output.validate_for_completion(
2661 expected_request_id,
2662 self.completed_chunk,
2663 vocabulary_size,
2664 )
2665 }
2666
2667 pub fn into_parts(self) -> (PlanRuntimePrefillOutput, PrefillChunk, PrefillChunk, u32) {
2668 (
2669 self.output,
2670 self.planned_chunk,
2671 self.completed_chunk,
2672 self.capacity_probe_count,
2673 )
2674 }
2675}
2676
2677pub enum PlanRuntimePrefillOutcome {
2678 Completed(PlanRuntimePrefillCompletion),
2679 Deferred(ExecutorExecutionDeferral),
2680}
2681
2682pub enum PlanRuntimeBatchPrefillOutcome {
2683 Completed(Vec<PlanRuntimePrefillCompletion>),
2684 NotSubmitted(ExecutorExecutionDeferral),
2685 Unsupported,
2686}
2687
2688pub struct ExecutorPrefillCompletion {
2694 output: PrefillOutput,
2695 planned_chunk: PrefillChunk,
2696 completed_chunk: PrefillChunk,
2697 capacity_probe_count: u32,
2698}
2699
2700impl ExecutorPrefillCompletion {
2701 pub fn new(
2702 output: PrefillOutput,
2703 planned_chunk: PrefillChunk,
2704 completed_chunk: PrefillChunk,
2705 capacity_probe_count: u32,
2706 ) -> Result<Self> {
2707 validate_prefill_completion_shape(planned_chunk, completed_chunk, capacity_probe_count)?;
2708 Ok(Self {
2709 output,
2710 planned_chunk,
2711 completed_chunk,
2712 capacity_probe_count,
2713 })
2714 }
2715
2716 pub fn exact(output: PrefillOutput, chunk: PrefillChunk) -> Self {
2717 Self {
2718 output,
2719 planned_chunk: chunk,
2720 completed_chunk: chunk,
2721 capacity_probe_count: 0,
2722 }
2723 }
2724
2725 pub const fn planned_chunk(&self) -> PrefillChunk {
2726 self.planned_chunk
2727 }
2728
2729 pub const fn completed_chunk(&self) -> PrefillChunk {
2730 self.completed_chunk
2731 }
2732
2733 pub const fn capacity_probe_count(&self) -> u32 {
2734 self.capacity_probe_count
2735 }
2736
2737 pub fn into_parts(self) -> (PrefillOutput, PrefillChunk, PrefillChunk, u32) {
2738 (
2739 self.output,
2740 self.planned_chunk,
2741 self.completed_chunk,
2742 self.capacity_probe_count,
2743 )
2744 }
2745}
2746
2747fn validate_prefill_completion_shape(
2748 planned_chunk: PrefillChunk,
2749 completed_chunk: PrefillChunk,
2750 capacity_probe_count: u32,
2751) -> Result<()> {
2752 if completed_chunk.tokens_processed() != planned_chunk.tokens_processed()
2753 || completed_chunk.total_prompt_tokens() != planned_chunk.total_prompt_tokens()
2754 || completed_chunk.tokens_to_process() > planned_chunk.tokens_to_process()
2755 {
2756 return Err(ferrum_types::FerrumError::internal(
2757 "completed prefill chunk is not a non-empty prefix of its planned chunk",
2758 ));
2759 }
2760 if completed_chunk != planned_chunk && capacity_probe_count == 0 {
2761 return Err(ferrum_types::FerrumError::internal(
2762 "partial prefill completion requires a failed capacity probe",
2763 ));
2764 }
2765 Ok(())
2766}
2767
2768pub enum ExecutorPrefillOutcome {
2769 Completed(ExecutorPrefillCompletion),
2770 Deferred(ExecutorExecutionDeferral),
2771}
2772
2773pub enum ExecutorBatchPrefillOutcome {
2781 Completed(Vec<ExecutorPrefillCompletion>),
2782 NotSubmitted(ExecutorExecutionDeferral),
2783 Unsupported,
2784}
2785
2786#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
2791#[serde(rename_all = "snake_case")]
2792pub enum ExecutorPrefillMaintenanceStage {
2793 LogicalCapacity,
2794 PhysicalBacking,
2795}
2796
2797#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2800#[serde(tag = "source", rename_all = "snake_case")]
2801pub enum ExecutorPrefillMaintenanceBlocker {
2802 Capacity {
2803 domain_id: Option<u32>,
2804 kind: crate::vnext::CapacityShortfallKind,
2805 requested: u64,
2806 available: u64,
2807 current_total: u64,
2808 maximum_total: u64,
2809 },
2810 Backing {
2811 pool_id: String,
2812 domain_id: u32,
2813 lifetime: crate::vnext::DynamicBackingClaimScope,
2814 reason: crate::vnext::DynamicBackingDeferralReason,
2815 requested_bytes: u64,
2816 free_bytes: u64,
2817 largest_contiguous_bytes: u64,
2818 },
2819}
2820
2821#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2827pub struct ExecutorPrefillMaintenanceDeferral {
2828 request_id: RequestId,
2829 observed: ExecutorAdmissionEpochs,
2830 wait_condition: crate::vnext::CapacityWaitCondition,
2831 stage: ExecutorPrefillMaintenanceStage,
2832 blockers: Vec<ExecutorPrefillMaintenanceBlocker>,
2833}
2834
2835impl ExecutorPrefillMaintenanceDeferral {
2836 pub fn new(
2837 request_id: RequestId,
2838 observed: ExecutorAdmissionEpochs,
2839 wait_condition: crate::vnext::CapacityWaitCondition,
2840 stage: ExecutorPrefillMaintenanceStage,
2841 blockers: Vec<ExecutorPrefillMaintenanceBlocker>,
2842 ) -> Result<Self> {
2843 if blockers.is_empty() {
2844 return Err(ferrum_types::FerrumError::request_validation(
2845 "executor prefill maintenance deferral requires at least one blocker",
2846 ));
2847 }
2848 if wait_condition.coordinator_id().get() != observed.coordinator_id.get() {
2849 return Err(ferrum_types::FerrumError::request_validation(
2850 "executor prefill maintenance wait condition belongs to a different coordinator",
2851 ));
2852 }
2853 Ok(Self {
2854 request_id,
2855 observed,
2856 wait_condition,
2857 stage,
2858 blockers,
2859 })
2860 }
2861
2862 pub fn from_admission(
2863 request_id: &RequestId,
2864 deferred: &crate::vnext::AdmissionDeferred,
2865 ) -> Result<Self> {
2866 if deferred.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2867 return Err(ferrum_types::FerrumError::internal(
2868 "logical prefill maintenance projection requires AwaitBackingGrowth",
2869 ));
2870 }
2871 let blockers = deferred
2872 .blockers()
2873 .iter()
2874 .map(|blocker| ExecutorPrefillMaintenanceBlocker::Capacity {
2875 domain_id: blocker.domain().map(|domain| domain.get()),
2876 kind: blocker.kind(),
2877 requested: blocker.requested().get(),
2878 available: blocker.available().get(),
2879 current_total: blocker.current_total().get(),
2880 maximum_total: blocker.maximum_total().get(),
2881 })
2882 .collect();
2883 Self::new(
2884 request_id.clone(),
2885 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2886 deferred.wait_condition().clone(),
2887 ExecutorPrefillMaintenanceStage::LogicalCapacity,
2888 blockers,
2889 )
2890 }
2891
2892 pub fn from_backing(
2893 request_id: &RequestId,
2894 deferred: &crate::vnext::DynamicBackingDeferred,
2895 ) -> Result<Self> {
2896 let blockers = deferred
2897 .blockers()
2898 .iter()
2899 .map(|blocker| ExecutorPrefillMaintenanceBlocker::Backing {
2900 pool_id: blocker.pool_id().as_str().to_string(),
2901 domain_id: blocker.domain_id().get(),
2902 lifetime: deferred.scope(),
2903 reason: blocker.reason(),
2904 requested_bytes: blocker.requested_bytes(),
2905 free_bytes: blocker.free_bytes(),
2906 largest_contiguous_bytes: blocker.largest_contiguous_bytes(),
2907 })
2908 .collect();
2909 Self::new(
2910 request_id.clone(),
2911 ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2912 deferred.wait_condition().clone(),
2913 ExecutorPrefillMaintenanceStage::PhysicalBacking,
2914 blockers,
2915 )
2916 }
2917
2918 pub fn request_id(&self) -> &RequestId {
2919 &self.request_id
2920 }
2921
2922 pub const fn observed(&self) -> ExecutorAdmissionEpochs {
2923 self.observed
2924 }
2925
2926 pub fn wait_condition(&self) -> &crate::vnext::CapacityWaitCondition {
2927 &self.wait_condition
2928 }
2929
2930 pub const fn stage(&self) -> ExecutorPrefillMaintenanceStage {
2931 self.stage
2932 }
2933
2934 pub fn blockers(&self) -> &[ExecutorPrefillMaintenanceBlocker] {
2935 &self.blockers
2936 }
2937}
2938
2939#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2941#[serde(tag = "outcome", rename_all = "snake_case")]
2942pub enum ExecutorPrefillMaintenanceOutcome {
2943 NoLongerPending,
2946 RetryAdmission { current: ExecutorAdmissionEpochs },
2951 WaitForRelease {
2955 current: ExecutorAdmissionEpochs,
2956 wait_condition: crate::vnext::CapacityWaitCondition,
2957 pressure: crate::vnext::DynamicBackingPressure,
2958 },
2959 Maintained {
2962 current: ExecutorAdmissionEpochs,
2963 pools_grown: usize,
2964 allocated_bytes: u64,
2965 pools_reclaimed: usize,
2966 chunks_reclaimed: usize,
2967 reclaimed_bytes: u64,
2968 rebalance: Option<crate::vnext::DynamicPoolRebalanceReceipt>,
2971 },
2972}
2973
2974#[derive(Debug, Clone)]
2980pub enum ExecutorPrefillAdmissionDecision {
2981 Admitted(ExecutorPrefillAdmissionReceipt),
2982 Deferred(crate::vnext::AdmissionDeferred),
2983 MaintenanceDeferred(ExecutorPrefillMaintenanceDeferral),
2984 PermanentRejected(crate::vnext::AdmissionRejected),
2985}
2986
2987#[async_trait]
2989pub trait ModelExecutor: Send + Sync {
2990 fn info(&self) -> &ModelInfo;
2992
2993 fn execution_resource_authority(&self) -> ExecutionResourceAuthority {
2999 ExecutionResourceAuthority::LegacyEngine
3000 }
3001
3002 fn admission_limits(&self) -> Result<Option<ExecutorAdmissionLimits>> {
3006 Ok(None)
3007 }
3008
3009 fn resolved_model_plan(&self) -> Option<&crate::vnext::ResolvedModelPlan> {
3014 None
3015 }
3016
3017 fn plan_runtime_resource_snapshot(&self) -> Result<Option<PlanRuntimeResourceSnapshot>> {
3021 Ok(None)
3022 }
3023
3024 fn supports_native_unified_decode(&self) -> bool {
3033 false
3034 }
3035
3036 fn kv_capacity(&self) -> Option<usize> {
3039 None
3040 }
3041
3042 fn attach_execution_event_sink(&self, _sink: Arc<dyn crate::vnext::ExecutionEventSink>) {}
3048
3049 fn execution_capacity_epochs(&self) -> Result<Option<ExecutorAdmissionEpochs>> {
3053 Ok(None)
3054 }
3055
3056 fn write_execution_capacity_snapshot(
3061 &self,
3062 availability: &mut Vec<crate::vnext::CapacityAvailabilityEpoch>,
3063 ) -> Result<Option<ExecutorAdmissionEpochs>> {
3064 availability.clear();
3065 self.execution_capacity_epochs()
3066 }
3067
3068 fn register_execution_capacity_waiter(
3076 &self,
3077 _observed: &crate::vnext::CapacityWaitCondition,
3078 ) -> Result<Option<ExecutorCapacityWaitRegistration>> {
3079 Ok(None)
3080 }
3081
3082 fn try_admit_prefill(
3086 &self,
3087 _input: ExecutorPrefillAdmission<'_>,
3088 ) -> Result<ExecutorPrefillAdmissionDecision> {
3089 Err(ferrum_types::FerrumError::unsupported(
3090 "plan-runtime prefill admission is not implemented",
3091 ))
3092 }
3093
3094 fn cancel_prefill_admission(&self, _request_id: &RequestId) -> bool {
3098 false
3099 }
3100
3101 fn write_execution_capacity_release_sources(
3110 &self,
3111 _preemption: &ExecutorExecutionCapacityPreemption,
3112 sources: &mut Vec<crate::vnext::CapacityAvailabilitySource>,
3113 ) -> Result<bool> {
3114 sources.clear();
3115 Ok(false)
3116 }
3117
3118 async fn preempt_execution_capacity(
3125 &self,
3126 _preemption: ExecutorExecutionCapacityPreemption,
3127 ) -> Result<ExecutorExecutionCapacityPreemptionReceipt> {
3128 Err(FerrumError::unsupported(
3129 "request-scoped execution-capacity preemption is not implemented",
3130 ))
3131 }
3132
3133 fn maintain_prefill_backing(
3138 &self,
3139 _request_id: &RequestId,
3140 ) -> Result<ExecutorPrefillMaintenanceOutcome> {
3141 Err(ferrum_types::FerrumError::unsupported(
3142 "plan-runtime prefill backing maintenance is not implemented",
3143 ))
3144 }
3145
3146 fn reserve_kv_slots(&self, _requests: &[KvSlotRequest]) -> Result<Option<KvSlotReservation>> {
3153 Ok(None)
3154 }
3155
3156 fn kv_slot_capacity_snapshot(&self) -> Option<KvSlotCapacitySnapshot> {
3160 None
3161 }
3162
3163 fn recurrent_state_spec(
3170 &self,
3171 _request_id: &RequestId,
3172 _input_tokens: &[TokenId],
3173 ) -> Result<Option<RecurrentStateSpec>> {
3174 Ok(None)
3175 }
3176
3177 async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput>;
3179
3180 async fn prefill_with_capacity(&self, input: &PrefillInput) -> Result<ExecutorPrefillOutcome> {
3183 let output = self.prefill(input).await?;
3184 let chunk = match input.chunk {
3185 Some(chunk) => chunk,
3186 None => PrefillChunk::new(0, input.sequence_length(), input.sequence_length())?,
3187 };
3188 Ok(ExecutorPrefillOutcome::Completed(
3189 ExecutorPrefillCompletion::exact(output, chunk),
3190 ))
3191 }
3192
3193 async fn batch_prefill(&self, inputs: &[PrefillInput]) -> Result<Vec<PrefillOutput>> {
3206 let mut outputs = Vec::with_capacity(inputs.len());
3207 for input in inputs {
3208 outputs.push(self.prefill(input).await?);
3209 }
3210 Ok(outputs)
3211 }
3212
3213 async fn batch_prefill_with_capacity(
3220 &self,
3221 _inputs: &[PrefillInput],
3222 ) -> Result<ExecutorBatchPrefillOutcome> {
3223 Ok(ExecutorBatchPrefillOutcome::Unsupported)
3224 }
3225
3226 async fn plan_runtime_prefill_with_capacity(
3232 &self,
3233 _input: &PlanRuntimePrefillInput,
3234 ) -> Result<PlanRuntimePrefillOutcome> {
3235 Err(FerrumError::unsupported(
3236 "tensor-free plan-runtime prefill is not implemented",
3237 ))
3238 }
3239
3240 async fn plan_runtime_batch_prefill_with_capacity(
3246 &self,
3247 _inputs: &[PlanRuntimePrefillInput],
3248 ) -> Result<PlanRuntimeBatchPrefillOutcome> {
3249 Ok(PlanRuntimeBatchPrefillOutcome::Unsupported)
3250 }
3251
3252 fn discard_plan_runtime_prefill(&self, authority: PlanRuntimePrefillAuthority) -> Result<()> {
3258 self.release_cache(&authority.kv_cache().cache_id());
3259 Ok(())
3260 }
3261
3262 async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput>;
3264
3265 async fn batch_decode(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>> {
3276 let mut outputs = Vec::with_capacity(inputs.len());
3277 for input in inputs {
3278 outputs.push(self.decode(input).await?);
3279 }
3280 Ok(outputs)
3281 }
3282
3283 async fn batch_decode_with_capacity(
3289 &self,
3290 inputs: &[DecodeInput],
3291 ) -> Result<ExecutorBatchDecodeOutcome> {
3292 self.batch_decode(inputs)
3293 .await
3294 .map(ExecutorBatchDecodeOutcome::Completed)
3295 }
3296
3297 async fn plan_runtime_batch_decode_with_capacity(
3305 &self,
3306 _inputs: &[PlanRuntimeDecodeInput],
3307 ) -> Result<PlanRuntimeBatchDecodeOutcome> {
3308 Err(FerrumError::unsupported(
3309 "tensor-free plan-runtime batch decode is not implemented",
3310 ))
3311 }
3312
3313 async fn unified_decode(&self, _batch: &UnifiedBatch) -> Result<Vec<Option<Vec<f32>>>> {
3334 Err(ferrum_types::FerrumError::unsupported(
3335 "unified_decode not implemented for this executor",
3336 ))
3337 }
3338
3339 async fn forward(&self, _input: &TensorRef) -> Result<TensorRef> {
3341 Err(ferrum_types::FerrumError::unsupported(
3343 "Full forward pass not supported by this executor",
3344 ))
3345 }
3346
3347 async fn truncate_kv(
3353 &self,
3354 _kv_cache: &std::sync::Arc<dyn crate::KvCacheHandle>,
3355 _new_len: usize,
3356 ) -> Result<()> {
3357 Ok(())
3358 }
3359
3360 async fn forward_verify(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>> {
3370 let mut out = Vec::with_capacity(inputs.len());
3371 for input in inputs {
3372 out.push(self.decode(input).await?);
3373 }
3374 Ok(out)
3375 }
3376
3377 fn capabilities(&self) -> ExecutorCapabilities;
3379
3380 fn status(&self) -> ExecutorStatus;
3382
3383 fn cache_metrics_snapshot(&self) -> Option<serde_json::Value> {
3389 None
3390 }
3391
3392 fn lora_metrics_snapshot(&self) -> Option<serde_json::Value> {
3394 None
3395 }
3396
3397 async fn prepare_startup(&self) -> Result<()> {
3405 Ok(())
3406 }
3407
3408 async fn warmup(&mut self) -> Result<()> {
3410 Ok(())
3412 }
3413
3414 async fn shutdown(&mut self) -> Result<()> {
3416 Ok(())
3418 }
3419
3420 fn complete_cache(&self, completion: ExecutorSequenceCompletion) -> Result<()> {
3427 self.release_cache(completion.cache_id());
3428 Ok(())
3429 }
3430
3431 fn release_cache(&self, _cache_id: &str) {
3437 }
3439}
3440
3441#[derive(Debug, Clone, Serialize, Deserialize)]
3443pub struct ExecutorCapabilities {
3444 pub max_batch_size: usize,
3446 pub max_sequence_length: usize,
3448 pub attention_mechanisms: Vec<AttentionType>,
3450 pub supports_dynamic_batching: bool,
3452 pub supports_continuous_batching: bool,
3454 pub supports_speculative_decoding: bool,
3456 pub supports_tensor_parallelism: bool,
3458 pub supports_pipeline_parallelism: bool,
3460 pub supported_dtypes: Vec<ferrum_types::DataType>,
3462 pub supported_devices: Vec<ferrum_types::Device>,
3464 pub memory_requirements: MemoryRequirements,
3466}
3467
3468#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
3470pub enum AttentionType {
3471 MultiHead,
3473 MultiQuery,
3475 GroupedQuery,
3477 Flash,
3479 Paged,
3481 SlidingWindow,
3483}
3484
3485#[derive(Debug, Clone, Serialize, Deserialize)]
3487pub struct MemoryRequirements {
3488 pub parameter_memory: u64,
3490 pub activation_memory_per_token: usize,
3492 pub kv_cache_memory_per_token: usize,
3494 pub overhead_memory: u64,
3496}
3497
3498impl MemoryRequirements {
3499 pub fn calculate_total_memory(
3501 &self,
3502 batch_size: usize,
3503 sequence_length: usize,
3504 num_layers: usize,
3505 ) -> u64 {
3506 let activation_mem =
3507 (self.activation_memory_per_token * batch_size * sequence_length) as u64;
3508 let kv_cache_mem =
3509 (self.kv_cache_memory_per_token * batch_size * sequence_length * num_layers) as u64;
3510
3511 self.parameter_memory + activation_mem + kv_cache_mem + self.overhead_memory
3512 }
3513}
3514
3515#[derive(Debug, Clone, Serialize, Deserialize)]
3517pub struct ExecutorStatus {
3518 pub state: ExecutorState,
3520 pub is_ready: bool,
3522 pub current_batch_size: usize,
3524 pub prefill_operations: u64,
3526 pub decode_operations: u64,
3528 pub avg_prefill_time_ms: f64,
3530 pub avg_decode_time_ms: f64,
3532 pub memory_usage: ExecutorMemoryUsage,
3534 #[serde(skip)]
3536 pub last_operation: Option<std::time::Instant>,
3537}
3538
3539#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
3541pub enum ExecutorState {
3542 Initializing,
3544 Ready,
3546 Busy,
3548 Error,
3550 Shutdown,
3552}
3553
3554#[derive(Debug, Clone, Serialize, Deserialize)]
3556pub struct ExecutorMemoryUsage {
3557 pub allocated_bytes: usize,
3559 pub used_bytes: usize,
3561 pub peak_bytes: usize,
3563 pub utilization_percent: f32,
3565}
3566
3567#[async_trait]
3569pub trait BatchModelExecutor: ModelExecutor {
3570 async fn batch_prefill(&self, inputs: &[PrefillInput]) -> Result<Vec<PrefillOutput>>;
3572
3573 async fn batch_decode(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>>;
3575
3576 fn optimal_batch_size(&self) -> usize;
3578
3579 fn supports_batch_size(&self, batch_size: usize) -> bool;
3581}
3582
3583#[async_trait]
3585pub trait SpeculativeExecutor: ModelExecutor {
3586 async fn speculative_decode(
3588 &self,
3589 input: &DecodeInput,
3590 draft_tokens: &[ferrum_types::TokenId],
3591 acceptance_threshold: f32,
3592 ) -> Result<SpeculativeDecodeOutput>;
3593}
3594
3595#[derive(Debug, Clone)]
3597pub struct SpeculativeDecodeOutput {
3598 pub accepted_tokens: Vec<ferrum_types::TokenId>,
3600 pub next_logits: TensorRef,
3602 pub kv_cache: Arc<dyn KvCacheHandle>,
3604 pub acceptance_count: usize,
3606}
3607
3608#[async_trait]
3610pub trait ModelExecutorFactory: Send + Sync {
3611 async fn create_executor(&self, config: &ExecutorConfig) -> Result<Box<dyn ModelExecutor>>;
3613
3614 async fn create_batch_executor(
3616 &self,
3617 config: &ExecutorConfig,
3618 ) -> Result<Box<dyn BatchModelExecutor>>;
3619
3620 fn supported_types(&self) -> Vec<ExecutorType>;
3622
3623 fn validate_config(&self, config: &ExecutorConfig) -> Result<()>;
3625}
3626
3627#[derive(Debug, Clone, Serialize, Deserialize)]
3629pub struct ExecutorConfig {
3630 pub model_info: ModelInfo,
3632 pub device: ferrum_types::Device,
3634 pub dtype: ferrum_types::DataType,
3636 pub max_batch_size: usize,
3638 pub max_sequence_length: usize,
3640 pub attention_config: ExecutorAttentionConfig,
3642 pub memory_config: ExecutorMemoryConfig,
3644 pub optimization_config: OptimizationConfig,
3646 pub executor_options: HashMap<String, serde_json::Value>,
3648}
3649
3650#[derive(Debug, Clone, Serialize, Deserialize)]
3656pub struct ExecutorAttentionConfig {
3657 pub attention_type: AttentionType,
3659 pub enable_flash_attention: bool,
3661 pub enable_paged_attention: bool,
3663 pub block_size: Option<usize>,
3665 pub sliding_window_size: Option<usize>,
3667}
3668
3669#[derive(Debug, Clone, Serialize, Deserialize)]
3671pub struct ExecutorMemoryConfig {
3672 pub enable_memory_pooling: bool,
3674 pub memory_pool_size: Option<usize>,
3676 pub enable_kv_cache_sharing: bool,
3678 pub max_memory_usage: f32,
3680}
3681
3682#[derive(Debug, Clone, Serialize, Deserialize)]
3684pub struct OptimizationConfig {
3685 pub enable_cuda_graphs: bool,
3687 pub enable_kernel_fusion: bool,
3689 pub enable_mixed_precision: bool,
3691 pub optimization_level: u8,
3693 pub custom_flags: HashMap<String, bool>,
3695}
3696
3697#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
3699pub enum ExecutorType {
3700 Sequential,
3702 Batch,
3704 ContinuousBatch,
3706 Speculative,
3708 PipelineParallel,
3710 TensorParallel,
3712}
3713
3714#[derive(Debug, Clone, Serialize, Deserialize)]
3716pub struct ExecutorMetrics {
3717 pub total_operations: u64,
3719 pub prefill_operations: u64,
3721 pub decode_operations: u64,
3723 pub avg_prefill_latency: f64,
3725 pub avg_decode_latency: f64,
3727 pub p95_prefill_latency: f64,
3729 pub p95_decode_latency: f64,
3731 pub throughput_tps: f64,
3733 pub memory_efficiency: f32,
3735 pub batch_utilization: f32,
3737}
3738
3739pub trait ExecutorRegistry: Send + Sync {
3741 fn register(&mut self, name: &str, executor: Box<dyn ModelExecutor>) -> Result<()>;
3743
3744 fn get(&self, name: &str) -> Option<&dyn ModelExecutor>;
3746
3747 fn remove(&mut self, name: &str) -> Option<Box<dyn ModelExecutor>>;
3749
3750 fn list_names(&self) -> Vec<String>;
3752
3753 fn get_metrics(&self, name: &str) -> Option<ExecutorMetrics>;
3755}