File size: 2,809 Bytes
2abfccb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import cProfile
import pstats
import io
import functools
import os
from collections import defaultdict
from distutils.util import strtobool
PETREL_PROFILE_ENV = os.getenv('PETREL_PROFILE', 'False')
try:
ENABLE_PROFILE = strtobool(PETREL_PROFILE_ENV)
except ValueError:
raise ValueError(
f'invalid value of environment variable PETREL_PROFILE: {PETREL_PROFILE_ENV}')
PROFILE_COUNT_ENV = os.getenv('PETREL_PROFILE_COUNT', 1000)
try:
PROFILE_COUNT = int(PROFILE_COUNT_ENV)
except ValueError:
raise ValueError(
f'invalid value of environment variable PETREL_PROFILE_COUNT: {PROFILE_COUNT_ENV}')
WORKER_LOOP_PROFILE_COUNT_ENV = os.getenv(
'PETREL_WORKER_LOOP_PROFILE_COUNT', 250)
try:
WORKER_LOOP_PROFILE_COUNT = int(WORKER_LOOP_PROFILE_COUNT_ENV)
except ValueError:
raise ValueError(
f'invalid value of environment variable PETREL_WORKER_LOOP_PROFILE_COUNT: {WORKER_LOOP_PROFILE_COUNT_ENV}')
def print_stats(prof, name, sortby='cumulative'):
s = io.StringIO()
if name:
s.write(f'\nProfile of function {name}:\n')
s.write(f'pid: {os.getpid()}\n')
ps = pstats.Stats(prof, stream=s).sort_stats(sortby)
ps.print_stats()
print(s.getvalue())
def profile_helper(func, name, count):
if not ENABLE_PROFILE:
return func
prof = cProfile.Profile()
call_count = 0
if not name:
try:
name = func.__name__
except AttributeError:
pass
@functools.wraps(func)
def wrapper(*args, **kwargs):
nonlocal prof
nonlocal call_count
try:
return prof.runcall(func, *args, **kwargs)
finally:
call_count += 1
if call_count == count:
print_stats(prof, name)
call_count = 0
prof = cProfile.Profile()
return wrapper
def profileit(*args, name=None, count=PROFILE_COUNT):
if args:
assert len(args) == 1 and callable(args[0])
return profile_helper(args[0], name, count)
else:
return functools.partial(profileit, name=name, count=count)
def wrap_with_stat_qsize(queue, cb, name, count=PROFILE_COUNT):
if not ENABLE_PROFILE:
return cb
cb_count = 0
qsize_dict = defaultdict(lambda: 0)
qsize_list = []
@functools.wraps(cb)
def wrapper(*args, **kwargs):
nonlocal cb_count
cb_count += 1
qsize = queue.qsize()
qsize_dict[qsize] += 1
qsize_list.append(qsize)
try:
return cb(*args, **kwargs)
finally:
if cb_count == count:
print('pid', os.getpid(), name, qsize_dict, '\n', qsize_list)
cb_count = 0
qsize_dict.clear()
qsize_list.clear()
return wrapper
|