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

@@ -746,6 +746,10 @@ cdef _create_metadata(header, peak_memory):
         PythonAllocatorType.PYTHON_ALLOCATOR_MALLOC: "malloc",
         PythonAllocatorType.PYTHON_ALLOCATOR_OTHER: "unknown",
     }
+    file_format = (
+        FileFormat.ALL_ALLOCATIONS if header["file_format"] == _FileFormat.ALL_ALLOCATIONS
+        else FileFormat.AGGREGATED_ALLOCATIONS
+    )
     return Metadata(
         start_time=millis_to_dt(stats["start_time"]),
         end_time=millis_to_dt(stats["end_time"]),
@@ -757,6 +761,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=file_format,
     )
 
 

src/memray/_metadata.py

@@ -14,3 +14,4 @@ class Metadata:
     python_allocator: str
     has_native_traces: bool
     trace_python_allocators: bool
+    file_format: int

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 = 1e7
+    if (
+        reader.metadata.file_format == FileFormat.ALL_ALLOCATIONS
+        and result_path.stat().st_size > FILE_SIZE_LIMIT
+    ):
+        pprint(
+            ":warning: [bold yellow] The file is large and may take a long time to "
+            "process [/] :warning:\n\n"
+            "Consider using the `--aggregated` option when running the memray command "
+            "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,8 @@ def write_report(
             if reader.metadata.has_native_traces:
                 warn_if_not_enough_symbols()
 
+            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
 
@@ -67,6 +68,7 @@ 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()
+            warn_if_file_is_not_aggregated_and_is_too_big(reader, result_path)
 
             if args.temporary_allocation_threshold >= 0:
                 snapshot = iter(

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,7 @@ 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()
+            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, 7, 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, 7, 0)
+            else ANY,
             call().get_leaked_allocation_records(merge_threads=merge_threads),
             call().get_memory_snapshots(),
         ]
@@ -227,6 +236,9 @@ def test_tracker_and_reporter_interactions_for_temporary_allocations(
         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, 7, 0)
+            else ANY,
             call().get_temporary_allocation_records(
                 threshold=3, merge_threads=merge_threads
             ),

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())