1mod pressure;
13
14#[cfg(test)]
15mod historical_replay_tests;
16
17use pressure::{
18 LogicalWorkFrontier, PressureCandidate, PressureCoordinator, PressureDecision,
19 PressureHoldStatus, PressureReleaseFenceDisposition,
20};
21pub use pressure::{
22 LogicalWorkGeneration, PressureEpisodeId, PressureEpisodeState, PressureHoldReleaseReason,
23 PressureInvariantViolation, PressureInvariantViolationClass, PressureTransition,
24 PressureTransitionKind, PressureTransitionOrdinal, PressureYieldKind, PressureYieldTransaction,
25};
26
27use crate::vnext::{
28 AdmissionDeferral, AdmissionProbeOutcome, AdmissionQueueEligibility, AdmissionQueueEvent,
29 AdmissionTickReceipt, AdmissionWakeEpochs, AdmissionWakeSnapshot, DynamicAdmissionQueue,
30 DynamicAdmissionQueuePolicy, WaitingAdmissionTicket,
31};
32use crate::{
33 BatchHint, BatchPlan, BatchResourceRequirements, PreemptionResult, PreemptionState,
34 ScheduledRequest, Scheduler,
35};
36use async_trait::async_trait;
37use ferrum_interfaces::model_executor::{
38 ExecutorExecutionMaintenanceRetry, ExecutorPrefillAdmissionReceipt,
39};
40use ferrum_interfaces::scheduler::SchedulerMetrics;
41use ferrum_interfaces::vnext::{
42 AdmissionRejected, CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
43};
44use ferrum_types::{
45 BatchId, FerrumError, InferenceRequest, InferenceResponse, Priority, RequestId, RequestState,
46 Result, SchedulerConfig, PROMPT_TOKENS_METADATA_KEY,
47};
48use indexmap::IndexMap;
49use parking_lot::{Mutex, RwLock};
50use serde::Serialize;
51use std::{
52 collections::{BTreeMap, HashMap, HashSet, VecDeque},
53 num::NonZeroU64,
54 sync::{
55 atomic::{AtomicBool, AtomicU64, AtomicU8, AtomicUsize, Ordering},
56 Arc,
57 },
58 time::Instant,
59};
60use tracing::{debug, info, warn};
61
62const NO_CAPACITY_BACKPRESSURE_LIMIT: usize = usize::MAX;
63const CAPACITY_DECODE_FREE_BLOCK_HEADROOM: usize = 1;
64const CAPACITY_MIXED_RECOMPUTE_FREE_BLOCK_HEADROOM: usize = 1;
65
66#[derive(Debug, Clone, Default, PartialEq, Eq)]
67struct ContinuousBatchRuntimeConfig {
68 prompt_token_estimate: bool,
69 prefill_first_until_active: Option<usize>,
70 prefill_step_chunk: Option<usize>,
71 active_decode_prefill_chunk: Option<usize>,
72 scheduler_none_prof: bool,
73}
74
75impl ContinuousBatchRuntimeConfig {
76 fn from_scheduler_config(config: &SchedulerConfig) -> Self {
77 Self {
78 prompt_token_estimate: config.prompt_token_estimate,
79 prefill_first_until_active: config.prefill_first_until_active,
80 prefill_step_chunk: config.prefill_step_chunk,
81 active_decode_prefill_chunk: config.active_decode_prefill_chunk,
82 scheduler_none_prof: config.scheduler_none_prof,
83 }
84 }
85}
86
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub enum RequestPhase {
90 Waiting,
92 Prefilling,
94 Decoding,
96 Completed,
98 Preempted,
100 Cancelled,
102 AdmissionFailed,
104}
105
106#[derive(Debug, Clone, PartialEq, Eq)]
112pub enum ExecutionCapacityAction {
113 Deferred {
114 count: usize,
115 },
116 YieldPlanned {
117 transaction: PressureYieldTransaction,
118 },
119 InvariantViolation {
120 violation: PressureInvariantViolation,
121 },
122}
123
124const EXECUTION_READINESS_PENDING: u8 = 0;
125const EXECUTION_READINESS_READY: u8 = 1;
126const EXECUTION_READINESS_FAILED: u8 = 2;
127const EXECUTION_READINESS_CANCELLED: u8 = 3;
128
129#[derive(Debug)]
130struct ExecutionReadinessState {
131 status: AtomicU8,
132}
133
134#[derive(Debug, Clone)]
138pub struct ExecutionReadinessWake {
139 ticket_id: NonZeroU64,
140 state: Arc<ExecutionReadinessState>,
141}
142
143impl ExecutionReadinessWake {
144 pub const fn ticket_id(&self) -> NonZeroU64 {
145 self.ticket_id
146 }
147
148 pub fn mark_ready(&self) -> bool {
149 self.transition(EXECUTION_READINESS_READY)
150 }
151
152 pub fn mark_failed(&self) -> bool {
153 self.transition(EXECUTION_READINESS_FAILED)
154 }
155
156 pub fn cancel(&self) -> bool {
157 self.transition(EXECUTION_READINESS_CANCELLED)
158 }
159
160 fn transition(&self, next: u8) -> bool {
161 self.state
162 .status
163 .compare_exchange(
164 EXECUTION_READINESS_PENDING,
165 next,
166 Ordering::Release,
167 Ordering::Acquire,
168 )
169 .is_ok()
170 }
171}
172
173#[derive(Debug, Clone)]
174struct ExecutionReadinessBlock {
175 ticket_id: NonZeroU64,
176 state: Arc<ExecutionReadinessState>,
177}
178
179impl ExecutionReadinessBlock {
180 fn status(&self) -> u8 {
181 self.state.status.load(Ordering::Acquire)
182 }
183
184 fn matches(&self, other: &Self) -> bool {
185 self.ticket_id == other.ticket_id && Arc::ptr_eq(&self.state, &other.state)
186 }
187}
188
189#[derive(Debug, Clone)]
190pub struct ExecutionReadinessDeferralReceipt {
191 deferred_count: usize,
192 wake: ExecutionReadinessWake,
193}
194
195impl ExecutionReadinessDeferralReceipt {
196 pub const fn deferred_count(&self) -> usize {
197 self.deferred_count
198 }
199
200 pub const fn wake(&self) -> &ExecutionReadinessWake {
201 &self.wake
202 }
203
204 pub fn into_wake(self) -> ExecutionReadinessWake {
205 self.wake
206 }
207}
208
209#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
212pub struct ExecutionMaintenanceRetryReceipt {
213 deferred_count: usize,
214 not_before_iteration: u64,
215 latest_capacity_epoch: u64,
216}
217
218impl ExecutionMaintenanceRetryReceipt {
219 pub const fn deferred_count(&self) -> usize {
220 self.deferred_count
221 }
222
223 pub const fn not_before_iteration(&self) -> u64 {
224 self.not_before_iteration
225 }
226
227 pub const fn latest_capacity_epoch(&self) -> u64 {
228 self.latest_capacity_epoch
229 }
230}
231
232#[derive(Debug, Clone, Copy, PartialEq, Eq)]
234pub enum ExecutionCapacityYieldDisposition {
235 ProgressOwnerResumable,
236 ProgressOwnerAdmissionPending,
237 SelfRecomputeQueued,
238 OwnerTerminal,
239}
240
241impl ExecutionCapacityYieldDisposition {
242 pub const fn as_str(self) -> &'static str {
243 match self {
244 Self::ProgressOwnerResumable => "progress_owner_resumable",
245 Self::ProgressOwnerAdmissionPending => "progress_owner_admission_pending",
246 Self::SelfRecomputeQueued => "self_recompute_queued",
247 Self::OwnerTerminal => "owner_terminal",
248 }
249 }
250}
251
252#[derive(Debug, Clone, PartialEq, Eq)]
254pub struct ExecutionCapacityPressureHoldReceipt {
255 episode_id: PressureEpisodeId,
256 transition_ordinal: PressureTransitionOrdinal,
257 request_id: RequestId,
258 progress_owner_id: RequestId,
259 progress_baseline: LogicalWorkGeneration,
260 progress_current: LogicalWorkGeneration,
261 waiting_ticket: u64,
262}
263
264impl ExecutionCapacityPressureHoldReceipt {
265 pub const fn episode_id(&self) -> PressureEpisodeId {
266 self.episode_id
267 }
268
269 pub const fn transition_ordinal(&self) -> PressureTransitionOrdinal {
270 self.transition_ordinal
271 }
272
273 pub const fn request_id(&self) -> &RequestId {
274 &self.request_id
275 }
276
277 pub const fn progress_owner_id(&self) -> &RequestId {
278 &self.progress_owner_id
279 }
280
281 pub const fn progress_baseline(&self) -> LogicalWorkGeneration {
282 self.progress_baseline
283 }
284
285 pub const fn progress_current(&self) -> LogicalWorkGeneration {
286 self.progress_current
287 }
288
289 pub const fn waiting_ticket(&self) -> u64 {
290 self.waiting_ticket
291 }
292}
293
294#[derive(Debug, Clone, PartialEq, Eq)]
303pub struct ExecutionCapacityYieldCompletion {
304 victim_requeued: bool,
305 installed_hold: Option<ExecutionCapacityPressureHoldReceipt>,
306 release_transition_ordinal: PressureTransitionOrdinal,
307 resumable_transition_ordinal: Option<PressureTransitionOrdinal>,
308 owner_admission_pending_transition_ordinal: Option<PressureTransitionOrdinal>,
309 closed_transition_ordinal: Option<PressureTransitionOrdinal>,
310 disposition: ExecutionCapacityYieldDisposition,
311}
312
313impl ExecutionCapacityYieldCompletion {
314 pub const fn victim_requeued(&self) -> bool {
315 self.victim_requeued
316 }
317
318 pub const fn installed_hold(&self) -> Option<&ExecutionCapacityPressureHoldReceipt> {
319 self.installed_hold.as_ref()
320 }
321
322 pub const fn progress_owner_resumable(&self) -> bool {
323 matches!(
324 self.disposition,
325 ExecutionCapacityYieldDisposition::ProgressOwnerResumable
326 )
327 }
328
329 pub const fn release_transition_ordinal(&self) -> PressureTransitionOrdinal {
330 self.release_transition_ordinal
331 }
332
333 pub const fn resumable_transition_ordinal(&self) -> Option<PressureTransitionOrdinal> {
334 self.resumable_transition_ordinal
335 }
336
337 pub const fn owner_admission_pending_transition_ordinal(
338 &self,
339 ) -> Option<PressureTransitionOrdinal> {
340 self.owner_admission_pending_transition_ordinal
341 }
342
343 pub const fn closed_transition_ordinal(&self) -> Option<PressureTransitionOrdinal> {
344 self.closed_transition_ordinal
345 }
346
347 pub const fn closed_reason(&self) -> Option<PressureHoldReleaseReason> {
348 match self.disposition {
349 ExecutionCapacityYieldDisposition::OwnerTerminal => {
350 Some(PressureHoldReleaseReason::OwnerTerminal)
351 }
352 ExecutionCapacityYieldDisposition::ProgressOwnerResumable
353 | ExecutionCapacityYieldDisposition::ProgressOwnerAdmissionPending
354 | ExecutionCapacityYieldDisposition::SelfRecomputeQueued => None,
355 }
356 }
357
358 pub const fn disposition(&self) -> ExecutionCapacityYieldDisposition {
359 self.disposition
360 }
361}
362
363#[derive(Debug, Clone, Default, PartialEq, Eq)]
370pub struct ExecutionCapacityReleaseSnapshot {
371 release_sources_by_request: HashMap<RequestId, Vec<CapacityAvailabilitySource>>,
372}
373
374impl ExecutionCapacityReleaseSnapshot {
375 pub fn new(
376 capabilities: impl IntoIterator<Item = (RequestId, Vec<CapacityAvailabilitySource>)>,
377 ) -> Self {
378 let mut release_sources_by_request = HashMap::new();
379 for (request_id, mut sources) in capabilities {
380 sources.sort_unstable();
381 sources.dedup();
382 if !sources.is_empty() {
383 release_sources_by_request.insert(request_id, sources);
384 }
385 }
386 Self {
387 release_sources_by_request,
388 }
389 }
390
391 fn can_advance(&self, request_id: &RequestId, condition: &CapacityWaitCondition) -> bool {
392 let Some(release_sources) = self.release_sources_by_request.get(request_id) else {
393 return false;
394 };
395 condition
396 .observed()
397 .iter()
398 .any(|observed| release_sources.binary_search(&observed.source()).is_ok())
399 }
400
401 pub fn has_external_releaser(
402 &self,
403 blocked_request_id: &RequestId,
404 condition: &CapacityWaitCondition,
405 ) -> bool {
406 self.release_sources_by_request.keys().any(|request_id| {
407 request_id != blocked_request_id && self.can_advance(request_id, condition)
408 })
409 }
410}
411
412#[derive(Debug, Clone, Copy, PartialEq, Eq)]
413struct ExecutionMaintenanceRetryTicket {
414 not_before_iteration: u64,
415 latest_capacity_epoch: u64,
416}
417
418#[derive(Debug, Clone)]
420pub struct ContinuousBatchRequest {
421 pub inner: ScheduledRequest,
423 pub phase: RequestPhase,
425 pub prefill_tokens: usize,
427 pub decode_tokens: usize,
429 logical_work_frontier: LogicalWorkFrontier,
431 pub kv_blocks: Vec<ferrum_types::BlockId>,
433 pub chunked_prefill: bool,
435 pub prefill_chunk_offset: usize,
437 pub prefill_execution_chunk_ceiling: Option<usize>,
440 pub last_iteration: u64,
442 pub prefill_time_ms: u64,
444 pub decode_time_ms: u64,
446 pub capacity_deferred_until_release_epoch: u64,
448 pub capacity_deferred_mixed_attempt_epoch: Option<u64>,
450 pub capacity_deferred_empty_retry_epoch: Option<u64>,
452 pub capacity_deferred_from_decode: bool,
454 pub waiting_admission_ticket: Option<WaitingAdmissionTicket>,
456 pub execution_capacity_deferral: Option<AdmissionDeferral>,
458 execution_readiness_block: Option<ExecutionReadinessBlock>,
462 execution_maintenance_retry: Option<ExecutionMaintenanceRetryTicket>,
464 last_execution_maintenance_capacity_epoch: Option<u64>,
466}
467
468impl ContinuousBatchRequest {
469 pub fn new(request: InferenceRequest) -> Self {
471 Self {
472 inner: ScheduledRequest::new(request),
473 phase: RequestPhase::Waiting,
474 prefill_tokens: 0,
475 decode_tokens: 0,
476 logical_work_frontier: LogicalWorkFrontier::default(),
477 kv_blocks: Vec::new(),
478 chunked_prefill: false,
479 prefill_chunk_offset: 0,
480 prefill_execution_chunk_ceiling: None,
481 last_iteration: 0,
482 prefill_time_ms: 0,
483 decode_time_ms: 0,
484 capacity_deferred_until_release_epoch: 0,
485 capacity_deferred_mixed_attempt_epoch: None,
486 capacity_deferred_empty_retry_epoch: None,
487 capacity_deferred_from_decode: false,
488 waiting_admission_ticket: None,
489 execution_capacity_deferral: None,
490 execution_readiness_block: None,
491 execution_maintenance_retry: None,
492 last_execution_maintenance_capacity_epoch: None,
493 }
494 }
495
496 pub fn total_tokens(&self) -> usize {
498 self.prefill_tokens + self.decode_tokens
499 }
500
501 pub fn is_active(&self) -> bool {
503 matches!(
504 self.phase,
505 RequestPhase::Prefilling | RequestPhase::Decoding
506 )
507 }
508
509 pub fn is_finished(&self) -> bool {
511 matches!(
512 self.phase,
513 RequestPhase::Completed | RequestPhase::Cancelled | RequestPhase::AdmissionFailed
514 )
515 }
516}
517
518pub type ExecutorAdmissionProbeOutcome =
519 AdmissionProbeOutcome<ExecutorPrefillAdmissionReceipt, AdmissionRejected, FerrumError>;
520
521#[derive(Debug, Clone, PartialEq, Eq)]
522pub enum ExecutorAdmissionQueueObservation {
523 PressureHoldReleased {
524 episode_id: PressureEpisodeId,
525 transition_ordinal: PressureTransitionOrdinal,
526 request_id: RequestId,
527 progress_owner_id: RequestId,
528 progress_baseline: LogicalWorkGeneration,
529 progress_current: LogicalWorkGeneration,
530 reason: PressureHoldReleaseReason,
531 previous_wait_condition: Option<CapacityWaitCondition>,
532 current_wait_condition: Option<CapacityWaitCondition>,
533 ticket: u64,
534 },
535 SkippedUnchanged {
536 request_id: RequestId,
537 ticket: u64,
538 deferral: AdmissionDeferral,
539 current: AdmissionWakeEpochs,
540 },
541 DecodeSkippedUnchanged {
542 request_id: RequestId,
543 deferral: AdmissionDeferral,
544 current: AdmissionWakeEpochs,
545 current_wait_sources: Vec<CapacityAvailabilityEpoch>,
546 },
547 DecodeResumed {
548 request_id: RequestId,
549 deferral: AdmissionDeferral,
550 current: AdmissionWakeEpochs,
551 current_wait_sources: Vec<CapacityAvailabilityEpoch>,
552 exact_source_changed: bool,
553 policy_epoch_changed: bool,
554 },
555 PrefillSkippedUnchanged {
556 request_id: RequestId,
557 deferral: AdmissionDeferral,
558 current: AdmissionWakeEpochs,
559 current_wait_sources: Vec<CapacityAvailabilityEpoch>,
560 },
561 PrefillResumed {
562 request_id: RequestId,
563 deferral: AdmissionDeferral,
564 current: AdmissionWakeEpochs,
565 current_wait_sources: Vec<CapacityAvailabilityEpoch>,
566 exact_source_changed: bool,
567 policy_epoch_changed: bool,
568 },
569}
570
571#[derive(Debug, Clone, Copy)]
572enum ExecutionCapacityQueuePhase {
573 Prefill,
574 Decode,
575}
576
577type ExecutorAdmissionQueueEvent = AdmissionQueueEvent<
578 ContinuousBatchRequest,
579 ExecutorPrefillAdmissionReceipt,
580 AdmissionRejected,
581 FerrumError,
582>;
583
584enum WaitingAdmissionMode<'a> {
585 Legacy,
586 Dynamic {
587 wake: AdmissionWakeSnapshot<'a>,
588 probe: &'a mut dyn FnMut(&InferenceRequest) -> ExecutorAdmissionProbeOutcome,
589 observer: Option<&'a mut dyn FnMut(ExecutorAdmissionQueueObservation)>,
590 },
591}
592
593impl<'a> WaitingAdmissionMode<'a> {
594 fn wake(&self) -> Option<AdmissionWakeSnapshot<'a>> {
595 match self {
596 Self::Legacy => None,
597 Self::Dynamic { wake, .. } => Some(*wake),
598 }
599 }
600
601 fn observe(&mut self, observation: ExecutorAdmissionQueueObservation) {
602 if let Self::Dynamic {
603 observer: Some(observer),
604 ..
605 } = self
606 {
607 observer(observation);
608 }
609 }
610
611 fn observes(&self) -> bool {
612 matches!(
613 self,
614 Self::Dynamic {
615 observer: Some(_),
616 ..
617 }
618 )
619 }
620}
621
622#[derive(Debug, Default)]
623struct DecodeQueueState {
624 requests: IndexMap<RequestId, ContinuousBatchRequest>,
625 selection_cursor: Option<RequestId>,
626}
627
628impl DecodeQueueState {
629 fn remove(&mut self, request_id: &RequestId) -> Option<ContinuousBatchRequest> {
630 let removed_index = self.requests.get_index_of(request_id)?;
631 let old_len = self.requests.len();
632 let cursor_was_removed = self.selection_cursor.as_ref() == Some(request_id);
633 let successor = if cursor_was_removed && old_len > 1 {
634 Some(
635 self.requests
636 .get_index((removed_index + 1) % old_len)
637 .expect("decode successor remains in bounds before removal")
638 .0
639 .clone(),
640 )
641 } else {
642 None
643 };
644
645 let removed = self.requests.swap_remove(request_id);
646 if self.requests.is_empty() {
647 self.selection_cursor = None;
648 } else if cursor_was_removed {
649 self.selection_cursor = successor;
650 } else if self
651 .selection_cursor
652 .as_ref()
653 .is_some_and(|cursor_id| !self.requests.contains_key(cursor_id))
654 {
655 self.selection_cursor = self.requests.get_index(0).map(|(id, _)| id.clone());
656 }
657 removed
658 }
659}
660
661#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
663pub struct ContinuousSchedulerTraceSnapshot {
664 pub current_iteration: u64,
665 pub waiting_queue_len: usize,
666 pub prefill_queue_len: usize,
667 pub decode_queue_len: usize,
668 pub decode_selection_cursor: Option<RequestId>,
669 pub preempted_queue_len: usize,
670 pub active_len: usize,
671 pub completed_total: u64,
672 pub failed_total: u64,
673 pub cancelled_total: u64,
674 pub preempted_total: u64,
675 pub admitted_total: u64,
676 pub capacity_deferred_total: u64,
677 pub capacity_backpressure_admit_limit: Option<usize>,
678 pub decode_capacity_backpressure_admit_limit: Option<usize>,
679 pub decode_execution_pressure_enforced: bool,
680 pub capacity_blocked_waiting_len: usize,
681 pub execution_capacity_blocked_prefill_len: usize,
682 pub execution_capacity_blocked_decode_len: usize,
683 pub execution_readiness_deferred_total: u64,
684 pub execution_readiness_blocked_prefill_len: usize,
685 pub execution_readiness_blocked_decode_len: usize,
686 pub capacity_release_epoch: u64,
687 pub capacity_mixed_recompute_epoch: u64,
688 pub capacity_mixed_recompute_blocked_until_epoch: u64,
689 pub capacity_mixed_recompute_required_blocks_per_slot: Option<usize>,
690 pub capacity_mixed_recompute_observed_free_blocks: Option<usize>,
691 pub legacy_waiting_admission_ticks: u64,
692 pub dynamic_admission_ticks: u64,
693 pub dynamic_admission_probes: u64,
694 pub dynamic_admission_skipped_unchanged: u64,
695 pub dynamic_admission_deferred: u64,
696 pub dynamic_backing_growth_requested: u64,
697 pub dynamic_admission_failed: u64,
698 pub pressure_episodes_created: u64,
699 pub pressure_episodes_merged: u64,
700 pub pressure_episode_bridges_deferred: u64,
701 pub pressure_active_episodes: usize,
702 pub pressure_pending_release_fences: usize,
703 pub pressure_candidate_scans: u64,
704 pub pressure_last_transition_ordinal: u64,
705 pub pressure_dropped_journal_entries: u64,
706}
707
708#[derive(Debug, Clone, Copy, PartialEq, Eq)]
713pub struct ContinuousSchedulerAdmissionCounts {
714 pub waiting_requests: usize,
715 pub active_prefill_sequences: usize,
716 pub active_decode_sequences: usize,
717}
718
719pub struct ContinuousBatchScheduler {
724 config: SchedulerConfig,
726
727 waiting_queue: RwLock<DynamicAdmissionQueue<ContinuousBatchRequest>>,
729
730 prefill_queue: RwLock<VecDeque<ContinuousBatchRequest>>,
732
733 decode_queue: RwLock<DecodeQueueState>,
735
736 preempted_requests: RwLock<HashMap<RequestId, ContinuousBatchRequest>>,
738
739 admission_failed_requests: RwLock<HashMap<RequestId, ContinuousBatchRequest>>,
741 admission_failures: Mutex<VecDeque<(RequestId, FerrumError)>>,
742 dynamic_admission_events: Mutex<Vec<ExecutorAdmissionQueueEvent>>,
743
744 request_index: RwLock<HashMap<RequestId, RequestPhase>>,
746
747 current_iteration: AtomicU64,
749
750 completed_counter: AtomicU64,
752 failed_counter: AtomicU64,
753 cancelled_counter: AtomicU64,
754 preempted_counter: AtomicU64,
755 admitted_counter: AtomicU64,
756 capacity_deferred_counter: AtomicU64,
757 execution_readiness_deferred_counter: AtomicU64,
758 next_execution_readiness_ticket: AtomicU64,
759 capacity_backpressure_limit: AtomicUsize,
760 decode_capacity_backpressure_limit: AtomicUsize,
761 decode_execution_pressure_enforced: AtomicBool,
762 decode_execution_recovery_release_epoch: AtomicU64,
763 decode_capacity_feedback_lock: Mutex<()>,
764 capacity_backpressure_iteration: AtomicU64,
765 capacity_release_epoch: AtomicU64,
766 capacity_mixed_recompute_epoch: AtomicU64,
767 capacity_mixed_recompute_blocked_until_epoch: AtomicU64,
768 capacity_mixed_recompute_required_blocks_per_slot: AtomicUsize,
769 capacity_mixed_recompute_observed_free_blocks: AtomicUsize,
770 total_wait_time_us: AtomicU64,
771 legacy_waiting_admission_ticks: AtomicU64,
772 dynamic_admission_ticks: AtomicU64,
773 dynamic_admission_probes: AtomicU64,
774 dynamic_admission_skipped_unchanged: AtomicU64,
775 dynamic_admission_deferred: AtomicU64,
776 dynamic_backing_growth_requested: AtomicU64,
777 dynamic_admission_failed: AtomicU64,
778
779 pressure_coordinator: Mutex<PressureCoordinator>,
781 pressure_active: AtomicBool,
783
784 start_time: Instant,
786
787 metrics_tracker: Arc<ContinuousBatchMetrics>,
789
790 cb_config: ContinuousBatchConfig,
792
793 runtime_config: ContinuousBatchRuntimeConfig,
795}
796
797#[derive(Debug, Clone)]
799pub struct ContinuousBatchConfig {
800 pub max_prefill_batch: usize,
802 pub max_decode_batch: usize,
804 pub enable_chunked_prefill: bool,
806 pub prefill_chunk_size: usize,
808 pub max_kv_blocks_per_request: usize,
810 pub enable_swapping: bool,
812 pub swap_priority_threshold: Priority,
814 pub target_iteration_time_ms: u64,
816}
817
818impl Default for ContinuousBatchConfig {
819 fn default() -> Self {
820 Self {
821 max_prefill_batch: 8,
822 max_decode_batch: 256,
823 enable_chunked_prefill: true,
824 prefill_chunk_size: 512,
825 max_kv_blocks_per_request: 1024,
826 enable_swapping: true,
827 swap_priority_threshold: Priority::Low,
828 target_iteration_time_ms: 50,
829 }
830 }
831}
832
833struct ContinuousBatchMetrics {
835 total_prefill_tokens: AtomicU64,
836 total_decode_tokens: AtomicU64,
837 total_prefill_time_ms: AtomicU64,
838 total_decode_time_ms: AtomicU64,
839 request_count: AtomicU64,
840 iteration_count: AtomicU64,
841}
842
843impl ContinuousBatchMetrics {
844 fn new() -> Self {
845 Self {
846 total_prefill_tokens: AtomicU64::new(0),
847 total_decode_tokens: AtomicU64::new(0),
848 total_prefill_time_ms: AtomicU64::new(0),
849 total_decode_time_ms: AtomicU64::new(0),
850 request_count: AtomicU64::new(0),
851 iteration_count: AtomicU64::new(0),
852 }
853 }
854
855 fn record_completion(&self, req: &ContinuousBatchRequest) {
856 self.total_prefill_tokens
857 .fetch_add(req.prefill_tokens as u64, Ordering::Relaxed);
858 self.total_decode_tokens
859 .fetch_add(req.decode_tokens as u64, Ordering::Relaxed);
860 self.total_prefill_time_ms
861 .fetch_add(req.prefill_time_ms, Ordering::Relaxed);
862 self.total_decode_time_ms
863 .fetch_add(req.decode_time_ms, Ordering::Relaxed);
864 self.request_count.fetch_add(1, Ordering::Relaxed);
865 }
866
867 fn record_iteration(&self) {
868 self.iteration_count.fetch_add(1, Ordering::Relaxed);
869 }
870}
871
872impl ContinuousBatchScheduler {
873 pub fn new(config: SchedulerConfig) -> Self {
875 Self::with_cb_config(config, ContinuousBatchConfig::default())
876 }
877
878 pub fn with_cb_config(config: SchedulerConfig, cb_config: ContinuousBatchConfig) -> Self {
880 info!(
881 "Creating continuous batch scheduler: max_prefill={}, max_decode={}",
882 cb_config.max_prefill_batch, cb_config.max_decode_batch
883 );
884 let runtime_config = ContinuousBatchRuntimeConfig::from_scheduler_config(&config);
885
886 Self {
887 config,
888 waiting_queue: RwLock::new(DynamicAdmissionQueue::new(
889 DynamicAdmissionQueuePolicy::default(),
890 )),
891 prefill_queue: RwLock::new(VecDeque::new()),
892 decode_queue: RwLock::new(DecodeQueueState::default()),
893 preempted_requests: RwLock::new(HashMap::new()),
894 admission_failed_requests: RwLock::new(HashMap::new()),
895 admission_failures: Mutex::new(VecDeque::new()),
896 dynamic_admission_events: Mutex::new(Vec::new()),
897 request_index: RwLock::new(HashMap::new()),
898 current_iteration: AtomicU64::new(0),
899 completed_counter: AtomicU64::new(0),
900 failed_counter: AtomicU64::new(0),
901 cancelled_counter: AtomicU64::new(0),
902 preempted_counter: AtomicU64::new(0),
903 admitted_counter: AtomicU64::new(0),
904 capacity_deferred_counter: AtomicU64::new(0),
905 execution_readiness_deferred_counter: AtomicU64::new(0),
906 next_execution_readiness_ticket: AtomicU64::new(1),
907 capacity_backpressure_limit: AtomicUsize::new(NO_CAPACITY_BACKPRESSURE_LIMIT),
908 decode_capacity_backpressure_limit: AtomicUsize::new(NO_CAPACITY_BACKPRESSURE_LIMIT),
909 decode_execution_pressure_enforced: AtomicBool::new(false),
910 decode_execution_recovery_release_epoch: AtomicU64::new(0),
911 decode_capacity_feedback_lock: Mutex::new(()),
912 capacity_backpressure_iteration: AtomicU64::new(u64::MAX),
913 capacity_release_epoch: AtomicU64::new(0),
914 capacity_mixed_recompute_epoch: AtomicU64::new(0),
915 capacity_mixed_recompute_blocked_until_epoch: AtomicU64::new(0),
916 capacity_mixed_recompute_required_blocks_per_slot: AtomicUsize::new(0),
917 capacity_mixed_recompute_observed_free_blocks: AtomicUsize::new(usize::MAX),
918 total_wait_time_us: AtomicU64::new(0),
919 legacy_waiting_admission_ticks: AtomicU64::new(0),
920 dynamic_admission_ticks: AtomicU64::new(0),
921 dynamic_admission_probes: AtomicU64::new(0),
922 dynamic_admission_skipped_unchanged: AtomicU64::new(0),
923 dynamic_admission_deferred: AtomicU64::new(0),
924 dynamic_backing_growth_requested: AtomicU64::new(0),
925 dynamic_admission_failed: AtomicU64::new(0),
926 pressure_coordinator: Mutex::new(PressureCoordinator::default()),
927 pressure_active: AtomicBool::new(false),
928 start_time: Instant::now(),
929 metrics_tracker: Arc::new(ContinuousBatchMetrics::new()),
930 cb_config,
931 runtime_config,
932 }
933 }
934
935 pub fn active_count(&self) -> usize {
937 self.prefill_queue.read().len() + self.decode_queue.read().requests.len()
938 }
939
940 pub fn waiting_count(&self) -> usize {
942 self.waiting_queue.read().len()
943 }
944
945 pub fn passive_capacity_wait_condition(
949 &self,
950 ) -> Result<Option<ferrum_interfaces::vnext::CapacityWaitCondition>> {
951 let mut conditions = Vec::new();
952 {
953 let prefill = self.prefill_queue.read();
954 for request in prefill.iter() {
955 if request
956 .execution_readiness_block
957 .as_ref()
958 .is_some_and(|block| block.status() != EXECUTION_READINESS_CANCELLED)
959 {
960 continue;
961 }
962 let Some(deferral) = request.execution_capacity_deferral.as_ref() else {
963 return Ok(None);
964 };
965 conditions.push(deferral.wait_condition().clone());
966 }
967 }
968 {
969 let decode = self.decode_queue.read();
970 for request in decode.requests.values() {
971 if request
972 .execution_readiness_block
973 .as_ref()
974 .is_some_and(|block| block.status() != EXECUTION_READINESS_CANCELLED)
975 {
976 continue;
977 }
978 let Some(deferral) = request.execution_capacity_deferral.as_ref() else {
979 return Ok(None);
980 };
981 conditions.push(deferral.wait_condition().clone());
982 }
983 }
984 let waiting_queue = self.waiting_queue.read();
985 let pressure_hold_is_active = |request: &ContinuousBatchRequest| {
986 self.pressure_active.load(Ordering::Acquire)
987 && matches!(
988 self.pressure_coordinator
989 .lock()
990 .hold_status(&request.inner.request.id),
991 PressureHoldStatus::Held { .. }
992 )
993 };
994 let waiting_count = waiting_queue
995 .iter()
996 .filter(|request| !pressure_hold_is_active(request))
997 .count();
998 let waiting = waiting_queue
999 .passive_wait_condition_for(|request| !pressure_hold_is_active(request))
1000 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
1001 drop(waiting_queue);
1002 if waiting_count > 0 && waiting.is_none() {
1003 return Ok(None);
1004 }
1005 if let Some(waiting) = waiting {
1006 conditions.push(waiting);
1007 }
1008 if conditions.is_empty() {
1009 return Ok(None);
1010 }
1011
1012 let coordinator = conditions[0].coordinator_id();
1013 let mut observed_by_source = BTreeMap::new();
1014 for condition in conditions {
1015 if condition.coordinator_id() != coordinator {
1016 return Err(FerrumError::scheduler(
1017 "passive capacity waits belong to different coordinators",
1018 ));
1019 }
1020 for observed in condition.observed() {
1021 observed_by_source
1022 .entry(observed.source())
1023 .and_modify(|epoch: &mut u64| *epoch = (*epoch).min(observed.epoch()))
1024 .or_insert(observed.epoch());
1025 }
1026 }
1027 let observed = observed_by_source
1028 .into_iter()
1029 .map(|(source, epoch)| {
1030 ferrum_interfaces::vnext::CapacityAvailabilityEpoch::new(source, epoch)
1031 .map_err(|error| FerrumError::scheduler(error.to_string()))
1032 })
1033 .collect::<Result<Vec<_>>>()?;
1034 let condition = ferrum_interfaces::vnext::CapacityWaitCondition::new(coordinator, observed)
1035 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
1036 let pressure = self.pressure_coordinator.lock();
1037 if pressure.has_pending_release_for(&condition) {
1038 return Ok(None);
1039 }
1040 if pressure.all_blocked_without_release_for(&condition) {
1041 return Err(FerrumError::scheduler(
1042 "capacity pressure contract reached all blocked frontiers without a pending release",
1043 ));
1044 }
1045 Ok(Some(condition))
1046 }
1047
1048 pub fn all_active_execution_readiness_blocked(&self) -> bool {
1052 let prefill = self.prefill_queue.read();
1053 let decode = self.decode_queue.read();
1054 let active = prefill.len() + decode.requests.len();
1055 active != 0
1056 && prefill
1057 .iter()
1058 .chain(decode.requests.values())
1059 .all(|request| {
1060 request
1061 .execution_readiness_block
1062 .as_ref()
1063 .is_some_and(|block| {
1064 matches!(
1065 block.status(),
1066 EXECUTION_READINESS_PENDING | EXECUTION_READINESS_FAILED
1067 )
1068 })
1069 })
1070 }
1071
1072 pub fn decoding_count(&self) -> usize {
1074 self.decode_queue.read().requests.len()
1075 }
1076
1077 pub fn prefilling_count(&self) -> usize {
1079 self.prefill_queue.read().len()
1080 }
1081
1082 pub fn trace_snapshot(&self) -> ContinuousSchedulerTraceSnapshot {
1084 self.trace_snapshot_with_prefill_read_observer(|| {})
1085 }
1086
1087 fn trace_snapshot_with_prefill_read_observer(
1088 &self,
1089 prefill_read_observer: impl FnOnce(),
1090 ) -> ContinuousSchedulerTraceSnapshot {
1091 let waiting_queue_len = self.waiting_queue.read().len();
1092 let (
1097 prefill_queue_len,
1098 execution_capacity_blocked_prefill_len,
1099 execution_readiness_blocked_prefill_len,
1100 ) = {
1101 let prefill_queue = self.prefill_queue.read();
1102 let counts = (
1103 prefill_queue.len(),
1104 prefill_queue
1105 .iter()
1106 .filter(|request| request.execution_capacity_deferral.is_some())
1107 .count(),
1108 prefill_queue
1109 .iter()
1110 .filter(|request| request.execution_readiness_block.is_some())
1111 .count(),
1112 );
1113 prefill_read_observer();
1114 counts
1115 };
1116 let (
1117 decode_queue_len,
1118 decode_selection_cursor,
1119 execution_capacity_blocked_decode_len,
1120 execution_readiness_blocked_decode_len,
1121 ) = {
1122 let decode_queue = self.decode_queue.read();
1123 (
1124 decode_queue.requests.len(),
1125 decode_queue.selection_cursor.clone(),
1126 decode_queue
1127 .requests
1128 .values()
1129 .filter(|request| request.execution_capacity_deferral.is_some())
1130 .count(),
1131 decode_queue
1132 .requests
1133 .values()
1134 .filter(|request| request.execution_readiness_block.is_some())
1135 .count(),
1136 )
1137 };
1138 let preempted_queue_len = self.preempted_requests.read().len();
1139 let pressure = self.pressure_coordinator.lock().stats();
1140 let (decode_capacity_backpressure_admit_limit, decode_execution_pressure_enforced) = {
1141 let _feedback = self.decode_capacity_feedback_lock.lock();
1142 (
1143 self.decode_capacity_backpressure_limit(),
1144 self.decode_execution_pressure_enforced
1145 .load(Ordering::Acquire),
1146 )
1147 };
1148
1149 ContinuousSchedulerTraceSnapshot {
1150 current_iteration: self.current_iteration.load(Ordering::Relaxed),
1151 waiting_queue_len,
1152 prefill_queue_len,
1153 decode_queue_len,
1154 decode_selection_cursor,
1155 preempted_queue_len,
1156 active_len: prefill_queue_len + decode_queue_len,
1157 completed_total: self.completed_counter.load(Ordering::Relaxed),
1158 failed_total: self.failed_counter.load(Ordering::Relaxed),
1159 cancelled_total: self.cancelled_counter.load(Ordering::Relaxed),
1160 preempted_total: self.preempted_counter.load(Ordering::Relaxed),
1161 admitted_total: self.admitted_counter.load(Ordering::Relaxed),
1162 capacity_deferred_total: self.capacity_deferred_counter.load(Ordering::Relaxed),
1163 capacity_backpressure_admit_limit: self.capacity_backpressure_admit_limit(),
1164 decode_capacity_backpressure_admit_limit,
1165 decode_execution_pressure_enforced,
1166 capacity_blocked_waiting_len: self.capacity_blocked_waiting_len(),
1167 execution_capacity_blocked_prefill_len,
1168 execution_capacity_blocked_decode_len,
1169 execution_readiness_deferred_total: self
1170 .execution_readiness_deferred_counter
1171 .load(Ordering::Relaxed),
1172 execution_readiness_blocked_prefill_len,
1173 execution_readiness_blocked_decode_len,
1174 capacity_release_epoch: self.capacity_release_epoch.load(Ordering::Relaxed),
1175 capacity_mixed_recompute_epoch: self
1176 .capacity_mixed_recompute_epoch
1177 .load(Ordering::Relaxed),
1178 capacity_mixed_recompute_blocked_until_epoch: self
1179 .capacity_mixed_recompute_blocked_until_epoch
1180 .load(Ordering::Relaxed),
1181 capacity_mixed_recompute_required_blocks_per_slot: match self
1182 .capacity_mixed_recompute_required_blocks_per_slot
1183 .load(Ordering::Relaxed)
1184 {
1185 0 => None,
1186 value => Some(value),
1187 },
1188 capacity_mixed_recompute_observed_free_blocks: match self
1189 .capacity_mixed_recompute_observed_free_blocks
1190 .load(Ordering::Relaxed)
1191 {
1192 usize::MAX => None,
1193 value => Some(value),
1194 },
1195 legacy_waiting_admission_ticks: self
1196 .legacy_waiting_admission_ticks
1197 .load(Ordering::Relaxed),
1198 dynamic_admission_ticks: self.dynamic_admission_ticks.load(Ordering::Relaxed),
1199 dynamic_admission_probes: self.dynamic_admission_probes.load(Ordering::Relaxed),
1200 dynamic_admission_skipped_unchanged: self
1201 .dynamic_admission_skipped_unchanged
1202 .load(Ordering::Relaxed),
1203 dynamic_admission_deferred: self.dynamic_admission_deferred.load(Ordering::Relaxed),
1204 dynamic_backing_growth_requested: self
1205 .dynamic_backing_growth_requested
1206 .load(Ordering::Relaxed),
1207 dynamic_admission_failed: self.dynamic_admission_failed.load(Ordering::Relaxed),
1208 pressure_episodes_created: pressure.episodes_created,
1209 pressure_episodes_merged: pressure.episodes_merged,
1210 pressure_episode_bridges_deferred: pressure.episode_bridges_deferred,
1211 pressure_active_episodes: pressure.active_episodes,
1212 pressure_pending_release_fences: pressure.pending_release_fences,
1213 pressure_candidate_scans: pressure.candidate_scans,
1214 pressure_last_transition_ordinal: pressure.last_transition_ordinal,
1215 pressure_dropped_journal_entries: pressure.dropped_journal_entries,
1216 }
1217 }
1218
1219 pub fn admission_phase_counts(&self) -> ContinuousSchedulerAdmissionCounts {
1221 let request_index = self.request_index.read();
1222 let mut counts = ContinuousSchedulerAdmissionCounts {
1223 waiting_requests: 0,
1224 active_prefill_sequences: 0,
1225 active_decode_sequences: 0,
1226 };
1227 for phase in request_index.values() {
1228 match phase {
1229 RequestPhase::Waiting => counts.waiting_requests += 1,
1230 RequestPhase::Prefilling => counts.active_prefill_sequences += 1,
1231 RequestPhase::Decoding => counts.active_decode_sequences += 1,
1232 RequestPhase::Completed
1233 | RequestPhase::Preempted
1234 | RequestPhase::Cancelled
1235 | RequestPhase::AdmissionFailed => {}
1236 }
1237 }
1238 counts
1239 }
1240
1241 pub fn trace_phase(&self, request_id: &RequestId) -> Option<RequestPhase> {
1243 self.request_index.read().get(request_id).copied()
1244 }
1245
1246 pub fn pressure_transition_journal(&self) -> Vec<PressureTransition> {
1248 self.pressure_coordinator.lock().journal()
1249 }
1250
1251 fn requeue_waiting_request(
1252 &self,
1253 waiting_queue: &mut DynamicAdmissionQueue<ContinuousBatchRequest>,
1254 request_index: &mut HashMap<RequestId, RequestPhase>,
1255 mut request: ContinuousBatchRequest,
1256 ) -> bool {
1257 let request_id = request.inner.request.id.clone();
1258 let Some(ticket) = request.waiting_admission_ticket else {
1259 let error = FerrumError::scheduler(format!(
1260 "request {request_id} lost its waiting admission identity"
1261 ));
1262 request.phase = RequestPhase::AdmissionFailed;
1263 request.inner.state = RequestState::Failed;
1264 request_index.insert(request_id.clone(), RequestPhase::AdmissionFailed);
1265 self.admission_failed_requests
1266 .write()
1267 .insert(request_id.clone(), request);
1268 self.admission_failures
1269 .lock()
1270 .push_back((request_id, error));
1271 self.dynamic_admission_failed
1272 .fetch_add(1, Ordering::Relaxed);
1273 return false;
1274 };
1275 let result = waiting_queue.requeue(ticket, request);
1276 match result {
1277 Ok(()) => {
1278 request_index.insert(request_id, RequestPhase::Waiting);
1279 true
1280 }
1281 Err((error, mut request)) => {
1282 let error = FerrumError::scheduler(error.to_string());
1283 request.phase = RequestPhase::AdmissionFailed;
1284 request.inner.state = RequestState::Failed;
1285 request_index.insert(request_id.clone(), RequestPhase::AdmissionFailed);
1286 self.admission_failed_requests
1287 .write()
1288 .insert(request_id.clone(), request);
1289 self.admission_failures
1290 .lock()
1291 .push_back((request_id, error));
1292 self.dynamic_admission_failed
1293 .fetch_add(1, Ordering::Relaxed);
1294 false
1295 }
1296 }
1297 }
1298
1299 fn promote_to_prefill_with_empty_retry(
1300 &self,
1301 request_id: &RequestId,
1302 empty_retry_epoch: Option<u64>,
1303 ) -> bool {
1304 let mut waiting_queue = self.waiting_queue.write();
1305 let mut prefill_queue = self.prefill_queue.write();
1306 let mut request_index = self.request_index.write();
1307
1308 let waiting_position =
1309 waiting_queue.position(|request| request.inner.request.id == *request_id);
1310 if let Some(pos) = waiting_position {
1311 let mut req = waiting_queue.remove(pos).unwrap();
1312 if let Some(epoch) = empty_retry_epoch {
1313 req.capacity_deferred_empty_retry_epoch = Some(epoch);
1314 }
1315 req.logical_work_frontier
1316 .begin_prefill(req.capacity_deferred_from_decode);
1317 req.phase = RequestPhase::Prefilling;
1318 req.inner.state = RequestState::Running;
1319 let started_at = chrono::Utc::now();
1320 let wait_us = started_at
1321 .signed_duration_since(req.inner.submitted_at)
1322 .num_microseconds()
1323 .unwrap_or(0)
1324 .max(0) as u64;
1325 req.inner.started_at = Some(started_at);
1326 self.total_wait_time_us
1327 .fetch_add(wait_us, Ordering::Relaxed);
1328 self.admitted_counter.fetch_add(1, Ordering::Relaxed);
1329
1330 request_index.insert(request_id.clone(), RequestPhase::Prefilling);
1331 prefill_queue.push_back(req);
1332 self.consume_pressure_hold(request_id);
1333
1334 debug!("Promoted request {} to prefill queue", request_id);
1335 true
1336 } else {
1337 false
1338 }
1339 }
1340
1341 fn promote_admitted_request(
1342 &self,
1343 mut request: ContinuousBatchRequest,
1344 receipt: &ExecutorPrefillAdmissionReceipt,
1345 ) {
1346 let request_id = request.inner.request.id.clone();
1347 if receipt.request_id != request_id {
1348 self.fail_typed_admission(
1349 request,
1350 FerrumError::scheduler(format!(
1351 "executor admission receipt belongs to {}, expected {}",
1352 receipt.request_id, request_id
1353 )),
1354 );
1355 return;
1356 }
1357 request
1358 .logical_work_frontier
1359 .begin_prefill(request.capacity_deferred_from_decode);
1360 request.phase = RequestPhase::Prefilling;
1361 request.inner.state = RequestState::Running;
1362 let started_at = chrono::Utc::now();
1363 let wait_us = started_at
1364 .signed_duration_since(request.inner.submitted_at)
1365 .num_microseconds()
1366 .unwrap_or(0)
1367 .max(0) as u64;
1368 request.inner.started_at = Some(started_at);
1369 self.total_wait_time_us
1370 .fetch_add(wait_us, Ordering::Relaxed);
1371 self.admitted_counter.fetch_add(1, Ordering::Relaxed);
1372 self.request_index
1373 .write()
1374 .insert(request_id.clone(), RequestPhase::Prefilling);
1375 self.prefill_queue.write().push_back(request);
1376 self.consume_pressure_hold(&request_id);
1377 debug!("Typed admission promoted request {} to prefill", request_id);
1378 }
1379
1380 fn fail_typed_admission(&self, mut request: ContinuousBatchRequest, error: FerrumError) {
1381 let request_id = request.inner.request.id.clone();
1382 request.logical_work_frontier.finish();
1383 self.record_pressure_frontier_terminal(&request_id);
1384 request.phase = RequestPhase::AdmissionFailed;
1385 request.inner.state = RequestState::Failed;
1386 self.request_index
1387 .write()
1388 .insert(request_id.clone(), RequestPhase::AdmissionFailed);
1389 self.admission_failed_requests
1390 .write()
1391 .insert(request_id.clone(), request);
1392 self.admission_failures
1393 .lock()
1394 .push_back((request_id, error));
1395 self.dynamic_admission_failed
1396 .fetch_add(1, Ordering::Relaxed);
1397 }
1398
1399 fn admit_waiting_dynamically(
1400 &self,
1401 maximum_probes: usize,
1402 maximum_admissions: usize,
1403 waiting_admission: &mut WaitingAdmissionMode<'_>,
1404 ) -> Result<AdmissionTickReceipt> {
1405 let WaitingAdmissionMode::Dynamic {
1406 wake,
1407 probe,
1408 observer,
1409 } = waiting_admission
1410 else {
1411 return Err(FerrumError::scheduler(
1412 "dynamic admission requires a typed wake/probe mode",
1413 ));
1414 };
1415 let wake = *wake;
1416 self.dynamic_admission_ticks.fetch_add(1, Ordering::Relaxed);
1417 let mut waiting = self.waiting_queue.write();
1418 let maximum_probes = waiting.len().min(maximum_probes);
1419 let mut events = self.dynamic_admission_events.lock();
1420 let observer = std::cell::RefCell::new(observer);
1421 let receipt = waiting
1422 .schedule_into_observed_with_eligibility(
1423 wake,
1424 maximum_probes,
1425 maximum_admissions,
1426 &mut events,
1427 |request, ticket| {
1428 let request_id = &request.inner.request.id;
1429 if !self.pressure_active.load(Ordering::Acquire) {
1430 return AdmissionQueueEligibility::Eligible;
1431 }
1432 let hold_status = {
1433 let coordinator = self.pressure_coordinator.lock();
1434 coordinator.hold_status(request_id)
1435 };
1436 match hold_status {
1437 PressureHoldStatus::Held { .. } => AdmissionQueueEligibility::Held,
1438 PressureHoldStatus::OwnerAdmissionEligible { .. } => {
1439 AdmissionQueueEligibility::Eligible
1440 }
1441 PressureHoldStatus::Released {
1442 episode_id,
1443 progress_owner_id,
1444 progress_baseline,
1445 progress_current,
1446 reason,
1447 ordinal,
1448 previous_wait_condition,
1449 current_wait_condition,
1450 } => {
1451 {
1452 let mut coordinator = self.pressure_coordinator.lock();
1453 if let Err(error) = coordinator.consume_released_hold(request_id) {
1454 warn!(
1455 request_id = %request_id,
1456 error = %error,
1457 "Pressure coordinator rejected terminal hold release"
1458 );
1459 return AdmissionQueueEligibility::Held;
1460 }
1461 self.pressure_active
1462 .store(coordinator.has_records(), Ordering::Release);
1463 }
1464 if let Some(observer) = observer.borrow_mut().as_deref_mut() {
1465 observer(ExecutorAdmissionQueueObservation::PressureHoldReleased {
1466 episode_id,
1467 transition_ordinal: ordinal,
1468 request_id: request_id.clone(),
1469 progress_owner_id,
1470 progress_baseline,
1471 progress_current,
1472 reason,
1473 previous_wait_condition,
1474 current_wait_condition,
1475 ticket: ticket.get(),
1476 });
1477 }
1478 AdmissionQueueEligibility::Eligible
1479 }
1480 PressureHoldStatus::None => AdmissionQueueEligibility::Eligible,
1481 }
1482 },
1483 |_request, _ticket| {},
1484 |request, ticket, deferral| {
1485 if let Some(observer) = observer.borrow_mut().as_deref_mut() {
1486 observer(ExecutorAdmissionQueueObservation::SkippedUnchanged {
1487 request_id: request.inner.request.id.clone(),
1488 ticket: ticket.get(),
1489 deferral: deferral.clone(),
1490 current: wake.epochs(),
1491 });
1492 }
1493 },
1494 |request| probe(&request.inner.request),
1495 )
1496 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
1497 drop(waiting);
1498
1499 self.dynamic_admission_probes
1500 .fetch_add(receipt.probed() as u64, Ordering::Relaxed);
1501 self.dynamic_admission_skipped_unchanged
1502 .fetch_add(receipt.skipped_unchanged() as u64, Ordering::Relaxed);
1503 self.dynamic_admission_deferred
1504 .fetch_add(receipt.deferred() as u64, Ordering::Relaxed);
1505 self.dynamic_backing_growth_requested
1506 .fetch_add(receipt.backing_growth_requested() as u64, Ordering::Relaxed);
1507
1508 for event in events.drain(..) {
1509 match event {
1510 AdmissionQueueEvent::Admitted {
1511 request, admission, ..
1512 } => self.promote_admitted_request(request, &admission),
1513 AdmissionQueueEvent::PermanentRejected {
1514 request, rejection, ..
1515 } => self.fail_typed_admission(
1516 request,
1517 FerrumError::request_validation(format!(
1518 "request cannot fit the vNext runtime: {rejection:?}"
1519 )),
1520 ),
1521 AdmissionQueueEvent::Faulted { request, error, .. } => {
1522 self.fail_typed_admission(request, error)
1523 }
1524 AdmissionQueueEvent::ContractFaulted { request, error, .. } => {
1525 self.fail_typed_admission(request, FerrumError::scheduler(error.to_string()))
1526 }
1527 AdmissionQueueEvent::PreemptionRequested { .. } => {}
1528 AdmissionQueueEvent::BackingGrowthRequested { .. } => {}
1529 }
1530 }
1531 Ok(receipt)
1532 }
1533
1534 pub fn take_admission_failures(&self) -> Vec<(RequestId, FerrumError)> {
1535 self.admission_failures.lock().drain(..).collect()
1536 }
1537
1538 pub fn fail_waiting_admission(&self, request_id: &RequestId, error: FerrumError) -> bool {
1542 let request = {
1543 let mut waiting = self.waiting_queue.write();
1544 waiting
1545 .position(|request| request.inner.request.id == *request_id)
1546 .and_then(|position| waiting.remove(position))
1547 };
1548 let Some(request) = request else {
1549 return false;
1550 };
1551 self.fail_typed_admission(request, error);
1552 true
1553 }
1554
1555 pub fn wait_for_release_after_backing_pressure(
1558 &self,
1559 request_id: &RequestId,
1560 observed: AdmissionWakeEpochs,
1561 wait_condition: ferrum_interfaces::vnext::CapacityWaitCondition,
1562 ) -> Result<bool> {
1563 self.waiting_queue
1564 .write()
1565 .wait_for_release_after_backing_pressure(
1566 |request| request.inner.request.id == *request_id,
1567 observed,
1568 wait_condition,
1569 )
1570 .map_err(|error| FerrumError::scheduler(error.to_string()))
1571 }
1572
1573 pub fn retry_after_backing_recheck(
1574 &self,
1575 request_id: &RequestId,
1576 observed: AdmissionWakeEpochs,
1577 ) -> Result<bool> {
1578 self.waiting_queue
1579 .write()
1580 .retry_after_backing_recheck(
1581 |request| request.inner.request.id == *request_id,
1582 observed,
1583 )
1584 .map_err(|error| FerrumError::scheduler(error.to_string()))
1585 }
1586
1587 pub fn next_batch_with_dynamic_admission(
1588 &self,
1589 hint: BatchHint,
1590 wake: AdmissionWakeSnapshot<'_>,
1591 probe: &mut dyn FnMut(&InferenceRequest) -> ExecutorAdmissionProbeOutcome,
1592 ) -> Result<Option<BatchPlan>> {
1593 self.create_iteration_batch_with_admission(
1594 hint,
1595 WaitingAdmissionMode::Dynamic {
1596 wake,
1597 probe,
1598 observer: None,
1599 },
1600 )
1601 }
1602
1603 pub fn next_batch_with_dynamic_admission_observed(
1604 &self,
1605 hint: BatchHint,
1606 wake: AdmissionWakeSnapshot<'_>,
1607 probe: &mut dyn FnMut(&InferenceRequest) -> ExecutorAdmissionProbeOutcome,
1608 observer: &mut dyn FnMut(ExecutorAdmissionQueueObservation),
1609 ) -> Result<Option<BatchPlan>> {
1610 self.create_iteration_batch_with_admission(
1611 hint,
1612 WaitingAdmissionMode::Dynamic {
1613 wake,
1614 probe,
1615 observer: Some(observer),
1616 },
1617 )
1618 }
1619
1620 pub fn prepare_dynamic_admission_observed(
1626 &self,
1627 maximum_admissions: usize,
1628 wake: AdmissionWakeSnapshot<'_>,
1629 probe: &mut dyn FnMut(&InferenceRequest) -> ExecutorAdmissionProbeOutcome,
1630 observer: &mut dyn FnMut(ExecutorAdmissionQueueObservation),
1631 ) -> Result<AdmissionTickReceipt> {
1632 let active_capacity = self
1633 .config
1634 .max_running_requests
1635 .saturating_sub(self.active_count());
1636 let decode_capacity = self
1637 .cb_config
1638 .max_decode_batch
1639 .saturating_sub(self.decoding_count());
1640 let available_slots = active_capacity.min(decode_capacity);
1641 let available_slots = self
1642 .capacity_backpressure_admit_limit()
1643 .map(|limit| available_slots.min(limit))
1644 .unwrap_or(available_slots)
1645 .min(maximum_admissions);
1646 let mut mode = WaitingAdmissionMode::Dynamic {
1647 wake,
1648 probe,
1649 observer: Some(observer),
1650 };
1651 self.admit_waiting_dynamically(available_slots, available_slots, &mut mode)
1652 }
1653
1654 pub fn fill_first_dynamic_admission_limit(&self, hint: &BatchHint, target: usize) -> usize {
1659 let mut remaining_tokens = hint.max_tokens;
1660 let prefill_step_chunk = self.runtime_config.prefill_step_chunk;
1661 let prefill_queue = self.prefill_queue.read();
1662 for request in prefill_queue.iter().take(hint.max_batch_size) {
1663 let tokens =
1664 self.prefill_budget_tokens(request, None, prefill_step_chunk, remaining_tokens);
1665 if tokens == 0 || tokens > remaining_tokens {
1666 return 0;
1667 }
1668 remaining_tokens -= tokens;
1669 }
1670 drop(prefill_queue);
1671
1672 let mut limit = target
1673 .saturating_sub(self.active_count())
1674 .min(hint.max_batch_size);
1675 if limit == 0 || remaining_tokens == 0 {
1676 return 0;
1677 }
1678 let waiting = self.waiting_queue.read();
1679 let mut admitted_tokens = 0usize;
1680 let mut admitted = 0usize;
1681 for request in waiting.iter() {
1682 if admitted >= limit {
1683 break;
1684 }
1685 let available = remaining_tokens.saturating_sub(admitted_tokens);
1686 let tokens = self.prefill_budget_tokens(request, None, prefill_step_chunk, available);
1687 if tokens == 0 || tokens > available {
1688 break;
1689 }
1690 admitted_tokens += tokens;
1691 admitted += 1;
1692 }
1693 limit = limit.min(admitted);
1694 limit
1695 }
1696
1697 fn capacity_backpressure_admit_limit(&self) -> Option<usize> {
1698 Self::read_backpressure_limit(&self.capacity_backpressure_limit)
1699 }
1700
1701 fn decode_capacity_backpressure_limit(&self) -> Option<usize> {
1702 Self::read_backpressure_limit(&self.decode_capacity_backpressure_limit)
1703 }
1704
1705 fn read_backpressure_limit(limit: &AtomicUsize) -> Option<usize> {
1706 let limit = limit.load(Ordering::Relaxed);
1707 if limit == NO_CAPACITY_BACKPRESSURE_LIMIT {
1708 None
1709 } else {
1710 Some(limit.max(1))
1711 }
1712 }
1713
1714 fn capacity_blocked_waiting_len(&self) -> usize {
1715 let has_active_requests = self.active_count() > 0;
1716 let release_epoch = self.capacity_release_epoch.load(Ordering::Relaxed);
1717 self.waiting_queue
1718 .read()
1719 .iter()
1720 .filter(|req| {
1721 req.capacity_deferred_until_release_epoch > release_epoch
1722 && (has_active_requests
1723 || req.capacity_deferred_empty_retry_epoch == Some(release_epoch))
1724 })
1725 .count()
1726 }
1727
1728 fn decode_capacity_deferred_backlog_len(&self) -> usize {
1729 let waiting = self
1730 .waiting_queue
1731 .read()
1732 .iter()
1733 .filter(|req| req.capacity_deferred_from_decode)
1734 .count();
1735 let prefilling = self
1736 .prefill_queue
1737 .read()
1738 .iter()
1739 .filter(|req| req.capacity_deferred_from_decode)
1740 .count();
1741 waiting + prefilling
1742 }
1743
1744 fn record_capacity_defer_feedback(&self, attempted_prefill_width: usize) {
1745 self.capacity_deferred_counter
1746 .fetch_add(1, Ordering::Relaxed);
1747
1748 let iteration = self.current_iteration.load(Ordering::Relaxed);
1749 let previous_iteration = self
1750 .capacity_backpressure_iteration
1751 .swap(iteration, Ordering::Relaxed);
1752 if previous_iteration == iteration {
1753 return;
1754 }
1755
1756 let max_running = self.config.max_running_requests.max(1);
1757 let proposed = attempted_prefill_width
1758 .max(1)
1759 .div_ceil(2)
1760 .max(1)
1761 .min(max_running);
1762 let _ = self.capacity_backpressure_limit.fetch_update(
1763 Ordering::Relaxed,
1764 Ordering::Relaxed,
1765 |current| {
1766 let current = if current == NO_CAPACITY_BACKPRESSURE_LIMIT {
1767 max_running
1768 } else {
1769 current.max(1).min(max_running)
1770 };
1771 let next = proposed.min(current).max(1);
1772 if next >= max_running {
1773 Some(NO_CAPACITY_BACKPRESSURE_LIMIT)
1774 } else {
1775 Some(next)
1776 }
1777 },
1778 );
1779 }
1780
1781 fn decode_capacity_pressure_limit(
1782 attempted_decode_width: usize,
1783 observed_free_blocks: Option<usize>,
1784 max_running: usize,
1785 ) -> usize {
1786 let attempted = attempted_decode_width.max(1).min(max_running);
1787 let half_width = attempted.div_ceil(2).max(1);
1788 let near_fit_width = observed_free_blocks
1789 .filter(|free_blocks| *free_blocks > 0)
1790 .map(|free_blocks| {
1791 let usable_free_blocks =
1792 free_blocks.saturating_sub(CAPACITY_DECODE_FREE_BLOCK_HEADROOM);
1793 usable_free_blocks
1794 .max(1)
1795 .min(attempted.saturating_sub(1).max(1))
1796 });
1797 near_fit_width
1798 .unwrap_or(half_width)
1799 .max(half_width)
1800 .min(max_running)
1801 }
1802
1803 pub fn record_decode_capacity_pressure(
1804 &self,
1805 attempted_decode_width: usize,
1806 observed_free_blocks: Option<usize>,
1807 ) {
1808 let _feedback = self.decode_capacity_feedback_lock.lock();
1809 self.record_decode_capacity_pressure_inner(attempted_decode_width, observed_free_blocks);
1810 }
1811
1812 fn record_decode_capacity_pressure_inner(
1813 &self,
1814 attempted_decode_width: usize,
1815 observed_free_blocks: Option<usize>,
1816 ) {
1817 let max_running = self.config.max_running_requests.max(1);
1818 let proposed = Self::decode_capacity_pressure_limit(
1819 attempted_decode_width,
1820 observed_free_blocks,
1821 max_running,
1822 );
1823 let _ = self.decode_capacity_backpressure_limit.fetch_update(
1824 Ordering::Relaxed,
1825 Ordering::Relaxed,
1826 |current| {
1827 let current = if current == NO_CAPACITY_BACKPRESSURE_LIMIT {
1828 max_running
1829 } else {
1830 current.max(1).min(max_running)
1831 };
1832 let next = proposed.min(current).max(1);
1833 if next >= max_running {
1834 Some(NO_CAPACITY_BACKPRESSURE_LIMIT)
1835 } else {
1836 Some(next)
1837 }
1838 },
1839 );
1840 }
1841
1842 pub fn record_decode_execution_capacity_pressure(&self, attempted_decode_width: usize) {
1843 let _feedback = self.decode_capacity_feedback_lock.lock();
1844 self.record_decode_capacity_pressure_inner(attempted_decode_width, None);
1845 self.decode_execution_recovery_release_epoch.store(
1846 self.capacity_release_epoch.load(Ordering::Relaxed),
1847 Ordering::Relaxed,
1848 );
1849 self.decode_execution_pressure_enforced
1850 .store(true, Ordering::Release);
1851 }
1852
1853 pub fn record_decode_execution_capacity_success(&self, successful_decode_width: usize) -> bool {
1859 if successful_decode_width == 0 {
1860 return false;
1861 }
1862 let _feedback = self.decode_capacity_feedback_lock.lock();
1863 if !self
1864 .decode_execution_pressure_enforced
1865 .load(Ordering::Acquire)
1866 {
1867 return false;
1868 }
1869
1870 let max_running = self.config.max_running_requests.max(1);
1871 let successful_decode_width = successful_decode_width.max(1).min(max_running);
1872 let release_epoch = self.capacity_release_epoch.load(Ordering::Relaxed);
1873 if release_epoch
1874 <= self
1875 .decode_execution_recovery_release_epoch
1876 .load(Ordering::Relaxed)
1877 {
1878 return false;
1879 }
1880 let relaxed = self
1881 .decode_capacity_backpressure_limit
1882 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
1883 if current == NO_CAPACITY_BACKPRESSURE_LIMIT {
1884 return Some(NO_CAPACITY_BACKPRESSURE_LIMIT);
1885 }
1886 let current = current.max(1).min(max_running);
1887 if successful_decode_width < current {
1888 return None;
1889 }
1890 let next = current.saturating_add(1).min(max_running);
1891 Some(if next >= max_running {
1892 NO_CAPACITY_BACKPRESSURE_LIMIT
1893 } else {
1894 next
1895 })
1896 })
1897 .is_ok();
1898
1899 if relaxed {
1900 self.decode_execution_recovery_release_epoch
1901 .store(release_epoch, Ordering::Relaxed);
1902 }
1903 if relaxed && self.decode_capacity_backpressure_limit().is_none() {
1904 self.decode_execution_pressure_enforced
1905 .store(false, Ordering::Release);
1906 }
1907 relaxed
1908 }
1909
1910 pub fn defer_prefill_for_execution_capacity(
1913 &self,
1914 request_id: &RequestId,
1915 deferral: AdmissionDeferral,
1916 release_snapshot: &ExecutionCapacityReleaseSnapshot,
1917 ) -> Result<ExecutionCapacityAction> {
1918 self.plan_execution_capacity_pressure(
1919 std::slice::from_ref(request_id),
1920 deferral,
1921 release_snapshot,
1922 )
1923 }
1924
1925 pub fn defer_decode_for_execution_capacity(
1928 &self,
1929 request_ids: &[RequestId],
1930 deferral: AdmissionDeferral,
1931 release_snapshot: &ExecutionCapacityReleaseSnapshot,
1932 ) -> Result<ExecutionCapacityAction> {
1933 self.plan_execution_capacity_pressure(request_ids, deferral, release_snapshot)
1934 }
1935
1936 pub fn defer_for_execution_readiness(
1941 &self,
1942 request_ids: &[RequestId],
1943 ) -> Result<ExecutionReadinessDeferralReceipt> {
1944 let requested = request_ids.iter().collect::<HashSet<_>>();
1945 if request_ids.is_empty() || requested.len() != request_ids.len() {
1946 return Err(FerrumError::scheduler(
1947 "execution readiness deferral requires a non-empty unique cohort",
1948 ));
1949 }
1950 let ticket_value = self
1951 .next_execution_readiness_ticket
1952 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
1953 current.checked_add(1)
1954 })
1955 .map_err(|_| FerrumError::scheduler("execution readiness ticket space exhausted"))?;
1956 let ticket_id = NonZeroU64::new(ticket_value)
1957 .ok_or_else(|| FerrumError::scheduler("execution readiness issued a zero ticket id"))?;
1958 let state = Arc::new(ExecutionReadinessState {
1959 status: AtomicU8::new(EXECUTION_READINESS_PENDING),
1960 });
1961 let block = ExecutionReadinessBlock {
1962 ticket_id,
1963 state: Arc::clone(&state),
1964 };
1965
1966 let mut prefill = self.prefill_queue.write();
1967 let mut decode = self.decode_queue.write();
1968 let existing = prefill
1969 .iter()
1970 .chain(decode.requests.values())
1971 .find(|request| {
1972 requested.contains(&request.inner.request.id)
1973 && request.execution_readiness_block.is_some()
1974 });
1975 if let Some(request) = existing {
1976 return Err(FerrumError::scheduler(format!(
1977 "request {} already owns an execution readiness ticket",
1978 request.inner.request.id
1979 )));
1980 }
1981 let mut installed = 0usize;
1982 for request in prefill.iter_mut().chain(decode.requests.values_mut()) {
1983 if requested.contains(&request.inner.request.id) {
1984 request.execution_readiness_block = Some(block.clone());
1985 installed += 1;
1986 }
1987 }
1988 if installed != request_ids.len() {
1989 Self::rollback_execution_readiness_install(&mut prefill, &mut decode, &block);
1990 return Err(FerrumError::scheduler(format!(
1991 "execution readiness retained {installed} of {} active frontiers",
1992 request_ids.len()
1993 )));
1994 }
1995 drop(decode);
1996 drop(prefill);
1997 self.execution_readiness_deferred_counter
1998 .fetch_add(installed as u64, Ordering::Relaxed);
1999 Ok(ExecutionReadinessDeferralReceipt {
2000 deferred_count: installed,
2001 wake: ExecutionReadinessWake { ticket_id, state },
2002 })
2003 }
2004
2005 fn rollback_execution_readiness_install(
2006 prefill: &mut VecDeque<ContinuousBatchRequest>,
2007 decode: &mut DecodeQueueState,
2008 block: &ExecutionReadinessBlock,
2009 ) {
2010 for request in prefill.iter_mut() {
2011 if request
2012 .execution_readiness_block
2013 .as_ref()
2014 .is_some_and(|installed| installed.matches(block))
2015 {
2016 request.execution_readiness_block = None;
2017 }
2018 }
2019 for request in decode.requests.values_mut() {
2020 if request
2021 .execution_readiness_block
2022 .as_ref()
2023 .is_some_and(|installed| installed.matches(block))
2024 {
2025 request.execution_readiness_block = None;
2026 }
2027 }
2028 }
2029
2030 pub fn defer_retry_after_execution_maintenance(
2035 &self,
2036 retry: &ExecutorExecutionMaintenanceRetry,
2037 ) -> Result<ExecutionMaintenanceRetryReceipt> {
2038 let progress = retry.progress();
2039 if progress.mutations().is_empty() {
2040 return Err(FerrumError::scheduler(
2041 "execution maintenance retry requires physical mutations",
2042 ));
2043 }
2044 self.defer_retry_after_execution_maintenance_epoch(
2045 retry.affected_request_ids(),
2046 progress.latest_capacity_epoch(),
2047 )
2048 }
2049
2050 fn defer_retry_after_execution_maintenance_epoch(
2051 &self,
2052 request_ids: &[RequestId],
2053 latest_capacity_epoch: u64,
2054 ) -> Result<ExecutionMaintenanceRetryReceipt> {
2055 if request_ids.is_empty() || latest_capacity_epoch == 0 {
2056 return Err(FerrumError::scheduler(
2057 "execution maintenance retry requires active requests and a physical mutation epoch",
2058 ));
2059 }
2060 let requested = request_ids.iter().cloned().collect::<HashSet<_>>();
2061 if requested.len() != request_ids.len() {
2062 return Err(FerrumError::scheduler(
2063 "execution maintenance retry contains duplicate request identities",
2064 ));
2065 }
2066
2067 let mut prefill = self.prefill_queue.write();
2071 let mut decode = self.decode_queue.write();
2072 let request_index = self.request_index.read();
2073 let validate = |request: &ContinuousBatchRequest| -> Result<()> {
2074 if request.execution_capacity_deferral.is_some()
2075 || request.execution_maintenance_retry.is_some()
2076 || request
2077 .last_execution_maintenance_capacity_epoch
2078 .is_some_and(|epoch| epoch >= latest_capacity_epoch)
2079 {
2080 return Err(FerrumError::scheduler(
2081 "execution maintenance retry reuses stale or concurrently blocked evidence",
2082 ));
2083 }
2084 Ok(())
2085 };
2086 for request_id in &requested {
2087 let prefill_matches = prefill
2088 .iter()
2089 .filter(|request| request.inner.request.id == *request_id)
2090 .count();
2091 let decode_request = decode.requests.get(request_id);
2092 let request = match request_index.get(request_id) {
2093 Some(RequestPhase::Prefilling)
2094 if prefill_matches == 1 && decode_request.is_none() =>
2095 {
2096 prefill
2097 .iter()
2098 .find(|request| request.inner.request.id == *request_id)
2099 .expect("validated prefill frontier remains locked")
2100 }
2101 Some(RequestPhase::Decoding)
2102 if prefill_matches == 0 && decode_request.is_some() =>
2103 {
2104 decode_request.expect("validated decode frontier remains locked")
2105 }
2106 _ => {
2107 return Err(FerrumError::scheduler(
2108 "execution maintenance retry lost an exact active logical frontier",
2109 ));
2110 }
2111 };
2112 validate(request)?;
2113 }
2114
2115 let not_before_iteration = self
2119 .current_iteration
2120 .load(Ordering::Relaxed)
2121 .saturating_add(1);
2122 let ticket = ExecutionMaintenanceRetryTicket {
2123 not_before_iteration,
2124 latest_capacity_epoch,
2125 };
2126 let mut deferred_count = 0;
2127 for request in prefill.iter_mut() {
2128 if requested.contains(&request.inner.request.id) {
2129 request.execution_maintenance_retry = Some(ticket);
2130 request.last_execution_maintenance_capacity_epoch = Some(latest_capacity_epoch);
2131 deferred_count += 1;
2132 }
2133 }
2134 for request in decode.requests.values_mut() {
2135 if requested.contains(&request.inner.request.id) {
2136 request.execution_maintenance_retry = Some(ticket);
2137 request.last_execution_maintenance_capacity_epoch = Some(latest_capacity_epoch);
2138 deferred_count += 1;
2139 }
2140 }
2141 if deferred_count != request_ids.len() {
2142 return Err(FerrumError::scheduler(format!(
2143 "execution maintenance retry retained {deferred_count} of {} active frontiers",
2144 request_ids.len()
2145 )));
2146 }
2147
2148 Ok(ExecutionMaintenanceRetryReceipt {
2149 deferred_count,
2150 not_before_iteration,
2151 latest_capacity_epoch,
2152 })
2153 }
2154
2155 fn plan_execution_capacity_pressure(
2156 &self,
2157 request_ids: &[RequestId],
2158 deferral: AdmissionDeferral,
2159 release_snapshot: &ExecutionCapacityReleaseSnapshot,
2160 ) -> Result<ExecutionCapacityAction> {
2161 if deferral.action() != ferrum_interfaces::vnext::DeferredAction::WaitForRelease {
2162 return Err(FerrumError::scheduler(
2163 "active execution-capacity deferral must wait for release",
2164 ));
2165 }
2166 let active_ids = {
2167 let request_index = self.request_index.read();
2168 request_ids
2169 .iter()
2170 .filter(|request_id| {
2171 matches!(
2172 request_index.get(*request_id),
2173 Some(RequestPhase::Prefilling | RequestPhase::Decoding)
2174 )
2175 })
2176 .cloned()
2177 .collect::<Vec<_>>()
2178 };
2179 if active_ids.is_empty() {
2180 return Ok(ExecutionCapacityAction::Deferred { count: 0 });
2181 }
2182
2183 let candidates =
2184 self.execution_capacity_candidates(release_snapshot, deferral.wait_condition());
2185 let decision = {
2186 let mut coordinator = self.pressure_coordinator.lock();
2187 let decision = coordinator
2188 .plan_failure(&active_ids, deferral.wait_condition(), &candidates)
2189 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
2190 self.pressure_active
2191 .store(coordinator.has_records(), Ordering::Release);
2192 decision
2193 };
2194
2195 match decision {
2196 PressureDecision::Deferred { count, .. } => {
2197 let installed =
2198 self.install_execution_capacity_deferral(&active_ids, &deferral, None);
2199 self.capacity_deferred_counter
2200 .fetch_add(installed as u64, Ordering::Relaxed);
2201 if installed != count {
2202 return Err(FerrumError::scheduler(format!(
2203 "execution-capacity deferral retained {installed} of {count} active frontiers"
2204 )));
2205 }
2206 Ok(ExecutionCapacityAction::Deferred { count: installed })
2207 }
2208 PressureDecision::YieldPlanned(transaction) => {
2209 let installed = self.install_execution_capacity_deferral(
2210 &active_ids,
2211 &deferral,
2212 Some(transaction.victim_request_id()),
2213 );
2214 self.capacity_deferred_counter
2215 .fetch_add(installed as u64, Ordering::Relaxed);
2216 Ok(ExecutionCapacityAction::YieldPlanned { transaction })
2217 }
2218 PressureDecision::InvariantViolation(violation) => {
2219 Ok(ExecutionCapacityAction::InvariantViolation { violation })
2220 }
2221 }
2222 }
2223
2224 fn execution_capacity_candidates(
2225 &self,
2226 release_snapshot: &ExecutionCapacityReleaseSnapshot,
2227 condition: &CapacityWaitCondition,
2228 ) -> Vec<PressureCandidate> {
2229 let mut candidates = Vec::new();
2230 {
2231 let prefill = self.prefill_queue.read();
2232 candidates.extend(prefill.iter().map(|request| {
2233 PressureCandidate {
2234 request_id: request.inner.request.id.clone(),
2235 work_kind: request.logical_work_frontier.work_kind(),
2236 priority: request.inner.request.priority,
2237 progress: request.logical_work_frontier.progress_generation(),
2238 recompute_cost: request.logical_work_frontier.recompute_cost(),
2239 advances_wait_source: release_snapshot
2240 .can_advance(&request.inner.request.id, condition),
2241 blocked_on: request
2242 .execution_capacity_deferral
2243 .as_ref()
2244 .map(|deferral| deferral.wait_condition().clone()),
2245 }
2246 }));
2247 }
2248 {
2249 let decode = self.decode_queue.read();
2250 candidates.extend(decode.requests.values().map(|request| {
2251 PressureCandidate {
2252 request_id: request.inner.request.id.clone(),
2253 work_kind: request.logical_work_frontier.work_kind(),
2254 priority: request.inner.request.priority,
2255 progress: request.logical_work_frontier.progress_generation(),
2256 recompute_cost: request.logical_work_frontier.recompute_cost(),
2257 advances_wait_source: release_snapshot
2258 .can_advance(&request.inner.request.id, condition),
2259 blocked_on: request
2260 .execution_capacity_deferral
2261 .as_ref()
2262 .map(|deferral| deferral.wait_condition().clone()),
2263 }
2264 }));
2265 }
2266 if self.pressure_active.load(Ordering::Acquire) {
2267 let waiting = self.waiting_queue.read();
2268 candidates.extend(waiting.iter().map(|request| {
2269 PressureCandidate {
2270 request_id: request.inner.request.id.clone(),
2271 work_kind: request.logical_work_frontier.work_kind(),
2272 priority: request.inner.request.priority,
2273 progress: request.logical_work_frontier.progress_generation(),
2274 recompute_cost: request.logical_work_frontier.recompute_cost(),
2275 advances_wait_source: false,
2276 blocked_on: request
2277 .execution_capacity_deferral
2278 .as_ref()
2279 .map(|deferral| deferral.wait_condition().clone()),
2280 }
2281 }));
2282 }
2283 candidates
2284 }
2285
2286 fn install_execution_capacity_deferral(
2287 &self,
2288 request_ids: &[RequestId],
2289 deferral: &AdmissionDeferral,
2290 yielding: Option<&RequestId>,
2291 ) -> usize {
2292 let requested = request_ids.iter().collect::<HashSet<_>>();
2293 let mut installed = 0usize;
2294 {
2295 let mut prefill = self.prefill_queue.write();
2296 for request in prefill.iter_mut() {
2297 let request_id = &request.inner.request.id;
2298 if requested.contains(request_id) && yielding != Some(request_id) {
2299 request.execution_capacity_deferral = Some(deferral.clone());
2300 installed += 1;
2301 }
2302 }
2303 }
2304 {
2305 let mut decode = self.decode_queue.write();
2306 for request in decode.requests.values_mut() {
2307 let request_id = &request.inner.request.id;
2308 if requested.contains(request_id) && yielding != Some(request_id) {
2309 request.execution_capacity_deferral = Some(deferral.clone());
2310 installed += 1;
2311 }
2312 }
2313 }
2314 installed
2315 }
2316
2317 fn relax_backpressure_limit(limit: &AtomicUsize, max_running: usize) {
2318 let current = limit.load(Ordering::Relaxed);
2319 if current == NO_CAPACITY_BACKPRESSURE_LIMIT {
2320 return;
2321 }
2322
2323 let current = current.max(1).min(max_running);
2324 let grown = current.saturating_mul(2).min(max_running);
2325 let next = if grown >= max_running {
2326 NO_CAPACITY_BACKPRESSURE_LIMIT
2327 } else {
2328 grown.max(1)
2329 };
2330 limit.store(next, Ordering::Relaxed);
2331 }
2332
2333 fn record_resource_progress(&self) {
2334 let max_running = self.config.max_running_requests.max(1);
2335 Self::relax_backpressure_limit(&self.capacity_backpressure_limit, max_running);
2336 let _feedback = self.decode_capacity_feedback_lock.lock();
2337 if !self
2338 .decode_execution_pressure_enforced
2339 .load(Ordering::Acquire)
2340 {
2341 Self::relax_backpressure_limit(&self.decode_capacity_backpressure_limit, max_running);
2342 }
2343 }
2344
2345 fn record_capacity_release_progress(&self) {
2346 self.capacity_release_epoch.fetch_add(1, Ordering::Relaxed);
2347 self.capacity_mixed_recompute_epoch
2348 .fetch_add(1, Ordering::Relaxed);
2349 self.capacity_mixed_recompute_required_blocks_per_slot
2350 .store(0, Ordering::Relaxed);
2351 self.capacity_mixed_recompute_observed_free_blocks
2352 .store(usize::MAX, Ordering::Relaxed);
2353 self.record_resource_progress();
2354 }
2355
2356 pub fn record_external_capacity_release(&self) {
2358 self.record_capacity_release_progress();
2359 }
2360
2361 fn record_capacity_recompute_progress(&self) {
2362 self.capacity_mixed_recompute_epoch
2363 .fetch_add(1, Ordering::Relaxed);
2364 }
2365
2366 fn capacity_mixed_recompute_usable_free_blocks(
2367 observed_free_blocks: usize,
2368 required_blocks_per_slot: usize,
2369 ) -> usize {
2370 if required_blocks_per_slot == 0 || observed_free_blocks == usize::MAX {
2371 return observed_free_blocks;
2372 }
2373 observed_free_blocks.saturating_sub(CAPACITY_MIXED_RECOMPUTE_FREE_BLOCK_HEADROOM)
2374 }
2375
2376 pub fn record_capacity_deferred_mixed_recompute_release_evidence(&self) {
2377 self.capacity_mixed_recompute_required_blocks_per_slot
2378 .store(0, Ordering::Relaxed);
2379 self.capacity_mixed_recompute_observed_free_blocks
2380 .store(usize::MAX, Ordering::Relaxed);
2381 self.record_capacity_recompute_progress();
2382 }
2383
2384 pub fn record_capacity_deferred_mixed_recompute_kv_capacity_snapshot(
2385 &self,
2386 free_blocks: usize,
2387 ) {
2388 self.capacity_mixed_recompute_observed_free_blocks
2389 .store(free_blocks, Ordering::Relaxed);
2390 let required_blocks_per_slot = self
2391 .capacity_mixed_recompute_required_blocks_per_slot
2392 .load(Ordering::Relaxed);
2393 let usable_free_blocks = Self::capacity_mixed_recompute_usable_free_blocks(
2394 free_blocks,
2395 required_blocks_per_slot,
2396 );
2397 if required_blocks_per_slot > 0 && usable_free_blocks >= required_blocks_per_slot {
2398 self.record_capacity_recompute_progress();
2399 }
2400 }
2401
2402 pub fn defer_capacity_deferred_mixed_recompute_until_release(&self) {
2409 self.capacity_mixed_recompute_required_blocks_per_slot
2410 .store(0, Ordering::Relaxed);
2411 self.capacity_mixed_recompute_observed_free_blocks
2412 .store(usize::MAX, Ordering::Relaxed);
2413 self.defer_capacity_deferred_mixed_recompute_until_kv_capacity(None, None, None);
2414 }
2415
2416 pub fn defer_capacity_deferred_mixed_recompute_until_kv_capacity(
2426 &self,
2427 required_admission_blocks: Option<usize>,
2428 observed_free_blocks: Option<usize>,
2429 attempted_prefill_width: Option<usize>,
2430 ) {
2431 let mixed_epoch = self.capacity_mixed_recompute_epoch.load(Ordering::Relaxed);
2432 let blocked_until = mixed_epoch.saturating_add(1);
2433 let mut required_blocks_per_slot_for_feedback = None;
2434 if let Some(required) = required_admission_blocks.filter(|required| *required > 0) {
2435 let width = attempted_prefill_width.unwrap_or(1).max(1);
2436 let required_blocks_per_slot = required.div_ceil(width).max(1);
2437 required_blocks_per_slot_for_feedback = Some(required_blocks_per_slot);
2438 self.capacity_mixed_recompute_required_blocks_per_slot
2439 .store(required_blocks_per_slot, Ordering::Relaxed);
2440 }
2441 if let Some(observed) = observed_free_blocks {
2442 self.capacity_mixed_recompute_observed_free_blocks
2443 .store(observed, Ordering::Relaxed);
2444 if let Some(required_blocks_per_slot) = required_blocks_per_slot_for_feedback {
2445 let usable_free_blocks = Self::capacity_mixed_recompute_usable_free_blocks(
2446 observed,
2447 required_blocks_per_slot,
2448 );
2449 if usable_free_blocks >= required_blocks_per_slot {
2450 self.record_capacity_recompute_progress();
2451 }
2452 }
2453 }
2454 let _ = self
2455 .capacity_mixed_recompute_blocked_until_epoch
2456 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
2457 Some(current.max(blocked_until))
2458 });
2459 }
2460
2461 pub fn defer_prefill_to_waiting(&self, request_id: &RequestId) -> bool {
2468 let mut prefill_queue = self.prefill_queue.write();
2469 let mut waiting_queue = self.waiting_queue.write();
2470 let mut request_index = self.request_index.write();
2471 let attempted_prefill_width = prefill_queue.len();
2472
2473 if let Some(pos) = prefill_queue
2474 .iter()
2475 .position(|r| r.inner.request.id == *request_id)
2476 {
2477 let mut req = prefill_queue.remove(pos).unwrap();
2478 req.phase = RequestPhase::Waiting;
2479 req.inner.state = RequestState::Waiting;
2480 req.inner.started_at = None;
2481 req.prefill_tokens = 0;
2482 req.kv_blocks.clear();
2483 req.chunked_prefill = false;
2484 req.prefill_chunk_offset = 0;
2485 req.prefill_execution_chunk_ceiling = None;
2486 req.logical_work_frontier.yield_for_recompute();
2487 req.capacity_deferred_until_release_epoch = self
2488 .capacity_release_epoch
2489 .load(Ordering::Relaxed)
2490 .saturating_add(1);
2491 if !self.decode_queue.read().requests.is_empty() {
2492 req.capacity_deferred_mixed_attempt_epoch =
2493 Some(self.capacity_mixed_recompute_epoch.load(Ordering::Relaxed));
2494 }
2495 req.last_iteration = self.current_iteration.load(Ordering::Relaxed);
2496 if !self.requeue_waiting_request(&mut waiting_queue, &mut request_index, req) {
2497 return false;
2498 }
2499 self.record_capacity_defer_feedback(attempted_prefill_width);
2500 debug!("Deferred prefill request {} back to waiting", request_id);
2501 true
2502 } else {
2503 false
2504 }
2505 }
2506
2507 pub fn defer_decode_to_waiting_for_capacity(
2514 &self,
2515 request_id: &RequestId,
2516 attempted_decode_width: usize,
2517 ) -> bool {
2518 self.defer_decode_to_waiting_for_capacity_with_pressure(
2519 request_id,
2520 attempted_decode_width,
2521 None,
2522 )
2523 }
2524
2525 pub fn defer_decode_to_waiting_for_capacity_with_pressure(
2526 &self,
2527 request_id: &RequestId,
2528 attempted_decode_width: usize,
2529 observed_free_blocks: Option<usize>,
2530 ) -> bool {
2531 self.defer_decode_to_waiting_for_capacity_inner(
2532 request_id,
2533 attempted_decode_width,
2534 observed_free_blocks,
2535 )
2536 }
2537
2538 pub fn arm_execution_capacity_yield(
2540 &self,
2541 transaction: &PressureYieldTransaction,
2542 ) -> Result<PressureTransitionOrdinal> {
2543 let ordinal = self
2544 .pressure_coordinator
2545 .lock()
2546 .arm_release_fence(transaction)
2547 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
2548 self.pressure_active.store(true, Ordering::Release);
2549 Ok(ordinal)
2550 }
2551
2552 pub fn complete_execution_capacity_yield(
2555 &self,
2556 transaction: &PressureYieldTransaction,
2557 attempted_decode_width: usize,
2558 observed_free_blocks: Option<usize>,
2559 ) -> Result<ExecutionCapacityYieldCompletion> {
2560 let request_id = transaction.victim_request_id();
2561 let victim_waiting_ticket = self.requeue_execution_capacity_victim(
2562 request_id,
2563 attempted_decode_width,
2564 observed_free_blocks,
2565 );
2566 let requeued = victim_waiting_ticket.is_some();
2567 let progress_owner_wait_condition =
2568 self.execution_capacity_wait_condition(transaction.progress_owner_id());
2569
2570 let (release_ordinal, disposition, installed_hold) = {
2571 let mut coordinator = self.pressure_coordinator.lock();
2572 if !requeued && transaction.kind() == PressureYieldKind::SelfRecompute {
2573 let _ = coordinator
2574 .record_terminal(request_id)
2575 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
2576 }
2577 let completion = coordinator
2578 .complete_release_fence(transaction, progress_owner_wait_condition.as_ref())
2579 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
2580 if !requeued && transaction.kind() == PressureYieldKind::PeerHandoff {
2581 let _ = coordinator
2582 .record_terminal(request_id)
2583 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
2584 }
2585 self.pressure_active
2586 .store(coordinator.has_records(), Ordering::Release);
2587 let installed_hold = match (coordinator.hold_status(request_id), victim_waiting_ticket)
2588 {
2589 (
2590 PressureHoldStatus::Held {
2591 episode_id,
2592 progress_owner_id,
2593 progress_baseline,
2594 progress_current,
2595 },
2596 Some(waiting_ticket),
2597 ) => Some(ExecutionCapacityPressureHoldReceipt {
2598 episode_id,
2599 transition_ordinal: completion.0,
2600 request_id: request_id.clone(),
2601 progress_owner_id,
2602 progress_baseline,
2603 progress_current,
2604 waiting_ticket,
2605 }),
2606 _ => None,
2607 };
2608 (completion.0, completion.1, installed_hold)
2609 };
2610 let (
2611 resumable_transition_ordinal,
2612 owner_admission_pending_transition_ordinal,
2613 closed_transition_ordinal,
2614 disposition,
2615 ) = match disposition {
2616 PressureReleaseFenceDisposition::Resumable(ordinal) => (
2617 Some(ordinal),
2618 None,
2619 None,
2620 ExecutionCapacityYieldDisposition::ProgressOwnerResumable,
2621 ),
2622 PressureReleaseFenceDisposition::OwnerAdmissionPending(ordinal) => (
2623 None,
2624 Some(ordinal),
2625 None,
2626 ExecutionCapacityYieldDisposition::ProgressOwnerAdmissionPending,
2627 ),
2628 PressureReleaseFenceDisposition::SelfRecomputeQueued(ordinal) => (
2629 None,
2630 None,
2631 Some(ordinal),
2632 ExecutionCapacityYieldDisposition::SelfRecomputeQueued,
2633 ),
2634 PressureReleaseFenceDisposition::Closed { ordinal, reason } => {
2635 let disposition = match reason {
2636 PressureHoldReleaseReason::OwnerTerminal => {
2637 ExecutionCapacityYieldDisposition::OwnerTerminal
2638 }
2639 };
2640 (None, None, Some(ordinal), disposition)
2641 }
2642 };
2643 Ok(ExecutionCapacityYieldCompletion {
2644 victim_requeued: requeued,
2645 installed_hold,
2646 release_transition_ordinal: release_ordinal,
2647 resumable_transition_ordinal,
2648 owner_admission_pending_transition_ordinal,
2649 closed_transition_ordinal,
2650 disposition,
2651 })
2652 }
2653
2654 fn execution_capacity_wait_condition(
2655 &self,
2656 request_id: &RequestId,
2657 ) -> Option<CapacityWaitCondition> {
2658 if let Some(condition) = self
2659 .prefill_queue
2660 .read()
2661 .iter()
2662 .find(|request| request.inner.request.id == *request_id)
2663 .and_then(|request| request.execution_capacity_deferral.as_ref())
2664 .map(|deferral| deferral.wait_condition().clone())
2665 {
2666 return Some(condition);
2667 }
2668 if let Some(condition) = self
2669 .decode_queue
2670 .read()
2671 .requests
2672 .get(request_id)
2673 .and_then(|request| request.execution_capacity_deferral.as_ref())
2674 .map(|deferral| deferral.wait_condition().clone())
2675 {
2676 return Some(condition);
2677 }
2678 self.waiting_queue
2679 .read()
2680 .iter()
2681 .find(|request| request.inner.request.id == *request_id)
2682 .and_then(|request| request.execution_capacity_deferral.as_ref())
2683 .map(|deferral| deferral.wait_condition().clone())
2684 }
2685
2686 pub fn abort_execution_capacity_yield(
2689 &self,
2690 transaction: &PressureYieldTransaction,
2691 victim_released: bool,
2692 attempted_decode_width: usize,
2693 observed_free_blocks: Option<usize>,
2694 ) -> Result<(bool, PressureTransitionOrdinal, PressureTransitionOrdinal)> {
2695 let (aborted_ordinal, closed_ordinal, participants) = {
2696 let mut coordinator = self.pressure_coordinator.lock();
2697 let (aborted, closed, participants) = coordinator
2698 .abort_yield(transaction)
2699 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
2700 self.pressure_active
2701 .store(coordinator.has_records(), Ordering::Release);
2702 (aborted, closed, participants)
2703 };
2704 let participants = participants.into_iter().collect::<HashSet<_>>();
2705 for request in self.prefill_queue.write().iter_mut() {
2706 if participants.contains(&request.inner.request.id) {
2707 request.execution_capacity_deferral = None;
2708 }
2709 }
2710 for request in self.decode_queue.write().requests.values_mut() {
2711 if participants.contains(&request.inner.request.id) {
2712 request.execution_capacity_deferral = None;
2713 }
2714 }
2715 let requeued = victim_released
2716 && self
2717 .requeue_execution_capacity_victim(
2718 transaction.victim_request_id(),
2719 attempted_decode_width,
2720 observed_free_blocks,
2721 )
2722 .is_some();
2723 Ok((requeued, aborted_ordinal, closed_ordinal))
2724 }
2725
2726 fn requeue_execution_capacity_victim(
2727 &self,
2728 request_id: &RequestId,
2729 attempted_decode_width: usize,
2730 observed_free_blocks: Option<usize>,
2731 ) -> Option<u64> {
2732 let request = {
2733 let mut prefill = self.prefill_queue.write();
2734 prefill
2735 .iter()
2736 .position(|request| request.inner.request.id == *request_id)
2737 .and_then(|position| prefill.remove(position))
2738 }
2739 .or_else(|| {
2740 let mut decode = self.decode_queue.write();
2741 decode.remove(request_id)
2742 });
2743
2744 let mut requeued_ticket = None;
2745 if let Some(mut request) = request {
2746 let waiting_ticket = request.waiting_admission_ticket;
2747 request.phase = RequestPhase::Waiting;
2748 request.inner.state = RequestState::Waiting;
2749 request.inner.started_at = None;
2750 request.prefill_tokens = 0;
2751 request.decode_tokens = 0;
2752 request.kv_blocks.clear();
2753 request.chunked_prefill = false;
2754 request.prefill_chunk_offset = 0;
2755 request.prefill_execution_chunk_ceiling = None;
2756 request.capacity_deferred_until_release_epoch = self
2757 .capacity_release_epoch
2758 .load(Ordering::Relaxed)
2759 .saturating_add(1);
2760 request.capacity_deferred_mixed_attempt_epoch = None;
2761 request.capacity_deferred_empty_retry_epoch = None;
2762 request.capacity_deferred_from_decode = true;
2763 request.execution_capacity_deferral = None;
2764 request.logical_work_frontier.yield_for_recompute();
2765 request.last_iteration = self.current_iteration.load(Ordering::Relaxed);
2766
2767 let mut waiting = self.waiting_queue.write();
2768 let mut request_index = self.request_index.write();
2769 if self.requeue_waiting_request(&mut waiting, &mut request_index, request) {
2770 requeued_ticket = waiting_ticket.map(|ticket| ticket.get());
2771 self.record_capacity_defer_feedback(attempted_decode_width.max(1));
2772 self.record_decode_capacity_pressure(
2773 attempted_decode_width.max(1),
2774 observed_free_blocks,
2775 );
2776 }
2777 }
2778
2779 requeued_ticket
2780 }
2781
2782 fn defer_decode_to_waiting_for_capacity_inner(
2783 &self,
2784 request_id: &RequestId,
2785 attempted_decode_width: usize,
2786 observed_free_blocks: Option<usize>,
2787 ) -> bool {
2788 let mut decode_queue = self.decode_queue.write();
2789 let mut waiting_queue = self.waiting_queue.write();
2790 let mut request_index = self.request_index.write();
2791
2792 if let Some(mut req) = decode_queue.remove(request_id) {
2793 req.phase = RequestPhase::Waiting;
2794 req.inner.state = RequestState::Waiting;
2795 req.inner.started_at = None;
2796 req.prefill_tokens = 0;
2797 req.decode_tokens = 0;
2798 req.kv_blocks.clear();
2799 req.chunked_prefill = false;
2800 req.prefill_chunk_offset = 0;
2801 req.prefill_execution_chunk_ceiling = None;
2802 req.capacity_deferred_until_release_epoch = self
2803 .capacity_release_epoch
2804 .load(Ordering::Relaxed)
2805 .saturating_add(1);
2806 req.capacity_deferred_mixed_attempt_epoch = None;
2807 req.capacity_deferred_empty_retry_epoch = None;
2808 req.capacity_deferred_from_decode = true;
2809 req.execution_capacity_deferral = None;
2810 req.logical_work_frontier.yield_for_recompute();
2811 req.last_iteration = self.current_iteration.load(Ordering::Relaxed);
2812 if !self.requeue_waiting_request(&mut waiting_queue, &mut request_index, req) {
2813 return false;
2814 }
2815 self.record_capacity_defer_feedback(attempted_decode_width.max(1));
2816 self.record_decode_capacity_pressure(
2817 attempted_decode_width.max(1),
2818 observed_free_blocks,
2819 );
2820 debug!("Deferred decode request {} back to waiting", request_id);
2821 true
2822 } else {
2823 false
2824 }
2825 }
2826
2827 fn promote_to_decode(&self, request_id: &RequestId) -> bool {
2829 let mut prefill_queue = self.prefill_queue.write();
2830 let mut decode_queue = self.decode_queue.write();
2831 let mut request_index = self.request_index.write();
2832
2833 if let Some(pos) = prefill_queue
2834 .iter()
2835 .position(|r| r.inner.request.id == *request_id)
2836 {
2837 let mut req = prefill_queue.remove(pos).unwrap();
2838 req.phase = RequestPhase::Decoding;
2839 req.capacity_deferred_until_release_epoch = 0;
2840 req.capacity_deferred_mixed_attempt_epoch = None;
2841 req.capacity_deferred_empty_retry_epoch = None;
2842 req.capacity_deferred_from_decode = false;
2843 req.execution_capacity_deferral = None;
2844 req.logical_work_frontier.begin_decode();
2845
2846 request_index.insert(request_id.clone(), RequestPhase::Decoding);
2847 decode_queue.requests.insert(request_id.clone(), req);
2848
2849 debug!("Promoted request {} to decode queue", request_id);
2850 true
2851 } else {
2852 false
2853 }
2854 }
2855
2856 fn initial_prefill_token_estimate(&self, req: &ContinuousBatchRequest) -> usize {
2857 if !self.runtime_config.prompt_token_estimate {
2858 return self.cb_config.prefill_chunk_size;
2859 }
2860
2861 self.prompt_token_estimate(req)
2862 .unwrap_or(self.cb_config.prefill_chunk_size)
2863 }
2864
2865 fn prompt_token_estimate(&self, req: &ContinuousBatchRequest) -> Option<usize> {
2866 req.inner
2867 .request
2868 .metadata
2869 .get(PROMPT_TOKENS_METADATA_KEY)
2870 .and_then(|v| v.as_u64())
2871 .map(|v| v as usize)
2872 .filter(|&v| v > 0)
2873 }
2874
2875 fn default_active_decode_prefill_chunk(&self) -> usize {
2876 self.cb_config.prefill_chunk_size.div_ceil(8).max(1)
2877 }
2878
2879 fn decode_pressure_prefill_cap_threshold(&self, hint: &BatchHint) -> usize {
2880 hint.max_batch_size
2881 .min(self.cb_config.max_decode_batch)
2882 .min(self.config.max_running_requests)
2883 .max(1)
2884 .div_ceil(2)
2885 .max(1)
2886 }
2887
2888 fn active_decode_prefill_chunk_for_iteration(
2889 &self,
2890 hint: &BatchHint,
2891 scheduled_decode_count: usize,
2892 ) -> Option<usize> {
2893 if scheduled_decode_count == 0 {
2894 return None;
2895 }
2896 if let Some(chunk) = self.runtime_config.active_decode_prefill_chunk {
2897 return Some(chunk);
2898 }
2899 let capacity_deferred_decode_backpressure = self.decode_capacity_deferred_backlog_len() > 0;
2900 if scheduled_decode_count < self.decode_pressure_prefill_cap_threshold(hint)
2901 && !capacity_deferred_decode_backpressure
2902 {
2903 return None;
2904 }
2905 Some(self.default_active_decode_prefill_chunk())
2906 }
2907
2908 fn effective_active_decode_prefill_chunk(
2909 &self,
2910 active_decode_prefill_chunk: Option<usize>,
2911 prefill_step_chunk: Option<usize>,
2912 ) -> Option<usize> {
2913 let chunk = active_decode_prefill_chunk?;
2914 Some(
2915 prefill_step_chunk
2916 .map(|step_chunk| step_chunk.min(chunk))
2917 .unwrap_or(chunk)
2918 .max(1),
2919 )
2920 }
2921
2922 fn active_decode_prefill_target_chunks(
2923 &self,
2924 hint: &BatchHint,
2925 scheduled_decode_count: usize,
2926 prefill_backlog: usize,
2927 ) -> usize {
2928 if scheduled_decode_count == 0 || prefill_backlog == 0 {
2929 return 0;
2930 }
2931
2932 let free_batch_slots = hint.max_batch_size.saturating_sub(scheduled_decode_count);
2933 if free_batch_slots == 0 {
2934 return 0;
2935 }
2936
2937 let max_mixed_prefill_chunks = self.cb_config.max_prefill_batch.div_ceil(2).max(1);
2942 free_batch_slots
2943 .min(max_mixed_prefill_chunks)
2944 .min(prefill_backlog)
2945 }
2946
2947 fn maybe_active_decode_prefill_chunk(
2948 &self,
2949 req: &ContinuousBatchRequest,
2950 active_decode_prefill_chunk: Option<usize>,
2951 ) -> Option<usize> {
2952 let chunk = active_decode_prefill_chunk?;
2953 if !req.chunked_prefill && self.decoding_count() == 0 {
2954 return None;
2955 }
2956 Some(chunk)
2957 }
2958
2959 fn remaining_prefill_tokens(&self, req: &ContinuousBatchRequest) -> usize {
2960 if req.prefill_tokens == 0 {
2961 self.initial_prefill_token_estimate(req)
2962 } else {
2963 req.prefill_tokens.saturating_sub(req.prefill_chunk_offset)
2964 }
2965 }
2966
2967 fn chunked_prefill_budget_tokens(&self, req: &ContinuousBatchRequest, chunk: usize) -> usize {
2968 let remaining = if req.prefill_tokens == 0 {
2969 self.prompt_token_estimate(req)
2970 .unwrap_or(self.cb_config.prefill_chunk_size)
2971 } else {
2972 req.prefill_tokens.saturating_sub(req.prefill_chunk_offset)
2973 };
2974 chunk.min(remaining).max(1)
2975 }
2976
2977 fn apply_prefill_execution_chunk_ceiling(req: &ContinuousBatchRequest, tokens: usize) -> usize {
2978 req.prefill_execution_chunk_ceiling
2979 .map(|ceiling| tokens.min(ceiling))
2980 .unwrap_or(tokens)
2981 .max(1)
2982 }
2983
2984 fn prefill_budget_tokens(
2985 &self,
2986 req: &ContinuousBatchRequest,
2987 active_decode_prefill_chunk: Option<usize>,
2988 prefill_step_chunk: Option<usize>,
2989 step_tokens_remaining: usize,
2990 ) -> usize {
2991 if step_tokens_remaining == 0 {
2992 return 0;
2993 }
2994 if let Some(chunk) =
2995 self.maybe_active_decode_prefill_chunk(req, active_decode_prefill_chunk)
2996 {
2997 let chunk = self
2998 .effective_active_decode_prefill_chunk(Some(chunk), prefill_step_chunk)
2999 .unwrap_or(chunk.max(1));
3000 let tokens = self
3001 .chunked_prefill_budget_tokens(req, chunk)
3002 .min(step_tokens_remaining)
3003 .max(1);
3004 return Self::apply_prefill_execution_chunk_ceiling(req, tokens);
3005 }
3006
3007 let remaining = self.remaining_prefill_tokens(req);
3008 if let Some(chunk) = prefill_step_chunk {
3009 let tokens = self
3010 .chunked_prefill_budget_tokens(req, chunk)
3011 .min(step_tokens_remaining)
3012 .max(1);
3013 return Self::apply_prefill_execution_chunk_ceiling(req, tokens);
3014 }
3015 let tokens = if self.cb_config.enable_chunked_prefill {
3016 remaining.min(step_tokens_remaining).max(1)
3017 } else {
3018 remaining.max(1)
3019 };
3020 Self::apply_prefill_execution_chunk_ceiling(req, tokens)
3021 }
3022
3023 fn active_decode_prefill_budget_tokens(
3024 &self,
3025 hint: &BatchHint,
3026 scheduled_decode_count: usize,
3027 active_decode_prefill_chunk: Option<usize>,
3028 prefill_step_chunk: Option<usize>,
3029 ) -> Option<usize> {
3030 let chunk = self.effective_active_decode_prefill_chunk(
3031 active_decode_prefill_chunk,
3032 prefill_step_chunk,
3033 )?;
3034 if scheduled_decode_count == 0 {
3035 return None;
3036 }
3037
3038 let remaining_step_tokens = hint.max_tokens.saturating_sub(scheduled_decode_count);
3039 let free_batch_slots = hint.max_batch_size.saturating_sub(scheduled_decode_count);
3040 if remaining_step_tokens == 0 || free_batch_slots == 0 {
3041 return Some(0);
3042 }
3043
3044 let prefill_backlog = self.prefilling_count().saturating_add(self.waiting_count());
3045 if prefill_backlog == 0 {
3046 return Some(0);
3047 }
3048
3049 let target_chunks =
3050 self.active_decode_prefill_target_chunks(hint, scheduled_decode_count, prefill_backlog);
3051
3052 Some(
3053 chunk
3054 .saturating_mul(target_chunks)
3055 .min(remaining_step_tokens),
3056 )
3057 }
3058
3059 fn active_decode_prefill_budget_chunks(
3060 &self,
3061 hint: &BatchHint,
3062 scheduled_decode_count: usize,
3063 active_decode_prefill_chunk: Option<usize>,
3064 ) -> Option<usize> {
3065 active_decode_prefill_chunk?;
3066 if scheduled_decode_count == 0 {
3067 return None;
3068 }
3069
3070 let remaining_step_tokens = hint.max_tokens.saturating_sub(scheduled_decode_count);
3071 let free_batch_slots = hint.max_batch_size.saturating_sub(scheduled_decode_count);
3072 if remaining_step_tokens == 0 || free_batch_slots == 0 {
3073 return Some(0);
3074 }
3075
3076 let prefill_backlog = self.prefilling_count().saturating_add(self.waiting_count());
3077 if prefill_backlog == 0 {
3078 return Some(0);
3079 }
3080
3081 Some(
3082 self.active_decode_prefill_target_chunks(hint, scheduled_decode_count, prefill_backlog)
3083 .min(remaining_step_tokens),
3084 )
3085 }
3086
3087 fn capacity_deferred_mixed_recompute_slot_budget(
3088 &self,
3089 active_decode_prefill_chunk: Option<usize>,
3090 prefill_step_chunk: Option<usize>,
3091 active_decode_prefill_tokens_remaining: Option<usize>,
3092 required_blocks_per_slot: usize,
3093 observed_free_blocks: usize,
3094 ) -> usize {
3095 let Some(tokens_remaining) = active_decode_prefill_tokens_remaining else {
3096 return 0;
3097 };
3098 if tokens_remaining == 0 {
3099 return 0;
3100 }
3101 let Some(chunk) = self
3102 .effective_active_decode_prefill_chunk(active_decode_prefill_chunk, prefill_step_chunk)
3103 else {
3104 return 0;
3105 };
3106 let token_budget_slots = tokens_remaining.div_ceil(chunk).max(1);
3107 if required_blocks_per_slot == 0 || observed_free_blocks == usize::MAX {
3108 return token_budget_slots;
3109 }
3110 let usable_free_blocks = Self::capacity_mixed_recompute_usable_free_blocks(
3111 observed_free_blocks,
3112 required_blocks_per_slot,
3113 );
3114 token_budget_slots.min(usable_free_blocks / required_blocks_per_slot)
3115 }
3116
3117 fn should_budget_capacity_deferred_mixed_recompute(
3118 req: &ContinuousBatchRequest,
3119 active_decode_prefill_chunk: Option<usize>,
3120 ) -> bool {
3121 req.capacity_deferred_until_release_epoch > 0 && active_decode_prefill_chunk.is_some()
3122 }
3123
3124 fn add_prefill_requests_to_batch(
3125 &self,
3126 iteration: u64,
3127 hint: &BatchHint,
3128 batch_requests: &mut Vec<ScheduledRequest>,
3129 total_tokens: &mut usize,
3130 scheduled_request_ids: &mut HashSet<RequestId>,
3131 active_decode_prefill_tokens_remaining: &mut Option<usize>,
3132 active_decode_prefill_chunks_remaining: &mut Option<usize>,
3133 capacity_deferred_mixed_recompute_slots_remaining: &mut Option<usize>,
3134 active_decode_prefill_chunk: Option<usize>,
3135 prefill_step_chunk: Option<usize>,
3136 waiting_admission: &mut WaitingAdmissionMode<'_>,
3137 _capacity_release_epoch: u64,
3138 capacity_mixed_recompute_epoch: u64,
3139 ) -> Result<()> {
3140 if batch_requests.len() >= hint.max_batch_size || *total_tokens >= hint.max_tokens {
3141 return Ok(());
3142 }
3143
3144 let mut prefill_queue = self.prefill_queue.write();
3145 for req in prefill_queue.iter_mut() {
3146 if batch_requests.len() >= hint.max_batch_size {
3147 break;
3148 }
3149 if scheduled_request_ids.contains(&req.inner.request.id) {
3150 continue;
3151 }
3152 if Self::execution_readiness_is_blocked(req) {
3153 continue;
3154 }
3155 if Self::execution_maintenance_retry_is_blocked(req, iteration)? {
3156 continue;
3157 }
3158 if Self::execution_capacity_is_blocked(
3159 req,
3160 waiting_admission,
3161 ExecutionCapacityQueuePhase::Prefill,
3162 )? {
3163 continue;
3164 }
3165 let budgeted_capacity_deferred = Self::should_budget_capacity_deferred_mixed_recompute(
3166 req,
3167 active_decode_prefill_chunk,
3168 );
3169 if budgeted_capacity_deferred
3170 && req.capacity_deferred_mixed_attempt_epoch == Some(capacity_mixed_recompute_epoch)
3171 {
3172 continue;
3173 }
3174 if budgeted_capacity_deferred
3175 && capacity_deferred_mixed_recompute_slots_remaining
3176 .as_ref()
3177 .copied()
3178 .unwrap_or(0)
3179 == 0
3180 {
3181 continue;
3182 }
3183 if active_decode_prefill_chunks_remaining
3184 .as_ref()
3185 .is_some_and(|remaining| *remaining == 0)
3186 {
3187 break;
3188 }
3189
3190 let mut step_tokens_remaining = hint.max_tokens.saturating_sub(*total_tokens);
3191 if let Some(remaining) = active_decode_prefill_tokens_remaining.as_ref() {
3192 step_tokens_remaining = step_tokens_remaining.min(*remaining);
3193 }
3194 let prefill_chunk_tokens = self.prefill_budget_tokens(
3195 req,
3196 active_decode_prefill_chunk,
3197 prefill_step_chunk,
3198 step_tokens_remaining,
3199 );
3200 if prefill_chunk_tokens == 0 {
3204 continue;
3205 }
3206 if let Some(remaining) = active_decode_prefill_tokens_remaining.as_mut() {
3207 if *remaining == 0 {
3208 break;
3209 }
3210 }
3211
3212 if *total_tokens + prefill_chunk_tokens <= hint.max_tokens {
3213 let mut scheduled = req.inner.clone();
3214 scheduled.tokens_processed = req.prefill_chunk_offset;
3215 scheduled.tokens_to_process = Some(prefill_chunk_tokens);
3216 req.logical_work_frontier
3217 .mark_scheduled(prefill_chunk_tokens);
3218 scheduled_request_ids.insert(scheduled.request.id.clone());
3219 batch_requests.push(scheduled);
3220 *total_tokens += prefill_chunk_tokens;
3221 if let Some(remaining) = active_decode_prefill_tokens_remaining.as_mut() {
3222 *remaining = remaining.saturating_sub(prefill_chunk_tokens);
3223 }
3224 if let Some(remaining) = active_decode_prefill_chunks_remaining.as_mut() {
3225 *remaining = remaining.saturating_sub(1);
3226 }
3227 if budgeted_capacity_deferred {
3228 req.capacity_deferred_mixed_attempt_epoch =
3229 Some(capacity_mixed_recompute_epoch);
3230 if let Some(remaining) =
3231 capacity_deferred_mixed_recompute_slots_remaining.as_mut()
3232 {
3233 *remaining = remaining.saturating_sub(1);
3234 }
3235 }
3236 }
3237 }
3238 Ok(())
3239 }
3240
3241 fn create_iteration_batch(&self, hint: BatchHint) -> Option<BatchPlan> {
3243 match self.create_iteration_batch_with_admission(hint, WaitingAdmissionMode::Legacy) {
3244 Ok(batch) => batch,
3245 Err(error) => {
3246 warn!("Legacy waiting admission failed: {}", error);
3247 None
3248 }
3249 }
3250 }
3251
3252 fn execution_capacity_is_blocked(
3253 req: &mut ContinuousBatchRequest,
3254 waiting_admission: &mut WaitingAdmissionMode<'_>,
3255 phase: ExecutionCapacityQueuePhase,
3256 ) -> Result<bool> {
3257 let Some(deferral) = req.execution_capacity_deferral.clone() else {
3258 return Ok(false);
3259 };
3260 let wake = waiting_admission.wake().ok_or_else(|| {
3261 FerrumError::scheduler(
3262 "typed execution capacity deferral reached a legacy scheduler tick",
3263 )
3264 })?;
3265 if deferral.observed().coordinator_id() != wake.epochs().coordinator_id()
3266 || deferral.wait_condition().coordinator_id().get()
3267 != wake.epochs().coordinator_id().get()
3268 {
3269 return Err(FerrumError::scheduler(
3270 "typed execution capacity deferral belongs to another coordinator",
3271 ));
3272 }
3273 let observed = deferral.observed();
3274 let current = wake.epochs();
3275 if current.release_epoch() < observed.release_epoch()
3276 || current.capacity_epoch() < observed.capacity_epoch()
3277 || current.policy_epoch() < observed.policy_epoch()
3278 {
3279 return Err(FerrumError::scheduler(
3280 "typed execution capacity audit epoch regressed",
3281 ));
3282 }
3283 let exact_source_changed = deferral
3284 .wait_condition()
3285 .changed_since(wake.availability())
3286 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
3287 let policy_epoch_changed = current.policy_epoch() != observed.policy_epoch();
3288 let current_wait_sources = waiting_admission.observes().then(|| {
3289 deferral
3290 .wait_condition()
3291 .observed()
3292 .iter()
3293 .map(|observed| {
3294 let index = wake
3295 .availability()
3296 .binary_search_by_key(&observed.source(), |entry| entry.source())
3297 .expect("validated wait source remains available");
3298 wake.availability()[index]
3299 })
3300 .collect::<Vec<_>>()
3301 });
3302 if !exact_source_changed && !policy_epoch_changed {
3303 if let Some(current_wait_sources) = current_wait_sources {
3304 let observation = match phase {
3305 ExecutionCapacityQueuePhase::Prefill => {
3306 ExecutorAdmissionQueueObservation::PrefillSkippedUnchanged {
3307 request_id: req.inner.request.id.clone(),
3308 deferral,
3309 current,
3310 current_wait_sources,
3311 }
3312 }
3313 ExecutionCapacityQueuePhase::Decode => {
3314 ExecutorAdmissionQueueObservation::DecodeSkippedUnchanged {
3315 request_id: req.inner.request.id.clone(),
3316 deferral,
3317 current,
3318 current_wait_sources,
3319 }
3320 }
3321 };
3322 waiting_admission.observe(observation);
3323 }
3324 return Ok(true);
3325 }
3326 if let Some(current_wait_sources) = current_wait_sources {
3327 let observation = match phase {
3328 ExecutionCapacityQueuePhase::Prefill => {
3329 ExecutorAdmissionQueueObservation::PrefillResumed {
3330 request_id: req.inner.request.id.clone(),
3331 deferral,
3332 current,
3333 current_wait_sources,
3334 exact_source_changed,
3335 policy_epoch_changed,
3336 }
3337 }
3338 ExecutionCapacityQueuePhase::Decode => {
3339 ExecutorAdmissionQueueObservation::DecodeResumed {
3340 request_id: req.inner.request.id.clone(),
3341 deferral,
3342 current,
3343 current_wait_sources,
3344 exact_source_changed,
3345 policy_epoch_changed,
3346 }
3347 }
3348 };
3349 waiting_admission.observe(observation);
3350 }
3351 req.execution_capacity_deferral = None;
3352 Ok(false)
3353 }
3354
3355 fn execution_readiness_is_blocked(req: &mut ContinuousBatchRequest) -> bool {
3356 let Some(block) = req.execution_readiness_block.as_ref() else {
3357 return false;
3358 };
3359 match block.status() {
3360 EXECUTION_READINESS_PENDING | EXECUTION_READINESS_FAILED => true,
3361 EXECUTION_READINESS_READY | EXECUTION_READINESS_CANCELLED => {
3362 req.execution_readiness_block = None;
3363 false
3364 }
3365 _ => true,
3366 }
3367 }
3368
3369 fn execution_maintenance_retry_is_blocked(
3370 req: &mut ContinuousBatchRequest,
3371 iteration: u64,
3372 ) -> Result<bool> {
3373 let Some(ticket) = req.execution_maintenance_retry else {
3374 return Ok(false);
3375 };
3376 if req.last_execution_maintenance_capacity_epoch != Some(ticket.latest_capacity_epoch) {
3377 return Err(FerrumError::scheduler(
3378 "execution maintenance retry ticket lost its mutation generation",
3379 ));
3380 }
3381 if iteration < ticket.not_before_iteration {
3382 return Ok(true);
3383 }
3384 req.execution_maintenance_retry = None;
3385 Ok(false)
3386 }
3387
3388 fn add_decode_requests_to_batch(
3389 &self,
3390 iteration: u64,
3391 hint: &BatchHint,
3392 batch_requests: &mut Vec<ScheduledRequest>,
3393 total_tokens: &mut usize,
3394 scheduled_request_ids: &mut HashSet<RequestId>,
3395 waiting_admission: &mut WaitingAdmissionMode<'_>,
3396 ) -> Result<()> {
3397 let has_deferred_recompute_backlog = self.decode_capacity_deferred_backlog_len() > 0;
3398 let (enforce_execution_backpressure, decode_capacity_backpressure_limit) = {
3399 let _feedback = self.decode_capacity_feedback_lock.lock();
3400 (
3401 self.decode_execution_pressure_enforced
3402 .load(Ordering::Acquire),
3403 self.decode_capacity_backpressure_limit(),
3404 )
3405 };
3406 let decode_batch_limit =
3407 if has_deferred_recompute_backlog && !enforce_execution_backpressure {
3408 hint.max_batch_size
3409 } else {
3410 decode_capacity_backpressure_limit
3411 .map(|limit| hint.max_batch_size.min(limit.max(1)))
3412 .unwrap_or(hint.max_batch_size)
3413 };
3414 let mut decode_queue = self.decode_queue.write();
3415 let decode_len = decode_queue.requests.len();
3416 if decode_len == 0 {
3417 decode_queue.selection_cursor = None;
3418 } else if decode_queue
3419 .selection_cursor
3420 .as_ref()
3421 .is_none_or(|cursor_id| !decode_queue.requests.contains_key(cursor_id))
3422 {
3423 decode_queue.selection_cursor =
3424 decode_queue.requests.get_index(0).map(|(id, _)| id.clone());
3425 }
3426 let start = decode_queue
3427 .selection_cursor
3428 .as_ref()
3429 .and_then(|cursor_id| decode_queue.requests.get_index_of(cursor_id))
3430 .unwrap_or(0);
3431 let mut next_cursor_index = start;
3432 let mut scheduled_count = 0usize;
3433 for offset in 0..decode_len {
3434 if batch_requests.len() >= decode_batch_limit || *total_tokens >= hint.max_tokens {
3435 break;
3436 }
3437 let index = (start + offset) % decode_len;
3438 let (_, req) = decode_queue
3439 .requests
3440 .get_index_mut(index)
3441 .expect("decode round-robin index remains in bounds");
3442 if scheduled_request_ids.contains(&req.inner.request.id) {
3443 continue;
3444 }
3445 if Self::execution_readiness_is_blocked(req) {
3446 continue;
3447 }
3448 if Self::execution_maintenance_retry_is_blocked(req, iteration)? {
3449 continue;
3450 }
3451 if Self::execution_capacity_is_blocked(
3452 req,
3453 waiting_admission,
3454 ExecutionCapacityQueuePhase::Decode,
3455 )? {
3456 continue;
3457 }
3458
3459 let mut scheduled = req.inner.clone();
3460 scheduled.tokens_processed = req.total_tokens();
3461 scheduled.tokens_to_process = Some(1);
3462 req.logical_work_frontier.mark_scheduled(1);
3463 scheduled_request_ids.insert(scheduled.request.id.clone());
3464 batch_requests.push(scheduled);
3465 *total_tokens += 1;
3466 scheduled_count += 1;
3467 next_cursor_index = (index + 1) % decode_len;
3468 }
3469 if scheduled_count > 0 {
3470 decode_queue.selection_cursor = decode_queue
3471 .requests
3472 .get_index(next_cursor_index)
3473 .map(|(id, _)| id.clone());
3474 }
3475 Ok(())
3476 }
3477
3478 fn create_iteration_batch_with_admission(
3479 &self,
3480 hint: BatchHint,
3481 mut waiting_admission: WaitingAdmissionMode<'_>,
3482 ) -> Result<Option<BatchPlan>> {
3483 let iteration = self.current_iteration.fetch_add(1, Ordering::Relaxed);
3484 self.metrics_tracker.record_iteration();
3485
3486 let mut batch_requests = Vec::new();
3487 let mut scheduled_request_ids = HashSet::new();
3488 let mut total_tokens = 0;
3489 let prefill_first_target = self
3490 .runtime_config
3491 .prefill_first_until_active
3492 .map(|target| {
3493 target
3494 .min(hint.max_batch_size)
3495 .min(self.cb_config.max_decode_batch)
3496 })
3497 .unwrap_or(0);
3498 let decoding_count = self.decoding_count();
3499 let active_count = self.active_count();
3500 let capacity_backpressure_active = self.capacity_backpressure_admit_limit().is_some();
3501 let skip_decode_for_prefill_first = prefill_first_target > 0
3502 && decoding_count < prefill_first_target
3503 && active_count < prefill_first_target
3504 && !(capacity_backpressure_active && decoding_count > 0)
3505 && (self.prefilling_count() > 0 || self.waiting_count() > 0);
3506 if !skip_decode_for_prefill_first {
3511 self.add_decode_requests_to_batch(
3512 iteration,
3513 &hint,
3514 &mut batch_requests,
3515 &mut total_tokens,
3516 &mut scheduled_request_ids,
3517 &mut waiting_admission,
3518 )?;
3519 }
3520 let scheduled_decode_count = batch_requests.len();
3521 let active_decode_prefill_chunk =
3522 self.active_decode_prefill_chunk_for_iteration(&hint, scheduled_decode_count);
3523 let prefill_step_chunk = self.runtime_config.prefill_step_chunk;
3524 let mut active_decode_prefill_tokens_remaining = self.active_decode_prefill_budget_tokens(
3525 &hint,
3526 scheduled_decode_count,
3527 active_decode_prefill_chunk,
3528 prefill_step_chunk,
3529 );
3530 let mut active_decode_prefill_chunks_remaining = self.active_decode_prefill_budget_chunks(
3531 &hint,
3532 scheduled_decode_count,
3533 active_decode_prefill_chunk,
3534 );
3535 let capacity_release_epoch = self.capacity_release_epoch.load(Ordering::Relaxed);
3536 let capacity_mixed_recompute_epoch =
3537 self.capacity_mixed_recompute_epoch.load(Ordering::Relaxed);
3538 let mixed_recompute_release_ready = self
3539 .capacity_mixed_recompute_blocked_until_epoch
3540 .load(Ordering::Relaxed)
3541 <= capacity_mixed_recompute_epoch;
3542 let mixed_recompute_required_blocks_per_slot = self
3543 .capacity_mixed_recompute_required_blocks_per_slot
3544 .load(Ordering::Relaxed);
3545 let mixed_recompute_observed_free_blocks = self
3546 .capacity_mixed_recompute_observed_free_blocks
3547 .load(Ordering::Relaxed);
3548 let mixed_recompute_kv_capacity_ready = mixed_recompute_required_blocks_per_slot == 0
3549 || Self::capacity_mixed_recompute_usable_free_blocks(
3550 mixed_recompute_observed_free_blocks,
3551 mixed_recompute_required_blocks_per_slot,
3552 ) >= mixed_recompute_required_blocks_per_slot;
3553 let allow_capacity_deferred_mixed_recompute = scheduled_decode_count > 0
3554 && active_decode_prefill_chunk.is_some()
3555 && active_decode_prefill_tokens_remaining.unwrap_or(0) > 0
3556 && mixed_recompute_release_ready
3557 && mixed_recompute_kv_capacity_ready;
3558 let mut capacity_deferred_mixed_recompute_slots_remaining =
3559 allow_capacity_deferred_mixed_recompute.then(|| {
3560 self.capacity_deferred_mixed_recompute_slot_budget(
3561 active_decode_prefill_chunk,
3562 prefill_step_chunk,
3563 active_decode_prefill_tokens_remaining,
3564 mixed_recompute_required_blocks_per_slot,
3565 mixed_recompute_observed_free_blocks,
3566 )
3567 });
3568
3569 self.add_prefill_requests_to_batch(
3578 iteration,
3579 &hint,
3580 &mut batch_requests,
3581 &mut total_tokens,
3582 &mut scheduled_request_ids,
3583 &mut active_decode_prefill_tokens_remaining,
3584 &mut active_decode_prefill_chunks_remaining,
3585 &mut capacity_deferred_mixed_recompute_slots_remaining,
3586 active_decode_prefill_chunk,
3587 prefill_step_chunk,
3588 &mut waiting_admission,
3589 capacity_release_epoch,
3590 capacity_mixed_recompute_epoch,
3591 )?;
3592
3593 let active_capacity = self
3595 .config
3596 .max_running_requests
3597 .saturating_sub(self.active_count());
3598 let decode_capacity = self
3599 .cb_config
3600 .max_decode_batch
3601 .saturating_sub(self.decoding_count());
3602 let available_slots = active_capacity.min(decode_capacity);
3603 let available_slots = self
3604 .capacity_backpressure_admit_limit()
3605 .map(|limit| available_slots.min(limit))
3606 .unwrap_or(available_slots);
3607 if matches!(&waiting_admission, WaitingAdmissionMode::Legacy) {
3608 self.legacy_waiting_admission_ticks
3609 .fetch_add(1, Ordering::Relaxed);
3610 let active_count_for_capacity_wait = self.active_count();
3611 let waiting_queue = self.waiting_queue.read();
3612 let mut requests_to_admit = Vec::new();
3613 let mut release_blocked_capacity_deferred_admissions = 0usize;
3614 for req in waiting_queue.iter() {
3615 if requests_to_admit.len() >= available_slots {
3616 break;
3617 }
3618 if self.pressure_active.load(Ordering::Acquire)
3619 && matches!(
3620 self.pressure_coordinator
3621 .lock()
3622 .hold_status(&req.inner.request.id),
3623 PressureHoldStatus::Held { .. }
3624 )
3625 {
3626 continue;
3627 }
3628 let budgeted_capacity_deferred =
3629 Self::should_budget_capacity_deferred_mixed_recompute(
3630 req,
3631 active_decode_prefill_chunk,
3632 );
3633 let release_ready =
3634 req.capacity_deferred_until_release_epoch <= capacity_release_epoch;
3635 let empty_scheduler_retry = active_count_for_capacity_wait == 0
3636 && !release_ready
3637 && req.capacity_deferred_empty_retry_epoch != Some(capacity_release_epoch);
3638 if release_ready && !budgeted_capacity_deferred {
3639 requests_to_admit.push((req.inner.request.id.clone(), None));
3640 } else if empty_scheduler_retry && !budgeted_capacity_deferred {
3641 requests_to_admit
3642 .push((req.inner.request.id.clone(), Some(capacity_release_epoch)));
3643 } else if req.capacity_deferred_mixed_attempt_epoch
3644 == Some(capacity_mixed_recompute_epoch)
3645 {
3646 continue;
3647 } else if allow_capacity_deferred_mixed_recompute
3648 && release_blocked_capacity_deferred_admissions
3649 < capacity_deferred_mixed_recompute_slots_remaining.unwrap_or(0)
3650 {
3651 requests_to_admit.push((req.inner.request.id.clone(), None));
3652 release_blocked_capacity_deferred_admissions += 1;
3653 }
3654 }
3655 drop(waiting_queue);
3656 for (req_id, empty_retry_epoch) in requests_to_admit {
3657 self.promote_to_prefill_with_empty_retry(&req_id, empty_retry_epoch);
3658 }
3659 } else {
3660 self.admit_waiting_dynamically(usize::MAX, available_slots, &mut waiting_admission)?;
3661 }
3662
3663 self.add_prefill_requests_to_batch(
3668 iteration,
3669 &hint,
3670 &mut batch_requests,
3671 &mut total_tokens,
3672 &mut scheduled_request_ids,
3673 &mut active_decode_prefill_tokens_remaining,
3674 &mut active_decode_prefill_chunks_remaining,
3675 &mut capacity_deferred_mixed_recompute_slots_remaining,
3676 active_decode_prefill_chunk,
3677 prefill_step_chunk,
3678 &mut waiting_admission,
3679 capacity_release_epoch,
3680 capacity_mixed_recompute_epoch,
3681 )?;
3682
3683 if skip_decode_for_prefill_first && batch_requests.is_empty() {
3688 self.add_decode_requests_to_batch(
3689 iteration,
3690 &hint,
3691 &mut batch_requests,
3692 &mut total_tokens,
3693 &mut scheduled_request_ids,
3694 &mut waiting_admission,
3695 )?;
3696 }
3697
3698 if self.runtime_config.scheduler_none_prof && !batch_requests.is_empty() {
3700 use std::sync::atomic::AtomicU64;
3701 static SOME_PROF_N: AtomicU64 = AtomicU64::new(0);
3702 let n = SOME_PROF_N.fetch_add(1, Ordering::Relaxed);
3703 if n.is_multiple_of(64) {
3704 let d_len = self.decode_queue.read().requests.len();
3705 let p_len = self.prefill_queue.read().len();
3706 let w_len = self.waiting_queue.read().len();
3707 eprintln!(
3708 "[sched-some] n={} returning_batch={} | decode_queue={} prefill_queue={} waiting_queue={}",
3709 n,
3710 batch_requests.len(),
3711 d_len,
3712 p_len,
3713 w_len,
3714 );
3715 }
3716 }
3717 if batch_requests.is_empty() {
3718 if self.runtime_config.scheduler_none_prof {
3720 use std::sync::atomic::AtomicU64;
3721 static NONE_PROF_N: AtomicU64 = AtomicU64::new(0);
3722 let n = NONE_PROF_N.fetch_add(1, Ordering::Relaxed);
3723 if n.is_multiple_of(512) {
3724 let d_len = self.decode_queue.read().requests.len();
3725 let p_len = self.prefill_queue.read().len();
3726 let w_len = self.waiting_queue.read().len();
3727 let d_count = self.decoding_count();
3728 eprintln!(
3729 "[sched-none] n={} decode_queue={} prefill_queue={} waiting_queue={} decoding_count={} hint.max_batch={}",
3730 n,
3731 d_len,
3732 p_len,
3733 w_len,
3734 d_count,
3735 hint.max_batch_size,
3736 );
3737 }
3738 }
3739 return Ok(None);
3740 }
3741
3742 let batch_id = BatchId::new();
3743 let max_seq_len = batch_requests
3744 .iter()
3745 .map(|r| r.request.sampling_params.max_tokens)
3746 .max()
3747 .unwrap_or(2048);
3748
3749 debug!(
3750 "Created iteration {} batch: {} requests, {} tokens",
3751 iteration,
3752 batch_requests.len(),
3753 total_tokens
3754 );
3755
3756 Ok(Some(BatchPlan {
3757 batch_id,
3758 requests: batch_requests,
3759 max_sequence_length: max_seq_len,
3760 estimated_time_ms: Some(self.cb_config.target_iteration_time_ms),
3761 resource_requirements: BatchResourceRequirements {
3762 gpu_memory: (total_tokens * 16) as u64,
3763 cpu_memory: (total_tokens * 4) as u64,
3764 kv_cache_blocks: total_tokens / 16,
3765 recurrent_state_bytes: 0,
3766 recurrent_state_slots: 0,
3767 compute_units: 1,
3768 },
3769 created_at: chrono::Utc::now(),
3770 }))
3771 }
3772
3773 pub fn mark_prefill_complete(&self, request_id: &RequestId, tokens: usize) {
3775 let mut prefill_queue = self.prefill_queue.write();
3776 let mut found = false;
3777 let mut progress = None;
3778 if let Some(pos) = prefill_queue
3779 .iter()
3780 .position(|r| r.inner.request.id == *request_id)
3781 {
3782 let req = &mut prefill_queue[pos];
3783 let delta = tokens.saturating_sub(req.prefill_chunk_offset);
3784 req.prefill_tokens = tokens;
3785 req.prefill_chunk_offset = tokens;
3786 req.chunked_prefill = false;
3787 req.logical_work_frontier.commit_prefill(tokens, delta);
3788 progress = Some(req.logical_work_frontier.progress_generation());
3789 found = true;
3790 }
3791 drop(prefill_queue);
3792
3793 self.promote_to_decode(request_id);
3795 if found {
3796 if let Some(progress) = progress {
3797 self.record_pressure_frontier_progress(request_id, progress);
3798 }
3799 self.record_resource_progress();
3800 }
3801 }
3802
3803 pub fn mark_prefill_chunk_processed(
3813 &self,
3814 request_id: &RequestId,
3815 total_prompt_tokens: usize,
3816 chunk_tokens: usize,
3817 ) -> bool {
3818 self.mark_prefill_chunk_processed_inner(request_id, total_prompt_tokens, chunk_tokens, None)
3819 }
3820
3821 pub fn mark_prefill_chunk_processed_with_capacity_feedback(
3824 &self,
3825 request_id: &RequestId,
3826 total_prompt_tokens: usize,
3827 planned_chunk_tokens: usize,
3828 completed_chunk_tokens: usize,
3829 ) -> Result<bool> {
3830 if completed_chunk_tokens == 0 || completed_chunk_tokens > planned_chunk_tokens {
3831 return Err(FerrumError::scheduler(
3832 "completed prefill prefix must be non-empty and no wider than its planned frontier",
3833 ));
3834 }
3835 Ok(self.mark_prefill_chunk_processed_inner(
3836 request_id,
3837 total_prompt_tokens,
3838 completed_chunk_tokens,
3839 Some((planned_chunk_tokens, completed_chunk_tokens)),
3840 ))
3841 }
3842
3843 fn mark_prefill_chunk_processed_inner(
3844 &self,
3845 request_id: &RequestId,
3846 total_prompt_tokens: usize,
3847 chunk_tokens: usize,
3848 execution_frontier_feedback: Option<(usize, usize)>,
3849 ) -> bool {
3850 let mut prefill_queue = self.prefill_queue.write();
3851 let mut fully_done = false;
3852 let mut made_progress = false;
3853 let mut progress = None;
3854 if let Some(pos) = prefill_queue
3855 .iter()
3856 .position(|r| r.inner.request.id == *request_id)
3857 {
3858 let req = &mut prefill_queue[pos];
3859 req.prefill_tokens = total_prompt_tokens;
3860 req.chunked_prefill = true;
3861 let previous_offset = req.prefill_chunk_offset;
3862 req.prefill_chunk_offset = req
3863 .prefill_chunk_offset
3864 .saturating_add(chunk_tokens)
3865 .min(total_prompt_tokens);
3866 let committed_tokens = req.prefill_chunk_offset.saturating_sub(previous_offset);
3867 req.logical_work_frontier
3868 .commit_prefill(req.prefill_chunk_offset, committed_tokens);
3869 if let Some((planned, completed)) = execution_frontier_feedback {
3870 if completed < planned {
3871 req.prefill_execution_chunk_ceiling = Some(
3872 req.prefill_execution_chunk_ceiling
3873 .map(|current| current.min(completed))
3874 .unwrap_or(completed),
3875 );
3876 } else if let Some(current) = req.prefill_execution_chunk_ceiling {
3877 req.prefill_execution_chunk_ceiling = Some(current.saturating_mul(2));
3878 }
3879 }
3880 progress = Some(req.logical_work_frontier.progress_generation());
3881 if committed_tokens > 0 {
3882 req.capacity_deferred_mixed_attempt_epoch = None;
3883 req.capacity_deferred_empty_retry_epoch = None;
3884 }
3885 fully_done = req.prefill_chunk_offset >= total_prompt_tokens;
3886 made_progress = committed_tokens > 0;
3887 }
3888 drop(prefill_queue);
3889
3890 if fully_done {
3891 self.promote_to_decode(request_id);
3892 }
3893 if made_progress {
3894 if let Some(progress) = progress {
3895 self.record_pressure_frontier_progress(request_id, progress);
3896 }
3897 }
3898 if made_progress && fully_done {
3899 self.record_resource_progress();
3900 }
3901 fully_done
3902 }
3903
3904 pub fn update_decode_progress(&self, request_id: &RequestId, tokens_generated: usize) {
3906 let mut decode_queue = self.decode_queue.write();
3907 let mut progress = None;
3908 if let Some(req) = decode_queue.requests.get_mut(request_id) {
3909 req.decode_tokens = tokens_generated;
3910 req.logical_work_frontier.commit_decode(tokens_generated);
3911 progress = Some(req.logical_work_frontier.progress_generation());
3912 req.last_iteration = self.current_iteration.load(Ordering::Relaxed);
3913 }
3914 drop(decode_queue);
3915 if let Some(progress) = progress {
3916 self.record_pressure_frontier_progress(request_id, progress);
3917 }
3918 }
3921
3922 fn record_pressure_frontier_progress(
3923 &self,
3924 request_id: &RequestId,
3925 progress: LogicalWorkGeneration,
3926 ) {
3927 if !self.pressure_active.load(Ordering::Acquire) {
3928 return;
3929 }
3930 let mut coordinator = self.pressure_coordinator.lock();
3931 if let Err(error) = coordinator.record_progress(request_id, progress) {
3932 warn!(
3933 request_id = %request_id,
3934 error = %error,
3935 "Pressure coordinator rejected logical frontier progress"
3936 );
3937 }
3938 self.pressure_active
3939 .store(coordinator.has_records(), Ordering::Release);
3940 }
3941
3942 fn record_pressure_frontier_terminal(&self, request_id: &RequestId) {
3943 if !self.pressure_active.load(Ordering::Acquire) {
3944 return;
3945 }
3946 let mut coordinator = self.pressure_coordinator.lock();
3947 if let Err(error) = coordinator.record_terminal(request_id) {
3948 warn!(
3949 request_id = %request_id,
3950 error = %error,
3951 "Pressure coordinator rejected logical frontier terminal state"
3952 );
3953 }
3954 self.pressure_active
3955 .store(coordinator.has_records(), Ordering::Release);
3956 }
3957
3958 fn consume_pressure_hold(&self, request_id: &RequestId) {
3959 if !self.pressure_active.load(Ordering::Acquire) {
3960 return;
3961 }
3962 let mut coordinator = self.pressure_coordinator.lock();
3963 if let Err(error) = coordinator.consume_released_hold(request_id) {
3964 warn!(
3965 request_id = %request_id,
3966 error = %error,
3967 "Pressure coordinator rejected admitted owner transition"
3968 );
3969 }
3970 self.pressure_active
3971 .store(coordinator.has_records(), Ordering::Release);
3972 }
3973}
3974
3975#[async_trait]
3976impl Scheduler for ContinuousBatchScheduler {
3977 async fn submit(&self, request: InferenceRequest) -> Result<RequestId> {
3978 let request_id = request.id.clone();
3979 debug!(
3980 "Submitting request {} to continuous batch scheduler",
3981 request_id
3982 );
3983
3984 let waiting_count = self.waiting_count();
3986 if waiting_count >= self.config.max_waiting_requests {
3987 warn!("Queue is full, rejecting request {}", request_id);
3988 return Err(FerrumError::scheduler(
3989 "Queue is full, cannot accept more requests",
3990 ));
3991 }
3992
3993 let cb_request = ContinuousBatchRequest::new(request);
3995
3996 let mut waiting_queue = self.waiting_queue.write();
3998 let queue_position = waiting_queue.len();
3999
4000 let mut req = cb_request;
4001 req.inner.queue_position = Some(queue_position);
4002
4003 let ticket = waiting_queue
4004 .enqueue(req)
4005 .map_err(|error| FerrumError::scheduler(error.to_string()))?;
4006 waiting_queue
4007 .request_mut(ticket)
4008 .expect("newly enqueued admission ticket remains present")
4009 .waiting_admission_ticket = Some(ticket);
4010
4011 self.request_index
4013 .write()
4014 .insert(request_id.clone(), RequestPhase::Waiting);
4015
4016 info!(
4017 "Request {} queued at position {}",
4018 request_id, queue_position
4019 );
4020 Ok(request_id)
4021 }
4022
4023 async fn next_batch(&self, hint: BatchHint) -> Option<BatchPlan> {
4024 self.create_iteration_batch(hint)
4025 }
4026
4027 async fn complete(&self, request_id: RequestId, response: &InferenceResponse) -> Result<()> {
4028 debug!("Completing request {}", request_id);
4029
4030 let mut decode_queue = self.decode_queue.write();
4032 if let Some(req) = decode_queue.remove(&request_id) {
4033 self.metrics_tracker.record_completion(&req);
4035
4036 match response.finish_reason {
4037 ferrum_types::FinishReason::EOS
4038 | ferrum_types::FinishReason::Stop
4039 | ferrum_types::FinishReason::Length => {
4040 self.completed_counter.fetch_add(1, Ordering::Relaxed);
4041 debug!("Request {} completed successfully", request_id);
4042 }
4043 _ => {
4044 self.failed_counter.fetch_add(1, Ordering::Relaxed);
4045 warn!(
4046 "Request {} completed with error: {:?}",
4047 request_id, response.finish_reason
4048 );
4049 }
4050 }
4051
4052 self.request_index.write().remove(&request_id);
4054 self.record_pressure_frontier_terminal(&request_id);
4055 self.record_capacity_release_progress();
4056
4057 Ok(())
4058 } else {
4059 let mut prefill_queue = self.prefill_queue.write();
4061 if let Some(pos) = prefill_queue
4062 .iter()
4063 .position(|r| r.inner.request.id == request_id)
4064 {
4065 prefill_queue.remove(pos);
4066 self.request_index.write().remove(&request_id);
4067 self.record_pressure_frontier_terminal(&request_id);
4068 match response.finish_reason {
4069 ferrum_types::FinishReason::EOS
4070 | ferrum_types::FinishReason::Stop
4071 | ferrum_types::FinishReason::Length => {
4072 self.completed_counter.fetch_add(1, Ordering::Relaxed);
4073 }
4074 _ => {
4075 self.failed_counter.fetch_add(1, Ordering::Relaxed);
4076 warn!(
4077 "Request {} completed with error during prefill: {:?}",
4078 request_id, response.finish_reason
4079 );
4080 }
4081 }
4082 self.record_capacity_release_progress();
4083 return Ok(());
4084 }
4085 drop(prefill_queue);
4086
4087 if self
4088 .admission_failed_requests
4089 .write()
4090 .remove(&request_id)
4091 .is_some()
4092 {
4093 self.request_index.write().remove(&request_id);
4094 self.failed_counter.fetch_add(1, Ordering::Relaxed);
4095 return Ok(());
4096 }
4097
4098 warn!("Attempted to complete unknown request: {}", request_id);
4099 Err(FerrumError::scheduler(format!(
4100 "Request {} not found in active queues",
4101 request_id
4102 )))
4103 }
4104 }
4105
4106 async fn cancel(&self, request_id: RequestId) -> Result<bool> {
4107 debug!("Cancelling request {}", request_id);
4108
4109 {
4111 let mut waiting_queue = self.waiting_queue.write();
4112 let waiting_position =
4113 waiting_queue.position(|request| request.inner.request.id == request_id);
4114 if let Some(pos) = waiting_position {
4115 waiting_queue.remove(pos);
4116 self.request_index.write().remove(&request_id);
4117 self.record_pressure_frontier_terminal(&request_id);
4118 self.cancelled_counter.fetch_add(1, Ordering::Relaxed);
4119 info!("Request {} cancelled from waiting queue", request_id);
4120 return Ok(true);
4121 }
4122 }
4123
4124 {
4126 let mut prefill_queue = self.prefill_queue.write();
4127 if let Some(pos) = prefill_queue
4128 .iter()
4129 .position(|r| r.inner.request.id == request_id)
4130 {
4131 prefill_queue.remove(pos);
4132 self.request_index.write().remove(&request_id);
4133 self.record_pressure_frontier_terminal(&request_id);
4134 self.cancelled_counter.fetch_add(1, Ordering::Relaxed);
4135 self.record_capacity_release_progress();
4136 warn!("Request {} cancelled during prefill", request_id);
4137 return Ok(true);
4138 }
4139 }
4140
4141 {
4143 let mut decode_queue = self.decode_queue.write();
4144 if decode_queue.remove(&request_id).is_some() {
4145 self.request_index.write().remove(&request_id);
4146 self.record_pressure_frontier_terminal(&request_id);
4147 self.cancelled_counter.fetch_add(1, Ordering::Relaxed);
4148 self.record_capacity_release_progress();
4149 warn!("Request {} cancelled during decode", request_id);
4150 return Ok(true);
4151 }
4152 drop(decode_queue);
4153 }
4154
4155 if self
4156 .admission_failed_requests
4157 .write()
4158 .remove(&request_id)
4159 .is_some()
4160 {
4161 self.request_index.write().remove(&request_id);
4162 self.cancelled_counter.fetch_add(1, Ordering::Relaxed);
4163 return Ok(true);
4164 }
4165
4166 warn!("Request {} not found for cancellation", request_id);
4167 Ok(false)
4168 }
4169
4170 async fn update_priority(&self, request_id: RequestId, priority: Priority) -> Result<()> {
4171 debug!(
4172 "Updating priority for request {} to {:?}",
4173 request_id, priority
4174 );
4175
4176 {
4178 let mut waiting_queue = self.waiting_queue.write();
4179 if let Some(req) =
4180 waiting_queue.find_mut(|request| request.inner.request.id == request_id)
4181 {
4182 req.inner.request.priority = priority;
4183 return Ok(());
4184 }
4185 }
4186
4187 {
4189 let mut prefill_queue = self.prefill_queue.write();
4190 if let Some(req) = prefill_queue
4191 .iter_mut()
4192 .find(|r| r.inner.request.id == request_id)
4193 {
4194 req.inner.request.priority = priority;
4195 return Ok(());
4196 }
4197 }
4198
4199 {
4201 let mut decode_queue = self.decode_queue.write();
4202 if let Some(req) = decode_queue.requests.get_mut(&request_id) {
4203 req.inner.request.priority = priority;
4204 return Ok(());
4205 }
4206 }
4207
4208 Ok(())
4209 }
4210
4211 fn metrics(&self) -> SchedulerMetrics {
4212 let waiting_count = self.waiting_count();
4213 let prefill_count = self.prefilling_count();
4214 let decode_count = self.decoding_count();
4215 let running_count = prefill_count + decode_count;
4216
4217 let completed_count = self.completed_counter.load(Ordering::Relaxed);
4218 let failed_count = self.failed_counter.load(Ordering::Relaxed);
4219 let cancelled_count = self.cancelled_counter.load(Ordering::Relaxed);
4220 let preempted_count = self.preempted_counter.load(Ordering::Relaxed);
4221 let admitted_count = self.admitted_counter.load(Ordering::Relaxed);
4222 let total_wait_time_us = self.total_wait_time_us.load(Ordering::Relaxed);
4223
4224 let uptime_secs = self.start_time.elapsed().as_secs_f64();
4225 let throughput = if uptime_secs > 0.0 {
4226 completed_count as f64 / uptime_secs
4227 } else {
4228 0.0
4229 };
4230
4231 let queue_utilization = waiting_count as f32 / self.config.max_waiting_requests as f32;
4232 let avg_wait_time_ms = if admitted_count > 0 {
4233 total_wait_time_us as f64 / admitted_count as f64 / 1000.0
4234 } else {
4235 0.0
4236 };
4237
4238 ferrum_types::SchedulerStats {
4239 waiting_requests: waiting_count,
4240 running_requests: running_count,
4241 preempted_requests: preempted_count as usize,
4242 completed_requests: completed_count,
4243 failed_requests: failed_count,
4244 cancelled_requests: cancelled_count,
4245 avg_wait_time_ms,
4246 avg_execution_time_ms: 0.0,
4247 throughput_rps: throughput,
4248 queue_utilization,
4249 }
4250 }
4251
4252 fn config(&self) -> &SchedulerConfig {
4253 &self.config
4254 }
4255
4256 fn request_state(&self, request_id: &RequestId) -> Option<RequestState> {
4257 self.request_index
4258 .read()
4259 .get(request_id)
4260 .copied()
4261 .map(|phase| match phase {
4262 RequestPhase::Waiting => RequestState::Waiting,
4263 RequestPhase::Prefilling | RequestPhase::Decoding => RequestState::Running,
4264 RequestPhase::Completed => RequestState::Completed,
4265 RequestPhase::Preempted => RequestState::Preempted,
4266 RequestPhase::Cancelled => RequestState::Cancelled,
4267 RequestPhase::AdmissionFailed => RequestState::Failed,
4268 })
4269 }
4270
4271 async fn preempt(&self, request_id: RequestId) -> Result<PreemptionResult> {
4272 if !self.cb_config.enable_swapping {
4273 return Err(FerrumError::unsupported("Swapping is not enabled"));
4274 }
4275
4276 debug!("Preempting request {}", request_id);
4277
4278 let mut decode_queue = self.decode_queue.write();
4280 if let Some(mut req) = decode_queue.remove(&request_id) {
4281 req.phase = RequestPhase::Preempted;
4282
4283 let state = PreemptionState {
4285 kv_cache_checkpoint: Vec::new(), tokens_processed: req.total_tokens(),
4287 generation_state: HashMap::new(),
4288 };
4289
4290 let freed_resources = req.inner.allocated_resources.clone();
4291
4292 self.preempted_requests
4294 .write()
4295 .insert(request_id.clone(), req);
4296 self.request_index
4297 .write()
4298 .insert(request_id, RequestPhase::Preempted);
4299 self.preempted_counter.fetch_add(1, Ordering::Relaxed);
4300 self.record_capacity_release_progress();
4301
4302 Ok(PreemptionResult {
4303 success: true,
4304 saved_state: Some(state),
4305 freed_resources,
4306 })
4307 } else {
4308 Err(FerrumError::scheduler(format!(
4309 "Request {} not found in decode queue",
4310 request_id
4311 )))
4312 }
4313 }
4314
4315 async fn resume(&self, request_id: RequestId) -> Result<()> {
4316 debug!("Resuming request {}", request_id);
4317
4318 let mut preempted = self.preempted_requests.write();
4319 if let Some(mut req) = preempted.remove(&request_id) {
4320 req.phase = RequestPhase::Decoding;
4321
4322 self.decode_queue
4323 .write()
4324 .requests
4325 .insert(request_id.clone(), req);
4326 self.request_index
4327 .write()
4328 .insert(request_id, RequestPhase::Decoding);
4329
4330 Ok(())
4331 } else {
4332 Err(FerrumError::scheduler(format!(
4333 "Request {} not found in preempted queue",
4334 request_id
4335 )))
4336 }
4337 }
4338}
4339
4340impl std::fmt::Debug for ContinuousBatchScheduler {
4341 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4342 f.debug_struct("ContinuousBatchScheduler")
4343 .field("waiting", &self.waiting_count())
4344 .field("prefilling", &self.prefilling_count())
4345 .field("decoding", &self.decoding_count())
4346 .field("iteration", &self.current_iteration.load(Ordering::Relaxed))
4347 .finish()
4348 }
4349}
4350
4351#[cfg(test)]
4356mod tests {
4357 use super::*;
4358 use ferrum_types::{ModelId, SamplingParams};
4359
4360 fn create_test_request(priority: Priority) -> InferenceRequest {
4361 InferenceRequest {
4362 id: RequestId::new(),
4363 prompt: "test".to_string(),
4364 model_id: ModelId::new("test-model"),
4365 sampling_params: SamplingParams::default(),
4366 stream: false,
4367 priority,
4368 client_id: None,
4369 session_id: None,
4370 created_at: chrono::Utc::now(),
4371 api_request: None,
4372 evidence_request: Default::default(),
4373 metadata: std::collections::HashMap::new(),
4374 }
4375 }
4376
4377 fn create_test_request_with_prompt_tokens(
4378 priority: Priority,
4379 prompt_tokens: usize,
4380 ) -> InferenceRequest {
4381 create_test_request(priority).with_metadata(
4382 PROMPT_TOKENS_METADATA_KEY,
4383 serde_json::Value::from(prompt_tokens as u64),
4384 )
4385 }
4386
4387 fn enqueue_waiting(scheduler: &ContinuousBatchScheduler, request: InferenceRequest) {
4388 let request_id = request.id.clone();
4389 let mut waiting = scheduler.waiting_queue.write();
4390 let ticket = waiting
4391 .enqueue(ContinuousBatchRequest::new(request))
4392 .unwrap();
4393 waiting
4394 .request_mut(ticket)
4395 .unwrap()
4396 .waiting_admission_ticket = Some(ticket);
4397 drop(waiting);
4398 scheduler
4399 .request_index
4400 .write()
4401 .insert(request_id, RequestPhase::Waiting);
4402 }
4403
4404 fn execution_capacity_release_snapshot<'a>(
4405 request_ids: impl IntoIterator<Item = &'a RequestId>,
4406 condition: &CapacityWaitCondition,
4407 ) -> ExecutionCapacityReleaseSnapshot {
4408 let sources = condition
4409 .observed()
4410 .iter()
4411 .map(|observed| observed.source())
4412 .collect::<Vec<_>>();
4413 ExecutionCapacityReleaseSnapshot::new(
4414 request_ids
4415 .into_iter()
4416 .map(|request_id| (request_id.clone(), sources.clone())),
4417 )
4418 }
4419
4420 #[tokio::test]
4421 async fn test_scheduler_creation() {
4422 let config = SchedulerConfig::default();
4423 let scheduler = ContinuousBatchScheduler::new(config);
4424 assert_eq!(scheduler.waiting_count(), 0);
4425 assert_eq!(scheduler.active_count(), 0);
4426 }
4427
4428 #[test]
4429 fn admission_phase_counts_are_mutually_exclusive_per_request() {
4430 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
4431 let phases = [
4432 RequestPhase::Waiting,
4433 RequestPhase::Prefilling,
4434 RequestPhase::Decoding,
4435 RequestPhase::Decoding,
4436 RequestPhase::Completed,
4437 RequestPhase::Cancelled,
4438 ];
4439 {
4440 let mut request_index = scheduler.request_index.write();
4441 for phase in phases {
4442 request_index.insert(RequestId::new(), phase);
4443 }
4444 }
4445
4446 let counts = scheduler.admission_phase_counts();
4447 assert_eq!(counts.waiting_requests, 1);
4448 assert_eq!(counts.active_prefill_sequences, 1);
4449 assert_eq!(counts.active_decode_sequences, 2);
4450 assert_eq!(
4451 counts.waiting_requests
4452 + counts.active_prefill_sequences
4453 + counts.active_decode_sequences,
4454 4
4455 );
4456 }
4457
4458 #[tokio::test]
4459 async fn test_submit_and_counts() {
4460 let config = SchedulerConfig::default();
4461 let scheduler = ContinuousBatchScheduler::new(config);
4462
4463 scheduler
4464 .submit(create_test_request(Priority::Normal))
4465 .await
4466 .unwrap();
4467 scheduler
4468 .submit(create_test_request(Priority::High))
4469 .await
4470 .unwrap();
4471
4472 assert_eq!(scheduler.waiting_count(), 2);
4473 assert_eq!(scheduler.active_count(), 0);
4474 }
4475
4476 #[tokio::test]
4477 async fn trace_snapshot_reports_queue_counters_and_phase() {
4478 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
4479 let request = create_test_request(Priority::Normal);
4480 let request_id = request.id.clone();
4481 scheduler.submit(request).await.unwrap();
4482
4483 let before = scheduler.trace_snapshot();
4484 assert_eq!(before.waiting_queue_len, 1);
4485 assert_eq!(before.active_len, 0);
4486 assert_eq!(before.admitted_total, 0);
4487 assert_eq!(before.capacity_release_epoch, 0);
4488 assert_eq!(before.capacity_mixed_recompute_epoch, 0);
4489 assert_eq!(before.capacity_mixed_recompute_blocked_until_epoch, 0);
4490 assert_eq!(
4491 before.capacity_mixed_recompute_required_blocks_per_slot,
4492 None
4493 );
4494 assert_eq!(before.capacity_mixed_recompute_observed_free_blocks, None);
4495 assert_eq!(before.decode_capacity_backpressure_admit_limit, None);
4496 assert_eq!(
4497 scheduler.trace_phase(&request_id),
4498 Some(RequestPhase::Waiting)
4499 );
4500
4501 let batch = scheduler.next_batch(BatchHint::simple(4)).await.unwrap();
4502 assert_eq!(batch.size(), 1);
4503
4504 let after = scheduler.trace_snapshot();
4505 assert_eq!(after.waiting_queue_len, 0);
4506 assert_eq!(after.prefill_queue_len, 1);
4507 assert_eq!(after.active_len, 1);
4508 assert_eq!(after.admitted_total, 1);
4509 assert_eq!(
4510 scheduler.trace_phase(&request_id),
4511 Some(RequestPhase::Prefilling)
4512 );
4513 }
4514
4515 #[test]
4516 fn trace_snapshot_releases_prefill_read_before_later_snapshot_work() {
4517 use std::sync::{mpsc::sync_channel, Arc, Barrier};
4518 use std::time::Duration;
4519
4520 const WRITER_WAIT: Duration = Duration::from_millis(250);
4521
4522 let scheduler = Arc::new(ContinuousBatchScheduler::new(SchedulerConfig::default()));
4523 let writer_scheduler = Arc::clone(&scheduler);
4524 let writer_start = Arc::new(Barrier::new(2));
4525 let writer_barrier = Arc::clone(&writer_start);
4526 let (attempting_tx, attempting_rx) = sync_channel(1);
4527 let (acquired_tx, acquired_rx) = sync_channel(1);
4528 let writer = std::thread::spawn(move || {
4529 writer_barrier.wait();
4530 attempting_tx.send(()).unwrap();
4531 let acquired = writer_scheduler
4532 .prefill_queue
4533 .try_write_for(WRITER_WAIT)
4534 .is_some();
4535 acquired_tx.send(acquired).unwrap();
4536 });
4537
4538 let snapshot = scheduler.trace_snapshot_with_prefill_read_observer(|| {
4539 writer_start.wait();
4540 attempting_rx
4541 .recv_timeout(Duration::from_secs(1))
4542 .expect("bounded writer must begin its fair-lock attempt");
4543 std::thread::sleep(Duration::from_millis(20));
4544 });
4545 let writer_acquired = acquired_rx
4546 .recv_timeout(Duration::from_secs(1))
4547 .expect("bounded writer must finish its fair-lock attempt");
4548 writer.join().expect("bounded snapshot writer must join");
4549
4550 assert!(
4551 writer_acquired,
4552 "trace snapshot must release its first prefill read before later snapshot work"
4553 );
4554 assert_eq!(snapshot.prefill_queue_len, 0);
4555 assert_eq!(snapshot.execution_capacity_blocked_prefill_len, 0);
4556 assert_eq!(snapshot.execution_readiness_blocked_prefill_len, 0);
4557 }
4558
4559 #[tokio::test]
4560 async fn execution_maintenance_retry_yields_one_iteration_without_opening_pressure() {
4561 let mut config = SchedulerConfig::default();
4562 config.max_running_requests = 2;
4563 let scheduler = ContinuousBatchScheduler::new(config);
4564 let maintained = create_test_request(Priority::Normal);
4565 let maintained_id = maintained.id.clone();
4566 let peer = create_test_request(Priority::Normal);
4567 let peer_id = peer.id.clone();
4568 scheduler.submit(maintained).await.unwrap();
4569 scheduler.submit(peer).await.unwrap();
4570
4571 let initial = scheduler
4572 .create_iteration_batch(BatchHint::simple(2))
4573 .unwrap();
4574 assert_eq!(initial.requests.len(), 2);
4575
4576 let receipt = scheduler
4577 .defer_retry_after_execution_maintenance_epoch(std::slice::from_ref(&maintained_id), 7)
4578 .unwrap();
4579 assert_eq!(receipt.deferred_count(), 1);
4580 assert_eq!(receipt.latest_capacity_epoch(), 7);
4581
4582 let fairness_iteration = scheduler
4583 .create_iteration_batch(BatchHint::simple(1))
4584 .unwrap();
4585 assert_eq!(fairness_iteration.requests.len(), 1);
4586 assert_eq!(fairness_iteration.requests[0].request.id, peer_id);
4587 let pressure = scheduler.trace_snapshot();
4588 assert_eq!(pressure.pressure_active_episodes, 0);
4589 assert_eq!(pressure.pressure_pending_release_fences, 0);
4590
4591 let retry_iteration = scheduler
4592 .create_iteration_batch(BatchHint::simple(1))
4593 .unwrap();
4594 assert_eq!(retry_iteration.requests.len(), 1);
4595 assert_eq!(retry_iteration.requests[0].request.id, maintained_id);
4596
4597 let replay = scheduler
4598 .defer_retry_after_execution_maintenance_epoch(std::slice::from_ref(&maintained_id), 7)
4599 .unwrap_err();
4600 assert!(replay
4601 .to_string()
4602 .contains("stale or concurrently blocked evidence"));
4603 }
4604
4605 #[tokio::test]
4606 async fn execution_maintenance_retry_is_atomic_when_queue_identity_is_inconsistent() {
4607 let mut config = SchedulerConfig::default();
4608 config.max_running_requests = 2;
4609 let scheduler = ContinuousBatchScheduler::new(config);
4610 let first = create_test_request(Priority::Normal);
4611 let first_id = first.id.clone();
4612 let second = create_test_request(Priority::Normal);
4613 let second_id = second.id.clone();
4614 scheduler.submit(first).await.unwrap();
4615 scheduler.submit(second).await.unwrap();
4616 scheduler
4617 .create_iteration_batch(BatchHint::simple(2))
4618 .unwrap();
4619
4620 let removed = {
4621 let mut prefill = scheduler.prefill_queue.write();
4622 let position = prefill
4623 .iter()
4624 .position(|request| request.inner.request.id == second_id)
4625 .unwrap();
4626 prefill.remove(position).unwrap()
4627 };
4628 let error = scheduler
4629 .defer_retry_after_execution_maintenance_epoch(&[first_id.clone(), second_id], 11)
4630 .unwrap_err();
4631 assert!(error
4632 .to_string()
4633 .contains("lost an exact active logical frontier"));
4634 let prefill = scheduler.prefill_queue.read();
4635 let first = prefill
4636 .iter()
4637 .find(|request| request.inner.request.id == first_id)
4638 .unwrap();
4639 assert!(first.execution_maintenance_retry.is_none());
4640 assert!(first.last_execution_maintenance_capacity_epoch.is_none());
4641 drop(prefill);
4642 drop(removed);
4643 }
4644
4645 #[tokio::test]
4646 async fn typed_dynamic_admission_defers_without_blocking_decode_or_smaller_work() {
4647 use ferrum_interfaces::vnext::{
4648 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
4649 DeferredAction,
4650 };
4651 use std::num::NonZeroU64;
4652
4653 let mut config = SchedulerConfig::default();
4654 config.max_running_requests = 2;
4655 let scheduler = ContinuousBatchScheduler::new(config);
4656 let large = create_test_request_with_prompt_tokens(Priority::Normal, 512);
4657 let large_id = large.id.clone();
4658 let small = create_test_request_with_prompt_tokens(Priority::Normal, 8);
4659 let small_id = small.id.clone();
4660 scheduler.submit(large).await.unwrap();
4661 scheduler.submit(small).await.unwrap();
4662
4663 let wake0 = AdmissionWakeEpochs::new(NonZeroU64::new(19).unwrap(), 0, 0, 0);
4664 let availability0 =
4665 [
4666 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 1)
4667 .unwrap(),
4668 ];
4669 let condition0 =
4670 CapacityWaitCondition::from_observation(19, availability0.to_vec()).unwrap();
4671 let mut first_probes = Vec::new();
4672 let mut first_probe = |request: &InferenceRequest| {
4673 first_probes.push(request.id.clone());
4674 if request.id == large_id {
4675 AdmissionProbeOutcome::Deferred(crate::vnext::AdmissionDeferral::new(
4676 DeferredAction::WaitForRelease,
4677 wake0,
4678 condition0.clone(),
4679 ))
4680 } else {
4681 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
4682 request_id: request.id.clone(),
4683 })
4684 }
4685 };
4686 let first = scheduler
4687 .next_batch_with_dynamic_admission(
4688 BatchHint::simple(2),
4689 AdmissionWakeSnapshot::new(wake0, &availability0),
4690 &mut first_probe,
4691 )
4692 .unwrap()
4693 .unwrap();
4694 assert_eq!(first_probes, vec![large_id.clone(), small_id.clone()]);
4695 assert_eq!(first.requests.len(), 1);
4696 assert_eq!(first.requests[0].request.id, small_id);
4697 assert_eq!(
4698 scheduler.trace_phase(&large_id),
4699 Some(RequestPhase::Waiting)
4700 );
4701
4702 scheduler.mark_prefill_complete(&small_id, 8);
4703 let mut unchanged_probe_count = 0;
4704 let mut unchanged_probe = |request: &InferenceRequest| {
4705 unchanged_probe_count += 1;
4706 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
4707 request_id: request.id.clone(),
4708 })
4709 };
4710 let mut observations = Vec::new();
4711 let unchanged = scheduler
4712 .next_batch_with_dynamic_admission_observed(
4713 BatchHint::simple(2),
4714 AdmissionWakeSnapshot::new(wake0, &availability0),
4715 &mut unchanged_probe,
4716 &mut |observation| observations.push(observation),
4717 )
4718 .unwrap()
4719 .unwrap();
4720 assert_eq!(unchanged_probe_count, 0);
4721 assert_eq!(unchanged.requests.len(), 1);
4722 assert_eq!(unchanged.requests[0].request.id, small_id);
4723 assert!(matches!(
4724 observations.as_slice(),
4725 [ExecutorAdmissionQueueObservation::SkippedUnchanged {
4726 request_id,
4727 deferral,
4728 current,
4729 ..
4730 }] if request_id == &large_id
4731 && deferral.observed() == wake0
4732 && *current == wake0
4733 ));
4734
4735 let wake1 = AdmissionWakeEpochs::new(NonZeroU64::new(19).unwrap(), 1, 0, 0);
4736 let availability1 =
4737 [
4738 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 2)
4739 .unwrap(),
4740 ];
4741 let mut released_probe_count = 0;
4742 let mut released_probe = |request: &InferenceRequest| {
4743 released_probe_count += 1;
4744 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
4745 request_id: request.id.clone(),
4746 })
4747 };
4748 let released = scheduler
4749 .next_batch_with_dynamic_admission(
4750 BatchHint::simple(2),
4751 AdmissionWakeSnapshot::new(wake1, &availability1),
4752 &mut released_probe,
4753 )
4754 .unwrap()
4755 .unwrap();
4756 assert_eq!(released_probe_count, 1);
4757 assert_eq!(released.requests.len(), 2);
4758 assert!(released
4759 .requests
4760 .iter()
4761 .any(|request| request.request.id == large_id));
4762 let trace = scheduler.trace_snapshot();
4763 assert_eq!(trace.legacy_waiting_admission_ticks, 0);
4764 assert_eq!(trace.dynamic_admission_probes, 3);
4765 assert_eq!(trace.dynamic_admission_skipped_unchanged, 1);
4766 assert_eq!(trace.dynamic_admission_deferred, 1);
4767 }
4768
4769 #[tokio::test]
4770 async fn open_pressure_episode_closes_after_resumed_decode_commits_progress() {
4771 use ferrum_interfaces::vnext::{
4772 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityDomainId,
4773 CapacityWaitCondition, DeferredAction,
4774 };
4775 use std::num::NonZeroU64;
4776
4777 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
4778 let blocked = create_test_request(Priority::Normal);
4779 let blocked_id = blocked.id.clone();
4780 let peer = create_test_request(Priority::Normal);
4781 let peer_id = peer.id.clone();
4782 scheduler.submit(blocked).await.unwrap();
4783 scheduler.submit(peer).await.unwrap();
4784
4785 let source = CapacityAvailabilitySource::Domain(CapacityDomainId::new(8).unwrap());
4786 let availability0 = [CapacityAvailabilityEpoch::new(source, 1).unwrap()];
4787 let wake0 = AdmissionWakeEpochs::new(NonZeroU64::new(29).unwrap(), 0, 0, 0);
4788 let admitted = scheduler
4789 .next_batch_with_dynamic_admission(
4790 BatchHint::simple(2),
4791 AdmissionWakeSnapshot::new(wake0, &availability0),
4792 &mut |request| {
4793 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
4794 request_id: request.id.clone(),
4795 })
4796 },
4797 )
4798 .unwrap()
4799 .unwrap();
4800 assert_eq!(admitted.size(), 2);
4801 scheduler.mark_prefill_complete(&blocked_id, 1);
4802 scheduler.mark_prefill_complete(&peer_id, 1);
4803
4804 let condition =
4805 CapacityWaitCondition::from_observation(29, availability0.to_vec()).unwrap();
4806 assert_eq!(
4807 scheduler
4808 .defer_decode_for_execution_capacity(
4809 std::slice::from_ref(&blocked_id),
4810 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake0, condition),
4811 &ExecutionCapacityReleaseSnapshot::default(),
4812 )
4813 .unwrap(),
4814 ExecutionCapacityAction::Deferred { count: 1 }
4815 );
4816 assert_eq!(scheduler.trace_snapshot().pressure_active_episodes, 1);
4817
4818 let availability1 = [CapacityAvailabilityEpoch::new(source, 2).unwrap()];
4819 let wake1 = AdmissionWakeEpochs::new(NonZeroU64::new(29).unwrap(), 1, 0, 0);
4820 let resumed = scheduler
4821 .next_batch_with_dynamic_admission(
4822 BatchHint::simple(2),
4823 AdmissionWakeSnapshot::new(wake1, &availability1),
4824 &mut |_| panic!("decode resume must not probe waiting admission"),
4825 )
4826 .unwrap()
4827 .expect("the exact-source wake must make the blocked decode schedulable");
4828 assert!(resumed
4829 .requests
4830 .iter()
4831 .any(|request| request.request.id == blocked_id));
4832 assert_eq!(
4833 scheduler.trace_snapshot().pressure_active_episodes,
4834 1,
4835 "an epoch change permits retry but is not execution-success evidence"
4836 );
4837
4838 scheduler.update_decode_progress(&blocked_id, 1);
4839 let snapshot = scheduler.trace_snapshot();
4840 assert_eq!(snapshot.pressure_active_episodes, 0);
4841 let journal = scheduler.pressure_transition_journal();
4842 let satisfied = journal
4843 .iter()
4844 .find(|transition| transition.kind() == PressureTransitionKind::WaitSatisfied)
4845 .unwrap();
4846 let closed = journal
4847 .iter()
4848 .find(|transition| transition.kind() == PressureTransitionKind::Closed)
4849 .unwrap();
4850 assert!(satisfied.ordinal() < closed.ordinal());
4851 }
4852
4853 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
4854 enum SchedulerReplayRole {
4855 Blocked,
4856 Runnable,
4857 }
4858
4859 #[derive(Debug, Clone, PartialEq, Eq)]
4860 struct SchedulerReplayBatchMember {
4861 role: SchedulerReplayRole,
4862 tokens_processed: usize,
4863 tokens_to_process: Option<usize>,
4864 }
4865
4866 #[derive(Debug, Clone, PartialEq, Eq)]
4867 struct SchedulerReplayProjection {
4868 batches: Vec<Vec<SchedulerReplayBatchMember>>,
4869 traces: Vec<ContinuousSchedulerTraceSnapshot>,
4870 observations: Vec<Vec<ExecutorAdmissionQueueObservation>>,
4871 journal: Vec<PressureTransition>,
4872 }
4873
4874 fn scheduler_replay_batch(
4875 batch: &BatchPlan,
4876 blocked_id: &RequestId,
4877 runnable_id: &RequestId,
4878 ) -> Vec<SchedulerReplayBatchMember> {
4879 batch
4880 .requests
4881 .iter()
4882 .map(|request| {
4883 let role = if request.request.id == *blocked_id {
4884 SchedulerReplayRole::Blocked
4885 } else if request.request.id == *runnable_id {
4886 SchedulerReplayRole::Runnable
4887 } else {
4888 panic!("unexpected scheduler replay request {}", request.request.id);
4889 };
4890 SchedulerReplayBatchMember {
4891 role,
4892 tokens_processed: request.tokens_processed,
4893 tokens_to_process: request.tokens_to_process,
4894 }
4895 })
4896 .collect()
4897 }
4898
4899 async fn collect_cross_phase_pressure_replay() -> SchedulerReplayProjection {
4900 use ferrum_interfaces::vnext::{
4901 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
4902 DeferredAction,
4903 };
4904 use std::num::NonZeroU64;
4905
4906 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
4907 let mut blocked = create_test_request(Priority::Normal);
4908 blocked.id = RequestId(uuid::Uuid::from_u128(1));
4909 let blocked_id = blocked.id.clone();
4910 let mut runnable = create_test_request(Priority::Normal);
4911 runnable.id = RequestId(uuid::Uuid::from_u128(2));
4912 let runnable_id = runnable.id.clone();
4913 let mut batches = Vec::new();
4914 let mut traces = Vec::new();
4915 let mut observation_batches = Vec::new();
4916 scheduler.submit(blocked).await.unwrap();
4917 scheduler.submit(runnable).await.unwrap();
4918
4919 let source = CapacityAvailabilitySource::ActiveSequenceSlots;
4920 let availability0 = [CapacityAvailabilityEpoch::new(source, 1).unwrap()];
4921 let wake0 = AdmissionWakeEpochs::new(NonZeroU64::new(29).unwrap(), 0, 0, 0);
4922 let admitted = scheduler
4923 .next_batch_with_dynamic_admission(
4924 BatchHint::simple(2),
4925 AdmissionWakeSnapshot::new(wake0, &availability0),
4926 &mut |request| {
4927 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
4928 request_id: request.id.clone(),
4929 })
4930 },
4931 )
4932 .unwrap()
4933 .unwrap();
4934 assert_eq!(admitted.size(), 2);
4935 batches.push(scheduler_replay_batch(&admitted, &blocked_id, &runnable_id));
4936 scheduler.mark_prefill_complete(&blocked_id, 1);
4937 scheduler.mark_prefill_complete(&runnable_id, 1);
4938 traces.push(scheduler.trace_snapshot());
4939
4940 let condition =
4941 CapacityWaitCondition::from_observation(29, availability0.to_vec()).unwrap();
4942 let deferral =
4943 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake0, condition.clone());
4944 assert_eq!(
4945 scheduler
4946 .defer_decode_for_execution_capacity(
4947 std::slice::from_ref(&blocked_id),
4948 deferral.clone(),
4949 &execution_capacity_release_snapshot([&blocked_id, &runnable_id], &condition,),
4950 )
4951 .unwrap(),
4952 ExecutionCapacityAction::Deferred { count: 1 }
4953 );
4954 traces.push(scheduler.trace_snapshot());
4955 let mut observations = Vec::new();
4956 let bypass = scheduler
4957 .next_batch_with_dynamic_admission_observed(
4958 BatchHint::simple(2),
4959 AdmissionWakeSnapshot::new(wake0, &availability0),
4960 &mut |_| panic!("no waiting admission probe is allowed"),
4961 &mut |observation| observations.push(observation),
4962 )
4963 .unwrap()
4964 .expect("runnable decode must bypass an unchanged blocked decode");
4965 assert_eq!(bypass.requests.len(), 1);
4966 assert_eq!(bypass.requests[0].request.id, runnable_id);
4967 batches.push(scheduler_replay_batch(&bypass, &blocked_id, &runnable_id));
4968 assert!(matches!(
4969 observations.as_slice(),
4970 [ExecutorAdmissionQueueObservation::DecodeSkippedUnchanged {
4971 request_id,
4972 current_wait_sources,
4973 ..
4974 }] if request_id == &blocked_id && current_wait_sources == &availability0
4975 ));
4976 observation_batches.push(observations.clone());
4977
4978 let unavailable_action = scheduler
4979 .defer_decode_for_execution_capacity(
4980 std::slice::from_ref(&runnable_id),
4981 deferral.clone(),
4982 &ExecutionCapacityReleaseSnapshot::default(),
4983 )
4984 .unwrap();
4985 assert!(matches!(
4986 unavailable_action,
4987 ExecutionCapacityAction::InvariantViolation {
4988 violation
4989 } if violation.class() == PressureInvariantViolationClass::NoReleasableFrontier
4990 ));
4991
4992 let action = scheduler
4993 .defer_decode_for_execution_capacity(
4994 std::slice::from_ref(&runnable_id),
4995 deferral,
4996 &execution_capacity_release_snapshot([&blocked_id, &runnable_id], &condition),
4997 )
4998 .unwrap();
4999 let ExecutionCapacityAction::YieldPlanned { transaction } = action else {
5000 panic!("the last runnable frontier must produce a pressure yield");
5001 };
5002 assert_eq!(transaction.victim_request_id(), &runnable_id);
5003 assert_eq!(transaction.progress_owner_id(), &blocked_id);
5004 let armed = scheduler
5005 .arm_execution_capacity_yield(&transaction)
5006 .unwrap();
5007 let completion = scheduler
5008 .complete_execution_capacity_yield(&transaction, 1, Some(0))
5009 .unwrap();
5010 let released = completion.release_transition_ordinal();
5011 let resumable = completion
5012 .resumable_transition_ordinal()
5013 .expect("live progress owner must become resumable");
5014 assert!(completion.victim_requeued());
5015 assert!(completion.progress_owner_resumable());
5016 assert!(completion.closed_transition_ordinal().is_none());
5017 assert!(transaction.planned_ordinal() < armed);
5018 assert!(armed < released);
5019 assert!(released < resumable);
5020 assert_eq!(
5021 scheduler
5022 .trace_snapshot()
5023 .execution_capacity_blocked_decode_len,
5024 1
5025 );
5026 assert_eq!(
5027 scheduler.trace_phase(&runnable_id),
5028 Some(RequestPhase::Waiting)
5029 );
5030 assert!(scheduler
5031 .passive_capacity_wait_condition()
5032 .unwrap()
5033 .is_some());
5034 traces.push(scheduler.trace_snapshot());
5035
5036 let availability1 = [CapacityAvailabilityEpoch::new(source, 2).unwrap()];
5037 let released = AdmissionWakeEpochs::new(NonZeroU64::new(29).unwrap(), 2, 0, 0);
5038 observations.clear();
5039 let resumed = scheduler
5040 .next_batch_with_dynamic_admission_observed(
5041 BatchHint::simple(2),
5042 AdmissionWakeSnapshot::new(released, &availability1),
5043 &mut |_| panic!("decode resume does not probe waiting admission"),
5044 &mut |observation| observations.push(observation),
5045 )
5046 .unwrap()
5047 .expect("relevant source change must resume the selected progress owner");
5048 assert_eq!(resumed.requests.len(), 1);
5049 assert_eq!(resumed.requests[0].request.id, blocked_id);
5050 batches.push(scheduler_replay_batch(&resumed, &blocked_id, &runnable_id));
5051 let resumed_ids = observations
5052 .iter()
5053 .filter_map(|observation| match observation {
5054 ExecutorAdmissionQueueObservation::DecodeResumed {
5055 request_id,
5056 current_wait_sources,
5057 exact_source_changed: true,
5058 policy_epoch_changed: false,
5059 ..
5060 } if current_wait_sources == &availability1 => Some(request_id.clone()),
5061 _ => None,
5062 })
5063 .collect::<HashSet<_>>();
5064 assert_eq!(resumed_ids, HashSet::from([blocked_id.clone()]));
5065 observation_batches.push(observations.clone());
5066 assert_eq!(
5067 scheduler
5068 .trace_snapshot()
5069 .execution_capacity_blocked_decode_len,
5070 0
5071 );
5072
5073 scheduler.update_decode_progress(&blocked_id, 1);
5074 observations.clear();
5075 let owner_only = scheduler
5076 .next_batch_with_dynamic_admission_observed(
5077 BatchHint::simple(2),
5078 AdmissionWakeSnapshot::new(released, &availability1),
5079 &mut |_| panic!("owner token progress must not probe a held victim"),
5080 &mut |observation| observations.push(observation),
5081 )
5082 .unwrap()
5083 .expect("progress owner must continue while the victim remains held");
5084 assert_eq!(owner_only.requests.len(), 1);
5085 assert_eq!(owner_only.requests[0].request.id, blocked_id);
5086 batches.push(scheduler_replay_batch(
5087 &owner_only,
5088 &blocked_id,
5089 &runnable_id,
5090 ));
5091 assert!(!observations.iter().any(|observation| matches!(
5092 observation,
5093 ExecutorAdmissionQueueObservation::PressureHoldReleased {
5094 request_id,
5095 ..
5096 } if request_id == &runnable_id
5097 )));
5098 observation_batches.push(observations.clone());
5099 assert_eq!(scheduler.trace_snapshot().pressure_active_episodes, 1);
5100 traces.push(scheduler.trace_snapshot());
5101
5102 let availability2 = [CapacityAvailabilityEpoch::new(source, 3).unwrap()];
5103 let wake2 = AdmissionWakeEpochs::new(NonZeroU64::new(29).unwrap(), 3, 0, 0);
5104 let owner_pressure =
5105 CapacityWaitCondition::from_observation(29, availability2.to_vec()).unwrap();
5106 let owner_deferral = AdmissionDeferral::new(
5107 DeferredAction::WaitForRelease,
5108 wake2,
5109 owner_pressure.clone(),
5110 );
5111 let owner_action = scheduler
5112 .defer_decode_for_execution_capacity(
5113 std::slice::from_ref(&blocked_id),
5114 owner_deferral,
5115 &execution_capacity_release_snapshot([&blocked_id], &owner_pressure),
5116 )
5117 .unwrap();
5118 let ExecutionCapacityAction::YieldPlanned {
5119 transaction: owner_transaction,
5120 } = owner_action
5121 else {
5122 panic!("the progressed owner must rotate at renewed pressure");
5123 };
5124 assert_eq!(owner_transaction.kind(), PressureYieldKind::PeerHandoff);
5125 assert_eq!(owner_transaction.progress_owner_id(), &runnable_id);
5126 assert_eq!(owner_transaction.victim_request_id(), &blocked_id);
5127 assert_eq!(
5128 owner_transaction.rotated_from_progress_owner_id(),
5129 Some(&blocked_id)
5130 );
5131 assert!(owner_transaction
5132 .rotated_from_progress_current()
5133 .is_some_and(
5134 |current| current > owner_transaction.rotated_from_progress_baseline().unwrap()
5135 ));
5136 scheduler
5137 .arm_execution_capacity_yield(&owner_transaction)
5138 .unwrap();
5139 let owner_completion = scheduler
5140 .complete_execution_capacity_yield(&owner_transaction, 1, Some(0))
5141 .unwrap();
5142 assert_eq!(
5143 owner_completion.disposition(),
5144 ExecutionCapacityYieldDisposition::ProgressOwnerAdmissionPending
5145 );
5146 assert!(owner_completion
5147 .owner_admission_pending_transition_ordinal()
5148 .is_some());
5149 assert!(owner_completion.closed_transition_ordinal().is_none());
5150 assert_eq!(scheduler.trace_snapshot().waiting_queue_len, 2);
5151
5152 let mut owner_recompute_probes = Vec::new();
5153 let rotated_owner = scheduler
5154 .next_batch_with_dynamic_admission(
5155 BatchHint::simple(2),
5156 AdmissionWakeSnapshot::new(wake2, &availability2),
5157 &mut |request| {
5158 owner_recompute_probes.push(request.id.clone());
5159 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5160 request_id: request.id.clone(),
5161 })
5162 },
5163 )
5164 .unwrap()
5165 .expect("the oldest held peer must be the only admission-eligible recompute");
5166 assert_eq!(owner_recompute_probes, vec![runnable_id.clone()]);
5167 assert_eq!(rotated_owner.requests.len(), 1);
5168 assert_eq!(rotated_owner.requests[0].request.id, runnable_id);
5169 batches.push(scheduler_replay_batch(
5170 &rotated_owner,
5171 &blocked_id,
5172 &runnable_id,
5173 ));
5174 assert!(matches!(
5175 scheduler
5176 .pressure_coordinator
5177 .lock()
5178 .hold_status(&blocked_id),
5179 PressureHoldStatus::Held { .. }
5180 ));
5181 let journal = scheduler.pressure_transition_journal();
5182 let admission_pending = journal
5183 .iter()
5184 .find(|event| event.kind() == PressureTransitionKind::OwnerAdmissionPending)
5185 .expect("owner rotation must publish admission-pending state");
5186 let owner_admitted = journal
5187 .iter()
5188 .find(|event| event.kind() == PressureTransitionKind::OwnerAdmitted)
5189 .expect("typed admission receipt must commit owner admission");
5190 assert!(admission_pending.ordinal() < owner_admitted.ordinal());
5191 scheduler.mark_prefill_complete(&runnable_id, 1);
5192
5193 let response = InferenceResponse {
5194 request_id: runnable_id.clone(),
5195 text: String::new(),
5196 tokens: Vec::new(),
5197 finish_reason: ferrum_types::FinishReason::Length,
5198 usage: ferrum_types::TokenUsage::new(0, 0),
5199 latency_ms: 0,
5200 created_at: chrono::Utc::now(),
5201 metadata: Default::default(),
5202 api_response: None,
5203 execution_evidence: None,
5204 };
5205 scheduler
5206 .complete(runnable_id.clone(), &response)
5207 .await
5208 .unwrap();
5209
5210 observations.clear();
5211 let callback_order = std::cell::RefCell::new(Vec::new());
5212 let admitted = scheduler
5213 .next_batch_with_dynamic_admission_observed(
5214 BatchHint::simple(2),
5215 AdmissionWakeSnapshot::new(wake2, &availability2),
5216 &mut |request| {
5217 callback_order.borrow_mut().push("admission_probe");
5218 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5219 request_id: request.id.clone(),
5220 })
5221 },
5222 &mut |observation| {
5223 if matches!(
5224 observation,
5225 ExecutorAdmissionQueueObservation::PressureHoldReleased { .. }
5226 ) {
5227 callback_order.borrow_mut().push("pressure_hold_released");
5228 }
5229 observations.push(observation);
5230 },
5231 )
5232 .unwrap()
5233 .expect("rotated owner terminal release must admit the prior owner");
5234 assert_eq!(
5235 callback_order.into_inner(),
5236 vec!["pressure_hold_released", "admission_probe"],
5237 "the trace capture boundary must observe causal hold release before re-admission"
5238 );
5239 assert_eq!(admitted.requests.len(), 1);
5240 assert_eq!(admitted.requests[0].request.id, blocked_id);
5241 batches.push(scheduler_replay_batch(&admitted, &blocked_id, &runnable_id));
5242 assert!(observations.iter().any(|observation| matches!(
5243 observation,
5244 ExecutorAdmissionQueueObservation::PressureHoldReleased {
5245 request_id,
5246 progress_owner_id,
5247 reason: PressureHoldReleaseReason::OwnerTerminal,
5248 ..
5249 } if request_id == &blocked_id && progress_owner_id == &runnable_id
5250 )));
5251 observation_batches.push(observations.clone());
5252 let journal = scheduler.pressure_transition_journal();
5253 assert!(journal
5254 .windows(2)
5255 .all(|pair| pair[0].ordinal() < pair[1].ordinal()));
5256 let final_trace = scheduler.trace_snapshot();
5257 assert_eq!(final_trace.pressure_active_episodes, 0);
5258 assert_eq!(final_trace.pressure_dropped_journal_entries, 0);
5259 traces.push(final_trace);
5260 SchedulerReplayProjection {
5261 batches,
5262 traces,
5263 observations: observation_batches,
5264 journal,
5265 }
5266 }
5267
5268 #[tokio::test]
5269 async fn cross_phase_pressure_yield_rotates_progressed_owner_to_oldest_held_peer() {
5270 let projection = collect_cross_phase_pressure_replay().await;
5271 assert_eq!(projection.batches.len(), 6);
5272 assert!(!projection.journal.is_empty());
5273 }
5274
5275 #[tokio::test]
5276 async fn cross_phase_scheduler_is_deterministic_one_hundred_of_one_hundred() {
5277 const REPLAY_COUNT: usize = 100;
5278 let expected = collect_cross_phase_pressure_replay().await;
5279 for ordinal in 1..REPLAY_COUNT {
5280 assert_eq!(
5281 collect_cross_phase_pressure_replay().await,
5282 expected,
5283 "scheduler execution {ordinal} diverged from execution 0"
5284 );
5285 }
5286 println!(
5287 "FERRUM G04 SCHEDULER DETERMINISM KEEP: deterministic_executions={REPLAY_COUNT}/{REPLAY_COUNT} batch_ticks={} trace_snapshots={} observation_batches={} journal_transitions={}",
5288 expected.batches.len(),
5289 expected.traces.len(),
5290 expected.observations.len(),
5291 expected.journal.len(),
5292 );
5293 }
5294
5295 async fn collect_constrained_decode_membership() -> Vec<Vec<RequestId>> {
5296 const REQUEST_COUNT: usize = 8;
5297 const BATCH_LIMIT: usize = 3;
5298 const ROUND_COUNT: usize = 4;
5299
5300 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
5301 max_running_requests: REQUEST_COUNT,
5302 ..SchedulerConfig::default()
5303 });
5304 let request_ids = (0..REQUEST_COUNT)
5305 .map(|ordinal| RequestId(uuid::Uuid::from_u128(0x100 + ordinal as u128)))
5306 .collect::<Vec<_>>();
5307 for request_id in &request_ids {
5308 let mut request = create_test_request(Priority::Normal);
5309 request.id = request_id.clone();
5310 scheduler.submit(request).await.unwrap();
5311 }
5312
5313 let initial = scheduler
5314 .create_iteration_batch(BatchHint::simple(REQUEST_COUNT))
5315 .expect("all constrained-decode fixtures must enter prefill");
5316 assert_eq!(
5317 initial
5318 .requests
5319 .iter()
5320 .map(|request| request.request.id.clone())
5321 .collect::<Vec<_>>(),
5322 request_ids,
5323 "initial admission must preserve the scheduler-owned queue order"
5324 );
5325 for request_id in &request_ids {
5326 scheduler.mark_prefill_complete(request_id, 1);
5327 }
5328
5329 let mut progress = HashMap::<RequestId, usize>::new();
5330 let mut memberships = Vec::new();
5331 for _ in 0..ROUND_COUNT {
5332 let batch = scheduler
5333 .create_iteration_batch(BatchHint::simple(BATCH_LIMIT))
5334 .expect("eligible decode work must fill every constrained batch");
5335 let membership = batch
5336 .requests
5337 .iter()
5338 .map(|request| request.request.id.clone())
5339 .collect::<Vec<_>>();
5340 assert_eq!(membership.len(), BATCH_LIMIT);
5341 assert_eq!(membership.iter().collect::<HashSet<_>>().len(), BATCH_LIMIT);
5342 for request_id in &membership {
5343 let generated = progress.entry(request_id.clone()).or_default();
5344 *generated += 1;
5345 scheduler.update_decode_progress(request_id, *generated);
5346 }
5347 memberships.push(membership);
5348 }
5349 assert!(
5350 request_ids
5351 .iter()
5352 .all(|request_id| progress.contains_key(request_id)),
5353 "round-robin decode selection must schedule every eligible frontier"
5354 );
5355 memberships
5356 }
5357
5358 #[tokio::test]
5359 async fn constrained_decode_membership_is_stable_and_fair_one_hundred_of_one_hundred() {
5360 const EXECUTION_COUNT: usize = 100;
5361 let expected = collect_constrained_decode_membership().await;
5362 let mut selection_counts = HashMap::<RequestId, usize>::new();
5363 for request_id in expected.iter().flatten() {
5364 *selection_counts.entry(request_id.clone()).or_default() += 1;
5365 }
5366 let min_count = selection_counts.values().copied().min().unwrap();
5367 let max_count = selection_counts.values().copied().max().unwrap();
5368 assert!(
5369 max_count - min_count <= 1,
5370 "constrained round-robin selection must remain balanced"
5371 );
5372 for ordinal in 1..EXECUTION_COUNT {
5373 assert_eq!(
5374 collect_constrained_decode_membership().await,
5375 expected,
5376 "constrained decode execution {ordinal} diverged from execution 0"
5377 );
5378 }
5379 println!(
5380 "FERRUM G04 CONSTRAINED DECODE DETERMINISM KEEP: deterministic_executions={EXECUTION_COUNT}/{EXECUTION_COUNT} requests=8 batch_limit=3 rounds=4"
5381 );
5382 }
5383
5384 #[tokio::test]
5385 async fn decode_cursor_keeps_stable_successor_after_swap_remove() {
5386 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
5387 max_running_requests: 4,
5388 ..SchedulerConfig::default()
5389 });
5390 let request_ids = activate_decode_requests(&scheduler, 4);
5391 let first = scheduler
5392 .create_iteration_batch(BatchHint::simple(3))
5393 .expect("first constrained decode batch must be scheduled");
5394 assert_eq!(
5395 first
5396 .requests
5397 .iter()
5398 .map(|request| request.request.id.clone())
5399 .collect::<Vec<_>>(),
5400 request_ids[..3]
5401 );
5402 assert_eq!(
5403 scheduler.trace_snapshot().decode_selection_cursor,
5404 Some(request_ids[3].clone())
5405 );
5406
5407 assert!(scheduler.cancel(request_ids[1].clone()).await.unwrap());
5408 assert_eq!(
5409 scheduler.trace_snapshot().decode_selection_cursor,
5410 Some(request_ids[3].clone()),
5411 "removing an earlier slot must not move the stable round-robin frontier"
5412 );
5413 let next = scheduler
5414 .create_iteration_batch(BatchHint::simple(1))
5415 .expect("cursor successor must remain runnable");
5416 assert_eq!(next.requests[0].request.id, request_ids[3]);
5417 }
5418
5419 #[tokio::test]
5420 async fn empty_decode_queue_resets_cursor_before_new_cohort() {
5421 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
5422 max_running_requests: 4,
5423 ..SchedulerConfig::default()
5424 });
5425 let old_ids = activate_decode_requests(&scheduler, 4);
5426 let first = scheduler
5427 .create_iteration_batch(BatchHint::simple(1))
5428 .expect("old cohort must establish a nonzero cursor");
5429 assert_eq!(first.requests[0].request.id, old_ids[0]);
5430 assert_eq!(
5431 scheduler.trace_snapshot().decode_selection_cursor,
5432 Some(old_ids[1].clone())
5433 );
5434 for request_id in old_ids {
5435 assert!(scheduler.cancel(request_id).await.unwrap());
5436 }
5437 assert_eq!(scheduler.trace_snapshot().decode_selection_cursor, None);
5438
5439 let new_ids = activate_decode_requests(&scheduler, 3);
5440 let new_first = scheduler
5441 .create_iteration_batch(BatchHint::simple(1))
5442 .expect("new cohort must start from its first admission");
5443 assert_eq!(new_first.requests[0].request.id, new_ids[0]);
5444 }
5445
5446 #[test]
5447 fn readiness_blocked_decode_rejoins_within_one_round_robin_rotation() {
5448 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
5449 max_running_requests: 4,
5450 ..SchedulerConfig::default()
5451 });
5452 let request_ids = activate_decode_requests(&scheduler, 4);
5453 let first = scheduler
5454 .create_iteration_batch(BatchHint::simple(1))
5455 .expect("first decode frontier must run");
5456 assert_eq!(first.requests[0].request.id, request_ids[0]);
5457 scheduler.update_decode_progress(&request_ids[0], 1);
5458
5459 let blocked = scheduler
5460 .defer_for_execution_readiness(std::slice::from_ref(&request_ids[1]))
5461 .unwrap();
5462 for expected_id in [&request_ids[2], &request_ids[3], &request_ids[0]] {
5463 let batch = scheduler
5464 .create_iteration_batch(BatchHint::simple(1))
5465 .expect("an eligible peer must keep decode progressing");
5466 assert_eq!(&batch.requests[0].request.id, expected_id);
5467 scheduler.update_decode_progress(expected_id, 2);
5468 }
5469
5470 assert!(blocked.wake().mark_ready());
5471 let resumed = scheduler
5472 .create_iteration_batch(BatchHint::simple(1))
5473 .expect("ready frontier must rejoin at its preserved cursor");
5474 assert_eq!(resumed.requests[0].request.id, request_ids[1]);
5475 }
5476
5477 #[tokio::test]
5478 async fn disjoint_release_footprints_self_recompute_without_stranded_episode() {
5479 use ferrum_interfaces::vnext::{
5480 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityDomainId,
5481 CapacityWaitCondition, DeferredAction,
5482 };
5483 use std::num::NonZeroU64;
5484
5485 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
5486 let first = create_test_request(Priority::Normal);
5487 let first_id = first.id.clone();
5488 let second = create_test_request(Priority::Normal);
5489 let second_id = second.id.clone();
5490 scheduler.submit(first).await.unwrap();
5491 scheduler.submit(second).await.unwrap();
5492
5493 let domain_two = CapacityAvailabilitySource::Domain(CapacityDomainId::new(2).unwrap());
5494 let domain_four = CapacityAvailabilitySource::Domain(CapacityDomainId::new(4).unwrap());
5495 let plan_budget = CapacityAvailabilitySource::PlanDeviceBudget;
5496 let availability = [
5497 CapacityAvailabilityEpoch::new(domain_two, 176).unwrap(),
5498 CapacityAvailabilityEpoch::new(domain_four, 136).unwrap(),
5499 CapacityAvailabilityEpoch::new(plan_budget, 1).unwrap(),
5500 ];
5501 let wake = AdmissionWakeEpochs::new(NonZeroU64::new(41).unwrap(), 0, 0, 0);
5502 scheduler
5503 .next_batch_with_dynamic_admission(
5504 BatchHint::simple(2),
5505 AdmissionWakeSnapshot::new(wake, &availability),
5506 &mut |request| {
5507 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5508 request_id: request.id.clone(),
5509 })
5510 },
5511 )
5512 .unwrap()
5513 .unwrap();
5514 scheduler.mark_prefill_complete(&first_id, 1);
5515 scheduler.mark_prefill_complete(&second_id, 1);
5516
5517 let wait_for = |source, epoch| {
5518 CapacityWaitCondition::from_observation(
5519 41,
5520 vec![
5521 CapacityAvailabilityEpoch::new(source, epoch).unwrap(),
5522 CapacityAvailabilityEpoch::new(plan_budget, 1).unwrap(),
5523 ],
5524 )
5525 .unwrap()
5526 };
5527 let first_wait = wait_for(domain_four, 136);
5528 let second_wait = wait_for(domain_two, 176);
5529 let release_snapshot = ExecutionCapacityReleaseSnapshot::new([
5530 (first_id.clone(), vec![domain_four]),
5531 (second_id.clone(), vec![domain_two]),
5532 ]);
5533
5534 let first_action = scheduler
5535 .defer_decode_for_execution_capacity(
5536 std::slice::from_ref(&first_id),
5537 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake, first_wait),
5538 &release_snapshot,
5539 )
5540 .unwrap();
5541 let ExecutionCapacityAction::YieldPlanned {
5542 transaction: first_transaction,
5543 } = first_action
5544 else {
5545 panic!("a disjoint runnable footprint must not suppress exact-source self recompute");
5546 };
5547 assert_eq!(first_transaction.kind(), PressureYieldKind::SelfRecompute);
5548 assert_eq!(first_transaction.victim_request_id(), &first_id);
5549 assert_eq!(
5550 scheduler.trace_snapshot().pressure_pending_release_fences,
5551 1
5552 );
5553 scheduler
5554 .arm_execution_capacity_yield(&first_transaction)
5555 .unwrap();
5556 assert!(scheduler
5557 .complete_execution_capacity_yield(&first_transaction, 1, None)
5558 .unwrap()
5559 .victim_requeued());
5560
5561 let second_action = scheduler
5562 .defer_decode_for_execution_capacity(
5563 std::slice::from_ref(&second_id),
5564 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake, second_wait),
5565 &release_snapshot,
5566 )
5567 .unwrap();
5568 let ExecutionCapacityAction::YieldPlanned {
5569 transaction: second_transaction,
5570 } = second_action
5571 else {
5572 panic!("the remaining exact-source owner must self recompute");
5573 };
5574 assert_eq!(second_transaction.kind(), PressureYieldKind::SelfRecompute);
5575 assert_eq!(second_transaction.victim_request_id(), &second_id);
5576 scheduler
5577 .arm_execution_capacity_yield(&second_transaction)
5578 .unwrap();
5579 assert!(scheduler
5580 .complete_execution_capacity_yield(&second_transaction, 1, None)
5581 .unwrap()
5582 .victim_requeued());
5583
5584 let snapshot = scheduler.trace_snapshot();
5585 assert_eq!(snapshot.pressure_active_episodes, 0);
5586 assert_eq!(snapshot.pressure_pending_release_fences, 0);
5587 assert_eq!(snapshot.waiting_queue_len, 2);
5588 assert_eq!(snapshot.execution_capacity_blocked_decode_len, 0);
5589 }
5590
5591 #[tokio::test]
5592 async fn release_fence_preserves_stable_owner_across_retargeted_condition() {
5593 use ferrum_interfaces::vnext::{
5594 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityDomainId,
5595 CapacityWaitCondition, DeferredAction,
5596 };
5597 use std::num::NonZeroU64;
5598
5599 let mut config = SchedulerConfig::default();
5600 config.max_running_requests = 2;
5601 let scheduler = ContinuousBatchScheduler::new(config);
5602 let owner = create_test_request(Priority::Normal);
5603 let owner_id = owner.id.clone();
5604 let victim = create_test_request(Priority::Normal);
5605 let victim_id = victim.id.clone();
5606 scheduler.submit(owner).await.unwrap();
5607 scheduler.submit(victim).await.unwrap();
5608
5609 let wait_for = |domain: u32, epoch: u64| {
5610 CapacityWaitCondition::from_observation(
5611 41,
5612 vec![
5613 CapacityAvailabilityEpoch::new(
5614 CapacityAvailabilitySource::Domain(CapacityDomainId::new(domain).unwrap()),
5615 epoch,
5616 )
5617 .unwrap(),
5618 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::PlanDeviceBudget, 1)
5619 .unwrap(),
5620 ],
5621 )
5622 .unwrap()
5623 };
5624 let original_wait = wait_for(4, 136);
5625 let retargeted_wait = wait_for(2, 178);
5626 let availability = [
5627 CapacityAvailabilityEpoch::new(
5628 CapacityAvailabilitySource::Domain(CapacityDomainId::new(2).unwrap()),
5629 178,
5630 )
5631 .unwrap(),
5632 CapacityAvailabilityEpoch::new(
5633 CapacityAvailabilitySource::Domain(CapacityDomainId::new(4).unwrap()),
5634 136,
5635 )
5636 .unwrap(),
5637 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::PlanDeviceBudget, 1)
5638 .unwrap(),
5639 ];
5640 let wake = AdmissionWakeEpochs::new(NonZeroU64::new(41).unwrap(), 0, 0, 0);
5641 scheduler
5642 .next_batch_with_dynamic_admission(
5643 BatchHint::simple(2),
5644 AdmissionWakeSnapshot::new(wake, &availability),
5645 &mut |request| {
5646 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5647 request_id: request.id.clone(),
5648 })
5649 },
5650 )
5651 .unwrap()
5652 .unwrap();
5653 scheduler.mark_prefill_complete(&owner_id, 1);
5654 scheduler.mark_prefill_complete(&victim_id, 1);
5655
5656 let original_deferral =
5657 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake, original_wait.clone());
5658 assert_eq!(
5659 scheduler
5660 .defer_decode_for_execution_capacity(
5661 std::slice::from_ref(&owner_id),
5662 original_deferral.clone(),
5663 &execution_capacity_release_snapshot([&owner_id, &victim_id], &original_wait,),
5664 )
5665 .unwrap(),
5666 ExecutionCapacityAction::Deferred { count: 1 }
5667 );
5668 let action = scheduler
5669 .defer_decode_for_execution_capacity(
5670 std::slice::from_ref(&victim_id),
5671 original_deferral,
5672 &execution_capacity_release_snapshot([&owner_id, &victim_id], &original_wait),
5673 )
5674 .unwrap();
5675 let ExecutionCapacityAction::YieldPlanned { transaction } = action else {
5676 panic!("all-blocked original domain must plan a typed yield");
5677 };
5678 assert_eq!(transaction.progress_owner_id(), &owner_id);
5679 assert_eq!(transaction.victim_request_id(), &victim_id);
5680 scheduler
5681 .arm_execution_capacity_yield(&transaction)
5682 .unwrap();
5683
5684 let retargeted_deferral = AdmissionDeferral::new(
5685 DeferredAction::WaitForRelease,
5686 wake,
5687 retargeted_wait.clone(),
5688 );
5689 assert_eq!(
5690 scheduler
5691 .defer_decode_for_execution_capacity(
5692 std::slice::from_ref(&owner_id),
5693 retargeted_deferral,
5694 &execution_capacity_release_snapshot(
5695 [&owner_id, &victim_id],
5696 &retargeted_wait,
5697 ),
5698 )
5699 .unwrap(),
5700 ExecutionCapacityAction::Deferred { count: 1 }
5701 );
5702
5703 let completion = scheduler
5704 .complete_execution_capacity_yield(&transaction, 1, Some(0))
5705 .unwrap();
5706 assert!(completion.victim_requeued());
5707 assert!(completion.progress_owner_resumable());
5708 assert!(completion.resumable_transition_ordinal().is_some());
5709 assert!(completion.closed_transition_ordinal().is_none());
5710 assert_eq!(completion.closed_reason(), None);
5711 assert!(matches!(
5712 scheduler
5713 .pressure_coordinator
5714 .lock()
5715 .hold_status(&victim_id),
5716 PressureHoldStatus::Held { .. }
5717 ));
5718 }
5719
5720 #[tokio::test]
5721 async fn pressure_yield_releases_after_phase_independent_owner_terminal() {
5722 use ferrum_interfaces::vnext::{
5723 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
5724 DeferredAction,
5725 };
5726 use std::num::NonZeroU64;
5727
5728 let mut config = SchedulerConfig::default();
5729 config.max_running_requests = 2;
5730 let scheduler = ContinuousBatchScheduler::new(config);
5731 let victim = create_test_request(Priority::Normal);
5732 let victim_id = victim.id.clone();
5733 let owner = create_test_request(Priority::Normal);
5734 let owner_id = owner.id.clone();
5735 scheduler.submit(victim).await.unwrap();
5736 scheduler.submit(owner).await.unwrap();
5737
5738 let source = CapacityAvailabilitySource::ActiveSequenceSlots;
5739 let availability = [CapacityAvailabilityEpoch::new(source, 1).unwrap()];
5740 let wake = AdmissionWakeEpochs::new(NonZeroU64::new(41).unwrap(), 0, 0, 0);
5741 scheduler
5742 .next_batch_with_dynamic_admission(
5743 BatchHint::simple(4),
5744 AdmissionWakeSnapshot::new(wake, &availability),
5745 &mut |request| {
5746 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5747 request_id: request.id.clone(),
5748 })
5749 },
5750 )
5751 .unwrap()
5752 .unwrap();
5753 scheduler.mark_prefill_complete(&victim_id, 1);
5754 scheduler.mark_prefill_complete(&owner_id, 1);
5755 let condition = CapacityWaitCondition::from_observation(41, availability.to_vec()).unwrap();
5756 let deferral =
5757 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake, condition.clone());
5758 assert_eq!(
5759 scheduler
5760 .defer_decode_for_execution_capacity(
5761 std::slice::from_ref(&victim_id),
5762 deferral.clone(),
5763 &execution_capacity_release_snapshot([&victim_id, &owner_id], &condition,),
5764 )
5765 .unwrap(),
5766 ExecutionCapacityAction::Deferred { count: 1 }
5767 );
5768 let action = scheduler
5769 .defer_decode_for_execution_capacity(
5770 std::slice::from_ref(&owner_id),
5771 deferral,
5772 &execution_capacity_release_snapshot([&victim_id, &owner_id], &condition),
5773 )
5774 .unwrap();
5775 let ExecutionCapacityAction::YieldPlanned { transaction } = action else {
5776 panic!("capacity pressure must plan a typed yield");
5777 };
5778 assert_eq!(transaction.progress_owner_id(), &victim_id);
5779 assert_eq!(transaction.victim_request_id(), &owner_id);
5780 scheduler
5781 .arm_execution_capacity_yield(&transaction)
5782 .unwrap();
5783 assert!(scheduler
5784 .complete_execution_capacity_yield(&transaction, 1, None)
5785 .unwrap()
5786 .victim_requeued());
5787 assert!(scheduler.cancel(victim_id.clone()).await.unwrap());
5788
5789 let mut probes = Vec::new();
5790 let mut observations = Vec::new();
5791 let released = scheduler
5792 .next_batch_with_dynamic_admission_observed(
5793 BatchHint::simple(1),
5794 AdmissionWakeSnapshot::new(wake, &availability),
5795 &mut |request| {
5796 probes.push(request.id.clone());
5797 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5798 request_id: request.id.clone(),
5799 })
5800 },
5801 &mut |observation| observations.push(observation),
5802 )
5803 .unwrap()
5804 .expect("owner terminal state must release the yielded frontier");
5805 assert_eq!(probes, vec![owner_id.clone()]);
5806 assert!(released
5807 .requests
5808 .iter()
5809 .any(|request| request.request.id == owner_id));
5810 assert!(observations.iter().any(|observation| matches!(
5811 observation,
5812 ExecutorAdmissionQueueObservation::PressureHoldReleased {
5813 request_id,
5814 progress_owner_id,
5815 reason: PressureHoldReleaseReason::OwnerTerminal,
5816 ..
5817 } if request_id == &owner_id && progress_owner_id == &victim_id
5818 )));
5819 }
5820
5821 #[tokio::test]
5822 async fn lone_active_decode_capacity_deferral_self_recomputes_to_release_its_source() {
5823 use ferrum_interfaces::vnext::{
5824 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
5825 DeferredAction,
5826 };
5827 use std::num::NonZeroU64;
5828
5829 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
5830 let request = create_test_request(Priority::Normal);
5831 let request_id = request.id.clone();
5832 scheduler.submit(request).await.unwrap();
5833
5834 let source = CapacityAvailabilitySource::ActiveSequenceSlots;
5835 let availability0 = [CapacityAvailabilityEpoch::new(source, 3).unwrap()];
5836 let wake0 = AdmissionWakeEpochs::new(NonZeroU64::new(37).unwrap(), 0, 0, 0);
5837 scheduler
5838 .next_batch_with_dynamic_admission(
5839 BatchHint::simple(1),
5840 AdmissionWakeSnapshot::new(wake0, &availability0),
5841 &mut |request| {
5842 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5843 request_id: request.id.clone(),
5844 })
5845 },
5846 )
5847 .unwrap()
5848 .expect("request must enter prefill before decode");
5849 scheduler.mark_prefill_complete(&request_id, 1);
5850
5851 let condition =
5852 CapacityWaitCondition::from_observation(37, availability0.to_vec()).unwrap();
5853 let deferral =
5854 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake0, condition.clone());
5855 let action = scheduler
5856 .defer_decode_for_execution_capacity(
5857 std::slice::from_ref(&request_id),
5858 deferral,
5859 &execution_capacity_release_snapshot([&request_id], &condition),
5860 )
5861 .unwrap();
5862 let ExecutionCapacityAction::YieldPlanned { transaction } = action else {
5863 panic!("a lone releasable decode must plan a typed self recompute");
5864 };
5865 assert_eq!(transaction.kind(), PressureYieldKind::SelfRecompute);
5866 assert_eq!(transaction.victim_request_id(), &request_id);
5867 assert_eq!(transaction.progress_owner_id(), &request_id);
5868
5869 scheduler
5870 .arm_execution_capacity_yield(&transaction)
5871 .unwrap();
5872 let completion = scheduler
5873 .complete_execution_capacity_yield(&transaction, 1, None)
5874 .unwrap();
5875 assert!(completion.victim_requeued());
5876 assert!(!completion.progress_owner_resumable());
5877 assert_eq!(
5878 completion.disposition(),
5879 ExecutionCapacityYieldDisposition::SelfRecomputeQueued
5880 );
5881 assert!(completion.closed_transition_ordinal().is_some());
5882 assert_eq!(completion.closed_reason(), None);
5883
5884 let snapshot = scheduler.trace_snapshot();
5885 assert_eq!(snapshot.decode_queue_len, 0);
5886 assert_eq!(snapshot.waiting_queue_len, 1);
5887 assert_eq!(snapshot.pressure_active_episodes, 0);
5888 assert_eq!(snapshot.pressure_pending_release_fences, 0);
5889 }
5890
5891 #[tokio::test]
5892 async fn active_prefill_capacity_deferral_retries_the_exact_chunk_after_source_change() {
5893 use ferrum_interfaces::vnext::{
5894 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
5895 DeferredAction,
5896 };
5897 use std::num::NonZeroU64;
5898
5899 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
5900 max_running_requests: 1,
5901 prefill_step_chunk: Some(3),
5902 ..SchedulerConfig::default()
5903 });
5904 let request = create_test_request_with_prompt_tokens(Priority::Normal, 8);
5905 let request_id = request.id.clone();
5906 scheduler.submit(request).await.unwrap();
5907
5908 let source = CapacityAvailabilitySource::ActiveSequenceSlots;
5909 let availability0 = [CapacityAvailabilityEpoch::new(source, 1).unwrap()];
5910 let wake0 = AdmissionWakeEpochs::new(NonZeroU64::new(31).unwrap(), 0, 0, 0);
5911 let first = scheduler
5912 .next_batch_with_dynamic_admission(
5913 BatchHint::simple(1),
5914 AdmissionWakeSnapshot::new(wake0, &availability0),
5915 &mut |request| {
5916 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
5917 request_id: request.id.clone(),
5918 })
5919 },
5920 )
5921 .unwrap()
5922 .unwrap();
5923 assert_eq!(first.requests[0].tokens_processed, 0);
5924 assert_eq!(first.requests[0].tokens_to_process, Some(3));
5925 assert!(!scheduler.mark_prefill_chunk_processed(&request_id, 8, 3));
5926
5927 let condition =
5928 CapacityWaitCondition::from_observation(31, availability0.to_vec()).unwrap();
5929 let deferral =
5930 AdmissionDeferral::new(DeferredAction::WaitForRelease, wake0, condition.clone());
5931 assert_eq!(
5932 scheduler
5933 .defer_prefill_for_execution_capacity(
5934 &request_id,
5935 deferral,
5936 &execution_capacity_release_snapshot([&request_id], &condition),
5937 )
5938 .unwrap(),
5939 ExecutionCapacityAction::Deferred { count: 1 }
5940 );
5941
5942 let mut observations = Vec::new();
5943 assert!(scheduler
5944 .next_batch_with_dynamic_admission_observed(
5945 BatchHint::simple(1),
5946 AdmissionWakeSnapshot::new(wake0, &availability0),
5947 &mut |_| panic!("active prefill must not re-enter admission"),
5948 &mut |observation| observations.push(observation),
5949 )
5950 .unwrap()
5951 .is_none());
5952 assert!(observations.iter().any(|observation| matches!(
5953 observation,
5954 ExecutorAdmissionQueueObservation::PrefillSkippedUnchanged {
5955 request_id: observed_id,
5956 current_wait_sources,
5957 ..
5958 } if observed_id == &request_id && current_wait_sources == &availability0
5959 )));
5960 assert_eq!(
5961 scheduler.passive_capacity_wait_condition().unwrap(),
5962 Some(condition)
5963 );
5964 assert_eq!(
5965 scheduler
5966 .trace_snapshot()
5967 .execution_capacity_blocked_prefill_len,
5968 1
5969 );
5970
5971 let availability1 = [CapacityAvailabilityEpoch::new(source, 2).unwrap()];
5972 let wake1 = AdmissionWakeEpochs::new(NonZeroU64::new(31).unwrap(), 1, 0, 0);
5973 observations.clear();
5974 let resumed = scheduler
5975 .next_batch_with_dynamic_admission_observed(
5976 BatchHint::simple(1),
5977 AdmissionWakeSnapshot::new(wake1, &availability1),
5978 &mut |_| panic!("active prefill resume must not re-enter admission"),
5979 &mut |observation| observations.push(observation),
5980 )
5981 .unwrap()
5982 .expect("source movement must resume the deferred prefill");
5983 assert_eq!(resumed.requests[0].request.id, request_id);
5984 assert_eq!(resumed.requests[0].tokens_processed, 3);
5985 assert_eq!(resumed.requests[0].tokens_to_process, Some(3));
5986 assert!(observations.iter().any(|observation| matches!(
5987 observation,
5988 ExecutorAdmissionQueueObservation::PrefillResumed {
5989 exact_source_changed: true,
5990 policy_epoch_changed: false,
5991 ..
5992 }
5993 )));
5994 assert_eq!(
5995 scheduler
5996 .trace_snapshot()
5997 .execution_capacity_blocked_prefill_len,
5998 0
5999 );
6000 }
6001
6002 #[tokio::test]
6003 async fn partial_prefill_completion_limits_the_next_scheduled_frontier() {
6004 use ferrum_interfaces::vnext::{CapacityAvailabilityEpoch, CapacityAvailabilitySource};
6005 use std::num::NonZeroU64;
6006
6007 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6008 max_running_requests: 1,
6009 prefill_step_chunk: Some(4),
6010 ..SchedulerConfig::default()
6011 });
6012 let request = create_test_request_with_prompt_tokens(Priority::Normal, 8);
6013 let request_id = request.id.clone();
6014 scheduler.submit(request).await.unwrap();
6015 let availability =
6016 [
6017 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 1)
6018 .unwrap(),
6019 ];
6020 let wake = AdmissionWakeSnapshot::new(
6021 AdmissionWakeEpochs::new(NonZeroU64::new(31).unwrap(), 0, 0, 0),
6022 &availability,
6023 );
6024
6025 let first = scheduler
6026 .next_batch_with_dynamic_admission(BatchHint::simple(1), wake, &mut |request| {
6027 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
6028 request_id: request.id.clone(),
6029 })
6030 })
6031 .unwrap()
6032 .unwrap();
6033 assert_eq!(first.requests[0].tokens_to_process, Some(4));
6034 assert!(!scheduler
6035 .mark_prefill_chunk_processed_with_capacity_feedback(&request_id, 8, 4, 2)
6036 .unwrap());
6037
6038 let second = scheduler
6039 .next_batch_with_dynamic_admission(BatchHint::simple(1), wake, &mut |_| {
6040 panic!("active prefill must not re-enter admission")
6041 })
6042 .unwrap()
6043 .unwrap();
6044 assert_eq!(second.requests[0].tokens_processed, 2);
6045 assert_eq!(second.requests[0].tokens_to_process, Some(2));
6046 assert!(!scheduler
6047 .mark_prefill_chunk_processed_with_capacity_feedback(&request_id, 8, 2, 2)
6048 .unwrap());
6049
6050 let third = scheduler
6051 .next_batch_with_dynamic_admission(BatchHint::simple(1), wake, &mut |_| {
6052 panic!("active prefill must not re-enter admission")
6053 })
6054 .unwrap()
6055 .unwrap();
6056 assert_eq!(third.requests[0].tokens_processed, 4);
6057 assert_eq!(third.requests[0].tokens_to_process, Some(4));
6058 }
6059
6060 #[tokio::test]
6061 async fn typed_wait_for_release_does_not_make_prefill_first_starve_decode() {
6062 use ferrum_interfaces::vnext::{
6063 CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
6064 DeferredAction,
6065 };
6066 use std::num::NonZeroU64;
6067
6068 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6069 max_running_requests: 2,
6070 prefill_first_until_active: Some(2),
6071 ..SchedulerConfig::default()
6072 });
6073 let blocked = create_test_request(Priority::Normal);
6074 let blocked_id = blocked.id.clone();
6075 let runnable = create_test_request(Priority::Normal);
6076 let runnable_id = runnable.id.clone();
6077 scheduler.submit(blocked).await.unwrap();
6078 scheduler.submit(runnable).await.unwrap();
6079
6080 let wake = AdmissionWakeEpochs::new(NonZeroU64::new(23).unwrap(), 7, 11, 0);
6081 let availability =
6082 [
6083 CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 19)
6084 .unwrap(),
6085 ];
6086 let condition = CapacityWaitCondition::from_observation(23, availability.to_vec()).unwrap();
6087 let first = scheduler
6088 .next_batch_with_dynamic_admission(
6089 BatchHint::simple(2),
6090 AdmissionWakeSnapshot::new(wake, &availability),
6091 &mut |request| {
6092 if request.id == blocked_id {
6093 AdmissionProbeOutcome::Deferred(crate::vnext::AdmissionDeferral::new(
6094 DeferredAction::WaitForRelease,
6095 wake,
6096 condition.clone(),
6097 ))
6098 } else {
6099 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
6100 request_id: request.id.clone(),
6101 })
6102 }
6103 },
6104 )
6105 .unwrap()
6106 .unwrap();
6107 assert_eq!(first.requests.len(), 1);
6108 assert_eq!(first.requests[0].request.id, runnable_id);
6109 scheduler.mark_prefill_complete(&runnable_id, 1);
6110
6111 let mut probes = 0;
6112 let mut observations = Vec::new();
6113 let unchanged = scheduler
6114 .next_batch_with_dynamic_admission_observed(
6115 BatchHint::simple(2),
6116 AdmissionWakeSnapshot::new(wake, &availability),
6117 &mut |request| {
6118 probes += 1;
6119 AdmissionProbeOutcome::Admitted(ExecutorPrefillAdmissionReceipt {
6120 request_id: request.id.clone(),
6121 })
6122 },
6123 &mut |observation| observations.push(observation),
6124 )
6125 .unwrap()
6126 .expect("unchanged capacity wait must not suppress runnable decode work");
6127
6128 assert_eq!(probes, 0);
6129 assert_eq!(unchanged.requests.len(), 1);
6130 assert_eq!(unchanged.requests[0].request.id, runnable_id);
6131 assert_eq!(unchanged.requests[0].tokens_to_process, Some(1));
6132 assert!(matches!(
6133 observations.as_slice(),
6134 [ExecutorAdmissionQueueObservation::SkippedUnchanged {
6135 request_id,
6136 current,
6137 ..
6138 }] if request_id == &blocked_id && *current == wake
6139 ));
6140 }
6141
6142 #[tokio::test]
6143 async fn defer_prefill_to_waiting_frees_active_slot_without_cancelling() {
6144 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
6145 let request = create_test_request(Priority::Normal);
6146 let request_id = request.id.clone();
6147 scheduler.submit(request).await.unwrap();
6148
6149 let batch = scheduler.next_batch(BatchHint::simple(4)).await.unwrap();
6150 assert_eq!(batch.size(), 1);
6151 let active = scheduler.trace_snapshot();
6152 assert_eq!(active.waiting_queue_len, 0);
6153 assert_eq!(active.prefill_queue_len, 1);
6154 assert_eq!(active.active_len, 1);
6155
6156 assert!(scheduler.defer_prefill_to_waiting(&request_id));
6157 let deferred = scheduler.trace_snapshot();
6158 assert_eq!(deferred.waiting_queue_len, 1);
6159 assert_eq!(deferred.prefill_queue_len, 0);
6160 assert_eq!(deferred.active_len, 0);
6161 assert_eq!(deferred.cancelled_total, 0);
6162 assert_eq!(
6163 scheduler.trace_phase(&request_id),
6164 Some(RequestPhase::Waiting)
6165 );
6166 assert_eq!(
6167 scheduler.request_state(&request_id),
6168 Some(RequestState::Waiting)
6169 );
6170 }
6171
6172 #[test]
6173 fn defer_prefill_to_waiting_resets_chunk_progress_after_capacity_loss() {
6174 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6175 max_running_requests: 1,
6176 prompt_token_estimate: true,
6177 ..SchedulerConfig::default()
6178 });
6179 let hint = BatchHint {
6180 max_batch_size: 1,
6181 max_tokens: 1024,
6182 target_latency_ms: None,
6183 available_memory: None,
6184 resource_constraints: Default::default(),
6185 };
6186
6187 let request = create_test_request_with_prompt_tokens(Priority::Normal, 128);
6188 let request_id = request.id.clone();
6189 enqueue_waiting(&scheduler, request);
6190
6191 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6192 assert_eq!(first_batch.requests[0].request.id, request_id);
6193 assert!(!scheduler.mark_prefill_chunk_processed(&request_id, 128, 64));
6194
6195 assert!(scheduler.defer_prefill_to_waiting(&request_id));
6196 let waiting = scheduler.waiting_queue.read();
6197 let deferred = waiting
6198 .iter()
6199 .find(|req| req.inner.request.id == request_id)
6200 .expect("request should be back in waiting queue");
6201 assert_eq!(deferred.prefill_tokens, 0);
6202 assert_eq!(deferred.prefill_chunk_offset, 0);
6203 assert!(!deferred.chunked_prefill);
6204 drop(waiting);
6205
6206 let retry_batch = scheduler.create_iteration_batch(hint).unwrap();
6207 assert_eq!(retry_batch.requests[0].request.id, request_id);
6208 assert_eq!(
6209 retry_batch.requests[0].tokens_to_process,
6210 Some(128),
6211 "released physical prefill state must be rebuilt from the start"
6212 );
6213 }
6214
6215 #[test]
6216 fn capacity_defer_halves_next_waiting_admission_width() {
6217 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6218 max_running_requests: 4,
6219 prompt_token_estimate: true,
6220 ..SchedulerConfig::default()
6221 });
6222 let hint = BatchHint {
6223 max_batch_size: 4,
6224 max_tokens: 1024,
6225 target_latency_ms: None,
6226 available_memory: None,
6227 resource_constraints: Default::default(),
6228 };
6229
6230 for _ in 0..4 {
6231 enqueue_waiting(
6232 &scheduler,
6233 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6234 );
6235 }
6236
6237 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6238 assert_eq!(first_batch.requests.len(), 4);
6239 let first_ids: Vec<_> = first_batch
6240 .requests
6241 .iter()
6242 .map(|request| request.request.id.clone())
6243 .collect();
6244 for request_id in &first_ids {
6245 assert!(scheduler.defer_prefill_to_waiting(request_id));
6246 }
6247
6248 let deferred = scheduler.trace_snapshot();
6249 assert_eq!(deferred.waiting_queue_len, 4);
6250 assert_eq!(deferred.active_len, 0);
6251 assert_eq!(deferred.capacity_deferred_total, 4);
6252 assert_eq!(deferred.capacity_backpressure_admit_limit, Some(2));
6253
6254 let second_batch = scheduler.create_iteration_batch(hint).unwrap();
6255 assert_eq!(
6256 second_batch.requests.len(),
6257 2,
6258 "capacity-deferred waiting requests should not be immediately re-admitted at the failed width"
6259 );
6260 let after = scheduler.trace_snapshot();
6261 assert_eq!(after.waiting_queue_len, 2);
6262 assert_eq!(after.prefill_queue_len, 2);
6263 assert_eq!(after.active_len, 2);
6264 assert_eq!(after.admitted_total, 6);
6265 assert_eq!(after.capacity_backpressure_admit_limit, Some(2));
6266 }
6267
6268 #[test]
6269 fn capacity_deferred_prefill_retries_once_without_release_then_parks() {
6270 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6271 max_running_requests: 1,
6272 prompt_token_estimate: true,
6273 ..SchedulerConfig::default()
6274 });
6275 let hint = BatchHint {
6276 max_batch_size: 1,
6277 max_tokens: 1024,
6278 target_latency_ms: None,
6279 available_memory: None,
6280 resource_constraints: Default::default(),
6281 };
6282 let request = create_test_request_with_prompt_tokens(Priority::Normal, 128);
6283 let request_id = request.id.clone();
6284 enqueue_waiting(&scheduler, request);
6285
6286 let first = scheduler.create_iteration_batch(hint.clone()).unwrap();
6287 assert_eq!(first.requests[0].request.id, request_id);
6288 assert!(scheduler.defer_prefill_to_waiting(&request_id));
6289
6290 let reduced_retry = scheduler.create_iteration_batch(hint.clone()).unwrap();
6291 assert_eq!(reduced_retry.requests[0].request.id, request_id);
6292 assert!(scheduler.defer_prefill_to_waiting(&request_id));
6293
6294 assert!(scheduler.create_iteration_batch(hint.clone()).is_none());
6295 let parked = scheduler.trace_snapshot();
6296 assert_eq!(parked.waiting_queue_len, 1);
6297 assert_eq!(parked.active_len, 0);
6298 assert_eq!(parked.capacity_blocked_waiting_len, 1);
6299 assert_eq!(parked.admitted_total, 2);
6300 assert_eq!(parked.capacity_deferred_total, 2);
6301
6302 scheduler.record_capacity_release_progress();
6303 let after_release = scheduler.create_iteration_batch(hint).unwrap();
6304 assert_eq!(after_release.requests[0].request.id, request_id);
6305 }
6306
6307 #[test]
6308 fn decode_capacity_defer_requeues_for_recompute_without_cancelling() {
6309 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6310 max_running_requests: 4,
6311 prompt_token_estimate: true,
6312 ..SchedulerConfig::default()
6313 });
6314 let hint = BatchHint {
6315 max_batch_size: 4,
6316 max_tokens: 1024,
6317 target_latency_ms: None,
6318 available_memory: None,
6319 resource_constraints: Default::default(),
6320 };
6321
6322 for _ in 0..4 {
6323 enqueue_waiting(
6324 &scheduler,
6325 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6326 );
6327 }
6328
6329 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6330 assert_eq!(first_batch.requests.len(), 4);
6331 let first_ids: Vec<_> = first_batch
6332 .requests
6333 .iter()
6334 .map(|request| request.request.id.clone())
6335 .collect();
6336 assert_eq!(first_ids.len(), 4);
6337 for request_id in &first_ids {
6338 scheduler.mark_prefill_complete(request_id, 128);
6339 }
6340 assert_eq!(scheduler.trace_snapshot().decode_queue_len, 4);
6341
6342 for request_id in &first_ids {
6343 assert!(scheduler.defer_decode_to_waiting_for_capacity(request_id, 4));
6344 }
6345
6346 let deferred = scheduler.trace_snapshot();
6347 assert_eq!(deferred.waiting_queue_len, 4);
6348 assert_eq!(deferred.decode_queue_len, 0);
6349 assert_eq!(deferred.active_len, 0);
6350 assert_eq!(deferred.cancelled_total, 0);
6351 assert_eq!(deferred.capacity_deferred_total, 4);
6352 assert_eq!(deferred.capacity_backpressure_admit_limit, Some(2));
6353 for request_id in &first_ids {
6354 assert_eq!(
6355 scheduler.trace_phase(request_id),
6356 Some(RequestPhase::Waiting)
6357 );
6358 assert_eq!(
6359 scheduler.request_state(request_id),
6360 Some(RequestState::Waiting)
6361 );
6362 }
6363
6364 let second_batch = scheduler.create_iteration_batch(hint).unwrap();
6365 assert_eq!(
6366 second_batch.requests.len(),
6367 2,
6368 "capacity-deferred decodes should recompute at a lower admission width"
6369 );
6370 let after = scheduler.trace_snapshot();
6371 assert_eq!(after.waiting_queue_len, 2);
6372 assert_eq!(after.prefill_queue_len, 2);
6373 assert_eq!(after.active_len, 2);
6374 assert_eq!(after.capacity_backpressure_admit_limit, Some(2));
6375 }
6376
6377 #[test]
6378 fn capacity_deferred_decode_recomputes_as_bounded_mixed_prefill_under_decode_pressure() {
6379 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6380 max_running_requests: 4,
6381 prompt_token_estimate: true,
6382 ..SchedulerConfig::default()
6383 });
6384 let hint = BatchHint {
6385 max_batch_size: 4,
6386 max_tokens: 1024,
6387 target_latency_ms: None,
6388 available_memory: None,
6389 resource_constraints: Default::default(),
6390 };
6391
6392 for _ in 0..4 {
6393 enqueue_waiting(
6394 &scheduler,
6395 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6396 );
6397 }
6398
6399 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6400 let first_ids: Vec<_> = first_batch
6401 .requests
6402 .iter()
6403 .map(|request| request.request.id.clone())
6404 .collect();
6405 for request_id in &first_ids {
6406 scheduler.mark_prefill_complete(request_id, 128);
6407 }
6408
6409 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
6410 let deferred = scheduler.trace_snapshot();
6411 assert_eq!(deferred.waiting_queue_len, 1);
6412 assert_eq!(deferred.decode_queue_len, 3);
6413 assert_eq!(deferred.active_len, 3);
6414 assert_eq!(deferred.capacity_blocked_waiting_len, 1);
6415 assert_eq!(deferred.capacity_backpressure_admit_limit, Some(2));
6416
6417 let decode_only = scheduler.create_iteration_batch(hint.clone()).unwrap();
6418 let scheduled_ids: HashSet<RequestId> = decode_only
6419 .requests
6420 .iter()
6421 .map(|request| request.request.id.clone())
6422 .collect();
6423 let scheduled_decodes = decode_only
6424 .requests
6425 .iter()
6426 .filter(|request| request.tokens_to_process == Some(1))
6427 .count();
6428 assert_eq!(
6429 scheduled_decodes, 3,
6430 "decode KV pressure should not globally cap decode-ready survivors while recompute runs"
6431 );
6432 assert!(
6433 !scheduler
6434 .trace_snapshot()
6435 .decode_execution_pressure_enforced
6436 );
6437 assert_eq!(decode_only.requests.len(), 4);
6438 assert!(
6439 scheduled_ids.contains(&first_ids[0]),
6440 "capacity-deferred recompute should use bounded mixed prefill budget under decode pressure"
6441 );
6442 let prefill_tokens = decode_only
6443 .requests
6444 .iter()
6445 .find(|request| request.request.id == first_ids[0])
6446 .and_then(|request| request.tokens_to_process);
6447 assert_eq!(
6448 prefill_tokens,
6449 Some(64),
6450 "the recompute prefill should still be capped by the mixed-prefill token budget"
6451 );
6452 assert_eq!(scheduler.trace_snapshot().capacity_blocked_waiting_len, 0);
6453 }
6454
6455 #[test]
6456 fn execution_capacity_pressure_caps_decode_survivors_while_recompute_runs() {
6457 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6458 max_running_requests: 4,
6459 prompt_token_estimate: true,
6460 ..SchedulerConfig::default()
6461 });
6462 let hint = BatchHint {
6463 max_batch_size: 4,
6464 max_tokens: 1024,
6465 target_latency_ms: None,
6466 available_memory: None,
6467 resource_constraints: Default::default(),
6468 };
6469
6470 for _ in 0..4 {
6471 enqueue_waiting(
6472 &scheduler,
6473 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6474 );
6475 }
6476
6477 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6478 let first_ids: Vec<_> = first_batch
6479 .requests
6480 .iter()
6481 .map(|request| request.request.id.clone())
6482 .collect();
6483 for request_id in &first_ids {
6484 scheduler.mark_prefill_complete(request_id, 128);
6485 }
6486
6487 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
6488 scheduler.record_decode_execution_capacity_pressure(3);
6489
6490 let bounded = scheduler.create_iteration_batch(hint).unwrap();
6491 let scheduled_decodes = bounded
6492 .requests
6493 .iter()
6494 .filter(|request| request.tokens_to_process == Some(1))
6495 .count();
6496 assert_eq!(
6497 scheduled_decodes, 2,
6498 "plan-runtime execution pressure must remain effective while a recompute backlog exists"
6499 );
6500 assert!(
6501 scheduler
6502 .trace_snapshot()
6503 .decode_execution_pressure_enforced
6504 );
6505 assert!(
6506 bounded.requests.iter().any(|request| {
6507 request.request.id == first_ids[0] && request.tokens_to_process == Some(64)
6508 }),
6509 "execution pressure must not block the bounded mixed recompute"
6510 );
6511
6512 let before_release = scheduler.trace_snapshot();
6513 assert_eq!(
6514 before_release.decode_capacity_backpressure_admit_limit,
6515 Some(2)
6516 );
6517 assert!(before_release.decode_execution_pressure_enforced);
6518
6519 assert!(!scheduler.record_decode_execution_capacity_success(2));
6520 assert_eq!(
6521 scheduler
6522 .trace_snapshot()
6523 .decode_capacity_backpressure_admit_limit,
6524 Some(2)
6525 );
6526
6527 scheduler.record_external_capacity_release();
6528 let still_bounded = scheduler.trace_snapshot();
6529 assert_eq!(
6530 still_bounded.decode_capacity_backpressure_admit_limit,
6531 Some(2)
6532 );
6533 assert!(still_bounded.decode_execution_pressure_enforced);
6534
6535 assert!(!scheduler.record_decode_execution_capacity_success(1));
6536 assert_eq!(
6537 scheduler
6538 .trace_snapshot()
6539 .decode_capacity_backpressure_admit_limit,
6540 Some(2)
6541 );
6542 assert!(scheduler.record_decode_execution_capacity_success(2));
6543 assert_eq!(
6544 scheduler
6545 .trace_snapshot()
6546 .decode_capacity_backpressure_admit_limit,
6547 Some(3)
6548 );
6549 assert!(!scheduler.record_decode_execution_capacity_success(3));
6550 assert_eq!(
6551 scheduler
6552 .trace_snapshot()
6553 .decode_capacity_backpressure_admit_limit,
6554 Some(3)
6555 );
6556 scheduler.record_external_capacity_release();
6557 assert!(scheduler.record_decode_execution_capacity_success(3));
6558 let recovered = scheduler.trace_snapshot();
6559 assert_eq!(recovered.decode_capacity_backpressure_admit_limit, None);
6560 assert!(!recovered.decode_execution_pressure_enforced);
6561 }
6562
6563 #[test]
6564 fn execution_capacity_pressure_survives_prefill_progress_and_no_backlog_selection() {
6565 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6566 max_running_requests: 16,
6567 prompt_token_estimate: true,
6568 ..SchedulerConfig::default()
6569 });
6570 let hint = BatchHint {
6571 max_batch_size: 16,
6572 max_tokens: 1024,
6573 target_latency_ms: None,
6574 available_memory: None,
6575 resource_constraints: Default::default(),
6576 };
6577 for _ in 0..8 {
6578 enqueue_waiting(
6579 &scheduler,
6580 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6581 );
6582 }
6583 let prefill = scheduler.create_iteration_batch(hint.clone()).unwrap();
6584 let request_ids = prefill
6585 .requests
6586 .iter()
6587 .map(|request| request.request.id.clone())
6588 .collect::<Vec<_>>();
6589 assert_eq!(request_ids.len(), 8);
6590
6591 scheduler.record_decode_execution_capacity_pressure(11);
6592 for request_id in &request_ids {
6593 scheduler.mark_prefill_complete(request_id, 128);
6594 }
6595 let after_prefill = scheduler.trace_snapshot();
6596 assert_eq!(
6597 after_prefill.decode_capacity_backpressure_admit_limit,
6598 Some(6)
6599 );
6600 assert!(after_prefill.decode_execution_pressure_enforced);
6601
6602 let decode = scheduler.create_iteration_batch(hint).unwrap();
6603 assert_eq!(decode.requests.len(), 6);
6604 let after_no_backlog_selection = scheduler.trace_snapshot();
6605 assert_eq!(
6606 after_no_backlog_selection.decode_capacity_backpressure_admit_limit,
6607 Some(6)
6608 );
6609 assert!(after_no_backlog_selection.decode_execution_pressure_enforced);
6610 }
6611
6612 #[test]
6613 fn capacity_deferred_decode_recompute_spends_available_mixed_slots() {
6614 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6615 max_running_requests: 8,
6616 prompt_token_estimate: true,
6617 ..SchedulerConfig::default()
6618 });
6619 let hint = BatchHint {
6620 max_batch_size: 8,
6621 max_tokens: 1024,
6622 target_latency_ms: None,
6623 available_memory: None,
6624 resource_constraints: Default::default(),
6625 };
6626
6627 for _ in 0..8 {
6628 enqueue_waiting(
6629 &scheduler,
6630 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6631 );
6632 }
6633
6634 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6635 let first_ids: Vec<_> = first_batch
6636 .requests
6637 .iter()
6638 .map(|request| request.request.id.clone())
6639 .collect();
6640 for request_id in &first_ids {
6641 scheduler.mark_prefill_complete(request_id, 128);
6642 }
6643
6644 for request_id in first_ids.iter().take(4) {
6645 assert!(scheduler.defer_decode_to_waiting_for_capacity(request_id, 8));
6646 }
6647 let deferred = scheduler.trace_snapshot();
6648 assert_eq!(deferred.waiting_queue_len, 4);
6649 assert_eq!(deferred.decode_queue_len, 4);
6650 assert_eq!(deferred.active_len, 4);
6651 assert_eq!(deferred.capacity_blocked_waiting_len, 4);
6652
6653 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
6654 let scheduled_deferred = mixed_batch
6655 .requests
6656 .iter()
6657 .filter(|request| first_ids[..4].contains(&request.request.id))
6658 .count();
6659 let scheduled_decodes = mixed_batch
6660 .requests
6661 .iter()
6662 .filter(|request| {
6663 first_ids[4..].contains(&request.request.id) && request.tokens_to_process == Some(1)
6664 })
6665 .count();
6666 let prefill_tokens: Vec<_> = mixed_batch
6667 .requests
6668 .iter()
6669 .filter(|request| first_ids[..4].contains(&request.request.id))
6670 .map(|request| request.tokens_to_process)
6671 .collect();
6672
6673 assert_eq!(scheduled_decodes, 4);
6674 assert_eq!(
6675 scheduled_deferred, 4,
6676 "bounded mixed recompute should spend available mixed-prefill slots"
6677 );
6678 assert_eq!(prefill_tokens, vec![Some(64), Some(64), Some(64), Some(64)]);
6679 assert_eq!(scheduler.trace_snapshot().capacity_blocked_waiting_len, 0);
6680 }
6681
6682 #[test]
6683 fn capacity_deferred_recompute_waits_after_no_progress_attempt() {
6684 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6685 max_running_requests: 4,
6686 prompt_token_estimate: true,
6687 ..SchedulerConfig::default()
6688 });
6689 let hint = BatchHint {
6690 max_batch_size: 4,
6691 max_tokens: 1024,
6692 target_latency_ms: None,
6693 available_memory: None,
6694 resource_constraints: Default::default(),
6695 };
6696
6697 for _ in 0..4 {
6698 enqueue_waiting(
6699 &scheduler,
6700 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6701 );
6702 }
6703
6704 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6705 let first_ids: Vec<_> = first_batch
6706 .requests
6707 .iter()
6708 .map(|request| request.request.id.clone())
6709 .collect();
6710 for request_id in &first_ids {
6711 scheduler.mark_prefill_complete(request_id, 128);
6712 }
6713
6714 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
6715
6716 let first_mixed = scheduler.create_iteration_batch(hint.clone()).unwrap();
6717 let first_scheduled_ids: HashSet<RequestId> = first_mixed
6718 .requests
6719 .iter()
6720 .map(|request| request.request.id.clone())
6721 .collect();
6722 assert!(
6723 first_scheduled_ids.contains(&first_ids[0]),
6724 "the first mixed iteration may spend its bounded recompute slot"
6725 );
6726 assert_eq!(scheduler.trace_snapshot().capacity_blocked_waiting_len, 0);
6727 assert_eq!(scheduler.trace_snapshot().prefill_queue_len, 1);
6728
6729 let no_progress_retry = scheduler.create_iteration_batch(hint.clone()).unwrap();
6730 let retry_ids: HashSet<RequestId> = no_progress_retry
6731 .requests
6732 .iter()
6733 .map(|request| request.request.id.clone())
6734 .collect();
6735 assert!(
6736 !retry_ids.contains(&first_ids[0]),
6737 "a release-blocked recompute must not be retried in the same release epoch without progress"
6738 );
6739 assert_eq!(
6740 no_progress_retry.requests.len(),
6741 3,
6742 "decode-ready survivors should continue at full width while the failed recompute waits"
6743 );
6744 assert_eq!(scheduler.trace_snapshot().prefill_queue_len, 1);
6745
6746 assert!(!scheduler.mark_prefill_chunk_processed(&first_ids[0], 128, 64));
6747 let after_progress = scheduler.create_iteration_batch(hint).unwrap();
6748 let progressed_tokens = after_progress
6749 .requests
6750 .iter()
6751 .find(|request| request.request.id == first_ids[0])
6752 .and_then(|request| request.tokens_to_process);
6753 assert_eq!(
6754 progressed_tokens,
6755 Some(64),
6756 "recorded prefill progress should make the next recompute chunk eligible again"
6757 );
6758 }
6759
6760 #[test]
6761 fn capacity_deferred_recompute_skips_marked_requests_without_blocking_later_candidates() {
6762 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6763 max_running_requests: 10,
6764 prompt_token_estimate: true,
6765 ..SchedulerConfig::default()
6766 });
6767 let init_hint = BatchHint {
6768 max_batch_size: 10,
6769 max_tokens: 2048,
6770 target_latency_ms: None,
6771 available_memory: None,
6772 resource_constraints: Default::default(),
6773 };
6774 let mixed_hint = BatchHint {
6775 max_batch_size: 8,
6776 max_tokens: 1024,
6777 target_latency_ms: None,
6778 available_memory: None,
6779 resource_constraints: Default::default(),
6780 };
6781
6782 for _ in 0..10 {
6783 enqueue_waiting(
6784 &scheduler,
6785 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6786 );
6787 }
6788
6789 let first_batch = scheduler.create_iteration_batch(init_hint.clone()).unwrap();
6790 let first_ids: Vec<_> = first_batch
6791 .requests
6792 .iter()
6793 .map(|request| request.request.id.clone())
6794 .collect();
6795 for request_id in &first_ids {
6796 scheduler.mark_prefill_complete(request_id, 128);
6797 }
6798
6799 for request_id in first_ids.iter().take(4) {
6800 assert!(scheduler.defer_decode_to_waiting_for_capacity(request_id, 10));
6801 }
6802
6803 let first_mixed = scheduler
6804 .create_iteration_batch(mixed_hint.clone())
6805 .unwrap();
6806 let first_mixed_ids: HashSet<RequestId> = first_mixed
6807 .requests
6808 .iter()
6809 .map(|request| request.request.id.clone())
6810 .collect();
6811 assert!(first_mixed_ids.contains(&first_ids[0]));
6812 assert!(first_mixed_ids.contains(&first_ids[1]));
6813 assert!(!first_mixed_ids.contains(&first_ids[2]));
6814 assert!(!first_mixed_ids.contains(&first_ids[3]));
6815 assert!(scheduler.defer_prefill_to_waiting(&first_ids[0]));
6816 assert!(scheduler.defer_prefill_to_waiting(&first_ids[1]));
6817
6818 let second_mixed = scheduler.create_iteration_batch(mixed_hint).unwrap();
6819 let second_mixed_ids: HashSet<RequestId> = second_mixed
6820 .requests
6821 .iter()
6822 .map(|request| request.request.id.clone())
6823 .collect();
6824 assert!(
6825 !second_mixed_ids.contains(&first_ids[0]),
6826 "the first failed recompute must still be skipped in the same release epoch"
6827 );
6828 assert!(
6829 !second_mixed_ids.contains(&first_ids[1]),
6830 "already failed blocked recomputes must not consume the later candidate's slot"
6831 );
6832 assert!(
6833 second_mixed_ids.contains(&first_ids[2]),
6834 "marked queue-head requests should not block a later untried recompute candidate"
6835 );
6836 assert_eq!(
6837 second_mixed.requests.len(),
6838 7,
6839 "six decode-ready survivors plus one later recompute should be scheduled after same-epoch failures"
6840 );
6841 }
6842
6843 #[tokio::test]
6844 async fn capacity_deferred_mixed_recompute_waits_after_capacity_feedback_until_release() {
6845 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6846 max_running_requests: 8,
6847 prompt_token_estimate: true,
6848 ..SchedulerConfig::default()
6849 });
6850 let hint = BatchHint {
6851 max_batch_size: 8,
6852 max_tokens: 1024,
6853 target_latency_ms: None,
6854 available_memory: None,
6855 resource_constraints: Default::default(),
6856 };
6857
6858 for _ in 0..8 {
6859 enqueue_waiting(
6860 &scheduler,
6861 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6862 );
6863 }
6864
6865 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6866 let first_ids: Vec<_> = first_batch
6867 .requests
6868 .iter()
6869 .map(|request| request.request.id.clone())
6870 .collect();
6871 for request_id in &first_ids {
6872 scheduler.mark_prefill_complete(request_id, 128);
6873 }
6874
6875 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
6876 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[1], 4));
6877
6878 let first_mixed = scheduler.create_iteration_batch(hint.clone()).unwrap();
6879 let first_mixed_ids: HashSet<RequestId> = first_mixed
6880 .requests
6881 .iter()
6882 .map(|request| request.request.id.clone())
6883 .collect();
6884 assert!(first_mixed_ids.contains(&first_ids[0]));
6885
6886 scheduler.defer_capacity_deferred_mixed_recompute_until_release();
6887 assert!(scheduler.defer_prefill_to_waiting(&first_ids[0]));
6888
6889 let blocked_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6890 let blocked_ids: HashSet<RequestId> = blocked_batch
6891 .requests
6892 .iter()
6893 .map(|request| request.request.id.clone())
6894 .collect();
6895 assert!(
6896 !blocked_ids.contains(&first_ids[0]) && !blocked_ids.contains(&first_ids[1]),
6897 "mixed recompute should wait after capacity feedback until a real release"
6898 );
6899 assert_eq!(
6900 blocked_batch.requests.len(),
6901 6,
6902 "decode-ready survivors should continue while blocked recomputes wait for capacity release"
6903 );
6904
6905 let response = InferenceResponse {
6906 request_id: first_ids[2].clone(),
6907 text: String::new(),
6908 tokens: Vec::new(),
6909 finish_reason: ferrum_types::FinishReason::Length,
6910 usage: ferrum_types::TokenUsage::new(0, 0),
6911 latency_ms: 0,
6912 created_at: chrono::Utc::now(),
6913 metadata: Default::default(),
6914 api_response: None,
6915 execution_evidence: None,
6916 };
6917 scheduler
6918 .complete(first_ids[2].clone(), &response)
6919 .await
6920 .unwrap();
6921
6922 let after_release = scheduler.create_iteration_batch(hint).unwrap();
6923 let after_release_ids: HashSet<RequestId> = after_release
6924 .requests
6925 .iter()
6926 .map(|request| request.request.id.clone())
6927 .collect();
6928 assert!(
6929 after_release_ids.contains(&first_ids[0]) || after_release_ids.contains(&first_ids[1]),
6930 "capacity release should reopen bounded mixed recompute scanning"
6931 );
6932 }
6933
6934 #[tokio::test]
6935 async fn capacity_deferred_mixed_recompute_resumes_after_decode_capacity_release() {
6936 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
6937 max_running_requests: 8,
6938 prompt_token_estimate: true,
6939 ..SchedulerConfig::default()
6940 });
6941 let hint = BatchHint {
6942 max_batch_size: 8,
6943 max_tokens: 1024,
6944 target_latency_ms: None,
6945 available_memory: None,
6946 resource_constraints: Default::default(),
6947 };
6948
6949 for _ in 0..8 {
6950 enqueue_waiting(
6951 &scheduler,
6952 create_test_request_with_prompt_tokens(Priority::Normal, 128),
6953 );
6954 }
6955
6956 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6957 let first_ids: Vec<_> = first_batch
6958 .requests
6959 .iter()
6960 .map(|request| request.request.id.clone())
6961 .collect();
6962 for request_id in &first_ids {
6963 scheduler.mark_prefill_complete(request_id, 128);
6964 }
6965
6966 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
6967 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[1], 4));
6968
6969 let first_mixed = scheduler.create_iteration_batch(hint.clone()).unwrap();
6970 let first_mixed_ids: HashSet<RequestId> = first_mixed
6971 .requests
6972 .iter()
6973 .map(|request| request.request.id.clone())
6974 .collect();
6975 assert!(first_mixed_ids.contains(&first_ids[0]) || first_mixed_ids.contains(&first_ids[1]));
6976
6977 scheduler.defer_capacity_deferred_mixed_recompute_until_release();
6978 assert!(scheduler.defer_prefill_to_waiting(&first_ids[0]));
6979
6980 let blocked_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
6981 let blocked_ids: HashSet<RequestId> = blocked_batch
6982 .requests
6983 .iter()
6984 .map(|request| request.request.id.clone())
6985 .collect();
6986 assert!(
6987 !blocked_ids.contains(&first_ids[0]) && !blocked_ids.contains(&first_ids[1]),
6988 "mixed recompute should stay blocked until fresh capacity evidence"
6989 );
6990
6991 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[2], 4));
6992 scheduler.record_capacity_deferred_mixed_recompute_release_evidence();
6993 let after_decode_defer = scheduler.trace_snapshot();
6994 assert_eq!(after_decode_defer.capacity_blocked_waiting_len, 2);
6995
6996 let after_physical_release = scheduler.create_iteration_batch(hint).unwrap();
6997 let recompute_ids: Vec<_> = after_physical_release
6998 .requests
6999 .iter()
7000 .filter(|request| {
7001 first_ids[..3].contains(&request.request.id) && request.tokens_to_process != Some(1)
7002 })
7003 .map(|request| request.request.id.clone())
7004 .collect();
7005 assert_eq!(
7006 recompute_ids.len(),
7007 2,
7008 "physical KV release should reopen older recomputes that fit active capacity"
7009 );
7010 assert!(recompute_ids.contains(&first_ids[0]));
7011 assert!(recompute_ids.contains(&first_ids[1]));
7012 assert!(
7013 !recompute_ids.contains(&first_ids[2]),
7014 "the just-deferred decode should wait behind older blocked recomputes when active capacity is limited"
7015 );
7016 }
7017
7018 #[tokio::test]
7019 async fn capacity_deferred_mixed_recompute_waits_until_kv_snapshot_has_required_free_blocks() {
7020 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7021 max_running_requests: 8,
7022 prompt_token_estimate: true,
7023 ..SchedulerConfig::default()
7024 });
7025 let hint = BatchHint {
7026 max_batch_size: 8,
7027 max_tokens: 1024,
7028 target_latency_ms: None,
7029 available_memory: None,
7030 resource_constraints: Default::default(),
7031 };
7032
7033 for _ in 0..8 {
7034 enqueue_waiting(
7035 &scheduler,
7036 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7037 );
7038 }
7039
7040 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7041 let first_ids: Vec<_> = first_batch
7042 .requests
7043 .iter()
7044 .map(|request| request.request.id.clone())
7045 .collect();
7046 for request_id in &first_ids {
7047 scheduler.mark_prefill_complete(request_id, 128);
7048 }
7049
7050 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
7051 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[1], 4));
7052
7053 let first_mixed = scheduler.create_iteration_batch(hint.clone()).unwrap();
7054 let first_mixed_ids: HashSet<RequestId> = first_mixed
7055 .requests
7056 .iter()
7057 .map(|request| request.request.id.clone())
7058 .collect();
7059 assert!(first_mixed_ids.contains(&first_ids[0]) || first_mixed_ids.contains(&first_ids[1]));
7060
7061 scheduler.defer_capacity_deferred_mixed_recompute_until_kv_capacity(
7062 Some(4),
7063 Some(0),
7064 Some(1),
7065 );
7066 let blocked_snapshot = scheduler.trace_snapshot();
7067 assert_eq!(
7068 blocked_snapshot.capacity_mixed_recompute_required_blocks_per_slot,
7069 Some(4)
7070 );
7071 assert_eq!(
7072 blocked_snapshot.capacity_mixed_recompute_observed_free_blocks,
7073 Some(0)
7074 );
7075 assert_eq!(
7076 blocked_snapshot.capacity_mixed_recompute_blocked_until_epoch,
7077 blocked_snapshot.capacity_mixed_recompute_epoch + 1
7078 );
7079 assert!(scheduler.defer_prefill_to_waiting(&first_ids[0]));
7080
7081 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[2], 4));
7082 scheduler.record_capacity_deferred_mixed_recompute_kv_capacity_snapshot(3);
7083
7084 let still_blocked = scheduler.create_iteration_batch(hint.clone()).unwrap();
7085 let still_blocked_ids: HashSet<RequestId> = still_blocked
7086 .requests
7087 .iter()
7088 .map(|request| request.request.id.clone())
7089 .collect();
7090 assert!(
7091 !still_blocked_ids.contains(&first_ids[0])
7092 && !still_blocked_ids.contains(&first_ids[1]),
7093 "insufficient paged-KV free blocks must not reopen failed mixed recompute"
7094 );
7095
7096 scheduler.record_capacity_deferred_mixed_recompute_kv_capacity_snapshot(4);
7097 let exact_without_headroom = scheduler.create_iteration_batch(hint.clone()).unwrap();
7098 let exact_without_headroom_ids: HashSet<RequestId> = exact_without_headroom
7099 .requests
7100 .iter()
7101 .map(|request| request.request.id.clone())
7102 .collect();
7103 assert!(
7104 !exact_without_headroom_ids.contains(&first_ids[0])
7105 && !exact_without_headroom_ids.contains(&first_ids[1]),
7106 "a KV snapshot must leave allocator headroom before reopening mixed recompute"
7107 );
7108
7109 scheduler.record_capacity_deferred_mixed_recompute_kv_capacity_snapshot(5);
7110 let reopened_snapshot = scheduler.trace_snapshot();
7111 assert_eq!(
7112 reopened_snapshot.capacity_mixed_recompute_observed_free_blocks,
7113 Some(5)
7114 );
7115 assert!(
7116 reopened_snapshot.capacity_mixed_recompute_epoch
7117 >= reopened_snapshot.capacity_mixed_recompute_blocked_until_epoch
7118 );
7119 let after_enough_free = scheduler.create_iteration_batch(hint).unwrap();
7120 let after_enough_ids: HashSet<RequestId> = after_enough_free
7121 .requests
7122 .iter()
7123 .map(|request| request.request.id.clone())
7124 .collect();
7125 assert!(
7126 after_enough_ids.contains(&first_ids[0]) || after_enough_ids.contains(&first_ids[1]),
7127 "mixed recompute should reopen once the model-owned KV snapshot reaches the failed admission need"
7128 );
7129 }
7130
7131 #[tokio::test]
7132 async fn capacity_deferred_mixed_recompute_reopens_from_capacity_feedback_when_fit() {
7133 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7134 max_running_requests: 8,
7135 prompt_token_estimate: true,
7136 ..SchedulerConfig::default()
7137 });
7138 let hint = BatchHint {
7139 max_batch_size: 8,
7140 max_tokens: 1024,
7141 target_latency_ms: None,
7142 available_memory: None,
7143 resource_constraints: Default::default(),
7144 };
7145
7146 for _ in 0..8 {
7147 enqueue_waiting(
7148 &scheduler,
7149 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7150 );
7151 }
7152
7153 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7154 let first_ids: Vec<_> = first_batch
7155 .requests
7156 .iter()
7157 .map(|request| request.request.id.clone())
7158 .collect();
7159 for request_id in &first_ids {
7160 scheduler.mark_prefill_complete(request_id, 128);
7161 }
7162
7163 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
7164
7165 let first_mixed = scheduler.create_iteration_batch(hint.clone()).unwrap();
7166 let first_mixed_ids: HashSet<RequestId> = first_mixed
7167 .requests
7168 .iter()
7169 .map(|request| request.request.id.clone())
7170 .collect();
7171 assert!(first_mixed_ids.contains(&first_ids[0]));
7172
7173 scheduler.defer_capacity_deferred_mixed_recompute_until_kv_capacity(
7174 Some(4),
7175 Some(0),
7176 Some(1),
7177 );
7178 assert!(scheduler.defer_prefill_to_waiting(&first_ids[0]));
7179
7180 scheduler.defer_capacity_deferred_mixed_recompute_until_kv_capacity(
7181 Some(4),
7182 Some(5),
7183 Some(1),
7184 );
7185
7186 let reopened = scheduler.create_iteration_batch(hint).unwrap();
7187 let reopened_ids: HashSet<RequestId> = reopened
7188 .requests
7189 .iter()
7190 .map(|request| request.request.id.clone())
7191 .collect();
7192 assert!(
7193 reopened_ids.contains(&first_ids[0]),
7194 "structured capacity feedback with enough free blocks should reopen a narrower recompute without waiting for a separate snapshot call"
7195 );
7196 }
7197
7198 #[tokio::test]
7199 async fn capacity_deferred_mixed_recompute_paces_width_by_kv_snapshot() {
7200 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7201 max_running_requests: 8,
7202 prompt_token_estimate: true,
7203 ..SchedulerConfig::default()
7204 });
7205 let hint = BatchHint {
7206 max_batch_size: 8,
7207 max_tokens: 1024,
7208 target_latency_ms: None,
7209 available_memory: None,
7210 resource_constraints: Default::default(),
7211 };
7212
7213 for _ in 0..8 {
7214 enqueue_waiting(
7215 &scheduler,
7216 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7217 );
7218 }
7219
7220 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7221 let first_ids: Vec<_> = first_batch
7222 .requests
7223 .iter()
7224 .map(|request| request.request.id.clone())
7225 .collect();
7226 for request_id in &first_ids {
7227 scheduler.mark_prefill_complete(request_id, 128);
7228 }
7229
7230 for request_id in first_ids.iter().take(4) {
7231 assert!(scheduler.defer_decode_to_waiting_for_capacity(request_id, 8));
7232 }
7233
7234 scheduler.defer_capacity_deferred_mixed_recompute_until_kv_capacity(
7235 Some(16),
7236 Some(0),
7237 Some(4),
7238 );
7239 scheduler.record_capacity_deferred_mixed_recompute_kv_capacity_snapshot(9);
7240
7241 let paced_mixed = scheduler.create_iteration_batch(hint).unwrap();
7242 let scheduled_recomputes = paced_mixed
7243 .requests
7244 .iter()
7245 .filter(|request| {
7246 first_ids[..4].contains(&request.request.id) && request.tokens_to_process != Some(1)
7247 })
7248 .count();
7249 let scheduled_decodes = paced_mixed
7250 .requests
7251 .iter()
7252 .filter(|request| {
7253 first_ids[4..].contains(&request.request.id) && request.tokens_to_process == Some(1)
7254 })
7255 .count();
7256
7257 assert_eq!(scheduled_decodes, 4);
7258 assert_eq!(
7259 scheduled_recomputes, 2,
7260 "free KV blocks should pace the number of reopened capacity-blocked recomputes"
7261 );
7262 assert_eq!(scheduler.trace_snapshot().capacity_blocked_waiting_len, 2);
7263 }
7264
7265 #[tokio::test]
7266 async fn release_ready_capacity_deferred_recompute_still_uses_kv_budget_under_decode_pressure()
7267 {
7268 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7269 max_running_requests: 8,
7270 prompt_token_estimate: true,
7271 ..SchedulerConfig::default()
7272 });
7273 let init_hint = BatchHint {
7274 max_batch_size: 8,
7275 max_tokens: 1024,
7276 target_latency_ms: None,
7277 available_memory: None,
7278 resource_constraints: Default::default(),
7279 };
7280
7281 for _ in 0..8 {
7282 enqueue_waiting(
7283 &scheduler,
7284 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7285 );
7286 }
7287
7288 let first_batch = scheduler.create_iteration_batch(init_hint.clone()).unwrap();
7289 let first_ids: Vec<_> = first_batch
7290 .requests
7291 .iter()
7292 .map(|request| request.request.id.clone())
7293 .collect();
7294 for request_id in &first_ids {
7295 scheduler.mark_prefill_complete(request_id, 128);
7296 }
7297
7298 for request_id in first_ids.iter().take(4) {
7299 assert!(scheduler.defer_decode_to_waiting_for_capacity(request_id, 8));
7300 }
7301
7302 scheduler.record_capacity_release_progress();
7303
7304 scheduler.defer_capacity_deferred_mixed_recompute_until_kv_capacity(
7305 Some(16),
7306 Some(0),
7307 Some(4),
7308 );
7309 scheduler.record_capacity_deferred_mixed_recompute_kv_capacity_snapshot(9);
7310
7311 let paced_mixed = scheduler.create_iteration_batch(init_hint).unwrap();
7312 let scheduled_recomputes = paced_mixed
7313 .requests
7314 .iter()
7315 .filter(|request| {
7316 first_ids[..4].contains(&request.request.id) && request.tokens_to_process != Some(1)
7317 })
7318 .count();
7319 let scheduled_decodes = paced_mixed
7320 .requests
7321 .iter()
7322 .filter(|request| {
7323 first_ids[4..].contains(&request.request.id) && request.tokens_to_process == Some(1)
7324 })
7325 .count();
7326
7327 assert_eq!(scheduled_decodes, 4);
7328 assert_eq!(
7329 scheduled_recomputes, 2,
7330 "release-ready recomputes under decode pressure must still obey the KV snapshot budget"
7331 );
7332 assert_eq!(scheduler.trace_snapshot().capacity_blocked_waiting_len, 0);
7333 assert_eq!(
7334 scheduler.trace_snapshot().waiting_queue_len,
7335 2,
7336 "the remaining capacity-deferred recomputes should stay waiting for a later KV window"
7337 );
7338 }
7339
7340 #[tokio::test]
7341 async fn capacity_deferred_decode_waits_for_release_without_bounded_mixed_budget() {
7342 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7343 max_running_requests: 4,
7344 prompt_token_estimate: true,
7345 ..SchedulerConfig::default()
7346 });
7347 let hint = BatchHint {
7348 max_batch_size: 4,
7349 max_tokens: 1024,
7350 target_latency_ms: None,
7351 available_memory: None,
7352 resource_constraints: Default::default(),
7353 };
7354
7355 for _ in 0..2 {
7356 enqueue_waiting(
7357 &scheduler,
7358 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7359 );
7360 }
7361
7362 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7363 let first_ids: Vec<_> = first_batch
7364 .requests
7365 .iter()
7366 .map(|request| request.request.id.clone())
7367 .collect();
7368 for request_id in &first_ids {
7369 scheduler.mark_prefill_complete(request_id, 128);
7370 }
7371
7372 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 4));
7373 let deferred = scheduler.trace_snapshot();
7374 assert_eq!(deferred.waiting_queue_len, 1);
7375 assert_eq!(deferred.decode_queue_len, 1);
7376 assert_eq!(deferred.active_len, 1);
7377 assert_eq!(deferred.capacity_blocked_waiting_len, 1);
7378
7379 let response = InferenceResponse {
7380 request_id: first_ids[1].clone(),
7381 text: String::new(),
7382 tokens: Vec::new(),
7383 finish_reason: ferrum_types::FinishReason::Length,
7384 usage: ferrum_types::TokenUsage::new(0, 0),
7385 latency_ms: 0,
7386 created_at: chrono::Utc::now(),
7387 metadata: Default::default(),
7388 api_response: None,
7389 execution_evidence: None,
7390 };
7391 scheduler
7392 .complete(first_ids[1].clone(), &response)
7393 .await
7394 .unwrap();
7395
7396 let after_release = scheduler.create_iteration_batch(hint).unwrap();
7397 let scheduled_ids: HashSet<RequestId> = after_release
7398 .requests
7399 .iter()
7400 .map(|request| request.request.id.clone())
7401 .collect();
7402 assert!(
7403 scheduled_ids.contains(&first_ids[0]),
7404 "a real capacity release should make the deferred recompute eligible again"
7405 );
7406 assert_eq!(scheduler.trace_snapshot().capacity_blocked_waiting_len, 0);
7407 }
7408
7409 #[tokio::test]
7410 async fn capacity_backpressure_survives_partial_prefill_and_waits_for_release() {
7411 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7412 max_running_requests: 4,
7413 prompt_token_estimate: true,
7414 ..SchedulerConfig::default()
7415 });
7416 let hint = BatchHint {
7417 max_batch_size: 4,
7418 max_tokens: 1024,
7419 target_latency_ms: None,
7420 available_memory: None,
7421 resource_constraints: Default::default(),
7422 };
7423
7424 for _ in 0..4 {
7425 enqueue_waiting(
7426 &scheduler,
7427 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7428 );
7429 }
7430
7431 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7432 let first_ids: Vec<RequestId> = first_batch
7433 .requests
7434 .iter()
7435 .map(|request| request.request.id.clone())
7436 .collect();
7437 for request in &first_batch.requests {
7438 assert!(scheduler.defer_prefill_to_waiting(&request.request.id));
7439 }
7440 assert_eq!(
7441 scheduler.trace_snapshot().capacity_backpressure_admit_limit,
7442 Some(2)
7443 );
7444
7445 let second_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7446 assert_eq!(second_batch.requests.len(), 2);
7447 let second_ids: HashSet<RequestId> = second_batch
7448 .requests
7449 .iter()
7450 .map(|request| request.request.id.clone())
7451 .collect();
7452 let progressed_id = second_batch.requests[0].request.id.clone();
7453 assert!(!scheduler.mark_prefill_chunk_processed(&progressed_id, 128, 1));
7454 assert_eq!(
7455 scheduler.trace_snapshot().capacity_backpressure_admit_limit,
7456 Some(2),
7457 "partial prefill progress still consumes capacity and should not reopen the failed width"
7458 );
7459 assert!(scheduler.mark_prefill_chunk_processed(&progressed_id, 128, 127));
7460 assert_eq!(
7461 scheduler.trace_snapshot().capacity_backpressure_admit_limit,
7462 None,
7463 "full prefill completion should relax the capacity backpressure window"
7464 );
7465
7466 let third_batch = scheduler.create_iteration_batch(hint).unwrap();
7467 assert_eq!(
7468 third_batch.requests.len(),
7469 2,
7470 "prefill completion may continue existing active work but must not release blocked waiting prefills"
7471 );
7472 let after = scheduler.trace_snapshot();
7473 assert_eq!(after.waiting_queue_len, 2);
7474 assert_eq!(after.active_len, 2);
7475
7476 let response = InferenceResponse {
7477 request_id: progressed_id.clone(),
7478 text: String::new(),
7479 tokens: Vec::new(),
7480 finish_reason: ferrum_types::FinishReason::Length,
7481 usage: ferrum_types::TokenUsage::new(0, 0),
7482 latency_ms: 0,
7483 created_at: chrono::Utc::now(),
7484 metadata: Default::default(),
7485 api_response: None,
7486 execution_evidence: None,
7487 };
7488 scheduler.complete(progressed_id, &response).await.unwrap();
7489
7490 let after_release_batch = scheduler
7491 .create_iteration_batch(BatchHint {
7492 max_batch_size: 4,
7493 max_tokens: 1024,
7494 target_latency_ms: None,
7495 available_memory: None,
7496 resource_constraints: Default::default(),
7497 })
7498 .unwrap();
7499 let after_release_ids: HashSet<RequestId> = after_release_batch
7500 .requests
7501 .iter()
7502 .map(|request| request.request.id.clone())
7503 .collect();
7504 assert!(
7505 first_ids
7506 .iter()
7507 .any(|id| !second_ids.contains(id) && after_release_ids.contains(id)),
7508 "actual request completion should release capacity and reopen blocked waiting prefills"
7509 );
7510 }
7511
7512 #[tokio::test]
7513 async fn capacity_backpressure_survives_cancel_without_token_progress() {
7514 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7515 max_running_requests: 4,
7516 prompt_token_estimate: true,
7517 ..SchedulerConfig::default()
7518 });
7519 let hint = BatchHint {
7520 max_batch_size: 4,
7521 max_tokens: 1024,
7522 target_latency_ms: None,
7523 available_memory: None,
7524 resource_constraints: Default::default(),
7525 };
7526
7527 for _ in 0..4 {
7528 enqueue_waiting(
7529 &scheduler,
7530 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7531 );
7532 }
7533
7534 let first_batch = scheduler.create_iteration_batch(hint).unwrap();
7535 let first_ids: Vec<_> = first_batch
7536 .requests
7537 .iter()
7538 .map(|request| request.request.id.clone())
7539 .collect();
7540 for request_id in &first_ids {
7541 assert!(scheduler.defer_prefill_to_waiting(request_id));
7542 }
7543 assert_eq!(
7544 scheduler.trace_snapshot().capacity_backpressure_admit_limit,
7545 Some(2)
7546 );
7547
7548 assert!(scheduler.cancel(first_ids[0].clone()).await.unwrap());
7549 let after_cancel = scheduler.trace_snapshot();
7550 assert_eq!(after_cancel.cancelled_total, 1);
7551 assert_eq!(
7552 after_cancel.capacity_backpressure_admit_limit,
7553 Some(2),
7554 "cancellation frees a slot but is not evidence that the failed admission width now fits"
7555 );
7556 }
7557
7558 #[tokio::test]
7559 async fn test_batch_creation() {
7560 let config = SchedulerConfig::default();
7561 let scheduler = ContinuousBatchScheduler::new(config);
7562
7563 for _ in 0..5 {
7565 scheduler
7566 .submit(create_test_request(Priority::Normal))
7567 .await
7568 .unwrap();
7569 }
7570
7571 let batch = scheduler.next_batch(BatchHint::simple(10)).await;
7573 assert!(batch.is_some());
7574
7575 assert!(scheduler.prefilling_count() > 0 || scheduler.decoding_count() > 0);
7577 }
7578
7579 #[test]
7580 fn prompt_token_metadata_expands_prefill_admission() {
7581 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7582 prompt_token_estimate: true,
7583 ..SchedulerConfig::default()
7584 });
7585
7586 for _ in 0..16 {
7587 enqueue_waiting(
7588 &scheduler,
7589 create_test_request_with_prompt_tokens(Priority::Normal, 256),
7590 );
7591 }
7592
7593 let hint = BatchHint {
7594 max_batch_size: 32,
7595 max_tokens: 2048,
7596 target_latency_ms: None,
7597 available_memory: None,
7598 resource_constraints: Default::default(),
7599 };
7600 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7601 assert_eq!(first_batch.requests.len(), 8);
7602
7603 for request in first_batch.requests {
7604 scheduler.mark_prefill_complete(&request.request.id, 256);
7605 }
7606
7607 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
7608 assert_eq!(mixed_batch.requests.len(), 16);
7609 assert_eq!(mixed_batch.resource_requirements.gpu_memory, 2048 * 16);
7610 }
7611
7612 #[test]
7613 fn prompt_token_metadata_can_be_disabled_for_prefill_admission() {
7614 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7615 prompt_token_estimate: false,
7616 ..SchedulerConfig::default()
7617 });
7618
7619 for _ in 0..16 {
7620 enqueue_waiting(
7621 &scheduler,
7622 create_test_request_with_prompt_tokens(Priority::Normal, 256),
7623 );
7624 }
7625
7626 let batch = scheduler
7627 .create_iteration_batch(BatchHint {
7628 max_batch_size: 32,
7629 max_tokens: 2048,
7630 target_latency_ms: None,
7631 available_memory: None,
7632 resource_constraints: Default::default(),
7633 })
7634 .unwrap();
7635 assert_eq!(batch.requests.len(), 4);
7636 }
7637
7638 #[test]
7639 fn scheduler_runtime_config_is_captured_at_construction() {
7640 let mut config = SchedulerConfig {
7641 prompt_token_estimate: true,
7642 ..SchedulerConfig::default()
7643 };
7644 let scheduler = ContinuousBatchScheduler::new(config.clone());
7645 config.prompt_token_estimate = false;
7646
7647 for _ in 0..16 {
7648 enqueue_waiting(
7649 &scheduler,
7650 create_test_request_with_prompt_tokens(Priority::Normal, 256),
7651 );
7652 }
7653
7654 let batch = scheduler
7655 .create_iteration_batch(BatchHint {
7656 max_batch_size: 32,
7657 max_tokens: 2048,
7658 target_latency_ms: None,
7659 available_memory: None,
7660 resource_constraints: Default::default(),
7661 })
7662 .unwrap();
7663 assert_eq!(batch.requests.len(), 8);
7664 }
7665
7666 #[test]
7667 fn max_running_requests_limits_waiting_admission() {
7668 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7669 max_running_requests: 1,
7670 prompt_token_estimate: true,
7671 ..SchedulerConfig::default()
7672 });
7673
7674 for _ in 0..3 {
7675 enqueue_waiting(
7676 &scheduler,
7677 create_test_request_with_prompt_tokens(Priority::Normal, 128),
7678 );
7679 }
7680
7681 let hint = BatchHint {
7682 max_batch_size: 8,
7683 max_tokens: 1024,
7684 target_latency_ms: None,
7685 available_memory: None,
7686 resource_constraints: Default::default(),
7687 };
7688 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7689 assert_eq!(first_batch.requests.len(), 1);
7690 assert_eq!(scheduler.prefilling_count(), 1);
7691 assert_eq!(scheduler.waiting_count(), 2);
7692
7693 let active_batch = scheduler.create_iteration_batch(hint).unwrap();
7694 assert_eq!(active_batch.requests.len(), 1);
7695 assert_eq!(
7696 active_batch.requests[0].request.id, first_batch.requests[0].request.id,
7697 "scheduler must not admit another waiting request while the active cap is full"
7698 );
7699 assert_eq!(scheduler.prefilling_count(), 1);
7700 assert_eq!(scheduler.waiting_count(), 2);
7701 }
7702
7703 #[test]
7704 fn newly_admitted_prefill_uses_remaining_budget_with_decode() {
7705 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7706 prompt_token_estimate: true,
7707 ..SchedulerConfig::default()
7708 });
7709 let hint = BatchHint {
7710 max_batch_size: 4,
7711 max_tokens: 4,
7712 target_latency_ms: None,
7713 available_memory: None,
7714 resource_constraints: Default::default(),
7715 };
7716
7717 let first = create_test_request_with_prompt_tokens(Priority::Normal, 2);
7718 let first_id = first.id.clone();
7719 enqueue_waiting(&scheduler, first);
7720 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7721 assert_eq!(first_batch.requests.len(), 1);
7722 scheduler.mark_prefill_complete(&first_id, 2);
7723
7724 let second = create_test_request_with_prompt_tokens(Priority::Normal, 2);
7725 let second_id = second.id.clone();
7726 enqueue_waiting(&scheduler, second);
7727
7728 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
7729 let ids: HashSet<RequestId> = mixed_batch
7730 .requests
7731 .iter()
7732 .map(|request| request.request.id.clone())
7733 .collect();
7734 assert_eq!(mixed_batch.requests.len(), 2);
7735 assert!(
7736 ids.contains(&first_id),
7737 "decode request should remain scheduled"
7738 );
7739 assert!(
7740 ids.contains(&second_id),
7741 "newly admitted prefill should use remaining same-iteration budget"
7742 );
7743 assert_eq!(mixed_batch.resource_requirements.gpu_memory, 3 * 16);
7744 }
7745
7746 #[test]
7747 fn default_scheduler_caps_mixed_prefill_only_under_decode_pressure() {
7748 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7749 max_running_requests: 8,
7750 prompt_token_estimate: true,
7751 ..SchedulerConfig::default()
7752 });
7753 let hint = BatchHint {
7754 max_batch_size: 8,
7755 max_tokens: 2048,
7756 target_latency_ms: None,
7757 available_memory: None,
7758 resource_constraints: Default::default(),
7759 };
7760
7761 let first = create_test_request_with_prompt_tokens(Priority::Normal, 256);
7762 let first_id = first.id.clone();
7763 enqueue_waiting(&scheduler, first);
7764 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7765 assert_eq!(first_batch.requests.len(), 1);
7766 scheduler.mark_prefill_complete(&first_id, 256);
7767
7768 for _ in 0..4 {
7769 enqueue_waiting(
7770 &scheduler,
7771 create_test_request_with_prompt_tokens(Priority::Normal, 256),
7772 );
7773 }
7774
7775 let low_decode_pressure = scheduler.create_iteration_batch(hint.clone()).unwrap();
7776 assert_eq!(
7777 low_decode_pressure.requests.len(),
7778 5,
7779 "small decode cohorts should use remaining token budget to build concurrency"
7780 );
7781 assert_eq!(
7782 low_decode_pressure.resource_requirements.gpu_memory,
7783 (1 + 4 * 256) * 16
7784 );
7785 assert_eq!(
7786 low_decode_pressure
7787 .requests
7788 .iter()
7789 .filter(|request| request.request.id != first_id)
7790 .map(|request| request.tokens_to_process)
7791 .collect::<Vec<_>>(),
7792 vec![Some(256), Some(256), Some(256), Some(256)]
7793 );
7794
7795 for request in low_decode_pressure
7796 .requests
7797 .iter()
7798 .filter(|request| request.request.id != first_id)
7799 {
7800 scheduler.mark_prefill_complete(&request.request.id, 256);
7801 }
7802 assert_eq!(scheduler.decoding_count(), 5);
7803
7804 for _ in 0..4 {
7805 enqueue_waiting(
7806 &scheduler,
7807 create_test_request_with_prompt_tokens(Priority::Normal, 256),
7808 );
7809 }
7810
7811 let high_decode_pressure = scheduler.create_iteration_batch(hint).unwrap();
7812 let prefill_tokens: Vec<_> = high_decode_pressure
7813 .requests
7814 .iter()
7815 .filter(|request| request.tokens_to_process != Some(1))
7816 .map(|request| request.tokens_to_process)
7817 .collect();
7818 assert_eq!(
7819 high_decode_pressure.requests.len(),
7820 8,
7821 "high decode pressure should admit bounded partial prefills up to available slots"
7822 );
7823 assert_eq!(prefill_tokens, vec![Some(64), Some(64), Some(64)]);
7824 assert_eq!(
7825 high_decode_pressure.resource_requirements.gpu_memory,
7826 (5 + 192) * 16
7827 );
7828 }
7829
7830 #[test]
7831 fn max_batched_tokens_limits_prefill_admission_by_prompt_tokens() {
7832 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7833 prompt_token_estimate: true,
7834 ..SchedulerConfig::default()
7835 });
7836
7837 for _ in 0..4 {
7838 enqueue_waiting(
7839 &scheduler,
7840 create_test_request_with_prompt_tokens(Priority::Normal, 256),
7841 );
7842 }
7843
7844 let batch = scheduler
7845 .create_iteration_batch(BatchHint {
7846 max_batch_size: 8,
7847 max_tokens: 512,
7848 target_latency_ms: None,
7849 available_memory: None,
7850 resource_constraints: Default::default(),
7851 })
7852 .unwrap();
7853
7854 assert_eq!(batch.requests.len(), 2);
7855 assert_eq!(batch.resource_requirements.gpu_memory, 512 * 16);
7856 assert_eq!(
7857 scheduler.prefilling_count(),
7858 4,
7859 "max_tokens limits the emitted iteration batch, not waiting-to-prefill promotion"
7860 );
7861 assert_eq!(scheduler.waiting_count(), 0);
7862 }
7863
7864 #[test]
7865 fn long_prefill_uses_remaining_step_budget_instead_of_fixed_chunk() {
7866 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7867 prompt_token_estimate: true,
7868 ..SchedulerConfig::default()
7869 });
7870
7871 for _ in 0..2 {
7872 enqueue_waiting(
7873 &scheduler,
7874 create_test_request_with_prompt_tokens(Priority::Normal, 1536),
7875 );
7876 }
7877
7878 let batch = scheduler
7879 .create_iteration_batch(BatchHint {
7880 max_batch_size: 8,
7881 max_tokens: 2048,
7882 target_latency_ms: None,
7883 available_memory: None,
7884 resource_constraints: Default::default(),
7885 })
7886 .unwrap();
7887
7888 assert_eq!(batch.requests.len(), 2);
7889 assert_eq!(
7890 batch
7891 .requests
7892 .iter()
7893 .map(|request| request.tokens_to_process)
7894 .collect::<Vec<_>>(),
7895 vec![Some(1536), Some(512)]
7896 );
7897 assert_eq!(batch.resource_requirements.gpu_memory, 2048 * 16);
7898 }
7899
7900 #[test]
7901 fn prefill_first_until_active_skips_early_decodes() {
7902 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7903 prefill_first_until_active: Some(4),
7904 ..SchedulerConfig::default()
7905 });
7906
7907 for _ in 0..3 {
7908 enqueue_waiting(&scheduler, create_test_request(Priority::Normal));
7909 }
7910
7911 let hint = BatchHint {
7912 max_batch_size: 8,
7913 max_tokens: 1024,
7914 target_latency_ms: None,
7915 available_memory: None,
7916 resource_constraints: Default::default(),
7917 };
7918 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7919 assert_eq!(first_batch.requests.len(), 2);
7920 let first_ids: Vec<RequestId> = first_batch
7921 .requests
7922 .iter()
7923 .map(|request| request.request.id.clone())
7924 .collect();
7925 for id in &first_ids {
7926 scheduler.mark_prefill_complete(id, 512);
7927 }
7928 assert_eq!(scheduler.decoding_count(), 2);
7929
7930 let second_batch = scheduler.create_iteration_batch(hint).unwrap();
7931 assert_eq!(second_batch.requests.len(), 1);
7932 assert!(
7933 second_batch
7934 .requests
7935 .iter()
7936 .all(|request| !first_ids.contains(&request.request.id)),
7937 "fill-first should schedule more prefills before decoding early requests"
7938 );
7939 }
7940
7941 #[test]
7942 fn prefill_first_until_active_resumes_decodes_at_active_target() {
7943 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7944 prefill_first_until_active: Some(4),
7945 ..SchedulerConfig::default()
7946 });
7947
7948 for _ in 0..4 {
7949 enqueue_waiting(&scheduler, create_test_request(Priority::Normal));
7950 }
7951
7952 let hint = BatchHint {
7953 max_batch_size: 8,
7954 max_tokens: 1024,
7955 target_latency_ms: None,
7956 available_memory: None,
7957 resource_constraints: Default::default(),
7958 };
7959 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
7960 assert_eq!(first_batch.requests.len(), 2);
7961 let first_ids: Vec<RequestId> = first_batch
7962 .requests
7963 .iter()
7964 .map(|request| request.request.id.clone())
7965 .collect();
7966 for id in &first_ids {
7967 scheduler.mark_prefill_complete(id, 512);
7968 }
7969
7970 assert_eq!(scheduler.decoding_count(), 2);
7971 assert_eq!(scheduler.prefilling_count(), 2);
7972 assert_eq!(scheduler.active_count(), 4);
7973
7974 let second_batch = scheduler.create_iteration_batch(hint).unwrap();
7975 let scheduled_decodes = second_batch
7976 .requests
7977 .iter()
7978 .filter(|request| {
7979 first_ids.contains(&request.request.id) && request.tokens_to_process == Some(1)
7980 })
7981 .count();
7982 assert_eq!(
7983 scheduled_decodes, 2,
7984 "fill-first must not starve decode once the active target is reached"
7985 );
7986 }
7987
7988 #[test]
7989 fn capacity_backpressure_disables_prefill_first_decode_skip() {
7990 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
7991 max_running_requests: 4,
7992 prompt_token_estimate: true,
7993 prefill_first_until_active: Some(4),
7994 ..SchedulerConfig::default()
7995 });
7996
7997 for _ in 0..4 {
7998 enqueue_waiting(
7999 &scheduler,
8000 create_test_request_with_prompt_tokens(Priority::Normal, 128),
8001 );
8002 }
8003
8004 let hint = BatchHint {
8005 max_batch_size: 4,
8006 max_tokens: 512,
8007 target_latency_ms: None,
8008 available_memory: None,
8009 resource_constraints: Default::default(),
8010 };
8011 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8012 assert_eq!(first_batch.requests.len(), 4);
8013 let first_ids: Vec<RequestId> = first_batch
8014 .requests
8015 .iter()
8016 .map(|request| request.request.id.clone())
8017 .collect();
8018 for id in first_ids.iter().take(3) {
8019 scheduler.mark_prefill_complete(id, 128);
8020 }
8021 assert!(scheduler.defer_prefill_to_waiting(&first_ids[3]));
8022
8023 let after_defer = scheduler.trace_snapshot();
8024 assert_eq!(after_defer.decode_queue_len, 3);
8025 assert_eq!(after_defer.waiting_queue_len, 1);
8026 assert_eq!(after_defer.active_len, 3);
8027 assert_eq!(after_defer.capacity_backpressure_admit_limit, Some(1));
8028
8029 let second_batch = scheduler.create_iteration_batch(hint).unwrap();
8030 let scheduled_decodes = second_batch
8031 .requests
8032 .iter()
8033 .filter(|request| {
8034 first_ids[..3].contains(&request.request.id) && request.tokens_to_process == Some(1)
8035 })
8036 .count();
8037 assert_eq!(
8038 scheduled_decodes, 3,
8039 "capacity backpressure must let decode-ready requests run instead of repeatedly admitting a capacity-blocked prefill"
8040 );
8041 assert_eq!(
8042 second_batch.requests.len(),
8043 3,
8044 "a capacity-blocked prefill must wait for capacity evidence instead of refilling an empty batch slot"
8045 );
8046 assert_eq!(scheduler.trace_snapshot().capacity_blocked_waiting_len, 1);
8047 }
8048
8049 #[tokio::test]
8050 async fn prefill_capacity_defer_waits_for_release_while_decode_active() {
8051 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8052 max_running_requests: 4,
8053 prompt_token_estimate: true,
8054 prefill_first_until_active: Some(4),
8055 ..SchedulerConfig::default()
8056 });
8057
8058 for _ in 0..4 {
8059 enqueue_waiting(
8060 &scheduler,
8061 create_test_request_with_prompt_tokens(Priority::Normal, 128),
8062 );
8063 }
8064
8065 let hint = BatchHint {
8066 max_batch_size: 4,
8067 max_tokens: 512,
8068 target_latency_ms: None,
8069 available_memory: None,
8070 resource_constraints: Default::default(),
8071 };
8072 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8073 assert_eq!(first_batch.requests.len(), 4);
8074 let first_ids: Vec<RequestId> = first_batch
8075 .requests
8076 .iter()
8077 .map(|request| request.request.id.clone())
8078 .collect();
8079 for id in first_ids.iter().take(3) {
8080 scheduler.mark_prefill_complete(id, 128);
8081 }
8082 assert!(scheduler.defer_prefill_to_waiting(&first_ids[3]));
8083
8084 let deferred = scheduler.trace_snapshot();
8085 assert_eq!(deferred.decode_queue_len, 3);
8086 assert_eq!(deferred.waiting_queue_len, 1);
8087 assert_eq!(deferred.capacity_blocked_waiting_len, 1);
8088
8089 let blocked_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8090 let blocked_ids: HashSet<RequestId> = blocked_batch
8091 .requests
8092 .iter()
8093 .map(|request| request.request.id.clone())
8094 .collect();
8095 assert!(
8096 !blocked_ids.contains(&first_ids[3]),
8097 "failed prefill must not be retried while only decode progress has happened"
8098 );
8099 assert_eq!(blocked_batch.requests.len(), 3);
8100
8101 let response = InferenceResponse {
8102 request_id: first_ids[0].clone(),
8103 text: String::new(),
8104 tokens: Vec::new(),
8105 finish_reason: ferrum_types::FinishReason::Length,
8106 usage: ferrum_types::TokenUsage::new(0, 0),
8107 latency_ms: 0,
8108 created_at: chrono::Utc::now(),
8109 metadata: Default::default(),
8110 api_response: None,
8111 execution_evidence: None,
8112 };
8113 scheduler
8114 .complete(first_ids[0].clone(), &response)
8115 .await
8116 .unwrap();
8117
8118 let after_release = scheduler.create_iteration_batch(hint).unwrap();
8119 let after_release_ids: HashSet<RequestId> = after_release
8120 .requests
8121 .iter()
8122 .map(|request| request.request.id.clone())
8123 .collect();
8124 assert!(
8125 after_release_ids.contains(&first_ids[3]),
8126 "real capacity release should make the blocked prefill eligible again"
8127 );
8128 }
8129
8130 #[test]
8131 fn decode_progress_does_not_relax_capacity_backpressure() {
8132 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8133 max_running_requests: 4,
8134 prompt_token_estimate: true,
8135 prefill_first_until_active: Some(4),
8136 ..SchedulerConfig::default()
8137 });
8138
8139 for _ in 0..4 {
8140 enqueue_waiting(
8141 &scheduler,
8142 create_test_request_with_prompt_tokens(Priority::Normal, 128),
8143 );
8144 }
8145
8146 let hint = BatchHint {
8147 max_batch_size: 4,
8148 max_tokens: 512,
8149 target_latency_ms: None,
8150 available_memory: None,
8151 resource_constraints: Default::default(),
8152 };
8153 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8154 assert_eq!(first_batch.requests.len(), 4);
8155 let first_ids: Vec<RequestId> = first_batch
8156 .requests
8157 .iter()
8158 .map(|request| request.request.id.clone())
8159 .collect();
8160 for id in first_ids.iter().take(3) {
8161 scheduler.mark_prefill_complete(id, 128);
8162 }
8163 assert!(scheduler.defer_prefill_to_waiting(&first_ids[3]));
8164 assert_eq!(
8165 scheduler.trace_snapshot().capacity_backpressure_admit_limit,
8166 Some(1)
8167 );
8168
8169 for id in first_ids.iter().take(3) {
8170 scheduler.update_decode_progress(id, 1);
8171 }
8172 assert_eq!(
8173 scheduler.trace_snapshot().capacity_backpressure_admit_limit,
8174 Some(1),
8175 "decode progress consumes KV capacity and must not reopen waiting admission"
8176 );
8177
8178 let second_batch = scheduler.create_iteration_batch(hint).unwrap();
8179 let scheduled_decodes = second_batch
8180 .requests
8181 .iter()
8182 .filter(|request| {
8183 first_ids[..3].contains(&request.request.id) && request.tokens_to_process == Some(1)
8184 })
8185 .count();
8186 assert_eq!(
8187 scheduled_decodes, 3,
8188 "capacity backpressure should keep fill-first from skipping decode after decode progress"
8189 );
8190 assert_eq!(
8191 second_batch.requests.len(),
8192 3,
8193 "decode progress alone must not make a capacity-limited prefill refill the remaining batch slot"
8194 );
8195 }
8196
8197 #[test]
8198 fn capacity_backpressure_keeps_decode_survivors_wide_after_decode_defer() {
8199 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8200 max_running_requests: 8,
8201 prompt_token_estimate: true,
8202 ..SchedulerConfig::default()
8203 });
8204 let hint = BatchHint {
8205 max_batch_size: 8,
8206 max_tokens: 1024,
8207 target_latency_ms: None,
8208 available_memory: None,
8209 resource_constraints: Default::default(),
8210 };
8211
8212 for _ in 0..8 {
8213 enqueue_waiting(
8214 &scheduler,
8215 create_test_request_with_prompt_tokens(Priority::Normal, 128),
8216 );
8217 }
8218
8219 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8220 let first_ids: Vec<_> = first_batch
8221 .requests
8222 .iter()
8223 .map(|request| request.request.id.clone())
8224 .collect();
8225 for request_id in &first_ids {
8226 scheduler.mark_prefill_complete(request_id, 128);
8227 }
8228
8229 assert!(scheduler.defer_decode_to_waiting_for_capacity(&first_ids[0], 8));
8230 assert_eq!(
8231 scheduler.trace_snapshot().capacity_backpressure_admit_limit,
8232 Some(4)
8233 );
8234
8235 let capped = scheduler.create_iteration_batch(hint).unwrap();
8236 let scheduled_decodes = capped
8237 .requests
8238 .iter()
8239 .filter(|request| request.tokens_to_process == Some(1))
8240 .count();
8241
8242 assert_eq!(
8243 scheduled_decodes, 7,
8244 "decode scheduling should keep decode-ready survivors wide after a decode KV failure"
8245 );
8246 }
8247
8248 #[test]
8249 fn decode_capacity_backpressure_uses_structured_free_blocks_when_nearly_fit() {
8250 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8251 max_running_requests: 16,
8252 prompt_token_estimate: true,
8253 ..SchedulerConfig::default()
8254 });
8255 let hint = BatchHint {
8256 max_batch_size: 16,
8257 max_tokens: 2048,
8258 target_latency_ms: None,
8259 available_memory: None,
8260 resource_constraints: Default::default(),
8261 };
8262
8263 for _ in 0..16 {
8264 enqueue_waiting(
8265 &scheduler,
8266 create_test_request_with_prompt_tokens(Priority::Normal, 128),
8267 );
8268 }
8269
8270 let first_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8271 let first_ids: Vec<_> = first_batch
8272 .requests
8273 .iter()
8274 .map(|request| request.request.id.clone())
8275 .collect();
8276 for request_id in &first_ids {
8277 scheduler.mark_prefill_complete(request_id, 128);
8278 }
8279
8280 scheduler.record_decode_capacity_pressure(16, Some(12));
8281 let capped = scheduler.create_iteration_batch(hint).unwrap();
8282 let scheduled_decodes = capped
8283 .requests
8284 .iter()
8285 .filter(|request| request.tokens_to_process == Some(1))
8286 .count();
8287
8288 assert_eq!(
8289 scheduled_decodes, 11,
8290 "near-fit KV pressure should cap to usable free blocks instead of blindly halving"
8291 );
8292 }
8293
8294 #[test]
8295 fn prefill_step_chunk_caps_prefill_first_batches() {
8296 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8297 prompt_token_estimate: true,
8298 prefill_first_until_active: Some(4),
8299 prefill_step_chunk: Some(128),
8300 ..SchedulerConfig::default()
8301 });
8302
8303 for _ in 0..4 {
8304 enqueue_waiting(
8305 &scheduler,
8306 create_test_request_with_prompt_tokens(Priority::Normal, 512),
8307 );
8308 }
8309
8310 let batch = scheduler
8311 .create_iteration_batch(BatchHint {
8312 max_batch_size: 8,
8313 max_tokens: 2048,
8314 target_latency_ms: None,
8315 available_memory: None,
8316 resource_constraints: Default::default(),
8317 })
8318 .unwrap();
8319
8320 assert_eq!(batch.requests.len(), 4);
8321 assert_eq!(
8322 batch
8323 .requests
8324 .iter()
8325 .map(|request| request.tokens_to_process)
8326 .collect::<Vec<_>>(),
8327 vec![Some(128), Some(128), Some(128), Some(128)]
8328 );
8329 assert_eq!(batch.resource_requirements.gpu_memory, (4 * 128) * 16);
8330 }
8331
8332 #[test]
8333 fn elastic_prefill_budget_uses_live_capacity_not_configured_concurrency() {
8334 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8335 max_running_requests: 16,
8336 prompt_token_estimate: true,
8337 prefill_first_until_active: Some(16),
8338 prefill_step_chunk: None,
8339 ..SchedulerConfig::default()
8340 });
8341
8342 enqueue_waiting(
8343 &scheduler,
8344 create_test_request_with_prompt_tokens(Priority::Normal, 64),
8345 );
8346
8347 let batch = scheduler
8348 .create_iteration_batch(BatchHint {
8349 max_batch_size: 16,
8350 max_tokens: 192,
8351 target_latency_ms: None,
8352 available_memory: None,
8353 resource_constraints: Default::default(),
8354 })
8355 .unwrap();
8356
8357 assert_eq!(batch.requests.len(), 1);
8358 assert_eq!(batch.requests[0].tokens_to_process, Some(64));
8359 assert_eq!(batch.resource_requirements.gpu_memory, 64 * 16);
8360 }
8361
8362 #[test]
8363 fn active_decode_prefill_chunk_only_caps_when_decode_is_active() {
8364 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8365 prompt_token_estimate: true,
8366 active_decode_prefill_chunk: Some(64),
8367 ..SchedulerConfig::default()
8368 });
8369 let hint = BatchHint {
8370 max_batch_size: 2,
8371 max_tokens: 512,
8372 target_latency_ms: None,
8373 available_memory: None,
8374 resource_constraints: Default::default(),
8375 };
8376
8377 let first = create_test_request_with_prompt_tokens(Priority::Normal, 256);
8378 let first_id = first.id.clone();
8379 enqueue_waiting(&scheduler, first);
8380 let initial_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8381 assert_eq!(initial_batch.requests.len(), 1);
8382 assert_eq!(initial_batch.resource_requirements.gpu_memory, 256 * 16);
8383 scheduler.mark_prefill_complete(&first_id, 256);
8384
8385 enqueue_waiting(
8386 &scheduler,
8387 create_test_request_with_prompt_tokens(Priority::Normal, 256),
8388 );
8389 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
8390 assert_eq!(mixed_batch.requests.len(), 2);
8391 assert_eq!(mixed_batch.resource_requirements.gpu_memory, (1 + 64) * 16);
8392 }
8393
8394 #[test]
8395 fn active_decode_prefill_chunk_caps_aggregate_mixed_prefill_tokens() {
8396 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8397 prompt_token_estimate: true,
8398 active_decode_prefill_chunk: Some(64),
8399 ..SchedulerConfig::default()
8400 });
8401 let hint = BatchHint {
8402 max_batch_size: 8,
8403 max_tokens: 2048,
8404 target_latency_ms: None,
8405 available_memory: None,
8406 resource_constraints: Default::default(),
8407 };
8408
8409 let first = create_test_request_with_prompt_tokens(Priority::Normal, 256);
8410 let first_id = first.id.clone();
8411 enqueue_waiting(&scheduler, first);
8412 let initial_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8413 assert_eq!(initial_batch.requests.len(), 1);
8414 scheduler.mark_prefill_complete(&first_id, 256);
8415
8416 for _ in 0..4 {
8417 enqueue_waiting(
8418 &scheduler,
8419 create_test_request_with_prompt_tokens(Priority::Normal, 256),
8420 );
8421 }
8422
8423 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
8424 assert_eq!(
8425 mixed_batch.requests.len(),
8426 5,
8427 "low decode pressure should admit more prefill chunks from batch headroom"
8428 );
8429 assert_eq!(mixed_batch.resource_requirements.gpu_memory, (1 + 256) * 16);
8430 assert_eq!(
8431 scheduler.prefilling_count(),
8432 4,
8433 "waiting requests may be promoted, but scheduling must respect the mixed-prefill budget"
8434 );
8435 }
8436
8437 #[test]
8438 fn active_decode_prefill_budget_scales_down_with_decode_pressure() {
8439 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8440 prompt_token_estimate: true,
8441 active_decode_prefill_chunk: Some(64),
8442 ..SchedulerConfig::default()
8443 });
8444 let hint = BatchHint {
8445 max_batch_size: 8,
8446 max_tokens: 2048,
8447 target_latency_ms: None,
8448 available_memory: None,
8449 resource_constraints: Default::default(),
8450 };
8451
8452 let mut decode_ids = Vec::new();
8453 for _ in 0..6 {
8454 let request = create_test_request_with_prompt_tokens(Priority::Normal, 128);
8455 decode_ids.push(request.id.clone());
8456 enqueue_waiting(&scheduler, request);
8457 }
8458 let initial_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8459 assert_eq!(initial_batch.requests.len(), 6);
8460 for id in &decode_ids {
8461 scheduler.mark_prefill_complete(id, 128);
8462 }
8463
8464 for _ in 0..4 {
8465 enqueue_waiting(
8466 &scheduler,
8467 create_test_request_with_prompt_tokens(Priority::Normal, 256),
8468 );
8469 }
8470
8471 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
8472 assert_eq!(
8473 mixed_batch.requests.len(),
8474 8,
8475 "high decode pressure should admit bounded partial prefills up to available slots"
8476 );
8477 assert_eq!(mixed_batch.resource_requirements.gpu_memory, (6 + 128) * 16);
8478 }
8479
8480 #[test]
8481 fn active_decode_prefill_budget_caps_small_final_chunks_by_count() {
8482 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8483 max_running_requests: 32,
8484 prompt_token_estimate: true,
8485 prefill_step_chunk: Some(6),
8486 ..SchedulerConfig::default()
8487 });
8488 let hint = BatchHint {
8489 max_batch_size: 32,
8490 max_tokens: 192,
8491 target_latency_ms: None,
8492 available_memory: None,
8493 resource_constraints: Default::default(),
8494 };
8495
8496 let mut decode_ids = Vec::new();
8497 for _ in 0..19 {
8498 let request = create_test_request_with_prompt_tokens(Priority::Normal, 1);
8499 decode_ids.push(request.id.clone());
8500 enqueue_waiting(&scheduler, request);
8501 }
8502 let initial_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8503 assert_eq!(initial_batch.requests.len(), 19);
8504 for id in &decode_ids {
8505 scheduler.mark_prefill_complete(id, 1);
8506 }
8507
8508 for _ in 0..13 {
8509 enqueue_waiting(
8510 &scheduler,
8511 create_test_request_with_prompt_tokens(Priority::Normal, 1),
8512 );
8513 }
8514
8515 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
8516 let prefill_tokens: Vec<_> = mixed_batch
8517 .requests
8518 .iter()
8519 .filter(|request| !decode_ids.contains(&request.request.id))
8520 .map(|request| request.tokens_to_process)
8521 .collect();
8522
8523 assert_eq!(
8524 mixed_batch.requests.len(),
8525 23,
8526 "small final prefill chunks must not bypass the mixed-prefill count budget"
8527 );
8528 assert_eq!(prefill_tokens, vec![Some(1), Some(1), Some(1), Some(1)]);
8529 }
8530
8531 #[test]
8532 fn active_decode_prefill_budget_uses_effective_step_chunk_for_aggregate_cap() {
8533 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig {
8534 max_running_requests: 32,
8535 prompt_token_estimate: true,
8536 active_decode_prefill_chunk: Some(8192),
8537 prefill_step_chunk: Some(64),
8538 ..SchedulerConfig::default()
8539 });
8540 let hint = BatchHint {
8541 max_batch_size: 32,
8542 max_tokens: 8192,
8543 target_latency_ms: None,
8544 available_memory: None,
8545 resource_constraints: Default::default(),
8546 };
8547
8548 let mut decode_ids = Vec::new();
8549 for _ in 0..7 {
8550 let request = create_test_request_with_prompt_tokens(Priority::Normal, 128);
8551 decode_ids.push(request.id.clone());
8552 enqueue_waiting(&scheduler, request);
8553 }
8554 let initial_batch = scheduler.create_iteration_batch(hint.clone()).unwrap();
8555 assert_eq!(initial_batch.requests.len(), 7);
8556 for id in &decode_ids {
8557 scheduler.mark_prefill_complete(id, 128);
8558 }
8559
8560 for _ in 0..25 {
8561 enqueue_waiting(
8562 &scheduler,
8563 create_test_request_with_prompt_tokens(Priority::Normal, 256),
8564 );
8565 }
8566
8567 let mixed_batch = scheduler.create_iteration_batch(hint).unwrap();
8568 let prefill_tokens: Vec<_> = mixed_batch
8569 .requests
8570 .iter()
8571 .filter(|request| request.tokens_to_process != Some(1))
8572 .map(|request| request.tokens_to_process)
8573 .collect();
8574 assert_eq!(
8575 mixed_batch.requests.len(),
8576 11,
8577 "large explicit active chunks must not bypass the prefill-step aggregate cap"
8578 );
8579 assert_eq!(prefill_tokens, vec![Some(64), Some(64), Some(64), Some(64)]);
8580 assert_eq!(mixed_batch.resource_requirements.gpu_memory, (7 + 256) * 16);
8581 }
8582
8583 #[tokio::test]
8584 async fn test_cancel_waiting() {
8585 let config = SchedulerConfig::default();
8586 let scheduler = ContinuousBatchScheduler::new(config);
8587
8588 let request = create_test_request(Priority::Normal);
8589 let id = request.id.clone();
8590 scheduler.submit(request).await.unwrap();
8591
8592 assert_eq!(scheduler.waiting_count(), 1);
8593
8594 let result = scheduler.cancel(id).await.unwrap();
8595 assert!(result);
8596 assert_eq!(scheduler.waiting_count(), 0);
8597 }
8598
8599 #[tokio::test]
8600 async fn test_metrics() {
8601 let config = SchedulerConfig::default();
8602 let scheduler = ContinuousBatchScheduler::new(config);
8603
8604 scheduler
8605 .submit(create_test_request(Priority::Normal))
8606 .await
8607 .unwrap();
8608
8609 let metrics = scheduler.metrics();
8610 assert_eq!(metrics.waiting_requests, 1);
8611 }
8612
8613 #[tokio::test]
8614 async fn metrics_track_queue_wait_time_on_admission() {
8615 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
8616 scheduler
8617 .submit(create_test_request(Priority::Normal))
8618 .await
8619 .unwrap();
8620 std::thread::sleep(std::time::Duration::from_millis(5));
8621
8622 let batch = scheduler.next_batch(BatchHint::simple(1)).await;
8623 assert!(batch.is_some());
8624
8625 let metrics = scheduler.metrics();
8626 assert_eq!(metrics.waiting_requests, 0);
8627 assert_eq!(metrics.running_requests, 1);
8628 assert!(
8629 metrics.avg_wait_time_ms >= 1.0,
8630 "expected non-zero wait time, got {}",
8631 metrics.avg_wait_time_ms
8632 );
8633 }
8634
8635 #[test]
8636 fn test_cb_request_states() {
8637 let request = create_test_request(Priority::Normal);
8638 let cb_req = ContinuousBatchRequest::new(request);
8639
8640 assert_eq!(cb_req.phase, RequestPhase::Waiting);
8641 assert!(!cb_req.is_active());
8642 assert!(!cb_req.is_finished());
8643 }
8644
8645 #[tokio::test]
8648 async fn chunked_prefill_advances_across_iterations() {
8649 let cb_cfg = ContinuousBatchConfig {
8650 enable_chunked_prefill: true,
8651 prefill_chunk_size: 128,
8652 ..ContinuousBatchConfig::default()
8653 };
8654 let scheduler =
8655 ContinuousBatchScheduler::with_cb_config(SchedulerConfig::default(), cb_cfg);
8656
8657 let request = create_test_request(Priority::Normal);
8658 let req_id = request.id.clone();
8659 scheduler.submit(request).await.unwrap();
8660
8661 let _ = scheduler.next_batch(BatchHint::simple(1024)).await;
8663 assert_eq!(scheduler.prefilling_count(), 1);
8664 assert_eq!(scheduler.decoding_count(), 0);
8665
8666 let done = scheduler.mark_prefill_chunk_processed(&req_id, 400, 128);
8669 assert!(!done, "first chunk should not finish prefill");
8670 assert_eq!(scheduler.prefilling_count(), 1);
8671 assert_eq!(scheduler.decoding_count(), 0);
8672
8673 let done = scheduler.mark_prefill_chunk_processed(&req_id, 400, 128);
8675 assert!(!done);
8676 assert_eq!(scheduler.prefilling_count(), 1);
8677 assert_eq!(scheduler.decoding_count(), 0);
8678
8679 let done = scheduler.mark_prefill_chunk_processed(&req_id, 400, 200);
8681 assert!(done, "last chunk should complete prefill");
8682 assert_eq!(scheduler.prefilling_count(), 0);
8683 assert_eq!(scheduler.decoding_count(), 1);
8684 }
8685
8686 #[tokio::test]
8690 async fn mark_prefill_complete_sets_offset_to_total() {
8691 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
8692 let request = create_test_request(Priority::Normal);
8693 let req_id = request.id.clone();
8694 scheduler.submit(request).await.unwrap();
8695 let _ = scheduler.next_batch(BatchHint::simple(1024)).await;
8696
8697 scheduler.mark_prefill_complete(&req_id, 256);
8698
8699 assert_eq!(scheduler.prefilling_count(), 0);
8700 assert_eq!(scheduler.decoding_count(), 1);
8701 }
8702
8703 fn activate_decode_requests(
8704 scheduler: &ContinuousBatchScheduler,
8705 count: usize,
8706 ) -> Vec<RequestId> {
8707 let mut request_ids = Vec::with_capacity(count);
8708 for _ in 0..count {
8709 let request = create_test_request_with_prompt_tokens(Priority::Normal, 1);
8710 request_ids.push(request.id.clone());
8711 enqueue_waiting(scheduler, request);
8712 }
8713 let batch = scheduler
8714 .create_iteration_batch(BatchHint::simple(count.max(1)))
8715 .expect("test requests should enter prefill");
8716 assert_eq!(batch.requests.len(), count);
8717 for request_id in &request_ids {
8718 scheduler.mark_prefill_complete(request_id, 1);
8719 }
8720 request_ids
8721 }
8722
8723 #[test]
8724 fn execution_readiness_blocks_only_the_exact_frontier() {
8725 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
8726 let request_ids = activate_decode_requests(&scheduler, 2);
8727 let receipt = scheduler
8728 .defer_for_execution_readiness(std::slice::from_ref(&request_ids[0]))
8729 .unwrap();
8730
8731 let batch = scheduler
8732 .create_iteration_batch(BatchHint::simple(2))
8733 .expect("unrelated decode must remain runnable");
8734 let scheduled = batch
8735 .requests
8736 .iter()
8737 .map(|request| &request.request.id)
8738 .collect::<Vec<_>>();
8739 assert!(!scheduled.contains(&&request_ids[0]));
8740 assert!(scheduled.contains(&&request_ids[1]));
8741 assert!(!scheduler.all_active_execution_readiness_blocked());
8742
8743 assert!(receipt.wake().mark_ready());
8744 let resumed = scheduler
8745 .create_iteration_batch(BatchHint::simple(2))
8746 .expect("ready ticket must authorize an exact reprobe");
8747 assert!(resumed
8748 .requests
8749 .iter()
8750 .any(|request| request.request.id == request_ids[0]));
8751 }
8752
8753 #[test]
8754 fn stale_execution_readiness_wake_cannot_unblock_replacement_ticket() {
8755 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
8756 let request_id = activate_decode_requests(&scheduler, 1).remove(0);
8757 let first = scheduler
8758 .defer_for_execution_readiness(std::slice::from_ref(&request_id))
8759 .unwrap();
8760 let stale = first.wake().clone();
8761 assert!(stale.mark_ready());
8762 assert!(scheduler
8763 .create_iteration_batch(BatchHint::simple(1))
8764 .is_some());
8765
8766 let replacement = scheduler
8767 .defer_for_execution_readiness(std::slice::from_ref(&request_id))
8768 .unwrap();
8769 assert_ne!(stale.ticket_id(), replacement.wake().ticket_id());
8770 assert!(!stale.mark_ready());
8771 assert!(scheduler
8772 .create_iteration_batch(BatchHint::simple(1))
8773 .is_none());
8774 assert!(scheduler.all_active_execution_readiness_blocked());
8775
8776 assert!(replacement.wake().mark_ready());
8777 assert!(scheduler
8778 .create_iteration_batch(BatchHint::simple(1))
8779 .is_some());
8780 }
8781
8782 #[test]
8783 fn failed_execution_readiness_ticket_remains_parked_for_terminal_owner() {
8784 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
8785 let request_id = activate_decode_requests(&scheduler, 1).remove(0);
8786 let receipt = scheduler
8787 .defer_for_execution_readiness(std::slice::from_ref(&request_id))
8788 .unwrap();
8789 assert!(receipt.wake().mark_failed());
8790
8791 assert!(scheduler
8792 .create_iteration_batch(BatchHint::simple(1))
8793 .is_none());
8794 assert!(scheduler.all_active_execution_readiness_blocked());
8795 let snapshot = scheduler.trace_snapshot();
8796 assert_eq!(snapshot.execution_readiness_deferred_total, 1);
8797 assert_eq!(snapshot.execution_readiness_blocked_decode_len, 1);
8798 }
8799
8800 #[test]
8801 fn execution_readiness_install_is_all_or_nothing() {
8802 let scheduler = ContinuousBatchScheduler::new(SchedulerConfig::default());
8803 let request_id = activate_decode_requests(&scheduler, 1).remove(0);
8804 let missing = RequestId::new();
8805 assert!(scheduler
8806 .defer_for_execution_readiness(&[request_id.clone(), missing])
8807 .is_err());
8808
8809 let batch = scheduler
8810 .create_iteration_batch(BatchHint::simple(1))
8811 .expect("failed cohort install must not retain a partial block");
8812 assert_eq!(batch.requests[0].request.id, request_id);
8813 }
8814}