bootstrap.py

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