Why does manipulating xarrays inside a function takes up so much memory permanently?

[mchoblet]

I'm new at using Xarray (using it inside jupyter notebooks). I love it and up to now everything has worked like a charm, except when I started to look at how much RAM is used by my functions (e.g. htop), which is confusing me and I haven't found anything related on the web..

I have a nc file with monthly data from which I want to compute yearly means, taking into account month lengths, masking nan values and also using specific months only, which requires the use of groupby and resample.
As I can see from using the memory profiler these operations temporarily take up ~15gm RAM, which as such is not a problem because I have 64gb RAM at hand.
Nonetheless it seems like some memory is blocked permanently, even though I call these methods inside a function. For the function below it blocks ~4gb of memory although the resulting xarray only has a size of ~440mb (55*10**6 float 64entries), with more complex operations more memory is permanently blocked (only freed when I restart my Kernel).
Explicitely using del , gc.collect() or Dataarray.close() inside the function did not change anything.

A basic function to compute a weighted yearly mean from monthly data looks like this:

import xarray as xr
#some dataarray for precipitation
test=xr.open_dataset(path)['pr']

def weighted_temporal_mean(ds):
    """
    Taken  from https://ncar.github.io/esds/posts/2021/yearly-averages-xarray/
    Compute yearly average from monthly data taking into account month length and 
    masking nan values
    """
    # Determine the month length
    month_length = ds.time.dt.days_in_month

    # Calculate the weights
    wgts = month_length.groupby("time.year") / month_length.groupby("time.year").sum()

    # Setup our masking for nan values
    cond = ds.isnull()
    ones = xr.where(cond, 0.0, 1.0)

    # Calculate the numerator
    obs_sum = (ds * wgts).resample(time="AS").sum(dim="time")

    # Calculate the denominator
    ones_out = (ones * wgts).resample(time="AS").sum(dim="time")

    # Return the weighted average
    return obs_sum / ones_out

wm=weighted_temporal_mean(test)
print("nbytes in MB:", wm.nbytes / (1024*1024))

Any idea how to ensure that the memory is freed up, or am I overlooking something and this behavior is actually expected?

Thank you!

(I also asked this question on stack exchange, as soon as I got a solution here or there I will add it to my post)

[mchoblet]

Got a tip on stack overflow: It seems like the groupby and resampling operations are working on my initial dataset ds. Deleting it afterwards frees up the space, thus I just insert a ds=ds.copy() at the beginning of the function to avoid the problem.
Of course I'm still open for suggestions regarding better workflow/optimized memory use for this type of tasks.

[max-sixty]

A couple of suggestions:

• What happens if you put the open_dataset inside the function? If you do that, then there is nothing that python can be holding onto, assuming the return value is small.
• Looking at htop is useful but not perfect — IIUC python will request the memory, and not give it back until it's requested back. (I am not the expert here but have had similar issues in the past). You could use something like memprofiler, IIRC.

[mchoblet]

Thank you Max for the suggestions! Indeed putting open_dataset inside is also a solution, although in my case I prefer handling the data separately and work on deep copies of my Datasets. I wasn't able to see where the memory went to via memprofiler, but that's probably due to me not being used to that tool. I guess that from now on I'll jest get used to the fact that groupby/resample manipulates the input array although its not obvious when I look at it as a simple cell output (type variable name).

[max-sixty]

Great, happy you solved the issue!