Tskit: Deleting tskit object in Python leaves a memory footprint

[brianzhang01]

I'm using psutil to track memory. I create a tree-sequence using msprime and then delete it using del, but I end up with a larger memory footprint than I started with. I repeat this 5 times. Code for the experiment can be found here.

Output on machine 1:

3.7.6 (default, Jan  8 2020, 19:59:22)
[GCC 7.3.0]
Linux-5.6.18-300.fc32.x86_64-x86_64-with-fedora-32-Thirty_Two
Starting memory: 53424128
Simulation with seed 1
Memory after creation: 180576256
Memory after deletion: 112984064
Simulation with seed 2
Memory after creation: 213544960
Memory after deletion: 129990656
Simulation with seed 3
Memory after creation: 225144832
Memory after deletion: 129990656
Simulation with seed 4
Memory after creation: 212770816
Memory after deletion: 130056192
Simulation with seed 5
Memory after creation: 211693568
Memory after deletion: 130056192

Output on machine 2:

3.6.8 (default, Sep  4 2019, 15:16:28)
[GCC 7.2.0]
Linux-3.10.0-957.10.1.el7.x86_64-x86_64-with-centos-7.6.1810-Core
Starting memory: 46424064
Simulation with seed 1
Memory after creation: 204124160
Memory after deletion: 91934720
Simulation with seed 2
Memory after creation: 203247616
Memory after deletion: 92504064
Simulation with seed 3
Memory after creation: 215572480
Memory after deletion: 92504064
Simulation with seed 4
Memory after creation: 203956224
Memory after deletion: 110862336
Simulation with seed 5
Memory after creation: 202358784
Memory after deletion: 109260800

On machine 1, if I create and destroy a large numpy array instead, the memory goes back to what it was at the start.

3.7.6 (default, Jan  8 2020, 19:59:22)
[GCC 7.3.0]
Linux-5.6.18-300.fc32.x86_64-x86_64-with-fedora-32-Thirty_Two
Starting memory: 53129216
Simulation with seed 1
Memory after creation: 852983808
Memory after deletion: 53166080
Simulation with seed 2
Memory after creation: 853118976
Memory after deletion: 53166080
Simulation with seed 3
Memory after creation: 853118976
Memory after deletion: 53166080
Simulation with seed 4
Memory after creation: 853118976
Memory after deletion: 53166080
Simulation with seed 5
Memory after creation: 853118976
Memory after deletion: 53166080

On machine 1, if I create a larger tree-sequence, the memory I am left with after deletion is larger than in the first case.

3.7.6 (default, Jan  8 2020, 19:59:22)
[GCC 7.3.0]
Linux-5.6.18-300.fc32.x86_64-x86_64-with-fedora-32-Thirty_Two
Starting memory: 52822016
Simulation with seed 1
Memory after creation: 323657728
Memory after deletion: 168804352
Simulation with seed 2
Memory after creation: 425463808
Memory after deletion: 217272320
Simulation with seed 3
Memory after creation: 430809088
Memory after deletion: 218316800
Simulation with seed 4
Memory after creation: 431714304
Memory after deletion: 232177664
Simulation with seed 5
Memory after creation: 423677952
Memory after deletion: 232177664

I am using tskit 0.2.2 on machine 1 and 0.2.1 on machine 2.

[molpopgen]

Have you tried forcing garbage collection for each loop iteration? Python has hard to intuit rules for when memory gets reclaimed, and you may need to force it to see an effect.

import gc

for i in x:
    # ...
    gc.collect()

[brianzhang01]

Here's the case of a smaller tree-sequence, with gc.collect() added right after the del and before the sleep

3.7.6 (default, Jan  8 2020, 19:59:22)
[GCC 7.3.0]
Linux-5.6.18-300.fc32.x86_64-x86_64-with-fedora-32-Thirty_Two
Starting memory: 53469184
Simulation with seed 1
Memory after creation: 180555776
Memory after deletion: 112963584
Simulation with seed 2
Memory after creation: 212213760
Memory after deletion: 132755456
Simulation with seed 3
Memory after creation: 223043584
Memory after deletion: 132755456
Simulation with seed 4
Memory after creation: 211750912
Memory after deletion: 132820992
Simulation with seed 5
Memory after creation: 211484672
Memory after deletion: 132820992

Here's the case for a larger tree-sequence

3.7.6 (default, Jan  8 2020, 19:59:22)
[GCC 7.3.0]
Linux-5.6.18-300.fc32.x86_64-x86_64-with-fedora-32-Thirty_Two
Starting memory: 53264384
Simulation with seed 1
Memory after creation: 324009984
Memory after deletion: 169156608
Simulation with seed 2
Memory after creation: 428179456
Memory after deletion: 219078656
Simulation with seed 3
Memory after creation: 430981120
Memory after deletion: 229109760
Simulation with seed 4
Memory after creation: 430477312
Memory after deletion: 229109760
Simulation with seed 5
Memory after creation: 426242048
Memory after deletion: 229109760

[molpopgen]

Thanks. I'm guessing that msprime is holding on to memory here in case another replicate is run. I guess it'd help to know a bit about what you're going for here?

[brianzhang01]

Sure, I'm running msprime to generate a tree sequence (moderate size, 20,000 haploids) and then using subprocess.run later on in a script (calling GCTA). I'm getting a memory crash while running GCTA and so I was trying to get back some memory by deleting the tskit object right before the GCTA call. I was surprised I wasn't getting back as much memory as I would have expected. For the time being, I've just requested more memory on my computing cluster, so hopefully that's good enough for now. GCTA is taking up the majority of the memory, it seems, followed by tskit and some numpy objects I have, which I've also deleted.

