Memory leak in daemon runners?

[giovannipizzi]

I have been running ~2000 (QE Relax) WorkChains with AiiDA 1.5.0.
Now, everything is finished, verdi process list is empty, and I have zero messages also in RabbitMQ (as a double check, see below).

However, my 8 workers are still using a lot of memory:

$ verdi daemon status
Profile: dispero2020
Daemon is running as PID 5782 since 2020-11-29 17:54:56
Active workers [8]:
  PID    MEM %    CPU %  started
-----  -------  -------  -------------------
 5786    2.965      0    2020-11-29 17:54:56
 5787    2.855      0    2020-11-29 17:54:56
 5788    2.907      0.2  2020-11-29 17:54:56
 5789    2.766      0    2020-11-29 17:54:56
 5790    2.908      0.1  2020-11-29 17:54:57
 5791    2.655      0    2020-11-29 17:54:57
 5792    2.749      0    2020-11-29 17:54:57
 5793    2.866      0    2020-11-29 17:54:57

Note that this means ~1.8GB RAM/worker, so a total of ~15GB used!

I initially reported this already in #4598 but I thought it was due to the overload described there.
Instead, this time everything went smoothly with no excepted jobs.

Therefore, I am assuming that this is a memory leak, with some resources not properly released.
Considering the size of the data, this is similar to the size of the corresponding file repository. Indeed, ArrayData should still have some 'caching' of the arrays in memory, so maybe an ArrayData node might keep in memory all arrays? Maybe this is the cause?
We might want to remove that caching, but still I think the daemon should explicitly delete or remove nodes from memory once they are not used, I believe. Maybe they remain because they stay in some DB session? (This is Django)

It would be good if someone could investigate this (maybe the next task for @chrisjsewell, but also feedback from @muhrin @sphuber @unkcpz is appreciated)

For completeness: as expected, if I stop the daemon and then restart it with verdi daemon start 8, I get a low memory usage:

$ verdi daemon status
Profile: dispero2020
Daemon is running as PID 56574 since 2020-12-01 13:00:33
Active workers [8]:
  PID    MEM %    CPU %  started
-----  -------  -------  -------------------
56578    0.145        0  2020-12-01 13:00:33
56579    0.144        0  2020-12-01 13:00:33
56580    0.144        0  2020-12-01 13:00:33
56581    0.144        0  2020-12-01 13:00:33
56582    0.145        0  2020-12-01 13:00:33
56583    0.144        0  2020-12-01 13:00:33
56584    0.144        0  2020-12-01 13:00:33
56585    0.145        0  2020-12-01 13:00:33

This is to show that there are no more messages (also before restarting the daemon)

[giovannipizzi]

Question for the engine-savvy people (@sphuber @muhrin @unkcpz) - I guess when a workflow starts, submits a subprocess, and then waits, I guess the whole state is kept in memory.
Would it be possible to "unload" it from memory, and keep only a minimal "ProcessMonitor" class that has minimal memory footprint, and can reinstantiate the full node only when a new message comes back from RMQ?
In any case, reloading should already be possible easily, since this is done when the daemon is restarted?

In this way, one can always keep the memory usage very low, and therefore it becomes very easy to increase a lot the number of slots per worker without risking to overload the machine.
What do you think? How feasible is this?

[sphuber]

It would definitely be possible but it would require a significant redesign and rewrite. The current idea is that the Process instance itself is responsible for listening to RPC and broadcast messages sent over RabbitMQ and when it receives one respond accordingly. The Runner is not doing much here, all it does is when a new process is submitted it serializes its state and sends the task to RabbitMQ. The Runner that is run in a daemon process has on top of that the ProcessLauncher listener that will actually pick up these process tasks from RabbitMQ and start running them.

