gh-154090: Store profiling mode and capture settings in binary profiles (GH-154105)

Author: pablogsal

Include/internal/pycore_global_objects_fini_generated.h

@@ -366,6 +366,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(canonical)
         STRUCT_FOR_ID(capath)
         STRUCT_FOR_ID(capitals)
+        STRUCT_FOR_ID(capture_features)
         STRUCT_FOR_ID(category)
         STRUCT_FOR_ID(cb_type)
         STRUCT_FOR_ID(certfile)

Include/internal/pycore_global_strings.h

@@ -84,15 +84,23 @@ with a single seek to `file_size - 32`, without first reading the header.
 |        |      |         | reserved)                              |
 |   12   |  8   | uint64  | Start timestamp (microseconds)         |
 |   20   |  8   | uint64  | Sample interval (microseconds)         |
-|   28   |  4   | uint32  | Total sample count                     |
-|   32   |  4   | uint32  | Thread count                           |
-|   36   |  8   | uint64  | String table offset                    |
-|   44   |  8   | uint64  | Frame table offset                     |
-|   52   |  4   | uint32  | Compression type (0=none, 1=zstd)      |
-|   56   |  8   | bytes   | Reserved (zero-filled)                 |
+|   28   |  8   | uint64  | Total sample count                     |
+|   36   |  4   | uint32  | Thread count                           |
+|   40   |  8   | uint64  | String table offset                    |
+|   48   |  8   | uint64  | Frame table offset                     |
+|   56   |  4   | uint32  | Compression type (0=none, 1=zstd)      |
+|   60   |  4   | uint32  | Profiling configuration bit field      |
 +--------+------+---------+----------------------------------------+
 ```
 
+The low three configuration bits contain the
+`_remote_debugging.PROFILING_MODE_*` value plus one. Zero means that the mode
+was not recorded, which is also the value in binaries written before this
+field was defined. Bit 3 indicates that capture features are known; when set,
+bits 4 through 8 respectively record `--all-threads`, `--native`, GC frames,
+`--opcodes`, and `--blocking`. Remaining bits are reserved for future capture
+features.
+
 The magic number `0x54414348` ("TACH" for Tachyon) identifies the file format
 and also serves as an **endianness marker**. When read on a system with
 different byte order than the writer, it appears as `0x48434154`. The reader
@@ -567,9 +575,10 @@ one write() call (or feeds through the compression stream).
 ## Future Considerations
 
 The optional profile-statistics block provides an extensible metadata area.
-The 16-byte checksum field in the footer is currently unused. The version
-field allows incompatible changes with graceful rejection. New compression
-types could be added (compression_type > 1).
+The Python-version field retains one reserved byte. The 16-byte checksum
+field in the footer is currently unused. The version field allows
+incompatible changes with graceful rejection. New compression types could
+be added (compression_type > 1).
 
 Any changes that alter the meaning of existing fields or the parsing logic
 should increment the version number to prevent older readers from

Include/internal/pycore_runtime_init_generated.h

@@ -10,6 +10,23 @@
 COMPRESSION_NONE = 0
 COMPRESSION_ZSTD = 1
 
+CAPTURE_FEATURES = {
+    "all_threads": 1 << 0,
+    "native": 1 << 1,
+    "gc": 1 << 2,
+    "opcodes": 1 << 3,
+    "blocking": 1 << 4,
+}
+
+
+def encode_capture_config(capture_config):
+    if capture_config is None:
+        return -1
+    return sum(
+        bit for name, bit in CAPTURE_FEATURES.items()
+        if capture_config.get(name, False)
+    )
+
 
 def _resolve_compression(compression):
     """Resolve compression type from string or int.
@@ -49,14 +66,17 @@ class BinaryCollector(Collector):
     """
 
     def __init__(self, filename, sample_interval_usec, *, skip_idle=False,
-                 compression='auto'):
+                 compression='auto', mode=None, capture_config=None):
         """Create a new binary collector.
 
         Args:
             filename: Path to output binary file
             sample_interval_usec: Sampling interval in microseconds
             skip_idle: If True, skip idle threads (not used in binary format)
             compression: 'auto', 'zstd', 'none', or int (0=none, 1=zstd)
+            mode: Profiling mode, or None if unknown
+            capture_config: Mapping of capture feature names to booleans, or
+                None if the capture configuration is unknown
         """
         self.filename = filename
         self.sample_interval_usec = sample_interval_usec
@@ -65,7 +85,10 @@ def __init__(self, filename, sample_interval_usec, *, skip_idle=False,
         compression_type = _resolve_compression(compression)
         start_time_us = int(time.monotonic() * 1_000_000)
         self._writer = _remote_debugging.BinaryWriter(
-            filename, sample_interval_usec, start_time_us, compression=compression_type
+            filename, sample_interval_usec, start_time_us,
+            compression=compression_type,
+            mode=-1 if mode is None else mode,
+            capture_features=encode_capture_config(capture_config),
         )
 
     def collect(self, stack_frames, timestamp_us=None):

Include/internal/pycore_unicodeobject_generated.h

@@ -6,6 +6,7 @@
 from .stack_collector import FlamegraphCollector, CollapsedStackCollector
 from .jsonl_collector import JsonlCollector
 from .pstats_collector import PstatsCollector
+from .binary_collector import CAPTURE_FEATURES
 
 
 class BinaryReader:
@@ -50,10 +51,21 @@ def get_info(self):
                 - string_count: Number of unique strings
                 - frame_count: Number of unique frames
                 - compression: Compression type used
+                - mode: Profiling mode, or None if not recorded
+                - capture_config: Capture feature mapping, or None if not
+                  recorded
         """
         if self._reader is None:
             raise RuntimeError("Reader not open. Use as context manager.")
-        return self._reader.get_info()
+        info = self._reader.get_info()
+        capture_features = info.pop("capture_features")
+        info["capture_config"] = (
+            None if capture_features is None else {
+                name: bool(capture_features & bit)
+                for name, bit in CAPTURE_FEATURES.items()
+            }
+        )
+        return info
 
     def replay_samples(self, collector, progress_callback=None):
         """Replay samples from binary file through a collector.
@@ -119,14 +131,16 @@ def convert_binary_to_format(input_file, output_file, output_format,
         elif output_format == 'gecko':
             collector = GeckoCollector(interval)
         elif output_format == "jsonl":
-            collector = JsonlCollector(interval)
+            collector = JsonlCollector(interval, mode=info.get("mode"))
         else:
             raise ValueError(f"Unknown output format: {output_format}")
 
         # Replay samples through collector
         count = reader.replay_samples(collector, progress_callback)
         if hasattr(collector, "set_replay_stats"):
             collector.set_replay_stats(info)
+        if hasattr(collector, "set_mode"):
+            collector.set_mode(info.get("mode"))
 
         # Export to target format
         collector.export(output_file)

InternalDocs/profiling_binary_format.md

@@ -629,9 +629,20 @@ def _sort_to_mode(sort_choice):
     }
     return sort_map.get(sort_choice, SORT_MODE_NSAMPLES)
 
+
+def _capture_config_from_args(args):
+    return {
+        "all_threads": args.all_threads,
+        "native": args.native,
+        "gc": args.gc,
+        "opcodes": args.opcodes,
+        "blocking": args.blocking,
+    }
+
+
 def _create_collector(format_type, sample_interval_usec, skip_idle, opcodes=False,
                       mode=None, output_file=None, compression='auto',
-                      diff_baseline=None):
+                      diff_baseline=None, capture_config=None):
     """Create the appropriate collector based on format type.
 
     Args:
@@ -645,6 +656,7 @@ def _create_collector(format_type, sample_interval_usec, skip_idle, opcodes=Fals
         output_file: Output file path (required for binary format)
         compression: Compression type for binary format ('auto', 'zstd', 'none')
         diff_baseline: Path to baseline binary file for differential flamegraph
+        capture_config: Capture feature mapping for binary profiles and diffs
 
     Returns:
         A collector instance of the appropriate type
@@ -661,15 +673,18 @@ def _create_collector(format_type, sample_interval_usec, skip_idle, opcodes=Fals
         return collector_class(
             sample_interval_usec,
             baseline_binary_path=diff_baseline,
-            skip_idle=skip_idle
+            skip_idle=skip_idle,
+            mode=mode,
+            capture_config=capture_config,
         )
 
     # Binary format requires output file and compression
     if format_type == "binary":
         if output_file is None:
             raise ValueError("Binary format requires an output file")
         return collector_class(output_file, sample_interval_usec, skip_idle=skip_idle,
-                              compression=compression)
+                              compression=compression, mode=mode,
+                              capture_config=capture_config)
 
     # Gecko format never skips idle (it needs both GIL and CPU data)
     # and is the only format that uses opcodes for interval markers
@@ -760,7 +775,9 @@ def _replay_with_reader(args, reader):
 
     collector = _create_collector(
         args.format, interval, skip_idle=False,
-        diff_baseline=args.diff_baseline
+        mode=info.get("mode"),
+        diff_baseline=args.diff_baseline,
+        capture_config=info.get("capture_config"),
     )
 
     def progress_callback(current, total):
@@ -778,6 +795,8 @@ def progress_callback(current, total):
     count = reader.replay_samples(collector, progress_callback)
     if hasattr(collector, "set_replay_stats"):
         collector.set_replay_stats(info)
+    if hasattr(collector, "set_mode"):
+        collector.set_mode(info.get("mode"))
     print()
 
     if args.format == "pstats":
@@ -791,7 +810,8 @@ def progress_callback(current, total):
             sort_mode = _sort_to_mode(sort_choice)
             collector.print_stats(
                 sort_mode, limit, not args.no_summary,
-                PROFILING_MODE_WALL
+                info.get("mode") if info.get("mode") is not None
+                else PROFILING_MODE_WALL
             )
     else:
         filename = (
@@ -1179,7 +1199,8 @@ def _handle_attach(args):
         args.format, args.sample_interval_usec, skip_idle, args.opcodes, mode,
         output_file=output_file,
         compression=getattr(args, 'compression', 'auto'),
-        diff_baseline=args.diff_baseline
+        diff_baseline=args.diff_baseline,
+        capture_config=_capture_config_from_args(args),
     )
 
     with _get_child_monitor_context(args, args.pid):
@@ -1286,7 +1307,8 @@ def _handle_run(args):
         args.format, args.sample_interval_usec, skip_idle, args.opcodes, mode,
         output_file=output_file,
         compression=getattr(args, 'compression', 'auto'),
-        diff_baseline=args.diff_baseline
+        diff_baseline=args.diff_baseline,
+        capture_config=_capture_config_from_args(args),
     )
 
     with _get_child_monitor_context(args, process.pid):

Lib/profiling/sampling/binary_collector.py

@@ -162,6 +162,8 @@ def set_replay_stats(self, info):
             missed_samples=info.get("missed_samples"),
             mode=self.stats.get("mode"),
         )
+    def set_mode(self, mode):
+        self.stats["mode"] = mode
 
     def export(self, filename):
         flamegraph_data = self._convert_to_flamegraph_format()
@@ -566,13 +568,16 @@ def _create_flamegraph_html(self, data):
 class DiffFlamegraphCollector(FlamegraphCollector):
     """Differential flamegraph collector that compares against a baseline binary profile."""
 
-    def __init__(self, sample_interval_usec, *, baseline_binary_path, skip_idle=False):
+    def __init__(self, sample_interval_usec, *, baseline_binary_path,
+                 skip_idle=False, mode=None, capture_config=None):
         super().__init__(sample_interval_usec, skip_idle=skip_idle)
         if not os.path.exists(baseline_binary_path):
             raise ValueError(f"Baseline file not found: {baseline_binary_path}")
         self.baseline_binary_path = baseline_binary_path
         self._baseline_collector = None
         self._elided_paths = set()
+        self.mode = mode
+        self.capture_config = capture_config
 
     def _load_baseline(self):
         """Load baseline profile from binary file."""
@@ -581,6 +586,32 @@ def _load_baseline(self):
         with BinaryReader(self.baseline_binary_path) as reader:
             info = reader.get_info()
 
+            baseline_mode = info.get("mode")
+            if (
+                baseline_mode is not None
+                and self.mode is not None
+                and baseline_mode != self.mode
+            ):
+                raise ValueError(
+                    "Baseline profiling mode does not match current mode"
+                )
+
+            baseline_config = info.get("capture_config")
+            if baseline_config is not None and self.capture_config is not None:
+                names = baseline_config.keys() | self.capture_config.keys()
+                mismatches = [
+                    name for name in names
+                    if baseline_config.get(name, False)
+                    != self.capture_config.get(name, False)
+                ]
+            else:
+                mismatches = []
+            if mismatches:
+                raise ValueError(
+                    "Baseline capture configuration does not match current "
+                    f"configuration: {', '.join(sorted(mismatches))}"
+                )
+
             baseline_collector = FlamegraphCollector(
                 sample_interval_usec=info['sample_interval_us'],
                 skip_idle=self.skip_idle