How to profile/benchmark Python ASYNCIO code?

Question

How do I profile/benchmark an assynchronous Python script (which uses ASYNCIO)?

I you would usualy do

totalMem = tracemalloc.get_traced_memory()[0]
totalTime = time.time()

retValue = myFunction()

totalTime = time.time() - totalTime 
totalMem = tracemalloc.get_traced_memory()[0] - totalMem 

This way I would save the total time taken by the function. I learned how to use decorators and I did just that - and dumped all stats into a text file for later analysis.

But, when you have ASYNCIO script, things get pretty different: the function will block while doing an "await aiohttpSession.get()", and control will go back to the event loop, which will run other functions.

This way, the elapsed time and changes in total allocated memory won't reveal anything, because I will have measured more than just that function.

The only way it would work would be something like

class MyTracer:
  def __init__(self):
    self.totalTime = 0
    self.totalMem = 0
    self.startTime = time.time()
    self.startMem = tracemalloc.get_traced_memory()[0]
  def stop(self):
    self.totalTime += time.time() - self.startTime
    self.totalMem += tracemalloc.get_traced_memory()[0] - self.startMem
  def start(self):
    self.startTime = time.time()
    self.startMem = tracemalloc.get_traced_memory()[0]

And now, somehow, insert it in the code:

def myFunction():

    tracer = MyTracer()

    session = aiohttp.ClientSession()

    # do something

    tracer.stop()
    # the time elapsed here, and the changes in the memory allocation, are not from the current function
    retValue = await(await session.get('https://hoochie-mama.org/cosmo-kramer',
        headers={
            'User-Agent': 'YoYo Mama! v3.0',
            'Cookies': 'those cookies are making me thirsty!',
            })).text()
    tracer.start()

    # do more things

    tracer.stop()

    # now "tracer" has the info about total time spent in this function, and the memory allocated by it
    # (the memory stats could be negative if the function releases more than allocates)

Is there a way to accomplish this, I mean, profile all my asyncio code without having to insert all this code? Or is there a module already capable of doing just that?

Answer 1

Check out Yappi profiler which has support for coroutine profiling. Their page on coroutine profiling describes the problem you're facing very clearly:

The main issue with coroutines is that, under the hood when a coroutine yields or in other words context switches, Yappi receives a return event just like we exit from the function. That means the time spent while the coroutine is in yield state does not get accumulated to the output. This is a problem especially for wall time as in wall time you want to see whole time spent in that function or coroutine. Another problem is call count. You see every time a coroutine yields, call count gets incremented since it is a regular function exit.

They also describe very high level how Yappi solves this problem:

With v1.2, Yappi corrects above issues with coroutine profiling. Under the hood, it differentiates the yield from real function exit and if wall time is selected as the clock_type it will accumulate the time and corrects the call count metric.