pydata / pydata/xarray

Uncontrolled memory growth in custom backend

Open
#9,576 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug topic-backends
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

What happened?

I wrote a custom backend. I'm using it to open a file, operate on the data, remove most of it from the Dataset using .isel, open the next, concatenate, and repeat. I noticed the memory used by the system grew significantly over time even though the size of the Dataset was approximately the same. I was able to reproduce the problem without most of this complexity.

I repeatedly created a dummy Dataset with 25 Variables and observed the number of objects with objgraph after each object creation. I see Variable instances continually increasing, even though I have del'd the Dataset after creating it. I think this suggests that something in xarray is not releasing the Dataset.

Variable               75       +25
ReferenceType        8511        +4
lock                   46        +4
dict                21283        +2
KeyedRef               22        +2
SerializableLock        9        +2
list                23256        +1
set                  1873        +1
method               1893        +1
SeedSequence            4        +1
---
Variable              100       +25
ReferenceType        8515        +4
lock                   50        +4
dict                21285        +2
KeyedRef               24        +2
SerializableLock       11        +2
list                23257        +1
set                  1874        +1
method               1894        +1
SeedSequence            5        +1

I picked a random Variable that was not released and printed the reference chain graph.

import random
objgraph.show_chain(objgraph.find_backref_chain(random.choice(objgraph.by_type('Variable')), objgraph.is_proper_module), filename='graph.png')

graph_variable

What did you expect to happen?

I expected the memory used for the Dataset to be released and garbage-collected. I expected the memory in use to plateau instead of grow.

Minimal Complete Verifiable Example
import os
os.system("mamba install --yes objgraph")

import io

import numpy as np
import xarray as xr
from xarray.backends import BackendArray, BackendEntrypoint
from xarray.backends.common import AbstractDataStore
from xarray.backends.file_manager import CachingFileManager
from xarray.backends.locks import SerializableLock, ensure_lock
from xarray.backends.store import StoreBackendEntrypoint
from xarray.core.indexing import (
    IndexingSupport,
    explicit_indexing_adapter,
)
from xarray.core.utils import Frozen, close_on_error


