Memory leak when trying to repeatedly make SOAP descriptors in a loop

[arthur-lin1027]

There seems to be a memory leak when descriptors in a long running loop - seems like python can't clear away previously created objects. While the use case in which this behavior manifests isn't necessarily a typical usecase (I'm testing out different coarse-grained configurations on the same underlying AA configuration), this behavior isn't correct.

If I run the MWE below, the memory will slowly balloon over time.

MWE:

from featomic.calculators import SoapPowerSpectrum
from ase.io import read

def soapify(frames, batchsize=1000, rcut=30, disable_progressbar=True):
    """ Splits frames into batches """
    HYPER_PARAMETERS = {
        "cutoff": {
            "radius": rcut,
            # "smoothing": {"type": "ShiftedCosine", "width": 0.5},
            "smoothing": {"type": "Step"},
        },
        "density": {
            "type": "Gaussian",
            "width": 5.0,
        },
        "basis": {
            "type": "TensorProduct",
            "max_angular": 4,
            "radial": {"type": "Gto", "max_radial": 6},
        },
    }
    calculator = SoapPowerSpectrum(**HYPER_PARAMETERS)
    soap_batches = []
    for i in tqdm(range(0, len(frames), batchsize), disable=disable_progressbar):
        desc = calculator.compute(frames[i:i+batchsize])
        desc = desc.keys_to_samples("center_type")
        desc = desc.keys_to_properties(["neighbor_1_type", "neighbor_2_type"])
        soap_batches.append(desc)
    return soap_batches

if __name__ == "__main__":
    frames = read("frames.xyz", ":")    # 10,000 frames, 175 atoms per frame.
    for i in range(500):
        desc = soapify(frames)
        ... # do a calculation with desc, but each iteration is independent from each other so desc should be overwritten and garbage collected.

[Luthaf]

Thanks for reporting this! How exactly are you checking memory growth? And on which OS are you running?

I can reproduce an increase in the max_rss with resource.getrusage, but this would also trigger on memory fragmentation. A leak detector shows nothing.

[arthur-lin1027]

Hi, I haven't used either resource.getrusage or a leak detector. I just monitor the memory usage using the top utility.

This long running loop causes a slow, menacing increase in memory usage that will cause my local workstation (Mac OS) to run out of memory (based on my workstation freezing as well as top results). This also occurs if I run on linux HPC context, where I will exceed the allocated job's allocated RAM, so this issue seems to occur regardless of OS.

RE your comment on memory fragmentation, would this occur for the repeated creation of any python object? Because if so, then it doesn't seem to be a featomic issue, and rather a python issue.

[arthur-lin1027]

Hi Guillaume, I think that within my code, there were probably some variables in my code that held onto parts of the TensorMap, and so the entire map was never able to be gc'ed. By assigning variables to copies of underlying arrays, later iterations could clear away the old TensorMaps safely.

[Luthaf]

This would explain the issue yes. Since the data is owned by the underlying C-level TensorMap, we need to keep it alive when you do something like this:

values = tensor.block(0).values
del tensor

The tensor map will only be freed when values is no longer needed. If you are only keeping small parts of the data, you can make a copy, this should also allow to free the TensorMap:

values = tensor.block(0).values[0, ...].copy()
del tensor

[arthur-lin1027]

OK great - this makes sense. Do you think that it's worth updating the metatensor or featomic documentation to encourage users to make copies if they only want particular portions of the TensorMap/TensorBlock? If so I can try to write up a PR with this explicitly mentioned.

[Luthaf]

This would make sense in metatensor's documentation, maybe the documentation for TensorBlock.values?