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
  10 import config
  11 import logging_utils
  12 from argparse_utils import ActionNoYes
  13
  14 # This module is commonly used by others in here and should avoid
  15 # taking any unnecessary dependencies back on them.
  16
  17
  18 logger = logging.getLogger(__name__)
  19
  20 cfg = config.add_commandline_args(
  21     f'Bootstrap ({__file__})',
  22     'Args related to python program bootstrapper and Swiss army knife',
  23 )
  24 cfg.add_argument(
  25     '--debug_unhandled_exceptions',
  26     action=ActionNoYes,
  27     default=False,
  28     help='Break into pdb on top level unhandled exceptions.',
  29 )
  30 cfg.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 cfg.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 cfg.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 cfg.add_argument(
  51     '--audit_import_events',
  52     action=ActionNoYes,
  53     default=False,
  54     help='Should we audit all import events?',
  55 )
  56 cfg.add_argument(
  57     '--run_profiler',
  58     action=ActionNoYes,
  59     default=False,
  60     help='Should we run cProfile on this code?',
  61 )
  62 cfg.add_argument(
  63     '--trace_memory',
  64     action=ActionNoYes,
  65     default=False,
  66     help='Should we record/report on memory utilization?',
  67 )
  68
  69 ORIGINAL_EXCEPTION_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     msg = f'Unhandled top level exception {exc_type}'
  80     logger.exception(msg)
  81     print(msg, file=sys.stderr)
  82     if issubclass(exc_type, KeyboardInterrupt):
  83         sys.__excepthook__(exc_type, exc_value, exc_tb)
  84         return
  85     else:
  86         if not sys.stderr.isatty() or not sys.stdin.isatty():
  87             # stdin or stderr is redirected, just do the normal thing
  88             ORIGINAL_EXCEPTION_HOOK(exc_type, exc_value, exc_tb)
  89         else:
  90             # a terminal is attached and stderr is not redirected, maybe debug.
  91             import traceback
  92
  93             traceback.print_exception(exc_type, exc_value, exc_tb)
  94             if config.config['debug_unhandled_exceptions']:
  95                 import pdb
  96
  97                 logger.info("Invoking the debugger...")
  98                 pdb.pm()
  99             else:
 100                 ORIGINAL_EXCEPTION_HOOK(exc_type, exc_value, exc_tb)
 101
 102
 103 class ImportInterceptor(importlib.abc.MetaPathFinder):
 104     def __init__(self):
 105         import collect.trie
 106
 107         self.module_by_filename_cache = {}
 108         self.repopulate_modules_by_filename()
 109         self.tree = collect.trie.Trie()
 110         self.tree_node_by_module = {}
 111
 112     def repopulate_modules_by_filename(self):
 113         self.module_by_filename_cache.clear()
 114         for mod in sys.modules:
 115             if hasattr(sys.modules[mod], '__file__'):
 116                 fname = getattr(sys.modules[mod], '__file__')
 117             else:
 118                 fname = 'unknown'
 119             self.module_by_filename_cache[fname] = mod
 120
 121     @staticmethod
 122     def should_ignore_filename(filename: str) -> bool:
 123         return 'importlib' in filename or 'six.py' in filename
 124
 125     def find_module(self, fullname, path):
 126         raise Exception("This method has been deprecated since Python 3.4, please upgrade.")
 127
 128     def find_spec(self, loaded_module, path=None, _=None):
 129         s = stack()
 130         for x in range(3, len(s)):
 131             filename = s[x].filename
 132             if ImportInterceptor.should_ignore_filename(filename):
 133                 continue
 134
 135             loading_function = s[x].function
 136             if filename in self.module_by_filename_cache:
 137                 loading_module = self.module_by_filename_cache[filename]
 138             else:
 139                 self.repopulate_modules_by_filename()
 140                 loading_module = self.module_by_filename_cache.get(filename, 'unknown')
 141
 142             path = self.tree_node_by_module.get(loading_module, [])
 143             path.extend([loaded_module])
 144             self.tree.insert(path)
 145             self.tree_node_by_module[loading_module] = path
 146
 147             msg = f'*** Import {loaded_module} from {filename}:{s[x].lineno} in {loading_module}::{loading_function}'
 148             logger.debug(msg)
 149             print(msg)
 150             return
 151         msg = f'*** Import {loaded_module} from ?????'
 152         logger.debug(msg)
 153         print(msg)
 154
 155     def invalidate_caches(self):
 156         pass
 157
 158     def find_importer(self, module: str):
 159         if module in self.tree_node_by_module:
 160             node = self.tree_node_by_module[module]
 161             return node
 162         return []
 163
 164
 165 # Audit import events?  Note: this runs early in the lifetime of the
 166 # process (assuming that import bootstrap happens early); config has
 167 # (probably) not yet been loaded or parsed the commandline.  Also,
 168 # some things have probably already been imported while we weren't
 169 # watching so this information may be incomplete.
 170 #
 171 # Also note: move bootstrap up in the global import list to catch
 172 # more import events and have a more complete record.
 173 IMPORT_INTERCEPTOR = None
 174 for arg in sys.argv:
 175     if arg == '--audit_import_events':
 176         IMPORT_INTERCEPTOR = ImportInterceptor()
 177         sys.meta_path.insert(0, IMPORT_INTERCEPTOR)
 178
 179
 180 def dump_all_objects() -> None:
 181     messages = {}
 182     all_modules = sys.modules
 183     for obj in object.__subclasses__():
 184         if not hasattr(obj, '__name__'):
 185             continue
 186         klass = obj.__name__
 187         if not hasattr(obj, '__module__'):
 188             continue
 189         class_mod_name = obj.__module__
 190         if class_mod_name in all_modules:
 191             mod = all_modules[class_mod_name]
 192             if not hasattr(mod, '__name__'):
 193                 mod_name = class_mod_name
 194             else:
 195                 mod_name = mod.__name__
 196             if hasattr(mod, '__file__'):
 197                 mod_file = mod.__file__
 198             else:
 199                 mod_file = 'unknown'
 200             if IMPORT_INTERCEPTOR is not None:
 201                 import_path = IMPORT_INTERCEPTOR.find_importer(mod_name)
 202             else:
 203                 import_path = 'unknown'
 204             msg = f'{class_mod_name}::{klass} ({mod_file})'
 205             if import_path != 'unknown' and len(import_path) > 0:
 206                 msg += f' imported by {import_path}'
 207             messages[f'{class_mod_name}::{klass}'] = msg
 208     for x in sorted(messages.keys()):
 209         logger.debug(messages[x])
 210         print(messages[x])
 211
 212
 213 def initialize(entry_point):
 214     """
 215     Remember to initialize config, initialize logging, set/log a random
 216     seed, etc... before running main.
 217
 218     """
 219
 220     @functools.wraps(entry_point)
 221     def initialize_wrapper(*args, **kwargs):
 222         # Hook top level unhandled exceptions, maybe invoke debugger.
 223         if sys.excepthook == sys.__excepthook__:
 224             sys.excepthook = handle_uncaught_exception
 225
 226         # Try to figure out the name of the program entry point.  Then
 227         # parse configuration (based on cmdline flags, environment vars
 228         # etc...)
 229         if '__globals__' in entry_point.__dict__ and '__file__' in entry_point.__globals__:
 230             config.parse(entry_point.__globals__['__file__'])
 231         else:
 232             config.parse(None)
 233
 234         if config.config['trace_memory']:
 235             import tracemalloc
 236
 237             tracemalloc.start()
 238
 239         # Initialize logging... and log some remembered messages from
 240         # config module.
 241         logging_utils.initialize_logging(logging.getLogger())
 242         config.late_logging()
 243
 244         # Maybe log some info about the python interpreter itself.
 245         logger.debug(
 246             'Platform: %s, maxint=0x%x, byteorder=%s',
 247             sys.platform, sys.maxsize, sys.byteorder
 248         )
 249         logger.debug('Python interpreter version: %s', sys.version)
 250         logger.debug('Python implementation: %s', sys.implementation)
 251         logger.debug('Python C API version: %s', sys.api_version)
 252         logger.debug('Python path: %s', sys.path)
 253
 254         # Log something about the site_config, many things use it.
 255         import site_config
 256
 257         logger.debug('Global site_config: %s', site_config.get_config())
 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('Starting %s (program entry point)', entry_point.__name__)
 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('%s (program entry point) returned %s.', entry_point.__name__, 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             if IMPORT_INTERCEPTOR is not None:
 311                 print(IMPORT_INTERCEPTOR.tree)
 312
 313         walltime = t()
 314         (utime, stime, cutime, cstime, elapsed_time) = os.times()
 315         logger.debug(
 316             '\n'
 317             'user: %.4fs\n'
 318             'system: %.4fs\n'
 319             'child user: %.4fs\n'
 320             'child system: %.4fs\n'
 321             'machine uptime: %.4fs\n'
 322             'walltime: %.4fs',
 323             utime, stime, cutime, cstime, elapsed_time, walltime
 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('Exit %s', ret)
 329         else:
 330             logger.debug('Exit %s', ret)
 331         sys.exit(ret)
 332
 333     return initialize_wrapper