To implement what you suggest, we would have to change this Process centralized concept to a Runner one, where it is the runner that subscribes to RPC and broadcast messages for all the processes that it is managing. It could theoretically load the Process instance from the corresponding serialized version stored on the node in the database, and forward the call and then destroy the instance when that is done. The problem is that we are not merely relying on messages from RabbitMQ. The Process itself is also scheduling callbacks on the event loop of the runner. If the Process instance is destroyed, these callbacks clearly cannot be executed.

[giovannipizzi]

I see. Still, this leaves us with a lot of slots taken when running many workflows, as we know, and potentially a large memory usage (even if this leak is fixed - this issue is indeed in the memory left over after the processes should not be in memory anymore). There should be a LightProcess superclass that implements only the methods to listen/schedule callbacks/..., from which Process inherits, and LightProcess has a method to return the corresponding "full" Process when needed? Does it make sense?

[giovannipizzi]

I confirm that the behaviour of #4598 (AiiDA starts having DuplicateSubscribers) happens when too much memory is used, the system starts putting programs in the disk swap, and everything becomes very slow (the whole machine, and then indirectly AiiDA itself starts choking).
For me, this is somewhere between 400 (OK, with ~1200 processes) and 800 (too much) submitted QERelaxWorkchains - I think we should try to reach higher numbers of managed processes!

