forge_orchestration/scheduler/
placement.rs

1//! Placement constraints and affinity rules
2//!
3//! Implements Kubernetes-style affinity/anti-affinity with extensions for:
4//! - GPU topology awareness
5//! - Network locality
6//! - Data locality
7
8use super::NodeResources;
9use serde::{Deserialize, Serialize};
10
11/// Placement constraint for workload scheduling
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub enum PlacementConstraint {
14    /// Node must have specific label
15    NodeSelector {
16        /// Label key the node must carry.
17        key: String,
18        /// Required value for that label.
19        value: String,
20    },
21    /// Node must be in specific zone
22    Zone(String),
23    /// Node must be in specific region
24    Region(String),
25    /// Node must have GPU
26    RequiresGpu,
27    /// Node must have specific GPU model
28    GpuModel(String),
29    /// Node must have minimum GPU memory
30    MinGpuMemory(u64),
31    /// Node must have tensor cores
32    RequiresTensorCores,
33    /// Node must have specific architecture
34    Architecture(String),
35    /// Custom constraint with expression
36    Expression(ConstraintExpression),
37}
38
39impl PlacementConstraint {
40    /// Check if node matches constraint
41    pub fn matches(&self, node: &NodeResources) -> bool {
42        match self {
43            PlacementConstraint::NodeSelector { key, value } => {
44                node.labels.get(key).map(|v| v == value).unwrap_or(false)
45            }
46            PlacementConstraint::Zone(zone) => {
47                node.labels.get("topology.kubernetes.io/zone")
48                    .or_else(|| node.labels.get("zone"))
49                    .map(|z| z == zone)
50                    .unwrap_or(false)
51            }
52            PlacementConstraint::Region(region) => {
53                node.labels.get("topology.kubernetes.io/region")
54                    .or_else(|| node.labels.get("region"))
55                    .map(|r| r == region)
56                    .unwrap_or(false)
57            }
58            PlacementConstraint::RequiresGpu => {
59                !node.gpus.is_empty() && node.gpus_available() > 0
60            }
61            PlacementConstraint::GpuModel(model) => {
62                node.gpus.iter().any(|g| g.model.contains(model))
63            }
64            PlacementConstraint::MinGpuMemory(min_mb) => {
65                node.gpus.iter().any(|g| g.memory_mb >= *min_mb)
66            }
67            PlacementConstraint::RequiresTensorCores => {
68                node.gpus.iter().any(|g| g.tensor_cores)
69            }
70            PlacementConstraint::Architecture(arch) => {
71                node.labels.get("kubernetes.io/arch")
72                    .or_else(|| node.labels.get("arch"))
73                    .map(|a| a == arch)
74                    .unwrap_or(false)
75            }
76            PlacementConstraint::Expression(expr) => {
77                expr.evaluate(node)
78            }
79        }
80    }
81}
82
83/// Constraint expression for complex matching
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct ConstraintExpression {
86    /// Key to match
87    pub key: String,
88    /// Operator
89    pub operator: ExpressionOperator,
90    /// Values to match against
91    pub values: Vec<String>,
92}
93
94/// Expression operator
95#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
96pub enum ExpressionOperator {
97    /// Key's value must be in values
98    In,
99    /// Key's value must not be in values
100    NotIn,
101    /// Key must exist
102    Exists,
103    /// Key must not exist
104    DoesNotExist,
105    /// Key's value must be greater than
106    Gt,
107    /// Key's value must be less than
108    Lt,
109}
110
111impl ConstraintExpression {
112    /// Evaluate expression against node
113    pub fn evaluate(&self, node: &NodeResources) -> bool {
114        let value = node.labels.get(&self.key);
115
116        match self.operator {
117            ExpressionOperator::In => {
118                value.map(|v| self.values.contains(v)).unwrap_or(false)
119            }
120            ExpressionOperator::NotIn => {
121                value.map(|v| !self.values.contains(v)).unwrap_or(true)
122            }
123            ExpressionOperator::Exists => value.is_some(),
124            ExpressionOperator::DoesNotExist => value.is_none(),
125            ExpressionOperator::Gt => {
126                if let (Some(v), Some(threshold)) = (value, self.values.first()) {
127                    v.parse::<i64>().ok()
128                        .zip(threshold.parse::<i64>().ok())
129                        .map(|(v, t)| v > t)
130                        .unwrap_or(false)
131                } else {
132                    false
133                }
134            }
135            ExpressionOperator::Lt => {
136                if let (Some(v), Some(threshold)) = (value, self.values.first()) {
137                    v.parse::<i64>().ok()
138                        .zip(threshold.parse::<i64>().ok())
139                        .map(|(v, t)| v < t)
140                        .unwrap_or(false)
141                } else {
142                    false
143                }
144            }
145        }
146    }
147}
148
149/// Affinity configuration
150#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct Affinity {
152    /// Node affinity rules
153    pub node_affinity: Option<NodeAffinity>,
154    /// Pod/workload affinity rules
155    pub workload_affinity: Option<WorkloadAffinity>,
156    /// Pod/workload anti-affinity rules
157    pub workload_anti_affinity: Option<WorkloadAffinity>,
158}
159
160impl Affinity {
161    /// Create empty affinity
162    pub fn new() -> Self {
163        Self {
164            node_affinity: None,
165            workload_affinity: None,
166            workload_anti_affinity: None,
167        }
168    }
169
170    /// Set node affinity
171    pub fn with_node_affinity(mut self, affinity: NodeAffinity) -> Self {
172        self.node_affinity = Some(affinity);
173        self
174    }
175
176    /// Set workload affinity
177    pub fn with_workload_affinity(mut self, affinity: WorkloadAffinity) -> Self {
178        self.workload_affinity = Some(affinity);
179        self
180    }
181
182    /// Set workload anti-affinity
183    pub fn with_workload_anti_affinity(mut self, affinity: WorkloadAffinity) -> Self {
184        self.workload_anti_affinity = Some(affinity);
185        self
186    }
187
188    /// Check if node matches affinity rules
189    pub fn matches(&self, node: &NodeResources) -> bool {
190        // Check node affinity
191        if let Some(node_affinity) = &self.node_affinity {
192            // Required rules must all match
193            for rule in &node_affinity.required {
194                if !rule.matches(node) {
195                    return false;
196                }
197            }
198        }
199
200        true
201    }
202}
203
204impl Default for Affinity {
205    fn default() -> Self {
206        Self::new()
207    }
208}
209
210/// Node affinity rules
211#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct NodeAffinity {
213    /// Required rules (must match)
214    pub required: Vec<AffinityRule>,
215    /// Preferred rules (soft preference)
216    pub preferred: Vec<WeightedAffinityRule>,
217}
218
219impl NodeAffinity {
220    /// Create new node affinity
221    pub fn new() -> Self {
222        Self {
223            required: Vec::new(),
224            preferred: Vec::new(),
225        }
226    }
227
228    /// Add required rule
229    pub fn require(mut self, rule: AffinityRule) -> Self {
230        self.required.push(rule);
231        self
232    }
233
234    /// Add preferred rule
235    pub fn prefer(mut self, weight: i32, rule: AffinityRule) -> Self {
236        self.preferred.push(WeightedAffinityRule { weight, rule });
237        self
238    }
239}
240
241impl Default for NodeAffinity {
242    fn default() -> Self {
243        Self::new()
244    }
245}
246
247/// Workload affinity/anti-affinity rules
248#[derive(Debug, Clone, Serialize, Deserialize)]
249pub struct WorkloadAffinity {
250    /// Required rules
251    pub required: Vec<WorkloadAffinityTerm>,
252    /// Preferred rules
253    pub preferred: Vec<WeightedWorkloadAffinityTerm>,
254}
255
256impl WorkloadAffinity {
257    /// Create new workload affinity
258    pub fn new() -> Self {
259        Self {
260            required: Vec::new(),
261            preferred: Vec::new(),
262        }
263    }
264}
265
266impl Default for WorkloadAffinity {
267    fn default() -> Self {
268        Self::new()
269    }
270}
271
272/// Workload affinity term
273#[derive(Debug, Clone, Serialize, Deserialize)]
274pub struct WorkloadAffinityTerm {
275    /// Label selector for matching workloads
276    pub label_selector: LabelSelector,
277    /// Topology key for co-location
278    pub topology_key: String,
279    /// Namespaces to consider
280    pub namespaces: Option<Vec<String>>,
281}
282
283/// Weighted workload affinity term
284#[derive(Debug, Clone, Serialize, Deserialize)]
285pub struct WeightedWorkloadAffinityTerm {
286    /// Weight (1-100)
287    pub weight: i32,
288    /// Affinity term
289    pub term: WorkloadAffinityTerm,
290}
291
292/// Label selector for matching
293#[derive(Debug, Clone, Serialize, Deserialize)]
294pub struct LabelSelector {
295    /// Match labels exactly
296    pub match_labels: std::collections::HashMap<String, String>,
297    /// Match expressions
298    pub match_expressions: Vec<ConstraintExpression>,
299}
300
301impl LabelSelector {
302    /// Create new label selector
303    pub fn new() -> Self {
304        Self {
305            match_labels: std::collections::HashMap::new(),
306            match_expressions: Vec::new(),
307        }
308    }
309
310    /// Add label match
311    pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
312        self.match_labels.insert(key.into(), value.into());
313        self
314    }
315
316    /// Add expression match
317    pub fn with_expression(mut self, expr: ConstraintExpression) -> Self {
318        self.match_expressions.push(expr);
319        self
320    }
321}
322
323impl Default for LabelSelector {
324    fn default() -> Self {
325        Self::new()
326    }
327}
328
329/// Affinity rule
330#[derive(Debug, Clone, Serialize, Deserialize)]
331pub struct AffinityRule {
332    /// Match expressions
333    pub match_expressions: Vec<ConstraintExpression>,
334}
335
336impl AffinityRule {
337    /// Create new affinity rule
338    pub fn new() -> Self {
339        Self {
340            match_expressions: Vec::new(),
341        }
342    }
343
344    /// Add expression
345    pub fn with_expression(mut self, expr: ConstraintExpression) -> Self {
346        self.match_expressions.push(expr);
347        self
348    }
349
350    /// Check if node matches rule
351    pub fn matches(&self, node: &NodeResources) -> bool {
352        self.match_expressions.iter().all(|expr| expr.evaluate(node))
353    }
354}
355
356impl Default for AffinityRule {
357    fn default() -> Self {
358        Self::new()
359    }
360}
361
362/// Weighted affinity rule
363#[derive(Debug, Clone, Serialize, Deserialize)]
364pub struct WeightedAffinityRule {
365    /// Weight (1-100)
366    pub weight: i32,
367    /// Rule
368    pub rule: AffinityRule,
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374    use crate::types::{NodeId, GpuResources};
375
376    #[test]
377    fn test_node_selector() {
378        let mut node = NodeResources::new(NodeId::new(), 4000, 8192);
379        node.labels.insert("env".to_string(), "production".to_string());
380
381        let constraint = PlacementConstraint::NodeSelector {
382            key: "env".to_string(),
383            value: "production".to_string(),
384        };
385
386        assert!(constraint.matches(&node));
387
388        let wrong_constraint = PlacementConstraint::NodeSelector {
389            key: "env".to_string(),
390            value: "staging".to_string(),
391        };
392
393        assert!(!wrong_constraint.matches(&node));
394    }
395
396    #[test]
397    fn test_gpu_constraints() {
398        let gpu = GpuResources::new(0, "NVIDIA A100", 40960)
399            .with_tensor_cores(true)
400            .with_compute_capability(8.0);
401
402        let node = NodeResources::new(NodeId::new(), 4000, 8192)
403            .with_gpu(gpu);
404
405        assert!(PlacementConstraint::RequiresGpu.matches(&node));
406        assert!(PlacementConstraint::RequiresTensorCores.matches(&node));
407        assert!(PlacementConstraint::GpuModel("A100".to_string()).matches(&node));
408        assert!(PlacementConstraint::MinGpuMemory(40000).matches(&node));
409        assert!(!PlacementConstraint::MinGpuMemory(50000).matches(&node));
410    }
411
412    #[test]
413    fn test_expression_operators() {
414        let mut node = NodeResources::new(NodeId::new(), 4000, 8192);
415        node.labels.insert("tier".to_string(), "frontend".to_string());
416        node.labels.insert("priority".to_string(), "10".to_string());
417
418        let in_expr = ConstraintExpression {
419            key: "tier".to_string(),
420            operator: ExpressionOperator::In,
421            values: vec!["frontend".to_string(), "backend".to_string()],
422        };
423        assert!(in_expr.evaluate(&node));
424
425        let gt_expr = ConstraintExpression {
426            key: "priority".to_string(),
427            operator: ExpressionOperator::Gt,
428            values: vec!["5".to_string()],
429        };
430        assert!(gt_expr.evaluate(&node));
431
432        let exists_expr = ConstraintExpression {
433            key: "tier".to_string(),
434            operator: ExpressionOperator::Exists,
435            values: vec![],
436        };
437        assert!(exists_expr.evaluate(&node));
438    }
439}
forge_orchestration/scheduler/placement.rs

forge_orchestration/scheduler/
placement.rs
