gh-138385: Sample all interpreters in the tachyon profiler (#138398)

Author: pablogsal

Lib/profiling/sampling/collector.py

@@ -9,3 +9,11 @@ def collect(self, stack_frames):
     @abstractmethod
     def export(self, filename):
         """Export collected data to a file."""
+
+    def _iter_all_frames(self, stack_frames):
+        """Iterate over all frame stacks from all interpreters and threads."""
+        for interpreter_info in stack_frames:
+            for thread_info in interpreter_info.threads:
+                frames = thread_info.frame_info
+                if frames:
+                    yield frames

Lib/profiling/sampling/pstats_collector.py

@@ -15,43 +15,33 @@ def __init__(self, sample_interval_usec):
             lambda: collections.defaultdict(int)
         )
 
-    def collect(self, stack_frames):
-        for thread_id, frames in stack_frames:
-            if not frames:
-                continue
+    def _process_frames(self, frames):
+        """Process a single thread's frame stack."""
+        if not frames:
+            return
 
-            # Process each frame in the stack to track cumulative calls
-            for frame in frames:
-                location = (frame.filename, frame.lineno, frame.funcname)
-                self.result[location]["cumulative_calls"] += 1
+        # Process each frame in the stack to track cumulative calls
+        for frame in frames:
+            location = (frame.filename, frame.lineno, frame.funcname)
+            self.result[location]["cumulative_calls"] += 1
 
-            # The top frame gets counted as an inline call (directly executing)
-            top_frame = frames[0]
-            top_location = (
-                top_frame.filename,
-                top_frame.lineno,
-                top_frame.funcname,
-            )
+        # The top frame gets counted as an inline call (directly executing)
+        top_location = (frames[0].filename, frames[0].lineno, frames[0].funcname)
+        self.result[top_location]["direct_calls"] += 1
 
-            self.result[top_location]["direct_calls"] += 1
+        # Track caller-callee relationships for call graph
+        for i in range(1, len(frames)):
+            callee_frame = frames[i - 1]
+            caller_frame = frames[i]
 
-            # Track caller-callee relationships for call graph
-            for i in range(1, len(frames)):
-                callee_frame = frames[i - 1]
-                caller_frame = frames[i]
+            callee = (callee_frame.filename, callee_frame.lineno, callee_frame.funcname)
+            caller = (caller_frame.filename, caller_frame.lineno, caller_frame.funcname)
 
-                callee = (
-                    callee_frame.filename,
-                    callee_frame.lineno,
-                    callee_frame.funcname,
-                )
-                caller = (
-                    caller_frame.filename,
-                    caller_frame.lineno,
-                    caller_frame.funcname,
-                )
+            self.callers[callee][caller] += 1
 
-                self.callers[callee][caller] += 1
+    def collect(self, stack_frames):
+        for frames in self._iter_all_frames(stack_frames):
+            self._process_frames(frames)
 
     def export(self, filename):
         self.create_stats()

Lib/profiling/sampling/stack_collector.py

@@ -9,16 +9,22 @@ def __init__(self):
         self.call_trees = []
         self.function_samples = collections.defaultdict(int)
 
+    def _process_frames(self, frames):
+        """Process a single thread's frame stack."""
+        if not frames:
+            return
+
+        # Store the complete call stack (reverse order - root first)
+        call_tree = list(reversed(frames))
+        self.call_trees.append(call_tree)
+
+        # Count samples per function
+        for frame in frames:
+            self.function_samples[frame] += 1
+
     def collect(self, stack_frames):
-        for thread_id, frames in stack_frames:
-            if frames:
-                # Store the complete call stack (reverse order - root first)
-                call_tree = list(reversed(frames))
-                self.call_trees.append(call_tree)
-
-                # Count samples per function
-                for frame in frames:
-                    self.function_samples[frame] += 1
+        for frames in self._iter_all_frames(stack_frames):
+            self._process_frames(frames)
 
 
 class CollapsedStackCollector(StackTraceCollector):