bootstrap.py

   1 #!/usr/bin/env python3
   2
   3 import functools
   4 import logging
   5 import os
   6 import sys
   7
   8 # This module is commonly used by others in here and should avoid
   9 # taking any unnecessary dependencies back on them.
  10
  11 from argparse_utils import ActionNoYes
  12 import config
  13 import logging_utils
  14
  15 logger = logging.getLogger(__name__)
  16
  17 args = config.add_commandline_args(
  18     f'Bootstrap ({__file__})',
  19     'Args related to python program bootstrapper and Swiss army knife')
  20 args.add_argument(
  21     '--debug_unhandled_exceptions',
  22     action=ActionNoYes,
  23     default=False,
  24     help='Break into pdb on top level unhandled exceptions.'
  25 )
  26 args.add_argument(
  27     '--show_random_seed',
  28     action=ActionNoYes,
  29     default=False,
  30     help='Should we display (and log.debug) the global random seed?'
  31 )
  32 args.add_argument(
  33     '--set_random_seed',
  34     type=int,
  35     nargs=1,
  36     default=None,
  37     metavar='SEED_INT',
  38     help='Override the global random seed with a particular number.'
  39 )
  40
  41 original_hook = sys.excepthook
  42
  43
  44 def handle_uncaught_exception(exc_type, exc_value, exc_tb):
  45     global original_hook
  46     msg = f'Unhandled top level exception {exc_type}'
  47     logger.exception(msg)
  48     print(msg, file=sys.stderr)
  49     if issubclass(exc_type, KeyboardInterrupt):
  50         sys.__excepthook__(exc_type, exc_value, exc_tb)
  51         return
  52     else:
  53         if (
  54                 not sys.stderr.isatty() or
  55                 not sys.stdin.isatty()
  56         ):
  57             # stdin or stderr is redirected, just do the normal thing
  58             original_hook(exc_type, exc_value, exc_tb)
  59         else:
  60             # a terminal is attached and stderr is not redirected, debug.
  61             if config.config['debug_unhandled_exceptions']:
  62                 import traceback
  63                 import pdb
  64                 traceback.print_exception(exc_type, exc_value, exc_tb)
  65                 logger.info("Invoking the debugger...")
  66                 pdb.pm()
  67             else:
  68                 original_hook(exc_type, exc_value, exc_tb)
  69
  70
  71 def initialize(entry_point):
  72
  73     """Remember to initialize config and logging before running main."""
  74     @functools.wraps(entry_point)
  75     def initialize_wrapper(*args, **kwargs):
  76         if sys.excepthook == sys.__excepthook__:
  77             sys.excepthook = handle_uncaught_exception
  78         if (
  79                 '__globals__' in entry_point.__dict__ and
  80                 '__file__' in entry_point.__globals__
  81         ):
  82             config.parse(entry_point.__globals__['__file__'])
  83         else:
  84             config.parse(None)
  85
  86         logging_utils.initialize_logging(logging.getLogger())
  87
  88         config.late_logging()
  89
  90         # Allow programs that don't bother to override the random seed
  91         # to be replayed via the commandline.
  92         import random
  93         random_seed = config.config['set_random_seed']
  94         if random_seed is not None:
  95             random_seed = random_seed[0]
  96         else:
  97             random_seed = int.from_bytes(os.urandom(4), 'little')
  98
  99         if config.config['show_random_seed']:
 100             msg = f'Global random seed is: {random_seed}'
 101             print(msg)
 102             logger.debug(msg)
 103         random.seed(random_seed)
 104
 105         logger.debug(f'Starting {entry_point.__name__} (program entry point)')
 106
 107         ret = None
 108         import stopwatch
 109         with stopwatch.Timer() as t:
 110             ret = entry_point(*args, **kwargs)
 111         logger.debug(
 112             f'{entry_point.__name__} (program entry point) returned {ret}.'
 113         )
 114
 115         walltime = t()
 116         (utime, stime, cutime, cstime, elapsed_time) = os.times()
 117         logger.debug('\n'
 118                      f'user: {utime}s\n'
 119                      f'system: {stime}s\n'
 120                      f'child user: {cutime}s\n'
 121                      f'child system: {cstime}s\n'
 122                      f'machine uptime: {elapsed_time}s\n'
 123                      f'walltime: {walltime}s')
 124         if ret is not None and ret != 0:
 125             logger.error(f'Exit {ret}')
 126         else:
 127             logger.debug(f'Exit {ret}')
 128         sys.exit(ret)
 129     return initialize_wrapper