Source

yt.profiling_support / profile_support.py

import time
import cProfile
import cProfile
import pkg_resources, os, sys, time
import numpy as na

from yt.mods import *
import yt.utilities.command_line as cl

try:
    import pyprof2html.core
except ImportError:
    pass

def get_version_string():
    yt_provider = pkg_resources.get_provider("yt")
    path = os.path.dirname(yt_provider.module_path)
    vstring = cl._get_hg_version(path)
    my_version = vstring.replace("(", "").replace(")","").replace(" ", "_").strip()
    return my_version

_my_version = get_version_string()

functions_to_profile = {}

def profile_test(func):
    functions_to_profile[func.func_name] = func
    return func

def run_all_tests(pf, ntimes = 2, output_html = True):
    from mpi4py import MPI
    rank = MPI.COMM_WORLD.rank
    size = MPI.COMM_WORLD.size
    td = None
    if rank == 0:
        td = time.strftime("%Y%m%d-%H%M")
    td = MPI.COMM_WORLD.bcast(td, root = 0)
    results = {}
    for fname in sorted(functions_to_profile):
        print "HANDLING %s on %s / %s" % (
            fname, rank, size)
        fnprefix = "profiles/%s_%s" % (fname, _my_version)
        times = []
        for n in xrange(ntimes):
            def new_func():
                func(pf)
            MPI.COMM_WORLD.Barrier()
            t1 = time.time()
            with parallel_profile(fnprefix) as ofn:
                functions_to_profile[fname](pf)
            MPI.COMM_WORLD.Barrier()
            t2 = time.time()
            times.append(t2-t1)
        results[fname] = (na.min(times), na.max(times), na.average(times))
        if MPI.COMM_WORLD.rank == 0:
            print "Results [%s]: %0.3e %0.3e %0.3e" % (
                fname, results[fname][0], results[fname][1], results[fname][2])
        if output_html:
            print "Writing output file on %03i / %03i" % (rank, size)
            # We step through it
            dname = "html/%s-%03i_%03i-%s/%s" % (td, size, rank, _my_version, fname)
            if rank == 0 and not os.path.exists(dname):
                os.makedirs(dname)
            MPI.COMM_WORLD.Barrier()
            # We assume that the final output is the only interesting one.
            ppw = pyprof2html.core.Converter(ofn)
            ppw.printout(filetype="html",
                  output_directory=dname,
                  output_htmlfile="index-%03i_%03i.html" % (size, rank),
                  functions_number=99999)
    return results