GH-141565: Add async code awareness to Tachyon (#141533)

Co-authored-by: Pablo Galindo Salgado <pablogsal@gmail.com>
This commit is contained in:
Savannah Ostrowski 2025-12-06 11:31:40 -08:00 committed by GitHub
parent 35142b18ae
commit 56a442d0d8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1355 additions and 83 deletions

View file

@ -195,6 +195,11 @@ def _add_sampling_options(parser):
dest="gc",
help='Don\'t include artificial "<GC>" frames to denote active garbage collection',
)
sampling_group.add_argument(
"--async-aware",
action="store_true",
help="Enable async-aware profiling (uses task-based stack reconstruction)",
)
def _add_mode_options(parser):
@ -205,7 +210,14 @@ def _add_mode_options(parser):
choices=["wall", "cpu", "gil"],
default="wall",
help="Sampling mode: wall (all samples), cpu (only samples when thread is on CPU), "
"gil (only samples when thread holds the GIL)",
"gil (only samples when thread holds the GIL). Incompatible with --async-aware",
)
mode_group.add_argument(
"--async-mode",
choices=["running", "all"],
default="running",
help='Async profiling mode: "running" (only running task) '
'or "all" (all tasks including waiting). Requires --async-aware',
)
@ -382,6 +394,27 @@ def _validate_args(args, parser):
"Live mode requires the curses module, which is not available."
)
# Async-aware mode is incompatible with --native, --no-gc, --mode, and --all-threads
if args.async_aware:
issues = []
if args.native:
issues.append("--native")
if not args.gc:
issues.append("--no-gc")
if hasattr(args, 'mode') and args.mode != "wall":
issues.append(f"--mode={args.mode}")
if hasattr(args, 'all_threads') and args.all_threads:
issues.append("--all-threads")
if issues:
parser.error(
f"Options {', '.join(issues)} are incompatible with --async-aware. "
"Async-aware profiling uses task-based stack reconstruction."
)
# --async-mode requires --async-aware
if hasattr(args, 'async_mode') and args.async_mode != "running" and not args.async_aware:
parser.error("--async-mode requires --async-aware to be enabled.")
# Live mode is incompatible with format options
if hasattr(args, 'live') and args.live:
if args.format != "pstats":
@ -570,6 +603,7 @@ def _handle_attach(args):
all_threads=args.all_threads,
realtime_stats=args.realtime_stats,
mode=mode,
async_aware=args.async_mode if args.async_aware else None,
native=args.native,
gc=args.gc,
)
@ -618,6 +652,7 @@ def _handle_run(args):
all_threads=args.all_threads,
realtime_stats=args.realtime_stats,
mode=mode,
async_aware=args.async_mode if args.async_aware else None,
native=args.native,
gc=args.gc,
)
@ -650,6 +685,7 @@ def _handle_live_attach(args, pid):
limit=20, # Default limit
pid=pid,
mode=mode,
async_aware=args.async_mode if args.async_aware else None,
)
# Sample in live mode
@ -660,6 +696,7 @@ def _handle_live_attach(args, pid):
all_threads=args.all_threads,
realtime_stats=args.realtime_stats,
mode=mode,
async_aware=args.async_mode if args.async_aware else None,
native=args.native,
gc=args.gc,
)
@ -689,6 +726,7 @@ def _handle_live_run(args):
limit=20, # Default limit
pid=process.pid,
mode=mode,
async_aware=args.async_mode if args.async_aware else None,
)
# Profile the subprocess in live mode
@ -700,6 +738,7 @@ def _handle_live_run(args):
all_threads=args.all_threads,
realtime_stats=args.realtime_stats,
mode=mode,
async_aware=args.async_mode if args.async_aware else None,
native=args.native,
gc=args.gc,
)