path: root/aergia/aergia.py

#!/usr/bin/env python3
'''
         _/_/                                  _/
      _/    _/    _/_/    _/  _/_/    _/_/_/        _/_/_/
     _/_/_/_/  _/_/_/_/  _/_/      _/    _/  _/  _/    _/
    _/    _/  _/        _/        _/    _/  _/  _/    _/
   _/    _/    _/_/_/  _/          _/_/_/  _/    _/_/_/
                                      _/
                                 _/_/

Copyright 2025 bdunahu

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.


Commentary:

   Aergia is a sampling based profiler based off of SCALENE
   by Emery Berger and the UMASS Plasma Lab
   (https://github.com/plasma-umass/scalene).

   It is not particularly informative, but unlike SCALENE
   or other sampling-based profilers I could find, reports
   the wall-time each asyncio await call spends idling.
   (yappi can profile asyncio, but only report time spent in
   each function. Instrumentation-profilers cannot do this
   without adding very large overhead).

   The goal behind Aergia is to eventually have these features,
   or similar, merged into SCALENE.


Code:
'''

from collections import defaultdict, namedtuple
from typing import Optional
import argparse
import asyncio
import signal
import sys
import threading
import time
import traceback
import gc

from types import AsyncGeneratorType

orig_thread_join = threading.Thread.join


def thread_join_replacement(
    self: threading.Thread, timeout: Optional[float] = None
) -> None:
    '''
    We replace threading.Thread.join with this method which always
    periodically yields.
    '''

    start_time = time.perf_counter()
    interval = sys.getswitchinterval()
    while self.is_alive():
        orig_thread_join(self, interval)
        # If a timeout was specified, check to see if it's expired.
        if timeout is not None:
            end_time = time.perf_counter()
            if end_time - start_time >= timeout:
                return None
    return None


threading.Thread.join = thread_join_replacement

# a tuple used as a key in the sample-dict
Sample = namedtuple('Sample', ['file', 'line', 'func'])


