diff --git a/news/455.feature.rst b/news/455.feature.rst new file mode 100644 index 0000000000..a6216da624 --- /dev/null +++ b/news/455.feature.rst @@ -0,0 +1 @@ +``memray attach`` now supports ``--aggregate`` to produce :ref:`aggregated capture files `. diff --git a/src/memray/commands/attach.py b/src/memray/commands/attach.py index 529dd5fd50..927fe338ef 100644 --- a/src/memray/commands/attach.py +++ b/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) diff --git a/tests/integration/test_attach.py b/tests/integration/test_attach.py index 520d06570c..f321087a53 100644 --- a/tests/integration/test_attach.py +++ b/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", ""] + +@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", ""] + + +@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", ""] diff --git a/tests/unit/test_attach.py b/tests/unit/test_attach.py new file mode 100644 index 0000000000..7ca44e4426 --- /dev/null +++ b/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