Add per-module allocation statistics to memray stats

[ivonastojanovic]

Issue number of the reported bug or feature request: #765

Describe your changes
Display the top N modules responsible for the most allocations in the memray stats reporter. For each allocation, walk the call stack to locate the nearest non-stdlib Python frame and extract its top-level module name. Aggregate allocation counts and allocated bytes by module, and include the top N results in both terminal and JSON output.

Testing performed
Unit tests + running memray stats on the example below:

"""Test script for memray stats module attribution."""

import json

import psutil
import rich.text
import yaml


def allocate_with_yaml():
    docs = [yaml.dump({"key": i, "nested": {"values": list(range(20))}}) for i in range(2000)]
    return [yaml.safe_load(d) for d in docs]


def allocate_with_psutil():
    procs = []
    for _ in range(100):
        procs.append(psutil.Process().memory_info())
        procs.append(psutil.Process().cpu_times())
    return procs


def allocate_with_rich():
    texts = [rich.text.Text(f"Hello world {i} " * 50) for i in range(3000)]
    return texts


def allocate_with_json():
    data = [json.dumps({"key": i, "values": list(range(50))}) for i in range(5000)]
    return [json.loads(d) for d in data]


def main():
    a = allocate_with_yaml()
    b = allocate_with_psutil()
    c = allocate_with_rich()
    d = allocate_with_json()
    print(f"Done: {len(a)} yaml, {len(b)} psutil, {len(c)} rich, {len(d)} json")


if __name__ == "__main__":
    main()

Output:

...

🥇 Top 5 allocating modules (by size):
        - yaml: 225.162MB total | 690815 allocations
        - test_modules: 37.624MB total | 43279 allocations
        - rich: 2.642MB total | 3000 allocations
        - psutil: 1.016MB total | 780 allocations
        - src: 352.635kB total | 199 allocations

🥇 Top 5 allocating modules (by number of allocations):
        - yaml: 225.162MB total | 690815 allocations
        - test_modules: 37.624MB total | 43279 allocations
        - rich: 2.642MB total | 3000 allocations
        - psutil: 1.016MB total | 780 allocations
        - src: 352.635kB total | 199 allocations

JSON Output:

