edq.cli.profile.flame

 1"""
 2Profile some Python target using Python's cprofile stdlib.
 3"""
 4
 5import argparse
 6import os
 7import shutil
 8import sys
 9
10import pyflame.sampler
11
12import edq.core.argparser
13import edq.util.dirent
14import edq.util.profile
15
16FLAMEGRAPH_URL: str = 'https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl'
17
18DEFAULT_FLAMEGRAPH_PATH: str = 'flamegraph.pl'
19DEFAULT_OUTPUT_PATH: str = 'flame.svg'
20DEFAULT_SAMPLE_INTERVAL_SECS: float = 0.000001
21
22def run_cli(args: argparse.Namespace) -> int:
23    """ Run the CLI. """
24
25    flamegraph_path = shutil.which(args.flamegraph_path)
26    if (flamegraph_path is None):
27        # Check for relative paths.
28        if (os.path.exists(args.flamegraph_path)):
29            flamegraph_path = os.path.abspath(args.flamegraph_path)
30        else:
31            raise ValueError(f"Failed to find flamegraph script ('{args.flamegraph_path}'), download from '{FLAMEGRAPH_URL}'.")
32
33    flamegraph_extra_args = [
34        '--width', '2000',
35    ]
36
37    sampler = pyflame.sampler.Sampler(args.sample_interval_secs)
38    edq.util.profile.exec_for_profile(args.run_target, imports = args.imports, is_module = args.module, module_args = args.module_args)
39    stack_counts = sampler.stop()
40
41    svg_text = pyflame.render.stack_counts_to_svg(stack_counts, flamegraph_path, flamegraph_extra_args)
42    edq.util.dirent.write_file(args.out_path, svg_text)
43
44    print(f"Wrote flamegraph to '{args.out_path}'.")
45
46    return 0
47
48def main() -> int:
49    """ Get a parser, parse the args, and call run. """
50    return run_cli(_get_parser().parse_args())
51
52def _get_parser() -> argparse.ArgumentParser:
53    """ Get the parser. """
54
55    parser = edq.core.argparser.get_default_parser(__doc__.strip())
56
57    parser.add_argument('run_target', metavar = 'RUN_TARGET',
58        action = 'store', type = str,
59        help = 'The code to profile.')
60
61    parser.add_argument('module_args', metavar = 'MODULE_ARGS',
62        action = 'store', type = str, nargs = '*',
63        help = 'Additional args to pass to the target if `--module` is used.')
64
65    parser.add_argument('--module', '-m', dest = 'module',
66        action = 'store_true', default = False,
67        help = 'Treat the provided run target as a module, and all other positional arguments are arguments to the module (default: %(default)s).')
68
69    parser.add_argument('--import', dest = 'imports',
70        action = 'append', default = [],
71        help = 'Import this module before running the command.')
72
73    parser.add_argument('--out', dest = 'out_path',
74        action = 'store', type = str, default = DEFAULT_OUTPUT_PATH,
75        help = 'The path where the output famegraph will be written (default: %(default)s).')
76
77    parser.add_argument('--flamegraph-path', dest = 'flamegraph_path',
78        action = 'store', type = str, default = DEFAULT_FLAMEGRAPH_PATH,
79        help = 'The path (or name in $PATH) of the flamegraph script (default: %(default)s).')
80
81    parser.add_argument('--sample-interval', dest = 'sample_interval_secs',
82        action = 'store', type = float, default = DEFAULT_SAMPLE_INTERVAL_SECS,
83        help = 'The number of seconds between samples (default: %(default)s).')
84
85    return parser
86
87if (__name__ == '__main__'):
88    sys.exit(main())