# propro is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# propro is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with propro. If not, see <http://www.gnu.org/licenses/>.
# System imports
from __future__ import print_function, division, absolute_import, unicode_literals
import time
import psutil
import subprocess
import threading
import os
import traceback
import sys
from collections import namedtuple
from datetime import datetime
from propro import format_date
__all__ = ["profile", "profile_pid", "profile_cmd"]
ProfileResult = namedtuple("ProfileResult", ("time", "rss_mem", "vms_mem", "cpu_load", "threads"))
SLEEP_TIME = 0.5
PLOTTING_FMTS = ("pdf", "png")
class TimeoutError(Exception):
pass
def _savetxt(filename, result):
with open(filename, "w") as fp:
for i in range(len(result.time)):
fp.write("\t".join((str(col[i]) for col in result)))
fp.write("\n")
def output(result, fmts, call, t0, baseline=None, save_fig=False):
fig = None
if baseline is not None:
result = ProfileResult(result.time,
[tmp - baseline[0] for tmp in result.rss_mem],
[tmp - baseline[1] for tmp in result.vms_mem],
result.cpu_load,
[tmp - baseline[3] for tmp in result.threads])
for fmt in set(fmts):
if fmt in PLOTTING_FMTS:
from propro import plotting
fig = plotting.profile_plot(result, call, t0, fmt, save_fig)
elif fmt == "txt":
_savetxt("propro_%s_%s.%s"%(call, format_date(t0),fmt), result)
return fig
def _measure(parent, interval=None):
processes = parent.children(recursive=True)
processes.append(parent)
rss_mem=0
vms_mem=0
cpu_percent=0
num_threads=0
for process in processes:
try:
mem_info = process.memory_info()
rss_mem += mem_info[0]
vms_mem += mem_info[1]
cpu_percent += parent.cpu_percent(interval=interval)
num_threads += process.num_threads()
except psutil.NoSuchProcess:
pass
return rss_mem, vms_mem, cpu_percent, num_threads
class Profiler(threading.Thread):
def __init__(self, pid, sleep_time=None):
if sleep_time is None:
sleep_time = SLEEP_TIME
self.sleep_time = sleep_time
self.process = psutil.Process(pid)
self._exception = None
self._cancelled = False
super(Profiler, self).__init__()
def _active(self):
return not self._cancelled and self.process.is_running()
def _profile(self):
times, rss_mems, vms_mems, cpu_percents, threads = [], [], [], [], []
try:
while self._active():
rss_mem, vms_mem, cpu_percent, num_threads = _measure(self.process)
times.append(time.time())
rss_mems.append(rss_mem)
vms_mems.append(vms_mem)
cpu_percents.append(cpu_percent)
threads.append(num_threads)
time.sleep(self.sleep_time)
except psutil.NoSuchProcess:
self.cancel()
except psutil.AccessDenied:
self.cancel()
return ProfileResult(times, rss_mems, vms_mems, cpu_percents, threads)
def run(self):
try:
self._result = self._profile()
except Exception:
exc_info = sys.exc_info()
self._exception = exc_info
def cancel(self):
self._cancelled = True
def exception(self, timeout=None):
self.join(timeout)
if self.isAlive():
raise TimeoutError("Call timed out after: "%timeout)
return self._exception
def result(self, timeout=None):
self.join(timeout)
if self.isAlive():
raise TimeoutError("Call timed out after: "%timeout)
return self._result
[docs]def profile_pid(pid, sample_rate=None, timeout=None):
"""
Profile a specific process id
:param pid: The process id to profile
:param sample_rate: (optional) Rate at which the process is being queried
:param timeout: (optional) Maximal time the process is being profiled
:returns ProfileResult: A `ProfileResult` namedtuple with the profiling result
"""
if not psutil.pid_exists(pid):
raise ValueError("Pid '%s' does not exist"%pid)
profiler = Profiler(pid, sample_rate)
profiler.start()
ex = profiler.exception(timeout)
if ex is not None:
traceback.print_exception(*ex)
raise ex
return profiler.result()
[docs]def profile_cmd(cmd, sample_rate=None, timeout=None):
"""
Profile a specific command
:param cmd: The command to profile
:param sample_rate: (optional) Rate at which the process is being queried
:param timeout: (optional) Maximal time the process is being profiled
:returns ProfileResult: A `ProfileResult` namedtuple with the profiling result
"""
process = subprocess.Popen(cmd, shell=True)
return profile_pid(process.pid, sample_rate, timeout)
[docs]class profile(object):
"""
Decorator to profile a function or method. The profile result is automatically written to disk.
:param sample_rate: (optional) Rate at which the process is being queried
:param timeout: (optional) Maximal time the process is being profiled
:param fmt: (optional) The desired output format. Can also be a tuple of formats. Supported: txt and any matplotlib fmt
:param callname: (optional) Name used for plot title and output file name
"""
def __init__(self, sample_rate=None, timeout=None, fmt="txt", callname=None):
self.sample_rate = sample_rate
self.timeout = timeout
if isinstance(fmt, basestring):
fmt = [fmt]
self.fmt = fmt
self.save_fig=True
self.callname = callname
def __call__(self, func):
def wrapper(*args, **kwargs):
pid = os.getpid()
profiler = Profiler(pid, self.sample_rate)
baseline = _measure(profiler.process)
try:
profiler.start()
t0 = datetime.now()
res = func(*args, **kwargs)
profiler.cancel()
ex = profiler.exception(self.timeout)
if ex is not None:
raise ex
prof_res = profiler.result(self.timeout)
if self.callname is None:
self.callname = func.func_name
self.fig = output(prof_res, self.fmt, self.callname, t0, baseline, self.save_fig)
return res
finally:
profiler.cancel()
return wrapper