Fix memory leak in Xarray computations using chunking (#694)

Author: sandorkertesz

docs/release_notes/version_0.14_updates.rst

@@ -2,6 +2,14 @@ Version 0.14 Updates
 /////////////////////////
 
 
+Version 0.14.1
+===============
+
+Fixes
++++++++++++++++++
+
+- Fixed issue when Xarray computations used excessive memory when the data was created with the Xarray engine using chunking (:pr:`694`).
+
 Version 0.14.0
 ===============
 

src/earthkit/data/readers/grib/virtual.py

@@ -10,7 +10,6 @@
 import logging
 from functools import cached_property
 
-from earthkit.data import from_source
 from earthkit.data.core.fieldlist import Field
 from earthkit.data.core.metadata import WrappedMetadata
 from earthkit.data.utils.dates import date_to_grib
@@ -120,22 +119,19 @@ def _metadata(self):
     def _values(self, dtype=None):
         return self._field._values(dtype=dtype)
 
-    @cached_property
+    @property
     def _field(self):
         if self.reference:
             return self.reference
         else:
-            p = self.owner.retriever.get(self.request)
-            return from_source("file", p, stream=True, read_all=True)[0]
+            return self.owner.retriever.get(self.request)[0]
 
 
 class VirtualGribFieldList(GribFieldList):
     def __init__(self, request_mapper, retriever):
         self.request_mapper = request_mapper
         self.retriever = retriever
 
-        path = self.retriever.get(self.request_mapper.request_at(0))
-        self.reference = from_source("file", path)[0]
         self._info_cache = {}
 
     def __len__(self):
@@ -144,6 +140,10 @@ def __len__(self):
     def mutate(self):
         return self
 
+    @cached_property
+    def reference(self):
+        return self.retriever.get(self.request_mapper.request_at(0))[0]
+
     def _getitem(self, n):
         if isinstance(n, int):
             if n < 0:

src/earthkit/data/sources/fdb.py

@@ -103,9 +103,9 @@ def __init__(self, fdb_kwargs):
         self.fdb_kwargs = fdb_kwargs
 
     def get(self, request):
-        fdb = pyfdb.FDB(**self.fdb_kwargs)
-        s = FDBFileSource(fdb, request)
-        return s.path
+        from . import from_source
+
+        return from_source("fdb", request, stream=True, read_all=True, **self.fdb_kwargs)
 
 
 class FDBRequestMapper(RequestMapper):

src/earthkit/data/utils/xarray/builder.py

@@ -8,6 +8,7 @@
 #
 
 import logging
+import threading
 from abc import ABCMeta
 from abc import abstractmethod
 
@@ -176,8 +177,7 @@ def __init__(self, tensor, dims, shape, xp, dtype, variable):
         if dtype is None:
             dtype = numpy.dtype("float64")
         self.dtype = xp.dtype(dtype)
-
-        self._var = variable
+        self.lock = threading.Lock()
 
     @property
     def nbytes(self):
@@ -205,39 +205,37 @@ def __getitem__(self, key: xarray.core.indexing.ExplicitIndexer):
         # patched in local copy for now, but could construct this ourself
 
     def _raw_indexing_method(self, key: tuple):
-        # must be threadsafe
-        # print("_var", self._var)
-        # print(f"dims: {self.dims} key: {key} shape: {self.shape}")
-        # isels = dict(zip(self.dims, key))
-        # r = self.ekds.isel(**isels)
-        # print(f"t-coords={self.tensor.user_coords}")
-        r = self.tensor[key]
-        # print(r.source.ls())
-        # print(f"r-shape: {r.user_shape}")
-
-        field_index = r.field_indexes(key)
-        # print(f"field.index={field_index} coords={r.user_coords}")
-        # result = r.to_numpy(index=field_index).squeeze()
-        result = r.to_numpy(index=field_index, dtype=self.dtype)
-
-        # ensure axes are squeezed when needed
-        singles = [i for i in list(range(len(r.user_shape))) if isinstance(key[i], int)]
-        if singles:
-            result = result.squeeze(axis=tuple(singles))
-
-        # print("result", result.shape)
-        # result = self.ekds.isel(**isels).to_numpy()
-
-        # print("result", result.shape)
-        # print(f"Loaded {self.xp.__name__} with shape: {result.shape}")
-
-        # Loading as numpy but then converting. This needs to be changed upstream (eccodes)
-        # to load directly into cupy.
-        # Maybe some incompatibilities when trying to copy from FFI to cupy directly
-        if self.xp and self.xp != numpy:
-            result = self.xp.asarray(result)
-
-        return result
+        with self.lock:
+            # print("_var", self._var)
+            # print(f"dims: {self.dims} key: {key} shape: {self.shape}")
+            # print(f"t-coords={self.tensor.user_coords}")
+            r = self.tensor[key]
+            # print(r.source.ls())
+            # print(f"r-shape: {r.user_shape}")
+
+            field_index = r.field_indexes(key)
+            # print(f"field.index={field_index} coords={r.user_coords}")
+            # result = r.to_numpy(index=field_index).squeeze()
+            result = r.to_numpy(index=field_index, dtype=self.dtype)
+
+            # ensure axes are squeezed when needed
+            singles = [i for i in list(range(len(r.user_shape))) if isinstance(key[i], int)]
+            if singles:
+                result = result.squeeze(axis=tuple(singles))
+
+            # print("result", result.shape)
+            # result = self.ekds.isel(**isels).to_numpy()
+
+            # print("result", result.shape)
+            # print(f"Loaded {self.xp.__name__} with shape: {result.shape}")
+
+            # Loading as numpy but then converting. This needs to be changed upstream (eccodes)
+            # to load directly into cupy.
+            # Maybe some incompatibilities when trying to copy from FFI to cupy directly
+            if self.xp and self.xp != numpy:
+                result = self.xp.asarray(result)
+
+            return result
 
 
 class BackendDataBuilder(metaclass=ABCMeta):

tests/xr_engine/test_xr_chunks.py

@@ -94,3 +94,33 @@ def test_xr_engine_chunk_2(_kwargs):
         r = ds["2t"].mean("valid_time").load()
 
         assert np.isclose(r.values.mean(), 275.9938876277779)
+
+
+@pytest.mark.skipif(True, reason="Needs to be fixed")
+@pytest.mark.cache
+@pytest.mark.parametrize(
+    "_kwargs",
+    [
+        {},
+        {"chunks": "auto"},
+        {"chunks": {"valid_time": 1}},
+        {"chunks": {"valid_time": 10}},
+        {"chunks": {"valid_time": (100, 200, 432), "latitude": (4, 5, 4), "longitude": (13, 3, 8)}},
+        {"chunks": -1},
+    ],
+)
+def test_xr_engine_chunk_3(_kwargs):
+    # in-memory fieldlist
+    ds_in = from_source(
+        "url",
+        earthkit_remote_test_data_file("test-data", "xr_engine", "date", "t2_1_year_hourly.grib"),
+        stream=True,
+        read_all=True,
+    )
+
+    ds = ds_in.to_xarray(time_dim_mode="valid_time", **_kwargs)
+    assert ds is not None
+
+    r = ds["2t"].mean("valid_time").load()
+
+    assert np.isclose(r.values.mean(), 275.9938876277779)