crosperf: Wait for CPU cooldown and setup governor.

Based on dut_config['cooldown_temp'] and dut_config['cooldown_time']
wait for CPU to cooldown to a temperature 'cooldown_temp' or until
timeout 'cooldown_time' expires.
While waiting change CPU governor to a power saving mode.

With 'cooldown_time' = 0 don't wait.

Setup CPU governor based on dut_config['governor'] for a benchmark run.

Additionally function helper DutWrapper is added in SuiteRunner
which simplifies remote DUT calls.

BUG=chromium:966514
TEST=Unitest and local HW tests passed.

Change-Id: I2b946b570e7a8d8c4e2321f8bb8de3be2c3162d0
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1790131
Tested-by: Denis Nikitin <denik@chromium.org>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
Reviewed-by: George Burgess <gbiv@chromium.org>

3 files changed, 409 insertions, 221 deletions

diff --git a/crosperf/schedv2.py b/crosperf/schedv2.py
index 1ef0b938..a7983d9f 100644
--- a/crosperf/schedv2.py
+++ b/crosperf/schedv2.py
@@ -198,6 +198,8 @@ class DutWorker(Thread):
     intel_pstate = self._sched.get_experiment().intel_pstate
     # Firstly, handle benchmarkruns that have cache hit.
     br = self._sched.get_cached_benchmark_run()
+    # Total wait time for cooling down.
+    total_waittime = 0
     while br:
       try:
         self._stat_annotation = 'finishing cached {}'.format(br)
@@ -236,7 +238,11 @@ class DutWorker(Thread):
 
           # Execute the br.
           self._execute_benchmark_run(br)
+          total_waittime += br.suite_runner.GetCooldownWaitTime()
+          br.suite_runner.ResetCooldownWaitTime()
     finally:
+      self._logger.LogOutput('Total wait time for cooldown: %d min'
+                             % (total_waittime / 60))
       self._stat_annotation = 'finished'
       # Thread finishes. Notify scheduler that I'm done.
       self._sched.dut_worker_finished(self)
diff --git a/crosperf/suite_runner.py b/crosperf/suite_runner.py
index d5d6c2c7..5b677949 100644
--- a/crosperf/suite_runner.py
+++ b/crosperf/suite_runner.py
@@ -5,6 +5,7 @@
 
 """SuiteRunner defines the interface from crosperf to test script."""
 
+from __future__ import division
 from __future__ import print_function
 
 import os
@@ -62,11 +63,48 @@ class SuiteRunner(object):
     self.log_level = log_level
     self._ce = cmd_exec or command_executer.GetCommandExecuter(
         self.logger, log_level=self.log_level)
+    # DUT command executer.
+    # Will be initialized and used within Run.
     self._ct = cmd_term or command_executer.CommandTerminator()
     self.enable_aslr = enable_aslr
     self.dut_config = dut_config
+    self.cooldown_wait_time = 0
+
+  def ResetCooldownWaitTime(self):
+    self.cooldown_wait_time = 0
+
+  def GetCooldownWaitTime(self):
+    return self.cooldown_wait_time
+
+  def DutWrapper(self, machine_name, chromeos_root):
+    """Wrap DUT parameters inside.
+
+    Eventially CommandExecuter will reqiure only one
+    argument - command.
+    """
+
+    def RunCommandOnDut(command, ignore_status=False):
+      ret, msg, err_msg = self._ce.CrosRunCommandWOutput(
+          command, machine=machine_name, chromeos_root=chromeos_root)
+
+      if ret:
+        err_msg = ('Command execution on DUT %s failed.\n'
+                   'Failing command: %s\n'
+                   'returned %d\n'
+                   'Error message: %s' % (machine_name, command, ret, err_msg))
+        if ignore_status:
+          self.logger.LogError(err_msg)
+        else:
+          self.logger.LogFatal(err_msg)
+
+      return ret, msg, err_msg
+
+    return RunCommandOnDut
 
   def Run(self, machine, label, benchmark, test_args, profiler_args):
+    if not label.skylab:
+      # Initialize command executer on DUT.
+      run_on_dut = self.DutWrapper(machine, label.chromeos_root)
     for i in range(0, benchmark.retries + 1):
       if label.skylab:
         # TODO: need to migrate DisableASLR and PinGovernorExecutionFrequencies
@@ -79,11 +117,36 @@ class SuiteRunner(object):
         # Unless the user turns on ASLR in the flag, we first disable ASLR
         # before running the benchmarks
         if not self.enable_aslr:
-          self.DisableASLR(machine, label.chromeos_root)
-        self.PinGovernorExecutionFrequencies(machine, label.chromeos_root)
-        self.SetupCpuUsage(machine, label.chromeos_root)
+          self.DisableASLR(run_on_dut)
+
+        if self.dut_config['cooldown_time']:
+          # Setup power conservative mode for effective cool down.
+          # Set ignore status since powersave may no be available
+          # on all platforms and we are going to handle it.
+          ret = self.SetCpuGovernor('powersave', run_on_dut, ignore_status=True)
+          if ret:
+            # "powersave" is not available, use "ondemand".
+            # Still not a fatal error if it fails.
+            ret = self.SetCpuGovernor(
+                'ondemand', run_on_dut, ignore_status=True)
+          # TODO(denik): Run comparison test for 'powersave' and 'ondemand'
+          # on scarlet and kevin64.
+          # We might have to consider reducing freq manually to the min
+          # if it helps to reduce waiting time.
+          self.cooldown_wait_time += self.WaitCooldown(run_on_dut)
+
+        # Setup CPU governor for the benchmark run.
+        # It overwrites the previous governor settings.
+        governor = self.dut_config['governor']
+        self.SetCpuGovernor(governor, run_on_dut, ignore_status=False)
+        self.DisableTurbo(run_on_dut)
+        self.SetupCpuUsage(run_on_dut)
+        # FIXME(denik): Currently we are not recovering the previous cpufreq
+        # settings since we do reboot/setup every time anyway.
+        # But it may change in the future and then we have to recover the
+        # settings.
         if benchmark.suite == 'telemetry_Crosperf':
