lib/cpupower_helper.py

# Copyright 2022 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""CPU Power related helpers."""

import contextlib
import logging
from pathlib import Path
from typing import Iterator, Optional

from chromite.lib import chromite_config
from chromite.lib import cros_build_lib
from chromite.lib import osutils


_AUTO_SET_GOV_CONTENT = (
    '# Delete this file to turn off automatic performance governor switch.\n')

_CPU_PATH = Path('/sys/devices/system/cpu')

# cpupower command with default options used to set scaling governor for CPUs
_CPUPOWER_CMD = ['cpupower', '-c', 'all', 'frequency-set', '-g']


def _FetchActiveCpuGovernor() -> Optional[str]:
  """Scans and returns the active governor if all cpus use it.

  Returns:
    The active governor or None if there is no single active governor.
  """
  cpufreq = _CPU_PATH / 'cpufreq'
  governors = {
      x.read_text(encoding='utf-8')
      for x in cpufreq.glob('policy*/scaling_governor')
  }
  if len(governors) != 1:
    logging.warning('Too many active CPU governors; refusing to use '
                    "'performance'.")
    return None
  return next(iter(governors))


def _AutoSetGovStickyConfigUpdate(perf_governor: bool, sticky: bool) -> None:
  """Create or remove config file based on performance governor sticky request.

  Args:
    perf_governor: If true, switch the CPU governors to performance.
    sticky: If true, cache the switch request for subsequent runs.
  """
  if not sticky:
    return

  if not perf_governor:
    logging.info('Future runs will respect --autosetgov')
    osutils.SafeUnlink(chromite_config.AUTO_SET_GOV_CONFIG)
  elif not chromite_config.AUTO_SET_GOV_CONFIG.exists():
    logging.info('Future runs will *always* use --autosetgov')
    chromite_config.initialize()
    chromite_config.AUTO_SET_GOV_CONFIG.write_text(
        _AUTO_SET_GOV_CONTENT, encoding='utf-8')


@contextlib.contextmanager
def ModifyCpuGovernor(perf_governor: bool, sticky: bool) -> Iterator[None]:
  """A context manager to switch the CPU governor policy to performance.

  This context manager will switch the CPU governors to the 'performance'
  governor and will revert the policy to the default governor when it exits.
  When the sticky option is selected, the governor will be automatically
  switched to performance for the subsequent runs. To clear this automatic
  switching, user must pass the perf_governor arg as false with sticky argument
  set to true.

  Args:
    perf_governor: If true, switch the CPU governors to performance.
    sticky: If true, cache the switch request for subsequent runs.

  Yields:
    Iterator.
  """

  _AutoSetGovStickyConfigUpdate(perf_governor, sticky)

  old_governor = None
  new_governor = None
  # If the config path is present we switch CPUs to performance governor.
  if chromite_config.AUTO_SET_GOV_CONFIG.exists():
    perf_governor = True
  if perf_governor:
    new_governor = 'performance'
  # Check if performance governor is supported by the CPU.
  governor_path = _CPU_PATH / 'cpu0' / 'cpufreq' / 'scaling_available_governors'
  try:
    governors = governor_path.read_text(encoding='utf-8').split()
  except FileNotFoundError as e:
    logging.debug('Error reading CPU scaling governor file: %s', e)
    governors = []

  try:
    if 'performance' in governors:
      old_governor = _FetchActiveCpuGovernor()

      # Switch the CPU governors to performance if they have the same active
      # governor and if the active governor is not performance.
      if new_governor and old_governor and new_governor != old_governor:
        logging.info('Temporarily setting CPU governors to %s', new_governor)
        perf_cmd = _CPUPOWER_CMD + [new_governor]
        try:
          cros_build_lib.sudo_run(
              perf_cmd, print_cmd=False, capture_output=True)
        except cros_build_lib.RunCommandError as e:
          logging.warning('Error switching CPU governors: %s', e)
      elif old_governor == 'powersave':
        logging.warning("Current CPU governor set to 'powersave' which can "
                        'slow down builds.')
        logging.warning('Use --autosetgov to automatically (and '
                        "temporarily) switch to 'performance'.")

    yield
  finally:
    if new_governor and old_governor and new_governor != old_governor:
      logging.info('Restoring CPU governors to %s', old_governor)
      restore_cmd = _CPUPOWER_CMD + [old_governor]
      try:
        cros_build_lib.sudo_run(
            restore_cmd, print_cmd=False, capture_output=True)
      except cros_build_lib.RunCommandError as e:
        logging.warning('Error restoring CPU governors: %s', e)