use std::collections::{BTreeSet, HashMap, HashSet};
use crate::network::metrics_parser::ParsedProcessRow;
pub const UNATTRIBUTED_USER: &str = "(unattributed)";
pub const UNATTRIBUTED_DISPLAY: &str = "?";
#[derive(Clone, Debug)]
pub struct GpuForAggregation {
pub host: String,
pub gpu_index: u32,
pub power_watts: f64,
}
#[derive(Clone, Debug)]
pub struct HostSnapshot {
pub host: String,
pub gpus: Vec<GpuForAggregation>,
pub processes: Vec<ParsedProcessRow>,
pub is_connected: bool,
}
#[derive(Clone, Debug)]
pub struct UserPerHost {
pub host: String,
pub gpu_indices: BTreeSet<u32>,
pub vram_bytes: u64,
pub power_watts: f64,
pub pid_count: usize,
pub top_command: String,
}
#[derive(Clone, Debug)]
pub struct UserAggregate {
pub user: String,
pub is_system: bool,
pub node_count: usize,
pub gpu_count: usize,
pub process_count: usize,
pub vram_bytes: u64,
pub power_watts: f64,
pub longest_seconds: u64,
pub top_command: String,
pub per_host: Vec<UserPerHost>,
}
#[derive(Clone, Debug, Default)]
pub struct UserAggregationResult {
pub users: Vec<UserAggregate>,
pub reporting_hosts: usize,
pub total_hosts: usize,
}
impl UserAggregationResult {
pub fn is_partial(&self) -> bool {
self.total_hosts > 0 && self.reporting_hosts < self.total_hosts
}
}
pub const SYSTEM_UID_THRESHOLD: u32 = 1000;
pub fn is_system_user(user: &str) -> bool {
if user == "root" {
return true;
}
if let Ok(uid) = user.parse::<u32>() {
return uid < SYSTEM_UID_THRESHOLD;
}
false
}
pub fn aggregate_users(snapshots: &[HostSnapshot]) -> UserAggregationResult {
let gpu_capacity_hint = snapshots.iter().map(|s| s.gpus.len()).sum::<usize>();
let mut total_vram_by_gpu: HashMap<(String, u32), u64> =
HashMap::with_capacity(gpu_capacity_hint);
for snap in snapshots {
for p in &snap.processes {
let entry = total_vram_by_gpu
.entry((snap.host.clone(), p.gpu_index))
.or_insert(0);
*entry = entry.saturating_add(p.gpu_memory_bytes);
}
}
let mut power_by_gpu: HashMap<(String, u32), f64> = HashMap::new();
for snap in snapshots {
for g in &snap.gpus {
power_by_gpu.insert((g.host.clone(), g.gpu_index), g.power_watts.max(0.0));
}
}
let mut user_scratch: HashMap<String, UserScratch> = HashMap::new();
let mut reporting_hosts: HashSet<&str> = HashSet::new();
for snap in snapshots {
if !snap.processes.is_empty() || snap.is_connected {
reporting_hosts.insert(snap.host.as_str());
}
for p in &snap.processes {
let canonical_user = if p.user.is_empty() {
UNATTRIBUTED_USER.to_string()
} else {
p.user.clone()
};
let scratch = user_scratch.entry(canonical_user).or_default();
scratch.absorb(p, snap.host.as_str());
}
}
let mut users: Vec<UserAggregate> = user_scratch
.into_iter()
.map(|(user, scratch)| scratch.finalize(user, &total_vram_by_gpu, &power_by_gpu))
.collect();
users.sort_by(|a, b| a.user.cmp(&b.user));
UserAggregationResult {
users,
reporting_hosts: reporting_hosts.len(),
total_hosts: snapshots.len(),
}
}
#[derive(Default, Clone)]
struct PerHostScratch {
gpu_indices: BTreeSet<u32>,
vram_bytes: u64,
pids: HashSet<u32>,
top_command_vram: u64,
top_command: String,
}
#[derive(Default)]
struct UserScratch {
touched_gpus: HashSet<(String, u32)>,
touched_pids: HashSet<(String, u32)>,
vram_bytes: u64,
vram_by_gpu: HashMap<(String, u32), u64>,
longest_seconds: u64,
top_command_vram: u64,
top_command: String,
per_host: HashMap<String, PerHostScratch>,
}
impl UserScratch {
fn absorb(&mut self, row: &ParsedProcessRow, host: &str) {
self.touched_gpus.insert((host.to_string(), row.gpu_index));
self.touched_pids.insert((host.to_string(), row.pid));
self.vram_bytes = self.vram_bytes.saturating_add(row.gpu_memory_bytes);
{
let entry = self
.vram_by_gpu
.entry((host.to_string(), row.gpu_index))
.or_insert(0);
*entry = entry.saturating_add(row.gpu_memory_bytes);
}
if row.start_time_seconds > self.longest_seconds {
self.longest_seconds = row.start_time_seconds;
}
if row.gpu_memory_bytes > self.top_command_vram {
self.top_command_vram = row.gpu_memory_bytes;
self.top_command = pick_display_command(row);
}
let ph = self.per_host.entry(host.to_string()).or_default();
ph.gpu_indices.insert(row.gpu_index);
ph.vram_bytes = ph.vram_bytes.saturating_add(row.gpu_memory_bytes);
ph.pids.insert(row.pid);
if row.gpu_memory_bytes > ph.top_command_vram {
ph.top_command_vram = row.gpu_memory_bytes;
ph.top_command = pick_display_command(row);
}
}
fn finalize(
self,
user: String,
total_vram_by_gpu: &HashMap<(String, u32), u64>,
power_by_gpu: &HashMap<(String, u32), f64>,
) -> UserAggregate {
let mut power_watts = 0.0_f64;
for ((host, gpu_index), user_vram) in &self.vram_by_gpu {
let total = total_vram_by_gpu
.get(&(host.clone(), *gpu_index))
.copied()
.unwrap_or(0);
if total == 0 {
continue;
}
let gpu_power = power_by_gpu
.get(&(host.clone(), *gpu_index))
.copied()
.unwrap_or(0.0);
let ratio = (*user_vram as f64) / (total as f64);
power_watts += (gpu_power * ratio).max(0.0);
}
if power_watts < 0.0 {
power_watts = 0.0;
}
let node_count: HashSet<&String> = self.touched_gpus.iter().map(|(h, _)| h).collect();
let is_system = is_system_user(&user);
let mut host_keys: Vec<String> = self.per_host.keys().cloned().collect();
host_keys.sort();
let per_host: Vec<UserPerHost> = host_keys
.into_iter()
.map(|host| {
let ph = self.per_host.get(&host).cloned().unwrap_or_default();
let mut host_power = 0.0_f64;
for g in &ph.gpu_indices {
let total = total_vram_by_gpu
.get(&(host.clone(), *g))
.copied()
.unwrap_or(0);
if total == 0 {
continue;
}
let user_vram = self
.vram_by_gpu
.get(&(host.clone(), *g))
.copied()
.unwrap_or(0);
let gpu_power = power_by_gpu
.get(&(host.clone(), *g))
.copied()
.unwrap_or(0.0);
host_power += (gpu_power * (user_vram as f64) / (total as f64)).max(0.0);
}
if host_power < 0.0 {
host_power = 0.0;
}
UserPerHost {
host,
gpu_indices: ph.gpu_indices,
vram_bytes: ph.vram_bytes,
power_watts: host_power,
pid_count: ph.pids.len(),
top_command: ph.top_command,
}
})
.collect();
UserAggregate {
user,
is_system,
node_count: node_count.len(),
gpu_count: self.touched_gpus.len(),
process_count: self.touched_pids.len(),
vram_bytes: self.vram_bytes,
power_watts,
longest_seconds: self.longest_seconds,
top_command: self.top_command,
per_host,
}
}
}
fn pick_display_command(row: &ParsedProcessRow) -> String {
if !row.command.is_empty() {
row.command.clone()
} else if !row.name.is_empty() {
row.name.clone()
} else {
String::new()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum UserSortKey {
#[default]
User,
Memory,
Power,
Nodes,
Longest,
}
pub fn sort_users(users: &mut [UserAggregate], key: UserSortKey) {
use std::cmp::Ordering;
users.sort_by(|a, b| {
let primary = match key {
UserSortKey::User => a.user.cmp(&b.user),
UserSortKey::Memory => b.vram_bytes.cmp(&a.vram_bytes),
UserSortKey::Power => b
.power_watts
.partial_cmp(&a.power_watts)
.unwrap_or(Ordering::Equal),
UserSortKey::Nodes => b.node_count.cmp(&a.node_count),
UserSortKey::Longest => b.longest_seconds.cmp(&a.longest_seconds),
};
primary.then_with(|| a.user.cmp(&b.user))
});
}
pub fn format_longest(seconds: u64) -> String {
if seconds == 0 {
return "—".to_string();
}
let days = seconds / 86_400;
let rem = seconds % 86_400;
let hours = rem / 3_600;
let minutes = (rem % 3_600) / 60;
let secs = rem % 60;
if days > 0 {
format!("{days}d {hours:02}:{minutes:02}:{secs:02}")
} else {
format!("{hours:02}:{minutes:02}:{secs:02}")
}
}
#[cfg(test)]
mod tests {
use super::*;
fn row(
pid: u32,
user: &str,
gpu_index: u32,
gpu_memory_bytes: u64,
start_time_seconds: u64,
command: &str,
) -> ParsedProcessRow {
ParsedProcessRow {
host: String::new(), pid,
user: user.to_string(),
command: command.to_string(),
name: command.to_string(),
gpu_index,
gpu_uuid: format!("GPU-{gpu_index}"),
gpu_memory_bytes,
cpu_pct_tenths: 0,
start_time_seconds,
}
}
fn gpu(host: &str, gpu_index: u32, power_watts: f64) -> GpuForAggregation {
GpuForAggregation {
host: host.to_string(),
gpu_index,
power_watts,
}
}
#[test]
fn empty_input_produces_empty_result() {
let result = aggregate_users(&[]);
assert!(result.users.is_empty());
assert!(!result.is_partial());
}
fn snap(
host: &str,
gpus: Vec<GpuForAggregation>,
processes: Vec<ParsedProcessRow>,
) -> HostSnapshot {
HostSnapshot {
host: host.to_string(),
gpus,
processes,
is_connected: true,
}
}
fn snap_disconnected(
host: &str,
gpus: Vec<GpuForAggregation>,
processes: Vec<ParsedProcessRow>,
) -> HostSnapshot {
HostSnapshot {
host: host.to_string(),
gpus,
processes,
is_connected: false,
}
}
#[test]
fn same_pid_on_two_hosts_counts_as_two_processes() {
let snapshots = vec![
snap(
"a",
vec![gpu("a", 0, 100.0)],
vec![row(42, "alice", 0, 1000, 100, "train")],
),
snap(
"b",
vec![gpu("b", 0, 100.0)],
vec![row(42, "alice", 0, 2000, 200, "train")],
),
];
let result = aggregate_users(&snapshots);
assert_eq!(result.users.len(), 1);
let u = &result.users[0];
assert_eq!(u.user, "alice");
assert_eq!(u.process_count, 2, "(host, pid) must disambiguate");
assert_eq!(u.node_count, 2);
assert_eq!(u.vram_bytes, 3000);
assert_eq!(u.longest_seconds, 200);
}
#[test]
fn root_is_marked_as_system_user() {
let snapshots = vec![snap(
"h",
vec![gpu("h", 0, 50.0)],
vec![row(1, "root", 0, 0, 10, "containerd-shim")],
)];
let result = aggregate_users(&snapshots);
assert_eq!(result.users.len(), 1);
assert!(result.users[0].is_system);
}
#[test]
fn numeric_system_uid_is_marked_as_system_user() {
assert!(is_system_user("0"));
assert!(is_system_user("999"));
assert!(!is_system_user("1000"));
assert!(!is_system_user("alice"));
}
#[test]
fn user_spanning_multiple_gpus_accumulates_correctly() {
let snapshots = vec![
snap(
"h1",
vec![gpu("h1", 0, 200.0), gpu("h1", 1, 300.0)],
vec![
row(100, "bob", 0, 1_000, 50, "a"),
row(100, "bob", 1, 2_000, 50, "a"),
],
),
snap(
"h2",
vec![gpu("h2", 0, 400.0)],
vec![row(200, "bob", 0, 5_000, 100, "b")],
),
];
let result = aggregate_users(&snapshots);
assert_eq!(result.users.len(), 1);
let u = &result.users[0];
assert_eq!(u.node_count, 2);
assert_eq!(u.gpu_count, 3, "two GPUs on h1 + one on h2");
assert_eq!(u.process_count, 2);
assert_eq!(u.vram_bytes, 8_000);
}
#[test]
fn oldest_start_time_wins_longest() {
let snapshots = vec![snap(
"h",
vec![gpu("h", 0, 10.0)],
vec![
row(1, "alice", 0, 10, 500, "a"),
row(2, "alice", 0, 20, 1_500_000, "b"),
row(3, "alice", 0, 30, 100, "c"),
],
)];
let result = aggregate_users(&snapshots);
assert_eq!(result.users[0].longest_seconds, 1_500_000);
}
#[test]
fn top_command_is_the_owner_of_the_largest_vram_row() {
let snapshots = vec![snap(
"h",
vec![gpu("h", 0, 10.0)],
vec![
row(1, "alice", 0, 1_000, 0, "small"),
row(2, "alice", 0, 9_000, 0, "big"),
row(3, "alice", 0, 5_000, 0, "medium"),
],
)];
let result = aggregate_users(&snapshots);
assert_eq!(result.users[0].top_command, "big");
}
#[test]
fn partial_coverage_is_detected() {
let snapshots = vec![
snap(
"a",
vec![gpu("a", 0, 100.0)],
vec![row(1, "alice", 0, 10, 0, "x")],
),
snap_disconnected("b", vec![gpu("b", 0, 100.0)], vec![]),
];
let result = aggregate_users(&snapshots);
assert!(result.is_partial());
assert_eq!(result.reporting_hosts, 1);
assert_eq!(result.total_hosts, 2);
}
#[test]
fn idle_connected_host_is_not_flagged_as_partial() {
let snapshots = vec![
snap(
"a",
vec![gpu("a", 0, 100.0)],
vec![row(1, "alice", 0, 10, 0, "x")],
),
snap("b", vec![gpu("b", 0, 100.0)], vec![]),
];
let result = aggregate_users(&snapshots);
assert!(
!result.is_partial(),
"idle-but-connected host must not trigger the partial chip"
);
assert_eq!(result.reporting_hosts, 2);
assert_eq!(result.total_hosts, 2);
}
#[test]
fn power_approximation_clamps_negatives_to_zero() {
let snapshots = vec![snap(
"h",
vec![gpu("h", 0, -50.0)],
vec![row(1, "alice", 0, 1000, 0, "x")],
)];
let result = aggregate_users(&snapshots);
assert_eq!(result.users[0].power_watts, 0.0);
}
#[test]
fn power_approximation_weights_by_vram_share() {
let snapshots = vec![snap(
"h",
vec![gpu("h", 0, 400.0)],
vec![
row(1, "alice", 0, 7_000, 0, "a"),
row(2, "bob", 0, 3_000, 0, "b"),
],
)];
let result = aggregate_users(&snapshots);
let alice = result.users.iter().find(|u| u.user == "alice").unwrap();
let bob = result.users.iter().find(|u| u.user == "bob").unwrap();
let total = alice.power_watts + bob.power_watts;
assert!((total - 400.0).abs() < 1e-6, "got {total}");
assert!(
(alice.power_watts - 280.0).abs() < 1e-6,
"got {}",
alice.power_watts
);
}
#[test]
fn missing_user_label_becomes_unattributed() {
let snapshots = vec![snap(
"h",
vec![gpu("h", 0, 100.0)],
vec![row(1, "", 0, 100, 0, "x")],
)];
let result = aggregate_users(&snapshots);
assert_eq!(result.users.len(), 1);
assert_eq!(result.users[0].user, UNATTRIBUTED_USER);
}
#[test]
fn per_host_breakdown_sums_to_the_aggregate_power() {
let snapshots = vec![
snap(
"h1",
vec![gpu("h1", 0, 200.0), gpu("h1", 1, 300.0)],
vec![
row(1, "alice", 0, 1_000, 0, "a"),
row(1, "alice", 1, 3_000, 0, "a"),
],
),
snap(
"h2",
vec![gpu("h2", 0, 400.0)],
vec![row(2, "alice", 0, 5_000, 0, "b")],
),
];
let result = aggregate_users(&snapshots);
let u = &result.users[0];
let per_host_sum: f64 = u.per_host.iter().map(|h| h.power_watts).sum();
assert!((per_host_sum - u.power_watts).abs() < 1e-6);
let mut hosts: Vec<&str> = u.per_host.iter().map(|h| h.host.as_str()).collect();
hosts.sort();
assert_eq!(hosts, vec!["h1", "h2"]);
}
#[test]
fn sort_users_by_memory_descending() {
let mut u1 = UserAggregate {
user: "alice".into(),
is_system: false,
node_count: 1,
gpu_count: 1,
process_count: 1,
vram_bytes: 100,
power_watts: 0.0,
longest_seconds: 0,
top_command: "".into(),
per_host: vec![],
};
let mut u2 = UserAggregate {
user: "bob".into(),
is_system: false,
node_count: 1,
gpu_count: 1,
process_count: 1,
vram_bytes: 300,
power_watts: 0.0,
longest_seconds: 0,
top_command: "".into(),
per_host: vec![],
};
u1.longest_seconds = 10;
u2.longest_seconds = 20;
let mut v = vec![u1, u2];
sort_users(&mut v, UserSortKey::Memory);
assert_eq!(v[0].user, "bob", "highest VRAM first");
}
#[test]
fn format_longest_renders_days_when_over_a_day() {
let s = format_longest(97_927);
assert!(s.contains("1d"), "expected days prefix, got {s}");
}
#[test]
fn format_longest_returns_dash_for_zero() {
assert_eq!(format_longest(0), "—");
}
#[test]
fn aggregation_handles_large_cluster_quickly() {
let users = ["alice", "bob", "carol", "dave"];
let mut snaps = Vec::with_capacity(100);
for h in 0..100 {
let host = format!("h{h}");
let gpus = (0..8)
.map(|i| GpuForAggregation {
host: host.clone(),
gpu_index: i,
power_watts: 200.0,
})
.collect();
let procs = (0..50)
.map(|p| ParsedProcessRow {
host: host.clone(),
pid: p as u32,
user: users[(p as usize) % users.len()].to_string(),
command: "python".to_string(),
name: "python".to_string(),
gpu_index: (p % 8) as u32,
gpu_uuid: format!("GPU-{h}-{}", p % 8),
gpu_memory_bytes: 1_000_000_000,
cpu_pct_tenths: 0,
start_time_seconds: (p * 7) as u64,
})
.collect();
snaps.push(HostSnapshot {
host,
gpus,
processes: procs,
is_connected: true,
});
}
let start = std::time::Instant::now();
let result = aggregate_users(&snaps);
let elapsed = start.elapsed();
assert_eq!(result.users.len(), users.len());
assert!(
elapsed.as_millis() < 500,
"aggregate_users took {elapsed:?} for 100 hosts × 50 procs"
);
}
#[test]
#[ignore = "adversarial stress — run manually with `cargo test -- --ignored`"]
fn aggregation_survives_adversarial_50k_per_host() {
let rows_per_host = 50_000;
let host_count = 10;
let users = ["alice", "bob", "carol", "dave", "eve", "frank"];
let mut snaps = Vec::with_capacity(host_count);
for h in 0..host_count {
let host = format!("h{h}");
let gpus = (0..8)
.map(|i| GpuForAggregation {
host: host.clone(),
gpu_index: i,
power_watts: 400.0,
})
.collect();
let procs = (0..rows_per_host)
.map(|p| ParsedProcessRow {
host: host.clone(),
pid: p as u32,
user: users[(p as usize) % users.len()].to_string(),
command: "python".to_string(),
name: "python".to_string(),
gpu_index: (p % 8) as u32,
gpu_uuid: format!("GPU-{h}-{}", p % 8),
gpu_memory_bytes: 1024,
cpu_pct_tenths: 0,
start_time_seconds: (p * 7) as u64,
})
.collect();
snaps.push(HostSnapshot {
host,
gpus,
processes: procs,
is_connected: true,
});
}
let start = std::time::Instant::now();
let result = aggregate_users(&snaps);
let elapsed = start.elapsed();
assert_eq!(result.users.len(), users.len());
assert!(
elapsed.as_secs() < 5,
"adversarial aggregation took {elapsed:?} for {host_count}×{rows_per_host} rows — possible O(n^2) regression"
);
}
}