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(
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 let assignments = assign_layers_chunked(&nodes, layers, tuning.max_layers_per_assignment)?;
446
447 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 let quantization_mode = select_quantization_mode(total_delivery_ratio);
472
473 Ok(PlacementPlan::new(assignments, quantization_mode))
474}
475
476pub fn rebalance_assignments(cluster: &ClusterState, plan: &mut PlacementPlan) -> bool {
478 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 continue;
486 }
487
488 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 potential_layers > 2 {
497 needs_rebalance = true;
498 break;
499 }
500 }
501 }
502 }
503
504 needs_rebalance
505}
506
507#[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 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 metrics.set_streaming_layers(start_layer, end_layer);
525
526 Some(metrics)
527}
528
529pub fn calculate_cluster_health(cluster: &ClusterState) -> (f32, usize, Vec<String>) {
531 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 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#[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; Self {
583 source_node: source,
584 target_node: target,
585 layers,
586 est_migration_time_ms: est_time,
587 }
588 }
589
590 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
609pub 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, 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 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 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 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 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 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}