Finding memory leaks in solara with Memray and objgraph

[maartenbreddels]

We noticed some memory leaks in solara, but the technique I've used to find them can be very useful in general for solara applications (or any Python program).

The first step is to isolate the issue. Let's first render a component, without any server:

import gc

import solara


@solara.component
def Page():
    for i in range(1000):
        solara.Button()


if __name__ == "__main__":
    box, rc = solara.render(Page(), handle_error=False)
    rc.close()

    del box
    del rc
    gc.collect()

We remove the reference to rc (the rendering context) and the container (box), and run the garbage collector. I didn't expect that t manually running the garbage collector was needed, but we do need it for memray 🤷, which probably means I don't fully understand the lifecycle of memray and the Python process it runs.

Now we run it under memray:

$ PYTHONMALLOC=malloc python -m memray run -f -o render1k.bin --trace-python-allocators issues/memleak.py
riting profile results into render1k.bin
[memray] Successfully generated profile results.

You can now generate reports from the stored allocation records.
Some example commands to generate reports:

/Users/maartenbreddels/miniconda3/envs/dev/bin/python -m memray flamegraph render1k.bin
# we execute the suggested command to generate the flamegraph, but add the --leaks argument 
$ memray flamegraph --leaks render1k.bin -f

See https://bloomberg.github.io/memray/flamegraph.html for the --leaks argument.

This gives us:

Where we see on the right that indeed, solara is not releasing some memory, zooming in we see this is using 11Mb of memory:

My guess here is that the rc (render context), has an extra reference to it. To check this, we use objgraph

import gc

import solara


@solara.component
def Page():
    for i in range(1000):
        solara.Button()


if __name__ == "__main__":
    import objgraph

    import solara

    box, rc = solara.render(Page(), handle_error=False)
    rc.close()
    objgraph.show_backrefs([rc], filename="solara-mem-leak.png")

    del rc
    del box
    gc.collect()

We run the program (with or without memray, doesn't matter), and inspect solara-mem-leak.png:

On the right, we see many internal references, but on the left we see the rc reference from our script and a reference from a thread-local object!.

Fixing this in Reacton (the underlying library of solara) in widgetti/reacton@8ebc3a4 we see no more significant memory usage due to solara:

Hope this helps people debug memory leak issues! Ideas or improvements on this workflow are welcome 🙏