This is not even an issue of a memory leak (or maybe it's partially related to it), but even just of memory consumption, with each worker quickly getting to >4GB of RAM needed, easily filling my (already powerful) workstation.

I think that it would be good at least to understand why so much memory is needed by the daemons, and fix it.
Suggestions on where to look? Maybe there is some logging that could be added to check what is using so much memory?
If useful I found this post, that beside other things points to the tracemalloc python module.
If someone can help in putting the logging in the right places, I can run in my setup to get the logging, so we get some understanding of the memory usage.

[giovannipizzi]

Since I had to run a bit more, I injected some tracemalloc code at the top of my PwBaseWorkChains.
I dumped (some GB!) of data to disk, and then I checked it with something like

import tracemalloc
before = tracemalloc.Snapshot.load("tracemalloc-dispero2020-20724-20201211-174035.233858.dump")
after = tracemalloc.Snapshot.load("tracemalloc-dispero2020-20724-20201211-174757.761099.dump")

tot_stats = after.compare_to(before, 'traceback')
stat = tot_stats[0]
print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
for line in stat.traceback.format():
    print(line)



tot_stats = after.compare_to(before, 'lineno')
for stat in tot_stats[:30]:
    print(stat)

(this is at the beginning of the run, and after submitting some hundreds of WorkChains). Then I had to stop because both my hard disk got full and I ended the memory, the computer started swapping and everything got slow...

Anyway the output is this:

/usr/lib/python3.7/functools.py:60: size=195 MiB (+195 MiB), count=1932601 (+1932466), average=106 B
/usr/lib/python3.7/threading.py:238: size=144 MiB (+143 MiB), count=471893 (+471672), average=319 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py:282: size=141 MiB (+141 MiB), count=1168787 (+1168460), average=126 B
/usr/lib/python3.7/threading.py:235: size=85.0 MiB (+84.9 MiB), count=474099 (+473825), average=188 B
/usr/lib/python3.7/types.py:265: size=81.4 MiB (+81.3 MiB), count=808426 (+808154), average=106 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/ioloop.py:939: size=79.5 MiB (+79.5 MiB), count=1110803 (+1110409), average=75 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/communications.py:61: size=76.8 MiB (+76.8 MiB), count=927804 (+927552), average=87 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py:292: size=72.4 MiB (+72.3 MiB), count=249462 (+249310), average=304 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/stack_context.py:273: size=66.8 MiB (+66.8 MiB), count=603948 (+603710), average=116 B
/usr/lib/python3.7/contextlib.py:82: size=53.1 MiB (+53.0 MiB), count=171140 (+170999), average=325 B
/usr/lib/python3.7/concurrent/futures/_base.py:314: size=50.2 MiB (+50.2 MiB), count=705055 (+704785), average=75 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/concurrent.py:169: size=45.8 MiB (+45.8 MiB), count=600387 (+600158), average=80 B
/usr/lib/python3.7/types.py:275: size=45.3 MiB (+45.3 MiB), count=282860 (+282778), average=168 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/communications.py:33: size=42.9 MiB (+42.9 MiB), count=387526 (+387375), average=116 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/communications.py:63: size=25.3 MiB (+25.3 MiB), count=474063 (+473795), average=56 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/communications.py:62: size=25.3 MiB (+25.3 MiB), count=474063 (+473795), average=56 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/ioloop.py:595: size=22.7 MiB (+22.7 MiB), count=425326 (+424990), average=56 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/futures.py:95: size=21.5 MiB (+21.5 MiB), count=156367 (+156306), average=144 B
/usr/lib/python3.7/functools.py:79: size=21.2 MiB (+21.2 MiB), count=278272 (+278272), average=80 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py:985: size=21.1 MiB (+21.1 MiB), count=131977 (+131888), average=168 B
/usr/lib/python3.7/weakref.py:409: size=21.1 MiB (+21.1 MiB), count=132124 (+132011), average=167 B
/usr/lib/python3.7/threading.py:231: size=18.1 MiB (+18.1 MiB), count=237047 (+236908), average=80 B
/usr/lib/python3.7/threading.py:86: size=18.1 MiB (+18.1 MiB), count=474058 (+473790), average=40 B
/usr/lib/python3.7/threading.py:227: size=18.1 MiB (+18.1 MiB), count=236909 (+236775), average=80 B
/usr/lib/python3.7/threading.py:222: size=18.0 MiB (+18.0 MiB), count=236372 (+236274), average=80 B
/usr/lib/python3.7/functools.py:54: size=17.7 MiB (+17.7 MiB), count=231388 (+231388), average=80 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/stack_context.py:269: size=16.8 MiB (+16.8 MiB), count=573367 (+573137), average=31 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py:246: size=16.7 MiB (+16.7 MiB), count=313039 (+312915), average=56 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py:1276: size=16.2 MiB (+16.2 MiB), count=236385 (+236367), average=72 B
/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/futures.py:93: size=14.5 MiB (+14.5 MiB), count=237037 (+236903), average=64 B

The first 30 entries already account for almost all memory of this worker (>1.5GB).

If I do

    tot_stats = after.compare_to(before, 'traceback')
    stat = tot_stats[0]
    print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
    for line in stat.traceback.format():
        print(line)

for the first one, I get:

1694391 memory blocks: 180566.0 KiB
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py", line 1003
    self.run()
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py", line 1069
    yielded = self.gen.send(value)
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/kiwipy/rmq/communicator.py", line 237
    msg[messages.BroadcastMessage.CORRELATION_ID])
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py", line 292
    result = func(*args, **kwargs)
  File "/usr/lib/python3.7/types.py", line 277
    coro = func(*args, **kwargs)
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/communications.py", line 62
    task_future = futures.create_task(msg_fn, loop)
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/plumpy/futures.py", line 95
    @gen.coroutine
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py", line 246
    return _make_coroutine_wrapper(func, replace_callback=True)
  File "/home/pizzi/.virtualenvs/aiida-prod/lib/python3.7/site-packages/tornado/gen.py", line 282
    @functools.wraps(wrapped)
  File "/usr/lib/python3.7/functools.py", line 60
    getattr(wrapper, attr).update(getattr(wrapped, attr, {}))

Can anyone with a bit of knowledge of the engine (@sphuber @unkcpz @muhrin) figure out if there is something obvious where we can spare some memory? Most of the places seems to be code coming from the engine, and not the nodes of AiiDA, if I'm not wrong.

[chrisjsewell]

I injected some tracemalloc code at the top of my PwBaseWorkChains

Out of interest, could you clarify what you did here? What exactly do you mean by "top", where exactly in the code does it go?

From an uneducated glimpse, it looks like tornado introduces a significant memory overhead, when wrapping a callback into a future.
Given that tornado is no longer used come v1.6, perhaps it would be good to compare this to running with develop, to see if this is still the case with asyncio?