-          self.DecreaseWaitTime(machine, label.chromeos_root)
+          self.DecreaseWaitTime(run_on_dut)
           ret_tup = self.Telemetry_Crosperf_Run(machine, label, benchmark,
                                                 test_args, profiler_args)
         else:
@@ -103,79 +166,76 @@ class SuiteRunner(object):
         break
     return ret_tup
 
-  def DisableASLR(self, machine_name, chromeos_root):
+  def DisableASLR(self, run_on_dut):
     disable_aslr = ('set -e && '
                     'stop ui; '
                     'if [[ -e /proc/sys/kernel/randomize_va_space ]]; then '
                     '  echo 0 > /proc/sys/kernel/randomize_va_space; '
                     'fi; '
                     'start ui ')
-    self.logger.LogOutput('Disable ASLR for %s' % machine_name)
-    self._ce.CrosRunCommand(
-        disable_aslr, machine=machine_name, chromeos_root=chromeos_root)
-
-  def PinGovernorExecutionFrequencies(self, machine_name, chromeos_root):
-    """Setup Intel CPU frequency.
+    if self.log_level == 'average':
+      self.logger.LogOutput('Disable ASLR.')
+    run_on_dut(disable_aslr)
 
-    Manages the cpu governor and other performance settings.
-    Includes support for setting cpu frequency to a static value.
-    """
-    # pyformat: disable
-    set_cpu_freq = (
-        # Disable Intel Opportunistic Processor
-        # Commented out because wrmsr requires kernel change
-        # to enable white-listed write access to msr 0x199.
-        # See Intel 64 and IA-32 Archtectures
-        # Software Developer's Manual, 14.3.2.2.
-        #'awk \'$1 ~ /^processor/ { print $NF }\' /proc/cpuinfo '
-        #'   | while read c; do '
-        # 'iotools wrmsr $c 0x199 $(printf "0x%x\n" $(( (1 << 32) '
-        # '   | $(iotools rdmsr $c 0x199) )));'
-        #'done;'
-        # Set up intel_pstate governor to performance if enabled.
+  def SetCpuGovernor(self, governor, run_on_dut, ignore_status=False):
+    set_gov_cmd = (
         'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
-        # Skip writing scaling_governor if cpu is not online.
+        # Skip writing scaling_governor if cpu is offline.
         ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
         '   && continue; '
-        # The cpu is online, can update.
         ' cd $f; '
         ' if [[ -e scaling_governor ]]; then '
-        '  echo performance > scaling_governor; fi; '
-        #
-        # Uncomment rest of lines to enable setting frequency by crosperf.
-        # It sets the cpu to the second highest supported frequency.
-        #'val=0; '
-        #'if [[ -e scaling_available_frequencies ]]; then '
-        # pylint: disable=line-too-long
-        #'  val=`cat scaling_available_frequencies | tr " " "\\n" | sort -n -b -r`; '
-        #'else '
-        #'  val=`cat scaling_max_freq | tr " " "\\n" | sort -n -b -r`; fi ; '
-        #'set -- $val; '
-        #'highest=$1; '
-        #'if [[ $# -gt 1 ]]; then '
-        #'  case $highest in *1000) highest=$2;; esac; '
-        #'fi ;'
-        #'echo $highest > scaling_max_freq; '
-        #'echo $highest > scaling_min_freq; '
-        'done; '
-        # Disable Turbo in Intel pstate driver
-        # no_turbo should follow governor setup.
-        # Otherwise it can be overwritten.
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    if self.log_level == 'average':
+      self.logger.LogOutput('Setup CPU Governor: %s.' % governor)
+    ret, _, _ = run_on_dut(set_gov_cmd % governor, ignore_status=ignore_status)
+    return ret
+
+  def DisableTurbo(self, run_on_dut):
+    dis_turbo_cmd = (
         'if [[ -e /sys/devices/system/cpu/intel_pstate/no_turbo ]]; then '
         '  if grep -q 0 /sys/devices/system/cpu/intel_pstate/no_turbo;  then '
         '    echo -n 1 > /sys/devices/system/cpu/intel_pstate/no_turbo; '
         '  fi; '
         'fi; ')
-    # pyformat: enable
     if self.log_level == 'average':
-      self.logger.LogOutput(
-          'Pinning governor execution frequencies for %s' % machine_name)
-    ret = self._ce.CrosRunCommand(
-        set_cpu_freq, machine=machine_name, chromeos_root=chromeos_root)
-    self.logger.LogFatalIf(
-        ret, 'Could not pin frequencies on machine: %s' % machine_name)
-
-  def SetupCpuUsage(self, machine_name, chromeos_root):
+      self.logger.LogOutput('Disable Turbo.')
+    run_on_dut(dis_turbo_cmd)
+
+  def WaitCooldown(self, run_on_dut):
+    waittime = 0
+    timeout_in_sec = int(self.dut_config['cooldown_time']) * 60
+    # Temperature from sensors come in uCelsius units.
+    temp_in_ucels = int(self.dut_config['cooldown_temp']) * 1000
+    sleep_interval = 30
+
+    # Wait until any of two events occurs:
+    # 1. CPU cools down to a specified temperature.
+    # 2. Timeout cooldown_time expires.
+    # For the case when targeted temperature is not reached within specified
+    # timeout the benchmark is going to start with higher initial CPU temp.
+    # In the worst case it may affect test results.
+    # TODO(denik): Report (or highlight) "high" CPU temperature in test results.
+    # "high" should be calculated based on empirical data per platform.
+    # Based on such reports we can adjust CPU configuration or
+    # cooldown limits accordingly.
+    while waittime < timeout_in_sec:
+      _, temp_output, _ = run_on_dut(
+          'cat /sys/class/thermal/thermal_zone*/temp')
+      if any(int(temp) > temp_in_ucels for temp in temp_output.split()):
+        time.sleep(sleep_interval)
+        waittime += sleep_interval
+      else:
+        # Reported temp numbers from all sensors do not exceed
+        # 'cooldown_temp'. Exit the loop.
+        break
+
+    self.logger.LogOutput(
+        'Cooldown wait time: %.1f min' % (float(waittime) / 60))
+    return waittime
+
+  def SetupCpuUsage(self, run_on_dut):
     """Setup CPU usage.
 
     Based on self.dut_config['cpu_usage'] configure CPU cores
@@ -184,15 +244,12 @@ class SuiteRunner(object):
 
     if (self.dut_config['cpu_usage'] == 'big_only' or
         self.dut_config['cpu_usage'] == 'little_only'):
-      ret, arch, _ = self._ce.CrosRunCommandWOutput(
-          'uname -m', machine=machine_name, chromeos_root=chromeos_root)
-      self.logger.LogFatalIf(
-          ret, 'Setup CPU failed. Could not retrieve architecture model.')
+      _, arch, _ = run_on_dut('uname -m')
 
       if arch.lower().startswith('arm') or arch.lower().startswith('aarch64'):
-        self.SetupArmCores(machine_name, chromeos_root)
+        self.SetupArmCores(run_on_dut)
 
-  def SetupArmCores(self, machine_name, chromeos_root):
+  def SetupArmCores(self, run_on_dut):
     """Setup ARM big/little cores."""
 
     # CPU implemeters/part numbers of big/LITTLE CPU.
@@ -228,10 +285,7 @@ class SuiteRunner(object):
     # CPU part        : 0xd03
     # CPU revision    : 4
 
-    ret, cpuinfo, _ = self._ce.CrosRunCommandWOutput(
-        'cat /proc/cpuinfo', machine=machine_name, chromeos_root=chromeos_root)
-    self.logger.LogFatalIf(ret,
-                           'Setup ARM CPU failed. Could not retrieve cpuinfo.')
+    _, cpuinfo, _ = run_on_dut('cat /proc/cpuinfo')
 
     # List of all CPU cores: 0, 1, ..
     proc_matches = re.findall(r'^processor\s*: (\d+)$', cpuinfo, re.MULTILINE)
@@ -278,12 +332,7 @@ class SuiteRunner(object):
             'echo 0 | tee /sys/devices/system/cpu/cpu{%s}/online' % ','.join(
                 sorted(cores_to_disable)))
 
-      ret = self._ce.CrosRunCommand(
-          '; '.join([cmd_enable_cores, cmd_disable_cores]),
-          machine=machine_name,
-          chromeos_root=chromeos_root)
-      self.logger.LogFatalIf(
-          ret, 'Setup ARM CPU failed. Could not retrieve cpuinfo.')
+      run_on_dut('; '.join([cmd_enable_cores, cmd_disable_cores]))
     else:
       # If there are no cores enabled by dut_config then configuration
       # is invalid for current platform and should be ignored.
@@ -295,20 +344,14 @@ class SuiteRunner(object):
           'Ignore ARM CPU setup and continue.' % (self.dut_config['cpu_usage'],
                                                   dut_big_cores, dut_lit_cores))
 
-  def DecreaseWaitTime(self, machine_name, chromeos_root):
+  def DecreaseWaitTime(self, run_on_dut):
     """Change the ten seconds wait time for pagecycler to two seconds."""
     FILE = '/usr/local/telemetry/src/tools/perf/page_sets/page_cycler_story.py'
-    ret = self._ce.CrosRunCommand(
-        'ls ' + FILE, machine=machine_name, chromeos_root=chromeos_root)
-    self.logger.LogFatalIf(
-        ret, 'Could not find {} on machine: {}'.format(FILE, machine_name))
+    ret = run_on_dut('ls ' + FILE)
 
     if not ret:
       sed_command = 'sed -i "s/_TTI_WAIT_TIME = 10/_TTI_WAIT_TIME = 2/g" '
-      ret = self._ce.CrosRunCommand(
-          sed_command + FILE, machine=machine_name, chromeos_root=chromeos_root)
-      self.logger.LogFatalIf(
-          ret, 'Could not modify {} on machine: {}'.format(FILE, machine_name))
+      run_on_dut(sed_command + FILE)
 
   def RestartUI(self, machine_name, chromeos_root):
     command = 'stop ui; sleep 5; start ui'
@@ -371,7 +414,6 @@ class SuiteRunner(object):
     t = 0
     RETRY_LIMIT = 5
     while t < RETRY_LIMIT:
-      self.logger.LogOutput('Downloading test results.')
       t += 1
       status = self._ce.RunCommand(ls_command, print_to_console=False)
       if status == 0:
@@ -503,7 +545,7 @@ class SuiteRunner(object):
     # process group.
     chrome_root_options = ('--no-ns-pid '
                            '--chrome_root={} --chrome_root_mount={} '
-                           "FEATURES=\"-usersandbox\" "
+                           'FEATURES="-usersandbox" '
                            'CHROME_ROOT={}'.format(label.chrome_src,
                                                    CHROME_MOUNT_DIR,
                                                    CHROME_MOUNT_DIR))
diff --git a/crosperf/suite_runner_unittest.py b/crosperf/suite_runner_unittest.py
index c5a4c181..6be3865f 100755
--- a/crosperf/suite_runner_unittest.py
+++ b/crosperf/suite_runner_unittest.py
@@ -140,9 +140,6 @@ class SuiteRunnerTest(unittest.TestCase):
   def __init__(self, *args, **kwargs):
     super(SuiteRunnerTest, self).__init__(*args, **kwargs)
     self.call_test_that_run = False
-    self.disable_aslr_args = []
-    self.pin_governor_args = []
-    self.setup_cpu_usage_args = []
     self.skylab_run_args = []
     self.test_that_args = []
     self.telemetry_run_args = []
@@ -150,8 +147,6 @@ class SuiteRunnerTest(unittest.TestCase):
     self.call_skylab_run = False
     self.call_telemetry_crosperf_run = False
     self.call_disable_aslr = False
-    self.call_pin_governor = False
-    self.call_setup_cpu_usage = False
 
   def setUp(self):
     self.runner = suite_runner.SuiteRunner(
@@ -168,29 +163,17 @@ class SuiteRunnerTest(unittest.TestCase):
   def test_run(self):
 
     def reset():
-      self.call_pin_governor = False
-      self.call_setup_cpu_usage = False
       self.call_test_that_run = False
       self.call_skylab_run = False
       self.call_telemetry_crosperf_run = False
-      self.pin_governor_args = []
-      self.setup_cpu_usage_args = []
+      self.call_disable_aslr = False
       self.skylab_run_args = []
       self.test_that_args = []
       self.telemetry_run_args = []
       self.telemetry_crosperf_args = []
 
-    def FakeDisableASLR(machine, chroot):
+    def FakeDisableASLR(runner):
       self.call_disable_aslr = True
-      self.disable_aslr_args = [machine, chroot]
-
-    def FakePinGovernor(machine, chroot):
-      self.call_pin_governor = True
-      self.pin_governor_args = [machine, chroot]
-
-    def FakeSetupCpuUsage(machine, chroot):
-      self.call_setup_cpu_usage = True
-      self.setup_cpu_usage_args = [machine, chroot]
 
     def FakeSkylabRun(test_label, benchmark, test_args, profiler_args):
       self.skylab_run_args = [test_label, benchmark, test_args, profiler_args]
@@ -213,12 +196,20 @@ class SuiteRunnerTest(unittest.TestCase):
       self.call_test_that_run = True
       return 'Ran FakeTestThatRun'
 
+    def FakeRunner(command, ignore_status=False):
+      return 0, '', ''
+
     self.runner.DisableASLR = FakeDisableASLR
-    self.runner.PinGovernorExecutionFrequencies = FakePinGovernor
-    self.runner.SetupCpuUsage = FakeSetupCpuUsage
     self.runner.Skylab_Run = FakeSkylabRun
     self.runner.Telemetry_Crosperf_Run = FakeTelemetryCrosperfRun
     self.runner.Test_That_Run = FakeTestThatRun
+    self.runner.SetupCpuUsage = mock.Mock()
+    self.runner.DutWrapper = mock.Mock(return_value=FakeRunner)
+    self.runner.DisableTurbo = mock.Mock()
+    self.runner.SetCpuGovernor = mock.Mock()
+    self.runner.WaitCooldown = mock.Mock(return_value=0)
+    self.runner.dut_config['cooldown_time'] = 0
+    self.runner.dut_config['governor'] = 'fake_governor'
     machine = 'fake_machine'
     test_args = ''
     profiler_args = ''
@@ -228,47 +219,122 @@ class SuiteRunnerTest(unittest.TestCase):
     self.runner.Run(machine, self.mock_label, self.telemetry_bench, test_args,
                     profiler_args)
     self.assertFalse(self.call_disable_aslr)
-    self.assertFalse(self.call_pin_governor)
-    self.assertFalse(self.call_setup_cpu_usage)
     self.assertTrue(self.call_skylab_run)
     self.assertFalse(self.call_test_that_run)
     self.assertFalse(self.call_telemetry_crosperf_run)
     self.assertEqual(self.skylab_run_args,
                      [self.mock_label, self.telemetry_bench, '', ''])
+    self.runner.SetupCpuUsage.assert_not_called()
+    self.runner.DutWrapper.assert_not_called()
+    self.runner.SetCpuGovernor.assert_not_called()
+    self.runner.DisableTurbo.assert_not_called()
+    self.runner.WaitCooldown.assert_not_called()
     self.mock_label.skylab = False
 
     reset()
     self.runner.Run(machine, self.mock_label, self.test_that_bench, test_args,
                     profiler_args)
     self.assertTrue(self.call_disable_aslr)
-    self.assertTrue(self.call_pin_governor)
-    self.assertTrue(self.call_setup_cpu_usage)
-    self.assertEqual(self.setup_cpu_usage_args,
-                     [machine, self.mock_label.chromeos_root])
     self.assertTrue(self.call_test_that_run)
     self.assertFalse(self.call_telemetry_crosperf_run)
     self.assertEqual(
         self.test_that_args,
         ['fake_machine', self.mock_label, self.test_that_bench, '', ''])
+    self.runner.SetupCpuUsage.assert_called_once_with(FakeRunner)
+    self.runner.DutWrapper.assert_called_once_with(
+        machine, self.mock_label.chromeos_root)
+    self.runner.SetCpuGovernor.assert_called_once_with(
+        'fake_governor', FakeRunner, ignore_status=False)
+    self.runner.DisableTurbo.assert_called_once_with(FakeRunner)
+    self.runner.WaitCooldown.assert_not_called()
 
     reset()
     self.runner.Run(machine, self.mock_label, self.telemetry_crosperf_bench,
                     test_args, profiler_args)
     self.assertTrue(self.call_disable_aslr)
-    self.assertTrue(self.call_pin_governor)
-    self.assertTrue(self.call_setup_cpu_usage)
     self.assertFalse(self.call_test_that_run)
     self.assertTrue(self.call_telemetry_crosperf_run)
     self.assertEqual(self.telemetry_crosperf_args, [
         'fake_machine', self.mock_label, self.telemetry_crosperf_bench, '', ''
     ])
+    self.runner.DutWrapper.assert_called_with(machine,
+                                              self.mock_label.chromeos_root)
+
+  def test_run_with_cooldown(self):
+
+    def FakeRunner(command, ignore_status=False):
+      return 0, '', ''
+
+    self.runner.DisableASLR = mock.Mock()
+    self.runner.DutWrapper = mock.Mock(return_value=FakeRunner)
+    self.runner.DisableTurbo = mock.Mock()
+    self.runner.SetCpuGovernor = mock.Mock()
+    self.runner.SetupCpuUsage = mock.Mock()
+    self.runner.WaitCooldown = mock.Mock(return_value=0)
+    self.runner.Telemetry_Crosperf_Run = mock.Mock(return_value=(0, '', ''))
+    self.runner.dut_config['cooldown_time'] = 10
+    self.runner.dut_config['governor'] = 'fake_governor'
+
+    self.runner.Run('fake_machine', self.mock_label,
+                    self.telemetry_crosperf_bench, '', '')
+
+    self.runner.WaitCooldown.assert_called_once_with(FakeRunner)
+    self.runner.DisableASLR.assert_called_once()
+    self.runner.Telemetry_Crosperf_Run.assert_called_once()
+    self.runner.DisableTurbo.assert_called_once_with(FakeRunner)
+    self.runner.SetupCpuUsage.assert_called_once_with(FakeRunner)
+    self.runner.SetCpuGovernor.assert_called()
+    self.assertGreater(self.runner.SetCpuGovernor.call_count, 1)
+    self.assertEqual(
+        self.runner.SetCpuGovernor.call_args,
+        mock.call('fake_governor', FakeRunner, ignore_status=False))
 
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_disable_aslr(self, mock_cros_runcmd):
-    self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.runner.DisableASLR('lumpy1.cros', '/tmp/chromeos')
-    self.assertEqual(mock_cros_runcmd.call_count, 1)
-    cmd = mock_cros_runcmd.call_args_list[0][0]
+  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
+  def test_dut_wrapper(self, mock_cros_runcmd):
+    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd
+    mock_cros_runcmd.return_value = (0, '', '')
+    run_on_dut = self.runner.DutWrapper('lumpy.cros2', '/tmp/chromeos')
+    mock_cros_runcmd.assert_not_called()
+    run_on_dut('run command;')
+    mock_cros_runcmd.assert_called_once_with(
+        'run command;', chromeos_root='/tmp/chromeos', machine='lumpy.cros2')
+
+  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
+  def test_dut_wrapper_fatal_error(self, mock_cros_runcmd):
+    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd
+    # Command returns error 1.
+    mock_cros_runcmd.return_value = (1, '', 'Error!')
+    run_on_dut = self.runner.DutWrapper('lumpy.cros2', '/tmp/chromeos')
+    mock_cros_runcmd.assert_not_called()
+    run_on_dut('run command;')
+    mock_cros_runcmd.assert_called_once_with(
+        'run command;', chromeos_root='/tmp/chromeos', machine='lumpy.cros2')
+    # Error status causes log fatal.
+    self.assertEqual(
+        self.mock_logger.method_calls[-1],
+        mock.call.LogFatal('Command execution on DUT lumpy.cros2 failed.\n'
+                           'Failing command: run command;\nreturned 1\n'
+                           'Error message: Error!'))
+
+  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
+  def test_dut_wrapper_ignore_error(self, mock_cros_runcmd):
+    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd
+    # Command returns error 1.
+    mock_cros_runcmd.return_value = (1, '', 'Error!')
+    run_on_dut = self.runner.DutWrapper('lumpy.cros2', '/tmp/chromeos')
+    run_on_dut('run command;', ignore_status=True)
+    mock_cros_runcmd.assert_called_once_with(
+        'run command;', chromeos_root='/tmp/chromeos', machine='lumpy.cros2')
+    # Error status is not fatal. LogError records the error message.
+    self.assertEqual(
+        self.mock_logger.method_calls[-1],
+        mock.call.LogError('Command execution on DUT lumpy.cros2 failed.\n'
+                           'Failing command: run command;\nreturned 1\n'
+                           'Error message: Error!'))
+
+  def test_disable_aslr(self):
+    run_on_dut = mock.Mock()
+    self.runner.DisableASLR(run_on_dut)
     # pyformat: disable
     set_cpu_cmd = ('set -e && '
                    'stop ui; '
@@ -276,147 +342,221 @@ class SuiteRunnerTest(unittest.TestCase):
                    '  echo 0 > /proc/sys/kernel/randomize_va_space; '
                    'fi; '
                    'start ui ')
-    # pyformat: enable
-    self.assertEqual(cmd, (set_cpu_cmd,))
+    run_on_dut.assert_called_once_with(set_cpu_cmd)
 
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_pin_governor_execution_frequencies(self, mock_cros_runcmd):
-    self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.runner.PinGovernorExecutionFrequencies('lumpy1.cros', '/tmp/chromeos')
-    self.assertEqual(mock_cros_runcmd.call_count, 1)
-    cmd = mock_cros_runcmd.call_args_list[0][0]
-    # pyformat: disable
+  def test_set_cpu_governor(self):
+    dut_runner = mock.Mock(return_value=(0, '', ''))
+    self.runner.SetCpuGovernor('new_governor', dut_runner, ignore_status=False)
+    set_cpu_cmd = (
+        'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
+        # Skip writing scaling_governor if cpu is not online.
+        ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
+        '   && continue; '
+        ' cd $f; '
+        ' if [[ -e scaling_governor ]]; then '
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    dut_runner.assert_called_once_with(
+        set_cpu_cmd % 'new_governor', ignore_status=False)
+
+  def test_set_cpu_governor_propagate_error(self):
+    dut_runner = mock.Mock(return_value=(1, '', 'Error.'))
+    self.runner.SetCpuGovernor('non-exist_governor', dut_runner)
+    set_cpu_cmd = (
+        'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
+        # Skip writing scaling_governor if cpu is not online.
+        ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
+        '   && continue; '
+        ' cd $f; '
+        ' if [[ -e scaling_governor ]]; then '
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    # By default error status is fatal.
+    dut_runner.assert_called_once_with(
+        set_cpu_cmd % 'non-exist_governor', ignore_status=False)
+
+  def test_set_cpu_governor_ignore_status(self):
+    dut_runner = mock.Mock(return_value=(1, '', 'Error.'))
+    ret_code = self.runner.SetCpuGovernor(
+        'non-exist_governor', dut_runner, ignore_status=True)
     set_cpu_cmd = (
         'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
         # Skip writing scaling_governor if cpu is not online.
         ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
         '   && continue; '
-        # The cpu is online, can update.
         ' cd $f; '
         ' if [[ -e scaling_governor ]]; then '
-        '  echo performance > scaling_governor; fi; '
-        'done; '
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    dut_runner.assert_called_once_with(
+        set_cpu_cmd % 'non-exist_governor', ignore_status=True)
+
+  def test_disable_turbo(self):
+    dut_runner = mock.Mock(return_value=(0, '', ''))
+    self.runner.DisableTurbo(dut_runner)
+    set_cpu_cmd = (
         # Disable Turbo in Intel pstate driver
         'if [[ -e /sys/devices/system/cpu/intel_pstate/no_turbo ]]; then '
         '  if grep -q 0 /sys/devices/system/cpu/intel_pstate/no_turbo;  then '
         '    echo -n 1 > /sys/devices/system/cpu/intel_pstate/no_turbo; '
         '  fi; '
         'fi; ')
-    # pyformat: enable
-    self.assertEqual(cmd, (set_cpu_cmd,))
+    dut_runner.assert_called_once_with(set_cpu_cmd)
 
   @mock.patch.object(suite_runner.SuiteRunner, 'SetupArmCores')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  def test_setup_cpu_usage_little_on_arm(self, mock_cros_runcmd_wout,
-                                         mock_setup_arm):
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
+  def test_setup_cpu_usage_little_on_arm(self, mock_setup_arm):
     self.runner.SetupArmCores = mock_setup_arm
-    mock_cros_runcmd_wout.return_value = (0, 'armv7l', '')
+    dut_runner = mock.Mock(return_value=(0, 'armv7l', ''))
     self.runner.dut_config['cpu_usage'] = 'little_only'
-    self.runner.SetupCpuUsage('remote.cros', '/tmp/chromeos')
-    self.assertEqual(mock_setup_arm.call_count, 1)
+    self.runner.SetupCpuUsage(dut_runner)
+    self.runner.SetupArmCores.assert_called_once_with(dut_runner)
 
   @mock.patch.object(suite_runner.SuiteRunner, 'SetupArmCores')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  def test_setup_cpu_usage_big_on_aarch64(self, mock_cros_runcmd_wout,
-                                          mock_setup_arm):
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
+  def test_setup_cpu_usage_big_on_aarch64(self, mock_setup_arm):
     self.runner.SetupArmCores = mock_setup_arm
-    mock_cros_runcmd_wout.return_value = (0, 'aarch64', '')
+    dut_runner = mock.Mock(return_value=(0, 'aarch64', ''))
     self.runner.dut_config['cpu_usage'] = 'big_only'
-    self.runner.SetupCpuUsage('remote.cros', '/tmp/chromeos')
-    self.assertEqual(mock_setup_arm.call_count, 1)
+    self.runner.SetupCpuUsage(dut_runner)
+    self.runner.SetupArmCores.assert_called_once_with(dut_runner)
 
   @mock.patch.object(suite_runner.SuiteRunner, 'SetupArmCores')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  def test_setup_cpu_usage_all_on_intel(self, mock_cros_runcmd_wout,
-                                        mock_setup_arm):
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
+  def test_setup_cpu_usage_big_on_intel(self, mock_setup_arm):
     self.runner.SetupArmCores = mock_setup_arm
-    mock_cros_runcmd_wout.return_value = (0, 'x86_64', '')
-    self.runner.dut_config['cpu_usage'] = 'all'
-    self.runner.SetupCpuUsage('remote.cros', '/tmp/chromeos')
-    # Check that SetupArmCores not called.
-    self.assertEqual(mock_setup_arm.call_count, 0)
-
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_setup_arm_cores_big_on_big_little(self, mock_cros_runcmd,
-                                             mock_cros_runcmd_wout):
-    self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
-    mock_cros_runcmd_wout.return_value = (0, BIG_LITTLE_CPUINFO, '')
+    dut_runner = mock.Mock(return_value=(0, 'x86_64', ''))
     self.runner.dut_config['cpu_usage'] = 'big_only'
-    self.runner.SetupArmCores('remote.cros', '/tmp/chromeos')
-    self.assertEqual(mock_cros_runcmd.call_args_list[0][0],
-                     ('echo 1 | tee /sys/devices/system/cpu/cpu{2}/online; '
-                      'echo 0 | tee /sys/devices/system/cpu/cpu{0,1}/online',))
+    self.runner.SetupCpuUsage(dut_runner)
+    # Check that SetupArmCores not called with invalid setup.
+    self.runner.SetupArmCores.assert_not_called()
 
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_setup_arm_cores_little_on_big_little(self, mock_cros_runcmd,
-                                                mock_cros_runcmd_wout):
-    self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
-    mock_cros_runcmd_wout.return_value = (0, BIG_LITTLE_CPUINFO, '')
+  @mock.patch.object(suite_runner.SuiteRunner, 'SetupArmCores')
+  def test_setup_cpu_usage_all_on_intel(self, mock_setup_arm):
+    self.runner.SetupArmCores = mock_setup_arm
+    dut_runner = mock.Mock(return_value=(0, 'x86_64', ''))
+    self.runner.dut_config['cpu_usage'] = 'all'
+    self.runner.SetupCpuUsage(dut_runner)
+    # Check that SetupArmCores not called in general case.
+    self.runner.SetupArmCores.assert_not_called()
+
+  def test_setup_arm_cores_big_on_big_little(self):
+    dut_runner = mock.Mock(
+        side_effect=[(0, BIG_LITTLE_CPUINFO, ''),
+                     (0, '', '')])
+    self.runner.dut_config['cpu_usage'] = 'big_only'
+    self.runner.SetupArmCores(dut_runner)
+    dut_runner.assert_called_with(
+        'echo 1 | tee /sys/devices/system/cpu/cpu{2}/online; '
+        'echo 0 | tee /sys/devices/system/cpu/cpu{0,1}/online')
+
+  def test_setup_arm_cores_little_on_big_little(self):
+    dut_runner = mock.Mock(
+        side_effect=[(0, BIG_LITTLE_CPUINFO, ''),
+                     (0, '', '')])
     self.runner.dut_config['cpu_usage'] = 'little_only'
-    self.runner.SetupArmCores('remote.cros', '/tmp/chromeos')
-    self.assertEqual(mock_cros_runcmd.call_args_list[0][0],
-                     ('echo 1 | tee /sys/devices/system/cpu/cpu{0,1}/online; '
-                      'echo 0 | tee /sys/devices/system/cpu/cpu{2}/online',))
-
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_setup_arm_cores_invalid_config(self, mock_cros_runcmd,
-                                          mock_cros_runcmd_wout):
-    self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
-    mock_cros_runcmd_wout.return_value = (0, LITTLE_ONLY_CPUINFO, '')
+    self.runner.SetupArmCores(dut_runner)
+    dut_runner.assert_called_with(
+        'echo 1 | tee /sys/devices/system/cpu/cpu{0,1}/online; '
+        'echo 0 | tee /sys/devices/system/cpu/cpu{2}/online')
+
+  def test_setup_arm_cores_invalid_config(self):
+    dut_runner = mock.Mock(
+        side_effect=[(0, LITTLE_ONLY_CPUINFO, ''),
+                     (0, '', '')])
     self.runner.dut_config['cpu_usage'] = 'big_only'
-    self.runner.SetupArmCores('remote.cros', '/tmp/chromeos')
-    # Check that CrosRun is not called when trying
+    self.runner.SetupArmCores(dut_runner)
+    # Check that setup command is not sent when trying
     # to use 'big_only' on a platform with all little cores.
-    self.assertEqual(mock_cros_runcmd.call_count, 0)
+    dut_runner.assert_called_once_with('cat /proc/cpuinfo')
 
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_setup_arm_cores_not_big_little(self, mock_cros_runcmd,
-                                          mock_cros_runcmd_wout):
-    self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
-    mock_cros_runcmd_wout.return_value = (0, NOT_BIG_LITTLE_CPUINFO, '')
+  def test_setup_arm_cores_not_big_little(self):
+    dut_runner = mock.Mock(
+        side_effect=[(0, NOT_BIG_LITTLE_CPUINFO, ''),
+                     (0, '', '')])
     self.runner.dut_config['cpu_usage'] = 'big_only'
-    self.runner.SetupArmCores('remote.cros', '/tmp/chromeos')
-    # Check that CrosRun is not called when trying
+    self.runner.SetupArmCores(dut_runner)
+    # Check that setup command is not sent when trying
     # to use 'big_only' on a platform w/o support of big/little.
-    self.assertEqual(mock_cros_runcmd.call_count, 0)
+    dut_runner.assert_called_once_with('cat /proc/cpuinfo')
 
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_setup_arm_cores_unsupported_cpu_usage(self, mock_cros_runcmd,
-                                                 mock_cros_runcmd_wout):
-    self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd_wout
-    mock_cros_runcmd_wout.return_value = (0, BIG_LITTLE_CPUINFO, '')
+  def test_setup_arm_cores_unsupported_cpu_usage(self):
+    dut_runner = mock.Mock(
+        side_effect=[(0, BIG_LITTLE_CPUINFO, ''),
+                     (0, '', '')])
     self.runner.dut_config['cpu_usage'] = 'exclusive_cores'
-    self.runner.SetupArmCores('remote.cros', '/tmp/chromeos')
-    # Check that CrosRun is not called when trying to use
+    self.runner.SetupArmCores(dut_runner)
+    # Check that setup command is not sent when trying to use
     # 'exclusive_cores' on ARM CPU setup.
-    self.assertEqual(mock_cros_runcmd.call_count, 0)
+    dut_runner.assert_called_once_with('cat /proc/cpuinfo')
 
-  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
-  def test_reboot_machine(self, mock_cros_runcmd):
+  def test_wait_cooldown_nowait(self):
+    dut_runner = mock.Mock(return_value=(0, '39000', ''))
+    self.runner.dut_config['cooldown_time'] = 10
+    self.runner.dut_config['cooldown_temp'] = 40
+    self.runner.WaitCooldown(dut_runner)
+    # Send command to DUT only once to check temperature
+    # and make sure it does not exceed the threshold.
+    dut_runner.assert_called_once()
 
-    def FakePinGovernor(machine_name, chromeos_root):
-      if machine_name or chromeos_root:
-        pass
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_needwait_once(self, mock_sleep):
+    """Wait one iteration for cooldown.
+
+    Set large enough timeout and changing temperature
+    output. Make sure it exits when expected value
+    received.
+    Expect that WaitCooldown check temp twice.
+    """
+    mock_sleep.return_value = 0
+    dut_runner = mock.Mock(side_effect=[(0, '41000', ''), (0, '39999', '')])
+    self.runner.dut_config['cooldown_time'] = 100
+    self.runner.dut_config['cooldown_temp'] = 40
+    self.runner.WaitCooldown(dut_runner)
+    dut_runner.assert_called()
+    self.assertEqual(dut_runner.call_count, 2)
 
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_needwait(self, mock_sleep):
+    """Test exit by timeout.
+
+    Send command to DUT checking the temperature and
+    check repeatedly until timeout goes off.
+    Output from temperature sensor never changes.
+    """
+    mock_sleep.return_value = 0
+    dut_runner = mock.Mock(return_value=(0, '41000', ''))
+    self.runner.dut_config['cooldown_time'] = 60
+    self.runner.dut_config['cooldown_temp'] = 40
+    self.runner.WaitCooldown(dut_runner)
+    dut_runner.assert_called()
+    self.assertGreater(dut_runner.call_count, 2)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_needwait_multtemp(self, mock_sleep):
+    """Wait until all temps go down.
+
+    Set large enough timeout and changing temperature
+    output. Make sure it exits when expected value
+    for all temperatures received.
+    Expect 3 checks.
+    """
+    mock_sleep.return_value = 0
+    dut_runner = mock.Mock(side_effect=[
+        (0, '41000\n20000\n30000\n45000', ''),
+        (0, '39000\n20000\n30000\n41000', ''),
+        (0, '39000\n20000\n30000\n31000', '')])
+    self.runner.dut_config['cooldown_time'] = 100
+    self.runner.dut_config['cooldown_temp'] = 40
+    self.runner.WaitCooldown(dut_runner)
+    dut_runner.assert_called()
+    self.assertEqual(dut_runner.call_count, 3)
+
+  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
+  def test_restart_ui(self, mock_cros_runcmd):
     self.mock_cmd_exec.CrosRunCommand = mock_cros_runcmd
-    self.runner.PinGovernorExecutionFrequencies = FakePinGovernor
     self.runner.RestartUI('lumpy1.cros', '/tmp/chromeos')
-    self.assertEqual(mock_cros_runcmd.call_count, 1)
-    self.assertEqual(mock_cros_runcmd.call_args_list[0][0],
-                     ('stop ui; sleep 5; start ui',))
+    mock_cros_runcmd.assert_called_once_with(
+        'stop ui; sleep 5; start ui',
+        chromeos_root='/tmp/chromeos',
+        machine='lumpy1.cros')
 
   @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommand')
   @mock.patch.object(command_executer.CommandExecuter,