gh-154090: Preserve modes across replay outputs

Author: pablogsal

InternalDocs/profiling_binary_format.md

@@ -533,11 +533,10 @@ one write() call (or feeds through the compression stream).
 
 ## Future Considerations
 
-The format reserves space for future extensions. The 12 reserved bytes in
-the header could hold additional metadata. 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

Lib/profiling/sampling/binary_reader.py

@@ -120,7 +120,7 @@ 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}")
 

Lib/profiling/sampling/cli.py

@@ -761,6 +761,7 @@ def _replay_with_reader(args, reader):
 
     collector = _create_collector(
         args.format, interval, skip_idle=False,
+        mode=info.get("mode"),
         diff_baseline=args.diff_baseline
     )
 

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -1598,8 +1598,10 @@ def test_timestamp_preservation_with_rle(self):
 class TestBinaryReplayToJsonl(BinaryFormatTestBase):
     """Tests for binary -> JSONL replay via convert_binary_to_format."""
 
-    def _replay_to_jsonl(self, samples, interval=1000):
-        bin_path = self.create_binary_file(samples, interval=interval)
+    def _replay_to_jsonl(self, samples, interval=1000, mode=None):
+        bin_path = self.create_binary_file(
+            samples, interval=interval, mode=mode
+        )
         with tempfile.NamedTemporaryFile(suffix=".jsonl", delete=False) as f:
             jsonl_path = f.name
         self.temp_files.append(jsonl_path)
@@ -1629,6 +1631,16 @@ def test_binary_replay_to_jsonl_basic(self):
         self.assertEqual(len(frame_defs), 1)
         self.assertEqual(frame_defs[0]["line"], 99)
 
+    def test_binary_replay_to_jsonl_preserves_mode(self):
+        frame = make_frame("hot.py", 99, "hot_func")
+        records = self._replay_to_jsonl(
+            [[make_interpreter(0, [make_thread(1, [frame])])]],
+            mode=PROFILING_MODE_CPU,
+        )
+
+        meta = next(record for record in records if record["type"] == "meta")
+        self.assertEqual(meta["mode"], "cpu")
+
     def test_binary_replay_to_jsonl_rle_weight_propagation(self):
         """RLE-batched identical samples land as a single agg entry with the right total."""
         frame = make_frame("rle.py", 42, "rle_func")

Lib/test/test_profiling/test_sampling_profiler/test_cli.py

@@ -29,6 +29,7 @@
     _create_collector,
     _generate_output_filename,
     _handle_output,
+    _replay_with_reader,
     main,
 )
 from profiling.sampling.constants import (
@@ -964,6 +965,39 @@ def test_cli_replay_reader_errors_exit_cleanly(self):
             "Error: Unsupported format version 2",
         )
 
+    def test_cli_replay_propagates_recorded_mode(self):
+        reader = mock.MagicMock()
+        reader.get_info.return_value = {
+            "sample_interval_us": 1000,
+            "sample_count": 0,
+            "compression_type": 0,
+            "mode": PROFILING_MODE_CPU,
+        }
+        reader.replay_samples.return_value = 0
+        collector = mock.MagicMock()
+        collector.export.return_value = True
+        args = SimpleNamespace(
+            format="diff_flamegraph",
+            input_file="current.bin",
+            diff_baseline="baseline.bin",
+            outfile="diff.html",
+            browser=False,
+        )
+
+        with mock.patch(
+            "profiling.sampling.cli._create_collector",
+            return_value=collector,
+        ) as create_collector:
+            _replay_with_reader(args, reader)
+
+        create_collector.assert_called_once_with(
+            "diff_flamegraph",
+            1000,
+            skip_idle=False,
+            mode=PROFILING_MODE_CPU,
+            diff_baseline="baseline.bin",
+        )
+
     def test_cli_jsonl_format_mutually_exclusive_with_pstats(self):
         """--jsonl and --pstats cannot be combined (mutually exclusive group)."""
         with (