class Aergia(object):

    # a key-value pair where keys represent frame metadata (see
    # Aergia.frame_to_string) and values represent number of times
    # sampled.
    samples = defaultdict(lambda: 0)
    # number of times samples have been collected
    total_samples = 0
    # the (ideal) interval between samples
    signal_interval = 0.0

    # the current task for the loop being processed
    current_task = None

    # if we should profile currently running tasks
    do_profile_current = False

    @staticmethod
    def __init__(signal_interval, do_profile_current):
        Aergia.signal_interval = signal_interval
        Aergia.do_profile_current = do_profile_current

    @staticmethod
    def start():
        '''Turns on asyncio debug mode and sets up our signals.

        Debug mode must be on by default to avoid losing samples.
        Debug mode is required to view the current coroutine being waited on
        in `Aergia._get_idle_task_frames'. The TimerHandler object otherwise
        does not keep track of a _source_traceback.
        '''
        signal.signal(signal.SIGALRM,
                      Aergia._idle_signal_handler)
        signal.setitimer(signal.ITIMER_REAL,
                         Aergia.signal_interval,
                         Aergia.signal_interval)

    @staticmethod
    def stop():
        '''Stops the profiler.'''
        signal.setitimer(signal.ITIMER_REAL, 0)

    @staticmethod
    def clear():
        Aergia.total_samples = 0
        Aergia.samples = defaultdict(lambda: 0)

    @staticmethod
    def get_samples():
        '''Returns the profiling results.'''
        return Aergia.samples

    @staticmethod
    def print_samples():
        '''Pretty-print profiling results.'''
        if Aergia.total_samples > 0:
            print(f"{'FILE':<30} {'FUNC':<30}"
                  f" {'PERC':<8} {'(ACTUAL -> SEC)':<10}")
            for key in Aergia._sort_samples(Aergia.samples):
                Aergia.print_sample(key)
        else:
            print("No samples were gathered. If you *are* using concurrency, "
                  "this is likely a bug and you may run Aergia again.")

    @staticmethod
    def print_sample(key):
        '''Pretty-print a single sample.'''
        sig_intv = Aergia.signal_interval
        value = Aergia.samples[key]
        print(f"{Aergia._tuple_to_string(key)} {value * 100 / Aergia.total_samples:.3f}% "
              f" ({value:.3f} -> {value*sig_intv:.6f})")

    @staticmethod
    def _idle_signal_handler(sig, frame):
        '''Obtains and records which lines are currently being waited on.'''
        keys = Aergia._compute_frames_to_record()
        for key in keys:
            Aergia.samples[Aergia._frame_to_tuple(key)] += 1
            Aergia.total_samples += 1

    @staticmethod
    def _compute_frames_to_record():
        '''Collects all stack frames which are currently being awaited on
        during a given timestamp, as well as those which are currently
        executing.

        Note that we do NOT need to walk back up the call-stack to find
        which of the user's lines caused the await call. There is NEVER
        a previous frame, because idle frames aren't on the call stack!

        Luckily, the event loop and asyncio.all_tasks keeps track of
        what is running for us.'''
        loops = Aergia._get_event_loops()
        # idle tasks
        frames = Aergia._get_frames_from_loops(loops)
        # current running frames
        if Aergia.do_profile_current:
            frames += Aergia._get_frames_from_threads()
        return frames

    @staticmethod
    def _get_event_loops():
        '''Returns each thread's event loop, if it exists.'''
        loops = []
        for t in threading.enumerate():
            frame = sys._current_frames().get(t.ident)
            if frame:
                loop = Aergia._walk_back_until_loop(frame)
                if loop and loop not in loops:
                    loops.append(loop)
        return loops

    @staticmethod
    def _walk_back_until_loop(frame):
        '''Walks back the callstack until we are in a method named '_run_once'.
        If this is ever true, we assume we are in an Asyncio event loop method,
        and check to see if the 'self' variable is indeed and instance of
        AbstractEventLoop. Return this variable if true.'''
        while frame:
            if frame.f_code.co_name == '_run_once' and 'self' in frame.f_locals:
                loop = frame.f_locals['self']
                if isinstance(loop, asyncio.AbstractEventLoop):
                    return loop
            else:
                frame = frame.f_back
        return None

    @staticmethod
    def _get_frames_from_loops(loops):
        '''Given LOOPS, returns a flat list of frames.'''
        return [
            frames for loop in loops
            for frames in Aergia._get_idle_task_frames(loop)
        ]

    @staticmethod
    def _get_frames_from_threads():
        frames = [sys._current_frames().get(t.ident, None)
                  for t in threading.enumerate()]
        # process frames to remove those we do not track
        new_frames = []
        for f in frames:
            if f is None:
                continue
            fname = f.f_code.co_filename
            while not Aergia._should_trace(fname):
                # walk the stack backwards until we hit a frame that is one
                # we should trace.
                if f:
                    f = f.f_back
                else:
                    break
                if f:
                    fname = f.f_code.co_filename
            if f:
                new_frames.append(f)
        return new_frames

    @staticmethod
    def _frame_to_tuple(frame):
        '''Given a frame, constructs a sample key for tallying lines.'''
        co = frame.f_code
        func_name = co.co_name
        line_no = frame.f_lineno
        filename = co.co_filename
        return Sample(filename, line_no, func_name)

    @staticmethod
    def _tuple_to_string(sample):
        '''Given a namedtuple corresponding to a sample key,
        pretty-prints a frame as a function/file name and a line number.'''
        filename = \
            (sample.file if len(sample.file) <= 25 else sample.file[-25:])
        return f"{filename}:{sample.line}".ljust(30) + f"{sample.func:30}"

    @staticmethod
    def _sort_samples(sample_dict):
        '''Returns SAMPLE_DICT in descending order by number of samples.'''
        return {k: v for k, v in sorted(sample_dict.items(),
                                        key=lambda item: item[1],
                                        reverse=True)}

    @staticmethod
    def _get_idle_task_frames(loop):
        '''Given an asyncio event loop, returns the list of idle task frames.
        We only care about idle task frames, as running tasks are already
        included elsewhere.'''
        idle = []

        # set this when we start processing a loop.
        # it is required later, but I only want to set it once.
        Aergia.current_task = asyncio.current_task(loop)

        for task in asyncio.all_tasks(loop):
            if not Aergia._should_trace_task(task):
                continue

            coro = task.get_coro()

            frame = Aergia._get_deepest_traceable_frame(coro)
            if frame:
                idle.append(frame)

        return idle

    @staticmethod
    def _get_deepest_traceable_frame(coro):
        '''Get the deepest frame of coro we care to trace.
        This is possible because each corooutine keeps a reference to the
        coroutine it is waiting on.

        Note that it cannot be the case that a task is suspended in a frame
        that does not belong to a coroutine, asyncio is very particular about
        that! This is also why we only track idle tasks this way.'''
        curr = coro
        deepest_frame = None
        while curr:
            frame = getattr(curr, 'cr_frame', None)

            if not frame:
                curr = Aergia._search_future(curr)
                if isinstance(curr, AsyncGeneratorType):
                    frame = getattr(curr, 'ag_frame', None)
                else:
                    break

            if Aergia._should_trace(frame.f_code.co_filename):
                deepest_frame = frame

            if isinstance(curr, AsyncGeneratorType):
                curr = getattr(curr, 'ag_await', None)
            else:
                curr = getattr(curr, 'cr_await', None)

        # if this task is found to point to another task we're profiling,
        # then we will get the deepest frame later and should return nothing.
        if isinstance(curr, list) and any(
                Aergia._should_trace_task(task)
                for task in curr
        ):
            return None

        return deepest_frame

    @staticmethod
    def _search_future(future):
        '''Given an awaitable which is not a coroutine, assume it is a future
        and attempt to find references to tasks or async generators.'''
        awaitable = None
        if not isinstance(future, asyncio.Future):
            # TODO some wrappers like _asyncio.FutureIter,
            # async_generator_asend get caught here, I am not sure if a more
            # robust approach is necessary

            # can gc be avoided here?
            refs = gc.get_referents(future)
            if refs:
                awaitable = refs[0]

        # this is specific to gathering futures, i.e., gather statement.
        # Other cases may need to be added.
        if isinstance(awaitable, asyncio.Future):
            return getattr(awaitable, '_children', [])

        # if this is not AsyncGeneratorType, it is ignored
        return awaitable

    @staticmethod
    def _should_trace_task(task):
        '''Returns FALSE if TASK is uninteresting to the user.

        A task is interesting if it is not CURRENT_TASK, if it has actually
        started executing, and if a child task did not originate from it.
        '''
        if not isinstance(task, asyncio.Task):
            return False

        # the task is not idle
        if task == Aergia.current_task:
            return False

        coro = task.get_coro()

        # the task hasn't even run yet
        # assumes that all started tasks are sitting at an await
        # statement.
        # if this isn't the case, the associated coroutine will
        # be 'waiting' on the coroutine declaration. No! Bad!
        if getattr(coro, 'cr_frame', None) is None or \
           getattr(coro, 'cr_await', None) is None:
            return False

        frame = getattr(coro, 'cr_frame', None)

        return Aergia._should_trace(frame.f_code.co_filename)

    @staticmethod
    def _should_trace(filename):
        '''Returns FALSE if filename is uninteresting to the user.
        Don't depend on this. It's good enough for testing.'''
        # FIXME Assume GuixSD. Makes filtering easy
        if not filename:
            return False
        if '/gnu/store' in filename:
            return False
        if 'site-packages' in filename:
            return False
        if 'propcache' in filename:
            return False
        if '.pyx' in filename:
            return False
        if filename[0] == '<':
            return False
        if 'aergia.py' in filename:
            return False
        return True

    @staticmethod
    def _gettime():
        '''returns the wallclock time'''
        return time.process_time()


