AFDO-Bisect: Implement state-saving

This CL adds state-saving functionality to this AFDO profile analysis
tool. It introduces two new flags: state_file, and no_resume. The former
allows the user to specify a specific state file for initial
state-loading, and state-writing on each iteration. If this flag is not
provided, a default file in the current working directory is used. The
no_resume flag allows the user to disable initial state loading, beginning a
fresh run of the script.

TEST=All added tests pass
BUG=None

Change-Id: Ib288180dc6ca15d47e2720f222905e2aa6b964a0
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1708215
Reviewed-by: Caroline Tice <cmtice@chromium.org>
Tested-by: Emma Vukelj <emmavukelj@google.com>

3 files changed, 227 insertions, 108 deletions

diff --git a/afdo_tools/bisection/afdo_prof_analysis.py b/afdo_tools/bisection/afdo_prof_analysis.py
index 35e109a0..095adf0c 100755
--- a/afdo_tools/bisection/afdo_prof_analysis.py
+++ b/afdo_tools/bisection/afdo_prof_analysis.py
@@ -17,6 +17,7 @@ from __future__ import print_function
 
 from absl import app
 from absl import flags
+from datetime import date
 from enum import IntEnum
 from tempfile import mkstemp
 
@@ -41,11 +42,21 @@ flags.DEFINE_string(
     'AFDO profile, returns GOOD/BAD/SKIP')
 flags.DEFINE_string('analysis_output_file', None,
                     'File to output JSON results to')
-flags.DEFINE_integer('seed', None, 'Integer specifying seed for randomness')
+flags.DEFINE_string(
+    'state_file', '%s/afdo_analysis_state.json' % os.getcwd(),
+    'File path containing state to load from initially, and '
+    'will be overwritten with new state on each iteration')
+flags.DEFINE_boolean(
+    'no_resume', False, 'If enabled, no initial state will be '
+    'loaded and the program will run from the beginning')
+flags.DEFINE_boolean(
+    'remove_state_on_completion', False, 'If enabled, state '
+    'file will be removed once profile analysis is completed')
+flags.DEFINE_float('seed', None, 'Float specifying seed for randomness')
 FLAGS = flags.FLAGS
 
 
-class status_enum(IntEnum):
+class StatusEnum(IntEnum):
   """Enum of valid statuses returned by profile decider."""
   GOOD_STATUS = 0
   BAD_STATUS = 1
@@ -53,7 +64,7 @@ class status_enum(IntEnum):
   PROBLEM_STATUS = 127
 
 
-statuses = status_enum.__members__.values()
+statuses = StatusEnum.__members__.values()
 
 _NUM_RUNS_RANGE_SEARCH = 20  # how many times range search should run its algo
 
@@ -76,22 +87,52 @@ def prof_to_tmp(prof):
   return temp_path
 
 
-def generate_decider():
-  """create the decider function with counter for total number of calls
-
-  generate_decider is a function which returns an inner function (and the
-  inner one does the work of running the external decider). This has been
-  structured as a function which returns a function so we can keep track of the
-  global/total number of runs as a member of the outer function, as opposed to
-  using a global variable.
-  """
-
-  # Inner functions can only modify, and not rebind, nonlocal variables
-  # which is why this uses a list to represent a single integer
-  _num_runs = [0]
-
-  def run(prof, increment_counter=True):
+class DeciderState(object):
+  """Class for the external decider."""
+
+  def __init__(self, state_file):
+    self.accumulated_results = []  # over this run of the script
+    self.saved_results = []  # imported from a previous run of this script
+    self.state_file = state_file
+    self.seed = FLAGS.seed or time.time()
+
+  def load_state(self):
+    if not os.path.exists(self.state_file):
+      logging.info('State file %s is empty, starting from beginning',
+                   self.state_file)
+      return
+
+    with open(self.state_file) as f:
+      try:
+        data = json.load(f)
+      except:
+        raise ValueError('Provided state file %s to resume from does not'
+                         ' contain a valid JSON.' % self.state_file)
+
+    if 'seed' not in data or 'accumulated_results' not in data:
+      raise ValueError('Provided state file %s to resume from does not contain'
+                       ' the correct information' % self.state_file)
+
+    self.seed = data['seed']
+    self.saved_results = data['accumulated_results']
+    logging.info('Restored state from %s...', self.state_file)
+
+  def save_state(self):
+    state = {'seed': self.seed, 'accumulated_results': self.accumulated_results}
+    fd, tmp_file = mkstemp()
+    with open(tmp_file, 'w') as f:
+      json.dump(state, f, indent=2)
+    os.close(fd)
+    os.rename(tmp_file, self.state_file)
+    logging.info('Logged state to %s...', self.state_file)
+
+  def run(self, prof, save_run=True):
     """Run the external deciding script on the given profile."""