# This backend creates a random 500MB variable and stores it into a Dataset 25 times
class memtest_DataStore(AbstractDataStore):
    """Store for binary data"""

    def __init__(
        self,
        filename,
        lock=None,
        num_bytes=500_000_000,
        **kwargs,
    ):
        if lock is None:
            lock = SerializableLock()
        self.lock = ensure_lock(lock)
        # Create a lock for the file manager
        self._fm_lock = SerializableLock()
        self._manager = CachingFileManager(
            self._open_file,
            filename,
            lock=self._fm_lock,
            kwargs=kwargs,
        )

        self.num_bytes = num_bytes
        self._num_vars = 25
        self._rng = np.random.default_rng()
        self._xr_obj = None

    def _open_file(self, filename, **kwargs):
        b = self._rng.random(self.num_bytes // 8, dtype=np.float64)
        file_obj = io.BytesIO(b, **kwargs)

        return file_obj

    def _create_dataset(self):
        with self._manager.acquire_context(needs_lock=True) as file_obj:
            file_obj.seek(0)
            fltarr = np.frombuffer(file_obj.getvalue(), np.float64)

        xr_obj = xr.Dataset()
        # Create coordinates
        xr_obj["index"] = (
            ("index",),
            range(len(fltarr)),
        )
        # Create several variables so we can see the change
        for i in range(self._num_vars):
            xr_obj[f"var{i}"] = (
                ("index",),
                fltarr,
            )

        return xr_obj

    @property
    def ds(self):
        if not self._xr_obj:
            self._xr_obj = self._create_dataset()

        return self._xr_obj

    def open_store_variable(self, name, var):
        return self.ds[name].variable

    def get_variables(self):
        return Frozen(self.ds.variables)

    def get_attrs(self):
        return Frozen(self.ds.attrs)

    def get_dimensions(self):
        return Frozen(self.ds.dims)

    def close(self):
        self._manager.close()


class memtest_backend(BackendEntrypoint):
    description = "Memory test backend"

    def open_dataset(
        self,
        filename_or_obj,
        *,
        drop_variables=None,
        lock=None,
        num_bytes=500_000_000,
    ):
        store = memtest_DataStore(
            filename_or_obj,
            lock=lock,
            num_bytes=num_bytes,
        )

        store_entrypoint = StoreBackendEntrypoint()
        with close_on_error(store):
            ds = store_entrypoint.open_dataset(
                store,
                drop_variables=drop_variables,
            )

        return ds

# Loop to test memory growth
# This results in about 10GB of additional memory used by the end of the process
import xarray as xr
import objgraph

objgraph.show_growth(limit=1)
for i in range(10):
    ds = xr.open_dataset(i, engine=memtest_backend)
    # Just to test that something's still holding on to the data
    del ds
    objgraph.show_growth()
    print('---')
MVCE confirmation
  • Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
  • Complete example — the example is self-contained, including all data and the text of any traceback.
  • Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
  • New issue — a search of GitHub Issues suggests this is not a duplicate.
  • Recent environment — the issue occurs with the latest version of xarray and its dependencies.
Relevant log output
function    18610    +18610
list                 15403    +12798
function             25617     +7007
dict                 13365     +6005
tuple                13360     +5284
CompositeUnit         2954     +2954
PrefixUnit            1526     +1526
ReferenceType         4284     +1181
cell                  4966     +1112
type                  2397      +792
getset_descriptor     3374      +628
---
Variable                         50       +25
ReferenceType                  4298       +14
builtin_function_or_method     2537       +10
lock                             34        +4
dict                          13367        +2
KeyedRef                         19        +2
SerializableLock                  7        +2
list                          15404        +1
set                             952        +1
method                          369        +1
---
Variable               75       +25
ReferenceType        4302        +4
lock                   38        +4
dict                13369        +2
KeyedRef               21        +2
SerializableLock        9        +2
list                15405        +1
set                   953        +1
method                370        +1
SeedSequence            4        +1
---
Variable              100       +25
ReferenceType        4306        +4
lock                   42        +4
dict                13371        +2
KeyedRef               23        +2
SerializableLock       11        +2
list                15406        +1
set                   954        +1
method                371        +1
SeedSequence            5        +1
---
Variable              125       +25
ReferenceType        4310        +4
lock                   46        +4
dict                13373        +2
KeyedRef               25        +2
SerializableLock       13        +2
list                15407        +1
set                   955        +1
method                372        +1
SeedSequence            6        +1
---
Variable              150       +25
ReferenceType        4314        +4
lock                   50        +4
dict                13375        +2
KeyedRef               27        +2
SerializableLock       15        +2
list                15408        +1
set                   956        +1
method                373        +1
SeedSequence            7        +1
---
Variable              175       +25
ReferenceType        4318        +4
lock                   54        +4
dict                13377        +2
KeyedRef               29        +2
SerializableLock       17        +2
list                15409        +1
set                   957        +1
method                374        +1
SeedSequence            8        +1
---
Variable              200       +25
ReferenceType        4322        +4
lock                   58        +4
dict                13379        +2
KeyedRef               31        +2
SerializableLock       19        +2
list                15410        +1
set                   958        +1
method                375        +1
SeedSequence            9        +1
---
Variable              225       +25
ReferenceType        4326        +4
lock                   62        +4
dict                13381        +2
KeyedRef               33        +2
SerializableLock       21        +2
list                15411        +1
set                   959        +1
method                376        +1
SeedSequence           10        +1
---
Variable              250       +25
ReferenceType        4330        +4
lock                   66        +4
dict                13383        +2
KeyedRef               35        +2
SerializableLock       23        +2
list                15412        +1
set                   960        +1
method                377        +1
SeedSequence           11        +1
---
Anything else we need to know?

This crashes the Binder notebook instance since it uses so much memory.

Environment

INSTALLED VERSIONS

commit: None
python: 3.12.6 | packaged by conda-forge | (main, Sep 30 2024, 18:08:52) [GCC 13.3.0]
python-bits: 64
OS: Linux
OS-release: 5.14.0-505.el9.x86_64
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.14.3
libnetcdf: 4.9.2

xarray: 2024.9.0
pandas: 2.2.3
numpy: 1.26.4
scipy: None
netCDF4: 1.7.1
pydap: None
h5netcdf: 1.2.0
h5py: 3.11.0
zarr: None
cftime: 1.6.4
nc_time_axis: None
iris: None
bottleneck: None
dask: 2023.9.3
distributed: 2023.9.3
matplotlib: 3.9.1
cartopy: None
seaborn: None
numbagg: None
fsspec: 2023.9.2
cupy: None
pint: None
sparse: None
flox: None
numpy_groupies: None
setuptools: 68.2.2
pip: 23.2.1
conda: None
pytest: None
mypy: None
IPython: 8.16.1
sphinx: None

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the MVCE with objgraph and tracing references through memtest_DataStore, BackendEntrypoint.open_dataset, StoreBackendEntrypoint, and CachingFileManager. Done means repeated open_dataset calls followed by deleting each Dataset no longer show unbounded Variable growth or memory use, with regression coverage for the reproduced behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.