WizardKit/scripts/wk/hw/cpu.py

"""WizardKit: CPU test functions"""
# vim: sts=2 sw=2 ts=2

import logging
import re
import subprocess

from typing import TextIO

from wk import ansi, exe
from wk.cfg.hw import CPU_FAILURE_TEMP
from wk.os.mac import set_fans as macos_set_fans
from wk.std import PLATFORM
from wk.ui import tmux


# STATIC VARIABLES
LOG = logging.getLogger(__name__)
SysbenchType = tuple[subprocess.Popen, TextIO]


# Functions
def check_cooling_results(test_obj, sensors, run_sysbench=False) -> None:
  """Check cooling results and update test_obj."""
  max_temp = sensors.cpu_max_temp()
  temp_labels = ['Idle', 'Max', 'Cooldown']
  if run_sysbench:
    temp_labels.append('Sysbench')

  # Check temps
  if not max_temp:
    test_obj.set_status('Unknown')
  elif max_temp >= CPU_FAILURE_TEMP:
    test_obj.failed = True
    test_obj.set_status('Failed')
  elif 'Aborted' not in test_obj.status:
    test_obj.passed = True
    test_obj.set_status('Passed')

  # Add temps to report
  for line in sensors.generate_report(*temp_labels, only_cpu=True):
    test_obj.report.append(f'  {line}')


def check_mprime_results(test_obj, working_dir) -> None:
  """Check mprime log files and update test_obj."""
  passing_lines = {}
  warning_lines = {}

  def _read_file(log_name) -> list[str]:
    """Read file and split into lines, returns list."""
    lines = []
    try:
      with open(f'{working_dir}/{log_name}', 'r', encoding='utf-8') as _f:
        lines = _f.readlines()
    except FileNotFoundError:
      # File may be missing on older systems
      lines = []

    return lines

  # results.txt (check if failed)
  for line in _read_file('results.txt'):
    line = line.strip()
    if re.search(r'(error|fail)', line, re.IGNORECASE):
      warning_lines[line] = None

  # prime.log (check if passed)
  for line in _read_file('prime.log'):
    line = line.strip()
    match = re.search(
      r'(completed.*(\d+) errors, (\d+) warnings)', line, re.IGNORECASE)
    if match:
      if int(match.group(2)) + int(match.group(3)) > 0:
        # Errors and/or warnings encountered
        warning_lines[match.group(1).capitalize()] = None
      else:
        # No errors/warnings
        passing_lines[match.group(1).capitalize()] = None

  # Update status
  if warning_lines:
    test_obj.failed = True
    test_obj.set_status('Failed')
  elif passing_lines and 'Aborted' not in test_obj.status:
    test_obj.passed = True
    test_obj.set_status('Passed')
  else:
    test_obj.set_status('Unknown')

  # Update report
  for line in passing_lines:
    test_obj.report.append(f'  {line}')
  for line in warning_lines:
    test_obj.report.append(ansi.color_string(f'  {line}', 'YELLOW'))
  if not (passing_lines or warning_lines):
    test_obj.report.append(ansi.color_string('  Unknown result', 'YELLOW'))


def start_mprime(working_dir, log_path) -> subprocess.Popen:
  """Start mprime and save filtered output to log, returns Popen object."""
  set_apple_fan_speed('max')
  proc_mprime = subprocess.Popen(
    ['mprime', '-t'],
    cwd=working_dir,
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    )
  proc_grep = subprocess.Popen(
    'grep --ignore-case --invert-match --line-buffered stress.txt'.split(),
    stdin=proc_mprime.stdout,
    stdout=subprocess.PIPE,
    )
  proc_mprime.stdout.close()
  save_nsbr = exe.NonBlockingStreamReader(proc_grep.stdout)
  exe.start_thread(
    save_nsbr.save_to_file,
    args=(proc_grep, log_path),
    )

  # Return objects
  return proc_mprime


def start_sysbench(sensors, sensors_out, log_path, pane) -> SysbenchType:
  """Start sysbench, returns tuple with Popen object and file handle."""
  set_apple_fan_speed('max')
  sysbench_cmd = [
    'sysbench',
    f'--threads={exe.psutil.cpu_count()}',
    '--cpu-max-prime=1000000000',
    'cpu',
    'run',
    ]

  # Restart background monitor for Sysbench
  sensors.stop_background_monitor()
  sensors.start_background_monitor(
    sensors_out,
    alt_max='Sysbench',
    thermal_action=('killall', 'sysbench', '-INT'),
    )

  # Update bottom pane
  tmux.respawn_pane(pane, watch_file=log_path, watch_cmd='tail')

  # Start sysbench
  filehandle_sysbench = open(
    log_path, 'a', encoding='utf-8',
    )
  proc_sysbench = exe.popen_program(sysbench_cmd, stdout=filehandle_sysbench)

  # Done
  return (proc_sysbench, filehandle_sysbench)


def set_apple_fan_speed(speed) -> None:
  """Set Apple fan speed."""
  cmd = None

  # Check
  if speed not in ('auto', 'max'):
    raise RuntimeError(f'Invalid speed {speed}')

  # Set cmd
  if PLATFORM == 'Darwin':
    try:
      macos_set_fans(speed)
    except (RuntimeError, ValueError, subprocess.CalledProcessError) as err:
      LOG.error('Failed to set fans to %s', speed)
      LOG.error('Error: %s', err)
      #ui.print_error(f'Failed to set fans to {speed}')
      #for line in str(err).splitlines():
      #  ui.print_warning(f'  {line.strip()}')
  elif PLATFORM == 'Linux':
    cmd = ['apple-fans', speed]
    exe.run_program(cmd, check=False)


def stop_mprime(proc_mprime) -> None:
  """Stop mprime gracefully, then forcefully as needed."""
  proc_mprime.terminate()
  try:
    proc_mprime.wait(timeout=5)
  except subprocess.TimeoutExpired:
    proc_mprime.kill()
  set_apple_fan_speed('auto')


def stop_sysbench(proc_sysbench, filehandle_sysbench) -> None:
  """Stop sysbench."""
  proc_sysbench.terminate()
  try:
    proc_sysbench.wait(timeout=5)
  except subprocess.TimeoutExpired:
    proc_sysbench.kill()
  filehandle_sysbench.flush()
  filehandle_sysbench.close()
  set_apple_fan_speed('auto')


if __name__ == '__main__':
  print("This file is not meant to be called directly.")