binary_search_state.py revision 649273e844d7e86a781cb9f7a870b504bc0febbe 
1#!/usr/bin/python2
2"""The binary search wrapper."""
3
4from __future__ import print_function
5
6import argparse
7import math
8import os
9import pickle
10import sys
11import tempfile
12
13# Programtically adding utils python path to PYTHONPATH
14if os.path.isabs(sys.argv[0]):
15  utils_pythonpath = os.path.abspath('{0}/..'.format(os.path.dirname(sys.argv[
16      0])))
17else:
18  wdir = os.getcwd()
19  utils_pythonpath = os.path.abspath('{0}/{1}/..'.format(wdir, os.path.dirname(
20      sys.argv[0])))
21sys.path.append(utils_pythonpath)
22# Now we do import from utils
23from utils import command_executer
24from utils import logger
25
26import binary_search_perforce
27
28STATE_FILE = '%s.state' % sys.argv[0]
29
30
31class BinarySearchState(object):
32  """The binary search state class."""
33
34  def __init__(self, get_initial_items, switch_to_good, switch_to_bad,
35               install_script, test_script, incremental, prune, iterations,
36               prune_iterations, verify_level, file_args, verbose):
37    self.get_initial_items = get_initial_items
38    self.switch_to_good = switch_to_good
39    self.switch_to_bad = switch_to_bad
40    self.install_script = install_script
41    self.test_script = test_script
42    self.incremental = incremental
43    self.prune = prune
44    self.iterations = iterations
45    self.prune_iterations = prune_iterations
46    self.verify_level = verify_level
47    self.file_args = file_args
48    self.verbose = verbose
49
50    self.l = logger.GetLogger()
51    self.ce = command_executer.GetCommandExecuter()
52
53    self.prune_cycles = 0
54    self.search_cycles = 0
55    self.bs = None
56    self.all_items = None
57    self.PopulateItemsUsingCommand(self.get_initial_items)
58    self.currently_good_items = set([])
59    self.currently_bad_items = set([])
60    self.found_items = set([])
61
62  def SwitchToGood(self, item_list):
63    if self.incremental:
64      self.l.LogOutput('Incremental set. Wanted to switch %s to good' %
65                       str(item_list), print_to_console=self.verbose)
66      incremental_items = [
67          item for item in item_list if item not in self.currently_good_items
68      ]
69      item_list = incremental_items
70      self.l.LogOutput('Incremental set. Actually switching %s to good' %
71                       str(item_list), print_to_console=self.verbose)
72
73    if not item_list:
74      return
75
76    self.l.LogOutput('Switching %s to good' % str(item_list),
77                     print_to_console=self.verbose)
78    self.RunSwitchScript(self.switch_to_good, item_list)
79    self.currently_good_items = self.currently_good_items.union(set(item_list))
80    self.currently_bad_items.difference_update(set(item_list))
81
82  def SwitchToBad(self, item_list):
83    if self.incremental:
84      self.l.LogOutput('Incremental set. Wanted to switch %s to bad' %
85                       str(item_list), print_to_console=self.verbose)
86      incremental_items = [
87          item for item in item_list if item not in self.currently_bad_items
88      ]
89      item_list = incremental_items
90      self.l.LogOutput('Incremental set. Actually switching %s to bad' %
91                       str(item_list), print_to_console=self.verbose)
92
93    if not item_list:
94      return
95
96    self.l.LogOutput('Switching %s to bad' % str(item_list),
97                     print_to_console=self.verbose)
98    self.RunSwitchScript(self.switch_to_bad, item_list)
99    self.currently_bad_items = self.currently_bad_items.union(set(item_list))
100    self.currently_good_items.difference_update(set(item_list))
101
102  def RunSwitchScript(self, switch_script, item_list):
103    if self.file_args:
104      temp_file = tempfile.mktemp()
105      f = open(temp_file, 'wb')
106      f.write('\n'.join(item_list))
107      f.close()
108      command = '%s %s' % (switch_script, temp_file)
109    else:
110      command = '%s %s' % (switch_script, ' '.join(item_list))
111    ret, _, _ = self.ce.RunCommandWExceptionCleanup(
112        command, print_to_console=self.verbose)
113    assert ret == 0, 'Switch script %s returned %d' % (switch_script, ret)
114
115  def TestScript(self):
116    command = self.test_script
117    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
118    return ret
119
120  def InstallScript(self):
121    if not self.install_script:
122      return 0
123
124    command = self.install_script
125    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
126    return ret
127
128  def DoVerify(self):
129    self.l.LogOutput('VERIFICATION')
130    self.l.LogOutput('Beginning %d tests to verify good/bad sets\n')
131    for _ in range(int(self.verify_level)):
132      self.l.LogOutput('Resetting all items to good to verify.')
133      self.SwitchToGood(self.all_items)
134      status = self.InstallScript()
135      assert status == 0, 'When reset_to_good, install should succeed.'
136      status = self.TestScript()
137      assert status == 0, 'When reset_to_good, status should be 0.'
138
139      self.l.LogOutput('Resetting all items to bad to verify.')
140      self.SwitchToBad(self.all_items)
141      status = self.InstallScript()
142      assert status == 0, 'When reset_to_bad, install should succeed.'
143      status = self.TestScript()
144      assert status == 1, 'When reset_to_bad, status should be 1.'
145
146  def DoSearch(self):
147    num_bad_items_history = []
148    self.prune_cycles = 0
149    while (True and
150           len(self.all_items) > 1 and
151           self.prune_cycles < self.prune_iterations):
152      self.prune_cycles += 1
153      terminated = self.DoBinarySearch()
154      if not terminated:
155        break
156      if not self.prune:
157        self.l.LogOutput('Not continuning further, --prune is not set')
158        break
159      # Prune is set.
160      prune_index = self.bs.current
161
162      if prune_index == len(self.all_items) - 1:
163        self.l.LogOutput('First bad item is the last item. Breaking.')
164        self.l.LogOutput('Only bad item is: %s' % self.all_items[-1])
165        break
166
167      num_bad_items = len(self.all_items) - prune_index
168      num_bad_items_history.append(num_bad_items)
169
170      if (num_bad_items_history[-num_bad_items:] ==
171          [num_bad_items for _ in range(num_bad_items)]):
172        self.l.LogOutput('num_bad_items_history: %s for past %d iterations. '
173                         'Breaking.' % (str(num_bad_items_history),
174                                        num_bad_items))
175        self.l.LogOutput('Bad items are: %s' %
176                         ' '.join(self.all_items[prune_index:]))
177        break
178
179      new_all_items = list(self.all_items)
180      # Move prune item to the end of the list.
181      new_all_items.append(new_all_items.pop(prune_index))
182      self.found_items.add(new_all_items[-1])
183
184      if prune_index:
185        new_all_items = new_all_items[prune_index - 1:]
186
187      self.l.LogOutput('Old list: %s. New list: %s' % (str(self.all_items),
188                                                       str(new_all_items)),
189                       print_to_console=self.verbose)
190
191      # FIXME: Do we need to Convert the currently good items to bad
192      self.PopulateItemsUsingList(new_all_items)
193
194  def DoBinarySearch(self):
195    self.search_cycles = 0
196    terminated = False
197    while self.search_cycles < self.iterations and not terminated:
198      self.OutputProgress()
199      self.search_cycles += 1
200      [bad_items, good_items] = self.GetNextItems()
201
202      # TODO: bad_items should come first.
203      self.SwitchToGood(good_items)
204      self.SwitchToBad(bad_items)
205      status = self.InstallScript()
206      if status == 0:
207        status = self.TestScript()
208      else:
209        # Install script failed, treat as skipped item
210        status = 2
211      terminated = self.bs.SetStatus(status)
212
213      if terminated:
214        self.l.LogOutput('Terminated!')
215    if not terminated:
216      self.l.LogOutput('Ran out of iterations searching...')
217    self.l.LogOutput(str(self), print_to_console=self.verbose)
218    return terminated
219
220  def PopulateItemsUsingCommand(self, command):
221    ce = command_executer.GetCommandExecuter()
222    _, out, _ = ce.RunCommandWExceptionCleanup(command,
223                                               return_output=True,
224                                               print_to_console=self.verbose)
225    all_items = out.split()
226    self.PopulateItemsUsingList(all_items)
227
228  def PopulateItemsUsingList(self, all_items):
229    self.all_items = all_items
230    self.bs = binary_search_perforce.BinarySearcher(logger_to_set=self.l)
231    self.bs.SetSortedList(self.all_items)
232
233  def SaveState(self):
234    self.l = None
235    self.ce = None
236    # TODO Implement save/restore
237    ###    return
238    f = open(STATE_FILE, 'wb')
239    pickle.dump(self, f)
240    f.close()
241
242  @classmethod
243  def LoadState(cls):
244    if not os.path.isfile(STATE_FILE):
245      return None
246    return pickle.load(file(STATE_FILE))
247
248  def GetNextItems(self):
249    border_item = self.bs.GetNext()
250    index = self.all_items.index(border_item)
251
252    next_bad_items = self.all_items[:index + 1]
253    next_good_items = self.all_items[index + 1:]
254
255    return [next_bad_items, next_good_items]
256
257  def OutputProgress(self):
258    out = ('\n********** PROGRESS **********\n'
259           'Search %d of estimated %d.\n'
260           'Prune %d of max %d.\n'
261           'Current bad items found:\n'
262           '%s\n'
263           '******************************')
264    out = out % (self.search_cycles + 1,
265                 math.ceil(math.log(len(self.all_items), 2)),
266                 self.prune_cycles,
267                 self.prune_iterations,
268                 str(self.found_items))
269
270    self.l.LogOutput(out)
271
272  def __str__(self):
273    ret = ''
274    ret += 'all: %s\n' % str(self.all_items)
275    ret += 'currently_good: %s\n' % str(self.currently_good_items)
276    ret += 'currently_bad: %s\n' % str(self.currently_bad_items)
277    ret += str(self.bs)
278    return ret
279
280
281def _CanonicalizeScript(script_name):
282  script_name = os.path.expanduser(script_name)
283  if not script_name.startswith('/'):
284    return os.path.join('.', script_name)
285
286
287def Main(argv):
288  """The main function."""
289  # Common initializations
290
291  parser = argparse.ArgumentParser()
292  parser.add_argument('-n',
293                      '--iterations',
294                      dest='iterations',
295                      help='Number of iterations to try in the search.',
296                      default=50)
297  parser.add_argument('-i',
298                      '--get_initial_items',
299                      dest='get_initial_items',
300                      help=('Script to run to get the initial objects. '
301                            'If your script requires user input '
302                            'the --verbose option must be used'))
303  parser.add_argument('-g',
304                      '--switch_to_good',
305                      dest='switch_to_good',
306                      help=('Script to run to switch to good. '
307                            'If your switch script requires user input '
308                            'the --verbose option must be used'))
309  parser.add_argument('-b',
310                      '--switch_to_bad',
311                      dest='switch_to_bad',
312                      help=('Script to run to switch to bad. '
313                            'If your switch script requires user input '
314                            'the --verbose option must be used'))
315  parser.add_argument('-I',
316                      '--install_script',
317                      dest='install_script',
318                      default=None,
319                      help=('Optional script to perform building, flashing, '
320                            'and other setup before the test script runs.'))
321  parser.add_argument('-t',
322                      '--test_script',
323                      dest='test_script',
324                      help=('Script to run to test the '
325                            'output after packages are built.'))
326  parser.add_argument('-p',
327                      '--prune',
328                      dest='prune',
329                      action='store_true',
330                      default=False,
331                      help=('Script to run to test the output after '
332                            'packages are built.'))
333  parser.add_argument('-c',
334                      '--noincremental',
335                      dest='noincremental',
336                      action='store_true',
337                      default=False,
338                      help='Do not propagate good/bad changes incrementally.')
339  parser.add_argument('-f',
340                      '--file_args',
341                      dest='file_args',
342                      action='store_true',
343                      default=False,
344                      help='Use a file to pass arguments to scripts.')
345  parser.add_argument('-v',
346                      '--verify_level',
347                      dest='verify_level',
348                      default=1,
349                      help=('Check binary search assumptions N times '
350                            'before starting.'))
351  parser.add_argument('-N',
352                      '--prune_iterations',
353                      dest='prune_iterations',
354                      help='Number of prune iterations to try in the search.',
355                      default=100)
356  parser.add_argument('-V',
357                      '--verbose',
358                      dest='verbose',
359                      action='store_true',
360                      help='Print full output to console.')
361
362  logger.GetLogger().LogOutput(' '.join(argv))
363  options = parser.parse_args(argv)
364
365  if not (options.get_initial_items and options.switch_to_good and
366          options.switch_to_bad and options.test_script):
367    parser.print_help()
368    return 1
369
370  iterations = int(options.iterations)
371  switch_to_good = _CanonicalizeScript(options.switch_to_good)
372  switch_to_bad = _CanonicalizeScript(options.switch_to_bad)
373  install_script = options.install_script
374  if install_script:
375    install_script = _CanonicalizeScript(options.install_script)
376  test_script = _CanonicalizeScript(options.test_script)
377  get_initial_items = _CanonicalizeScript(options.get_initial_items)
378  prune = options.prune
379  prune_iterations = options.prune_iterations
380  verify_level = options.verify_level
381  file_args = options.file_args
382  verbose = options.verbose
383  binary_search_perforce.verbose = verbose
384
385  if options.noincremental:
386    incremental = False
387  else:
388    incremental = True
389
390  try:
391    bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad,
392                            install_script, test_script, incremental, prune,
393                            iterations, prune_iterations, verify_level,
394                            file_args, verbose)
395    bss.DoVerify()
396    bss.DoSearch()
397
398  except (KeyboardInterrupt, SystemExit):
399    print('C-c pressed')
400    bss.SaveState()
401  return 0
402
403
404if __name__ == '__main__':
405  sys.exit(Main(sys.argv[1:]))
406