[molpopgen]

So that can be tricky to handle in Python. If a package is holding onto memory, then you need to get that package out of memory, which is hard in general. (See here.)

[jeromekelleher]

That's interesting @brianzhang01, I'm also surprised that Python is hanging on to so much memory. I am certain that there's no memory leaks in msprime/tskit C APIs, as we check that pretty thoroughly, and I'm fairly sure that the Python/C layer doesn't leak memory either (we have a script where we run the tests over and over, and watch the memory footprint).

So, it feels to me like this either a Python thing, or just the behaviour of malloc on the different machines. In general, processes are pretty reluctant to hand memory back to the OS after sbrk has been called.

Have you tried running tracemalloc on this?

[grahamgower]

I bet its a memory fragmentation issue.
https://stackoverflow.com/a/1316799/9500949

[brianzhang01]

Thanks all for the comments. I'm treating this issue less as "msprime / tskit is broken, please fix" and more as gathering information about the behaviour. As I commented earlier, I have a temporary workaround which is to just request more memory from my computing cluster. I suggest we can close the issue once we have a good enough idea of what's going on.

I will spend some time today doing some analysis on this. I plan to run tracemalloc and report the results. I also thought I would write out a few large tree sequences to disk, and then use a separate script and tskit.load to sequentially read them back in and delete each one. That way, I can isolate whether "msprime is holding on to memory" or whether it's a more fundamental tskit issue.

Are there any other suggested experiments I run? I read through the memory fragmentation post but it's not clear to me how we would go about diagnosing that as the issue. But maybe tracemalloc will give useful information.

Here's a few Google results that look interesting:

• effbot.org post on free lists, xrange vs. range
• Python 3.8 feature: Enhance tracemalloc to trace properly free lists (I'll try running tracemalloc in both Python 3.8 and previous)
• Ongoing Python bug report: Excessive memory use or memory fragmentation when unpickling many small objects
• Dropbox's engineering manager's blog post, no clear solutions
• StackOverflow post with an actual solution: setting the MALLOC_MMAP_THRESHOLD_ and MALLOC_ARENA_MAX environment variables. I hope to try this as well.

Also, I think what we are talking about is not the Python standard itself, but behaviour that arises using the CPython implementation.

[brianzhang01]

So, reporting back from investigation (sorry for the delay), things went better than expected! A simple test shows that this is indeed memory fragmentation, and better yet, that it can be fixed!

I followed the model in this StackOverflow post to try and cut down on memory fragmentation. Here is the comparison. With no command-line parameters, we have the same behaviour as before, where the process is holding on to memory:

MALLOC_MMAP_THRESHOLD_=None MALLOC_MMAP_MAX_=None MALLOC_ARENA_MAX=None
Starting memory: 54284288
Simulation with seed 1
Memory after creation: 179482624
Memory after deletion: 111890432
Simulation with seed 2
Memory after creation: 214134784
Memory after deletion: 130527232
Simulation with seed 3
Memory after creation: 224059392
Memory after deletion: 130527232
Simulation with seed 4
Memory after creation: 214118400
Memory after deletion: 130592768
Simulation with seed 5
Memory after creation: 213041152
Memory after deletion: 130592768

By instead running MALLOC_MMAP_THRESHOLD_=8192 python tskit_memory.py, now after deletion the memory goes back down to what it was at the start.

MALLOC_MMAP_THRESHOLD_=8192 MALLOC_MMAP_MAX_=None MALLOC_ARENA_MAX=None
Starting memory: 54013952
Simulation with seed 1
Memory after creation: 111599616
Memory after deletion: 54886400
Simulation with seed 2
Memory after creation: 113770496
Memory after deletion: 54886400
Simulation with seed 3
Memory after creation: 111620096
Memory after deletion: 54886400
Simulation with seed 4
Memory after creation: 112820224
Memory after deletion: 54951936
Simulation with seed 5
Memory after creation: 112136192
Memory after deletion: 54951936

The post recommended some other parameters (which I printed above), but it seems like this one parameter suffices in this case.

I said I would "write out a few large tree sequences to disk, and then use a separate script and tskit.load to sequentially read them back in and delete each one. That way, I can isolate whether "msprime is holding on to memory" or whether it's a more fundamental tskit issue." I didn't end up doing this.

I did try to do tracemalloc, printing out the top 10 allocation lines. Nothing is surprising, 9 of the top 10 hits end up being the lines from this function.

[ Top 10 ]
tskit_memory.py:27: size=1856 B, count=2, average=928 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1417: size=984 B, count=2, average=492 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1415: size=968 B, count=2, average=484 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1413: size=968 B, count=2, average=484 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1416: size=952 B, count=2, average=476 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1412: size=952 B, count=2, average=476 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1419: size=936 B, count=2, average=468 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1414: size=936 B, count=2, average=468 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/provenance.py:64: size=912 B, count=5, average=182 B
/homes/bzhang/.local/lib/python3.7/site-packages/tskit/tables.py:1418: size=904 B, count=2, average=452 B

One thing I didn't discuss earlier: I had came across a few other methods of profiling memory using pympler and resource. However pympler seems to only give the Python size of objects, whereas the method I used using psutil tracked very closely with what I observed with top on my system. resource.getrusage(resource.RUSAGE_SELF).ru_maxrss (what is used in the StackOverflow post) also records process-wide memory, but it seems to only give the maximum memory consumption across the lifetime of the process, which is not granular enough for our purposes.

A dump of my newer code is appended to the previous gist. I think we can close the issue.

[jeromekelleher]

Excellent! Thanks for the detailed report @brianzhang01, I'm sure this will be useful to other people in the future.