Python profiling to chrome://tracing

python_profiling.py

#
# Example of how to profile a Python app with multiple processes
# by logging events and opening the resulting trace file in chrome://tracing.
#

# pip install multiprocessing_logging

from functools import wraps
import json
import logging
from multiprocessing import Pool
import os
import random
import time
import threading

from multiprocessing_logging import install_mp_handler

# we want to be able to log from multiple processes
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger()
install_mp_handler()

# separate logger that only stores events into a file
prof_logger = logging.getLogger('profiling')
# not not propagate to root logger, we want to store this separately
prof_logger.propagate = False
handler = logging.FileHandler('profiling.log')
handler.setFormatter(logging.Formatter("%(message)s"))
prof_logger.addHandler(handler)
install_mp_handler(prof_logger)

def log_profile(category=None):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            # format compatible with chrome://tracing
            # info: https://www.gamasutra.com/view/news/176420/Indepth_Using_Chrometracing_to_view_your_inline_profiling_data.php
            base_info = {
                'name': f.__name__,
                'pid': os.getpid(),
                'tid': threading.current_thread().ident,
                'cat': category,
            }
            def log_event(**kwargs):
                prof_logger.debug(json.dumps(kwargs))
                   
            def time_usec():
                return int(round(1e6 * time.time()))

            start_time = time_usec()
            log_event(ph='B', ts=start_time, **base_info)

            result = f(*args, **kwargs)

            end_time = time_usec()
            duration = end_time - start_time
            # TODO: duration could possibly be computed afterwards (if we can pair the events correctly)
            log_event(ph='E', ts=end_time, duration=duration, **base_info)

            return result
        return wrapper
    return decorator

def convert_log_to_trace(log_file, trace_file):
    with open(trace_file, 'wt') as output, open(log_file, 'rt') as input:
        events = [json.loads(line) for line in input]
        json.dump({"traceEvents": events}, output)

@log_profile('compute')
def foo(i):
    time.sleep(random.uniform(.5, 1.))
    result = i**2
    logger.info('foo(): %s -> %s', i, result)
    return result

pool = Pool(4)
pool.map(foo, range(40))
pool.close()

convert_log_to_trace('profiling.log', 'profiling_trace.json')

# open `profiling_trace.json` in chrome://tracing -> Load