Add --aggregate option to memray attach

Add --aggregate option which allows user to request aggregated mode for
in-memory aggregation.

Signed-off-by: Ivona Stojanovic <[email protected]>

news/455.feature.rst

@@ -0,0 +1 @@
+``memray attach`` now supports ``--aggregate`` to produce :ref:`aggregated capture files <aggregated capture files>`.

src/memray/commands/attach.py

@@ -241,6 +241,13 @@ def prepare_parser(self, parser: argparse.ArgumentParser) -> None:
             default=False,
         )
 
+        parser.add_argument(
+            "--aggregate",
+            help="Write aggregated stats to the output file instead of all allocations",
+            action="store_true",
+            default=False,
+        )
+
         parser.add_argument(
             "--native",
             help="Track native (C/C++) stack frames as well",
@@ -332,11 +339,21 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
             live_port = _get_free_port()
             destination = memray.SocketDestination(server_port=live_port)
 
+        if args.aggregate and not args.output:
+            parser.error("Can't use aggregated mode without an output file.")
+
+        file_format = (
+            "file_format=memray.FileFormat.AGGREGATED_ALLOCATIONS"
+            if args.aggregate
+            else ""
+        )
+
         tracker_call = (
             f"memray.Tracker(destination=memray.{destination!r},"
             f" native_traces={args.native},"
             f" follow_fork={args.follow_fork},"
-            f" trace_python_allocators={args.trace_python_allocators})"
+            f" trace_python_allocators={args.trace_python_allocators},"
+            f"{file_format})"
         )
 
         server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

tests/integration/test_attach.py

@@ -34,7 +34,7 @@ def bar():
 
 def baz():
     allocator = MemoryAllocator()
-    allocator.valloc(1024)
+    allocator.valloc(50 * 1024 * 1024)
     allocator.free()
 
 
@@ -52,13 +52,27 @@ def baz():
 """
 
 
-@pytest.mark.parametrize("method", ["lldb", "gdb"])
-def test_basic_attach(tmp_path, method):
-    if not debugger_available(method):
-        pytest.skip(f"a supported {method} debugger isn't installed")
+def generate_command(method, output, aggregate):
+    cmd = [
+        sys.executable,
+        "-m",
+        "memray",
+        "attach",
+        "--verbose",
+        "--force",
+        "--method",
+        method,
+        "-o",
+        str(output),
+    ]
 
-    # GIVEN
-    output = tmp_path / "test.bin"
+    if aggregate:
+        cmd.append("--aggregate")
+
+    return cmd
+
+
+def run_process(cmd):
     tracked_process = subprocess.Popen(
         [sys.executable, "-uc", PROGRAM],
         stdin=subprocess.PIPE,
@@ -71,22 +85,12 @@ def test_basic_attach(tmp_path, method):
     assert tracked_process.stdout is not None
 
     assert tracked_process.stdout.readline() == "ready\n"
-    attach_cmd = [
-        sys.executable,
-        "-m",
-        "memray",
-        "attach",
-        "--verbose",
-        "--method",
-        method,
-        "-o",
-        str(output),
-        str(tracked_process.pid),
-    ]
+
+    cmd.append(str(tracked_process.pid))
 
     # WHEN
     try:
-        subprocess.check_output(attach_cmd, stderr=subprocess.STDOUT, text=True)
+        subprocess.check_output(cmd, stderr=subprocess.STDOUT, text=True)
     except subprocess.CalledProcessError as exc:
         if "Couldn't write extended state status" in exc.output:
             # https://bugs.debian.org/cgi-bin/bugreport.cgi?bug=898048
@@ -103,14 +107,56 @@ def test_basic_attach(tmp_path, method):
     assert "" == tracked_process.stdout.read()
     assert tracked_process.returncode == 0
 
-    reader = FileReader(output)
-    records = list(reader.get_allocation_records())
-    vallocs = [
+
+def get_call_stack(allocation):
+    return [f[0] for f in allocation.stack_trace()]
+
+
+def get_relevant_vallocs(records):
+    return [
         record
         for record in filter_relevant_allocations(records)
         if record.allocator == AllocatorType.VALLOC
     ]
 
-    (valloc,) = vallocs
-    functions = [f[0] for f in valloc.stack_trace()]
-    assert functions == ["valloc", "baz", "bar", "foo", "<module>"]
+
+@pytest.mark.parametrize("method", ["lldb", "gdb"])
+def test_basic_attach(tmp_path, method):
+    if not debugger_available(method):
+        pytest.skip(f"a supported {method} debugger isn't installed")
+
+    # GIVEN
+    output = tmp_path / "test.bin"
+    attach_cmd = generate_command(method, output, aggregate=False)
+
+    # WHEN
+    run_process(attach_cmd)
+
+    # THEN
+    reader = FileReader(output)
+    (valloc,) = get_relevant_vallocs(reader.get_allocation_records())
+    assert get_call_stack(valloc) == ["valloc", "baz", "bar", "foo", "<module>"]
+
+
+@pytest.mark.parametrize("method", ["lldb", "gdb"])
+def test_aggregated_attach(tmp_path, method):
+    if not debugger_available(method):
+        pytest.skip(f"a supported {method} debugger isn't installed")
+
+    # GIVEN
+    output = tmp_path / "test.bin"
+    attach_cmd = generate_command(method, output, aggregate=True)
+
+    # WHEN
+    run_process(attach_cmd)
+
+    # THEN
+    reader = FileReader(output)
+    with pytest.raises(
+        NotImplementedError,
+        match="Can't get all allocations from a pre-aggregated capture file.",
+    ):
+        list(reader.get_allocation_records())
+
+    (valloc,) = get_relevant_vallocs(reader.get_high_watermark_allocation_records())
+    assert get_call_stack(valloc) == ["valloc", "baz", "bar", "foo", "<module>"]

tests/unit/test_attach.py

@@ -0,0 +1,21 @@
+from unittest.mock import patch
+
+import pytest
+
+from memray.commands import main
+
+
+@patch("memray.commands.attach.debugger_available")
+class TestAttachSubCommand:
+    def test_memray_attach_aggregated_without_output_file(
+        self, is_debugger_available_mock, capsys
+    ):
+        # GIVEN
+        is_debugger_available_mock.return_value = True
+
+        # WHEN
+        with pytest.raises(SystemExit):
+            main(["attach", "--aggregate", "1234"])
+
+        captured = capsys.readouterr()
+        assert "Can't use aggregated mode without an output file." in captured.err