gh-138709: Implement CPU time profiling in profiling.sample (#138710)

Author: pablogsal

Include/internal/pycore_global_objects_fini_generated.h

@@ -758,6 +758,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(size)
         STRUCT_FOR_ID(sizehint)
         STRUCT_FOR_ID(skip_file_prefixes)
+        STRUCT_FOR_ID(skip_non_matching_threads)
         STRUCT_FOR_ID(sleep)
         STRUCT_FOR_ID(sock)
         STRUCT_FOR_ID(sort)

Include/internal/pycore_global_strings.h

@@ -1,5 +1,17 @@
 from abc import ABC, abstractmethod
 
+# Enums are slow
+THREAD_STATE_RUNNING = 0
+THREAD_STATE_IDLE = 1
+THREAD_STATE_GIL_WAIT = 2
+THREAD_STATE_UNKNOWN = 3
+
+STATUS = {
+    THREAD_STATE_RUNNING: "running",
+    THREAD_STATE_IDLE: "idle",
+    THREAD_STATE_GIL_WAIT: "gil_wait",
+    THREAD_STATE_UNKNOWN: "unknown",
+}
 
 class Collector(ABC):
     @abstractmethod
@@ -10,10 +22,12 @@ def collect(self, stack_frames):
     def export(self, filename):
         """Export collected data to a file."""
 
-    def _iter_all_frames(self, stack_frames):
+    def _iter_all_frames(self, stack_frames, skip_idle=False):
         """Iterate over all frame stacks from all interpreters and threads."""
         for interpreter_info in stack_frames:
             for thread_info in interpreter_info.threads:
+                if skip_idle and thread_info.status != THREAD_STATE_RUNNING:
+                    continue
                 frames = thread_info.frame_info
                 if frames:
                     yield frames

Include/internal/pycore_runtime_init_generated.h

@@ -5,7 +5,7 @@
 
 
 class PstatsCollector(Collector):
-    def __init__(self, sample_interval_usec):
+    def __init__(self, sample_interval_usec, *, skip_idle=False):
         self.result = collections.defaultdict(
             lambda: dict(total_rec_calls=0, direct_calls=0, cumulative_calls=0)
         )
@@ -14,6 +14,7 @@ def __init__(self, sample_interval_usec):
         self.callers = collections.defaultdict(
             lambda: collections.defaultdict(int)
         )
+        self.skip_idle = skip_idle
 
     def _process_frames(self, frames):
         """Process a single thread's frame stack."""
@@ -40,7 +41,7 @@ def _process_frames(self, frames):
             self.callers[callee][caller] += 1
 
     def collect(self, stack_frames):
-        for frames in self._iter_all_frames(stack_frames):
+        for frames in self._iter_all_frames(stack_frames, skip_idle=self.skip_idle):
             self._process_frames(frames)
 
     def export(self, filename):

Include/internal/pycore_unicodeobject_generated.h

@@ -15,6 +15,21 @@
 from .stack_collector import CollapsedStackCollector, FlamegraphCollector
 
 _FREE_THREADED_BUILD = sysconfig.get_config_var("Py_GIL_DISABLED") is not None
+
+# Profiling mode constants
+PROFILING_MODE_WALL = 0
+PROFILING_MODE_CPU = 1
+PROFILING_MODE_GIL = 2
+
+
+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,
+    }
+    return mode_map[mode_string]
 _HELP_DESCRIPTION = """Sample a process's stack frames and generate profiling data.
 Supports the following target modes:
   - -p PID: Profile an existing process by PID
@@ -120,18 +135,18 @@ def _run_with_sync(original_cmd):
 
 
 class SampleProfiler:
-    def __init__(self, pid, sample_interval_usec, all_threads):
+    def __init__(self, pid, sample_interval_usec, all_threads, *, mode=PROFILING_MODE_WALL):
         self.pid = pid
         self.sample_interval_usec = sample_interval_usec
         self.all_threads = all_threads
         if _FREE_THREADED_BUILD:
             self.unwinder = _remote_debugging.RemoteUnwinder(
-                self.pid, all_threads=self.all_threads
+                self.pid, all_threads=self.all_threads, mode=mode
             )
         else:
             only_active_threads = bool(self.all_threads)
             self.unwinder = _remote_debugging.RemoteUnwinder(
-                self.pid, only_active_thread=only_active_threads
+                self.pid, only_active_thread=only_active_threads, mode=mode
             )
         # Track sample intervals and total sample count
         self.sample_intervals = deque(maxlen=100)
@@ -596,21 +611,25 @@ def sample(
     show_summary=True,
     output_format="pstats",
     realtime_stats=False,
+    mode=PROFILING_MODE_WALL,
 ):
     profiler = SampleProfiler(
-        pid, sample_interval_usec, all_threads=all_threads
+        pid, sample_interval_usec, all_threads=all_threads, mode=mode
     )
     profiler.realtime_stats = realtime_stats
 
+    # Determine skip_idle for collector compatibility
+    skip_idle = mode != PROFILING_MODE_WALL
+
     collector = None
     match output_format:
         case "pstats":
-            collector = PstatsCollector(sample_interval_usec)
+            collector = PstatsCollector(sample_interval_usec, skip_idle=skip_idle)
         case "collapsed":
-            collector = CollapsedStackCollector()
+            collector = CollapsedStackCollector(skip_idle=skip_idle)
             filename = filename or f"collapsed.{pid}.txt"
         case "flamegraph":
-            collector = FlamegraphCollector()
+            collector = FlamegraphCollector(skip_idle=skip_idle)
             filename = filename or f"flamegraph.{pid}.html"
         case _:
             raise ValueError(f"Invalid output format: {output_format}")
@@ -661,6 +680,8 @@ def wait_for_process_and_sample(pid, sort_value, args):
     if not filename and args.format == "collapsed":
         filename = f"collapsed.{pid}.txt"
 
+    mode = _parse_mode(args.mode)
+
     sample(
         pid,
         sort=sort_value,
@@ -672,6 +693,7 @@ def wait_for_process_and_sample(pid, sort_value, args):
         show_summary=not args.no_summary,
         output_format=args.format,
         realtime_stats=args.realtime_stats,
+        mode=mode,
     )
 
 
@@ -726,6 +748,15 @@ def main():
         help="Print real-time sampling statistics (Hz, mean, min, max, stdev) during profiling",
     )
 
+    # Mode options
+    mode_group = parser.add_argument_group("Mode options")
+    mode_group.add_argument(
+        "--mode",
+        choices=["wall", "cpu", "gil"],
+        default="wall",
+        help="Sampling mode: wall (all threads), cpu (only CPU-running threads), gil (only GIL-holding threads)",
+    )
+
     # Output format selection
     output_group = parser.add_argument_group("Output options")
     output_format = output_group.add_mutually_exclusive_group()
@@ -850,6 +881,8 @@ def main():
     elif target_count > 1:
         parser.error("only one target type can be specified: -p/--pid, -m/--module, or script")
 
+    mode = _parse_mode(args.mode)
+
     if args.pid:
         sample(
             args.pid,
@@ -862,6 +895,7 @@ def main():
             show_summary=not args.no_summary,
             output_format=args.format,
             realtime_stats=args.realtime_stats,
+            mode=mode,
         )
     elif args.module or args.args:
         if args.module:

Lib/profiling/sampling/collector.py

@@ -11,8 +11,11 @@
 
 
 class StackTraceCollector(Collector):
-    def collect(self, stack_frames):
-        for frames in self._iter_all_frames(stack_frames):
+    def __init__(self, *, skip_idle=False):
+        self.skip_idle = skip_idle
+
+    def collect(self, stack_frames, skip_idle=False):
+        for frames in self._iter_all_frames(stack_frames, skip_idle=skip_idle):
             if not frames:
                 continue
             self.process_frames(frames)
@@ -22,7 +25,8 @@ def process_frames(self, frames):
 
 
 class CollapsedStackCollector(StackTraceCollector):
-    def __init__(self):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         self.stack_counter = collections.Counter()
 
     def process_frames(self, frames):
@@ -46,7 +50,8 @@ def export(self, filename):
 
 
 class FlamegraphCollector(StackTraceCollector):
-    def __init__(self):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         self.stats = {}
         self._root = {"samples": 0, "children": {}}
         self._total_samples = 0