A timing context for asyncio

asyncio_timing.py

# See the coresponding stackoverflow post:
# https://stackoverflow.com/a/34827291/2846140

import time
import asyncio
import selectors
import contextlib


class TimedSelector(selectors.DefaultSelector):

    select_time = 0.

    def reset_select_time(self):
        self.select_time = 0.

    def select(self, timeout=None):
        if timeout <= 0:
            return super().select(timeout)
        start = time.time()
        try:
            return super().select(timeout)
        finally:
            self.select_time += time.time() - start


class TimedEventLoopPolicy(asyncio.DefaultEventLoopPolicy):

    def new_event_loop(self):
        selector = TimedSelector()
        return asyncio.DefaultEventLoopPolicy._loop_factory(selector)


@contextlib.contextmanager
def print_timing():
    asyncio.get_event_loop()._selector.reset_select_time()
    real_time = time.time()
    process_time = time.process_time()
    yield
    real_time = time.time() - real_time
    cpu_time = time.process_time() - process_time
    select_time = asyncio.get_event_loop()._selector.select_time
    other_io_time = max(0., real_time - cpu_time - select_time)
    print(f"CPU time:      {cpu_time:.3f} s")
    print(f"Select time:   {select_time:.3f} s")
    print(f"Other IO time: {other_io_time:.3f} s")
    print(f"Real time:     {real_time:.3f} s")


# Testing

async def main():
    print("~ Correct IO management ~")
    with print_timing():
        await asyncio.sleep(1)
        sum(range(10**6))
    print()

    print("~ Incorrect IO management ~")
    with print_timing():
        time.sleep(0.2)
        await asyncio.sleep(0.8)
        sum(range(10**6))
    print()


if __name__ == "__main__":
    asyncio.set_event_loop_policy(TimedEventLoopPolicy())
    asyncio.run(main())

@vxgmichel Any idea how to track "select_time" per task?

I'd like to implement task_perf_counter() for asyncio as I was able to for Trio.
https://gist.github.com/belm0/4fdb51ce04c5c26f01d58913ae5c43da

@belm0 See the end of this revision, it is a task-based approach that I removed from this answer because it broke with python 3.7.

I'd be extremely interested in something similar for trio though. I'll start with your gist, but do you have more resources about detecting incrorrect IO management on trio?

@vxgmichel Thank you, I'll see if I can find some approach along those lines that works with the current API. Otherwise, once I've published some interesting tools showing what can be done with Trio, we can bring it up to the asyncio lead.

Just to be clear, Trio is a different async framework than asyncio, so if you application is already on the latter it would not be a trivial change.

By incorrect IO management, you mean freezing the event loop by making blocking calls from async tasks? There is an idea for a blocked-task watchdog (python-trio/trio#591), or else Trio has an [Instrument API (https://trio.readthedocs.io/en/latest/reference-hazmat.html#instrumentation) which can be use to target tasks with the largest step time, etc.

@belm0 Interesting, thanks!

Just to be clear, Trio is a different async framework than asyncio, so if you application is already on the latter it would not be a trivial change.

Yes I'm working on a trio app at the moment :)

I ended writing an instrument for my own purposes: https://gist.github.com/vxgmichel/2317738886c75e9626d92c326196789f

It would be great if you could add a license for this

@cganterh It's all MIT, go for it 😃

pygimp gimpfu

Hi @vxgmichel , I add a decorator version and it passes the test:
https://github.com/gongyisheng/playground/blob/main/python/asyncio/timing.py
It provides another option than contextlib.contextmanager