Referring to tables in TableCollection leaks memory

[jeromekelleher]

See analysis in #2424 for details, and e.g., #2427.

[jeromekelleher]

TODO: add some reproducible example here using mprof

[jeromekelleher]

Let's mimic the memory leak we see in ts.site(position=x) by accessing the relevant bits of the tables in a loop:

import msprime
import time

before = time.perf_counter()
ts = msprime.sim_ancestry(
    10000,
    sequence_length=1000000,
    population_size=10_000,
    recombination_rate=1e-8,
    random_seed=1234,
)
ts = msprime.sim_mutations(ts, rate=1e-7, random_seed=1)
print(ts)
duration = time.perf_counter() - before
print(f"Simulation of {ts.num_trees} trees done after {duration:.2f} seconds")

for site in ts.sites():
    # pos = ts.tables.sites.position
    tables = ts.tables
    if site.id >= 1000:
        break

First, just accessing ts.tables like here:

As it should be

[jeromekelleher]

Now, change it around so we do pos=ts.tables.sites.position:

We're using waaaaay more memory, and it looks like garbage collection isn't recovering it all.

[jeromekelleher]

Accessing ts.tables.sites gives us:

🤔

[jeromekelleher]

It seems like there must be some circular references down in where we do the allocation of the Table instances. It's all quite complicated and twisty down there though, so I'm at a bit of a loss as to what's going on here.

[molpopgen]

We've got a couple of these piling up. Relevant to tskit-dev/msprime#1946?

[jeromekelleher]

Yep, that's probably the same thing, thanks @molpopgen

[jeromekelleher]

Note: #2080 is relevant here, because we create a new instance of the high-level table class each time tables.nodes is accessed.

[jeromekelleher]

It seems this is partially at least because of how we're doing metadata decoding. If I run the code accessing ts.tables.sites above (and therefore creating a new instance of SiteTable each time because of #2080) then we get memory usage going up to 1GB like above. However, if I comment out this line here:

tskit/python/tskit/tables.py

Line 721 in c12c160

self.row_class = row_class

Usage goes down to a nice flat 60MiB. So it seems like we've created a circular reference of some sort by wrapping the row_class, and so we leak some memory every time we access one of the tables.