Reducing memory usage

[ddundo]

@stephankramer @jwallwork23 could you please take a look at the simple example below? I tried to put it together to show the memory usage by the process before and after adapting the mesh. It really piles up when we have multiple meshes.

Edit: I cleaned up the example. Based on this I think there really is a memory leak. I am investigating but I've never used memory profilers before so it might take a bit. I will bring it up on the Friday meeting.

import gc, os, psutil
from firedrake import *
from animate.metric import RiemannianMetric

process = psutil.Process(os.getpid())
def mem(i): print(f"Memory usage {i}: {process.memory_full_info().uss / 1024**2:.0f} MB")

def get_metric():
    mesh = UnitSquareMesh(1000, 100)
    x, y = SpatialCoordinate(mesh)
    f = Function(FunctionSpace(mesh, "CG", 1)).interpolate(cos(x*y))

    P1_ten = TensorFunctionSpace(mesh, "CG", 1)
    metric = RiemannianMetric(P1_ten)
    metric.set_parameters({"dm_plex_metric_target_complexity": 10000})
    metric.compute_hessian(f)
    metric.normalise()
    
    return metric

mem(0)
metric = get_metric()
mem(1)

gc.collect()
mem(2)

del metric
gc.collect()
mem(3)

This prints:

Memory usage 0: 117 MB
Memory usage 1: 790 MB
Memory usage 2: 508 MB
Memory usage 3: 508 MB

[ddundo]

My workaround so far: use subprocess to adapt and checkpoint adapted_mesh. Then load it in the original script. This works well but it's a pain. But it might help the investigation of what we can get rid of and get this sorted properly.

Checkpointing the adapted_mesh above and then loading it in a separate script:

import os
import psutil
from firedrake import *

process = psutil.Process(os.getpid())

ram_usage = lambda: print("RAM usage", process.memory_info().rss / 1024**2)

ram_usage().  # 201 MB
with CheckpointFile(..., "r") as afile:
    adapted_mesh = afile.load_mesh("adapted_mesh")
ram_usage().  # 227 MB

[ddundo]

And an even smaller example, with just firedrake:

import gc, os, psutil
from firedrake import *

process = psutil.Process(os.getpid())
def mem(i): print(f"Memory usage {i}: {process.memory_full_info().uss / 1024**2:.0f} MB")

mem(0)

mesh = UnitSquareMesh(1000, 100)
x, y = SpatialCoordinate(mesh)
f = Function(FunctionSpace(mesh, "CG", 1)).interpolate(x*y)
mem(1)

gc.collect()
mem(2)

del f, x, y, mesh
gc.collect()
mem(3)

Which prints

Memory usage 0: 117 MB
Memory usage 1: 210 MB
Memory usage 2: 210 MB
Memory usage 3: 210 MB

And with a finer mesh UnitSquareMesh(1000, 1000) I get

Memory usage 0: 117 MB
Memory usage 1: 1024 MB
Memory usage 2: 1024 MB
Memory usage 3: 181 MB

which is completely unintuitive to me.