Profile a module or script with sampling profiler

Add `-m` and `filename` arguments to the sampling profiler to launch the
specified Python program in a subprocess and start profiling it.
Previously only a PID was accepted, this can now be done by passing
`-p PID`.

Author: lkollar

Lib/profile/sample.py

@@ -2,6 +2,7 @@
 import _remote_debugging
 import os
 import pstats
+import subprocess
 import statistics
 import sys
 import sysconfig
@@ -542,46 +543,66 @@ def main():
     parser = argparse.ArgumentParser(
         description=(
             "Sample a process's stack frames and generate profiling data.\n"
-            "Supports two output formats:\n"
-            "  - pstats: Detailed profiling statistics with sorting options\n"
-            "  - collapsed: Stack traces for generating flamegraphs\n"
+            "Supports the following target modes:\n"
+            "  - -p PID: Profile an existing process by PID\n"
+            "  - -m MODULE [ARGS...]: Profile a module as python -m module ... \n"
+            "  - filename [ARGS...]: Profile the specified script by running it in a subprocess\n"
             "\n"
             "Examples:\n"
             "  # Profile process 1234 for 10 seconds with default settings\n"
-            "  python -m profile.sample 1234\n"
+            "  python -m profile.sample -p 1234\n"
+            "\n"
+            "  # Profile a script by running it in a subprocess\n"
+            "  python -m profile.sample myscript.py arg1 arg2\n"
+            "\n"
+            "  # Profile a module by running it as python -m module in a subprocess\n"
+            "  python -m profile.sample -m mymodule arg1 arg2\n"
             "\n"
             "  # Profile with custom interval and duration, save to file\n"
-            "  python -m profile.sample -i 50 -d 30 -o profile.stats 1234\n"
+            "  python -m profile.sample -i 50 -d 30 -o profile.stats -p 1234\n"
             "\n"
             "  # Generate collapsed stacks for flamegraph\n"
-            "  python -m profile.sample --collapsed 1234\n"
+            "  python -m profile.sample --collapsed -p 1234\n"
             "\n"
             "  # Profile all threads, sort by total time\n"
-            "  python -m profile.sample -a --sort-tottime 1234\n"
+            "  python -m profile.sample -a --sort-tottime -p 1234\n"
             "\n"
             "  # Profile for 1 minute with 1ms sampling interval\n"
-            "  python -m profile.sample -i 1000 -d 60 1234\n"
+            "  python -m profile.sample -i 1000 -d 60 -p 1234\n"
             "\n"
             "  # Show only top 20 functions sorted by direct samples\n"
-            "  python -m profile.sample --sort-nsamples -l 20 1234\n"
+            "  python -m profile.sample --sort-nsamples -l 20 -p 1234\n"
             "\n"
             "  # Profile all threads and save collapsed stacks\n"
-            "  python -m profile.sample -a --collapsed -o stacks.txt 1234\n"
+            "  python -m profile.sample -a --collapsed -o stacks.txt -p 1234\n"
             "\n"
             "  # Profile with real-time sampling statistics\n"
-            "  python -m profile.sample --realtime-stats 1234\n"
+            "  python -m profile.sample --realtime-stats -p 1234\n"
             "\n"
             "  # Sort by sample percentage to find most sampled functions\n"
-            "  python -m profile.sample --sort-sample-pct 1234\n"
+            "  python -m profile.sample --sort-sample-pct -p 1234\n"
             "\n"
             "  # Sort by cumulative samples to find functions most on call stack\n"
-            "  python -m profile.sample --sort-nsamples-cumul 1234"
+            "  python -m profile.sample --sort-nsamples-cumul -p 1234\n"
         ),
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
 
-    # Required arguments
-    parser.add_argument("pid", type=int, help="Process ID to sample")
+    # Target selection
+    target_group = parser.add_mutually_exclusive_group(required=True)
+    target_group.add_argument(
+        "-p", "--pid", type=int, help="Process ID to sample"
+    )
+    target_group.add_argument(
+        "-m", "--module",
+        nargs=argparse.REMAINDER,
+        help="Run and profile a module as python -m module [ARGS...]"
+    )
+    target_group.add_argument(
+        "script",
+        nargs=argparse.REMAINDER,
+        help="Script to run and profile, with optional arguments"
+    )
 
     # Sampling options
     sampling_group = parser.add_argument_group("Sampling configuration")
@@ -712,19 +733,59 @@ def main():
 
     sort_value = args.sort if args.sort is not None else 2
 
-    sample(
-        args.pid,
-        sample_interval_usec=args.interval,
-        duration_sec=args.duration,
-        filename=args.outfile,
-        all_threads=args.all_threads,
-        limit=args.limit,
-        sort=sort_value,
-        show_summary=not args.no_summary,
-        output_format=args.format,
-        realtime_stats=args.realtime_stats,
-    )
+    if not(args.pid or args.module or args.script):
+        parser.error(
+            "You must specify either a process ID (-p), a module (-m), or a script to run."
+        )
 
+    if args.pid:
+        sample(
+            args.pid,
+            sample_interval_usec=args.interval,
+            duration_sec=args.duration,
+            filename=args.outfile,
+            all_threads=args.all_threads,
+            limit=args.limit,
+            sort=sort_value,
+            show_summary=not args.no_summary,
+            output_format=args.format,
+            realtime_stats=args.realtime_stats,
+        )
+    elif args.module or args.script:
+        if args.module:
+            cmd = [sys.executable, "-m", *args.module]
+        else:
+            cmd = [sys.executable, *args.script]
+
+        process = subprocess.Popen(cmd)
+
+        try:
+            exit_code = process.wait(timeout=0.1)
+            sys.exit(exit_code)
+        except subprocess.TimeoutExpired:
+            pass
+
+        try:
+            sample(
+                process.pid,
+                sort=sort_value,
+                sample_interval_usec=args.interval,
+                duration_sec=args.duration,
+                filename=args.outfile,
+                all_threads=args.all_threads,
+                limit=args.limit,
+                show_summary=not args.no_summary,
+                output_format=args.format,
+                realtime_stats=args.realtime_stats,
+            )
+        finally:
+            if process.poll() is None:
+                process.terminate()
+                try:
+                    process.wait(timeout=2)
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                    process.wait()
 
 if __name__ == "__main__":
     main()