profiling_decorator.py

"""
For tracemalloc see: https://github.com/wyplay/pytracemalloc

Quick apply guide (ubuntu)::

    apt-get source python2.7
    cd python2.7-2.7.3
    patch -p1 < ../pytracemalloc/python2.7_track_free_list.patch
    debuild -us -uc
    cd ..
    sudo dpkg -i python2.7-minimal_2.7.3-0ubuntu3.1_amd64.deb


"""
def logged_task(func=None, profile=False, **options):
    if profile:
        from cProfile import Profile
        from pstats import Stats
        try:
            import pyprof2calltree
        except ImportError:
            pyprof2calltree = None

    def decorator(func):
        @wraps(func)
        def logged_task_wrapped(self, *args, **kwargs):
            logger.info("Starting task %s with context:%s", func.__name__, self.request)
            stime = time.time()
            try:
                if profile == 'memory':
                    import tracemalloc
                    tracemalloc.enable()
                    top = tracemalloc.DisplayTop(5000, file=open('/tmp/task-profile-%s-%s.memory' % (func.__name__, self.request.id), "w"))
                    top.show_lineno = True
                    try:
                        return func(*args, **kwargs)
                    finally:
                        top.display()
                    #    snapshot = tracemalloc.Snapshot.create()
                    #    snapshot.write( ... )
                if profile is True or profile == 'both':
                    try:
                        p = Profile()
                        return p.runcall(func, *args, **kwargs)
                    finally:
                        with file('/tmp/task-profile-%s-%s' % (func.__name__, self.request.id), 'w') as stream:
                            stats = Stats(p, stream=stream)
                            if pyprof2calltree:
                                with file('/tmp/task-profile-%s-%s.grind' % (func.__name__, self.request.id), 'wb') as grind:
                                    conv = pyprof2calltree.CalltreeConverter(stats)
                                    conv.output(grind)
                            stats.sort_stats('cumulative', 'calls')
                            stats.print_stats()
                else:
                    return func(*args, **kwargs)


            except Exception:
                logger.exception(
                    "Unexpected exception in task %s with id: %s, run time: %.2fsec.",
                    func.__name__, self.request.id, time.time() - stime
                )
                raise
            finally:
                logger.info(
                    "Completed task %s with id: %s, run time: %.2fsec.",
                    func.__name__, self.request.id, time.time() - stime
                )
        return task(func, run=logged_task_wrapped, **options)
    if func is None:
        return decorator
    else:
        return decorator(func)