use crate::accelerator::ExecutionBackend;
use crate::cluster::{ClusterState, NodeMetrics};
use crate::host::{AccelerationMode, RuntimeProfile};
use crate::system_profile::SystemProfile;
use crate::watcher::ProfileChange;
use std::sync::Arc;
#[derive(Clone, Copy, Debug)]
pub struct LoadBalanceConfig {
pub lock_timeout_us: u64,
pub min_load_threshold: f32,
pub max_layers_per_assignment: usize,
pub max_concurrent_rebalances: usize,
}
impl Default for LoadBalanceConfig {
fn default() -> Self {
Self {
lock_timeout_us: 1000, min_load_threshold: 1.25,
max_layers_per_assignment: 100,
max_concurrent_rebalances: 1,
}
}
}
impl LoadBalanceConfig {
pub fn autotuned(profile: &RuntimeProfile) -> Self {
let base_timeout = match profile.acceleration_mode {
AccelerationMode::Gpu => 500,
AccelerationMode::Avx512 => 750,
_ => 1000,
};
let min_load_threshold = match profile.acceleration_mode {
AccelerationMode::Gpu => 1.15,
AccelerationMode::Avx512 => 1.05,
_ => 1.02,
};
let max_layers_per_assignment = profile.recommended_workers.max(1) * 2;
Self {
lock_timeout_us: base_timeout,
min_load_threshold,
max_layers_per_assignment,
max_concurrent_rebalances: profile.recommended_workers.clamp(1, 8),
}
}
}
#[derive(Clone, Debug)]
pub struct TensorSlice {
pub layer_range: (usize, usize), pub size_gb: f32,
pub num_weights: u32,
}
impl TensorSlice {
pub fn new(layer_range: (usize, usize), size_gb: f32) -> Self {
Self {
layer_range,
size_gb,
num_weights: 0,
}
}
}
#[derive(Clone, Debug)]
pub struct LoadDistributionPlan {
pub distributions: Vec<(String, Vec<TensorSlice>)>,
pub total_layers: usize,
pub participating_nodes: Vec<String>,
}
impl LoadDistributionPlan {
pub fn new(distributions: Vec<(String, Vec<TensorSlice>)>, total_layers: usize) -> Self {
let participating_nodes = distributions
.iter()
.map(|(node_id, _)| node_id.clone())
.collect();
Self {
distributions,
total_layers,
participating_nodes,
}
}
pub fn summary(&self) -> String {
let mut output = String::from("Load Distribution Plan\n");
output.push_str("======================\n\n");
for (node_id, slices) in &self.distributions {
output.push_str(&format!("Node: {node_id}\n"));
let total_size_gb: f32 = slices.iter().map(|s| s.size_gb).sum();
output.push_str(&format!(" Total size: {total_size_gb:.1} GB\n"));
output.push_str(&format!(
" Layers: {}-{}\n",
slices.first().map(|s| s.layer_range.0).unwrap_or(0),
slices.last().map(|s| s.layer_range.1).unwrap_or(0)
));
output.push('\n');
}
output.push_str(&format!("Total layers: {}\n", self.total_layers));
output.push_str(&format!("Nodes: {}\n", self.participating_nodes.join(", ")));
output
}
}
#[derive(Clone, Debug)]
pub struct LoadBalancer {
cluster: Arc<ClusterState>,
config: Arc<std::sync::Mutex<LoadBalanceConfig>>,
}
impl LoadBalancer {
pub fn new(cluster: Arc<ClusterState>, config: LoadBalanceConfig) -> Self {
Self {
cluster,
config: Arc::new(std::sync::Mutex::new(config)),
}
}
pub fn with_runtime_profile(cluster: Arc<ClusterState>, profile: &RuntimeProfile) -> Self {
Self::new(cluster, LoadBalanceConfig::autotuned(profile))
}
pub fn config(&self) -> LoadBalanceConfig {
self.config
.lock()
.ok()
.as_deref()
.copied()
.unwrap_or_default()
}
pub fn distribute_layers(
&self,
layers: &[crate::planning::LayerSpec],
) -> Result<LoadDistributionPlan, String> {
let nodes_snapshot = self.cluster.nodes_snapshot();
if nodes_snapshot.is_empty() {
return Err("no nodes available".into());
}
let mut sorted_nodes: Vec<&crate::protocol::NodeResources> =
nodes_snapshot.iter().collect();
sorted_nodes.sort_by(|a, b| {
b.vram_gb
.partial_cmp(&a.vram_gb)
.unwrap_or(std::cmp::Ordering::Equal)
});
let is_sorted = layers.windows(2).all(|w| w[0].index <= w[1].index);
let sorted_layers: std::borrow::Cow<'_, [crate::planning::LayerSpec]> = if is_sorted {
std::borrow::Cow::Borrowed(layers)
} else {
let mut all_layers = layers.to_vec();
all_layers.sort_by_key(|l| l.index);
std::borrow::Cow::Owned(all_layers)
};
let total_layer_count = sorted_layers.len();
let mut distributions = Vec::with_capacity(sorted_nodes.len());
let mut current_layer_idx = 0usize;
for node in &sorted_nodes {
if current_layer_idx >= sorted_layers.len() {
break;
}
let mut used_vram = 0.0f32;
let start_idx = current_layer_idx;
let mut end_idx = current_layer_idx;
for layer in &sorted_layers[current_layer_idx..] {
if used_vram + layer.vram_gb > node.vram_gb {
break;
}
end_idx += 1;
used_vram += layer.vram_gb;
}
if end_idx > start_idx {
let slices: Vec<TensorSlice> = sorted_layers[start_idx..end_idx]
.iter()
.map(|l| TensorSlice::new((l.index, l.index + 1), l.vram_gb))
.collect();
distributions.push((node.id.clone(), slices));
current_layer_idx = end_idx;
}
}
if current_layer_idx >= sorted_layers.len() {
Ok(LoadDistributionPlan::new(distributions, total_layer_count))
} else {
Err(format!(
"insufficient VRAM: {} layers remain",
sorted_layers.len() - current_layer_idx
))
}
}
pub fn distribute_layers_with_runtime_profile(
&self,
layers: &[crate::planning::LayerSpec],
profile: &RuntimeProfile,
) -> Result<LoadDistributionPlan, String> {
let cfg = self.config();
let plan = self.distribute_layers(layers)?;
let backend = ExecutionBackend::from_runtime_profile(profile);
let vector_bias = (backend.vector_width_bits / 128).max(1);
let max_layers_per_slice = cfg
.max_layers_per_assignment
.min(profile.recommended_workers.max(1).saturating_mul(2))
.min(backend.preferred_batch_size / vector_bias)
.max(1);
Ok(chunk_distribution_plan(plan, max_layers_per_slice))
}
pub fn rebalance(&self) -> bool {
let cfg = self.config();
let active_nodes = self.cluster.active_nodes();
if active_nodes.is_empty() {
return false;
}
let mut min_available = f32::MAX;
let mut max_available = 0.0_f32;
for node in &active_nodes {
min_available = min_available.min(node.available_vram_gb.max(0.0));
max_available = max_available.max(node.available_vram_gb.max(0.0));
}
let skew_ratio = if min_available <= f32::EPSILON {
if max_available > 0.0 {
f32::INFINITY
} else {
1.0
}
} else {
max_available / min_available
};
if skew_ratio >= cfg.min_load_threshold {
tracing::info!(
"Detected load skew ratio {:.2} (threshold {:.2})",
skew_ratio,
cfg.min_load_threshold
);
return true;
}
false
}
pub fn shed_load(&self) -> Vec<(String, String)> {
let cfg = self.config();
let mut transfers: Vec<(String, String)> = Vec::new();
let overloaded_nodes: Vec<_> = self
.cluster
.active_nodes()
.into_iter()
.filter(|m| m.available_vram_gb < m.total_vram_gb * 0.2)
.collect();
let underloaded_nodes: Vec<_> = self
.cluster
.active_nodes()
.into_iter()
.filter(|m| m.available_vram_gb > m.total_vram_gb * 0.5)
.collect();
for overloaded in &overloaded_nodes {
if transfers.len() >= cfg.max_concurrent_rebalances {
break;
}
if let Some(target) = self.find_best_target(underloaded_nodes.as_slice(), overloaded) {
transfers.push((overloaded.name.clone(), target.name.clone()));
}
}
transfers
}
fn find_best_target(
&self,
targets: &[NodeMetrics],
source: &NodeMetrics,
) -> Option<NodeMetrics> {
let mut best_target: Option<&NodeMetrics> = None;
let mut max_available = 0.0f32;
for target in targets {
if target.name == source.name {
continue;
}
if target.available_vram_gb > max_available {
max_available = target.available_vram_gb;
best_target = Some(target);
}
}
best_target.cloned()
}
pub fn reconfigure_from_system_profile(&self, profile: &SystemProfile) {
let rp: RuntimeProfile = profile.into();
if let Ok(mut guard) = self.config.lock() {
*guard = LoadBalanceConfig::autotuned(&rp);
}
}
pub fn subscribe_to_watcher(
self: &Arc<Self>,
watcher: &crate::watcher::SystemProfileWatcher,
) -> tokio::task::JoinHandle<()> {
let this = self.clone();
let mut rx = watcher.subscribe();
tokio::spawn(async move {
while let Ok(event) = rx.recv().await {
if let ProfileChange::Updated(profile) = event {
this.reconfigure_from_system_profile(&profile);
tracing::info!(
"load balancer re-tuned from profile: {} workers, {:?} accel",
profile.recommended_workers,
profile.acceleration_mode,
);
}
}
})
}
pub fn distribute_with_deadlock_prevention(
&self,
layers: &[crate::planning::LayerSpec],
) -> Result<LoadDistributionPlan, String> {
let cfg = self.config();
let start_time = std::time::Instant::now();
while start_time.elapsed().as_micros() < cfg.lock_timeout_us as u128 {
match self.distribute_layers(layers) {
Ok(plan) => return Ok(plan),
Err(_) => {
std::thread::sleep(std::time::Duration::from_millis(10));
}
}
}
Err("deadlock prevention timeout".into())
}
}
fn chunk_distribution_plan(
plan: LoadDistributionPlan,
max_layers_per_slice: usize,
) -> LoadDistributionPlan {
let slice_limit = max_layers_per_slice.max(1);
let distributions = plan
.distributions
.into_iter()
.map(|(node_id, slices)| {
let estimated_chunks = slices
.iter()
.map(|slice| {
let total_layers = slice.layer_range.1.saturating_sub(slice.layer_range.0);
total_layers.div_ceil(slice_limit).max(1)
})
.sum();
let mut chunked = Vec::with_capacity(estimated_chunks);
for slice in slices {
let total_layers = slice.layer_range.1.saturating_sub(slice.layer_range.0);
if total_layers <= slice_limit {
chunked.push(slice);
continue;
}
let avg_size = if total_layers == 0 {
0.0
} else {
slice.size_gb / total_layers as f32
};
let mut start = slice.layer_range.0;
while start < slice.layer_range.1 {
let end = (start + slice_limit).min(slice.layer_range.1);
let layers_in_chunk = end - start;
let mut chunk =
TensorSlice::new((start, end), avg_size * layers_in_chunk as f32);
chunk.num_weights = slice.num_weights;
chunked.push(chunk);
start = end;
}
}
(node_id, chunked)
})
.collect();
LoadDistributionPlan::new(distributions, plan.total_layers)
}
#[derive(Clone, Debug, Default)]
pub struct LoadStats {
pub total_layers_distributed: usize,
pub total_vram_used_gb: f32,
pub avg_load_balance_ratio: f32,
pub rebalancing_count: usize,
balance_ratio_initialized: bool,
}
impl LoadStats {
pub fn new() -> Self {
Self::default()
}
pub fn record_distribution(&mut self, num_layers: usize, vram_gb: f32) {
self.total_layers_distributed += num_layers;
self.total_vram_used_gb += vram_gb;
}
pub fn update_balance_ratio(&mut self, ratio: f32) {
if !self.balance_ratio_initialized {
self.avg_load_balance_ratio = ratio;
self.balance_ratio_initialized = true;
} else {
self.avg_load_balance_ratio = self.avg_load_balance_ratio * 0.9 + ratio * 0.1;
}
}
pub fn record_rebalancing(&mut self) {
self.rebalancing_count += 1;
}
pub fn report(&self) -> String {
format!(
"Load Statistics\n\
==========\n\
Total layers distributed: {}\n\
Total VRAM used: {:.1} GB\n\
Avg load balance ratio: {:.2}\n\
Rebalancing operations: {}",
self.total_layers_distributed,
self.total_vram_used_gb,
self.avg_load_balance_ratio,
self.rebalancing_count
)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cluster::ClusterState;
use crate::host::AccelerationMode;
use crate::protocol::NodeResources;
use std::sync::Arc;
fn sample_layers(count: usize, vram_gb: f32) -> Vec<crate::planning::LayerSpec> {
(0..count)
.map(|index| crate::planning::LayerSpec {
index,
vram_gb,
num_weights: 0,
})
.collect()
}
fn sample_nodes(count: usize, base_vram: f32) -> Vec<NodeResources> {
(0..count)
.map(|i| {
NodeResources::new(
format!("node-{i}"),
base_vram + (i as f32 * 6.0),
64.0,
"8.9".to_string(),
None,
)
})
.collect()
}
#[test]
fn load_balancer_distributes_layers() {
let cluster = ClusterState::new();
for node in sample_nodes(2, 24.0) {
cluster.register(node);
}
let layers = sample_layers(33, 1.0);
let balancer = LoadBalancer::new(Arc::new(cluster), LoadBalanceConfig::default());
let plan = balancer.distribute_layers(&layers).unwrap();
assert_eq!(plan.total_layers, 33);
}
#[test]
fn load_balancer_generates_summary() {
let cluster = ClusterState::new();
for node in sample_nodes(2, 24.0) {
cluster.register(node);
}
let layers = sample_layers(33, 1.0);
let balancer = LoadBalancer::new(Arc::new(cluster), LoadBalanceConfig::default());
let plan = balancer.distribute_layers(&layers).unwrap();
let summary = plan.summary();
assert!(summary.contains("Total layers: 33"));
}
#[test]
fn load_stats_records_distribution() {
let mut stats = LoadStats::new();
stats.record_distribution(24, 24.0);
stats.record_distribution(9, 9.0);
assert_eq!(stats.total_layers_distributed, 33);
assert!((stats.total_vram_used_gb - 33.0).abs() < 0.01);
}
#[test]
fn load_stats_updates_balance_ratio() {
let mut stats = LoadStats::new();
stats.update_balance_ratio(0.95);
assert!((stats.avg_load_balance_ratio - 0.95).abs() < 1e-6);
stats.update_balance_ratio(0.85);
assert!((stats.avg_load_balance_ratio - 0.94).abs() < 1e-6);
}
#[test]
fn load_stats_first_balance_ratio_not_diluted_by_prior_distribution() {
let mut stats = LoadStats::new();
stats.record_distribution(24, 24.0);
stats.update_balance_ratio(0.95);
assert!(
(stats.avg_load_balance_ratio - 0.95).abs() < 1e-6,
"first balance ratio sample must be set directly, got {}",
stats.avg_load_balance_ratio
);
}
#[test]
fn load_stats_reports() {
let mut stats = LoadStats::new();
stats.record_distribution(24, 24.0);
stats.update_balance_ratio(0.95);
let report = stats.report();
assert!(report.contains("Total layers distributed: 24"));
}
#[test]
fn autotuned_config_reflects_runtime_profile() {
let profile = RuntimeProfile {
node_resources: NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
logical_cores: 16,
recommended_workers: 6,
acceleration_mode: AccelerationMode::Gpu,
gpu_backend: crate::host::GpuBackend::Cuda,
xdp_supported: true,
detection_source: String::from("test"),
probe_mode: crate::host::ProbeMode::Fast,
};
let config = LoadBalanceConfig::autotuned(&profile);
assert_eq!(config.lock_timeout_us, 500);
assert_eq!(config.max_layers_per_assignment, 12);
assert_eq!(config.max_concurrent_rebalances, 6);
assert!(config.min_load_threshold > 1.0);
}
#[test]
fn autotuned_cpu_threshold_stays_above_one() {
let profile = RuntimeProfile {
node_resources: NodeResources::new("node-a", 0.0, 16.0, "generic", None),
logical_cores: 4,
recommended_workers: 2,
acceleration_mode: AccelerationMode::Generic,
gpu_backend: crate::host::GpuBackend::Cpu,
xdp_supported: false,
detection_source: String::from("test"),
probe_mode: crate::host::ProbeMode::Fast,
};
let config = LoadBalanceConfig::autotuned(&profile);
assert!(
config.min_load_threshold > 1.0,
"CPU-tier threshold {} must exceed skew_ratio's mathematical floor of 1.0",
config.min_load_threshold
);
}
#[test]
fn rebalance_does_not_always_trigger_on_balanced_cpu_cluster() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 0.0, 16.0, "generic", None));
cluster.register(NodeResources::new("node-b", 0.0, 16.0, "generic", None));
let profile = RuntimeProfile {
node_resources: NodeResources::new("node-a", 0.0, 16.0, "generic", None),
logical_cores: 4,
recommended_workers: 2,
acceleration_mode: AccelerationMode::Generic,
gpu_backend: crate::host::GpuBackend::Cpu,
xdp_supported: false,
detection_source: String::from("test"),
probe_mode: crate::host::ProbeMode::Fast,
};
let balancer = LoadBalancer::new(cluster, LoadBalanceConfig::autotuned(&profile));
assert!(
!balancer.rebalance(),
"a perfectly balanced (zero-VRAM) CPU cluster must not be reported as needing rebalance"
);
}
#[test]
fn runtime_profile_chunks_large_distribution_slices() {
let plan = LoadDistributionPlan::new(
vec![("node-a".into(), vec![TensorSlice::new((0, 10), 10.0)])],
10,
);
let chunked = chunk_distribution_plan(plan, 4);
assert_eq!(chunked.distributions[0].1.len(), 3);
assert_eq!(chunked.distributions[0].1[0].layer_range, (0, 4));
assert_eq!(chunked.distributions[0].1[2].layer_range, (8, 10));
}
#[test]
fn shed_load_moves_from_low_headroom_to_high_headroom_nodes() {
let cluster = ClusterState::new();
cluster.register(NodeResources::new("node-hot", 24.0, 64.0, "8.9", None));
cluster.register(NodeResources::new("node-cool", 24.0, 64.0, "8.9", None));
cluster.get_metrics_mut("node-hot", |m| {
m.available_vram_gb = 2.0;
m.total_vram_gb = 24.0;
});
cluster.get_metrics_mut("node-cool", |m| {
m.available_vram_gb = 20.0;
m.total_vram_gb = 24.0;
});
let balancer = LoadBalancer::new(Arc::new(cluster), LoadBalanceConfig::default());
let transfers = balancer.shed_load();
assert_eq!(transfers.len(), 1);
assert_eq!(transfers[0].0, "node-hot");
assert_eq!(transfers[0].1, "node-cool");
}
}