ferrum_interfaces/
scheduler.rs

1//! Unified scheduler interface with resource awareness and SLA support
2//!
3//! This module provides the unified scheduler interface that replaces the
4//! conflicting scheduler definitions in the original codebase.
5
6use async_trait::async_trait;
7use ferrum_types::{
8    BatchId, InferenceRequest, InferenceResponse, Priority, RequestId, RequestState, Result,
9    SchedulerConfig as TypesSchedulerConfig, SchedulerStats,
10};
11use serde::{Deserialize, Serialize};
12use std::{collections::HashMap, time::Duration};
13
14mod prefix_restore;
15pub use prefix_restore::PreparedPrefixRestore;
16
17/// Main scheduler trait for request management and batching
18#[async_trait]
19pub trait Scheduler: Send + Sync {
20    /// Submit new inference request
21    async fn submit(&self, request: InferenceRequest) -> Result<RequestId>;
22
23    /// Get next batch of requests to execute
24    async fn next_batch(&self, hint: BatchHint) -> Option<BatchPlan>;
25
26    /// Mark request as completed
27    async fn complete(&self, request_id: RequestId, response: &InferenceResponse) -> Result<()>;
28
29    /// Cancel pending request
30    async fn cancel(&self, request_id: RequestId) -> Result<bool>;
31
32    /// Update request priority
33    async fn update_priority(&self, request_id: RequestId, priority: Priority) -> Result<()>;
34
35    /// Get scheduler metrics
36    fn metrics(&self) -> SchedulerMetrics;
37
38    /// Get scheduler configuration
39    fn config(&self) -> &TypesSchedulerConfig;
40
41    /// Get current request state if the request is tracked by scheduler.
42    fn request_state(&self, request_id: &RequestId) -> Option<RequestState> {
43        let _ = request_id;
44        None
45    }
46
47    /// Reserve the current admitted, unscheduled prefill incarnation before
48    /// starting an asynchronous state restore. Unsupported schedulers return
49    /// `None`. Dropping the preparation must release its scheduling hold
50    /// without advancing progress; no scheduler lock may span device work.
51    fn prepare_prefix_restore(
52        &self,
53        _request_id: &RequestId,
54        _expected_offset: usize,
55        _prompt_tokens: usize,
56    ) -> Result<Option<PreparedPrefixRestore>> {
57        Ok(None)
58    }
59
60    /// Conditionally publish the boundary returned by a completed restore.
61    /// Implementations must revalidate the preparation's admission and work
62    /// identities, and require `expected_offset < restored_boundary < prompt`.
63    /// This is state progress, not model execution or capacity-fit feedback.
64    fn commit_prefix_restored(
65        &self,
66        _prepared: PreparedPrefixRestore,
67        _restored_boundary: usize,
68    ) -> Result<()> {
69        Err(ferrum_types::FerrumError::unsupported(
70            "Prefix state restoration is not supported by this scheduler",
71        ))
72    }
73
74    /// Preempt running request (if supported)
75    async fn preempt(&self, _request_id: RequestId) -> Result<PreemptionResult> {
76        // Default implementation: preemption not supported
77        Err(ferrum_types::FerrumError::unsupported(
78            "Preemption not supported",
79        ))
80    }
81
82    /// Resume preempted request
83    async fn resume(&self, _request_id: RequestId) -> Result<()> {
84        // Default implementation: resumption not supported
85        Err(ferrum_types::FerrumError::unsupported(
86            "Resumption not supported",
87        ))
88    }
89}
90
91/// Batch hint for scheduler optimization
92#[derive(Debug, Clone)]
93pub struct BatchHint {
94    /// Maximum batch size
95    pub max_batch_size: usize,
96    /// Maximum total tokens in batch
97    pub max_tokens: usize,
98    /// Target latency for batch formation
99    pub target_latency_ms: Option<u64>,
100    /// Available memory for batch
101    pub available_memory: Option<u64>,
102    /// Resource constraints
103    pub resource_constraints: ResourceConstraints,
104}
105
106impl BatchHint {
107    /// Create simple batch hint with size limit
108    pub fn simple(max_batch_size: usize) -> Self {
109        Self {
110            max_batch_size,
111            max_tokens: max_batch_size * 2048, // Default reasonable token limit
112            target_latency_ms: None,
113            available_memory: None,
114            resource_constraints: ResourceConstraints::default(),
115        }
116    }
117}
118
119/// Resource constraints for scheduling
120#[derive(Debug, Clone, Serialize, Deserialize, Default)]
121pub struct ResourceConstraints {
122    /// Maximum GPU memory usage
123    pub max_gpu_memory: Option<u64>,
124    /// Maximum CPU memory usage
125    pub max_cpu_memory: Option<u64>,
126    /// Maximum recurrent-state memory usage
127    pub max_recurrent_state_bytes: Option<u64>,
128    /// Maximum recurrent-state slots
129    pub max_recurrent_state_slots: Option<usize>,
130    /// Maximum compute units
131    pub max_compute_units: Option<usize>,
132    /// Required device types
133    pub required_devices: Vec<ferrum_types::Device>,
134}
135
136/// Batch execution plan
137#[derive(Debug, Clone)]
138pub struct BatchPlan {
139    /// Unique batch identifier
140    pub batch_id: BatchId,
141    /// Requests included in this batch
142    pub requests: Vec<ScheduledRequest>,
143    /// Maximum sequence length in batch
144    pub max_sequence_length: usize,
145    /// Estimated execution time
146    pub estimated_time_ms: Option<u64>,
147    /// Resource requirements
148    pub resource_requirements: BatchResourceRequirements,
149    /// Batch creation timestamp
150    pub created_at: chrono::DateTime<chrono::Utc>,
151}
152
153impl BatchPlan {
154    /// Get total number of tokens in batch
155    pub fn total_tokens(&self) -> usize {
156        self.requests
157            .iter()
158            .map(|req| {
159                req.tokens_to_process
160                    .unwrap_or(req.request.sampling_params.max_tokens)
161            })
162            .sum()
163    }
164
165    /// Get batch size
166    pub fn size(&self) -> usize {
167        self.requests.len()
168    }
169
170    /// Check if batch is empty
171    pub fn is_empty(&self) -> bool {
172        self.requests.is_empty()
173    }
174
175    /// Get highest priority in batch
176    pub fn max_priority(&self) -> Priority {
177        self.requests
178            .iter()
179            .map(|req| req.request.priority)
180            .max()
181            .unwrap_or(Priority::Low)
182    }
183}
184
185/// Scheduled request with additional metadata
186#[derive(Debug, Clone)]
187pub struct ScheduledRequest {
188    /// Original inference request
189    pub request: InferenceRequest,
190    /// Current scheduling state
191    pub state: RequestState,
192    /// Queue position when waiting
193    pub queue_position: Option<usize>,
194    /// Estimated wait time
195    pub estimated_wait_time: Option<Duration>,
196    /// Number of tokens processed so far
197    pub tokens_processed: usize,
198    /// Number of tokens the engine should process for this request in this batch.
199    ///
200    /// `None` preserves legacy schedulers that did not carry per-request
201    /// token budgets.
202    pub tokens_to_process: Option<usize>,
203    /// Allocated resources
204    pub allocated_resources: AllocatedResources,
205    /// Request submission time
206    pub submitted_at: chrono::DateTime<chrono::Utc>,
207    /// Request start time (when moved from waiting to running)
208    pub started_at: Option<chrono::DateTime<chrono::Utc>>,
209}
210
211impl ScheduledRequest {
212    /// Create new scheduled request
213    pub fn new(request: InferenceRequest) -> Self {
214        Self {
215            request,
216            state: RequestState::Waiting,
217            queue_position: None,
218            estimated_wait_time: None,
219            tokens_processed: 0,
220            tokens_to_process: None,
221            allocated_resources: AllocatedResources::default(),
222            submitted_at: chrono::Utc::now(),
223            started_at: None,
224        }
225    }
226
227    /// Get request age since submission
228    pub fn age(&self) -> Duration {
229        (chrono::Utc::now() - self.submitted_at)
230            .to_std()
231            .unwrap_or_default()
232    }
233
234    /// Get processing time (if started)
235    pub fn processing_time(&self) -> Option<Duration> {
236        self.started_at
237            .map(|start| (chrono::Utc::now() - start).to_std().unwrap_or_default())
238    }
239}
240
241/// Allocated resources for a request
242#[derive(Debug, Clone, Default)]
243pub struct AllocatedResources {
244    /// KV cache blocks allocated
245    pub kv_cache_blocks: Vec<ferrum_types::BlockId>,
246    /// GPU memory allocated (bytes)
247    pub gpu_memory: u64,
248    /// CPU memory allocated (bytes)
249    pub cpu_memory: u64,
250    /// Recurrent-state memory allocated (bytes)
251    pub recurrent_state_bytes: u64,
252    /// Recurrent-state slots allocated
253    pub recurrent_state_slots: usize,
254    /// Compute units reserved
255    pub compute_units: usize,
256}
257
258/// Resource requirements for batch execution
259#[derive(Debug, Clone, Default)]
260pub struct BatchResourceRequirements {
261    /// Required GPU memory
262    pub gpu_memory: u64,
263    /// Required CPU memory
264    pub cpu_memory: u64,
265    /// Required KV cache blocks
266    pub kv_cache_blocks: usize,
267    /// Required recurrent-state memory
268    pub recurrent_state_bytes: u64,
269    /// Required recurrent-state slots
270    pub recurrent_state_slots: usize,
271    /// Required compute units
272    pub compute_units: usize,
273}
274
275/// Preemption result
276#[derive(Debug, Clone)]
277pub struct PreemptionResult {
278    /// Whether preemption was successful
279    pub success: bool,
280    /// Saved state for resumption (if any)
281    pub saved_state: Option<PreemptionState>,
282    /// Resources freed by preemption
283    pub freed_resources: AllocatedResources,
284}
285
286/// State saved during preemption
287#[derive(Debug, Clone)]
288pub struct PreemptionState {
289    /// KV cache checkpoint
290    pub kv_cache_checkpoint: Vec<u8>,
291    /// Number of tokens processed
292    pub tokens_processed: usize,
293    /// Generation state
294    pub generation_state: HashMap<String, serde_json::Value>,
295}
296
297/// Scheduler configuration
298pub type SchedulerConfig = TypesSchedulerConfig;
299
300/// Scheduling policies
301#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
302pub enum SchedulingPolicy {
303    /// First-Come-First-Served
304    FCFS,
305    /// Priority-based scheduling
306    Priority,
307    /// Fair-share scheduling
308    FairShare,
309    /// Shortest-Job-First
310    SJF,
311    /// Resource-aware scheduling
312    ResourceAware,
313    /// SLA-driven scheduling
314    SlaAware,
315}
316
317/// Fair share configuration
318#[derive(Debug, Clone, Serialize, Deserialize)]
319pub struct FairShareConfig {
320    /// Share weights per client
321    pub client_shares: HashMap<String, f32>,
322    /// Default share for unspecified clients
323    pub default_share: f32,
324    /// Share enforcement strictness (0.0 - 1.0)
325    pub enforcement_strictness: f32,
326}
327
328/// SLA configuration
329#[derive(Debug, Clone, Serialize, Deserialize)]
330pub struct SlaConfig {
331    /// Enable SLA enforcement
332    pub enabled: bool,
333    /// Default SLA for requests without specific SLA
334    pub default_sla: SlaRequirements,
335    /// Per-client SLA overrides
336    pub client_slas: HashMap<String, SlaRequirements>,
337}
338
339/// SLA requirements
340#[derive(Debug, Clone, Serialize, Deserialize)]
341pub struct SlaRequirements {
342    /// Maximum latency (P95)
343    pub max_latency_p95_ms: u64,
344    /// Maximum latency (P99)
345    pub max_latency_p99_ms: u64,
346    /// Minimum throughput
347    pub min_throughput_rps: f32,
348    /// Availability requirement
349    pub availability_percent: f32,
350}
351
352/// Resource limits
353#[derive(Debug, Clone, Serialize, Deserialize, Default)]
354pub struct ResourceLimits {
355    /// Maximum total GPU memory
356    pub max_gpu_memory: Option<u64>,
357    /// Maximum total CPU memory
358    pub max_cpu_memory: Option<u64>,
359    /// Maximum KV cache blocks
360    pub max_kv_cache_blocks: Option<usize>,
361    /// Maximum recurrent-state memory
362    pub max_recurrent_state_bytes: Option<u64>,
363    /// Maximum recurrent-state slots
364    pub max_recurrent_state_slots: Option<usize>,
365    /// Per-client resource limits
366    pub per_client_limits: HashMap<String, ClientResourceLimits>,
367}
368
369/// Per-client resource limits
370#[derive(Debug, Clone, Serialize, Deserialize)]
371pub struct ClientResourceLimits {
372    /// Max concurrent requests per client
373    pub max_concurrent_requests: usize,
374    /// Max GPU memory per client
375    pub max_gpu_memory: Option<u64>,
376    /// Max recurrent-state memory per client
377    pub max_recurrent_state_bytes: Option<u64>,
378    /// Max requests per minute
379    pub max_requests_per_minute: Option<u32>,
380}
381
382pub type SchedulerMetrics = SchedulerStats;
383
384/// Advanced scheduler capabilities
385#[async_trait]
386pub trait AdvancedScheduler: Scheduler {
387    /// Enable resource-aware scheduling
388    async fn enable_resource_awareness(&mut self, config: ResourceAwarenessConfig) -> Result<()>;
389
390    /// Set custom admission policy
391    async fn set_admission_policy(&mut self, policy: Box<dyn AdmissionPolicy>) -> Result<()>;
392
393    /// Configure dynamic batching
394    async fn configure_dynamic_batching(&mut self, config: DynamicBatchingConfig) -> Result<()>;
395
396    /// Get detailed queue analysis
397    fn queue_analysis(&self) -> QueueAnalysis;
398
399    /// Simulate scheduling for capacity planning
400    async fn simulate_load(
401        &self,
402        workload: &SimulatedWorkload,
403    ) -> Result<SchedulingSimulationResult>;
404}
405
406/// Resource awareness configuration
407#[derive(Debug, Clone, Serialize, Deserialize)]
408pub struct ResourceAwarenessConfig {
409    /// Enable memory-aware scheduling
410    pub enable_memory_awareness: bool,
411    /// Enable compute-aware scheduling
412    pub enable_compute_awareness: bool,
413    /// Resource prediction horizon
414    pub prediction_horizon_ms: u64,
415    /// Resource safety margin (0.0 - 1.0)
416    pub safety_margin: f32,
417}
418
419/// Admission policy for request acceptance
420pub trait AdmissionPolicy: Send + Sync {
421    /// Decide whether to admit a request
422    fn should_admit(
423        &self,
424        request: &InferenceRequest,
425        current_metrics: &SchedulerMetrics,
426    ) -> AdmissionDecision;
427
428    /// Get policy name
429    fn name(&self) -> &str;
430}
431
432/// Admission decision
433#[derive(Debug, Clone)]
434pub enum AdmissionDecision {
435    /// Accept the request
436    Accept,
437    /// Reject the request with reason
438    Reject(String),
439    /// Accept but suggest delay
440    AcceptWithDelay(Duration),
441}
442
443/// Dynamic batching configuration
444#[derive(Debug, Clone, Serialize, Deserialize)]
445pub struct DynamicBatchingConfig {
446    /// Minimum batch size
447    pub min_batch_size: usize,
448    /// Maximum batch size
449    pub max_batch_size: usize,
450    /// Batch formation timeout
451    pub batch_timeout_ms: u64,
452    /// Enable adaptive batch sizing
453    pub enable_adaptive_sizing: bool,
454    /// Target batch utilization
455    pub target_utilization: f32,
456}
457
458/// Queue analysis results
459#[derive(Debug, Clone)]
460pub struct QueueAnalysis {
461    /// Queue depth over time
462    pub queue_depth_history: Vec<(chrono::DateTime<chrono::Utc>, usize)>,
463    /// Wait time distribution
464    pub wait_time_distribution: WaitTimeDistribution,
465    /// Request pattern analysis
466    pub request_patterns: RequestPatternAnalysis,
467    /// Bottleneck identification
468    pub bottlenecks: Vec<BottleneckAnalysis>,
469}
470
471/// Wait time distribution
472#[derive(Debug, Clone)]
473pub struct WaitTimeDistribution {
474    /// P50 wait time
475    pub p50_ms: f64,
476    /// P95 wait time
477    pub p95_ms: f64,
478    /// P99 wait time
479    pub p99_ms: f64,
480    /// Maximum wait time
481    pub max_ms: f64,
482    /// Average wait time
483    pub mean_ms: f64,
484}
485
486/// Request pattern analysis
487#[derive(Debug, Clone)]
488pub struct RequestPatternAnalysis {
489    /// Peak request times
490    pub peak_times: Vec<chrono::DateTime<chrono::Utc>>,
491    /// Request rate trend
492    pub rate_trend: RateTrend,
493    /// Seasonality patterns
494    pub seasonality: SeasonalityPattern,
495}
496
497/// Request rate trend
498#[derive(Debug, Clone, Copy)]
499pub enum RateTrend {
500    Increasing,
501    Decreasing,
502    Stable,
503    Volatile,
504}
505
506/// Seasonality patterns
507#[derive(Debug, Clone)]
508pub struct SeasonalityPattern {
509    /// Hourly patterns
510    pub hourly_pattern: Vec<f32>,
511    /// Daily patterns  
512    pub daily_pattern: Vec<f32>,
513    /// Weekly patterns
514    pub weekly_pattern: Vec<f32>,
515}
516
517/// Bottleneck analysis
518#[derive(Debug, Clone)]
519pub struct BottleneckAnalysis {
520    /// Bottleneck type
521    pub bottleneck_type: BottleneckType,
522    /// Severity (0.0 - 1.0)
523    pub severity: f32,
524    /// Description
525    pub description: String,
526    /// Suggested mitigation
527    pub mitigation: String,
528}
529
530/// Types of bottlenecks
531#[derive(Debug, Clone, Copy)]
532pub enum BottleneckType {
533    /// Memory bottleneck
534    Memory,
535    /// Compute bottleneck
536    Compute,
537    /// I/O bottleneck
538    IO,
539    /// Scheduling bottleneck
540    Scheduling,
541    /// Network bottleneck
542    Network,
543}
544
545/// Simulated workload for capacity planning
546#[derive(Debug, Clone)]
547pub struct SimulatedWorkload {
548    /// Request arrival pattern
549    pub arrival_pattern: ArrivalPattern,
550    /// Request size distribution
551    pub size_distribution: SizeDistribution,
552    /// Simulation duration
553    pub duration_seconds: u64,
554}
555
556/// Request arrival patterns
557#[derive(Debug, Clone)]
558pub enum ArrivalPattern {
559    /// Constant rate
560    Constant { rate_rps: f32 },
561    /// Poisson process
562    Poisson { lambda: f32 },
563    /// Bursty pattern
564    Bursty {
565        burst_rate: f32,
566        quiet_rate: f32,
567        burst_duration_s: f32,
568    },
569    /// Seasonal pattern
570    Seasonal {
571        base_rate: f32,
572        peaks: Vec<(f32, f32)>,
573    }, // (time, multiplier)
574}
575
576/// Request size distribution
577#[derive(Debug, Clone)]
578pub enum SizeDistribution {
579    /// Fixed size
580    Fixed { tokens: usize },
581    /// Uniform distribution
582    Uniform {
583        min_tokens: usize,
584        max_tokens: usize,
585    },
586    /// Normal distribution
587    Normal { mean: f32, std_dev: f32 },
588    /// Log-normal distribution
589    LogNormal { mu: f32, sigma: f32 },
590}
591
592/// Scheduling simulation results
593#[derive(Debug, Clone)]
594pub struct SchedulingSimulationResult {
595    /// Total requests processed
596    pub total_requests: u64,
597    /// Successful requests
598    pub successful_requests: u64,
599    /// Failed/rejected requests
600    pub failed_requests: u64,
601    /// Average latency
602    pub avg_latency_ms: f64,
603    /// P95 latency
604    pub p95_latency_ms: f64,
605    /// P99 latency
606    pub p99_latency_ms: f64,
607    /// Throughput achieved
608    pub throughput_rps: f32,
609    /// Resource utilization (optional placeholder)
610    pub resource_utilization: Option<ResourceStats>,
611    /// Predicted bottlenecks
612    pub bottlenecks: Vec<BottleneckAnalysis>,
613}
614
615#[derive(Debug, Clone, Serialize, Deserialize, Default)]
616pub struct ResourceStats {
617    pub gpu_memory_bytes: Option<u64>,
618    pub cpu_memory_bytes: Option<u64>,
619    pub compute_utilization: Option<f32>,
620}
ferrum_interfaces/scheduler.rs

ferrum_interfaces/
scheduler.rs
