Skip to main content

ghostlink_core/
planning.rs

1//! Greedy Layer Assignment with Fault Tolerance and Adaptive Quantization
2//!
3//! This module provides:
4//! - Sequential greedy layer splitting across nodes based on VRAM capacity
5//! - Adaptive quantization trigger (select_quantization_mode)
6//! - Load balancing and fault detection integration
7
8use crate::accelerator::ExecutionBackend;
9use crate::cluster::ClusterState;
10use crate::cluster::NodeStatus;
11use crate::host::{AccelerationMode, RuntimeProfile};
12use crate::protocol::NodeResources;
13
14/// Delivery ratio thresholds for adaptive quantization
15pub const DELIVERY_RATIO_INT8_THRESHOLD: f32 = 0.95;
16pub const DELIVERY_RATIO_INT4_THRESHOLD: f32 = 0.80;
17
18/// Layer specification with VRAM requirements
19#[derive(Clone, Debug, PartialEq)]
20pub struct LayerSpec {
21    /// Layer index (0-based)
22    pub index: usize,
23    /// VRAM required in GB
24    pub vram_gb: f32,
25    /// Number of weights in the layer
26    pub num_weights: u32,
27}
28
29impl Default for LayerSpec {
30    fn default() -> Self {
31        Self {
32            index: 0,
33            vram_gb: 1.0,
34            num_weights: 0,
35        }
36    }
37}
38
39/// Layer assignment to a specific node
40#[derive(Clone, Debug, PartialEq)]
41pub struct LayerAssignment {
42    /// Node ID
43    pub node_id: String,
44    /// Start layer index (inclusive)
45    pub start_layer: usize,
46    /// End layer index (exclusive)
47    pub end_layer: usize,
48    /// VRAM used on this node
49    pub used_vram_gb: f32,
50    /// Number of layers assigned
51    pub num_layers: usize,
52}
53
54impl LayerAssignment {
55    /// Create new layer assignment
56    pub fn new(node_id: String, start_layer: usize, end_layer: usize, vram_gb: f32) -> Self {
57        Self {
58            node_id,
59            start_layer,
60            end_layer,
61            used_vram_gb: vram_gb,
62            num_layers: end_layer - start_layer,
63        }
64    }
65
66    /// Get average VRAM per layer
67    pub fn avg_vram_per_layer(&self) -> f32 {
68        if self.num_layers == 0 {
69            0.0
70        } else {
71            self.used_vram_gb / self.num_layers as f32
72        }
73    }
74}
75
76/// Quantization mode enumeration
77#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
78pub enum QuantizationMode {
79    /// No quantization (full precision)
80    #[default]
81    None,
82    /// 8-bit quantization
83    Int8,
84    /// 4-bit quantization
85    Int4,
86}
87
88/// Layer placement plan across nodes
89#[derive(Clone, Debug, Default)]
90pub struct PlacementPlan {
91    /// Assignments per node
92    pub assignments: Vec<LayerAssignment>,
93    /// Selected quantization mode
94    pub quantization_mode: QuantizationMode,
95    /// Total layers assigned
96    pub total_layers: usize,
97    /// Nodes participating in plan
98    pub participating_nodes: Vec<String>,
99}
100
101impl PlacementPlan {
102    /// Create new placement plan
103    pub fn new(assignments: Vec<LayerAssignment>, quantization_mode: QuantizationMode) -> Self {
104        let participating_nodes = assignments.iter().map(|a| a.node_id.clone()).collect();
105        let total_layers = assignments.iter().map(|a| a.num_layers).sum();
106
107        Self {
108            assignments,
109            quantization_mode,
110            total_layers,
111            participating_nodes,
112        }
113    }
114
115    /// Get human-readable plan summary
116    pub fn summary(&self) -> String {
117        let mode_str = match self.quantization_mode {
118            QuantizationMode::None => "Full Precision",
119            QuantizationMode::Int8 => "8-bit Quantized",
120            QuantizationMode::Int4 => "4-bit Quantized",
121        };
122
123        format!(
124            "Placement Plan ({})\n\
125             =================\n\
126             Total layers: {}\n\
127             Quantization: {}\n\
128             Nodes: {}\n",
129            mode_str,
130            self.total_layers,
131            match self.quantization_mode {
132                QuantizationMode::None => "Full Precision".to_string(),
133                QuantizationMode::Int8 => "8-bit Quantized".to_string(),
134                QuantizationMode::Int4 => "4-bit Quantized".to_string(),
135            },
136            self.participating_nodes.join(", ")
137        )
138    }
139}
140
141/// Runtime-aware planning hints derived from host auto-detection.
142#[derive(Clone, Copy, Debug, PartialEq, Eq)]
143pub struct PlanningTuning {
144    /// Preferred maximum chunk size to expose parallel work to workers.
145    pub max_layers_per_assignment: usize,
146}
147
148impl PlanningTuning {
149    /// Derive planning hints from the detected runtime profile.
150    pub fn from_runtime_profile(profile: &RuntimeProfile, total_layers: usize) -> Self {
151        let backend = ExecutionBackend::from_runtime_profile(profile);
152        let worker_count = backend.worker_count.max(1);
153        let accelerator_bias = match profile.acceleration_mode {
154            AccelerationMode::Gpu => 2,
155            AccelerationMode::Avx512 => 1,
156            _ => 0,
157        };
158        let vector_bias = (backend.vector_width_bits / 256).max(1);
159        let target_chunks = (worker_count + accelerator_bias + vector_bias - 1).max(1);
160        let chunk_size = if total_layers == 0 {
161            1
162        } else {
163            total_layers.div_ceil(target_chunks).max(1)
164        };
165
166        Self {
167            max_layers_per_assignment: chunk_size,
168        }
169    }
170}
171
172/// Select quantization mode based on cluster health metrics
173pub fn select_quantization_mode(delivery_ratio: f32) -> QuantizationMode {
174    if delivery_ratio >= DELIVERY_RATIO_INT8_THRESHOLD {
175        QuantizationMode::None
176    } else if delivery_ratio >= DELIVERY_RATIO_INT4_THRESHOLD {
177        QuantizationMode::Int8
178    } else {
179        QuantizationMode::Int4
180    }
181}
182
183/// Assign layers sequentially across nodes based on VRAM capacity
184pub fn assign_layers_sequentially(
185    nodes: &[NodeResources],
186    layers: &[LayerSpec],
187) -> Result<Vec<LayerAssignment>, String> {
188    if nodes.is_empty() {
189        return Err("at least one node is required".into());
190    }
191    if layers.is_empty() {
192        return Ok(Vec::new());
193    }
194
195    let mut assignments = Vec::new();
196    let mut current_node_index = 0usize;
197    let mut remaining_capacity = nodes[0].vram_gb;
198    let mut current_assignment: Option<LayerAssignment> = None;
199
200    for layer in layers {
201        while layer.vram_gb > remaining_capacity {
202            // Need to move to next node
203            if let Some(assignment) = current_assignment.take() {
204                assignments.push(assignment);
205            }
206
207            current_node_index += 1;
208            if current_node_index >= nodes.len() {
209                return Err(format!(
210                    "insufficient cluster VRAM for layer {} (needs {:.2} GB)",
211                    layer.index, layer.vram_gb
212                ));
213            }
214            remaining_capacity = nodes[current_node_index].vram_gb;
215        }
216
217        // Assign layer to current node
218        remaining_capacity -= layer.vram_gb;
219
220        match current_assignment.as_mut() {
221            Some(assignment) => {
222                assignment.end_layer = layer.index + 1;
223                assignment.used_vram_gb += layer.vram_gb;
224                assignment.num_layers += 1;
225            }
226            None => {
227                current_assignment = Some(LayerAssignment::new(
228                    nodes[current_node_index].id.clone(),
229                    layer.index,
230                    layer.index + 1,
231                    layer.vram_gb,
232                ));
233            }
234        }
235    }
236
237    // Finalize last assignment
238    if let Some(assignment) = current_assignment {
239        assignments.push(assignment);
240    }
241
242    Ok(assignments)
243}
244
245/// Split large node assignments into smaller contiguous chunks for worker-level parallelism.
246pub fn chunk_assignments_for_workers(
247    assignments: Vec<LayerAssignment>,
248    max_layers_per_assignment: usize,
249) -> Vec<LayerAssignment> {
250    let chunk_size = max_layers_per_assignment.max(1);
251
252    let needs_chunking = assignments.iter().any(|a| a.num_layers > chunk_size);
253    if !needs_chunking {
254        return assignments;
255    }
256
257    let mut chunked = Vec::with_capacity(assignments.len() * 2);
258
259    for assignment in assignments {
260        if assignment.num_layers <= chunk_size {
261            chunked.push(assignment);
262            continue;
263        }
264
265        let avg_vram = assignment.avg_vram_per_layer();
266        let mut start_layer = assignment.start_layer;
267
268        while start_layer < assignment.end_layer {
269            let end_layer = (start_layer + chunk_size).min(assignment.end_layer);
270            let num_layers = end_layer - start_layer;
271            chunked.push(LayerAssignment::new(
272                assignment.node_id.clone(),
273                start_layer,
274                end_layer,
275                avg_vram * num_layers as f32,
276            ));
277            start_layer = end_layer;
278        }
279    }
280
281    chunked
282}
283
284/// Assign layers sequentially across nodes, directly producing chunked assignments
285/// sized for worker-level parallelism. Avoids creating intermediate full-size
286/// `LayerAssignment` objects and their associated allocation/clone overhead.
287fn assign_layers_chunked(
288    nodes: &[NodeResources],
289    layers: &[LayerSpec],
290    max_layers_per_assignment: usize,
291) -> Result<Vec<LayerAssignment>, String> {
292    if nodes.is_empty() {
293        return Err("at least one node is required".into());
294    }
295    if layers.is_empty() {
296        return Ok(Vec::new());
297    }
298
299    let chunk_size = max_layers_per_assignment.max(1);
300    let est_assignments = (layers.len().div_ceil(chunk_size)).max(nodes.len());
301    let mut assignments = Vec::with_capacity(est_assignments);
302    let mut node_idx = 0usize;
303    let mut remaining_vram = nodes[0].vram_gb;
304    // Index of the first layer in the current chunk
305    let mut chunk_start = 0usize;
306    let mut chunk_vram = 0.0f32;
307
308    for (i, layer) in layers.iter().enumerate() {
309        // Move to next node if current node is out of capacity
310        while layer.vram_gb > remaining_vram {
311            if i > chunk_start {
312                assignments.push(LayerAssignment::new(
313                    nodes[node_idx].id.clone(),
314                    chunk_start,
315                    i,
316                    chunk_vram,
317                ));
318            }
319            node_idx += 1;
320            if node_idx >= nodes.len() {
321                return Err(format!(
322                    "insufficient cluster VRAM for layer {} (needs {:.2} GB)",
323                    layer.index, layer.vram_gb
324                ));
325            }
326            remaining_vram = nodes[node_idx].vram_gb;
327            chunk_start = i;
328            chunk_vram = 0.0;
329        }
330
331        remaining_vram -= layer.vram_gb;
332        chunk_vram += layer.vram_gb;
333
334        // Flush chunk if it reached the maximum chunk size
335        let chunk_len = i + 1 - chunk_start;
336        if chunk_len >= chunk_size {
337            assignments.push(LayerAssignment::new(
338                nodes[node_idx].id.clone(),
339                chunk_start,
340                i + 1,
341                chunk_vram,
342            ));
343            chunk_start = i + 1;
344            chunk_vram = 0.0;
345        }
346    }
347
348    // Finalize last chunk. Checked by pending-layer range, not accumulated
349    // VRAM: a trailing run of zero-VRAM layers (e.g. bias-only layers) would
350    // otherwise leave `chunk_vram == 0.0` and be silently dropped from the
351    // plan even though `chunk_start..layers.len()` is non-empty.
352    if chunk_start < layers.len() {
353        assignments.push(LayerAssignment::new(
354            nodes[node_idx].id.clone(),
355            chunk_start,
356            layers.len(),
357            chunk_vram,
358        ));
359    }
360
361    Ok(assignments)
362}
363
364/// Assign layers using runtime auto-detection to expose worker-parallel chunks.
365pub fn assign_layers_with_runtime_profile(
366    nodes: &[NodeResources],
367    layers: &[LayerSpec],
368    profile: &RuntimeProfile,
369) -> Result<Vec<LayerAssignment>, String> {
370    let tuning = PlanningTuning::from_runtime_profile(profile, layers.len());
371    assign_layers_chunked(nodes, layers, tuning.max_layers_per_assignment)
372}
373
374/// Assign layers with fault tolerance and load balancing
375pub fn assign_layers_with_fault_tolerance(
376    cluster: &ClusterState,
377    layers: &[LayerSpec],
378) -> Result<PlacementPlan, String> {
379    let nodes = cluster.nodes_snapshot();
380
381    if nodes.is_empty() {
382        return Err("no nodes available".into());
383    }
384
385    // First pass: greedy assignment
386    let assignments = assign_layers_sequentially(&nodes, layers)?;
387
388    // Calculate average delivery ratio across all nodes in a single lock pass
389    // (avoids the clone/allocation of active_nodes() and the lock contention
390    // of taking it twice). Also avoids a real divide-by-zero/TOCTOU hazard:
391    // calling `active_nodes()` twice (once for the sum, once for the count)
392    // risked dividing by a different node count than was summed if a node's
393    // status changed between calls, and an empty active set (e.g. all nodes
394    // failed but still registered) would divide by zero.
395    let (sum_delivery, active_count) = {
396        let metrics = cluster
397            .metrics
398            .lock()
399            .unwrap_or_else(|poison| poison.into_inner());
400        let mut sum = 0.0f32;
401        let mut count = 0usize;
402        for metric in metrics.values() {
403            if metric.status == NodeStatus::Active {
404                sum += metric.delivery_ratio;
405                count += 1;
406            }
407        }
408        (sum, count)
409    };
410
411    // Fall back to 0.0 (not 1.0) when there are zero active nodes to sample:
412    // this selects the most conservative quantization mode below. Assuming a
413    // perfect 1.0 ratio when we have literally no corroborating health data
414    // is optimistic in exactly the scenario (all nodes failed/degraded) where
415    // it's least likely to be true.
416    let total_delivery_ratio = if active_count > 0 {
417        sum_delivery / active_count as f32
418    } else {
419        0.0
420    };
421
422    // Select quantization mode based on health
423    let quantization_mode = select_quantization_mode(total_delivery_ratio);
424
425    Ok(PlacementPlan::new(assignments, quantization_mode))
426}
427
428/// Assign layers with fault tolerance and runtime-aware chunking.
429///
430/// Optimized to directly call `assign_layers_chunked`, bypassing intermediate
431/// `LayerAssignment` allocations and `chunk_assignments_for_workers` overhead.
432pub fn assign_layers_with_fault_tolerance_and_runtime(
433    cluster: &ClusterState,
434    layers: &[LayerSpec],
435    profile: &RuntimeProfile,
436) -> Result<PlacementPlan, String> {
437    let nodes = cluster.nodes_snapshot();
438
439    if nodes.is_empty() {
440        return Err("no nodes available".into());
441    }
442
443    let tuning = PlanningTuning::from_runtime_profile(profile, layers.len());
444    // Directly run the chunked layout computation
445    let assignments = assign_layers_chunked(&nodes, layers, tuning.max_layers_per_assignment)?;
446
447    // Calculate average delivery ratio across all nodes in a single lock pass
448    let (sum_delivery, active_count) = {
449        let metrics = cluster
450            .metrics
451            .lock()
452            .unwrap_or_else(|poison| poison.into_inner());
453        let mut sum = 0.0f32;
454        let mut count = 0usize;
455        for metric in metrics.values() {
456            if metric.status == NodeStatus::Active {
457                sum += metric.delivery_ratio;
458                count += 1;
459            }
460        }
461        (sum, count)
462    };
463
464    let total_delivery_ratio = if active_count > 0 {
465        sum_delivery / active_count as f32
466    } else {
467        0.0
468    };
469
470    // Select quantization mode based on health
471    let quantization_mode = select_quantization_mode(total_delivery_ratio);
472
473    Ok(PlacementPlan::new(assignments, quantization_mode))
474}
475
476/// Update layer assignments based on node health metrics
477pub fn rebalance_assignments(cluster: &ClusterState, plan: &mut PlacementPlan) -> bool {
478    // Check if any active node has high available VRAM
479    let mut needs_rebalance = false;
480
481    for assignment in &mut plan.assignments {
482        if let Some(metrics) = cluster.get_metrics(&assignment.node_id) {
483            if metrics.status != NodeStatus::Active {
484                // Skip failed nodes
485                continue;
486            }
487
488            // Check if this node can take more layers
489            let available = metrics.available_vram_gb;
490            let avg_layer_size = assignment.avg_vram_per_layer();
491
492            if available > 0.0 && avg_layer_size > 0.0 {
493                let potential_layers = (available / avg_layer_size) as usize;
494
495                // If node has significant capacity, mark for rebalancing
496                if potential_layers > 2 {
497                    needs_rebalance = true;
498                    break;
499                }
500            }
501        }
502    }
503
504    needs_rebalance
505}
506
507/// Simulate layer streaming on a node with metrics updates
508#[cfg(test)]
509pub fn simulate_layer_streaming(
510    node_id: &str,
511    cluster: &ClusterState,
512    start_layer: usize,
513    end_layer: usize,
514) -> Option<crate::cluster::NodeMetrics> {
515    let mut metrics = cluster.get_metrics(node_id)?;
516
517    // Record VRAM usage
518    let num_layers = end_layer - start_layer;
519    let avg_vram = (metrics.available_vram_gb + metrics.vram_gb) / 2.0;
520    let vram_per_layer = avg_vram / num_layers as f32;
521    metrics.record_vram_usage(vram_per_layer * num_layers as f32);
522
523    // Set streaming layers
524    metrics.set_streaming_layers(start_layer, end_layer);
525
526    Some(metrics)
527}
528
529/// Calculate network health across the cluster
530pub fn calculate_cluster_health(cluster: &ClusterState) -> (f32, usize, Vec<String>) {
531    // Acquire the lock on metrics map exactly ONCE to avoid redundant mutex lock/unlock and clones
532    let metrics = cluster
533        .metrics
534        .lock()
535        .unwrap_or_else(|poison| poison.into_inner());
536
537    if metrics.is_empty() {
538        return (0.0, 0, vec![]);
539    }
540
541    let mut sum_delivery_ratio = 0.0f32;
542    let mut active_count = 0usize;
543    let mut failed_nodes = Vec::new();
544
545    // Iterate once over all live node metrics to aggregate statistics and collect failed nodes
546    for (node_id, metric) in metrics.iter() {
547        match metric.status {
548            NodeStatus::Active => {
549                sum_delivery_ratio += metric.delivery_ratio;
550                active_count += 1;
551            }
552            NodeStatus::Failed => {
553                failed_nodes.push(node_id.clone());
554            }
555            _ => {}
556        }
557    }
558
559    let avg_delivery_ratio = if active_count > 0 {
560        sum_delivery_ratio / active_count as f32
561    } else {
562        0.0
563    };
564
565    let failed_count = failed_nodes.len();
566
567    (avg_delivery_ratio, failed_count, failed_nodes)
568}
569
570/// Dynamic migration planner for moving layers between nodes under load.
571#[derive(Debug, Clone)]
572pub struct MigrationPlanner {
573    pub source_node: String,
574    pub target_node: String,
575    pub layers: Vec<usize>,
576    pub est_migration_time_ms: f32,
577}
578
579impl MigrationPlanner {
580    pub fn new(source: String, target: String, layers: Vec<usize>) -> Self {
581        let est_time = layers.len() as f32 * 45.0; // Baseline 45ms per layer migration
582        Self {
583            source_node: source,
584            target_node: target,
585            layers,
586            est_migration_time_ms: est_time,
587        }
588    }
589
590    /// Generate a safe handoff sequence for a migration.
591    pub fn generate_handoff_plan(&self) -> Vec<String> {
592        let mut steps = Vec::new();
593        let target_node = &self.target_node;
594        let source_node = &self.source_node;
595        steps.push(format!(
596            "PREPARE: Target node {target_node} allocating VRAM"
597        ));
598        steps.push(format!(
599            "STREAM: Moving layers {:?} to {target_node}",
600            self.layers
601        ));
602        steps.push(format!("VERIFY: Integrity check on {target_node}"));
603        steps.push(format!("COMMIT: Switch routing to {target_node}"));
604        steps.push(format!("CLEANUP: Free VRAM on source node {source_node}"));
605        steps
606    }
607}
608
609/// Trigger mechanism for dynamic rebalancing based on load/health.
610pub struct RebalanceTrigger {
611    pub min_imbalance_ratio: f32,
612    pub max_p95_latency_ms: f32,
613}
614
615impl Default for RebalanceTrigger {
616    fn default() -> Self {
617        Self {
618            min_imbalance_ratio: 0.25, // 25% drift
619            max_p95_latency_ms: 25.0,
620        }
621    }
622}
623
624impl RebalanceTrigger {
625    pub fn evaluate(
626        &self,
627        cluster: &ClusterState,
628        current_plan: &PlacementPlan,
629    ) -> Option<MigrationPlanner> {
630        let active_nodes = cluster.active_nodes();
631        if active_nodes.len() < 2 {
632            return None;
633        }
634
635        // Detect high-load nodes and high-headroom nodes
636        let mut overtaxed = None;
637        let mut available = None;
638
639        for metrics in active_nodes {
640            if metrics.delivery_ratio < 0.90 || metrics.avg_latency_us > 15.0 {
641                overtaxed = Some(metrics.name.clone());
642            } else if metrics.available_vram_gb > 8.0 {
643                available = Some(metrics.name.clone());
644            }
645        }
646
647        if let (Some(source), Some(target)) = (overtaxed, available) {
648            if source == target {
649                return None;
650            }
651
652            // Find layers to move from overtaxed node
653            let layers_to_move = current_plan
654                .assignments
655                .iter()
656                .find(|a| a.node_id == source)
657                .map(|a| (a.start_layer..a.end_layer).take(2).collect::<Vec<_>>())
658                .unwrap_or_default();
659
660            if !layers_to_move.is_empty() {
661                return Some(MigrationPlanner::new(source, target, layers_to_move));
662            }
663        }
664
665        None
666    }
667}
668
669#[cfg(test)]
670mod tests {
671    use super::*;
672    use crate::host::AccelerationMode;
673
674    fn sample_layers(count: usize, vram_gb: f32) -> Vec<LayerSpec> {
675        (0..count)
676            .map(|index| LayerSpec {
677                index,
678                vram_gb,
679                num_weights: 0,
680            })
681            .collect()
682    }
683
684    #[test]
685    fn greedily_places_layers_across_nodes() {
686        let nodes = vec![
687            NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
688            NodeResources::new("node-b", 12.0, 32.0, "8.6", None),
689        ];
690
691        let assignments = assign_layers_sequentially(&nodes, &sample_layers(33, 1.0)).unwrap();
692
693        assert_eq!(
694            assignments,
695            vec![
696                LayerAssignment {
697                    node_id: "node-a".into(),
698                    start_layer: 0,
699                    end_layer: 24,
700                    used_vram_gb: 24.0,
701                    num_layers: 24,
702                },
703                LayerAssignment {
704                    node_id: "node-b".into(),
705                    start_layer: 24,
706                    end_layer: 33,
707                    used_vram_gb: 9.0,
708                    num_layers: 9,
709                }
710            ]
711        );
712    }
713
714    #[test]
715    fn reports_insufficient_capacity() {
716        let nodes = vec![NodeResources::new("node-a", 2.0, 64.0, "8.9", None)];
717        let error = assign_layers_sequentially(&nodes, &sample_layers(3, 1.0)).unwrap_err();
718
719        assert!(error.contains("insufficient cluster VRAM"));
720    }
721
722    #[test]
723    fn selects_quantization_mode_from_delivery_ratio() {
724        assert_eq!(select_quantization_mode(0.98), QuantizationMode::None);
725        assert_eq!(select_quantization_mode(0.90), QuantizationMode::Int8);
726        assert_eq!(select_quantization_mode(0.75), QuantizationMode::Int4);
727    }
728
729    #[test]
730    fn placement_plan_summary() {
731        let plan = PlacementPlan::new(
732            vec![
733                LayerAssignment::new("node-a".into(), 0, 24, 24.0),
734                LayerAssignment::new("node-b".into(), 24, 33, 9.0),
735            ],
736            QuantizationMode::None,
737        );
738
739        let summary = plan.summary();
740        assert!(summary.contains("Total layers: 33"));
741        assert!(summary.contains("Full Precision"));
742    }
743
744    #[test]
745    fn rebalance_assignments_detects_capacity() {
746        let mut plan = PlacementPlan::new(
747            vec![LayerAssignment::new("node-a".into(), 0, 24, 24.0)],
748            QuantizationMode::None,
749        );
750
751        // Would need actual cluster state to test rebalancing
752        assert!(!rebalance_assignments(&ClusterState::default(), &mut plan));
753    }
754
755    #[test]
756    fn runtime_profile_chunks_large_assignments() {
757        let assignments = vec![LayerAssignment::new("node-a".into(), 0, 12, 12.0)];
758        let chunked = chunk_assignments_for_workers(assignments, 5);
759
760        assert_eq!(chunked.len(), 3);
761        assert_eq!(chunked[0].start_layer, 0);
762        assert_eq!(chunked[0].end_layer, 5);
763        assert_eq!(chunked[2].start_layer, 10);
764        assert_eq!(chunked[2].end_layer, 12);
765    }
766
767    #[test]
768    fn runtime_profile_assignment_keeps_trailing_zero_vram_layers() {
769        // Regression: assign_layers_chunked finalized its last pending chunk
770        // by checking `chunk_vram > 0.0`, so a trailing run of zero-VRAM
771        // layers left a non-empty pending chunk with chunk_vram == 0.0 and
772        // silently dropped those layers from the plan.
773        let nodes = vec![NodeResources::new("node-a", 24.0, 64.0, "8.9", None)];
774        let layers = sample_layers(3, 0.0);
775        let profile = RuntimeProfile {
776            node_resources: NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
777            logical_cores: 4,
778            recommended_workers: 1,
779            acceleration_mode: AccelerationMode::Generic,
780            gpu_backend: crate::host::GpuBackend::Cpu,
781            xdp_supported: false,
782            detection_source: String::from("test"),
783            probe_mode: crate::host::ProbeMode::Fast,
784        };
785
786        let assignments = assign_layers_with_runtime_profile(&nodes, &layers, &profile).unwrap();
787        let total: usize = assignments.iter().map(|a| a.num_layers).sum();
788        assert_eq!(total, 3, "all zero-VRAM layers must still be assigned");
789    }
790
791    #[test]
792    fn fault_tolerance_handles_all_failed_nodes_without_panicking() {
793        // Regression: dividing active_nodes' summed delivery ratio by
794        // active_nodes().len() produced NaN (or a TOCTOU-inconsistent
795        // average) when every registered node's metrics were marked Failed.
796        let cluster = ClusterState::new();
797        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
798        cluster.mark_failed("node-a");
799
800        let plan = assign_layers_with_fault_tolerance(&cluster, &sample_layers(4, 1.0)).unwrap();
801        assert_eq!(plan.quantization_mode, QuantizationMode::Int4);
802    }
803
804    #[test]
805    fn planning_tuning_scales_chunk_size_with_workers() {
806        let profile = RuntimeProfile {
807            node_resources: NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
808            logical_cores: 16,
809            recommended_workers: 8,
810            acceleration_mode: AccelerationMode::Gpu,
811            gpu_backend: crate::host::GpuBackend::Cuda,
812            xdp_supported: true,
813            detection_source: String::from("test"),
814            probe_mode: crate::host::ProbeMode::Fast,
815        };
816
817        let tuning = PlanningTuning::from_runtime_profile(&profile, 40);
818        assert!(tuning.max_layers_per_assignment <= 4);
819    }
820}