1use crate::accelerator::ExecutionBackend;
9use crate::cluster::ClusterState;
10use crate::cluster::NodeStatus;
11use crate::host::{AccelerationMode, RuntimeProfile};
12use crate::protocol::NodeResources;
13
14pub const DELIVERY_RATIO_INT8_THRESHOLD: f32 = 0.95;
16pub const DELIVERY_RATIO_INT4_THRESHOLD: f32 = 0.80;
17
18#[derive(Clone, Debug, PartialEq)]
20pub struct LayerSpec {
21 pub index: usize,
23 pub vram_gb: f32,
25 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#[derive(Clone, Debug, PartialEq)]
41pub struct LayerAssignment {
42 pub node_id: String,
44 pub start_layer: usize,
46 pub end_layer: usize,
48 pub used_vram_gb: f32,
50 pub num_layers: usize,
52}
53
54impl LayerAssignment {
55 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 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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
78pub enum QuantizationMode {
79 #[default]
81 None,
82 Int8,
84 Int4,
86}
87
88#[derive(Clone, Debug, Default)]
90pub struct PlacementPlan {
91 pub assignments: Vec<LayerAssignment>,
93 pub quantization_mode: QuantizationMode,
95 pub total_layers: usize,
97 pub participating_nodes: Vec<String>,
99}
100
101impl PlacementPlan {
102 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 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
143pub struct PlanningTuning {
144 pub max_layers_per_assignment: usize,
146}
147
148impl PlanningTuning {
149 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
172pub 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
183pub 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 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 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 if let Some(assignment) = current_assignment {
239 assignments.push(assignment);
240 }
241
242 Ok(assignments)
243}
244
245pub 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
284fn 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 let mut chunk_start = 0usize;
306 let mut chunk_vram = 0.0f32;
307
308 for (i, layer) in layers.iter().enumerate() {
309 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 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 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
364pub 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
374pub 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 let assignments = assign_layers_sequentially(&nodes, layers)?;
387
388 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 let total_delivery_ratio = if active_count > 0 {
417 sum_delivery / active_count as f32
418 } else {
419 0.0
420 };
421
422 let quantization_mode = select_quantization_mode(total_delivery_ratio);
424
425 Ok(PlacementPlan::new(assignments, quantization_mode))
426}
427
428pub 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
453pub fn rebalance_assignments(cluster: &ClusterState, plan: &mut PlacementPlan) -> bool {
455 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 continue;
463 }
464
465 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 potential_layers > 2 {
474 needs_rebalance = true;
475 break;
476 }
477 }
478 }
479 }
480
481 needs_rebalance
482}
483
484#[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 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 metrics.set_streaming_layers(start_layer, end_layer);
502
503 Some(metrics)
504}
505
506pub fn calculate_cluster_health(cluster: &ClusterState) -> (f32, usize, Vec<String>) {
508 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 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#[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; Self {
560 source_node: source,
561 target_node: target,
562 layers,
563 est_migration_time_ms: est_time,
564 }
565 }
566
567 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
586pub 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, 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 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 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 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 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 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}