Skip to content
Prev Previous commit
Next Next commit
Add async benchmark modes
  • Loading branch information
pablogsal committed May 12, 2026
commit f7fe3beda1d911455fe36ec50901a7e745219548
100 changes: 91 additions & 9 deletions Tools/inspection/benchmark_external_inspection.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,45 @@ def create_threads(n):
time.sleep(0.05)
'''

ASYNC_CODE = '''\
import asyncio
import contextlib
import math

def compute_slice(seed):
result = 0.0
for i in range(2000):
result += math.sin(seed + i) * math.sqrt(i + 1)
return result

async def leaf_task(seed):
total = 0.0
while True:
total += compute_slice(seed)
await asyncio.sleep(0)

async def parent_task(seed):
child = asyncio.create_task(leaf_task(seed + 1000), name=f"leaf-{seed}")
try:
while True:
compute_slice(seed)
await asyncio.sleep(0.001)
finally:
child.cancel()
with contextlib.suppress(asyncio.CancelledError):
await child

async def main():
tasks = [
asyncio.create_task(parent_task(i), name=f"parent-{i}")
for i in range(8)
]
await asyncio.gather(*tasks)

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

CODE_EXAMPLES = {
"basic": {
"code": CODE,
Expand All @@ -164,22 +203,44 @@ def create_threads(n):
"code": CODE_WITH_TONS_OF_THREADS,
"description": "Tons of threads doing mixed CPU/IO work",
},
"asyncio": {
"code": ASYNC_CODE,
"description": "Asyncio tasks with active and awaited coroutine chains",
},
}

OPERATIONS = {
"stack_trace": {
"method": "get_stack_trace",
"label": "get_stack_trace()",
},
"async_stack_trace": {
"method": "get_async_stack_trace",
"label": "get_async_stack_trace()",
},
"all_awaited_by": {
"method": "get_all_awaited_by",
"label": "get_all_awaited_by()",
},
}


def benchmark(unwinder, duration_seconds=10, blocking=False):
def benchmark(unwinder, duration_seconds=10, blocking=False, operation="stack_trace"):
"""Benchmark mode - measure raw sampling speed for specified duration"""
sample_count = 0
fail_count = 0
total_work_time = 0.0
start_time = time.perf_counter()
end_time = start_time + duration_seconds
total_attempts = 0
operation_info = OPERATIONS[operation]
operation_method = getattr(unwinder, operation_info["method"])

colors = get_colors(can_colorize())

print(
f"{colors.BOLD_BLUE}Benchmarking sampling speed for {duration_seconds} seconds...{colors.RESET}"
f"{colors.BOLD_BLUE}Benchmarking {operation_info['label']} speed "
f"for {duration_seconds} seconds...{colors.RESET}"
)

try:
Expand All @@ -190,8 +251,8 @@ def benchmark(unwinder, duration_seconds=10, blocking=False):
if blocking:
unwinder.pause_threads()
try:
stack_trace = unwinder.get_stack_trace()
if stack_trace:
sample = operation_method()
if sample:
sample_count += 1
finally:
if blocking:
Expand Down Expand Up @@ -239,6 +300,7 @@ def benchmark(unwinder, duration_seconds=10, blocking=False):
(sample_count / total_attempts) * 100 if total_attempts > 0 else 0
),
"total_work_time": total_work_time,
"operation": operation_info["label"],
"avg_work_time_us": (
(total_work_time / total_attempts) * 1e6 if total_attempts > 0 else 0
),
Expand All @@ -252,7 +314,7 @@ def print_benchmark_results(results):
colors = get_colors(can_colorize())

print(f"\n{colors.BOLD_GREEN}{'='*60}{colors.RESET}")
print(f"{colors.BOLD_GREEN}get_stack_trace() Benchmark Results{colors.RESET}")
print(f"{colors.BOLD_GREEN}{results['operation']} Benchmark Results{colors.RESET}")
print(f"{colors.BOLD_GREEN}{'='*60}{colors.RESET}")

# Basic statistics
Expand Down Expand Up @@ -329,6 +391,8 @@ def parse_arguments():
%(prog)s -d 60 # Run basic benchmark for 60 seconds
%(prog)s --code deep_static # Run deep static call stack benchmark
%(prog)s --code deep_static -d 30 # Run deep static benchmark for 30 seconds
%(prog)s --operation async_stack_trace
%(prog)s --operation all_awaited_by

Available code examples:
{examples_desc}
Expand All @@ -348,8 +412,15 @@ def parse_arguments():
"--code",
"-c",
choices=list(CODE_EXAMPLES.keys()),
default="basic",
help="Code example to benchmark (default: basic)",
default=None,
help="Code example to benchmark (default: basic, or asyncio for async operations)",
)

parser.add_argument(
"--operation",
choices=list(OPERATIONS.keys()),
default="stack_trace",
help="Remote unwinder operation to benchmark (default: stack_trace)",
)

parser.add_argument(
Expand All @@ -365,7 +436,10 @@ def parse_arguments():
help="Stop all threads before sampling for consistent snapshots",
)

return parser.parse_args()
args = parser.parse_args()
if args.code is None:
args.code = "asyncio" if args.operation != "stack_trace" else "basic"
return args


def create_target_process(temp_file, code_example="basic"):
Expand Down Expand Up @@ -420,6 +494,9 @@ def main():
print(
f"{colors.CYAN}Benchmark Duration:{colors.RESET} {colors.YELLOW}{args.duration}{colors.RESET} seconds"
)
print(
f"{colors.CYAN}Operation:{colors.RESET} {colors.GREEN}{OPERATIONS[args.operation]['label']}{colors.RESET}"
)
print(
f"{colors.CYAN}Blocking Mode:{colors.RESET} {colors.GREEN if args.blocking else colors.YELLOW}{'enabled' if args.blocking else 'disabled'}{colors.RESET}"
)
Expand Down Expand Up @@ -451,7 +528,12 @@ def main():
unwinder = _remote_debugging.RemoteUnwinder(
process.pid, cache_frames=True, **kwargs
)
results = benchmark(unwinder, duration_seconds=args.duration, blocking=args.blocking)
results = benchmark(
unwinder,
duration_seconds=args.duration,
blocking=args.blocking,
operation=args.operation,
)
finally:
cleanup_process(process, temp_file_path)

Expand Down