+    if self.saved_results and save_run:
+      result = StatusEnum(self.saved_results.pop(0))
+      self.accumulated_results.append(result)
+      return result
+
     filename = prof_to_tmp(prof)
 
     try:
@@ -100,18 +141,17 @@ def generate_decider():
       os.remove(filename)
 
     if return_code in statuses:
-      if increment_counter:
-        _num_runs[0] += 1
-
-      status = status_enum(return_code)
-      logging.info('Run %d of external script %s returned %s', _num_runs[0],
-                   FLAGS.external_decider, status.name)
+      status = StatusEnum(return_code)
+      if save_run:
+        self.accumulated_results.append(status.value)
+      logging.info('Run %d of external script %s returned %s',
+                   len(self.accumulated_results), FLAGS.external_decider,
+                   status.name)
+      self.save_state()
       return status
     raise ValueError(
         'Provided external script had unexpected return code %d' % return_code)
 
-  return run
-
 
 def bisect_profiles(decider, good, bad, common_funcs, lo, hi):
   """Recursive function which bisects good and bad profiles.
@@ -148,21 +188,21 @@ def bisect_profiles(decider, good, bad, common_funcs, lo, hi):
   for func in common_funcs[mid:hi]:
     mid_hi_prof[func] = bad[func]
 
-  lo_mid_verdict = decider(lo_mid_prof)
-  mid_hi_verdict = decider(mid_hi_prof)
+  lo_mid_verdict = decider.run(lo_mid_prof)
+  mid_hi_verdict = decider.run(mid_hi_prof)
 
-  if lo_mid_verdict == status_enum.BAD_STATUS:
+  if lo_mid_verdict == StatusEnum.BAD_STATUS:
     result = bisect_profiles(decider, good, bad, common_funcs, lo, mid)
     results['individuals'].extend(result['individuals'])
     results['ranges'].extend(result['ranges'])
-  if mid_hi_verdict == status_enum.BAD_STATUS:
+  if mid_hi_verdict == StatusEnum.BAD_STATUS:
     result = bisect_profiles(decider, good, bad, common_funcs, mid, hi)
     results['individuals'].extend(result['individuals'])
     results['ranges'].extend(result['ranges'])
 
   # neither half is bad -> the issue is caused by several things occuring
   # in conjunction, and this combination crosses 'mid'
-  if lo_mid_verdict == mid_hi_verdict == status_enum.GOOD_STATUS:
+  if lo_mid_verdict == mid_hi_verdict == StatusEnum.GOOD_STATUS:
     problem_range = range_search(decider, good, bad, common_funcs, lo, hi)
     if problem_range:
       logging.info('Found %s as a problematic combination of profiles',
@@ -178,9 +218,9 @@ def bisect_profiles_wrapper(decider, good, bad, perform_check=True):
   # Validate good and bad profiles are such, otherwise bisection reports noise
   # Note that while decider is a random mock, these assertions may fail.
   if perform_check:
-    if decider(good, increment_counter=False) != status_enum.GOOD_STATUS:
+    if decider.run(good, save_run=False) != StatusEnum.GOOD_STATUS:
       raise ValueError("Supplied good profile is not actually GOOD")
-    if decider(bad, increment_counter=False) != status_enum.BAD_STATUS:
+    if decider.run(bad, save_run=False) != StatusEnum.BAD_STATUS:
       raise ValueError("Supplied bad profile is not actually BAD")
 
   common_funcs = sorted(func for func in good if func in bad)
@@ -223,13 +263,13 @@ def range_search(decider, good, bad, common_funcs, lo, hi):
 
     for func in funcs[lo:mid]:
       good_copy[func] = bad[func]
-    verdict = decider(good_copy)
+    verdict = decider.run(good_copy)
 
     # reset for next iteration
     for func in funcs:
       good_copy[func] = good[func]
 
-    if verdict == status_enum.BAD_STATUS:
+    if verdict == StatusEnum.BAD_STATUS:
       return find_upper_border(good_copy, funcs, lo, mid, mid)
     return find_upper_border(good_copy, funcs, mid, hi, last_bad_val)
 
@@ -241,13 +281,13 @@ def range_search(decider, good, bad, common_funcs, lo, hi):
 
     for func in funcs[lo:mid]:
       good_copy[func] = good[func]
-    verdict = decider(good_copy)
+    verdict = decider.run(good_copy)
 
     # reset for next iteration
     for func in funcs:
       good_copy[func] = bad[func]
 
-    if verdict == status_enum.BAD_STATUS:
+    if verdict == StatusEnum.BAD_STATUS:
       return find_lower_border(good_copy, funcs, mid, hi, lo)
     return find_lower_border(good_copy, funcs, lo, mid, last_bad_val)
 
@@ -299,7 +339,7 @@ def check_good_not_bad(decider, good, bad):
   for func in good:
     if func not in bad:
       bad_copy[func] = good[func]
-  return decider(bad_copy) == status_enum.GOOD_STATUS
+  return decider.run(bad_copy) == StatusEnum.GOOD_STATUS
 
 
 def check_bad_not_good(decider, good, bad):
@@ -308,33 +348,46 @@ def check_bad_not_good(decider, good, bad):
   for func in bad:
     if func not in good:
       good_copy[func] = bad[func]
-  return decider(good_copy) == status_enum.BAD_STATUS
+  return decider.run(good_copy) == StatusEnum.BAD_STATUS
 
 
 def main(_):
+
+  if not FLAGS.no_resume and FLAGS.seed:  # conflicting seeds
+    raise RuntimeError('Ambiguous seed value; do not resume from existing '
+                       'state and also specify seed by command line flag')
+
+  decider = DeciderState(FLAGS.state_file)
+  if not FLAGS.no_resume:
+    decider.load_state()
+  random.seed(decider.seed)
+
   with open(FLAGS.good_prof) as good_f:
     good_items = parse_afdo(good_f)
   with open(FLAGS.bad_prof) as bad_f:
     bad_items = parse_afdo(bad_f)
 
-  seed = FLAGS.seed
-  if not FLAGS.seed:
-    seed = time.time()
-  random.seed(seed)
-
-  decider = generate_decider()
   bisect_results = bisect_profiles_wrapper(decider, good_items, bad_items)
   gnb_result = check_good_not_bad(decider, good_items, bad_items)
   bng_result = check_bad_not_good(decider, good_items, bad_items)
 
   results = {
-      'seed': seed,
+      'seed': decider.seed,
       'bisect_results': bisect_results,
       'good_only_functions': gnb_result,
       'bad_only_functions': bng_result
   }
   with open(FLAGS.analysis_output_file, 'wb') as f:
     json.dump(results, f, indent=2)
+  if FLAGS.remove_state_on_completion:
+    os.remove(FLAGS.state_file)
+    logging.info('Removed state file %s following completion of script...',
+                 FLAGS.state_file)
+  else:
+    completed_state_file = '%s.completed.%s' % (FLAGS.state_file,
+                                                str(date.today()))
+    os.rename(FLAGS.state_file, completed_state_file)
+    logging.info('Stored completed state file as %s...', completed_state_file)
   return results
 
 
diff --git a/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py b/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
index e757cb11..3f7c844a 100755
--- a/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
+++ b/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
@@ -9,9 +9,11 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+from datetime import date
 
 import afdo_prof_analysis as analysis
 
+import json
 import shutil
 import tempfile
 import os
@@ -21,57 +23,104 @@ import unittest
 class AfdoProfAnalysisE2ETest(unittest.TestCase):
   """Class for end-to-end testing of AFDO Profile Analysis"""
 
-  def test_afdo_prof_analysis(self):
-    # nothing significant about the values, just easier to remember even vs odd
-    good_prof = {
-        'func_a': ':1\n 1: 3\n 3: 5\n 5: 7\n',
-        'func_b': ':3\n 3: 5\n 5: 7\n 7: 9\n',
-        'func_c': ':5\n 5: 7\n 7: 9\n 9: 11\n',
-        'func_d': ':7\n 7: 9\n 9: 11\n 11: 13\n',
-        'good_func_a': ':11\n',
-        'good_func_b': ':13\n'
-    }
-
-    bad_prof = {
-        'func_a': ':2\n 2: 4\n 4: 6\n 6: 8\n',
-        'func_b': ':4\n 4: 6\n 6: 8\n 8: 10\n',
-        'func_c': ':6\n 6: 8\n 8: 10\n 10: 12\n',
-        'func_d': ':8\n 8: 10\n 10: 12\n 12: 14\n',
-        'bad_func_a': ':12\n',
-        'bad_func_b': ':14\n'
-    }
+  # nothing significant about the values, just easier to remember even vs odd
+  good_prof = {
+      'func_a': ':1\n 1: 3\n 3: 5\n 5: 7\n',
+      'func_b': ':3\n 3: 5\n 5: 7\n 7: 9\n',
+      'func_c': ':5\n 5: 7\n 7: 9\n 9: 11\n',
+      'func_d': ':7\n 7: 9\n 9: 11\n 11: 13\n',
+      'good_func_a': ':11\n',
+      'good_func_b': ':13\n'
+  }
+
+  bad_prof = {
+      'func_a': ':2\n 2: 4\n 4: 6\n 6: 8\n',
+      'func_b': ':4\n 4: 6\n 6: 8\n 8: 10\n',
+      'func_c': ':6\n 6: 8\n 8: 10\n 10: 12\n',
+      'func_d': ':8\n 8: 10\n 10: 12\n 12: 14\n',
+      'bad_func_a': ':12\n',
+      'bad_func_b': ':14\n'
+  }
+
+  expected = {
+      'good_only_functions': False,
+      'bad_only_functions': True,
+      'bisect_results': {
+          'ranges': [],
+          'individuals': ['func_a']
+      }
+  }
 
+  def test_afdo_prof_analysis(self):
     # Individual issues take precedence by nature of our algos
     # so first, that should be caught
-    expected = {
-        'good_only_functions': False,
-        'bad_only_functions': True,
-        'bisect_results': {
-            'ranges': [],
-            'individuals': ['func_a']
-        }
-    }
-    self.run_check(good_prof, bad_prof, expected)
+    good = self.good_prof.copy()
+    bad = self.bad_prof.copy()
+    self.run_check(good, bad, self.expected)
 
     # Now remove individuals and exclusively BAD, and check that range is caught
-    bad_prof['func_a'] = good_prof['func_a']
-    bad_prof.pop('bad_func_a')
-    bad_prof.pop('bad_func_b')
-    expected['bad_only_functions'] = False
-    expected['bisect_results'] = {
+    bad['func_a'] = good['func_a']
+    bad.pop('bad_func_a')
+    bad.pop('bad_func_b')
+
+    expected_cp = self.expected.copy()
+    expected_cp['bad_only_functions'] = False
+    expected_cp['bisect_results'] = {
         'individuals': [],
         'ranges': [['func_b', 'func_c', 'func_d']]
     }
 
-    self.run_check(good_prof, bad_prof, expected)
+    self.run_check(good, bad, expected_cp)
 
-  def run_check(self, good_prof, bad_prof, expected):
+  def test_afdo_prof_state(self):
+    """Verifies that saved state is correct replication."""
     temp_dir = tempfile.mkdtemp()
-
-    def cleanup():
-      shutil.rmtree(temp_dir, ignore_errors=True)
-
-    self.addCleanup(cleanup)
+    self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
+
+    good = self.good_prof.copy()
+    bad = self.bad_prof.copy()
+    # add more functions to data
+    for x in range(400):
+      good['func_%d' % x] = ''
+      bad['func_%d' % x] = ''
+
+    fd_first, first_result = tempfile.mkstemp(dir=temp_dir)
+    os.close(fd_first)
+    fd_state, state_file = tempfile.mkstemp(dir=temp_dir)
+    os.close(fd_state)
+    self.run_check(
+        self.good_prof,
+        self.bad_prof,
+        self.expected,
+        state_file=state_file,
+        out_file=first_result)
+
+    fd_second, second_result = tempfile.mkstemp(dir=temp_dir)
+    os.close(fd_second)
+    completed_state_file = '%s.completed.%s' % (state_file, str(date.today()))
+    self.run_check(
+        self.good_prof,
+        self.bad_prof,
+        self.expected,
+        state_file=completed_state_file,
+        no_resume=False,
+        out_file=second_result)
+
+    with open(first_result) as f:
+      initial_run = json.load(f)
+    with open(second_result) as f:
+      loaded_run = json.load(f)
+    self.assertEqual(initial_run, loaded_run)
+
+  def run_check(self,
+                good_prof,
+                bad_prof,
+                expected,
+                state_file=None,
+                no_resume=True,
+                out_file=None):
+    temp_dir = tempfile.mkdtemp()
+    self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
 
     good_prof_file = '%s/%s' % (temp_dir, 'good_prof.txt')
     bad_prof_file = '%s/%s' % (temp_dir, 'bad_prof.txt')
@@ -84,11 +133,14 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
 
     analysis.FLAGS.good_prof = good_prof_file
     analysis.FLAGS.bad_prof = bad_prof_file
+    if state_file:
+      analysis.FLAGS.state_file = state_file
+    analysis.FLAGS.no_resume = no_resume
+    analysis.FLAGS.analysis_output_file = out_file or '/dev/null'
 
     dir_path = os.path.dirname(os.path.realpath(__file__))  # dir of this file
     external_script = '%s/e2e_external.sh' % (dir_path)
     analysis.FLAGS.external_decider = external_script
-    analysis.FLAGS.analysis_output_file = '/dev/null'
 
     actual = analysis.main(None)
     actual.pop('seed')  # nothing to check
diff --git a/afdo_tools/bisection/afdo_prof_analysis_test.py b/afdo_tools/bisection/afdo_prof_analysis_test.py
index ba1b3ab0..ae87a1f9 100755
--- a/afdo_tools/bisection/afdo_prof_analysis_test.py
+++ b/afdo_tools/bisection/afdo_prof_analysis_test.py
@@ -39,15 +39,18 @@ class AfdoProfAnalysisTest(unittest.TestCase):
   def test_bisect_profiles(self):
 
     # mock run of external script with arbitrarily-chosen bad profile vals
-    # increment_counter specified and unused b/c afdo_prof_analysis.py
+    # save_run specified and unused b/c afdo_prof_analysis.py
     # will call with argument explicitly specified
     # pylint: disable=unused-argument
-    def decider(prof, increment_counter=True):
-      if '1' in prof['func_a'] or '3' in prof['func_b']:
-        return analysis.status_enum.BAD_STATUS
-      return analysis.status_enum.GOOD_STATUS
+    class DeciderClass(object):
+      """Class for this tests's decider."""
 
