bootstrap.py

   1 #!/usr/bin/env python3
   2
   3 import functools
   4 import importlib
   5 import logging
   6 import os
   7 import sys
   8 from inspect import stack
   9 from typing import List
  10
  11 import config
  12 import logging_utils
  13 from argparse_utils import ActionNoYes
  14
  15 # This module is commonly used by others in here and should avoid
  16 # taking any unnecessary dependencies back on them.
  17
  18
  19 logger = logging.getLogger(__name__)
  20
  21 args = config.add_commandline_args(
  22     f'Bootstrap ({__file__})',
  23     'Args related to python program bootstrapper and Swiss army knife',
  24 )
  25 args.add_argument(
  26     '--debug_unhandled_exceptions',
  27     action=ActionNoYes,
  28     default=False,
  29     help='Break into pdb on top level unhandled exceptions.',
  30 )
  31 args.add_argument(
  32     '--show_random_seed',
  33     action=ActionNoYes,
  34     default=False,
  35     help='Should we display (and log.debug) the global random seed?',
  36 )
  37 args.add_argument(
  38     '--set_random_seed',
  39     type=int,
  40     nargs=1,
  41     default=None,
  42     metavar='SEED_INT',
  43     help='Override the global random seed with a particular number.',
  44 )
  45 args.add_argument(
  46     '--dump_all_objects',
  47     action=ActionNoYes,
  48     default=False,
  49     help='Should we dump the Python import tree before main?',
  50 )
  51 args.add_argument(
  52     '--audit_import_events',
  53     action=ActionNoYes,
  54     default=False,
  55     help='Should we audit all import events?',
  56 )
  57 args.add_argument(
  58     '--run_profiler',
  59     action=ActionNoYes,
  60     default=False,
  61     help='Should we run cProfile on this code?',
  62 )
  63 args.add_argument(
  64     '--trace_memory',
  65     action=ActionNoYes,
  66     default=False,
  67     help='Should we record/report on memory utilization?',
  68 )
  69
  70 original_hook = sys.excepthook
  71
  72
  73 def handle_uncaught_exception(exc_type, exc_value, exc_tb):
  74     """
  75     Top-level exception handler for exceptions that make it past any exception
  76     handlers in the python code being run.  Logs the error and stacktrace then
  77     maybe attaches a debugger.
  78
  79     """
  80     global original_hook
  81     msg = f'Unhandled top level exception {exc_type}'
  82     logger.exception(msg)
  83     print(msg, file=sys.stderr)
  84     if issubclass(exc_type, KeyboardInterrupt):
  85         sys.__excepthook__(exc_type, exc_value, exc_tb)
  86         return
  87     else:
  88         if not sys.stderr.isatty() or not sys.stdin.isatty():
  89             # stdin or stderr is redirected, just do the normal thing
  90             original_hook(exc_type, exc_value, exc_tb)
  91         else:
  92             # a terminal is attached and stderr is not redirected, maybe debug.
  93             import traceback
  94
  95             traceback.print_exception(exc_type, exc_value, exc_tb)
  96             if config.config['debug_unhandled_exceptions']:
  97                 import pdb
  98
  99                 logger.info("Invoking the debugger...")
 100                 pdb.pm()
 101             else:
 102                 original_hook(exc_type, exc_value, exc_tb)
 103
 104
 105 class ImportInterceptor(importlib.abc.MetaPathFinder):
 106     def __init__(self):
 107         import collect.trie
 108
 109         self.module_by_filename_cache = {}
 110         self.repopulate_modules_by_filename()
 111         self.tree = collect.trie.Trie()
 112         self.tree_node_by_module = {}
 113
 114     def repopulate_modules_by_filename(self):
 115         self.module_by_filename_cache.clear()
 116         for mod in sys.modules:
 117             if hasattr(sys.modules[mod], '__file__'):
 118                 fname = getattr(sys.modules[mod], '__file__')
 119             else:
 120                 fname = 'unknown'
 121             self.module_by_filename_cache[fname] = mod
 122
 123     def should_ignore_filename(self, filename: str) -> bool:
 124         return 'importlib' in filename or 'six.py' in filename
 125
 126     def find_module(self, fullname, path):
 127         raise Exception("This method has been deprecated since Python 3.4, please upgrade.")
 128
 129     def find_spec(self, loaded_module, path=None, target=None):
 130         s = stack()
 131         for x in range(3, len(s)):
 132             filename = s[x].filename
 133             if self.should_ignore_filename(filename):
 134                 continue
 135
 136             loading_function = s[x].function
 137             if filename in self.module_by_filename_cache:
 138                 loading_module = self.module_by_filename_cache[filename]
 139             else:
 140                 self.repopulate_modules_by_filename()
 141                 loading_module = self.module_by_filename_cache.get(filename, 'unknown')
 142
 143             path = self.tree_node_by_module.get(loading_module, [])
 144             path.extend([loaded_module])
 145             self.tree.insert(path)
 146             self.tree_node_by_module[loading_module] = path
 147
 148             msg = f'*** Import {loaded_module} from {filename}:{s[x].lineno} in {loading_module}::{loading_function}'
 149             logger.debug(msg)
 150             print(msg)
 151             return
 152         msg = f'*** Import {loaded_module} from ?????'
 153         logger.debug(msg)
 154         print(msg)
 155
 156     def invalidate_caches(self):
 157         pass
 158
 159     def find_importer(self, module: str):
 160         if module in self.tree_node_by_module:
 161             node = self.tree_node_by_module[module]
 162             return node
 163         return []
 164
 165
 166 # Audit import events?  Note: this runs early in the lifetime of the
 167 # process (assuming that import bootstrap happens early); config has
 168 # (probably) not yet been loaded or parsed the commandline.  Also,
 169 # some things have probably already been imported while we weren't
 170 # watching so this information may be incomplete.
 171 #
 172 # Also note: move bootstrap up in the global import list to catch
 173 # more import events and have a more complete record.
 174 import_interceptor = None
 175 for arg in sys.argv:
 176     if arg == '--audit_import_events':
 177         import_interceptor = ImportInterceptor()
 178         sys.meta_path.insert(0, import_interceptor)
 179
 180
 181 def dump_all_objects() -> None:
 182     global import_interceptor
 183     messages = {}
 184     all_modules = sys.modules
 185     for obj in object.__subclasses__():
 186         if not hasattr(obj, '__name__'):
 187             continue
 188         klass = obj.__name__
 189         if not hasattr(obj, '__module__'):
 190             continue
 191         class_mod_name = obj.__module__
 192         if class_mod_name in all_modules:
 193             mod = all_modules[class_mod_name]
 194             if not hasattr(mod, '__name__'):
 195                 mod_name = class_mod_name
 196             else:
 197                 mod_name = mod.__name__
 198             if hasattr(mod, '__file__'):
 199                 mod_file = mod.__file__
 200             else:
 201                 mod_file = 'unknown'
 202             if import_interceptor is not None:
 203                 import_path = import_interceptor.find_importer(mod_name)
 204             else:
 205                 import_path = 'unknown'
 206             msg = f'{class_mod_name}::{klass} ({mod_file})'
 207             if import_path != 'unknown' and len(import_path) > 0:
 208                 msg += f' imported by {import_path}'
 209             messages[f'{class_mod_name}::{klass}'] = msg
 210     for x in sorted(messages.keys()):
 211         logger.debug(messages[x])
 212         print(messages[x])
 213
 214
 215 def initialize(entry_point):
 216     """
 217     Remember to initialize config, initialize logging, set/log a random
 218     seed, etc... before running main.
 219
 220     """
 221
 222     @functools.wraps(entry_point)
 223     def initialize_wrapper(*args, **kwargs):
 224         # Hook top level unhandled exceptions, maybe invoke debugger.
 225         if sys.excepthook == sys.__excepthook__:
 226             sys.excepthook = handle_uncaught_exception
 227
 228         # Try to figure out the name of the program entry point.  Then
 229         # parse configuration (based on cmdline flags, environment vars
 230         # etc...)
 231         if '__globals__' in entry_point.__dict__ and '__file__' in entry_point.__globals__:
 232             config.parse(entry_point.__globals__['__file__'])
 233         else:
 234             config.parse(None)
 235
 236         if config.config['trace_memory']:
 237             import tracemalloc
 238
 239             tracemalloc.start()
 240
 241         # Initialize logging... and log some remembered messages from
 242         # config module.
 243         logging_utils.initialize_logging(logging.getLogger())
 244         config.late_logging()
 245
 246         # Maybe log some info about the python interpreter itself.
 247         logger.debug(
 248             f'Platform: {sys.platform}, maxint=0x{sys.maxsize:x}, byteorder={sys.byteorder}'
 249         )
 250         logger.debug(f'Python interpreter version: {sys.version}')
 251         logger.debug(f'Python implementation: {sys.implementation}')
 252         logger.debug(f'Python C API version: {sys.api_version}')
 253         logger.debug(f'Python path: {sys.path}')
 254
 255         # Allow programs that don't bother to override the random seed
 256         # to be replayed via the commandline.
 257         import random
 258
 259         random_seed = config.config['set_random_seed']
 260         if random_seed is not None:
 261             random_seed = random_seed[0]
 262         else:
 263             random_seed = int.from_bytes(os.urandom(4), 'little')
 264
 265         if config.config['show_random_seed']:
 266             msg = f'Global random seed is: {random_seed}'
 267             logger.debug(msg)
 268             print(msg)
 269         random.seed(random_seed)
 270
 271         # Do it, invoke the user's code.  Pay attention to how long it takes.
 272         logger.debug(f'Starting {entry_point.__name__} (program entry point)')
 273         ret = None
 274         import stopwatch
 275
 276         if config.config['run_profiler']:
 277             import cProfile
 278             from pstats import SortKey
 279
 280             with stopwatch.Timer() as t:
 281                 cProfile.runctx(
 282                     "ret = entry_point(*args, **kwargs)",
 283                     globals(),
 284                     locals(),
 285                     None,
 286                     SortKey.CUMULATIVE,
 287                 )
 288         else:
 289             with stopwatch.Timer() as t:
 290                 ret = entry_point(*args, **kwargs)
 291
 292         logger.debug(f'{entry_point.__name__} (program entry point) returned {ret}.')
 293
 294         if config.config['trace_memory']:
 295             snapshot = tracemalloc.take_snapshot()
 296             top_stats = snapshot.statistics('lineno')
 297             print()
 298             print("--trace_memory's top 10 memory using files:")
 299             for stat in top_stats[:10]:
 300                 print(stat)
 301
 302         if config.config['dump_all_objects']:
 303             dump_all_objects()
 304
 305         if config.config['audit_import_events']:
 306             global import_interceptor
 307             if import_interceptor is not None:
 308                 print(import_interceptor.tree)
 309
 310         walltime = t()
 311         (utime, stime, cutime, cstime, elapsed_time) = os.times()
 312         logger.debug(
 313             '\n'
 314             f'user: {utime}s\n'
 315             f'system: {stime}s\n'
 316             f'child user: {cutime}s\n'
 317             f'child system: {cstime}s\n'
 318             f'machine uptime: {elapsed_time}s\n'
 319             f'walltime: {walltime}s'
 320         )
 321
 322         # If it doesn't return cleanly, call attention to the return value.
 323         if ret is not None and ret != 0:
 324             logger.error(f'Exit {ret}')
 325         else:
 326             logger.debug(f'Exit {ret}')
 327         sys.exit(ret)
 328
 329     return initialize_wrapper