forge_orchestration/scheduler/
mod.rs

1//! Distributed Scheduler for Forge Orchestration
2//!
3//! A full-featured scheduler comparable to Kubernetes, with:
4//! - Bin-packing and spread scheduling algorithms
5//! - Affinity/anti-affinity rules
6//! - Resource-aware placement (CPU, memory, GPU)
7//! - Preemption and priority-based scheduling
8//! - Topology-aware scheduling for NUMA/GPU locality
9
10pub mod algorithms;
11pub mod deadline;
12pub mod gang;
13pub mod optimized;
14pub mod placement;
15pub mod preemption;
16pub mod queue;
17pub mod reconcile;
18pub mod sim;
19
20use std::collections::HashMap;
21use std::sync::Arc;
22use parking_lot::RwLock;
23use serde::{Deserialize, Serialize};
24use tracing::{debug, info, warn};
25
26use crate::types::{NodeId, GpuResources};
27
28pub use algorithms::{SchedulingAlgorithm, BinPackScheduler, SpreadScheduler, GpuLocalityScheduler, LearnedScheduler, SchedulingFeedback};
29pub use deadline::{DeadlineEntry, DeadlineQueue, Eligibility, MissPolicy, TickDeadlineScheduler, TickOutcome};
30pub use gang::{GangDecision, GangReservation, GangScheduler};
31pub use reconcile::{Assignment, MetricsSource, ReconcileReport, Reconciler, TaskStatus};
32pub use optimized::{OptimizedScheduler, WorkloadBatch, SchedulerStats, ClusterUtilization, FFDBinPacker};
33pub use placement::{PlacementConstraint, AffinityRule, Affinity};
34pub use preemption::{PreemptionPolicy, PriorityClass};
35pub use queue::{SchedulingQueue, QueuedWorkload};
36pub use sim::{AgentPolicy, CoPlacement, GangGroup, MemberRole, Region3D, SimCell, SimMember, SimWorld};
37
38/// Resource requirements for a workload
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct ResourceRequirements {
41    /// CPU cores requested (millicores, 1000 = 1 core)
42    pub cpu_millis: u64,
43    /// Memory requested in MB
44    pub memory_mb: u64,
45    /// GPU count requested
46    pub gpu_count: u32,
47    /// GPU memory required per GPU in MB
48    pub gpu_memory_mb: u64,
49    /// Ephemeral storage in MB
50    pub storage_mb: u64,
51    /// Network bandwidth in Mbps
52    pub network_mbps: u32,
53}
54
55impl Default for ResourceRequirements {
56    fn default() -> Self {
57        Self {
58            cpu_millis: 100,
59            memory_mb: 128,
60            gpu_count: 0,
61            gpu_memory_mb: 0,
62            storage_mb: 0,
63            network_mbps: 0,
64        }
65    }
66}
67
68impl ResourceRequirements {
69    /// Create new resource requirements
70    pub fn new() -> Self {
71        Self::default()
72    }
73
74    /// Set CPU requirement
75    pub fn cpu(mut self, millis: u64) -> Self {
76        self.cpu_millis = millis;
77        self
78    }
79
80    /// Set memory requirement
81    pub fn memory(mut self, mb: u64) -> Self {
82        self.memory_mb = mb;
83        self
84    }
85
86    /// Set GPU requirements
87    pub fn gpu(mut self, count: u32, memory_mb: u64) -> Self {
88        self.gpu_count = count;
89        self.gpu_memory_mb = memory_mb;
90        self
91    }
92}
93
94/// Node resources and capacity
95#[derive(Debug, Clone, Serialize, Deserialize)]
96pub struct NodeResources {
97    /// Node identifier
98    pub node_id: NodeId,
99    /// Total CPU capacity (millicores)
100    pub cpu_capacity: u64,
101    /// Allocated CPU (millicores)
102    pub cpu_allocated: u64,
103    /// Total memory capacity (MB)
104    pub memory_capacity: u64,
105    /// Allocated memory (MB)
106    pub memory_allocated: u64,
107    /// GPU resources
108    pub gpus: Vec<GpuResources>,
109    /// GPUs allocated (by device ID)
110    pub gpus_allocated: Vec<u32>,
111    /// Node labels for affinity matching
112    pub labels: HashMap<String, String>,
113    /// Node taints
114    pub taints: Vec<Taint>,
115    /// Is node schedulable
116    pub schedulable: bool,
117    /// Node conditions
118    pub conditions: Vec<NodeCondition>,
119}
120
121impl NodeResources {
122    /// Create new node resources
123    pub fn new(node_id: NodeId, cpu_capacity: u64, memory_capacity: u64) -> Self {
124        Self {
125            node_id,
126            cpu_capacity,
127            cpu_allocated: 0,
128            memory_capacity,
129            memory_allocated: 0,
130            gpus: Vec::new(),
131            gpus_allocated: Vec::new(),
132            labels: HashMap::new(),
133            taints: Vec::new(),
134            schedulable: true,
135            conditions: Vec::new(),
136        }
137    }
138
139    /// Add GPU to node
140    pub fn with_gpu(mut self, gpu: GpuResources) -> Self {
141        self.gpus.push(gpu);
142        self
143    }
144
145    /// Add label
146    pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
147        self.labels.insert(key.into(), value.into());
148        self
149    }
150
151    /// Add taint
152    pub fn with_taint(mut self, taint: Taint) -> Self {
153        self.taints.push(taint);
154        self
155    }
156
157    /// Available CPU
158    pub fn cpu_available(&self) -> u64 {
159        self.cpu_capacity.saturating_sub(self.cpu_allocated)
160    }
161
162    /// Available memory
163    pub fn memory_available(&self) -> u64 {
164        self.memory_capacity.saturating_sub(self.memory_allocated)
165    }
166
167    /// Available GPUs
168    pub fn gpus_available(&self) -> usize {
169        self.gpus.len() - self.gpus_allocated.len()
170    }
171
172    /// Check if node can fit workload
173    pub fn can_fit(&self, req: &ResourceRequirements) -> bool {
174        if !self.schedulable {
175            return false;
176        }
177
178        if self.cpu_available() < req.cpu_millis {
179            return false;
180        }
181
182        if self.memory_available() < req.memory_mb {
183            return false;
184        }
185
186        if req.gpu_count > 0 {
187            let available_gpus: Vec<_> = self.gpus.iter()
188                .filter(|g| !self.gpus_allocated.contains(&g.device_id))
189                .filter(|g| g.available_memory_mb() >= req.gpu_memory_mb)
190                .collect();
191            
192            if available_gpus.len() < req.gpu_count as usize {
193                return false;
194            }
195        }
196
197        true
198    }
199
200    /// Allocate resources for workload
201    pub fn allocate(&mut self, req: &ResourceRequirements) -> bool {
202        if !self.can_fit(req) {
203            return false;
204        }
205
206        self.cpu_allocated += req.cpu_millis;
207        self.memory_allocated += req.memory_mb;
208
209        // Allocate GPUs
210        for _ in 0..req.gpu_count {
211            if let Some(gpu) = self.gpus.iter()
212                .find(|g| !self.gpus_allocated.contains(&g.device_id) 
213                    && g.available_memory_mb() >= req.gpu_memory_mb)
214            {
215                self.gpus_allocated.push(gpu.device_id);
216            }
217        }
218
219        true
220    }
221
222    /// Release resources
223    pub fn release(&mut self, req: &ResourceRequirements, gpu_ids: &[u32]) {
224        self.cpu_allocated = self.cpu_allocated.saturating_sub(req.cpu_millis);
225        self.memory_allocated = self.memory_allocated.saturating_sub(req.memory_mb);
226        self.gpus_allocated.retain(|id| !gpu_ids.contains(id));
227    }
228}
229
230/// Node taint for scheduling constraints
231#[derive(Debug, Clone, Serialize, Deserialize)]
232pub struct Taint {
233    /// Taint key
234    pub key: String,
235    /// Taint value
236    pub value: String,
237    /// Taint effect
238    pub effect: TaintEffect,
239}
240
241/// Taint effect
242#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
243pub enum TaintEffect {
244    /// Do not schedule new workloads
245    NoSchedule,
246    /// Prefer not to schedule
247    PreferNoSchedule,
248    /// Evict existing workloads
249    NoExecute,
250}
251
252/// Node condition
253#[derive(Debug, Clone, Serialize, Deserialize)]
254pub struct NodeCondition {
255    /// Condition type
256    pub condition_type: String,
257    /// Condition status
258    pub status: bool,
259    /// Last transition time
260    pub last_transition: chrono::DateTime<chrono::Utc>,
261    /// Reason for condition
262    pub reason: Option<String>,
263    /// Human-readable message
264    pub message: Option<String>,
265}
266
267/// Workload to be scheduled
268#[derive(Debug, Clone, Serialize, Deserialize)]
269pub struct Workload {
270    /// Unique workload ID
271    pub id: String,
272    /// Workload name
273    pub name: String,
274    /// Namespace
275    pub namespace: String,
276    /// Resource requirements
277    pub resources: ResourceRequirements,
278    /// Priority (higher = more important)
279    pub priority: i32,
280    /// Priority class name
281    pub priority_class: Option<String>,
282    /// Placement constraints
283    pub constraints: Vec<PlacementConstraint>,
284    /// Node affinity rules
285    pub affinity: Option<Affinity>,
286    /// Tolerations for taints
287    pub tolerations: Vec<Toleration>,
288    /// Preemption policy
289    pub preemption_policy: PreemptionPolicy,
290    /// Creation timestamp
291    pub created_at: chrono::DateTime<chrono::Utc>,
292}
293
294impl Workload {
295    /// Create a new workload
296    pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
297        Self {
298            id: id.into(),
299            name: name.into(),
300            namespace: "default".to_string(),
301            resources: ResourceRequirements::default(),
302            priority: 0,
303            priority_class: None,
304            constraints: Vec::new(),
305            affinity: None,
306            tolerations: Vec::new(),
307            preemption_policy: PreemptionPolicy::PreemptLowerPriority,
308            created_at: chrono::Utc::now(),
309        }
310    }
311
312    /// Set resource requirements
313    pub fn with_resources(mut self, resources: ResourceRequirements) -> Self {
314        self.resources = resources;
315        self
316    }
317
318    /// Set priority
319    pub fn with_priority(mut self, priority: i32) -> Self {
320        self.priority = priority;
321        self
322    }
323
324    /// Add constraint
325    pub fn with_constraint(mut self, constraint: PlacementConstraint) -> Self {
326        self.constraints.push(constraint);
327        self
328    }
329
330    /// Set affinity
331    pub fn with_affinity(mut self, affinity: Affinity) -> Self {
332        self.affinity = Some(affinity);
333        self
334    }
335
336    /// Add toleration
337    pub fn with_toleration(mut self, toleration: Toleration) -> Self {
338        self.tolerations.push(toleration);
339        self
340    }
341}
342
343/// Toleration for node taints
344#[derive(Debug, Clone, Serialize, Deserialize)]
345pub struct Toleration {
346    /// Key to match
347    pub key: Option<String>,
348    /// Operator for matching
349    pub operator: TolerationOperator,
350    /// Value to match
351    pub value: Option<String>,
352    /// Effect to tolerate
353    pub effect: Option<TaintEffect>,
354    /// Toleration seconds for NoExecute
355    pub toleration_seconds: Option<u64>,
356}
357
358/// Toleration operator
359#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
360pub enum TolerationOperator {
361    /// Key must equal value
362    Equal,
363    /// Key must exist
364    Exists,
365}
366
367/// Scheduling decision
368#[derive(Debug, Clone)]
369pub struct SchedulingDecision {
370    /// Workload ID
371    pub workload_id: String,
372    /// Selected node
373    pub node_id: Option<NodeId>,
374    /// Score for the selected node
375    pub score: f64,
376    /// Reason for decision
377    pub reason: String,
378    /// Preempted workloads (if any)
379    pub preempted: Vec<String>,
380    /// Scheduling latency
381    pub latency_ms: u64,
382}
383
384/// The main scheduler
385pub struct Scheduler {
386    /// Registered nodes
387    nodes: Arc<RwLock<HashMap<NodeId, NodeResources>>>,
388    /// Scheduling queue
389    queue: Arc<SchedulingQueue>,
390    /// Scheduling algorithm
391    algorithm: Arc<dyn SchedulingAlgorithm + Send + Sync>,
392    /// Scheduled workloads (workload_id -> node_id)
393    assignments: Arc<RwLock<HashMap<String, NodeId>>>,
394    /// Priority classes
395    priority_classes: Arc<RwLock<HashMap<String, PriorityClass>>>,
396}
397
398impl Scheduler {
399    /// Create a new scheduler with default algorithm
400    pub fn new() -> Self {
401        Self {
402            nodes: Arc::new(RwLock::new(HashMap::new())),
403            queue: Arc::new(SchedulingQueue::new()),
404            algorithm: Arc::new(BinPackScheduler::new()),
405            assignments: Arc::new(RwLock::new(HashMap::new())),
406            priority_classes: Arc::new(RwLock::new(HashMap::new())),
407        }
408    }
409
410    /// Create with specific algorithm
411    pub fn with_algorithm<A: SchedulingAlgorithm + Send + Sync + 'static>(algorithm: A) -> Self {
412        Self {
413            nodes: Arc::new(RwLock::new(HashMap::new())),
414            queue: Arc::new(SchedulingQueue::new()),
415            algorithm: Arc::new(algorithm),
416            assignments: Arc::new(RwLock::new(HashMap::new())),
417            priority_classes: Arc::new(RwLock::new(HashMap::new())),
418        }
419    }
420
421    /// Register a node
422    pub fn register_node(&self, node: NodeResources) {
423        info!(node_id = %node.node_id, "Registering node");
424        self.nodes.write().insert(node.node_id, node);
425    }
426
427    /// Unregister a node
428    pub fn unregister_node(&self, node_id: &NodeId) {
429        info!(node_id = %node_id, "Unregistering node");
430        self.nodes.write().remove(node_id);
431    }
432
433    /// Update node resources
434    pub fn update_node(&self, node: NodeResources) {
435        self.nodes.write().insert(node.node_id, node);
436    }
437
438    /// Get node count
439    pub fn node_count(&self) -> usize {
440        self.nodes.read().len()
441    }
442
443    /// Submit workload for scheduling
444    pub fn submit(&self, workload: Workload) {
445        debug!(workload_id = %workload.id, "Submitting workload");
446        self.queue.enqueue(workload);
447    }
448
449    /// Schedule next workload from queue
450    pub fn schedule_next(&self) -> Option<SchedulingDecision> {
451        let workload = self.queue.dequeue()?;
452        Some(self.schedule(&workload))
453    }
454
455    /// Schedule a specific workload
456    pub fn schedule(&self, workload: &Workload) -> SchedulingDecision {
457        let start = std::time::Instant::now();
458        let nodes = self.nodes.read();
459
460        // Filter nodes that can fit the workload
461        let candidates: Vec<_> = nodes.values()
462            .filter(|n| n.can_fit(&workload.resources))
463            .filter(|n| self.check_constraints(workload, n))
464            .filter(|n| self.check_taints(workload, n))
465            .collect();
466
467        if candidates.is_empty() {
468            // Try preemption
469            drop(nodes);
470            if let Some(decision) = self.try_preemption(workload) {
471                return decision;
472            }
473
474            return SchedulingDecision {
475                workload_id: workload.id.clone(),
476                node_id: None,
477                score: 0.0,
478                reason: "No suitable nodes available".to_string(),
479                preempted: Vec::new(),
480                latency_ms: start.elapsed().as_millis() as u64,
481            };
482        }
483
484        // Score and select best node
485        let scored: Vec<_> = candidates.iter()
486            .map(|n| (n, self.algorithm.score(workload, n)))
487            .collect();
488
489        let (best_node, score) = scored.iter()
490            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
491            .map(|(n, s)| (*n, *s))
492            .unwrap();
493
494        let node_id = best_node.node_id;
495        drop(nodes);
496
497        // Allocate resources
498        if let Some(node) = self.nodes.write().get_mut(&node_id) {
499            node.allocate(&workload.resources);
500        }
501
502        self.assignments.write().insert(workload.id.clone(), node_id);
503
504        info!(
505            workload_id = %workload.id,
506            node_id = %node_id,
507            score = score,
508            "Workload scheduled"
509        );
510
511        SchedulingDecision {
512            workload_id: workload.id.clone(),
513            node_id: Some(node_id),
514            score,
515            reason: "Scheduled successfully".to_string(),
516            preempted: Vec::new(),
517            latency_ms: start.elapsed().as_millis() as u64,
518        }
519    }
520
521    /// Check placement constraints
522    fn check_constraints(&self, workload: &Workload, node: &NodeResources) -> bool {
523        for constraint in &workload.constraints {
524            if !constraint.matches(node) {
525                return false;
526            }
527        }
528
529        // Check affinity
530        if let Some(affinity) = &workload.affinity {
531            if !affinity.matches(node) {
532                return false;
533            }
534        }
535
536        true
537    }
538
539    /// Check if workload tolerates node taints
540    fn check_taints(&self, workload: &Workload, node: &NodeResources) -> bool {
541        for taint in &node.taints {
542            let tolerated = workload.tolerations.iter().any(|t| {
543                // Check key match
544                let key_matches = t.key.as_ref().map(|k| k == &taint.key).unwrap_or(true);
545                
546                // Check operator
547                let value_matches = match t.operator {
548                    TolerationOperator::Exists => true,
549                    TolerationOperator::Equal => {
550                        t.value.as_ref().map(|v| v == &taint.value).unwrap_or(false)
551                    }
552                };
553
554                // Check effect
555                let effect_matches = t.effect.map(|e| e == taint.effect).unwrap_or(true);
556
557                key_matches && value_matches && effect_matches
558            });
559
560            if !tolerated && taint.effect == TaintEffect::NoSchedule {
561                return false;
562            }
563        }
564
565        true
566    }
567
568    /// Try to preempt lower priority workloads
569    fn try_preemption(&self, workload: &Workload) -> Option<SchedulingDecision> {
570        if workload.preemption_policy == PreemptionPolicy::Never {
571            return None;
572        }
573
574        // Find workloads that can be preempted
575        let assignments = self.assignments.read();
576        let mut nodes = self.nodes.write();
577
578        for (node_id, node) in nodes.iter_mut() {
579            // Find lower priority workloads on this node
580            let preemptable: Vec<_> = assignments.iter()
581                .filter(|(_, n)| *n == node_id)
582                .map(|(w, _)| w.clone())
583                .collect();
584
585            // Simulate releasing resources
586            // In a real implementation, we'd track workload resources per node
587            if node.can_fit(&workload.resources) || !preemptable.is_empty() {
588                // For now, just return that preemption is possible
589                warn!(
590                    workload_id = %workload.id,
591                    node_id = %node_id,
592                    "Preemption would be required"
593                );
594            }
595        }
596
597        None
598    }
599
600    /// Get workload assignment
601    pub fn get_assignment(&self, workload_id: &str) -> Option<NodeId> {
602        self.assignments.read().get(workload_id).copied()
603    }
604
605    /// Release workload resources
606    pub fn release(&self, workload_id: &str, resources: &ResourceRequirements, gpu_ids: &[u32]) {
607        if let Some(node_id) = self.assignments.write().remove(workload_id) {
608            if let Some(node) = self.nodes.write().get_mut(&node_id) {
609                node.release(resources, gpu_ids);
610            }
611        }
612    }
613
614    /// Get queue length
615    pub fn queue_len(&self) -> usize {
616        self.queue.len()
617    }
618
619    /// Register priority class
620    pub fn register_priority_class(&self, class: PriorityClass) {
621        self.priority_classes.write().insert(class.name.clone(), class);
622    }
623}
624
625impl Default for Scheduler {
626    fn default() -> Self {
627        Self::new()
628    }
629}
630
631#[cfg(test)]
632mod tests {
633    use super::*;
634
635    #[test]
636    fn test_node_resources() {
637        let mut node = NodeResources::new(NodeId::new(), 4000, 8192);
638        
639        let req = ResourceRequirements::new().cpu(1000).memory(2048);
640        assert!(node.can_fit(&req));
641        
642        assert!(node.allocate(&req));
643        assert_eq!(node.cpu_available(), 3000);
644        assert_eq!(node.memory_available(), 6144);
645    }
646
647    #[test]
648    fn test_scheduler_basic() {
649        let scheduler = Scheduler::new();
650        
651        let node = NodeResources::new(NodeId::new(), 4000, 8192);
652        scheduler.register_node(node);
653        
654        let workload = Workload::new("w1", "test")
655            .with_resources(ResourceRequirements::new().cpu(1000).memory(1024));
656        
657        let decision = scheduler.schedule(&workload);
658        assert!(decision.node_id.is_some());
659    }
660
661    #[test]
662    fn test_scheduler_no_capacity() {
663        let scheduler = Scheduler::new();
664        
665        let node = NodeResources::new(NodeId::new(), 1000, 1024);
666        scheduler.register_node(node);
667        
668        let workload = Workload::new("w1", "test")
669            .with_resources(ResourceRequirements::new().cpu(2000).memory(2048));
670        
671        let decision = scheduler.schedule(&workload);
672        assert!(decision.node_id.is_none());
673    }
674}
forge_orchestration/scheduler/mod.rs

forge_orchestration/scheduler/
mod.rs
