|
2 | 2 | import _remote_debugging
|
3 | 3 | import os
|
4 | 4 | import pstats
|
| 5 | +import subprocess |
5 | 6 | import statistics
|
6 | 7 | import sys
|
7 | 8 | import sysconfig
|
@@ -542,46 +543,66 @@ def main():
|
542 | 543 | parser = argparse.ArgumentParser(
|
543 | 544 | description=(
|
544 | 545 | "Sample a process's stack frames and generate profiling data.\n"
|
545 |
| - "Supports two output formats:\n" |
546 |
| - " - pstats: Detailed profiling statistics with sorting options\n" |
547 |
| - " - collapsed: Stack traces for generating flamegraphs\n" |
| 546 | + "Supports the following target modes:\n" |
| 547 | + " - -p PID: Profile an existing process by PID\n" |
| 548 | + " - -m MODULE [ARGS...]: Profile a module as python -m module ... \n" |
| 549 | + " - filename [ARGS...]: Profile the specified script by running it in a subprocess\n" |
548 | 550 | "\n"
|
549 | 551 | "Examples:\n"
|
550 | 552 | " # Profile process 1234 for 10 seconds with default settings\n"
|
551 |
| - " python -m profile.sample 1234\n" |
| 553 | + " python -m profile.sample -p 1234\n" |
| 554 | + "\n" |
| 555 | + " # Profile a script by running it in a subprocess\n" |
| 556 | + " python -m profile.sample myscript.py arg1 arg2\n" |
| 557 | + "\n" |
| 558 | + " # Profile a module by running it as python -m module in a subprocess\n" |
| 559 | + " python -m profile.sample -m mymodule arg1 arg2\n" |
552 | 560 | "\n"
|
553 | 561 | " # Profile with custom interval and duration, save to file\n"
|
554 |
| - " python -m profile.sample -i 50 -d 30 -o profile.stats 1234\n" |
| 562 | + " python -m profile.sample -i 50 -d 30 -o profile.stats -p 1234\n" |
555 | 563 | "\n"
|
556 | 564 | " # Generate collapsed stacks for flamegraph\n"
|
557 |
| - " python -m profile.sample --collapsed 1234\n" |
| 565 | + " python -m profile.sample --collapsed -p 1234\n" |
558 | 566 | "\n"
|
559 | 567 | " # Profile all threads, sort by total time\n"
|
560 |
| - " python -m profile.sample -a --sort-tottime 1234\n" |
| 568 | + " python -m profile.sample -a --sort-tottime -p 1234\n" |
561 | 569 | "\n"
|
562 | 570 | " # Profile for 1 minute with 1ms sampling interval\n"
|
563 |
| - " python -m profile.sample -i 1000 -d 60 1234\n" |
| 571 | + " python -m profile.sample -i 1000 -d 60 -p 1234\n" |
564 | 572 | "\n"
|
565 | 573 | " # Show only top 20 functions sorted by direct samples\n"
|
566 |
| - " python -m profile.sample --sort-nsamples -l 20 1234\n" |
| 574 | + " python -m profile.sample --sort-nsamples -l 20 -p 1234\n" |
567 | 575 | "\n"
|
568 | 576 | " # Profile all threads and save collapsed stacks\n"
|
569 |
| - " python -m profile.sample -a --collapsed -o stacks.txt 1234\n" |
| 577 | + " python -m profile.sample -a --collapsed -o stacks.txt -p 1234\n" |
570 | 578 | "\n"
|
571 | 579 | " # Profile with real-time sampling statistics\n"
|
572 |
| - " python -m profile.sample --realtime-stats 1234\n" |
| 580 | + " python -m profile.sample --realtime-stats -p 1234\n" |
573 | 581 | "\n"
|
574 | 582 | " # Sort by sample percentage to find most sampled functions\n"
|
575 |
| - " python -m profile.sample --sort-sample-pct 1234\n" |
| 583 | + " python -m profile.sample --sort-sample-pct -p 1234\n" |
576 | 584 | "\n"
|
577 | 585 | " # Sort by cumulative samples to find functions most on call stack\n"
|
578 |
| - " python -m profile.sample --sort-nsamples-cumul 1234" |
| 586 | + " python -m profile.sample --sort-nsamples-cumul -p 1234\n" |
579 | 587 | ),
|
580 | 588 | formatter_class=argparse.RawDescriptionHelpFormatter,
|
581 | 589 | )
|
582 | 590 |
|
583 |
| - # Required arguments |
584 |
| - parser.add_argument("pid", type=int, help="Process ID to sample") |
| 591 | + # Target selection |
| 592 | + target_group = parser.add_mutually_exclusive_group(required=True) |
| 593 | + target_group.add_argument( |
| 594 | + "-p", "--pid", type=int, help="Process ID to sample" |
| 595 | + ) |
| 596 | + target_group.add_argument( |
| 597 | + "-m", "--module", |
| 598 | + nargs=argparse.REMAINDER, |
| 599 | + help="Run and profile a module as python -m module [ARGS...]" |
| 600 | + ) |
| 601 | + target_group.add_argument( |
| 602 | + "script", |
| 603 | + nargs=argparse.REMAINDER, |
| 604 | + help="Script to run and profile, with optional arguments" |
| 605 | + ) |
585 | 606 |
|
586 | 607 | # Sampling options
|
587 | 608 | sampling_group = parser.add_argument_group("Sampling configuration")
|
@@ -712,19 +733,59 @@ def main():
|
712 | 733 |
|
713 | 734 | sort_value = args.sort if args.sort is not None else 2
|
714 | 735 |
|
715 |
| - sample( |
716 |
| - args.pid, |
717 |
| - sample_interval_usec=args.interval, |
718 |
| - duration_sec=args.duration, |
719 |
| - filename=args.outfile, |
720 |
| - all_threads=args.all_threads, |
721 |
| - limit=args.limit, |
722 |
| - sort=sort_value, |
723 |
| - show_summary=not args.no_summary, |
724 |
| - output_format=args.format, |
725 |
| - realtime_stats=args.realtime_stats, |
726 |
| - ) |
| 736 | + if not(args.pid or args.module or args.script): |
| 737 | + parser.error( |
| 738 | + "You must specify either a process ID (-p), a module (-m), or a script to run." |
| 739 | + ) |
727 | 740 |
|
| 741 | + if args.pid: |
| 742 | + sample( |
| 743 | + args.pid, |
| 744 | + sample_interval_usec=args.interval, |
| 745 | + duration_sec=args.duration, |
| 746 | + filename=args.outfile, |
| 747 | + all_threads=args.all_threads, |
| 748 | + limit=args.limit, |
| 749 | + sort=sort_value, |
| 750 | + show_summary=not args.no_summary, |
| 751 | + output_format=args.format, |
| 752 | + realtime_stats=args.realtime_stats, |
| 753 | + ) |
| 754 | + elif args.module or args.script: |
| 755 | + if args.module: |
| 756 | + cmd = [sys.executable, "-m", *args.module] |
| 757 | + else: |
| 758 | + cmd = [sys.executable, *args.script] |
| 759 | + |
| 760 | + process = subprocess.Popen(cmd) |
| 761 | + |
| 762 | + try: |
| 763 | + exit_code = process.wait(timeout=0.1) |
| 764 | + sys.exit(exit_code) |
| 765 | + except subprocess.TimeoutExpired: |
| 766 | + pass |
| 767 | + |
| 768 | + try: |
| 769 | + sample( |
| 770 | + process.pid, |
| 771 | + sort=sort_value, |
| 772 | + sample_interval_usec=args.interval, |
| 773 | + duration_sec=args.duration, |
| 774 | + filename=args.outfile, |
| 775 | + all_threads=args.all_threads, |
| 776 | + limit=args.limit, |
| 777 | + show_summary=not args.no_summary, |
| 778 | + output_format=args.format, |
| 779 | + realtime_stats=args.realtime_stats, |
| 780 | + ) |
| 781 | + finally: |
| 782 | + if process.poll() is None: |
| 783 | + process.terminate() |
| 784 | + try: |
| 785 | + process.wait(timeout=2) |
| 786 | + except subprocess.TimeoutExpired: |
| 787 | + process.kill() |
| 788 | + process.wait() |
728 | 789 |
|
729 | 790 | if __name__ == "__main__":
|
730 | 791 | main()
|
0 commit comments