Add a warning to promote --aggregated if files are too big

Users tend to run memray over long running programs but they don't tend
to know about the --aggregated option, while most of the cases they
don't need the extra features that runnin with the default mode has.

Signed-off-by: Pablo Galindo <[email protected]>

Author: pablogsal

src/memray/_memray.pyx

@@ -757,6 +757,7 @@ cdef _create_metadata(header, peak_memory):
         python_allocator=allocator_id_to_name[header["python_allocator"]],
         has_native_traces=header["native_traces"],
         trace_python_allocators=header["trace_python_allocators"],
+        file_format=FileFormat(header["file_format"]),
     )
 
 

src/memray/_metadata.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from datetime import datetime
 
+from ._memray import FileFormat
+
 
 @dataclass
 class Metadata:
@@ -14,3 +16,4 @@ class Metadata:
     python_allocator: str
     has_native_traces: bool
     trace_python_allocators: bool
+    file_format: FileFormat

src/memray/commands/common.py

@@ -20,6 +20,7 @@
 from memray import FileReader
 from memray import MemorySnapshot
 from memray._errors import MemrayCommandError
+from memray._memray import FileFormat
 from memray._memray import SymbolicSupport
 from memray._memray import TemporalAllocationRecord
 from memray._memray import get_symbolic_support
@@ -77,6 +78,24 @@ def warn_if_not_enough_symbols() -> None:
         return
 
 
+def warn_if_file_is_not_aggregated_and_is_too_big(
+    reader: FileReader, result_path: Path
+) -> None:
+    FILE_SIZE_LIMIT = 10 * 1000 * 1000
+    if (
+        reader.metadata.file_format == FileFormat.ALL_ALLOCATIONS
+        and result_path.stat().st_size > FILE_SIZE_LIMIT
+    ):
+        pprint(
+            ":warning: [bold yellow] This capture file is large and may take a long"
+            " time to process [/] :warning:\n\n"
+            "Next time, consider using the `--aggregate` option to `memray run` to"
+            " reduce the size of the file.\n"
+            "Check https://bloomberg.github.io/memray/run.html#aggregated-capture-files"
+            " for more information.\n"
+        )
+
+
 class HighWatermarkCommand:
     def __init__(
         self,
@@ -140,6 +159,9 @@ def write_report(
             if reader.metadata.has_native_traces:
                 warn_if_not_enough_symbols()
 
+            if not temporal and not temporary_allocation_threshold >= 0:
+                warn_if_file_is_not_aggregated_and_is_too_big(reader, result_path)
+
             if temporal:
                 assert self.temporal_reporter_factory is not None
                 if show_memory_leaks:

src/memray/commands/summary.py

@@ -5,6 +5,7 @@
 
 from memray import FileReader
 from memray._errors import MemrayCommandError
+from memray.commands.common import warn_if_file_is_not_aggregated_and_is_too_big
 from memray.commands.common import warn_if_not_enough_symbols
 from memray.reporters.summary import SummaryReporter
 
@@ -68,6 +69,9 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
             if reader.metadata.has_native_traces:
                 warn_if_not_enough_symbols()
 
+            if not args.temporary_allocation_threshold >= 0:
+                warn_if_file_is_not_aggregated_and_is_too_big(reader, result_path)
+
             if args.temporary_allocation_threshold >= 0:
                 snapshot = iter(
                     reader.get_temporary_allocation_records(

src/memray/commands/tree.py

@@ -5,6 +5,7 @@
 
 from memray import FileReader
 from memray._errors import MemrayCommandError
+from memray.commands.common import warn_if_file_is_not_aggregated_and_is_too_big
 from memray.commands.common import warn_if_not_enough_symbols
 from memray.reporters.tree import TreeReporter
 
@@ -56,6 +57,10 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
             reader = FileReader(os.fspath(args.results), report_progress=True)
             if reader.metadata.has_native_traces:
                 warn_if_not_enough_symbols()
+
+            if not args.temporary_allocation_threshold >= 0:
+                warn_if_file_is_not_aggregated_and_is_too_big(reader, result_path)
+
             if args.temporary_allocation_threshold >= 0:
                 snapshot = iter(
                     reader.get_temporary_allocation_records(

tests/unit/test_highwatermark_command.py

@@ -1,5 +1,7 @@
 import os
+import sys
 from pathlib import Path
+from unittest.mock import ANY
 from unittest.mock import MagicMock
 from unittest.mock import Mock
 from unittest.mock import call
@@ -8,6 +10,7 @@
 import pytest
 
 from memray._errors import MemrayCommandError
+from memray._memray import FileFormat
 from memray.commands.common import HighWatermarkCommand
 
 
@@ -165,6 +168,9 @@ def test_tracker_and_reporter_interactions_for_peak(self, tmp_path, merge_thread
         calls = [
             call(os.fspath(result_path), report_progress=True),
             call().metadata.has_native_traces.__bool__(),
+            call().metadata.file_format.__eq__(FileFormat.ALL_ALLOCATIONS)
+            if sys.version_info >= (3, 8, 0)
+            else ANY,
             call().get_high_watermark_allocation_records(merge_threads=merge_threads),
             call().get_memory_snapshots(),
         ]
@@ -195,6 +201,9 @@ def test_tracker_and_reporter_interactions_for_leak(self, tmp_path, merge_thread
         calls = [
             call(os.fspath(result_path), report_progress=True),
             call().metadata.has_native_traces.__bool__(),
+            call().metadata.file_format.__eq__(FileFormat.ALL_ALLOCATIONS)
+            if sys.version_info >= (3, 8, 0)
+            else ANY,
             call().get_leaked_allocation_records(merge_threads=merge_threads),
             call().get_memory_snapshots(),
         ]

tests/unit/test_stats_reporter.py

@@ -9,6 +9,7 @@
 import pytest
 
 from memray import AllocatorType as AT
+from memray._memray import FileFormat
 from memray._metadata import Metadata
 from memray._stats import Stats
 from memray.reporters.stats import StatsReporter
@@ -97,6 +98,7 @@ def fake_stats():
             python_allocator="pymalloc",
             has_native_traces=False,
             trace_python_allocators=True,
+            file_format=FileFormat.ALL_ALLOCATIONS,
         ),
         total_num_allocations=20,
         total_memory_allocated=sum(mem_allocation_list),
@@ -435,6 +437,7 @@ def test_stats_output_json(fake_stats, tmp_path):
             "python_allocator": "pymalloc",
             "has_native_traces": False,
             "trace_python_allocators": True,
+            "file_format": 0,
         },
     }
     actual = json.loads(output_file.read_text())