use rustc_hash::FxHashMap;
use crate::graph::{
AffinityRule, DomainAssignment, DomainSnapshot, GraphError,
PlannedTopology, QueueDistribution, QueueStrategy,
RuntimeGraph, TopologyPlanner,
};
#[derive(Debug)]
pub struct ExtremePlanner;
#[derive(Debug, Clone)]
struct NumaTopology {
cpu_cores: u32,
memory_mb: u32,
nic_queues: u32,
}
#[derive(Debug, Clone, Copy, PartialEq)]
struct OptimizationScore {
numa_affinity: f64,
cache_localness: f64,
bandwidth_optimality: f64,
load_balance: f64,
resource_efficiency: f64,
total: f64,
}
impl OptimizationScore {
fn new(numa: f64, cache: f64, bw: f64, lb: f64, re: f64) -> Self {
let total = numa * 0.30 + cache * 0.25 + bw * 0.20 + lb * 0.15 + re * 0.10;
Self {
numa_affinity: numa,
cache_localness: cache,
bandwidth_optimality: bw,
load_balance: lb,
resource_efficiency: re,
total,
}
}
}
impl ExtremePlanner {
pub fn plan<const MAX_NODES: usize, const MAX_MAPPINGS: usize>(
&self,
graph: &RuntimeGraph<MAX_NODES, MAX_MAPPINGS>,
snapshots: &[DomainSnapshot],
) -> Result<PlannedTopology, GraphError> {
let numa_map = self.build_numa_topology(graph);
let mut best_assignments: Vec<DomainAssignment> = Vec::new();
let mut best_distributions: Vec<QueueDistribution> = Vec::new();
let mut best_score = OptimizationScore::new(0.0, 0.0, 0.0, 0.0, 0.0);
let mut any_succeeded = false;
let mut last_err: Option<GraphError> = None;
for strategy in [
QueueStrategy::RssHash,
QueueStrategy::RoundRobin,
QueueStrategy::NumaLocal,
] {
match self.try_strategy(graph, snapshots, strategy, &numa_map) {
Ok((assignments, distributions, score)) => {
any_succeeded = true;
if score.total > best_score.total {
best_score = score;
best_assignments = assignments;
best_distributions = distributions;
}
}
Err(e) => {
tracing::warn!("extreme_planner: strategy {strategy:?} failed, trying next: {e}");
last_err = Some(e);
}
}
}
if !any_succeeded {
return Err(last_err.unwrap_or_else(|| {
GraphError::InvalidTopology("extreme_planner: all queue strategies failed")
}));
}
Ok(PlannedTopology {
domain_assignments: best_assignments,
queue_distributions: best_distributions,
})
}
fn try_strategy<const MAX_NODES: usize, const MAX_MAPPINGS: usize>(
&self,
graph: &RuntimeGraph<MAX_NODES, MAX_MAPPINGS>,
snapshots: &[DomainSnapshot],
strategy: QueueStrategy,
numa_map: &FxHashMap<u32, NumaTopology>,
) -> Result<(Vec<DomainAssignment>, Vec<QueueDistribution>, OptimizationScore), GraphError> {
let planner = TopologyPlanner;
let topology = planner.plan(snapshots, graph, strategy)?;
let score = self.evaluate_score(
&topology.domain_assignments,
&topology.queue_distributions,
graph,
numa_map,
);
Ok((topology.domain_assignments, topology.queue_distributions, score))
}
fn build_numa_topology<const MAX_NODES: usize, const MAX_MAPPINGS: usize>(
&self,
graph: &RuntimeGraph<MAX_NODES, MAX_MAPPINGS>,
) -> FxHashMap<u32, NumaTopology> {
let mut map: FxHashMap<u32, NumaTopology> = FxHashMap::default();
for node in graph.iter_nodes() {
if let Some(top) = map.get_mut(&node.numa_node) {
top.cpu_cores = top.cpu_cores.saturating_add(node.cpu_cores);
top.memory_mb = top.memory_mb.saturating_add(node.memory_mb);
top.nic_queues = top.nic_queues.saturating_add(node.nic_queues);
} else {
map.insert(
node.numa_node,
NumaTopology {
cpu_cores: node.cpu_cores,
memory_mb: node.memory_mb,
nic_queues: node.nic_queues,
},
);
}
}
map
}
fn evaluate_score<const MAX_NODES: usize, const MAX_MAPPINGS: usize>(
&self,
assignments: &[DomainAssignment],
distributions: &[QueueDistribution],
graph: &RuntimeGraph<MAX_NODES, MAX_MAPPINGS>,
numa_map: &FxHashMap<u32, NumaTopology>,
) -> OptimizationScore {
let numa_affinity = self.score_numa_affinity(assignments);
let cache_localness = self.score_cache_localness(assignments);
let bandwidth_optimality = self.score_bandwidth(assignments, distributions);
let load_balance = self.score_load_balance(assignments);
let resource_efficiency = self.score_resource_efficiency(assignments, graph, numa_map);
OptimizationScore::new(
numa_affinity,
cache_localness,
bandwidth_optimality,
load_balance,
resource_efficiency,
)
}
fn score_numa_affinity(&self, assignments: &[DomainAssignment]) -> f64 {
let mut total_workers = 0usize;
let mut numa_local_workers = 0usize;
for assignment in assignments {
let mut numa_groups: FxHashMap<u32, usize> = FxHashMap::default();
for worker in &assignment.workers {
*numa_groups.entry(worker.numa_node).or_insert(0) += 1;
}
total_workers += assignment.workers.len();
if let Some(max) = numa_groups.values().max() {
numa_local_workers += *max;
}
}
if total_workers == 0 {
return 0.0;
}
numa_local_workers as f64 / total_workers as f64
}
fn score_cache_localness(&self, assignments: &[DomainAssignment]) -> f64 {
let mut score = 0.0f64;
let mut count = 0usize;
for assignment in assignments {
if assignment.workers.is_empty() {
continue;
}
let mut distinct_numa: Vec<u32> = Vec::new();
for worker in &assignment.workers {
if !distinct_numa.contains(&worker.numa_node) {
distinct_numa.push(worker.numa_node);
}
}
score += 1.0 / distinct_numa.len() as f64;
count += 1;
}
if count == 0 {
return 0.0;
}
score / count as f64
}
fn score_bandwidth(
&self,
_assignments: &[DomainAssignment],
distributions: &[QueueDistribution],
) -> f64 {
let mut score = 0.0f64;
let mut count = 0usize;
for dist in distributions {
for mapping in &dist.mappings {
if matches!(mapping.affinity, AffinityRule::NumaLocal) {
score += 1.0;
} else if matches!(mapping.affinity, AffinityRule::CpuLocal) {
score += 0.8;
} else {
score += 0.5;
}
count += 1;
}
}
if count == 0 {
return 0.0;
}
score / count as f64
}
fn score_load_balance(&self, assignments: &[DomainAssignment]) -> f64 {
let mut worker_counts: Vec<usize> = assignments
.iter()
.map(|a| a.workers.len())
.collect();
if worker_counts.is_empty() {
return 0.0;
}
worker_counts.sort();
let min = worker_counts.first().copied().unwrap_or(0) as f64;
let max = worker_counts.last().copied().unwrap_or(0) as f64;
if max == 0.0 {
return 0.0;
}
min / max
}
fn score_resource_efficiency<const MAX_NODES: usize, const MAX_MAPPINGS: usize>(
&self,
assignments: &[DomainAssignment],
graph: &RuntimeGraph<MAX_NODES, MAX_MAPPINGS>,
_numa_map: &FxHashMap<u32, NumaTopology>,
) -> f64 {
let mut total_cpu = 0.0f64;
let mut total_mem = 0.0f64;
let mut used_cpu = 0.0f64;
let mut used_mem = 0.0f64;
for assignment in assignments {
let (per_cpu, per_mem, _) = assignment.domain.default_requirements();
used_cpu += (assignment.workers.len() as f64) * (per_cpu as f64);
used_mem += (assignment.workers.len() as f64) * (per_mem as f64);
let mut seen_nodes: Vec<u64> = Vec::new();
for worker in &assignment.workers {
if seen_nodes.contains(&worker.node_id) {
continue;
}
seen_nodes.push(worker.node_id);
if let Some(node) = graph.find_node(worker.node_id) {
total_cpu += node.cpu_cores as f64;
total_mem += node.memory_mb as f64;
}
}
}
if total_cpu == 0.0 || total_mem == 0.0 {
return 0.0;
}
let cpu_util = (used_cpu / total_cpu).min(1.0);
let mem_util = (used_mem / total_mem).min(1.0);
(cpu_util + mem_util) / 2.0
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::graph::{ExecutionDomain, ResourceNode, WorkerAssignment};
fn make_graph() -> RuntimeGraph<16, 64> {
let mut graph = RuntimeGraph::new();
for i in 0..6 {
graph
.add_node(ResourceNode::new(
i,
ExecutionDomain::DataPlane,
16,
16384,
8,
))
.unwrap();
}
for i in 6..8 {
graph
.add_node(ResourceNode::new(
i,
ExecutionDomain::ControlPlane,
4,
8192,
4,
))
.unwrap();
}
graph
}
#[test]
fn test_extreme_planner_basic() {
let graph = make_graph();
let planner = ExtremePlanner;
let snapshots = vec![
DomainSnapshot::from_domain(ExecutionDomain::DataPlane, 3),
DomainSnapshot::from_domain(ExecutionDomain::ControlPlane, 1),
];
let topology = planner.plan(&graph, &snapshots).unwrap();
assert_eq!(topology.domain_count(), 2);
assert!(topology.total_workers() >= 4);
}
#[test]
fn test_extreme_planner_score_comparison() {
let graph = make_graph();
let planner = ExtremePlanner;
let snapshots = vec![DomainSnapshot::from_domain(ExecutionDomain::DataPlane, 2)];
let topology = planner.plan(&graph, &snapshots).unwrap();
assert_eq!(topology.domain_count(), 1);
let dp = topology
.find_assignment(ExecutionDomain::DataPlane)
.unwrap();
assert_eq!(dp.workers.len(), 2);
}
#[test]
fn test_numa_topology_build() {
let graph = make_graph();
let planner = ExtremePlanner;
let numa_map = planner.build_numa_topology(&graph);
assert!(!numa_map.is_empty());
}
#[test]
fn test_optimization_scores() {
let score = OptimizationScore::new(0.9, 0.8, 0.85, 0.95, 0.9);
assert!(score.total > 0.0);
assert!(score.numa_affinity > 0.0);
}
#[test]
fn test_score_numa_affinity_perfect() {
let assignments = vec![DomainAssignment {
domain: ExecutionDomain::DataPlane,
workers: vec![
WorkerAssignment {
worker_id: 0,
node_id: 1,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
},
WorkerAssignment {
worker_id: 1,
node_id: 2,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
},
],
}];
let planner = ExtremePlanner;
let score = planner.score_numa_affinity(&assignments);
assert!((score - 1.0).abs() < 0.001);
}
#[test]
fn test_score_load_balance() {
let assignments = vec![
DomainAssignment {
domain: ExecutionDomain::DataPlane,
workers: vec![WorkerAssignment {
worker_id: 0,
node_id: 0,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
}],
},
DomainAssignment {
domain: ExecutionDomain::ControlPlane,
workers: vec![WorkerAssignment {
worker_id: 0,
node_id: 0,
domain: ExecutionDomain::ControlPlane,
numa_node: 0,
queue_ids: Vec::new(),
}],
},
];
let planner = ExtremePlanner;
let score = planner.score_load_balance(&assignments);
assert!((score - 1.0).abs() < 0.001);
}
#[test]
fn test_empty_assignments_score() {
let planner = ExtremePlanner;
let score = planner.score_numa_affinity(&[]);
assert_eq!(score, 0.0);
let score = planner.score_load_balance(&[]);
assert_eq!(score, 0.0);
}
#[test]
fn test_score_cache_localness_numa_concentration() {
let planner = ExtremePlanner;
let concentrated = vec![DomainAssignment {
domain: ExecutionDomain::DataPlane,
workers: vec![
WorkerAssignment {
worker_id: 0,
node_id: 1,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
},
WorkerAssignment {
worker_id: 1,
node_id: 2,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
},
],
}];
assert!((planner.score_cache_localness(&concentrated) - 1.0).abs() < 0.001);
let scattered = vec![DomainAssignment {
domain: ExecutionDomain::DataPlane,
workers: vec![
WorkerAssignment {
worker_id: 0,
node_id: 1,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
},
WorkerAssignment {
worker_id: 1,
node_id: 2,
domain: ExecutionDomain::DataPlane,
numa_node: 1,
queue_ids: Vec::new(),
},
],
}];
assert!((planner.score_cache_localness(&scattered) - 0.5).abs() < 0.001);
}
#[test]
fn test_score_resource_efficiency_real_utilization() {
let planner = ExtremePlanner;
let mut graph: RuntimeGraph<4, 8> = RuntimeGraph::new();
graph
.add_node(ResourceNode::new(1, ExecutionDomain::DataPlane, 8, 8192, 8))
.unwrap();
let full = vec![DomainAssignment {
domain: ExecutionDomain::DataPlane,
workers: vec![
WorkerAssignment {
worker_id: 0,
node_id: 1,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
},
WorkerAssignment {
worker_id: 1,
node_id: 1,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
},
],
}];
let numa_map = planner.build_numa_topology(&graph);
let score = planner.score_resource_efficiency(&full, &graph, &numa_map);
assert!((score - 1.0).abs() < 0.001, "满载效率应为 1.0,实际 {score}");
let half = vec![DomainAssignment {
domain: ExecutionDomain::DataPlane,
workers: vec![WorkerAssignment {
worker_id: 0,
node_id: 1,
domain: ExecutionDomain::DataPlane,
numa_node: 0,
queue_ids: Vec::new(),
}],
}];
let score = planner.score_resource_efficiency(&half, &graph, &numa_map);
assert!((score - 0.5).abs() < 0.001, "半载效率应为 0.5,实际 {score}");
}
#[test]
fn test_extreme_planner_queue_ids_unique_across_strategies() {
let graph = make_graph();
let planner = ExtremePlanner;
let snapshots = vec![DomainSnapshot::from_domain(ExecutionDomain::DataPlane, 3)];
let topology = planner.plan(&graph, &snapshots).unwrap();
for dist in &topology.queue_distributions {
let mut ids: Vec<u32> = dist.mappings.iter().map(|m| m.queue_id).collect();
ids.sort_unstable();
ids.dedup();
assert_eq!(ids.len(), dist.mappings.len());
}
}
}