Worker managed cpu utilization

crates/hyperqueue/src/dashboard/ui/screens/cluster/worker/cpu_util_table.rs

@@ -1,14 +1,16 @@
 use crate::common::format::human_size;
+use crate::dashboard::ui::screens::cluster::worker::UtilizationRenderMode;
 use ratatui::layout::{Constraint, Rect};
 use ratatui::style::Style;
 use ratatui::widgets::{Cell, Row, Table};
 use std::cmp;
 use tako::hwstats::MemoryStats;
+use tako::resources::ResourceIndex;
 
 use crate::dashboard::ui::styles;
 use crate::dashboard::ui::terminal::DashboardFrame;
 use crate::dashboard::ui::widgets::progressbar::{
-    ProgressPrintStyle, get_progress_bar_color, render_progress_bar_at,
+    ProgressPrintStyle, get_progress_bar_color, get_progress_bar_cpu_color, render_progress_bar_at,
 };
 use crate::dashboard::utils::calculate_average;
 
@@ -19,6 +21,8 @@ const CPU_METER_WIDTH: u8 = CPU_METER_PROGRESSBAR_WIDTH + 4;
 pub fn render_cpu_util_table(
     cpu_util_list: &[f64],
     mem_util: &MemoryStats,
+    used_cpus: &[ResourceIndex],
+    util_render_mode: &UtilizationRenderMode,
     rect: Rect,
     frame: &mut DashboardFrame,
     table_style: Style,
@@ -31,26 +35,32 @@ pub fn render_cpu_util_table(
     let width = constraints.len();
     let height = (cpu_util_list.len() as f64 / width as f64).ceil() as usize;
 
-    let mut rows: Vec<Vec<(f64, usize)>> = vec![vec![]; height];
+    let mut rows: Vec<Vec<(f64, usize, bool)>> = vec![vec![]; height];
     for (position, &cpu_util) in cpu_util_list.iter().enumerate() {
         let row = position % height;
-        rows[row].push((cpu_util, position));
+        let used = used_cpus.contains(&ResourceIndex::new(position as u32));
+        rows[row].push((cpu_util, position, used));
     }
 
     let rows: Vec<Row> = rows
         .into_iter()
         .map(|targets| {
             let columns: Vec<Cell> = targets
                 .into_iter()
-                .map(|(cpu_util, position)| {
+                .map(|(cpu_util, position, used)| {
                     let progress = cpu_util / 100.00;
+                    let style = match util_render_mode {
+                        UtilizationRenderMode::Global => get_progress_bar_color(progress),
+                        UtilizationRenderMode::Worker => get_progress_bar_cpu_color(progress, used),
+                    };
+
                     Cell::from(render_progress_bar_at(
                         Some(format!("{position:>3} ")),
                         progress,
                         CPU_METER_PROGRESSBAR_WIDTH,
                         ProgressPrintStyle::default(),
                     ))
-                    .style(get_progress_bar_color(progress))
+                    .style(style)
                 })
                 .collect();
             Row::new(columns)

crates/hyperqueue/src/dashboard/ui/screens/cluster/worker/mod.rs

@@ -12,6 +12,7 @@ use crate::dashboard::ui::widgets::text::draw_text;
 use crossterm::event::{KeyCode, KeyEvent};
 use ratatui::layout::{Constraint, Direction, Layout, Rect};
 use tako::hwstats::MemoryStats;
+use tako::resources::{CPU_RESOURCE_NAME, ResourceIndex};
 use tako::{JobTaskId, WorkerId};
 
 mod cpu_util_table;
@@ -26,6 +27,7 @@ pub struct WorkerDetail {
     worker_tasks_table: TasksTable,
 
     utilization: Option<Utilization>,
+    utilization_render_mode: UtilizationRenderMode,
 }
 
 impl Default for WorkerDetail {
@@ -36,13 +38,36 @@ impl Default for WorkerDetail {
             worker_config_table: Default::default(),
             worker_tasks_table: TasksTable::non_interactive(),
             utilization: None,
+            utilization_render_mode: UtilizationRenderMode::Global,
+        }
+    }
+}
+
+enum UtilizationRenderMode {
+    Global,
+    Worker,
+}
+
+impl UtilizationRenderMode {
+    fn next(&mut self) {
+        *self = match self {
+            UtilizationRenderMode::Global => UtilizationRenderMode::Worker,
+            UtilizationRenderMode::Worker => UtilizationRenderMode::Global,
+        }
+    }
+
+    fn next_text(&self) -> &str {
+        match self {
+            UtilizationRenderMode::Global => "Show worker CPU utilization",
+            UtilizationRenderMode::Worker => "Show global CPU utilization",
         }
     }
 }
 
 struct Utilization {
     cpu: Vec<f64>,
     memory: MemoryStats,
+    used_cpus: Vec<ResourceIndex>,
 }
 
 impl WorkerDetail {
@@ -66,12 +91,24 @@ impl WorkerDetail {
             frame,
             style_header_text(),
         );
-        draw_text("<backspace>: Back", layout.footer, frame, style_footer());
+
+        draw_text(
+            format!(
+                "<backspace>: Back <c>: {}",
+                self.utilization_render_mode.next_text()
+            )
+            .as_str(),
+            layout.footer,
+            frame,
+            style_footer(),
+        );
 
         if let Some(util) = &self.utilization {
             render_cpu_util_table(
                 &util.cpu,
                 &util.memory,
+                &util.used_cpus,
+                &self.utilization_render_mode,
                 layout.current_utilization,
                 frame,
                 table_style_deselected(),
@@ -95,21 +132,50 @@ impl WorkerDetail {
         if let Some(worker_id) = self.worker_id {
             self.utilization_history.update(data, worker_id);
 
-            if let Some((cpu_util, mem_util)) = data
+            if let Some(overview) = data
                 .workers()
                 .query_worker_overview_at(worker_id, data.current_time())
-                .and_then(|overview| overview.item.hw_state.as_ref())
-                .map(|hw_state| {
-                    (
-                        &hw_state.state.cpu_usage.cpu_per_core_percent_usage,
-                        &hw_state.state.memory_usage,
-                    )
-                })
             {
-                self.utilization = Some(Utilization {
-                    cpu: cpu_util.iter().map(|&v| v as f64).collect(),
-                    memory: mem_util.clone(),
-                });
+                let worker_used_cpus: Vec<ResourceIndex> = match self.utilization_render_mode {
+                    UtilizationRenderMode::Worker => overview
+                        .item
+                        .running_tasks
+                        .iter()
+                        .flat_map(|(_id, task_resource_alloc)| {
+                            task_resource_alloc
+                                .resources
+                                .iter()
+                                .filter_map(|resource_alloc| {
+                                    if resource_alloc.resource == CPU_RESOURCE_NAME {
+                                        Some(
+                                            resource_alloc
+                                                .indices
+                                                .iter()
+                                                .map(|(index, _)| *index),
+                                        )
+                                    } else {
+                                        None
+                                    }
+                                })
+                        })
+                        .flatten()
+                        .collect(),
+                    UtilizationRenderMode::Global => vec![],
+                };
+
+                if let Some(hw_state) = overview.item.hw_state.as_ref() {
+                    self.utilization = Some(Utilization {
+                        cpu: hw_state
+                            .state
+                            .cpu_usage
+                            .cpu_per_core_percent_usage
+                            .iter()
+                            .map(|&v| v as f64)
+                            .collect(),
+                        memory: hw_state.state.memory_usage.clone(),
+                        used_cpus: worker_used_cpus,
+                    })
+                }
             }
 
             let tasks_info: Vec<(JobTaskId, &TaskInfo)> =
@@ -127,6 +193,7 @@ impl WorkerDetail {
     pub fn handle_key(&mut self, key: KeyEvent) {
         match key.code {
             KeyCode::Backspace => self.worker_tasks_table.clear_selection(),
+            KeyCode::Char('c') => self.utilization_render_mode.next(),
             _ => self.worker_tasks_table.handle_key(key),
         }
     }

crates/hyperqueue/src/dashboard/ui/widgets/progressbar.rs

@@ -30,6 +30,23 @@ pub fn get_progress_bar_color(progress: f64) -> Style {
     }
 }
 
+pub fn get_progress_bar_cpu_color(progress: f64, used: bool) -> Style {
+    let (used_color, unused_color) = match progress {
+        p if p <= GREEN_THRESHOLD => (Color::Green, Color::Cyan),
+        p if p <= YELLOW_THRESHOLD => (Color::Yellow, Color::LightBlue),
+        _ => (Color::Red, Color::Blue),
+    };
+
+    let color = if used { used_color } else { unused_color };
+
+    Style {
+        fg: Some(color),
+        bg: None,
+        add_modifier: Modifier::empty(),
+        sub_modifier: Modifier::empty(),
+    }
+}
+
 /**
  * Creates a string progress bar for 0 < progress < 1
  */