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.
429pub fn assign_layers_with_fault_tolerance_and_runtime(
430    cluster: &ClusterState,
431    layers: &[LayerSpec],
432    profile: &RuntimeProfile,
433) -> Result<PlacementPlan, String> {
434    let mut plan = assign_layers_with_fault_tolerance(cluster, layers)?;
435    let tuning = PlanningTuning::from_runtime_profile(profile, layers.len());
436    plan.assignments = chunk_assignments_for_workers(
437        std::mem::take(&mut plan.assignments),
438        tuning.max_layers_per_assignment,
439    );
440    plan.total_layers = plan
441        .assignments
442        .iter()
443        .map(|assignment| assignment.num_layers)
444        .sum();
445    plan.participating_nodes = plan
446        .assignments
447        .iter()
448        .map(|assignment| assignment.node_id.clone())
449        .collect();
450    Ok(plan)
451}
452
453/// Update layer assignments based on node health metrics
454pub fn rebalance_assignments(cluster: &ClusterState, plan: &mut PlacementPlan) -> bool {
455    // Check if any active node has high available VRAM
456    let mut needs_rebalance = false;
457
458    for assignment in &mut plan.assignments {
459        if let Some(metrics) = cluster.get_metrics(&assignment.node_id) {
460            if metrics.status != NodeStatus::Active {
461                // Skip failed nodes
462                continue;
463            }
464
465            // Check if this node can take more layers
466            let available = metrics.available_vram_gb;
467            let avg_layer_size = assignment.avg_vram_per_layer();
468
469            if available > 0.0 && avg_layer_size > 0.0 {
470                let potential_layers = (available / avg_layer_size) as usize;
471
472                // If node has significant capacity, mark for rebalancing
473                if potential_layers > 2 {
474                    needs_rebalance = true;
475                    break;
476                }
477            }
478        }
479    }
480
481    needs_rebalance
482}
483
484/// Simulate layer streaming on a node with metrics updates
485#[cfg(test)]
486pub fn simulate_layer_streaming(
487    node_id: &str,
488    cluster: &ClusterState,
489    start_layer: usize,
490    end_layer: usize,
491) -> Option<crate::cluster::NodeMetrics> {
492    let mut metrics = cluster.get_metrics(node_id)?;
493
494    // Record VRAM usage
495    let num_layers = end_layer - start_layer;
496    let avg_vram = (metrics.available_vram_gb + metrics.vram_gb) / 2.0;
497    let vram_per_layer = avg_vram / num_layers as f32;
498    metrics.record_vram_usage(vram_per_layer * num_layers as f32);
499
500    // Set streaming layers
501    metrics.set_streaming_layers(start_layer, end_layer);
502
503    Some(metrics)
504}
505
506/// Calculate network health across the cluster
507pub fn calculate_cluster_health(cluster: &ClusterState) -> (f32, usize, Vec<String>) {
508    // Acquire the lock on metrics map exactly ONCE to avoid redundant mutex lock/unlock and clones
509    let metrics = cluster
510        .metrics
511        .lock()
512        .unwrap_or_else(|poison| poison.into_inner());
513
514    if metrics.is_empty() {
515        return (0.0, 0, vec![]);
516    }
517
518    let mut sum_delivery_ratio = 0.0f32;
519    let mut active_count = 0usize;
520    let mut failed_nodes = Vec::new();
521
522    // Iterate once over all live node metrics to aggregate statistics and collect failed nodes
523    for (node_id, metric) in metrics.iter() {
524        match metric.status {
525            NodeStatus::Active => {
526                sum_delivery_ratio += metric.delivery_ratio;
527                active_count += 1;
528            }
529            NodeStatus::Failed => {
530                failed_nodes.push(node_id.clone());
531            }
532            _ => {}
533        }
534    }
535
536    let avg_delivery_ratio = if active_count > 0 {
537        sum_delivery_ratio / active_count as f32
538    } else {
539        0.0
540    };
541
542    let failed_count = failed_nodes.len();
543
544    (avg_delivery_ratio, failed_count, failed_nodes)
545}
546
547/// Dynamic migration planner for moving layers between nodes under load.
548#[derive(Debug, Clone)]
549pub struct MigrationPlanner {
550    pub source_node: String,
551    pub target_node: String,
552    pub layers: Vec<usize>,
553    pub est_migration_time_ms: f32,
554}
555
556impl MigrationPlanner {
557    pub fn new(source: String, target: String, layers: Vec<usize>) -> Self {
558        let est_time = layers.len() as f32 * 45.0; // Baseline 45ms per layer migration
559        Self {
560            source_node: source,
561            target_node: target,
562            layers,
563            est_migration_time_ms: est_time,
564        }
565    }
566
567    /// Generate a safe handoff sequence for a migration.
568    pub fn generate_handoff_plan(&self) -> Vec<String> {
569        let mut steps = Vec::new();
570        let target_node = &self.target_node;
571        let source_node = &self.source_node;
572        steps.push(format!(
573            "PREPARE: Target node {target_node} allocating VRAM"
574        ));
575        steps.push(format!(
576            "STREAM: Moving layers {:?} to {target_node}",
577            self.layers
578        ));
579        steps.push(format!("VERIFY: Integrity check on {target_node}"));
580        steps.push(format!("COMMIT: Switch routing to {target_node}"));
581        steps.push(format!("CLEANUP: Free VRAM on source node {source_node}"));
582        steps
583    }
584}
585
586/// Trigger mechanism for dynamic rebalancing based on load/health.
587pub struct RebalanceTrigger {
588    pub min_imbalance_ratio: f32,
589    pub max_p95_latency_ms: f32,
590}
591
592impl Default for RebalanceTrigger {
593    fn default() -> Self {
594        Self {
595            min_imbalance_ratio: 0.25, // 25% drift
596            max_p95_latency_ms: 25.0,
597        }
598    }
599}
600
601impl RebalanceTrigger {
602    pub fn evaluate(
603        &self,
604        cluster: &ClusterState,
605        current_plan: &PlacementPlan,
606    ) -> Option<MigrationPlanner> {
607        let active_nodes = cluster.active_nodes();
608        if active_nodes.len() < 2 {
609            return None;
610        }
611
612        // Detect high-load nodes and high-headroom nodes
613        let mut overtaxed = None;
614        let mut available = None;
615
616        for metrics in active_nodes {
617            if metrics.delivery_ratio < 0.90 || metrics.avg_latency_us > 15.0 {
618                overtaxed = Some(metrics.name.clone());
619            } else if metrics.available_vram_gb > 8.0 {
620                available = Some(metrics.name.clone());
621            }
622        }
623
624        if let (Some(source), Some(target)) = (overtaxed, available) {
625            if source == target {
626                return None;
627            }
628
629            // Find layers to move from overtaxed node
630            let layers_to_move = current_plan
631                .assignments
632                .iter()
633                .find(|a| a.node_id == source)
634                .map(|a| (a.start_layer..a.end_layer).take(2).collect::<Vec<_>>())
635                .unwrap_or_default();
636
637            if !layers_to_move.is_empty() {
638                return Some(MigrationPlanner::new(source, target, layers_to_move));
639            }
640        }
641
642        None
643    }
644}
645
646#[cfg(test)]
647mod tests {
648    use super::*;
649    use crate::host::AccelerationMode;
650
651    fn sample_layers(count: usize, vram_gb: f32) -> Vec<LayerSpec> {
652        (0..count)
653            .map(|index| LayerSpec {
654                index,
655                vram_gb,
656                num_weights: 0,
657            })
658            .collect()
659    }
660
661    #[test]
662    fn greedily_places_layers_across_nodes() {
663        let nodes = vec![
664            NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
665            NodeResources::new("node-b", 12.0, 32.0, "8.6", None),
666        ];
667
668        let assignments = assign_layers_sequentially(&nodes, &sample_layers(33, 1.0)).unwrap();
669
670        assert_eq!(
671            assignments,
672            vec![
673                LayerAssignment {
674                    node_id: "node-a".into(),
675                    start_layer: 0,
676                    end_layer: 24,
677                    used_vram_gb: 24.0,
678                    num_layers: 24,
679                },
680                LayerAssignment {
681                    node_id: "node-b".into(),
682                    start_layer: 24,
683                    end_layer: 33,
684                    used_vram_gb: 9.0,
685                    num_layers: 9,
686                }
687            ]
688        );
689    }
690
691    #[test]
692    fn reports_insufficient_capacity() {
693        let nodes = vec![NodeResources::new("node-a", 2.0, 64.0, "8.9", None)];
694        let error = assign_layers_sequentially(&nodes, &sample_layers(3, 1.0)).unwrap_err();
695
696        assert!(error.contains("insufficient cluster VRAM"));
697    }
698
699    #[test]
700    fn selects_quantization_mode_from_delivery_ratio() {
701        assert_eq!(select_quantization_mode(0.98), QuantizationMode::None);
702        assert_eq!(select_quantization_mode(0.90), QuantizationMode::Int8);
703        assert_eq!(select_quantization_mode(0.75), QuantizationMode::Int4);
704    }
705
706    #[test]
707    fn placement_plan_summary() {
708        let plan = PlacementPlan::new(
709            vec![
710                LayerAssignment::new("node-a".into(), 0, 24, 24.0),
711                LayerAssignment::new("node-b".into(), 24, 33, 9.0),
712            ],
713            QuantizationMode::None,
714        );
715
716        let summary = plan.summary();
717        assert!(summary.contains("Total layers: 33"));
718        assert!(summary.contains("Full Precision"));
719    }
720
721    #[test]
722    fn rebalance_assignments_detects_capacity() {
723        let mut plan = PlacementPlan::new(
724            vec![LayerAssignment::new("node-a".into(), 0, 24, 24.0)],
725            QuantizationMode::None,
726        );
727
728        // Would need actual cluster state to test rebalancing
729        assert!(!rebalance_assignments(&ClusterState::default(), &mut plan));
730    }
731
732    #[test]
733    fn runtime_profile_chunks_large_assignments() {
734        let assignments = vec![LayerAssignment::new("node-a".into(), 0, 12, 12.0)];
735        let chunked = chunk_assignments_for_workers(assignments, 5);
736
737        assert_eq!(chunked.len(), 3);
738        assert_eq!(chunked[0].start_layer, 0);
739        assert_eq!(chunked[0].end_layer, 5);
740        assert_eq!(chunked[2].start_layer, 10);
741        assert_eq!(chunked[2].end_layer, 12);
742    }
743
744    #[test]
745    fn runtime_profile_assignment_keeps_trailing_zero_vram_layers() {
746        // Regression: assign_layers_chunked finalized its last pending chunk
747        // by checking `chunk_vram > 0.0`, so a trailing run of zero-VRAM
748        // layers left a non-empty pending chunk with chunk_vram == 0.0 and
749        // silently dropped those layers from the plan.
750        let nodes = vec![NodeResources::new("node-a", 24.0, 64.0, "8.9", None)];
751        let layers = sample_layers(3, 0.0);
752        let profile = RuntimeProfile {
753            node_resources: NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
754            logical_cores: 4,
755            recommended_workers: 1,
756            acceleration_mode: AccelerationMode::Generic,
757            gpu_backend: crate::host::GpuBackend::Cpu,
758            xdp_supported: false,
759            detection_source: String::from("test"),
760            probe_mode: crate::host::ProbeMode::Fast,
761        };
762
763        let assignments = assign_layers_with_runtime_profile(&nodes, &layers, &profile).unwrap();
764        let total: usize = assignments.iter().map(|a| a.num_layers).sum();
765        assert_eq!(total, 3, "all zero-VRAM layers must still be assigned");
766    }
767
768    #[test]
769    fn fault_tolerance_handles_all_failed_nodes_without_panicking() {
770        // Regression: dividing active_nodes' summed delivery ratio by
771        // active_nodes().len() produced NaN (or a TOCTOU-inconsistent
772        // average) when every registered node's metrics were marked Failed.
773        let cluster = ClusterState::new();
774        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
775        cluster.mark_failed("node-a");
776
777        let plan = assign_layers_with_fault_tolerance(&cluster, &sample_layers(4, 1.0)).unwrap();
778        assert_eq!(plan.quantization_mode, QuantizationMode::Int4);
779    }
780
781    #[test]
782    fn planning_tuning_scales_chunk_size_with_workers() {
783        let profile = RuntimeProfile {
784            node_resources: NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
785            logical_cores: 16,
786            recommended_workers: 8,
787            acceleration_mode: AccelerationMode::Gpu,
788            gpu_backend: crate::host::GpuBackend::Cuda,
789            xdp_supported: true,
790            detection_source: String::from("test"),
791            probe_mode: crate::host::ProbeMode::Fast,
792        };
793
794        let tuning = PlanningTuning::from_runtime_profile(&profile, 40);
795        assert!(tuning.max_layers_per_assignment <= 4);
796    }
797}