[sphuber]

What @chrisjsewell said about the fact that tornado is a goner in 1.6.0 and so may change things quite a bit. On top of that, it could also be worth taking a look in kiwipy. We can see in the last stack trace that the Communicator is involved. This is a single permanent instance per runner, that manages connections with RabbitMQ. When a new process is launched, broadcast and rpc subscribers are created and registered with this communicator. It might very well be that after the process is done, the subscribers are maybe not removed or if that, maybe the Communicator keeps a reference in memory, slowly accumulating dead space as processes are loaded and completed.

[giovannipizzi]

Hi both. I'm aware that tornado is gone now; I fear that (even if the code has changed) some of these "leaks" might still be in and might be also an indirect cause of the issues in #4595. I am running in production and since develop still has some issues for me (#4595) I had to continue working with 1.5, and so I started this investigation - I thought it was due to other parts of AiIDA that didn't change; this last report of mine above confirms that indeed it's better to investigate on develop.

I have some runs to finish in the next few weeks that will fill my RAM, and also have to free up some disk space before I can do more tests, so won't be able to test this super soon.
But I attach below what I did so someone else maybe can try? It should be easy, just run ~400+ workflows (but I guess you see the issue even with less - actually good to check maybe with less, to see which memory is left once all workflows are completed; I had to do my test and compare the memory only while wfs were still running, since then my machine choked) - I'm running the PwRelaxWorkchain, one could run that with low convergence cutoffs, it does not matter if the calculations are not converged. Or even use some other simple "test" workchain, it shouldn't matter.
I was running with 8 daemon workers, and I had increased at some point the number of slots to 500.

NOTE: this will create a lot of files and a possibly tens of GBs, and also the memory usage will increase by 2x or more because of tracemalloc - so maybe better to start conservative, with e.g. 50 work chains only, and make sure there is enough disk space.

What I did was to insert this code at the end of the setup step of the PwBaseWorkChain:

# Injecting memory monitoring code
import os
import datetime
import tracemalloc

from aiida.manage.configuration.settings import AIIDA_CONFIG_FOLDER
from aiida import get_profile

if not tracemalloc.is_tracing():
    tracemalloc.start(10)
snapshot = tracemalloc.take_snapshot()
snapshot_fname = f"{AIIDA_CONFIG_FOLDER}/daemon/log/tracemalloc-{get_profile().name}-{os.getpid()}-{datetime.datetime.now().strftime('%Y%m%d-%H%M%S.%f')}.dump"
snapshot.dump(snapshot_fname)                

