From 1ede6ad93bfc9d33869683987cff3b08e4f187ec Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 7 Oct 2023 19:34:55 -0700
Subject: [PATCH] Export CPU temperature graph

---
 scripts/wk/graph.py    | 75 +++++++++++++++++++++++++++++++++++++++++-
 scripts/wk/hw/cpu.py   | 10 +++++-
 scripts/wk/hw/diags.py |  9 ++++-
 3 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/scripts/wk/graph.py b/scripts/wk/graph.py
index 842833a9..91ade848 100644
--- a/scripts/wk/graph.py
+++ b/scripts/wk/graph.py
@@ -45,8 +45,80 @@ THRESH_GREAT =        750 * 1024**2
 
 
 # Functions
+def export_cpu_graph(cpu_description, log_dir, sensor_history):
+  """Exports PNG graph using matplotlib."""
+  lines = {}
+  offset_labels = []
+  out_path = pathlib.Path(f'{log_dir}/cpu_tests.png')
+  run_averages = {}
+
+  # Safety check
+  if not sensor_history:
+    raise RuntimeError('No sensor_data available.')
+
+  # Prep
+  offset = 0
+  for run_label, sensor_data in sensor_history:
+    all_run_temps = []
+    offset_labels.append((offset, run_label))
+    run_length = 0
+
+    for adapter in sensor_data.get('CPUTemps', {}).values():
+      for source, data in adapter.items():
+        y_values = data['Temps']
+        all_run_temps.extend(y_values)
+        run_length = max(run_length, len(y_values))
+        if source not in lines:
+          lines[source] = []
+        lines[source].extend(y_values)
+    try:
+      run_averages[run_label] = {
+        'Start': offset,
+        'End': offset+run_length,
+        'Temp': sum(all_run_temps) / len(all_run_temps),
+        }
+    except ZeroDivisionError:
+      # Ignore
+      pass
+    offset += run_length
+
+  # Build graph
+  _, ax = pyplot.subplots(
+    dpi=72,
+    figsize=list(x*2 for x in GRAPH_FIGURE_SIZE),
+    layout='constrained',
+    )
+  ax.set_title(cpu_description)
+  for label, data in lines.items():
+    ax.plot(data, label=label)
+  #prev_label = 'Idle' # Always skip Idle
+  #for offset, label in offset_labels:
+  #  if label == prev_label:
+  #    continue
+  #  color = 'r' if label in ('Prime95', 'Sysbench') else 'b'
+  #  if label == 'Cooldown':
+  #    label = ''
+  #  ax.axvline(x=offset, color=color, label=label)
+  #  #ax.axvline(x=offset, color=color)
+  #  prev_label = label
+  for run_label, data in run_averages.items():
+    if run_label not in ('Prime95', 'Sysbench'):
+      continue
+    ax.axhline(
+      y = data['Temp'],
+      color = 'orange' if run_label == 'Sysbench' else 'red',
+      label = f'{run_label} (Avg)',
+      linestyle = '--',
+      )
+  ax.legend()
+  pyplot.savefig(out_path)
+
+  # Done
+  return out_path
+
+
 def export_io_graph(disk, log_dir, read_rates):
-  """Exports PNG graph using gnuplot."""
+  """Exports PNG graph using matplotlib."""
 
   # Safety check
   if not read_rates:
@@ -74,6 +146,7 @@ def export_io_graph(disk, log_dir, read_rates):
   ax.legend()
   pyplot.savefig(out_path)
 
+  # Done
   return out_path
 
 
diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py
index 74708d60..8d495ff5 100644
--- a/scripts/wk/hw/cpu.py
+++ b/scripts/wk/hw/cpu.py
@@ -9,6 +9,7 @@ from typing import TextIO
 
 from wk import exe
 from wk.cfg.hw import CPU_TEMPS
+from wk.graph import export_cpu_graph
 from wk.os.mac import set_fans as macos_set_fans
 from wk.std import PLATFORM
 from wk.ui import ansi
@@ -20,7 +21,7 @@ SysbenchType = tuple[subprocess.Popen, TextIO]
 
 
 # Functions
-def check_cooling_results(sensors, test_object) -> None:
+def check_cooling_results(cpu_description, log_dir, sensors, test_object) -> None:
   """Check cooling result via sensor data."""
   idle_temp = sensors.get_cpu_temp('Idle')
   cooldown_temp = sensors.get_cpu_temp('Cooldown')
@@ -84,6 +85,13 @@ def check_cooling_results(sensors, test_object) -> None:
       *report_labels, only_cpu=True, include_avg_for=average_labels):
     test_object.report.append(f'  {line}')
 
+  # Export graph
+  export_cpu_graph(
+    cpu_description = cpu_description,
+    log_dir = log_dir,
+    sensor_history = sensors.history,
+    )
+
 
 def check_mprime_results(test_obj, working_dir) -> None:
   """Check mprime log files and update test_obj."""
diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py
index caef5f4d..1bebbb20 100644
--- a/scripts/wk/hw/diags.py
+++ b/scripts/wk/hw/diags.py
@@ -413,7 +413,12 @@ def cpu_test_cooling(state: State, test_object, test_mode=False) -> None:
   if test_object.disabled:
     return
 
-  hw_cpu.check_cooling_results(state.sensors, test_object)
+  hw_cpu.check_cooling_results(
+    state.system.cpu_description,
+    state.log_dir,
+    state.sensors,
+    test_object,
+    )
   state.update_progress_file()
 
 
@@ -461,6 +466,7 @@ def cpu_test_mprime(state: State, test_object, test_mode=False) -> None:
     aborted = True
 
   # Stop Prime95
+  state.sensors.save_average_temps(temp_label='Cooldown', seconds=5)
   hw_cpu.stop_mprime(proc)
 
   # Get cooldown temp
@@ -532,6 +538,7 @@ def cpu_test_sysbench(state: State, test_object, test_mode=False) -> None:
     LOG.error('Failed to find sysbench process', exc_info=True)
   except KeyboardInterrupt:
     aborted = True
+  state.sensors.save_average_temps(temp_label='Cooldown', seconds=5)
   hw_cpu.stop_sysbench(proc, filehandle)
 
   # Get cooldown temp