...
],
"top_allocations_by_module": [
    {
      "module": "yaml",
      "num_allocations": 690815,
      "total_bytes": 225161828
    },
    {
      "module": "test_modules",
      "num_allocations": 43279,
      "total_bytes": 37623961
    },
    {
      "module": "rich",
      "num_allocations": 3000,
      "total_bytes": 2641500
    },
    {
      "module": "psutil",
      "num_allocations": 780,
      "total_bytes": 1016183
    },
    {
      "module": "src",
      "num_allocations": 199,
      "total_bytes": 352635
    }
  ],
  "top_allocations_by_module_by_count": [
    {
      "module": "yaml",
      "num_allocations": 690815,
      "total_bytes": 225161828
    },
    {
      "module": "test_modules",
      "num_allocations": 43279,
      "total_bytes": 37623961
    },
    {
      "module": "rich",
      "num_allocations": 3000,
      "total_bytes": 2641500
    },
    {
      "module": "psutil",
      "num_allocations": 780,
      "total_bytes": 1016183
    },
    {
      "module": "src",
      "num_allocations": 199,
      "total_bytes": 352635
    }
  ],
  "metadata": {
  ...

Additional context
Add any other context about your contribution here.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 98.95288% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.33%. Comparing base (4e49d28) to head (f65c72a).
⚠️ Report is 27 commits behind head on main.

Files with missing lines	Patch %	Lines
src/memray/reporters/module_tools.py	98.21%	1 Missing ⚠️
tests/integration/test_tracking.py	96.15%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #938      +/-   ##
==========================================
- Coverage   92.38%   92.33%   -0.06%     
==========================================
  Files          99      101       +2     
  Lines       11785    12748     +963     
  Branches      429      453      +24     
==========================================
+ Hits        10888    11771     +883     
- Misses        897      977      +80     

Flag	Coverage Δ
cpp	92.33% <98.95%> (-0.06%)	⬇️
python_and_cython	92.33% <98.95%> (-0.06%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[grantlouisherman]

Hey @ivonastojanovic do you have a timeline of when this might land? I started looking at #766 and I can just rebase onto this but its always easier if this is merged into main. Obviously no rush.

[godlygeek]

Do we actually need to check any of this? I think that allocation.frame_index != 0 implies AllocatorType(<int>allocation.allocator) not in (AllocatorType.FREE, AllocatorType.MUNMAP, AllocatorType.PYMALLOC_FREE, ObjectTrackingEvent.OBJECT_DESTROYED) - if we have a non-zero frame_index then we have a stack to look at, and if we have a stack to look at it must not be a deallocation, because we don't record stacks for deallocations.

And indeed, if I change this to if True: the tests all still pass - so I'm pretty sure we don't need this if at all, we just need to make sure to correctly handle the degenerate case where Py_GetStackFrame returns [] - and it seems like we do.

[ivonastojanovic]

Oh you're totally right, we don't need this check then

[ivonastojanovic]

@godlygeek Updated the tests, we can’t remove the if allocation.frame_index != 0 check because deallocations are now being counted under __main__.

[godlygeek]

Ah, actually, we do want to explicitly check for deallocations rather than empty stacks, so that allocations with no Python stack can be attributed to the tracking root as well instead of being skipped (due to having allocation.frame_index == 0)

[godlygeek]

Hm. I think that this shows that our fallback is wrong. I think if all of the allocations come from the stdlib, we should report the module as <root> or __main__ or <string> or something like that... something called the stdlib, after all. We might have trouble figuring out what to call it, but I don't think we should ignore it entirely...

[ivonastojanovic]

What do you think about __main__ then?

[godlygeek]

__main__ is misleading since it may not be from the main module. I'm gonna flip this to do <root> instead, which matches what we use for the root of our flame graphs.

[godlygeek]

This looks pretty good. I have a lot of comments, but they're almost all pretty minor stuff.

[godlygeek]

I pushed some small fixes. At this point this PR looks semantically correct to me, but it's very slow when working with large numbers of allocations, even after I attempted to add some caching to it. I'm thinking we're going to need to push some of what's currently happening in the Python layer into the Cython or C++ layer to get this to be fast enough...

Running:

memray run --trace-python-allocators --no-compress -fo output.bin -m test test_set

to create a capture file and then analyzing it with

time memray stats -f --json output.bin

takes 8 seconds on the main branch for me, and 25 seconds on this PR branch even after adding some caching, so memray stats is spending around 2/3rds of its time attributing allocations to the closest non-stdlib module with what's in this PR now.

[godlygeek]

Granted, that's a particularly stressful stress test - since the python -m test test suite is part of the stdlib, it has to walk further up the stack than might normally be expected to find non-stdlib modules to attribute allocations to.

[godlygeek]

I have a few different ideas about how to make this fast enough. Roughly from easiest to hardest:

1. Add a cache to extract_module_name_and_type so that it doesn't repeatedly compute the answer for the same file name (this is probably easiest to do by factoring extract_module_name_and_type and get_module_for_stack into methods of a class that holds the cache and the path info as member variables)
2. Expose a new RecordReader::getCallerStackId() that takes a frame index representing a call stack, and returns the frame index representing the call stack with all but the most recent frame. Then, change compute_statistics to iterate over each Python stack one frame at a time, stopping when we reach the first non-stdlib frame (saving us the work of creating Python objects representing frames we never need to look at because they're below the topmost non-stdlib frame, and more importantly allowing us to memoize the results for intermediate stacks and avoid recomputing the answer for each unique call stack that calls into stdlib code from the same initial call stack)
3. Move everything from module_tools.py into cdef functions in _memray.pyx and try to get a bit of speedup by avoiding str and pathlib.Path operations in favor of faster std::string or std::filesystem operations

Or of course, we could just accept that it makes things take 3x as long, and put it behind a flag so people need to opt into the slower but more informative behavior.

I lean towards trying both (1) and (2) though, and seeing how far those two changes get us. I'm betting that they should be enough...