-    results = analysis.bisect_profiles_wrapper(decider, self.good_items,
+      def run(self, prof, save_run=False):
+        if '1' in prof['func_a'] or '3' in prof['func_b']:
+          return analysis.StatusEnum.BAD_STATUS
+        return analysis.StatusEnum.GOOD_STATUS
+
+    results = analysis.bisect_profiles_wrapper(DeciderClass(), self.good_items,
                                                self.bad_items)
     self.assertEqual(results['individuals'], sorted(['func_a', 'func_b']))
     self.assertEqual(results['ranges'], [])
@@ -57,10 +60,13 @@ class AfdoProfAnalysisTest(unittest.TestCase):
     # arbitrarily chosen functions whose values in the bad profile constitute
     # a problematic pair
     # pylint: disable=unused-argument
-    def decider(prof, increment_counter=True):
-      if '1' in prof['func_a'] and '3' in prof['func_b']:
-        return analysis.status_enum.BAD_STATUS
-      return analysis.status_enum.GOOD_STATUS
+    class DeciderClass(object):
+      """Class for this tests's decider."""
+
+      def run(self, prof, save_run=False):
+        if '1' in prof['func_a'] and '3' in prof['func_b']:
+          return analysis.StatusEnum.BAD_STATUS
+        return analysis.StatusEnum.GOOD_STATUS
 
     # put the problematic combination in separate halves of the common funcs
     # so that non-bisecting search is invoked for its actual use case
@@ -70,7 +76,7 @@ class AfdoProfAnalysisTest(unittest.TestCase):
     common_funcs.remove('func_b')
     common_funcs.append('func_b')
 
-    problem_range = analysis.range_search(decider, self.good_items,
+    problem_range = analysis.range_search(DeciderClass(), self.good_items,
                                           self.bad_items, common_funcs, 0,
                                           len(common_funcs))
 
@@ -80,25 +86,33 @@ class AfdoProfAnalysisTest(unittest.TestCase):
     func_in_good = 'func_c'
 
     # pylint: disable=unused-argument
-    def decider(prof, increment_counter=True):
-      if func_in_good in prof:
-        return analysis.status_enum.GOOD_STATUS
-      return analysis.status_enum.BAD_STATUS
+    class DeciderClass(object):
+      """Class for this tests's decider."""
+
+      def run(self, prof, save_run=False):
+        if func_in_good in prof:
+          return analysis.StatusEnum.GOOD_STATUS
+        return analysis.StatusEnum.BAD_STATUS
 
     self.assertTrue(
-        analysis.check_good_not_bad(decider, self.good_items, self.bad_items))
+        analysis.check_good_not_bad(DeciderClass(), self.good_items,
+                                    self.bad_items))
 
   def test_check_bad_not_good(self):
     func_in_bad = 'func_d'
 
     # pylint: disable=unused-argument
-    def decider(prof, increment_counter=True):
-      if func_in_bad in prof:
-        return analysis.status_enum.BAD_STATUS
-      return analysis.status_enum.GOOD_STATUS
+    class DeciderClass(object):
+      """Class for this tests's decider."""
+
+      def run(self, prof, save_run=False):
+        if func_in_bad in prof:
+          return analysis.StatusEnum.BAD_STATUS
+        return analysis.StatusEnum.GOOD_STATUS
 
     self.assertTrue(
-        analysis.check_bad_not_good(decider, self.good_items, self.bad_items))
+        analysis.check_bad_not_good(DeciderClass(), self.good_items,
+                                    self.bad_items))
 
 
 if __name__ == '__main__':