#!/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 os import signal import sys import threading import time import traceback 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 # 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. A task is considered 'idle' if it is pending and not the current task.''' idle = [] coros = Aergia._get_traceable_coros(loop) for coro in coros: frame = Aergia._get_deepest_traceable_frame(coro) # if it's a future, if frame: idle.append(frame) # handle async generators # ideally, we would access these from _get_deepest_traceable_frame # doing it this way causes us to assign generator time to the # coroutine that called this generator in _get_deepest_traceable_frame for ag in loop._asyncgens: f = getattr(ag, 'ag_frame', None) if f and Aergia._should_trace(f.f_code.co_filename): idle.append(f) return idle @staticmethod def _get_traceable_coros(loop): '''Given event loop LOOP, returns a list of coroutines corresponding to each traceable task. A task is traceable if it is not the current task, if it has actually started executing, and if a child task did not originate from it.''' current = asyncio.current_task(loop) coros = dict() p_tasks = set() for task in asyncio.all_tasks(loop): # the task is not idle if task == current: continue 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: continue frame = getattr(coro, 'cr_frame', None) # for methods like gather, tasks create new tasks. keep # track of these tasks so we don't count them twice. parent = Aergia._get_parent_task(task) if parent is not None: p_tasks.add(parent) coros[coro] = (frame.f_code.co_filename, frame.f_lineno, frame.f_code.co_name) return [ coro for coro, frame_summary in coros.items() if frame_summary not in p_tasks ] @staticmethod def _get_parent_task(task): '''Given TASK, returns the first frame in the source traceback which is tracable. This only works if the event loop has debug mode turned on. If this is not the case, returns None.''' tb = task._source_traceback # true when debug mode isn't turned on if tb is None: return None for frame in tb[::-1]: if Aergia._should_trace(frame.filename): return (frame.filename, frame.lineno, frame.name) return None @staticmethod def _get_deepest_traceable_frame(coro): '''Get the deepest frame of coro we care to trace. Luckily for us, coroutines keep track of the coroutine object which they are waiting on (cr_await). We simply need to trace that down as far as possible to get the actual line we are 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! If cr_await is None, then the task was created but never started. (since once a task is started, it either starts waiting on another coroutine, or finishes immediately. Note when you call task_create, asyncio schedules it for the next event loop iteration, so the code does not actually run immediately). These cases were already filtered out in `_get_idle_task_frames'.''' curr = coro deepest_frame = None while curr: frame = getattr(curr, 'cr_frame', None) if frame and Aergia._should_trace(frame.f_code.co_filename): deepest_frame = frame curr = getattr(curr, 'cr_await', None) return deepest_frame @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 os.environ["PYTHONASYNCIODEBUG"] = "1" 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()