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     """
  46     Top-level exception handler for exceptions that make it past any exception
  47     handlers in the python code being run.  Logs the error and stacktrace then
  48     maybe attaches a debugger.
  49     """
  50     global original_hook
  51     msg = f'Unhandled top level exception {exc_type}'
  52     logger.exception(msg)
  53     print(msg, file=sys.stderr)
  54     if issubclass(exc_type, KeyboardInterrupt):
  55         sys.__excepthook__(exc_type, exc_value, exc_tb)
  56         return
  57     else:
  58         if (
  59                 not sys.stderr.isatty() or
  60                 not sys.stdin.isatty()
  61         ):
  62             # stdin or stderr is redirected, just do the normal thing
  63             original_hook(exc_type, exc_value, exc_tb)
  64         else:
  65             # a terminal is attached and stderr is not redirected, maybe debug.
  66             import traceback
  67             traceback.print_exception(exc_type, exc_value, exc_tb)
  68             if config.config['debug_unhandled_exceptions']:
  69                 import pdb
  70                 logger.info("Invoking the debugger...")
  71                 pdb.pm()
  72             else:
  73                 original_hook(exc_type, exc_value, exc_tb)
  74
  75
  76 def initialize(entry_point):
  77     """
  78     Remember to initialize config, initialize logging, set/log a random
  79     seed, etc... before running main.
  80
  81     """
  82     @functools.wraps(entry_point)
  83     def initialize_wrapper(*args, **kwargs):
  84
  85         # Hook top level unhandled exceptions, maybe invoke debugger.
  86         if sys.excepthook == sys.__excepthook__:
  87             sys.excepthook = handle_uncaught_exception
  88
  89         # Try to figure out the name of the program entry point.  Then
  90         # parse configuration (based on cmdline flags, environment vars
  91         # etc...)
  92         if (
  93                 '__globals__' in entry_point.__dict__ and
  94                 '__file__' in entry_point.__globals__
  95         ):
  96             config.parse(entry_point.__globals__['__file__'])
  97         else:
  98             config.parse(None)
  99
 100         # Initialize logging... and log some remembered messages from
 101         # config module.
 102         logging_utils.initialize_logging(logging.getLogger())
 103         config.late_logging()
 104
 105         # Allow programs that don't bother to override the random seed
 106         # to be replayed via the commandline.
 107         import random
 108         random_seed = config.config['set_random_seed']
 109         if random_seed is not None:
 110             random_seed = random_seed[0]
 111         else:
 112             random_seed = int.from_bytes(os.urandom(4), 'little')
 113
 114         if config.config['show_random_seed']:
 115             msg = f'Global random seed is: {random_seed}'
 116             print(msg)
 117             logger.debug(msg)
 118         random.seed(random_seed)
 119
 120         # Do it, invoke the user's code.  Pay attention to how long it takes.
 121         logger.debug(f'Starting {entry_point.__name__} (program entry point)')
 122         ret = None
 123         import stopwatch
 124         with stopwatch.Timer() as t:
 125             ret = entry_point(*args, **kwargs)
 126         logger.debug(
 127             f'{entry_point.__name__} (program entry point) returned {ret}.'
 128         )
 129
 130         walltime = t()
 131         (utime, stime, cutime, cstime, elapsed_time) = os.times()
 132         logger.debug('\n'
 133                      f'user: {utime}s\n'
 134                      f'system: {stime}s\n'
 135                      f'child user: {cutime}s\n'
 136                      f'child system: {cstime}s\n'
 137                      f'machine uptime: {elapsed_time}s\n'
 138                      f'walltime: {walltime}s')
 139
 140         # If it doesn't return cleanly, call attention to the return value.
 141         if ret is not None and ret != 0:
 142             logger.error(f'Exit {ret}')
 143         else:
 144             logger.debug(f'Exit {ret}')
 145         sys.exit(ret)
 146     return initialize_wrapper