(I'm not sure if the parameter 10 is needed in tracemalloc.start - maybe the default of 1 is OK).

As you see, is creating a dump file at each run of a new workflow, in the daemon/log folder, with an appropriate - hopefully unique - name. This also includes the PID so one can distinguish different workflows.
At the end, when everything is done, to see the final memory usage, one can just run another workchain to get one more dump and see the leftover memory.

To analyze the dumps, I use something like this.

1. I go in the folder, I do ls possibly filtering only the files of one specific worker
2. I pick the before and after file I want to compare (e.g. one more or less at the beginning, and one more or less at the end - I think the file size is in some way proportional to the allocated memory.
3. I run something like this in a python shell:

   import tracemalloc
   
   before = tracemalloc.Snapshot.load('<BEFORE_FNAME>.dump')
   after = tracemalloc.Snapshot.load('<AFTER_FNAME>.dump')
   
   top_stats = after.compare_to(before, 'lineno')
   print("[ Top 30 differences ]")
   for stat in top_stats[:30]:
       print(stat)

This gives the top 30 differences.

If one wants the traceback of e.g. the very first entry, one could do instead:

    tot_stats = after.compare_to(before, 'traceback')
    stat = tot_stats[0]
    print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
    for line in stat.traceback.format():
        print(line)

[ltalirz]

We confirm the issue of excessive memory use by the engine in AiiDA v1.5.2 (not yet on develop), running on python 3.6.9:

 verdi daemon status
Profile: reference_set_mofs
Daemon is running as PID 13422 since 2020-12-20 17:30:54
Active workers [3]:
  PID    MEM %    CPU %  started
-----  -------  -------  -------------------
32525   31.261     18.3  2020-12-21 15:56:54
32526   20.591      7.7  2020-12-21 15:56:54
32527   21.895      9.4  2020-12-21 15:56:54

That is three daemon workers together using ~74% of 32GB, i.e. 24GB of RAM!
A total crash of the machine we observed earlier was therefore, in retrospect, likely caused by excessive swapping as mentioned also by @giovannipizzi .

I attach the output of pip list
pip.txt

In this state AiiDA is essentially unusable for our high-throughput work.
It would be great to know:

1. If there are any concrete ideas for what may be causing this issue
2. What we can do to help pinpoint the issue
3. If there is an older AiiDA version that does not suffer from this issue so that people can continue doing production work

[giovannipizzi]

Some short answers:

1. no more than my report above (it seems to be in the engine, anyway - loading 1000 nodes only takes a few MB)
2. anything is appreciated. I doubt it's something that we can solve in a few days, but A) having more detailed reports of the memory usage to understand where the problem is, and B) testing earlier versions of AiiDA with a similar throughput to see if the issue was already there in AiIDA 1.0.x, 1.1.x, 1.2.x, 1.3.x, 1.4.x (maybe going with a "bisection algorithm"). I would start with the oldest version where the workflows would still work, and try to apply a similar "pressure", with similar workflows - it might well be that it always was like this; but it this is not the case then we're lucky and this can help find the commit where the problem was added
3. See my answer above (short answer - I don't know)

[atztogo]

Please ignore my comment below if I misunderstand this issue.

Probably this issue hit me. I ran with aiida-vasp and aiida-phonopy, supercell phonon (phono3py) calculation with 780 supercells. Submission of supercell vasp calculations are done at this point, https://github.com/aiida-phonopy/aiida-phonopy/blob/13efc549647f595b4f76ab455691cd95dc06d05d/aiida_phonopy/workflows/phono3py.py#L165-L198. With the latest develop, a few aiida-vasp calculations could not be retrieved (probably due to lack of memory of postgresql etc). I rerun the same phono3py calculation three times, but all similarly failed. With v1.5.2 (from git tag), memory usage is still a lot, but all the calculations went well, but maybe just lucky.

[unkcpz]

Think I might narrow down the place where this issue happening.

how to easily reproduce the issue

By running the script which modified from .ci/test_daemon.py with submitting more (16 in my case) NestedWorkChain into the running daemon. The memory usage is shown by daemon status quickly increase from 1.13% to 3.4% and will not get memory released after all the process finished. (My laptop has 8G mem, and I use only one daemon in the test)

The crux

I roll the version back to v1.2.1 the issue is not there anymore but happened again in v1.3.0. The most important changes related to engine added in v1.3.0 is 1) LoopCommunicator is introduced to hold all message transmission with RabbitMQ. 2) Reinstate broadcast subscription of workchain for child processes

The two commits are tightly related since the second depends on the first implement and the LoopCommunicator is used to realize the Runner.call_on_process_finish now also registers the broadcast subscriber which is the main purpose of the second commit.

In my test, the mem show in daemon status after running the test script:

• At f9d66b4: 1.3% (80M+)
• At 56a4cc8: 2.4% (160M+)
• At 1af57cc: 3.4% (240M+)

However, I have no idea how to fix this, manually calling the garbage collector to free the memory in the callback of Runner.call_on_process_finish seems has no effect.
The good thing is I have a repeatable test case, so will make a full-scale memory profiling in the next few days. Any head-ups and ideas are welcome before I do that. @giovannipizzi @sphuber

Sorry to bother you with this. Happy Christmas!