Source code for schrodinger.utils.profiling

"""
Profiling and timing decorators and context managers.

For more information on profiling python code, see also
https://confluence.schrodinger.com/display/SOFTWAREDEV/Profiling
"""
import atexit
import contextlib
import cProfile
import logging
from functools import wraps
from statistics import fmean
from statistics import quantiles
from statistics import stdev
from time import perf_counter
from time import strftime
from typing import List


[docs]class Timer: """ Track and report wall time elapsed during a function or code block execution. Can be used as a function decorator or context manager. Both the context manager and decorator will output time-related messages to the log at the specified logging level. The context manager can be assigned to a variable, which can be accessed inside the context to get elapsed time and log checkpoints. It can also be inspected after the context is closed to get total time elapsed. Elapsed time does not update after context is closed (until a new context is opened). Before we get into examples, let's set up our imports and logger boilerplate:: >>> from schrodinger.utils import log >>> from schrodinger.utils.profiling import Timer >>> >>> logger = log.get_output_logger(__name__) >>> logger.setLevel(log.DEBUG) Now, to use this class as a context manager:: >>> with Timer(activity='do foo', logger=logger, level=log.DEBUG) as t: ... for _ in range(100): ... # Do something expensive here ... pass ... # Write a checkpoint message to log ... checkpoint = t.checkpoint('first thing done') ... for _ in range(100): ... # Do something else here ... pass Timer - do foo - Starting - ... Timer - do foo - Checkpoint - first thing done - ... s Timer - do foo - Completed - ... s Alternatively, to use the class as a decorator:: >>> @Timer(logger=logger, level=log.INFO) ... def do_foo(): ... for _ in range(100): ... pass >>> do_foo() Timer - do_foo - Starting - ... Timer - do_foo - Completed - ... s :cvar TIME_FORMAT: The time format string used by `strftime` when reporting the time that measurement began. :cvar PRECISION: The precision to be used when reporting elapsed seconds. """ TIME_FORMAT: str = '%Y-%m-%d %H:%M:%S' PRECISION: int = 4
[docs] def __init__(self, *, activity=None, logger=None, level=logging.INFO): """ :param activity: A description of the activity to be timed. If this class is used as a decorator and `activity` is NoneType, the function name will be used. :type activity: Optional[str] :param logger: A logger to receive timing messages. If none is provided, messages will be printed to stdout using the builtin `print` function. To silence messages entirely, use a logger with a NullHandler. :type logger: Optional[logging.Logger] :param int level: The logger level to use for timing messages. """ self.activity = activity self.logger = logger self.level = level self._checkpoint = 0 self._start_time = None self._elapsed = None
def __call__(self, func): if self.activity is None: self.activity = func.__name__ # Just like contextlib.ContextDecorator @wraps(func) def inner(*args, **kwds): with self: return func(*args, **kwds) return inner @property def elapsed(self) -> float: """ Fractional seconds elapsed since context was entered. If context manager has already exited, this is the total execution time of the code block inside the context manager. :return: Fractional seconds elapsed since the start time. """ if self._elapsed is not None: return self._elapsed if self._start_time is not None: return perf_counter() - self._start_time raise RuntimeError('Nothing has been timed.')
[docs] def checkpoint(self, description=None) -> float: """ Write a checkpoint to the log and get elapsed time. Be aware that logging can add non-negligible time to a code block's execution. Empirically, each call to this function appears add about 0.5ms. :param description: A description of the checkpoint. If not supplied, the sequential checkpoint number will be used. :return: Fractional seconds elapsed since the start time. """ elapsed = self.elapsed if description is None: description = self._checkpoint self._checkpoint += 1 self._log('Checkpoint', description, self._format_seconds(elapsed)) return elapsed
def _generate_log_entry(self, items): cls_name = self.__class__.__name__ activity = self.activity or '' return ' - '.join(map(str, [cls_name, activity, *items])) def _log(self, *items): entry = self._generate_log_entry(items) if self.logger: self.logger.log(self.level, entry) else: print(entry) def _format_seconds(self, seconds): return f'{seconds:.{self.PRECISION}f} s' def __enter__(self): self._log('Starting', strftime(self.TIME_FORMAT)) self._elapsed = None self._start_time = perf_counter() return self def __exit__(self, *exc): self._elapsed = perf_counter() - self._start_time self._log('Completed', self._format_seconds(self._elapsed)) return False
[docs]class RecordingTimer(Timer): """ Decorator that will print the duration of the execution of decorated function or method to the log every time it's called. Will also report summary statistics of all call times at program exit, much like a flat profiler. """ _timings: List[float] def __call__(self, func): if self.activity is None: self.activity = func.__name__ self._timings = [] @wraps(func) def inner(*args, **kwargs): with self: retval = func(*args, **kwargs) self._timings.append(self.elapsed) return retval atexit.register(self._summary_report) return inner def _compute_metrics(self): data = self._timings if len(data) < 2: return () q25, q50, q75 = quantiles(data, method='inclusive') mean = fmean(data) std_dev = stdev(data, xbar=mean) metrics = ( ('Sum', sum(data)), ('Mean', mean), ('StdDev', std_dev), ('Min', min(data)), ('Q25', q25), ('Q50', q50), ('Q75', q75), ('Max', max(data)), ) return metrics def _summary_report(self): self._log('Summary', 'N', len(self._timings)) for name, value in self._compute_metrics(): self._log('Summary', name, self._format_seconds(value))
[docs]def profile_call(func, *args, profile_filename=None, **kwargs): """ Profile a single function call. :param func: The function to profile. All arguments to `profile_call` other than `profile_filename` will be passed to this function. :type func: Callable :param profile_filename: The name of the file to save the profiling data to. Note that this argument is keyword-only and is required. :type profile_filename: str :return: The value returned by `func`. :rtype: object """ if profile_filename is None: raise ValueError("No filename given.") profiler = cProfile.Profile() retval = profiler.runcall(func, *args, **kwargs) profiler.dump_stats(profile_filename) return retval
[docs]class Profiler(contextlib.ContextDecorator): """ Profile a block of code or function. This class can be used as either a decorator or a context manager. Measurements are cumulative. Thus if the same instance is used multiple times, the stats from all previous uses of that instance will be included in the profiling summary output. To "use an instance multiple times" in this context means either repeated calls to a decorated function or using the same instance to enter several contexts with repeated uses of the `with` keyword. """
[docs] def __init__(self, outfile=None): """ :param outfile: The file name to be used for the profiling summary output (traditionally '.prof' extension). If no outfile is provided, output will be printed to stdout. Note that if the same instance is used multiple times, the file will be updated with the latest values and the previous output will be overwritten (see class documentation for details). :type outfile: Optional[str] """ self.outfile = outfile self.profiler = cProfile.Profile()
def __enter__(self): self.profiler.enable() return self.profiler def __exit__(self, *exc): self.profiler.disable() if self.outfile is not None: self.profiler.dump_stats(self.outfile) else: self.profiler.print_stats() return False