Repository: ecmwf/earthkit-data
Pull request #669 (2025-04-03 10:49:39)

Fix crash when chunks used in Xarray created from GRIB data
This PR fixes the following issue:

Assume the input grib contains a large number or large fields and we want to use chunking (dask) to process it as an xarray. Then the following computation crashes due to a multithreading problem:

```python
ds_in = from_source("file", "t2.grib")
ds = ds_in.to_xarray(time_dim_mode="valid_time", chunks="auto")

# this line crashes
avg = ds["2t"].mean("valid_time").load()
```

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/669?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 90.18%. Comparing base [(`66cddef`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/66cddef0de9c444561a4b74d43cb1f375e4e28a9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) to head [(`ae2b3c6`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/ae2b3c6e1346a8bf600bd053e56dbd008c1fa0f2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).
> Report is 17 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #669      +/-   ##
==========================================
+ Coverage   90.02%   90.18%   +0.15%     
==========================================
  Files         158      159       +1     
  Lines       12086    12160      +74     
  Branches      593      595       +2     
==========================================
+ Hits        10881    10967      +86     
+ Misses       1023     1011      -12     
  Partials      182      182              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/669?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>

Patch:
diff --git a/docs/release_notes/version_0.13_updates.rst b/docs/release_notes/version_0.13_updates.rst
index aebc4bd..b17ee63 100644
--- a/docs/release_notes/version_0.13_updates.rst
+++ b/docs/release_notes/version_0.13_updates.rst
@@ -2,6 +2,15 @@ Version 0.13 Updates
 /////////////////////////
 
 
+Version 0.13.7
+===============
+
+Fixes
+++++++++++
+
+- Fixed issue when Xarray created from GRIB with chunks failed in computations  (:pr:`669`).
+
+
 Version 0.13.6
 ===============
 
diff --git a/src/earthkit/data/readers/grib/index/__init__.py b/src/earthkit/data/readers/grib/index/__init__.py
index 36b509d..99ce578 100644
--- a/src/earthkit/data/readers/grib/index/__init__.py
+++ b/src/earthkit/data/readers/grib/index/__init__.py
@@ -10,6 +10,7 @@
 import logging
 import math
 import os
+import threading
 from abc import abstractmethod
 from collections import defaultdict
 
@@ -244,6 +245,7 @@ class GribFieldManager:
     def __init__(self, policy, owner):
         self.policy = policy
         self.cache = None
+        self.lock = threading.Lock()
 
         if self.policy == "persistent":
             from lru import LRU
@@ -262,13 +264,14 @@ class GribFieldManager:
 
     def field(self, n, create):
         if self.cache is not None:
-            if n in self.cache:
-                return self.cache[n]
-            else:
-                field = create(n)
-                self._field_created()
-                self.cache[n] = field
-                return field
+            with self.lock:
+                if n in self.cache:
+                    return self.cache[n]
+                else:
+                    field = create(n)
+                    self._field_created()
+                    self.cache[n] = field
+                    return field
         else:
             self._field_created()
             return create(n)
@@ -287,10 +290,12 @@ class GribFieldManager:
 
 
 class GribHandleManager:
+    # TODO: split into policies
     def __init__(self, policy, cache_size):
         self.policy = policy
         self.max_cache_size = cache_size
         self.cache = None
+        self.lock = threading.Lock()
 
         if self.policy == "cache":
             if self.max_cache_size > 0:
@@ -313,17 +318,21 @@ class GribHandleManager:
     def handle(self, field, create):
         if self.policy == "cache":
             key = (field.path, field._offset)
-            if key in self.cache:
-                return self.cache[key]
-            else:
-                handle = create()
-                self._handle_created()
-                self.cache[key] = handle
-                return handle
+            with self.lock:
+                if key in self.cache:
+                    return self.cache[key]
+                else:
+                    handle = create()
+                    self._handle_created()
+                    self.cache[key] = handle
+                    return handle
         elif self.policy == "persistent":
             if field._handle is None:
-                field._handle = create()
-                self._handle_created()
+                with self.lock:
+                    if field._handle is None:
+                        field._handle = create()
+                        self._handle_created()
+                    return field._handle
             return field._handle
         elif self.policy == "temporary":
             self._handle_created()
@@ -382,7 +391,6 @@ class GribFieldListInFiles(GribFieldList):
         return field
 
     def _getitem(self, n):
-        # TODO: check if we need a mutex here
         if isinstance(n, int):
             if n < 0:
                 n += len(self)


Test patch:
diff --git a/tests/xr_engine/test_xr_chunks.py b/tests/xr_engine/test_xr_chunks.py
new file mode 100644
index 0000000..79143d7
--- /dev/null
+++ b/tests/xr_engine/test_xr_chunks.py
@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+
+import numpy as np
+import pytest
+
+from earthkit.data import config
+from earthkit.data import from_source
+from earthkit.data.testing import earthkit_remote_test_data_file
+
+
+@pytest.mark.long_test
+@pytest.mark.cache
+@pytest.mark.parametrize(
+    "field_policy",
+    [
+        {"grib-field-policy": "persistent"},
+        {"grib-field-policy": "temporary"},
+    ],
+)
+@pytest.mark.parametrize(
+    "handle_policy",
+    [
+        {"grib-handle-policy": "persistent"},
+        {"grib-handle-policy": "temporary"},
+        {"grib-handle-policy": "cache", "grib-handle-cache-size": 1},
+        {"grib-handle-policy": "cache", "grib-handle-cache-size": 10},
+        {"grib-handle-policy": "cache", "grib-handle-cache-size": 100},
+        {"grib-handle-policy": "cache", "grib-handle-cache-size": 300},
+    ],
+)
+@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_1(field_policy, handle_policy, _kwargs):
+    with config.temporary(**field_policy, **handle_policy):
+        ds_in = from_source(
+            "url", earthkit_remote_test_data_file("test-data", "xr_engine", "date", "t2_1_year_hourly.grib")
+        )
+
+        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)
+
+
+# This test is a copy of the previous one, but only using the default config
+@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_2(_kwargs):
+    # the default settings
+    field_policy = {"grib-field-policy": "persistent"}
+    handle_policy = {"grib-handle-policy": "cache", "grib-handle-cache-size": 1}
+
+    with config.temporary(**field_policy, **handle_policy):
+        ds_in = from_source(
+            "url", earthkit_remote_test_data_file("test-data", "xr_engine", "date", "t2_1_year_hourly.grib")
+        )
+
+        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)