From 901eb883a7b8fa0477f49eb60fecdea77ae4abfc Mon Sep 17 00:00:00 2001 From: youknowdot Date: Sun, 25 Oct 2026 03:16:32 -0700 Subject: [PATCH] Add canonical sampling profiler CLI for offline replay (#9358) Import the CLI, entry point, error type, coordinator dependency, and CLI tests with scripts/update_lib migrate from CPython 3.15.0 at d220409a13b77de7a6c30abe16ea90db01528edc. This frontend follows the existing offline collector/codec support. Live attach, run, and dump remain unsupported and unvalidated. Assisted-by: OpenAI Codex: model version unavailable --- Lib/profiling/sampling/__main__.py | 73 + Lib/profiling/sampling/_sync_coordinator.py | 256 +++ Lib/profiling/sampling/cli.py | 1462 +++++++++++++++++ Lib/profiling/sampling/errors.py | 19 + .../test_sampling_profiler/test_cli.py | 1083 ++++++++++++ 5 files changed, 2893 insertions(+) create mode 100644 Lib/profiling/sampling/__main__.py create mode 100644 Lib/profiling/sampling/_sync_coordinator.py create mode 100644 Lib/profiling/sampling/cli.py create mode 100644 Lib/profiling/sampling/errors.py create mode 100644 Lib/test/test_profiling/test_sampling_profiler/test_cli.py diff --git a/Lib/profiling/sampling/__main__.py b/Lib/profiling/sampling/__main__.py new file mode 100644 index 000000000..a45b645ea --- /dev/null +++ b/Lib/profiling/sampling/__main__.py @@ -0,0 +1,73 @@ +"""Run the sampling profiler from the command line.""" + +import sys + +MACOS_PERMISSION_ERROR = """\ +๐Ÿ”’ Permission Error: Unable to access process memory on macOS + +Tachyon needs elevated permissions to profile processes. Try one of these solutions: + +1. Try running again with elevated permissions by running 'sudo -E !!'. + +2. If targeting system Python processes: + Note: Apple's System Integrity Protection (SIP) may block access to system + Python binaries. Consider using a user-installed Python instead. +""" + +LINUX_PERMISSION_ERROR = """ +๐Ÿ”’ Tachyon was unable to access process memory. This could be because tachyon +has insufficient privileges (the required capability is CAP_SYS_PTRACE). +Unprivileged processes cannot trace processes that they cannot send signals +to or those running set-user-ID/set-group-ID programs, for security reasons. + +If your uid matches the uid of the target process you want to analyze, you +can do one of the following to get 'ptrace' scope permissions: + +* If you are running inside a Docker container, you need to make sure you + start the container using the '--cap-add=SYS_PTRACE' or '--privileged' + command line arguments. Notice that this may not be enough if you are not + running as 'root' inside the Docker container as you may need to disable + hardening (see next points). + +* Try running again with elevated permissions by running 'sudo -E !!'. + +* You can disable kernel hardening for the current session temporarily (until + a reboot happens) by running 'echo 0 | sudo tee /proc/sys/kernel/yama/ptrace_scope'. +""" + +WINDOWS_PERMISSION_ERROR = """ +๐Ÿ”’ Tachyon requires administrator rights to access process memory on Windows. +Please run your command prompt as Administrator and try again. +""" + +GENERIC_PERMISSION_ERROR = """ +๐Ÿ”’ Tachyon was unable to access the target process due to operating +system restrictions or missing privileges. +""" + +from .cli import main +from .errors import SamplingUnknownProcessError, SamplingModuleNotFoundError, SamplingScriptNotFoundError + +def handle_permission_error(): + """Handle PermissionError by displaying appropriate error message.""" + if sys.platform == "darwin": + print(MACOS_PERMISSION_ERROR, file=sys.stderr) + elif sys.platform.startswith("linux"): + print(LINUX_PERMISSION_ERROR, file=sys.stderr) + elif sys.platform.startswith("win"): + print(WINDOWS_PERMISSION_ERROR, file=sys.stderr) + else: + print(GENERIC_PERMISSION_ERROR, file=sys.stderr) + sys.exit(1) + +if __name__ == '__main__': + try: + main() + except PermissionError: + handle_permission_error() + except SamplingUnknownProcessError as err: + print(f"Tachyon cannot find the process: {err}", file=sys.stderr) + sys.exit(1) + except (SamplingModuleNotFoundError, SamplingScriptNotFoundError) as err: + print(f"Tachyon cannot find the target: {err}", file=sys.stderr) + sys.exit(1) diff --git a/Lib/profiling/sampling/_sync_coordinator.py b/Lib/profiling/sampling/_sync_coordinator.py new file mode 100644 index 000000000..a1cce314b --- /dev/null +++ b/Lib/profiling/sampling/_sync_coordinator.py @@ -0,0 +1,256 @@ +""" +Internal synchronization coordinator for the sample profiler. + +This module is used internally by the sample profiler to coordinate +the startup of target processes. It should not be called directly by users. +""" + +import importlib.util +import os +import sys +import socket +import runpy +import time +import types +from typing import List, NoReturn + + +class CoordinatorError(Exception): + """Base exception for coordinator errors.""" + pass + + +class ArgumentError(CoordinatorError): + """Raised when invalid arguments are provided.""" + pass + + +class SyncError(CoordinatorError): + """Raised when synchronization with profiler fails.""" + pass + + +class TargetError(CoordinatorError): + """Raised when target execution fails.""" + pass + + +def _validate_arguments(args: List[str]) -> tuple[int, str, List[str]]: + """ + Validate and parse command line arguments. + + Args: + args: Command line arguments including script name + + Returns: + Tuple of (sync_port, working_directory, target_args) + + Raises: + ArgumentError: If arguments are invalid + """ + if len(args) < 4: + raise ArgumentError( + "Insufficient arguments. Expected: [args...]" + ) + + try: + sync_port = int(args[1]) + if not (1 <= sync_port <= 65535): + raise ValueError("Port out of range") + except ValueError as e: + raise ArgumentError(f"Invalid sync port '{args[1]}': {e}") from e + + cwd = args[2] + if not os.path.isdir(cwd): + raise ArgumentError(f"Working directory does not exist: {cwd}") + + target_args = args[3:] + if not target_args: + raise ArgumentError("No target specified") + + return sync_port, cwd, target_args + + +# Constants for socket communication +_MAX_RETRIES = 3 +_INITIAL_RETRY_DELAY_SEC = 0.1 +_SOCKET_TIMEOUT_SEC = 2.0 +_READY_MESSAGE = b"ready" + + +def _signal_readiness(sync_port: int) -> None: + """ + Signal readiness to the profiler via TCP socket. + + Args: + sync_port: Port number where profiler is listening + + Raises: + SyncError: If unable to signal readiness + """ + last_error = None + + for attempt in range(_MAX_RETRIES): + try: + # Use context manager for automatic cleanup + with socket.create_connection(("127.0.0.1", sync_port), timeout=_SOCKET_TIMEOUT_SEC) as sock: + sock.send(_READY_MESSAGE) + return + except (socket.error, OSError) as e: + last_error = e + if attempt < _MAX_RETRIES - 1: + # Exponential backoff before retry + time.sleep(_INITIAL_RETRY_DELAY_SEC * (2 ** attempt)) + + # If we get here, all retries failed + raise SyncError(f"Failed to signal readiness after {_MAX_RETRIES} attempts: {last_error}") from last_error + + +def _setup_environment(cwd: str) -> None: + """ + Set up the execution environment. + + Args: + cwd: Working directory to change to + + Raises: + TargetError: If unable to set up environment + """ + try: + os.chdir(cwd) + except OSError as e: + raise TargetError(f"Failed to change to directory {cwd}: {e}") from e + + # Add current directory to sys.path if not present (for module imports) + if cwd not in sys.path: + sys.path.insert(0, cwd) + + +def _execute_module(module_name: str, module_args: List[str]) -> None: + """ + Execute a Python module. + + Args: + module_name: Name of the module to execute + module_args: Arguments to pass to the module + + Raises: + TargetError: If module cannot be found + """ + # Replace sys.argv to match how Python normally runs modules + # When running 'python -m module args', sys.argv is ["__main__.py", "args"] + sys.argv = ["__main__.py"] + module_args + + try: + runpy.run_module(module_name, run_name="__main__", alter_sys=True) + except ImportError as e: + raise TargetError(f"Module '{module_name}' not found: {e}") from e + # Let other exceptions (including SystemExit) propagate naturally + # so Python prints the full traceback to stderr + + +def _execute_script(script_path: str, script_args: List[str], cwd: str) -> None: + """ + Execute a Python script. + + Args: + script_path: Path to the script to execute + script_args: Arguments to pass to the script + cwd: Current working directory for path resolution + + Raises: + TargetError: If script execution fails + """ + # Make script path absolute if it isn't already + if not os.path.isabs(script_path): + script_path = os.path.join(cwd, script_path) + + if not os.path.isfile(script_path): + raise TargetError(f"Script not found: {script_path}") + + # Replace sys.argv to match original script call + sys.argv = [script_path] + script_args + + try: + with open(script_path, 'rb') as f: + source_code = f.read() + + except FileNotFoundError as e: + raise TargetError(f"Script file not found: {script_path}") from e + except PermissionError as e: + raise TargetError(f"Permission denied reading script: {script_path}") from e + + main_module = types.ModuleType("__main__") + main_module.__file__ = script_path + main_module.__builtins__ = __builtins__ + # gh-140729: Create a __mp_main__ module to allow pickling + sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module + + try: + code = compile(source_code, script_path, 'exec', module='__main__') + except SyntaxError as e: + raise TargetError(f"Syntax error in script {script_path}: {e}") from e + + # Execute the script - let exceptions propagate naturally so Python + # prints the full traceback to stderr + exec(code, main_module.__dict__) + + +def main() -> NoReturn: + """ + Main coordinator function. + + This function coordinates the startup of a target Python process + with the sample profiler by signaling when the process is ready + to be profiled. + """ + # Phase 1: Parse arguments and set up environment + # Errors here are coordinator errors, not script errors + try: + # Parse and validate arguments + sync_port, cwd, target_args = _validate_arguments(sys.argv) + + # Set up execution environment + _setup_environment(cwd) + + # Determine execution type and validate target exists + is_module = target_args[0] == "-m" + if is_module: + if len(target_args) < 2: + raise ArgumentError("Module name required after -m") + module_name = target_args[1] + module_args = target_args[2:] + + if importlib.util.find_spec(module_name) is None: + raise TargetError(f"Module not found: {module_name}") + else: + script_path = target_args[0] + script_args = target_args[1:] + # Match the path resolution logic in _execute_script + check_path = script_path if os.path.isabs(script_path) else os.path.join(cwd, script_path) + if not os.path.isfile(check_path): + raise TargetError(f"Script not found: {script_path}") + + # Signal readiness to profiler + _signal_readiness(sync_port) + + except CoordinatorError as e: + print(f"Profiler coordinator error: {e}", file=sys.stderr) + sys.exit(1) + except KeyboardInterrupt: + print("Interrupted", file=sys.stderr) + sys.exit(1) + + # Phase 2: Execute the target script/module + # Let exceptions propagate naturally so Python prints full tracebacks + if is_module: + _execute_module(module_name, module_args) + else: + _execute_script(script_path, script_args, cwd) + + # Normal exit + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/Lib/profiling/sampling/cli.py b/Lib/profiling/sampling/cli.py new file mode 100644 index 000000000..67adf4c93 --- /dev/null +++ b/Lib/profiling/sampling/cli.py @@ -0,0 +1,1462 @@ +"""Command-line interface for the sampling profiler.""" + +import argparse +import importlib.util +import locale +import os +import re +import selectors +import socket +import subprocess +import sys +import time +import webbrowser +from contextlib import nullcontext + +from .errors import SamplingUnknownProcessError, SamplingModuleNotFoundError, SamplingScriptNotFoundError +from .sample import sample, sample_live, dump_stack, _is_process_running +from .dump import print_stack_dump +from .pstats_collector import PstatsCollector +from .stack_collector import CollapsedStackCollector, FlamegraphCollector, DiffFlamegraphCollector +from .heatmap_collector import HeatmapCollector +from .gecko_collector import GeckoCollector +from .jsonl_collector import JsonlCollector +from .binary_collector import BinaryCollector +from .binary_reader import BinaryReader +from .constants import ( + MICROSECONDS_PER_SECOND, + PROFILING_MODE_ALL, + PROFILING_MODE_WALL, + PROFILING_MODE_CPU, + PROFILING_MODE_GIL, + PROFILING_MODE_EXCEPTION, + SORT_MODE_NSAMPLES, + SORT_MODE_TOTTIME, + SORT_MODE_CUMTIME, + SORT_MODE_SAMPLE_PCT, + SORT_MODE_CUMUL_PCT, + SORT_MODE_NSAMPLES_CUMUL, +) + +try: + from ._child_monitor import ChildProcessMonitor +except ImportError: + # _remote_debugging module not available on this platform (e.g., WASI) + ChildProcessMonitor = None + +try: + from .live_collector import LiveStatsCollector +except ImportError: + LiveStatsCollector = None + + +class CustomFormatter( + argparse.ArgumentDefaultsHelpFormatter, + argparse.RawDescriptionHelpFormatter, +): + """Custom formatter that combines default values display with raw description formatting.""" + pass + + +class DiffFlamegraphAction(argparse.Action): + """Custom action for --diff-flamegraph that sets both format and baseline path.""" + def __call__(self, parser, namespace, values, option_string=None): + namespace.format = 'diff_flamegraph' + namespace.diff_baseline = values + + +_HELP_DESCRIPTION = """Sample a process's stack frames and generate profiling data. + +Examples: + # Run and profile a script + `python -m profiling.sampling run script.py arg1 arg2` + + # Attach to a running process + `python -m profiling.sampling attach 1234` + + # Dump a running process's current stack + `python -m profiling.sampling dump 1234` + + # Live interactive mode for a script + `python -m profiling.sampling run --live script.py` + + # Live interactive mode for a running process + `python -m profiling.sampling attach --live 1234` + +Use `python -m profiling.sampling --help` for command-specific help.""" + + +# Constants for socket synchronization +_SYNC_TIMEOUT_SEC = 5.0 +_PROCESS_KILL_TIMEOUT_SEC = 2.0 +_READY_MESSAGE = b"ready" +_RECV_BUFFER_SIZE = 1024 +_BINARY_PROFILE_HEADER_SIZE = 64 +_BINARY_PROFILE_MAGICS = (b"HCAT", b"TACH") + +# Format configuration +FORMAT_EXTENSIONS = { + "pstats": "pstats", + "collapsed": "txt", + "flamegraph": "html", + "diff_flamegraph": "html", + "gecko": "json", + "heatmap": "html", + "jsonl": "jsonl", + "binary": "bin", +} + +COLLECTOR_MAP = { + "pstats": PstatsCollector, + "collapsed": CollapsedStackCollector, + "flamegraph": FlamegraphCollector, + "diff_flamegraph": DiffFlamegraphCollector, + "gecko": GeckoCollector, + "heatmap": HeatmapCollector, + "jsonl": JsonlCollector, + "binary": BinaryCollector, +} + +BROWSER_COMPATIBLE_FORMATS = ("flamegraph", "diff_flamegraph", "heatmap") + + +def _setup_child_monitor(args, parent_pid): + # Build CLI args for child profilers (excluding --subprocesses to avoid recursion) + child_cli_args = _build_child_profiler_args(args) + + # Build output pattern + output_pattern = _build_output_pattern(args) + + return ChildProcessMonitor( + pid=parent_pid, + cli_args=child_cli_args, + output_pattern=output_pattern, + ) + + +def _get_child_monitor_context(args, pid): + if getattr(args, 'subprocesses', False): + return _setup_child_monitor(args, pid) + return nullcontext() + + +def _build_child_profiler_args(args): + child_args = [] + + # Sampling options + hz = MICROSECONDS_PER_SECOND // args.sample_interval_usec + child_args.extend(["-r", str(hz)]) + if args.duration is not None: + child_args.extend(["-d", str(args.duration)]) + if args.all_threads: + child_args.append("-a") + if args.realtime_stats: + child_args.append("--realtime-stats") + if args.native: + child_args.append("--native") + if not args.gc: + child_args.append("--no-gc") + if args.opcodes: + child_args.append("--opcodes") + if args.async_aware: + child_args.append("--async-aware") + async_mode = getattr(args, 'async_mode', 'running') + if async_mode != "running": + child_args.extend(["--async-mode", async_mode]) + + # Mode options + mode = getattr(args, 'mode', 'wall') + if mode != "wall": + child_args.extend(["--mode", mode]) + + # Format options (skip pstats as it's the default) + if args.format == "diff_flamegraph": + child_args.extend(["--diff-flamegraph", args.diff_baseline]) + elif args.format != "pstats": + child_args.append(f"--{args.format}") + + return child_args + + +def _build_output_pattern(args): + """Build output filename pattern for child profilers. + + The pattern uses {pid} as a placeholder which will be replaced with the + actual child PID using str.replace(), so user filenames with braces are safe. + """ + if args.outfile: + # User specified output - add PID to filename + base, ext = os.path.splitext(args.outfile) + if ext: + return f"{base}_{{pid}}{ext}" + else: + return f"{args.outfile}_{{pid}}" + else: + # Use default pattern based on format (consistent _ separator) + extension = FORMAT_EXTENSIONS.get(args.format, "txt") + if args.format == "heatmap": + return "heatmap_{pid}" + if args.format == "pstats": + # pstats defaults to stdout, but for subprocesses we need files + return "profile_{pid}.pstats" + return f"{args.format}_{{pid}}.{extension}" + + +def _parse_mode(mode_string): + """Convert mode string to mode constant.""" + mode_map = { + "wall": PROFILING_MODE_WALL, + "cpu": PROFILING_MODE_CPU, + "gil": PROFILING_MODE_GIL, + "exception": PROFILING_MODE_EXCEPTION, + } + return mode_map[mode_string] + + +def _check_process_died(process): + """Check if process died and raise an error with stderr if available.""" + if process.poll() is None: + return + + # Process died - try to get stderr for error message + stderr_msg = "" + if process.stderr: + try: + stderr_msg = process.stderr.read().decode().strip() + except (OSError, UnicodeDecodeError): + pass + + if stderr_msg: + raise RuntimeError(stderr_msg) + raise RuntimeError(f"Process exited with code {process.returncode}") + + +def _wait_for_ready_signal(sync_sock, process, timeout): + """Wait for the ready signal from the subprocess, checking for early death.""" + deadline = time.monotonic() + timeout + sel = selectors.DefaultSelector() + sel.register(sync_sock, selectors.EVENT_READ) + + try: + while True: + _check_process_died(process) + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise socket.timeout("timed out") + + if not sel.select(timeout=min(0.1, remaining)): + continue + + conn, _ = sync_sock.accept() + try: + ready_signal = conn.recv(_RECV_BUFFER_SIZE) + finally: + conn.close() + + if ready_signal != _READY_MESSAGE: + raise RuntimeError(f"Invalid ready signal received: {ready_signal!r}") + return + finally: + sel.close() + + +def _run_with_sync(original_cmd, suppress_output=False): + """Run a command with socket-based synchronization and return the process.""" + # Create a TCP socket for synchronization with better socket options + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sync_sock: + # Set SO_REUSEADDR to avoid "Address already in use" errors + sync_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sync_sock.bind(("127.0.0.1", 0)) # Let OS choose a free port + sync_port = sync_sock.getsockname()[1] + sync_sock.listen(1) + sync_sock.settimeout(_SYNC_TIMEOUT_SEC) + + # Get current working directory to preserve it + cwd = os.getcwd() + + # Build command using the sync coordinator + target_args = original_cmd[1:] # Remove python executable + cmd = ( + sys.executable, + "-m", + "profiling.sampling._sync_coordinator", + str(sync_port), + cwd, + ) + tuple(target_args) + + # Start the process with coordinator + # When suppress_output=True (live mode), capture stderr so we can + # report errors if the process dies before signaling ready. + # When suppress_output=False (normal mode), let stderr inherit so + # script errors print to the terminal. + popen_kwargs = {} + if suppress_output: + popen_kwargs["stdin"] = subprocess.DEVNULL + popen_kwargs["stdout"] = subprocess.DEVNULL + popen_kwargs["stderr"] = subprocess.PIPE + + process = subprocess.Popen(cmd, **popen_kwargs) + + try: + _wait_for_ready_signal(sync_sock, process, _SYNC_TIMEOUT_SEC) + except socket.timeout: + # If we timeout, kill the process and raise an error + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=_PROCESS_KILL_TIMEOUT_SEC) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + raise RuntimeError( + "Process failed to signal readiness within timeout" + ) + + return process + + +_RATE_PATTERN = re.compile(r''' + ^ # Start of string + ( # Group 1: The numeric value + \d+ # One or more digits (integer part) + (?:\.\d+)? # Optional: decimal point followed by digits + ) # Examples: "10", "0.5", "100.25" + ( # Group 2: Optional unit suffix + hz # "hz" - hertz + | khz # "khz" - kilohertz + | k # "k" - shorthand for kilohertz + )? # Suffix is optional (bare number = Hz) + $ # End of string + ''', re.VERBOSE | re.IGNORECASE) + + +def _parse_sampling_rate(rate_str: str) -> int: + """Parse sampling rate string to microseconds.""" + rate_str = rate_str.strip().lower() + + match = _RATE_PATTERN.match(rate_str) + if not match: + raise argparse.ArgumentTypeError( + f"Invalid sampling rate format: {rate_str}. " + "Expected: number followed by optional suffix (hz, khz, k) with no spaces (e.g., 10khz)" + ) + + number_part = match.group(1) + suffix = match.group(2) or '' + + # Determine multiplier based on suffix + suffix_map = { + 'hz': 1, + 'khz': 1000, + 'k': 1000, + } + multiplier = suffix_map.get(suffix, 1) + hz = float(number_part) * multiplier + if hz <= 0: + raise argparse.ArgumentTypeError(f"Sampling rate must be positive: {rate_str}") + + interval_usec = int(MICROSECONDS_PER_SECOND / hz) + if interval_usec < 1: + raise argparse.ArgumentTypeError(f"Sampling rate too high: {rate_str}") + + return interval_usec + + +def _add_sampling_options(parser): + """Add sampling configuration options to a parser.""" + sampling_group = parser.add_argument_group("Sampling configuration") + sampling_group.add_argument( + "-r", + "--sampling-rate", + type=_parse_sampling_rate, + default="1khz", + metavar="RATE", + dest="sample_interval_usec", + help="sampling rate (e.g., 10000, 10khz, 10k)", + ) + sampling_group.add_argument( + "-d", + "--duration", + type=int, + default=None, + metavar="SECONDS", + help="Sampling duration (default: run to completion)", + ) + sampling_group.add_argument( + "-a", + "--all-threads", + action="store_true", + help="Sample all threads in the process instead of just the main thread", + ) + sampling_group.add_argument( + "--realtime-stats", + action="store_true", + help="Print real-time sampling statistics (Hz, mean, min, max) during profiling", + ) + sampling_group.add_argument( + "--native", + action="store_true", + help='Include artificial `` frames to denote calls to non-Python code', + ) + sampling_group.add_argument( + "--no-gc", + action="store_false", + dest="gc", + help='Don\'t include artificial `` frames to denote active garbage collection', + ) + sampling_group.add_argument( + "--opcodes", + action="store_true", + help="Gather bytecode opcode information for instruction-level profiling " + "(shows which bytecode instructions are executing, including specializations).", + ) + sampling_group.add_argument( + "--async-aware", + action="store_true", + help="Enable async-aware profiling (uses task-based stack reconstruction)", + ) + sampling_group.add_argument( + "--subprocesses", + action="store_true", + help="Also profile subprocesses. Each subprocess gets its own profiler and output file.", + ) + sampling_group.add_argument( + "--blocking", + action="store_true", + help="Stop all threads in target process before sampling to get consistent snapshots. " + "Uses thread_suspend on macOS and ptrace on Linux. Adds overhead but ensures memory " + "reads are from a frozen state.", + ) + + +def _add_mode_options(parser): + """Add mode options to a parser.""" + mode_group = parser.add_argument_group("Mode options") + mode_group.add_argument( + "--mode", + choices=["wall", "cpu", "gil", "exception"], + default="wall", + help="Sampling mode: wall (all samples), cpu (only samples when thread is on CPU), " + "gil (only samples when thread holds the GIL), " + "exception (only samples when thread has an active exception). " + "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`', + ) + + +def _add_format_options(parser, include_compression=True, include_binary=True): + """Add output format options to a parser.""" + output_group = parser.add_argument_group("Output options") + format_group = output_group.add_mutually_exclusive_group() + format_group.add_argument( + "--pstats", + action="store_const", + const="pstats", + dest="format", + help="Generate pstats output (default)", + ) + format_group.add_argument( + "--collapsed", + action="store_const", + const="collapsed", + dest="format", + help="Generate collapsed stack traces for flamegraphs", + ) + format_group.add_argument( + "--flamegraph", + action="store_const", + const="flamegraph", + dest="format", + help="Generate interactive HTML flamegraph visualization", + ) + format_group.add_argument( + "--gecko", + action="store_const", + const="gecko", + dest="format", + help="Generate Gecko format for Firefox Profiler", + ) + format_group.add_argument( + "--heatmap", + action="store_const", + const="heatmap", + dest="format", + help="Generate interactive HTML heatmap visualization with line-level sample counts", + ) + format_group.add_argument( + "--diff-flamegraph", + metavar="BASELINE", + action=DiffFlamegraphAction, + help="Generate differential flamegraph comparing current profile to `BASELINE` binary file", + ) + format_group.add_argument( + "--jsonl", + action="store_const", + const="jsonl", + dest="format", + help="Generate newline-delimited JSON (JSONL) for programmatic consumers", + ) + if include_binary: + format_group.add_argument( + "--binary", + action="store_const", + const="binary", + dest="format", + help="Generate high-performance binary format (use `replay` command to convert)", + ) + parser.set_defaults(format="pstats", diff_baseline=None) + + if include_compression: + output_group.add_argument( + "--compression", + choices=["auto", "zstd", "none"], + default="auto", + help="Compression for binary format: auto (use zstd if available), zstd, none", + ) + + output_group.add_argument( + "-o", + "--output", + dest="outfile", + help="Output path (default: `stdout` for `pstats` text; with `-o`, `pstats` is binary). " + "Auto-generated for other formats. For heatmap: directory name (default: `heatmap_PID`)", + ) + output_group.add_argument( + "--browser", + action="store_true", + help=( + "Automatically open HTML output " + f"({', '.join('--' + f.replace('_', '-') for f in BROWSER_COMPATIBLE_FORMATS)}) " + "in browser. " + "When using `--subprocesses`, only the main process opens the browser" + ), + ) + + +def _add_pstats_options(parser): + """Add pstats-specific display options to a parser.""" + pstats_group = parser.add_argument_group("pstats format options") + pstats_group.add_argument( + "--sort", + choices=[ + "nsamples", + "tottime", + "cumtime", + "sample-pct", + "cumul-pct", + "nsamples-cumul", + "name", + ], + default=None, + help="Sort order for pstats output (default: nsamples)", + ) + pstats_group.add_argument( + "-l", + "--limit", + type=int, + default=None, + help="Limit the number of rows in the output (default: 15)", + ) + pstats_group.add_argument( + "--no-summary", + action="store_true", + help="Disable the summary section in the pstats output", + ) + + +def _add_dump_options(parser): + """Add one-shot stack dump options to a parser.""" + dump_group = parser.add_argument_group("Dump options") + dump_group.add_argument( + "-a", + "--all-threads", + action="store_true", + help="Dump all threads in the process instead of just the main thread", + ) + dump_group.add_argument( + "--native", + action="store_true", + help='Include artificial `` frames to denote calls to non-Python code', + ) + dump_group.add_argument( + "--no-gc", + action="store_false", + dest="gc", + help='Don\'t include artificial `` frames to denote active garbage collection', + ) + dump_group.add_argument( + "--opcodes", + action="store_true", + help="Show bytecode opcode names when available.", + ) + dump_group.add_argument( + "--async-aware", + action="store_true", + help="Enable async-aware stack reconstruction", + ) + dump_group.add_argument( + "--async-mode", + choices=["running", "all"], + default=argparse.SUPPRESS, + help='Async stack mode: "running" (only running task) ' + 'or "all" (all tasks including waiting, default for dump). ' + "Requires `--async-aware`", + ) + dump_group.add_argument( + "--blocking", + action="store_true", + help="Stop all threads in target process before dumping the stack.", + ) + + +def _sort_to_mode(sort_choice): + """Convert sort choice string to SORT_MODE constant.""" + sort_map = { + "nsamples": SORT_MODE_NSAMPLES, + "tottime": SORT_MODE_TOTTIME, + "cumtime": SORT_MODE_CUMTIME, + "sample-pct": SORT_MODE_SAMPLE_PCT, + "cumul-pct": SORT_MODE_CUMUL_PCT, + "nsamples-cumul": SORT_MODE_NSAMPLES_CUMUL, + "name": -1, + } + return sort_map.get(sort_choice, SORT_MODE_NSAMPLES) + + +def _capture_config_from_args(args): + return { + "all_threads": args.all_threads, + "native": args.native, + "gc": args.gc, + "opcodes": args.opcodes, + "blocking": args.blocking, + } + + +def _create_collector(format_type, sample_interval_usec, skip_idle, opcodes=False, + mode=None, output_file=None, compression='auto', + diff_baseline=None, capture_config=None): + """Create the appropriate collector based on format type. + + Args: + format_type: The output format ('pstats', 'collapsed', 'flamegraph', + 'gecko', 'heatmap', 'jsonl', 'binary', 'diff_flamegraph') + sample_interval_usec: Sampling interval in microseconds + skip_idle: Whether to skip idle samples + opcodes: Whether to collect opcode information (only used by gecko format + for creating interval markers in Firefox Profiler) + mode: Profiling mode for collectors that expose it in metadata + output_file: Output file path (required for binary format) + compression: Compression type for binary format ('auto', 'zstd', 'none') + diff_baseline: Path to baseline binary file for differential flamegraph + capture_config: Capture feature mapping for binary profiles and diffs + + Returns: + A collector instance of the appropriate type + """ + collector_class = COLLECTOR_MAP.get(format_type) + if collector_class is None: + raise ValueError(f"Unknown format: {format_type}") + + if format_type == "diff_flamegraph": + if diff_baseline is None: + raise ValueError("Differential flamegraph requires a baseline file") + if not os.path.exists(diff_baseline): + raise ValueError(f"Baseline file not found: {diff_baseline}") + return collector_class( + sample_interval_usec, + baseline_binary_path=diff_baseline, + skip_idle=skip_idle, + mode=mode, + capture_config=capture_config, + ) + + # Binary format requires output file and compression + if format_type == "binary": + if output_file is None: + raise ValueError("Binary format requires an output file") + return collector_class(output_file, sample_interval_usec, skip_idle=skip_idle, + compression=compression, mode=mode, + capture_config=capture_config) + + # Gecko format never skips idle (it needs both GIL and CPU data) + # and is the only format that uses opcodes for interval markers + if format_type == "gecko": + skip_idle = False + return collector_class(sample_interval_usec, skip_idle=skip_idle, opcodes=opcodes) + + if format_type == "jsonl": + return collector_class( + sample_interval_usec, skip_idle=skip_idle, mode=mode + ) + + return collector_class(sample_interval_usec, skip_idle=skip_idle) + + +def _generate_output_filename(format_type, pid): + """Generate output filename based on format and PID. + + Args: + format_type: The output format + pid: Process ID + + Returns: + Generated filename + """ + extension = FORMAT_EXTENSIONS.get(format_type, "txt") + # For heatmap, use cleaner directory name without extension + if format_type == "heatmap": + return f"heatmap_{pid}" + return f"{format_type}_{pid}.{extension}" + + +def _open_in_browser(path): + """Open a file or directory in the default web browser. + + Args: + path: File path or directory path to open + + For directories (heatmap), opens the index.html file inside. + """ + abs_path = os.path.abspath(path) + + # For heatmap directories, open the index.html file + if os.path.isdir(abs_path): + index_path = os.path.join(abs_path, 'index.html') + if os.path.exists(index_path): + abs_path = index_path + else: + print(f"Warning: Could not find index.html in {path}", file=sys.stderr) + return + + file_url = f"file://{abs_path}" + try: + webbrowser.open(file_url) + except Exception as e: + print(f"Warning: Could not open browser: {e}", file=sys.stderr) + + +def _validate_replay_input_file(filename): + """Validate that the replay input looks like a sampling binary profile.""" + try: + with open(filename, "rb") as file: + header = file.read(_BINARY_PROFILE_HEADER_SIZE) + except OSError as exc: + sys.exit(f"Error: Could not read input file {filename}: {exc}") + + if ( + len(header) < _BINARY_PROFILE_HEADER_SIZE + or header[:4] not in _BINARY_PROFILE_MAGICS + ): + sys.exit( + "Error: Input file is not a binary sampling profile. " + "The replay command only accepts files created with --binary" + ) + + +def _replay_with_reader(args, reader): + """Replay samples from an open binary reader.""" + info = reader.get_info() + interval = info['sample_interval_us'] + + print(f"Replaying {info['sample_count']} samples from {args.input_file}") + print(f" Sample interval: {interval} us") + print( + " Compression: " + f"{'zstd' if info.get('compression_type', 0) == 1 else 'none'}" + ) + + collector = _create_collector( + args.format, interval, skip_idle=False, + mode=info.get("mode"), + diff_baseline=args.diff_baseline, + capture_config=info.get("capture_config"), + ) + + def progress_callback(current, total): + if total > 0: + pct = current / total + bar_width = 40 + filled = int(bar_width * pct) + bar = 'โ–ˆ' * filled + 'โ–‘' * (bar_width - filled) + print( + f"\r [{bar}] {pct*100:5.1f}% ({current:,}/{total:,})", + end="", + flush=True, + ) + + count = reader.replay_samples(collector, progress_callback) + if hasattr(collector, "set_replay_stats"): + collector.set_replay_stats(info) + if hasattr(collector, "set_mode"): + collector.set_mode(info.get("mode")) + print() + + if args.format == "pstats": + if args.outfile: + collector.export(args.outfile) + else: + sort_choice = ( + args.sort if args.sort is not None else "nsamples" + ) + limit = args.limit if args.limit is not None else 15 + sort_mode = _sort_to_mode(sort_choice) + collector.print_stats( + sort_mode, limit, not args.no_summary, + info.get("mode") if info.get("mode") is not None + else PROFILING_MODE_WALL + ) + else: + filename = ( + args.outfile + or _generate_output_filename(args.format, os.getpid()) + ) + export_ok = collector.export(filename) + + # Auto-open browser for HTML output if --browser flag is set + if ( + export_ok + and args.format in BROWSER_COMPATIBLE_FORMATS + and getattr(args, 'browser', False) + ): + _open_in_browser(filename) + + print(f"Replayed {count} samples") + + +def _handle_output(collector, args, pid, mode): + """Handle output for the collector based on format and arguments. + + Args: + collector: The collector instance with profiling data + args: Parsed command-line arguments + pid: Process ID (for generating filenames) + mode: Profiling mode used + """ + if args.format == "binary": + # Binary format already wrote to file incrementally, just finalize + collector.export(None) + filename = collector.filename + print(f"Binary profile written to {filename} ({collector.total_samples} samples)") + elif args.format == "pstats": + if args.outfile: + # If outfile is a directory, generate filename inside it + if os.path.isdir(args.outfile): + filename = os.path.join(args.outfile, _generate_output_filename(args.format, pid)) + collector.export(filename) + else: + collector.export(args.outfile) + else: + # Print to stdout with defaults applied + sort_choice = args.sort if args.sort is not None else "nsamples" + limit = args.limit if args.limit is not None else 15 + sort_mode = _sort_to_mode(sort_choice) + collector.print_stats( + sort_mode, limit, not args.no_summary, mode + ) + else: + # Export to file + if args.outfile and os.path.isdir(args.outfile): + # If outfile is a directory, generate filename inside it + filename = os.path.join(args.outfile, _generate_output_filename(args.format, pid)) + else: + filename = args.outfile or _generate_output_filename(args.format, pid) + export_ok = collector.export(filename) + + # Auto-open browser for HTML output if --browser flag is set + if ( + export_ok + and args.format in BROWSER_COMPATIBLE_FORMATS + and getattr(args, 'browser', False) + ): + _open_in_browser(filename) + + +def _validate_args(args, parser): + """Validate format-specific options and live mode requirements. + + Args: + args: Parsed command-line arguments + parser: ArgumentParser instance for error reporting + """ + command = getattr(args, 'command', None) + + if command == "replay": + return + + # Check if live mode is available + if hasattr(args, 'live') and args.live and LiveStatsCollector is None: + parser.error( + "Live mode requires the curses module, which is not available." + ) + + # --subprocesses is incompatible with --live + if hasattr(args, 'subprocesses') and args.subprocesses: + if ChildProcessMonitor is None: + parser.error( + "--subprocesses is not available on this platform " + "(requires _remote_debugging module)." + ) + if hasattr(args, 'live') and args.live: + parser.error("--subprocesses is incompatible with --live mode.") + + # Async-aware mode is incompatible with options that need thread data. + if getattr(args, 'async_aware', False): + issues = [] + if getattr(args, 'native', False): + issues.append("--native") + if not getattr(args, 'gc', True): + issues.append("--no-gc") + if getattr(args, 'format', None) == "binary": + issues.append("--binary") + 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 when explicitly set + if not getattr(args, 'async_aware', False): + if command == "dump": + # dump uses SUPPRESS default, so attr only exists if user passed it + if hasattr(args, 'async_mode'): + parser.error("--async-mode requires --async-aware to be enabled.") + elif hasattr(args, 'async_mode') and args.async_mode != "running": + parser.error("--async-mode requires --async-aware to be enabled.") + + if command == "dump": + return + + # Warn about blocking mode with aggressive sampling intervals + if args.blocking and args.sample_interval_usec < 100: + print( + f"Warning: --blocking with a {args.sample_interval_usec} ยตs interval will stop all threads " + f"{1_000_000 // args.sample_interval_usec} times per second. " + "Consider using --sampling-rate 1khz or lower to reduce overhead.", + file=sys.stderr + ) + + # Live mode is incompatible with format options + if hasattr(args, 'live') and args.live: + if args.format != "pstats": + format_flag = f"--{args.format}" + parser.error( + f"--live is incompatible with {format_flag}. Live mode uses a TUI interface." + ) + + # Live mode is also incompatible with pstats-specific options + issues = [] + if args.sort is not None: + issues.append("--sort") + if args.limit is not None: + issues.append("--limit") + if args.no_summary: + issues.append("--no-summary") + + if issues: + parser.error( + f"Options {', '.join(issues)} are incompatible with --live. " + "Live mode uses a TUI interface with its own controls." + ) + return + + # Validate gecko mode doesn't use non-wall mode + if args.format == "gecko" and getattr(args, 'mode', 'wall') != "wall": + parser.error( + "--mode option is incompatible with --gecko. " + "Gecko format automatically includes both GIL-holding and CPU status analysis." + ) + + # Validate --opcodes is only used with compatible formats + opcodes_compatible_formats = ("live", "gecko", "flamegraph", "diff_flamegraph", "heatmap", "binary") + if getattr(args, 'opcodes', False) and args.format not in opcodes_compatible_formats: + parser.error( + f"--opcodes is only compatible with {', '.join('--' + f for f in opcodes_compatible_formats)}." + ) + + # Validate pstats-specific options are only used with pstats format + if args.format != "pstats": + issues = [] + if args.sort is not None: + issues.append("--sort") + if args.limit is not None: + issues.append("--limit") + if args.no_summary: + issues.append("--no-summary") + + if issues: + format_flag = f"--{args.format}" + parser.error( + f"Options {', '.join(issues)} are only valid with --pstats, not {format_flag}" + ) + + +def main(): + """Main entry point for the CLI.""" + # Set locale for number formatting, restore on exit + old_locale = locale.setlocale(locale.LC_ALL, None) + locale.setlocale(locale.LC_ALL, "") + try: + _main() + finally: + locale.setlocale(locale.LC_ALL, old_locale) + + +def _main(): + # Create the main parser + parser = argparse.ArgumentParser( + description=_HELP_DESCRIPTION, + formatter_class=CustomFormatter, + ) + + # Create subparsers for commands + subparsers = parser.add_subparsers( + dest="command", required=True, help="Command to run" + ) + + # === RUN COMMAND === + run_parser = subparsers.add_parser( + "run", + help="Run and profile a script or module", + formatter_class=CustomFormatter, + description="""Run and profile a Python script or module + +Examples: + # Run and profile a module + `python -m profiling.sampling run -m mymodule arg1 arg2` + + # Generate flamegraph from a script + `python -m profiling.sampling run --flamegraph -o output.html script.py` + + # Profile with custom rate and duration + `python -m profiling.sampling run -r 5khz -d 30 script.py` + + # Save collapsed stacks to file + `python -m profiling.sampling run --collapsed -o stacks.txt script.py` + + # Live interactive mode for a script + `python -m profiling.sampling run --live script.py`""", + ) + run_parser.add_argument( + "-m", + "--module", + action="store_true", + help="Run target as a module (like `python -m`)", + ) + run_parser.add_argument( + "target", + help="Script file or module name to profile", + ) + run_parser.add_argument( + "args", + nargs=argparse.REMAINDER, + help="Arguments to pass to the script or module", + ) + run_parser.add_argument( + "--live", + action="store_true", + help="Interactive TUI profiler (top-like interface, press 'q' to quit, 's' to cycle sort)", + ) + _add_sampling_options(run_parser) + _add_mode_options(run_parser) + _add_format_options(run_parser) + _add_pstats_options(run_parser) + + # === ATTACH COMMAND === + attach_parser = subparsers.add_parser( + "attach", + help="Attach to and profile a running process", + formatter_class=CustomFormatter, + description="""Attach to a running process and profile it + +Examples: + # Profile all threads, sort by total time + `python -m profiling.sampling attach -a --sort tottime 1234` + + # Live interactive mode for a running process + `python -m profiling.sampling attach --live 1234`""", + ) + attach_parser.add_argument( + "pid", + type=int, + help="Process ID to attach to", + ) + attach_parser.add_argument( + "--live", + action="store_true", + help="Interactive TUI profiler (top-like interface, press 'q' to quit, 's' to cycle sort)", + ) + _add_sampling_options(attach_parser) + _add_mode_options(attach_parser) + _add_format_options(attach_parser) + _add_pstats_options(attach_parser) + + # === DUMP COMMAND === + dump_parser = subparsers.add_parser( + "dump", + help="Dump a running process's current stack", + formatter_class=CustomFormatter, + description="""Dump a running process's current Python stack + +Examples: + # Dump the main thread stack + `python -m profiling.sampling dump 1234` + + # Dump all thread stacks + `python -m profiling.sampling dump -a 1234`""", + ) + dump_parser.add_argument( + "pid", + type=int, + help="Process ID to dump", + ) + _add_dump_options(dump_parser) + + # === REPLAY COMMAND === + replay_parser = subparsers.add_parser( + "replay", + help="Replay a binary profile and convert to another format", + formatter_class=CustomFormatter, + description="""Replay a binary profile file and convert to another format + +Examples: + # Convert binary to flamegraph + `python -m profiling.sampling replay --flamegraph -o output.html profile.bin` + + # Convert binary to pstats and print to stdout + `python -m profiling.sampling replay profile.bin` + + # Convert binary to gecko format + `python -m profiling.sampling replay --gecko -o profile.json profile.bin`""", + ) + replay_parser.add_argument( + "input_file", + help="Binary profile file to replay", + ) + _add_format_options(replay_parser, include_compression=False, include_binary=False) + _add_pstats_options(replay_parser) + + # Parse arguments + args = parser.parse_args() + + # Validate arguments + _validate_args(args, parser) + + # Command dispatch table + command_handlers = { + "run": _handle_run, + "attach": _handle_attach, + "dump": _handle_dump, + "replay": _handle_replay, + } + + # Execute the appropriate command + handler = command_handlers.get(args.command) + if handler: + handler(args) + else: + parser.error(f"Unknown command: {args.command}") + + +def _handle_attach(args): + """Handle the 'attach' command.""" + if not _is_process_running(args.pid): + raise SamplingUnknownProcessError(args.pid) + # Check if live mode is requested + if args.live: + _handle_live_attach(args, args.pid) + return + + # Use PROFILING_MODE_ALL for gecko format + mode = ( + PROFILING_MODE_ALL + if args.format == "gecko" + else _parse_mode(args.mode) + ) + + # Determine skip_idle based on mode + skip_idle = ( + mode != PROFILING_MODE_WALL if mode != PROFILING_MODE_ALL else False + ) + + output_file = None + if args.format == "binary": + output_file = args.outfile or _generate_output_filename(args.format, args.pid) + + # Create the appropriate collector + collector = _create_collector( + args.format, args.sample_interval_usec, skip_idle, args.opcodes, mode, + output_file=output_file, + compression=getattr(args, 'compression', 'auto'), + diff_baseline=args.diff_baseline, + capture_config=_capture_config_from_args(args), + ) + + with _get_child_monitor_context(args, args.pid): + collector = sample( + args.pid, + collector, + duration_sec=args.duration, + 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, + opcodes=args.opcodes, + blocking=args.blocking, + ) + _handle_output(collector, args, args.pid, mode) + + +def _handle_dump(args): + if not _is_process_running(args.pid): + raise SamplingUnknownProcessError(args.pid) + + # Async-aware reconstruction requires wall mode so every thread is sampled, + # not just those holding the GIL or on CPU. + mode = PROFILING_MODE_WALL if args.async_aware else PROFILING_MODE_ALL + async_mode = getattr(args, "async_mode", "all") if args.async_aware else None + try: + stack_frames = dump_stack( + args.pid, + all_threads=args.all_threads, + mode=mode, + async_aware=async_mode, + native=args.native, + gc=args.gc, + opcodes=args.opcodes, + blocking=args.blocking, + ) + except ProcessLookupError: + sys.exit( + f"No stack dump collected - process {args.pid} exited before " + "its stack could be read." + ) + + print_stack_dump(stack_frames, pid=args.pid) + + +def _handle_run(args): + """Handle the 'run' command.""" + # Validate target exists before launching subprocess + if args.module: + # Temporarily add cwd to sys.path so we can find modules in the + # current directory, matching the coordinator's behavior + cwd = os.getcwd() + added_cwd = False + if cwd not in sys.path: + sys.path.insert(0, cwd) + added_cwd = True + try: + if importlib.util.find_spec(args.target) is None: + raise SamplingModuleNotFoundError(args.target) + finally: + if added_cwd: + sys.path.remove(cwd) + else: + if not os.path.exists(args.target): + raise SamplingScriptNotFoundError(args.target) + + # Check if live mode is requested + if args.live: + _handle_live_run(args) + return + + # Build the command to run + if args.module: + cmd = (sys.executable, "-m", args.target, *args.args) + else: + cmd = (sys.executable, args.target, *args.args) + + # Run with synchronization + try: + process = _run_with_sync(cmd, suppress_output=False) + except RuntimeError as e: + sys.exit(f"Error: {e}") + + # Use PROFILING_MODE_ALL for gecko format + mode = ( + PROFILING_MODE_ALL + if args.format == "gecko" + else _parse_mode(args.mode) + ) + + # Determine skip_idle based on mode + skip_idle = ( + mode != PROFILING_MODE_WALL if mode != PROFILING_MODE_ALL else False + ) + + output_file = None + if args.format == "binary": + output_file = args.outfile or _generate_output_filename(args.format, process.pid) + + # Create the appropriate collector + collector = _create_collector( + args.format, args.sample_interval_usec, skip_idle, args.opcodes, mode, + output_file=output_file, + compression=getattr(args, 'compression', 'auto'), + diff_baseline=args.diff_baseline, + capture_config=_capture_config_from_args(args), + ) + + with _get_child_monitor_context(args, process.pid): + try: + collector = sample( + process.pid, + collector, + duration_sec=args.duration, + 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, + opcodes=args.opcodes, + blocking=args.blocking, + ) + _handle_output(collector, args, process.pid, mode) + finally: + # Terminate the main subprocess - child profilers finish when their + # target processes exit + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=_PROCESS_KILL_TIMEOUT_SEC) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +def _handle_live_attach(args, pid): + """Handle live mode for an existing process.""" + mode = _parse_mode(args.mode) + + # Determine skip_idle based on mode + skip_idle = mode != PROFILING_MODE_WALL + + # Create live collector with default settings + collector = LiveStatsCollector( + args.sample_interval_usec, + skip_idle=skip_idle, + sort_by="tottime", # Default initial sort + limit=20, # Default limit + pid=pid, + mode=mode, + opcodes=args.opcodes, + async_aware=args.async_mode if args.async_aware else None, + ) + + # Sample in live mode + sample_live( + pid, + collector, + duration_sec=args.duration, + 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, + opcodes=args.opcodes, + blocking=args.blocking, + ) + + +def _handle_live_run(args): + """Handle live mode for running a script/module.""" + # Build the command to run + if args.module: + cmd = (sys.executable, "-m", args.target, *args.args) + else: + cmd = (sys.executable, args.target, *args.args) + + # Run with synchronization, suppressing output for live mode + try: + process = _run_with_sync(cmd, suppress_output=True) + except RuntimeError as e: + sys.exit(f"Error: {e}") + + mode = _parse_mode(args.mode) + + # Determine skip_idle based on mode + skip_idle = mode != PROFILING_MODE_WALL + + # Create live collector with default settings + collector = LiveStatsCollector( + args.sample_interval_usec, + skip_idle=skip_idle, + sort_by="tottime", # Default initial sort + limit=20, # Default limit + pid=process.pid, + mode=mode, + opcodes=args.opcodes, + async_aware=args.async_mode if args.async_aware else None, + ) + + # Profile the subprocess in live mode + try: + sample_live( + process.pid, + collector, + duration_sec=args.duration, + 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, + opcodes=args.opcodes, + blocking=args.blocking, + ) + finally: + # Clean up the subprocess and get any error output + returncode = process.poll() + if returncode is None: + # Process still running - terminate it + process.terminate() + try: + process.wait(timeout=_PROCESS_KILL_TIMEOUT_SEC) + except subprocess.TimeoutExpired: + process.kill() + # Ensure process is fully terminated + process.wait() + # Read any stderr output (tracebacks, errors, etc.) + if process.stderr: + with process.stderr: + try: + stderr = process.stderr.read() + if stderr: + print(stderr.decode(), file=sys.stderr) + except (OSError, ValueError): + # Ignore errors if pipe is already closed + pass + + +def _handle_replay(args): + """Handle the 'replay' command - convert binary profile to another format.""" + if not os.path.exists(args.input_file): + sys.exit(f"Error: Input file not found: {args.input_file}") + + _validate_replay_input_file(args.input_file) + + try: + with BinaryReader(args.input_file) as reader: + _replay_with_reader(args, reader) + except (OSError, ValueError) as exc: + sys.exit(f"Error: {exc}") + + +if __name__ == "__main__": + main() diff --git a/Lib/profiling/sampling/errors.py b/Lib/profiling/sampling/errors.py new file mode 100644 index 000000000..0832ad2d4 --- /dev/null +++ b/Lib/profiling/sampling/errors.py @@ -0,0 +1,19 @@ +"""Custom exceptions for the sampling profiler.""" + +class SamplingProfilerError(Exception): + """Base exception for sampling profiler errors.""" + +class SamplingUnknownProcessError(SamplingProfilerError): + def __init__(self, pid): + self.pid = pid + super().__init__(f"Process with PID '{pid}' does not exist.") + +class SamplingScriptNotFoundError(SamplingProfilerError): + def __init__(self, script_path): + self.script_path = script_path + super().__init__(f"Script '{script_path}' not found.") + +class SamplingModuleNotFoundError(SamplingProfilerError): + def __init__(self, module_name): + self.module_name = module_name + super().__init__(f"Module '{module_name}' not found.") diff --git a/Lib/test/test_profiling/test_sampling_profiler/test_cli.py b/Lib/test/test_profiling/test_sampling_profiler/test_cli.py new file mode 100644 index 000000000..e5dcd72c2 --- /dev/null +++ b/Lib/test/test_profiling/test_sampling_profiler/test_cli.py @@ -0,0 +1,1083 @@ +"""Tests for sampling profiler CLI argument parsing and functionality.""" + +import io +import json +import os +import subprocess +import sys +import tempfile +import unittest +from types import SimpleNamespace +from unittest import mock + +try: + import _remote_debugging # noqa: F401 +except ImportError: + raise unittest.SkipTest( + "Test only runs when _remote_debugging is available" + ) + +from test.support import ( + force_not_colorized, + is_emscripten, + requires_remote_subprocess_debugging, +) + +from profiling.sampling.binary_reader import BinaryReader +from profiling.sampling.cli import ( + FORMAT_EXTENSIONS, + _create_collector, + _generate_output_filename, + _handle_output, + _replay_with_reader, + main, +) +from profiling.sampling.constants import ( + PROFILING_MODE_ALL, + PROFILING_MODE_CPU, + PROFILING_MODE_WALL, +) +from profiling.sampling.errors import SamplingScriptNotFoundError, SamplingModuleNotFoundError, SamplingUnknownProcessError +from profiling.sampling.jsonl_collector import JsonlCollector + +class TestSampleProfilerCLI(unittest.TestCase): + def _setup_sync_mocks(self, mock_socket, mock_popen): + """Helper to set up socket and process mocks for coordinator tests.""" + # Mock the sync socket with context manager support + mock_sock_instance = mock.MagicMock() + mock_sock_instance.getsockname.return_value = ("127.0.0.1", 12345) + + # Mock the connection with context manager support + mock_conn = mock.MagicMock() + mock_conn.recv.return_value = b"ready" + mock_conn.__enter__.return_value = mock_conn + mock_conn.__exit__.return_value = None + + # Mock accept() to return (connection, address) and support indexing + mock_accept_result = mock.MagicMock() + mock_accept_result.__getitem__.return_value = ( + mock_conn # [0] returns the connection + ) + mock_sock_instance.accept.return_value = mock_accept_result + + # Mock socket with context manager support + mock_sock_instance.__enter__.return_value = mock_sock_instance + mock_sock_instance.__exit__.return_value = None + mock_socket.return_value = mock_sock_instance + + # Mock the subprocess + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + return mock_process + + def _verify_coordinator_command(self, mock_popen, expected_target_args): + """Helper to verify the coordinator command was called correctly.""" + args, kwargs = mock_popen.call_args + coordinator_cmd = args[0] + self.assertEqual(coordinator_cmd[0], sys.executable) + self.assertEqual(coordinator_cmd[1], "-m") + self.assertEqual( + coordinator_cmd[2], "profiling.sampling._sync_coordinator" + ) + self.assertEqual(coordinator_cmd[3], "12345") # port + # cwd is coordinator_cmd[4] + self.assertEqual(coordinator_cmd[5:], expected_target_args) + + @unittest.skipIf(is_emscripten, "socket.SO_REUSEADDR does not exist") + @requires_remote_subprocess_debugging() + def test_cli_module_argument_parsing(self): + test_args = ["profiling.sampling.cli", "run", "-m", "mymodule"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + mock.patch("socket.socket") as mock_socket, + mock.patch("profiling.sampling.cli._wait_for_ready_signal"), + mock.patch("importlib.util.find_spec", return_value=True), + ): + self._setup_sync_mocks(mock_socket, mock_popen) + main() + + self._verify_coordinator_command(mock_popen, ("-m", "mymodule")) + # Verify sample was called once (exact arguments will vary with the new API) + mock_sample.assert_called_once() + + @unittest.skipIf(is_emscripten, "socket.SO_REUSEADDR does not exist") + @requires_remote_subprocess_debugging() + def test_cli_module_with_arguments(self): + test_args = [ + "profiling.sampling.cli", + "run", + "-m", + "mymodule", + "arg1", + "arg2", + "--flag", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + mock.patch("socket.socket") as mock_socket, + mock.patch("profiling.sampling.cli._wait_for_ready_signal"), + mock.patch("importlib.util.find_spec", return_value=True), + ): + self._setup_sync_mocks(mock_socket, mock_popen) + main() + + self._verify_coordinator_command( + mock_popen, ("-m", "mymodule", "arg1", "arg2", "--flag") + ) + mock_sample.assert_called_once() + + @unittest.skipIf(is_emscripten, "socket.SO_REUSEADDR does not exist") + def test_cli_script_argument_parsing(self): + test_args = ["profiling.sampling.cli", "run", "myscript.py"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + mock.patch("socket.socket") as mock_socket, + mock.patch("profiling.sampling.cli._wait_for_ready_signal"), + mock.patch("os.path.exists", return_value=True), + ): + self._setup_sync_mocks(mock_socket, mock_popen) + main() + + self._verify_coordinator_command(mock_popen, ("myscript.py",)) + mock_sample.assert_called_once() + + @unittest.skipIf(is_emscripten, "socket.SO_REUSEADDR does not exist") + def test_cli_script_with_arguments(self): + test_args = [ + "profiling.sampling.cli", + "run", + "myscript.py", + "arg1", + "arg2", + "--flag", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + mock.patch("socket.socket") as mock_socket, + mock.patch("profiling.sampling.cli._wait_for_ready_signal"), + mock.patch("os.path.exists", return_value=True), + ): + # Use the helper to set up mocks consistently + mock_process = self._setup_sync_mocks(mock_socket, mock_popen) + # Override specific behavior for this test + mock_process.wait.side_effect = [ + subprocess.TimeoutExpired(test_args, 0.1), + None, + ] + + main() + + # Verify the coordinator command was called + args, kwargs = mock_popen.call_args + coordinator_cmd = args[0] + self.assertEqual(coordinator_cmd[0], sys.executable) + self.assertEqual(coordinator_cmd[1], "-m") + self.assertEqual( + coordinator_cmd[2], "profiling.sampling._sync_coordinator" + ) + self.assertEqual(coordinator_cmd[3], "12345") # port + # cwd is coordinator_cmd[4] + self.assertEqual( + coordinator_cmd[5:], ("myscript.py", "arg1", "arg2", "--flag") + ) + + def test_cli_mutually_exclusive_pid_module(self): + # In new CLI, attach and run are separate subcommands, so this test + # verifies that mixing them causes an error + test_args = [ + "profiling.sampling.cli", + "attach", # attach subcommand uses PID + "12345", + "-m", # -m is only for run subcommand + "mymodule", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("unrecognized arguments", error_msg) + + def test_cli_mutually_exclusive_pid_script(self): + # In new CLI, you can't mix attach (PID) with run (script) + # This would be caught by providing a PID to run subcommand + test_args = ["profiling.sampling.cli", "run", "12345"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SamplingScriptNotFoundError) as cm, + ): + main() + + # Verify the error is about the non-existent script + self.assertIn("12345", str(cm.exception)) + + def test_cli_no_target_specified(self): + # In new CLI, must specify a subcommand + test_args = ["profiling.sampling.cli", "-d", "5"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("invalid choice", error_msg) + + @unittest.skipIf(is_emscripten, "socket.SO_REUSEADDR does not exist") + @requires_remote_subprocess_debugging() + def test_cli_module_with_profiler_options(self): + test_args = [ + "profiling.sampling.cli", + "run", + "-r", + "1000", + "-d", + "30", + "-a", + "--sort", + "tottime", + "-l", + "20", + "-m", + "mymodule", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + mock.patch("socket.socket") as mock_socket, + mock.patch("profiling.sampling.cli._wait_for_ready_signal"), + mock.patch("importlib.util.find_spec", return_value=True), + ): + self._setup_sync_mocks(mock_socket, mock_popen) + main() + + self._verify_coordinator_command(mock_popen, ("-m", "mymodule")) + mock_sample.assert_called_once() + + @unittest.skipIf(is_emscripten, "socket.SO_REUSEADDR does not exist") + def test_cli_script_with_profiler_options(self): + """Test script with various profiler options.""" + test_args = [ + "profiling.sampling.cli", + "run", + "-r", + "500", + "-d", + "60", + "--collapsed", + "-o", + "output.txt", + "myscript.py", + "scriptarg", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + mock.patch("socket.socket") as mock_socket, + mock.patch("profiling.sampling.cli._wait_for_ready_signal"), + mock.patch("os.path.exists", return_value=True), + ): + self._setup_sync_mocks(mock_socket, mock_popen) + main() + + self._verify_coordinator_command( + mock_popen, ("myscript.py", "scriptarg") + ) + # Verify profiler was called + mock_sample.assert_called_once() + + def test_cli_empty_module_name(self): + test_args = ["profiling.sampling.cli", "run", "-m"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("required: target", error_msg) # argparse error for missing positional arg + + @unittest.skipIf(is_emscripten, "socket.SO_REUSEADDR does not exist") + @requires_remote_subprocess_debugging() + def test_cli_long_module_option(self): + test_args = [ + "profiling.sampling.cli", + "run", + "-m", + "mymodule", + "arg1", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + mock.patch("socket.socket") as mock_socket, + mock.patch("profiling.sampling.cli._wait_for_ready_signal"), + mock.patch("importlib.util.find_spec", return_value=True), + ): + self._setup_sync_mocks(mock_socket, mock_popen) + main() + + self._verify_coordinator_command( + mock_popen, ("-m", "mymodule", "arg1") + ) + + def test_cli_complex_script_arguments(self): + test_args = [ + "profiling.sampling.cli", + "run", + "script.py", + "--input", + "file.txt", + "-v", + "--output=/tmp/out", + "positional", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + mock.patch( + "profiling.sampling.cli._run_with_sync" + ) as mock_run_with_sync, + mock.patch("os.path.exists", return_value=True), + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [ + subprocess.TimeoutExpired(test_args, 0.1), + None, + ] + mock_process.poll.return_value = None + mock_run_with_sync.return_value = mock_process + + main() + + mock_run_with_sync.assert_called_once_with( + ( + sys.executable, + "script.py", + "--input", + "file.txt", + "-v", + "--output=/tmp/out", + "positional", + ), + suppress_output=False + ) + + def test_cli_collapsed_format_validation(self): + """Test that CLI properly validates incompatible options with collapsed format.""" + test_cases = [ + # Test sort option is invalid with collapsed + ( + [ + "profiling.sampling.cli", + "attach", + "12345", + "--collapsed", + "--sort", + "tottime", # Changed from nsamples (default) to trigger validation + ], + "sort", + ), + # Test limit option is invalid with collapsed + ( + [ + "profiling.sampling.cli", + "attach", + "12345", + "--collapsed", + "-l", + "20", + ], + "limit", + ), + # Test no-summary option is invalid with collapsed + ( + [ + "profiling.sampling.cli", + "attach", + "12345", + "--collapsed", + "--no-summary", + ], + "summary", + ), + ] + + for test_args, expected_error_keyword in test_cases: + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + mock.patch("profiling.sampling.cli.sample"), # Prevent actual profiling + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error code + error_msg = mock_stderr.getvalue() + self.assertIn("error:", error_msg) + self.assertIn("only valid with --pstats", error_msg) + + def test_cli_default_collapsed_filename(self): + """Test that collapsed format gets a default filename when not specified.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--collapsed"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli._is_process_running", return_value=True), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + ): + main() + + # Check that sample was called (exact filename depends on implementation) + mock_sample.assert_called_once() + + def test_cli_custom_output_filenames(self): + """Test custom output filenames for both formats.""" + test_cases = [ + ( + [ + "profiling.sampling.cli", + "attach", + "12345", + "--pstats", + "-o", + "custom.pstats", + ], + "custom.pstats", + "pstats", + ), + ( + [ + "profiling.sampling.cli", + "attach", + "12345", + "--collapsed", + "-o", + "custom.txt", + ], + "custom.txt", + "collapsed", + ), + ] + + for test_args, expected_filename, expected_format in test_cases: + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli._is_process_running", return_value=True), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + ): + main() + + mock_sample.assert_called_once() + + def test_cli_missing_required_arguments(self): + """Test that CLI requires subcommand.""" + with ( + mock.patch("sys.argv", ["profiling.sampling.cli"]), + mock.patch("sys.stderr", io.StringIO()), + ): + with self.assertRaises(SystemExit): + main() + + def test_cli_mutually_exclusive_format_options(self): + """Test that pstats and collapsed options are mutually exclusive.""" + with ( + mock.patch( + "sys.argv", + [ + "profiling.sampling.cli", + "attach", + "12345", + "--pstats", + "--collapsed", + ], + ), + mock.patch("sys.stderr", io.StringIO()), + ): + with self.assertRaises(SystemExit): + main() + + def test_argument_parsing_basic(self): + test_args = ["profiling.sampling.cli", "attach", "12345"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli._is_process_running", return_value=True), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + ): + main() + + mock_sample.assert_called_once() + + def _run_dump_cli( + self, + *cli_args, + dump_return=None, + dump_side_effect=None, + process_running=True, + ): + """Run main() for a `dump` invocation, returning (mock_dump, mock_print).""" + argv = ["profiling.sampling.cli", "dump", *cli_args] + dump_kwargs = {} + if dump_side_effect is not None: + dump_kwargs["side_effect"] = dump_side_effect + else: + dump_kwargs["return_value"] = [] if dump_return is None else dump_return + with ( + mock.patch("sys.argv", argv), + mock.patch( + "profiling.sampling.cli._is_process_running", + return_value=process_running, + ), + mock.patch("profiling.sampling.cli.dump_stack", **dump_kwargs) as mock_dump_stack, + mock.patch("profiling.sampling.cli.print_stack_dump") as mock_print_stack_dump, + ): + main() + return mock_dump_stack, mock_print_stack_dump + + def _run_dump_cli_expecting_exit(self, *cli_args, capture="stderr"): + """Run main() for a `dump` invocation expected to SystemExit, return (cm, captured_text).""" + argv = ["profiling.sampling.cli", "dump", *cli_args] + buf = io.StringIO() + stream = "sys.stderr" if capture == "stderr" else "sys.stdout" + with ( + mock.patch("sys.argv", argv), + mock.patch(stream, buf), + self.assertRaises(SystemExit) as cm, + ): + main() + return cm, buf.getvalue() + + def test_cli_dump_subcommand(self): + stack_frames = [mock.sentinel.stack_frames] + mock_dump_stack, mock_print_stack_dump = self._run_dump_cli( + "12345", dump_return=stack_frames + ) + + mock_dump_stack.assert_called_once() + self.assertEqual(mock_dump_stack.call_args.args, (12345,)) + call_kwargs = mock_dump_stack.call_args.kwargs + self.assertFalse(call_kwargs["all_threads"]) + self.assertIsNone(call_kwargs["async_aware"]) + self.assertFalse(call_kwargs["native"]) + self.assertTrue(call_kwargs["gc"]) + self.assertFalse(call_kwargs["opcodes"]) + self.assertFalse(call_kwargs["blocking"]) + self.assertEqual(call_kwargs["mode"], PROFILING_MODE_ALL) + mock_print_stack_dump.assert_called_once_with(stack_frames, pid=12345) + + def test_cli_dump_subcommand_options(self): + mock_dump_stack, _ = self._run_dump_cli( + "-a", "--native", "--no-gc", "--opcodes", "--blocking", "12345" + ) + + call_kwargs = mock_dump_stack.call_args.kwargs + self.assertTrue(call_kwargs["all_threads"]) + self.assertTrue(call_kwargs["native"]) + self.assertFalse(call_kwargs["gc"]) + self.assertTrue(call_kwargs["opcodes"]) + self.assertTrue(call_kwargs["blocking"]) + self.assertEqual(call_kwargs["mode"], PROFILING_MODE_ALL) + + def test_cli_dump_rejects_mode_option(self): + cm, stderr = self._run_dump_cli_expecting_exit("12345", "--mode", "cpu") + self.assertEqual(cm.exception.code, 2) + self.assertIn("unrecognized arguments: --mode", stderr) + + def test_cli_dump_async_aware_defaults_to_all(self): + mock_dump_stack, _ = self._run_dump_cli("--async-aware", "12345") + call_kwargs = mock_dump_stack.call_args.kwargs + self.assertEqual(call_kwargs["async_aware"], "all") + self.assertEqual(call_kwargs["mode"], PROFILING_MODE_WALL) + + def test_cli_dump_async_mode_all_is_forwarded(self): + mock_dump_stack, _ = self._run_dump_cli( + "--async-aware", "--async-mode", "all", "12345" + ) + call_kwargs = mock_dump_stack.call_args.kwargs + self.assertEqual(call_kwargs["async_aware"], "all") + self.assertEqual(call_kwargs["mode"], PROFILING_MODE_WALL) + + def test_cli_dump_async_mode_running_is_forwarded(self): + mock_dump_stack, _ = self._run_dump_cli( + "--async-aware", "--async-mode", "running", "12345" + ) + call_kwargs = mock_dump_stack.call_args.kwargs + self.assertEqual(call_kwargs["async_aware"], "running") + self.assertEqual(call_kwargs["mode"], PROFILING_MODE_WALL) + + def test_cli_dump_async_mode_requires_async_aware(self): + cm, stderr = self._run_dump_cli_expecting_exit( + "--async-mode", "running", "12345" + ) + self.assertEqual(cm.exception.code, 2) + self.assertIn("--async-mode requires --async-aware", stderr) + + def test_cli_dump_rejects_async_aware_with_all_threads(self): + cm, stderr = self._run_dump_cli_expecting_exit( + "--async-aware", "-a", "12345" + ) + self.assertEqual(cm.exception.code, 2) + self.assertIn("--all-threads", stderr) + self.assertIn("incompatible with --async-aware", stderr) + + def test_cli_dump_rejects_async_aware_with_native(self): + cm, stderr = self._run_dump_cli_expecting_exit( + "--async-aware", "--native", "12345" + ) + self.assertEqual(cm.exception.code, 2) + self.assertIn("--native", stderr) + self.assertIn("incompatible with --async-aware", stderr) + + def test_cli_dump_rejects_async_aware_with_no_gc(self): + cm, stderr = self._run_dump_cli_expecting_exit( + "--async-aware", "--no-gc", "12345" + ) + self.assertEqual(cm.exception.code, 2) + self.assertIn("--no-gc", stderr) + self.assertIn("incompatible with --async-aware", stderr) + + def test_cli_dump_unknown_process(self): + with self.assertRaises(SamplingUnknownProcessError): + self._run_dump_cli("12345", process_running=False) + + def test_cli_dump_process_exits_before_snapshot(self): + with self.assertRaises(SystemExit) as cm: + self._run_dump_cli("12345", dump_side_effect=ProcessLookupError) + self.assertIn( + "No stack dump collected - process 12345 exited", + str(cm.exception.code), + ) + + @force_not_colorized + def test_cli_dump_help_lists_dump_options_without_mode(self): + cm, stdout = self._run_dump_cli_expecting_exit("--help", capture="stdout") + self.assertEqual(cm.exception.code, 0) + self.assertIn("--async-mode {running,all}", stdout) + self.assertIn("--opcodes", stdout) + self.assertNotIn("--mode {wall,cpu,gil,exception}", stdout) + + def test_sort_options(self): + sort_options = [ + ("nsamples", 0), + ("tottime", 1), + ("cumtime", 2), + ("sample-pct", 3), + ("cumul-pct", 4), + ("name", -1), + ] + + for option, expected_sort_value in sort_options: + test_args = ["profiling.sampling.cli", "attach", "12345", "--sort", option] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli._is_process_running", return_value=True), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + ): + main() + + mock_sample.assert_called_once() + mock_sample.reset_mock() + + def test_async_aware_flag_defaults_to_running(self): + """Test --async-aware flag enables async profiling with default 'running' mode.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli._is_process_running", return_value=True), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + ): + main() + + mock_sample.assert_called_once() + # Verify async_aware was passed with default "running" mode + call_kwargs = mock_sample.call_args[1] + self.assertEqual(call_kwargs.get("async_aware"), "running") + + def test_handle_output_browser_not_opened_when_export_fails(self): + for format_type in ("flamegraph", "diff_flamegraph", "heatmap"): + with self.subTest(format=format_type): + collector = mock.MagicMock() + collector.export.return_value = False + args = SimpleNamespace( + format=format_type, + outfile="profile.html", + browser=True, + ) + + with ( + mock.patch("profiling.sampling.cli.os.path.isdir", return_value=False), + mock.patch("profiling.sampling.cli._open_in_browser") as mock_open, + ): + _handle_output(collector, args, pid=12345, mode=0) + + collector.export.assert_called_once_with("profile.html") + mock_open.assert_not_called() + + def test_async_aware_with_async_mode_all(self): + """Test --async-aware with --async-mode all.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--async-mode", "all"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli._is_process_running", return_value=True), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + ): + main() + + mock_sample.assert_called_once() + call_kwargs = mock_sample.call_args[1] + self.assertEqual(call_kwargs.get("async_aware"), "all") + + def test_async_aware_default_is_none(self): + """Test async_aware defaults to None when --async-aware not specified.""" + test_args = ["profiling.sampling.cli", "attach", "12345"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profiling.sampling.cli._is_process_running", return_value=True), + mock.patch("profiling.sampling.cli.sample") as mock_sample, + ): + main() + + mock_sample.assert_called_once() + call_kwargs = mock_sample.call_args[1] + self.assertIsNone(call_kwargs.get("async_aware")) + + def test_async_mode_invalid_choice(self): + """Test --async-mode with invalid choice raises error.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--async-mode", "invalid"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()), + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + + def test_async_mode_requires_async_aware(self): + """Test --async-mode without --async-aware raises error.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-mode", "all"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("--async-mode requires --async-aware", error_msg) + + def test_async_aware_incompatible_with_native(self): + """Test --async-aware is incompatible with --native.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--native"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("--native", error_msg) + self.assertIn("incompatible with --async-aware", error_msg) + + def test_async_aware_incompatible_with_no_gc(self): + """Test --async-aware is incompatible with --no-gc.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--no-gc"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("--no-gc", error_msg) + self.assertIn("incompatible with --async-aware", error_msg) + + def test_async_aware_incompatible_with_both_native_and_no_gc(self): + """Test --async-aware is incompatible with both --native and --no-gc.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--native", "--no-gc"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("--native", error_msg) + self.assertIn("--no-gc", error_msg) + self.assertIn("incompatible with --async-aware", error_msg) + + def test_async_aware_incompatible_with_mode(self): + """Test --async-aware is incompatible with --mode (non-wall).""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--mode", "cpu"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("--mode=cpu", error_msg) + self.assertIn("incompatible with --async-aware", error_msg) + + def test_async_aware_incompatible_with_all_threads(self): + """Test --async-aware is incompatible with --all-threads.""" + test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--all-threads"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("--all-threads", error_msg) + self.assertIn("incompatible with --async-aware", error_msg) + + def test_async_aware_incompatible_with_binary(self): + """Test --async-aware is incompatible with --binary.""" + test_args = ["profiling.sampling.cli", "attach", "12345", + "--async-aware", "--binary"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("--binary", error_msg) + self.assertIn("incompatible with --async-aware", error_msg) + + @unittest.skipIf(is_emscripten, "subprocess not available") + def test_run_nonexistent_script_exits_cleanly(self): + """Test that running a non-existent script exits with a clean error.""" + with mock.patch("sys.argv", ["profiling.sampling.cli", "run", "/nonexistent/script.py"]): + with self.assertRaisesRegex(SamplingScriptNotFoundError, "Script '[\\w/.]+' not found."): + main() + + @unittest.skipIf(is_emscripten, "subprocess not available") + def test_run_nonexistent_module_exits_cleanly(self): + """Test that running a non-existent module exits with a clean error.""" + with mock.patch("sys.argv", ["profiling.sampling.cli", "run", "-m", "nonexistent_module_xyz"]): + with self.assertRaisesRegex(SamplingModuleNotFoundError, "Module '[\\w/.]+' not found."): + main() + + @unittest.skipIf(is_emscripten, "subprocess not available") + def test_cli_attach_nonexistent_pid(self): + fake_pid = "99999" + with mock.patch("sys.argv", ["profiling.sampling.cli", "attach", fake_pid]): + with self.assertRaises(SamplingUnknownProcessError) as cm: + main() + + self.assertIn(fake_pid, str(cm.exception)) + + def test_cli_replay_rejects_non_binary_profile(self): + with tempfile.TemporaryDirectory() as tempdir: + profile = os.path.join(tempdir, "output.prof") + with open(profile, "wb") as file: + file.write(b"not a binary sampling profile") + + with mock.patch("sys.argv", ["profiling.sampling.cli", "replay", profile]): + with self.assertRaises(SystemExit) as cm: + main() + + error = str(cm.exception) + self.assertIn("not a binary sampling profile", error) + self.assertIn("--binary", error) + + def test_cli_replay_reader_errors_exit_cleanly(self): + with tempfile.TemporaryDirectory() as tempdir: + profile = os.path.join(tempdir, "output.bin") + with open(profile, "wb") as file: + file.write(b"HCAT" + (b"\0" * 60)) + + with ( + mock.patch("sys.argv", ["profiling.sampling.cli", "replay", profile]), + mock.patch( + "profiling.sampling.cli.BinaryReader", + side_effect=ValueError("Unsupported format version 2"), + ), + ): + with self.assertRaises(SystemExit) as cm: + main() + + self.assertEqual( + str(cm.exception), + "Error: Unsupported format version 2", + ) + + def test_cli_replay_propagates_recorded_mode(self): + reader = mock.MagicMock() + reader.get_info.return_value = { + "sample_interval_us": 1000, + "sample_count": 0, + "compression_type": 0, + "mode": PROFILING_MODE_CPU, + "capture_config": {"all_threads": True}, + } + reader.replay_samples.return_value = 0 + collector = mock.MagicMock() + collector.export.return_value = True + args = SimpleNamespace( + format="diff_flamegraph", + input_file="current.bin", + diff_baseline="baseline.bin", + outfile="diff.html", + browser=False, + ) + + with mock.patch( + "profiling.sampling.cli._create_collector", + return_value=collector, + ) as create_collector: + _replay_with_reader(args, reader) + + create_collector.assert_called_once_with( + "diff_flamegraph", + 1000, + skip_idle=False, + mode=PROFILING_MODE_CPU, + diff_baseline="baseline.bin", + capture_config={"all_threads": True}, + ) + + def test_cli_jsonl_format_mutually_exclusive_with_pstats(self): + """--jsonl and --pstats cannot be combined (mutually exclusive group).""" + with ( + mock.patch( + "sys.argv", + [ + "profiling.sampling.cli", + "attach", + "12345", + "--jsonl", + "--pstats", + ], + ), + mock.patch("sys.stderr", io.StringIO()), + ): + with self.assertRaises(SystemExit): + main() + + def test_cli_jsonl_extension_in_format_extensions(self): + """FORMAT_EXTENSIONS maps 'jsonl' -> 'jsonl' so default filenames work.""" + self.assertEqual(FORMAT_EXTENSIONS["jsonl"], "jsonl") + self.assertEqual(_generate_output_filename("jsonl", 12345), "jsonl_12345.jsonl") + + def test_cli_jsonl_create_collector_propagates_mode(self): + """_create_collector('jsonl', ..., mode=X) lands X in the meta record.""" + collector = _create_collector( + "jsonl", + sample_interval_usec=1000, + skip_idle=False, + mode=PROFILING_MODE_CPU, + ) + self.assertIsInstance(collector, JsonlCollector) + + with tempfile.NamedTemporaryFile(suffix=".jsonl", delete=False) as f: + jsonl_path = f.name + self.addCleanup(os.unlink, jsonl_path) + collector.export(jsonl_path) + with open(jsonl_path, "r", encoding="utf-8") as f: + records = [json.loads(line) for line in f] + meta = next(r for r in records if r["type"] == "meta") + self.assertEqual(meta["mode"], "cpu") + + def test_cli_binary_create_collector_propagates_mode(self): + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + binary_path = f.name + self.addCleanup(os.unlink, binary_path) + collector = _create_collector( + "binary", + sample_interval_usec=1000, + skip_idle=True, + mode=PROFILING_MODE_CPU, + output_file=binary_path, + compression="none", + capture_config={"native": True}, + ) + collector.export(None) + + with BinaryReader(binary_path) as reader: + info = reader.get_info() + self.assertEqual(info["mode"], PROFILING_MODE_CPU) + self.assertTrue(info["capture_config"]["native"]) + + def test_cli_jsonl_rejects_opcodes_combination(self): + """--opcodes is incompatible with --jsonl per opcodes_compatible_formats.""" + test_args = [ + "profiling.sampling.cli", + "attach", + "12345", + "--jsonl", + "--opcodes", + ] + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + mock.patch("profiling.sampling.cli.sample"), + self.assertRaises(SystemExit) as cm, + ): + main() + + self.assertEqual(cm.exception.code, 2) + self.assertIn("--opcodes", mock_stderr.getvalue())