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

ferrum_interfaces/
scheduler.rs