the_globals = {
    '__name__': '__main__',
    '__doc__': None,
    '__package__': None,
    '__loader__': globals()['__loader__'],
    '__spec__': None,
    '__annotations__': {},
    '__builtins__': globals()['__builtins__'],
    '__file__': None,
    '__cached__': None,
}


if __name__ == "__main__":
    # parses CLI arguments and facilitates profiler runtime.
    parser = argparse.ArgumentParser(
        usage='%(prog)s [args] script [args]'
    )

    parser.add_argument('-i', '--interval',
                        help='The minimum amount of time inbetween \
                        samples in seconds.',
                        metavar='',
                        type=float,
                        default=0.01)
    parser.add_argument('-a', '--async-only',
                        help='Do not profile currently running tasks.',
                        action='store_true')
    parser.add_argument('script', help='A python script to run.')
    parser.add_argument('s_args', nargs=argparse.REMAINDER,
                        help='python script args')
    args = parser.parse_args()

    sys.argv = [args.script] + args.s_args
    try:
        with open(args.script, 'r', encoding='utf-8') as fp:
            code = compile(fp.read(), args.script, "exec")
            Aergia(args.interval, not args.async_only).start()
            exec(code, the_globals)
            Aergia.print_samples()
            Aergia.stop()
    except Exception:
        traceback.print_exc()