core.py

from __future__ import print_function

import argparse
import json
import os
import traceback
import logging
from time import time
from types import FunctionType
from collections import OrderedDict

import numpy as np
import matplotlib.pyplot as plt
from jinja2 import Template

import multiprocessing
import platform

log = logging.getLogger("run_benchmarks")

REPORT_FILENAME = 'report/index.html'
DATA_FILENAME = 'report/benchmark_results.json'

LOG_FORMAT = '%(asctime)s %(levelname)-8s %(name)-8s %(message)s'


# =========================
# BENCHMARK DISCOVERY
# =========================

def find_benchmarks(folders=None, platforms=None):
    """Collect benchmarks callable and shared environment initializers."""
    benchmark_groups = []

    if folders is None:
        here = os.path.dirname(os.path.abspath(__file__))
        folders = [
            f for f in os.listdir(here)
            if os.path.isdir(f) and
               os.path.exists(os.path.join(f, '__init__.py'))
        ]
        folders.sort()

    for folder in folders:
        group_name = os.path.basename(folder)
        collected_benchmarks = []
        modules_in_error = []

        pkg = __import__(group_name, fromlist="dummy")

        benchmark_groups.append(OrderedDict([
            ('name', group_name),
            ('make_env', getattr(pkg, 'make_env', None)),
            ('benchmarks', collected_benchmarks),
            ('import_errors', modules_in_error),
        ]))

        for module_filename in sorted(os.listdir(folder)):
            module_name, ext = os.path.splitext(module_filename)

            if ext not in ('.py', '.so', '.dll', '.pyx'):
                continue

            if not module_name.startswith(group_name + "_"):
                continue

            abs_module_name = "%s.%s" % (group_name, module_name)

            try:
                module = __import__(abs_module_name, fromlist="dummy")
            except Exception as e:
                modules_in_error.append({
                    'name': module_name,
                    'error': str(e),
                })
                module = None

            for benchmark in getattr(module, 'benchmarks', ()):
                if callable(benchmark):
                    collected_benchmarks.append(
                        (module_filename, benchmark.__name__, benchmark)
                    )

    return benchmark_groups


# =========================
# BENCHMARK EXECUTION
# =========================

def run_benchmark(name, func, args, kwargs, n_runs=5):
    def time_once():
        tic = time()
        func(*args, **kwargs)
        return time() - tic

    first_timing = time_once()

    if isinstance(func, FunctionType):
        cold = None
        warm = first_timing
        all_warm_timings = [first_timing]
    else:
        other_timings = []
        for _ in range(n_runs - 1):
            other_timings.append(time_once())

        cold = first_timing
        warm = np.min(other_timings)
        all_warm_timings = other_timings

    return OrderedDict([
        ('name', name),
        ('cold_time', cold),
        ('warm_time', warm),
        ('all_warm_times', all_warm_timings),
        ('std_warm_times', np.std(all_warm_timings)),
    ])


def run_benchmarks(folders=None, platforms=None):
    collected = find_benchmarks(folders=folders, platforms=platforms)
    results = []

    for group in collected:
        records = []
        make_env = group.get('make_env')
        args, kwargs = make_env() if make_env else ((), {})

        for module_source, name, func in group['benchmarks']:
            record = run_benchmark(name, func, args, kwargs)
            records.append(record)

        results.append({
            'group_name': group['name'],
            'records': records
        })

    return results


# =========================
# REPORT
# =========================

def build_report(bench_data,
                 report_filename=REPORT_FILENAME):

    with open(report_filename, "w") as f:
        f.write("Report Generated")


# =========================
# CLI
# =========================

def parse_args(args):
    parser = argparse.ArgumentParser()
    parser.add_argument('--folders', nargs='*', default=None)
    parser.add_argument('--platforms', nargs='*', default=None)
    return parser.parse_args(args)