Repository: ESMValGroup/ESMValCore

Issue:
Preprocessor function `anomalies` does not work well with Dask Distributed
When I run this recipe
```
preprocessors:
  test:
    anomalies:
      period: month
      reference:
        start_year: 1950
        start_month: 1
        start_day: 1
        end_year: 1979
        end_month: 12
        end_day: 31

diagnostics:
  AR6_Figure_9.3:
    variables:
      tos_ssp585:
        short_name: tos
        exp: ['historical', 'ssp585']
        project: CMIP6
        mip: Omon
        preprocessor: test
        timerange: '1850/2100'
    scripts: null

datasets:
  - {dataset: ACCESS-CM2, ensemble: r1i1p1f1, grid: gn}
```
using `~/.esmvaltool/dask.yml`:
```
cluster:
  n_workers: 3
  threads_per_worker: 1
  memory_limit: 4GiB
```
the following warnings are issued:
```
2023-08-24 09:38:16,787 - distributed.worker.memory - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker-memory.html#memory-not-released-back-to-the-os for more information. -- Unmanaged memory: 2.43 GiB -- Worker memory limit: 4.00 GiB
2023-08-24 09:38:21,367 - distributed.worker.memory - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker-memory.html#memory-not-released-back-to-the-os for more information. -- Unmanaged memory: 2.71 GiB -- Worker memory limit: 4.00 GiB
2023-08-24 09:38:21,602 - distributed.worker.memory - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker-memory.html#memory-not-released-back-to-the-os for more information. -- Unmanaged memory: 2.56 GiB -- Worker memory limit: 4.00 GiB
2023-08-24 09:38:26,036 - distributed.worker.memory - WARNING - gc.collect() took 1.092s. This is usually a sign that some tasks handle too many Python objects at the same time. Rechunking the work into smaller tasks might help.
2023-08-24 09:38:26,037 - distributed.worker.memory - WARNING - Worker is at 80% memory usage. Pausing worker.  Process memory: 3.21 GiB -- Worker memory limit: 4.00 GiB
```

A possible cause could be that the function `_compute_anomalies` is creating many tiny chunks because it is using a Python loop to go over small parts of the data here:
https://github.com/ESMValGroup/ESMValCore/blob/4fd8701efd9e618ade0b3ec83a3733274d4f6b25/esmvalcore/preprocessor/_time.py#L856-L865


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/ESMValGroup/ESMValCore/pull/2200?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup) Report
> Merging [#2200](https://app.codecov.io/gh/ESMValGroup/ESMValCore/pull/2200?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup) (a77773c) into [main](https://app.codecov.io/gh/ESMValGroup/ESMValCore/commit/09e20d41e4590f5b146672727964ab1d0849224b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup) (09e20d4) will **decrease** coverage by `0.01%`.
> The diff coverage is `90.90%`.

```diff
@@            Coverage Diff             @@
##             main    #2200      +/-   ##
==========================================
- Coverage   93.14%   93.13%   -0.01%     
==========================================
  Files         237      237              
  Lines       12829    12828       -1     
==========================================
- Hits        11949    11947       -2     
- Misses        880      881       +1     
```


| [Files Changed](https://app.codecov.io/gh/ESMValGroup/ESMValCore/pull/2200?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup) | Coverage Δ | |
|---|---|---|
| [esmvalcore/preprocessor/\_time.py](https://app.codecov.io/gh/ESMValGroup/ESMValCore/pull/2200?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup#diff-ZXNtdmFsY29yZS9wcmVwcm9jZXNzb3IvX3RpbWUucHk=) | `97.46% <90.90%> (-0.26%)` | :arrow_down: |


bouweandela: Thanks for reviewing! There are decent-looking tests here: https://github.com/ESMValGroup/ESMValCore/blob/10bd77633672f9120730362b1d77c69d62ce828a/tests/unit/preprocessor/_time/test_time.py#L1643-L1806 and I ran the recipe in https://github.com/ESMValGroup/ESMValTool/pull/3356 with it (these changes are needed to be able to run that recipe).
schlunma: Great, this can be merged :+1: 

Patch:
diff --git a/esmvalcore/_recipe/recipe.py b/esmvalcore/_recipe/recipe.py
index ed858e0c3..6f003e17e 100644
--- a/esmvalcore/_recipe/recipe.py
+++ b/esmvalcore/_recipe/recipe.py
@@ -213,9 +213,6 @@ def _get_default_settings(dataset):
 
     settings = {}
 
-    # Configure (deprecated, remove for v2.10.0) load callback
-    settings['load'] = {'callback': 'default'}
-
     if _derive_needed(dataset):
         settings['derive'] = {
             'short_name': facets['short_name'],
diff --git a/esmvalcore/dataset.py b/esmvalcore/dataset.py
index 35c858840..849a9a579 100644
--- a/esmvalcore/dataset.py
+++ b/esmvalcore/dataset.py
@@ -677,19 +677,15 @@ class Dataset:
         iris.cube.Cube
             An :mod:`iris` cube with the data corresponding the the dataset.
         """
-        return self._load_with_callback(callback='default')
-
-    def _load_with_callback(self, callback):
-        # TODO: Remove the callback argument for v2.10.0.
         input_files = list(self.files)
         for supplementary_dataset in self.supplementaries:
             input_files.extend(supplementary_dataset.files)
         esgf.download(input_files, self.session['download_dir'])
 
-        cube = self._load(callback)
+        cube = self._load()
         supplementary_cubes = []
         for supplementary_dataset in self.supplementaries:
-            supplementary_cube = supplementary_dataset._load(callback)
+            supplementary_cube = supplementary_dataset._load()
             supplementary_cubes.append(supplementary_cube)
 
         output_file = _get_output_file(self.facets, self.session.preproc_dir)
@@ -704,7 +700,7 @@ class Dataset:
 
         return cubes[0]
 
-    def _load(self, callback) -> Cube:
+    def _load(self) -> Cube:
         """Load self.files into an iris cube and return it."""
         if not self.files:
             lines = [
@@ -731,7 +727,6 @@ class Dataset:
             **self.facets,
         }
         settings['load'] = {
-            'callback': callback,
             'ignore_warnings': get_ignored_warnings(
                 self.facets['project'], 'load'
             ),
diff --git a/esmvalcore/preprocessor/__init__.py b/esmvalcore/preprocessor/__init__.py
index 736ed0215..36d8cc495 100644
--- a/esmvalcore/preprocessor/__init__.py
+++ b/esmvalcore/preprocessor/__init__.py
@@ -371,6 +371,10 @@ def preprocess(
     function = globals()[step]
     itype = _get_itype(step)
 
+    for item in items:
+        if isinstance(item, Cube) and item.has_lazy_data():
+            item.data = item.core_data().rechunk()
+
     result = []
     if itype.endswith('s'):
         result.append(_run_preproc_function(function, items, settings,
@@ -478,10 +482,7 @@ class PreprocessorFile(TrackedFile):
     def cubes(self):
         """Cubes."""
         if self._cubes is None:
-            callback = self.settings.get('load', {}).get('callback')
-            self._cubes = [
-                ds._load_with_callback(callback) for ds in self.datasets
-            ]
+            self._cubes = [ds.load() for ds in self.datasets]
         return self._cubes
 
     @cubes.setter
diff --git a/esmvalcore/preprocessor/_io.py b/esmvalcore/preprocessor/_io.py
index 62eb52853..596b3b331 100644
--- a/esmvalcore/preprocessor/_io.py
+++ b/esmvalcore/preprocessor/_io.py
@@ -6,7 +6,6 @@ import logging
 import os
 import shutil
 import warnings
-from collections.abc import Callable
 from itertools import groupby
 from pathlib import Path
 from typing import Optional
@@ -96,7 +95,7 @@ def _get_attr_from_field_coord(ncfield, coord_name, attr):
     return None
 
 
-def concatenate_callback(raw_cube, field, _):
+def _load_callback(raw_cube, field, _):
     """Use this callback to fix anything Iris tries to break."""
     # Remove attributes that cause issues with merging and concatenation
     _delete_attributes(
@@ -122,7 +121,6 @@ def _delete_attributes(iris_object, atts):
 
 def load(
     file: str | Path,
-    callback: Optional[Callable] = None,
     ignore_warnings: Optional[list[dict]] = None,
 ) -> CubeList:
     """Load iris cubes from string or Path objects.
@@ -131,11 +129,6 @@ def load(
     ----------
     file:
         File to be loaded. Could be string or POSIX Path object.
-    callback:
-        Callback function passed to :func:`iris.load_raw`.
-
-        .. deprecated:: 2.8.0
-            This argument will be removed in 2.10.0.
     ignore_warnings:
         Keyword arguments passed to :func:`warnings.filterwarnings` used to
         ignore warnings issued by :func:`iris.load_raw`. Each list element
@@ -152,14 +145,6 @@ def load(
         Cubes are empty.
 
     """
-    if not (callback is None or callback == 'default'):
-        msg = ("The argument `callback` has been deprecated in "
-               "ESMValCore version 2.8.0 and is scheduled for removal in "
-               "version 2.10.0.")
-        warnings.warn(msg, ESMValCoreDeprecationWarning)
-    if callback == 'default':
-        callback = concatenate_callback
-
     file = Path(file)
     logger.debug("Loading:\n%s", file)
 
@@ -191,7 +176,7 @@ def load(
         # warnings.filterwarnings
         # (see https://github.com/SciTools/cf-units/issues/240)
         with suppress_errors():
-            raw_cubes = iris.load_raw(file, callback=callback)
+            raw_cubes = iris.load_raw(file, callback=_load_callback)
     logger.debug("Done with loading %s", file)
 
     if not raw_cubes:
diff --git a/esmvalcore/preprocessor/_time.py b/esmvalcore/preprocessor/_time.py
index e49afd478..6c363379e 100644
--- a/esmvalcore/preprocessor/_time.py
+++ b/esmvalcore/preprocessor/_time.py
@@ -12,6 +12,7 @@ from typing import Iterable, Optional
 from warnings import filterwarnings
 
 import dask.array as da
+import dask.config
 import iris
 import iris.coord_categorisation
 import iris.exceptions
@@ -738,7 +739,6 @@ def climate_statistics(
     -------
     iris.cube.Cube
         Climate statistics cube.
-
     """
     original_dtype = cube.dtype
     period = period.lower()
@@ -820,7 +820,6 @@ def anomalies(
     -------
     iris.cube.Cube
         Anomalies cube.
-
     """
     if reference is None:
         reference_cube = cube
@@ -862,23 +861,31 @@ def anomalies(
     return cube
 
 
-def _compute_anomalies(cube, reference, period, seasons):
+def _compute_anomalies(
+    cube: Cube,
+    reference: Cube,
+    period: str,
+    seasons: Iterable[str],
+):
     cube_coord = _get_period_coord(cube, period, seasons)
     ref_coord = _get_period_coord(reference, period, seasons)
-
-    data = cube.core_data()
-    cube_time = cube.coord('time')
-    ref = {}
-    for ref_slice in reference.slices_over(ref_coord):
-        ref[ref_slice.coord(ref_coord).points[0]] = ref_slice.core_data()
-
-    cube_coord_dim = cube.coord_dims(cube_coord)[0]
-    slicer = [slice(None)] * len(data.shape)
-    new_data = []
-    for i in range(cube_time.shape[0]):
-        slicer[cube_coord_dim] = i
-        new_data.append(data[tuple(slicer)] - ref[cube_coord.points[i]])
-    data = da.stack(new_data, axis=cube_coord_dim)
+    indices = np.empty_like(cube_coord.points, dtype=np.int32)
+    for idx, point in enumerate(ref_coord.points):
+        indices = np.where(cube_coord.points == point, idx, indices)
+    ref_data = reference.core_data()
+    axis, = cube.coord_dims(cube_coord)
+    if cube.has_lazy_data() and reference.has_lazy_data():
+        # Rechunk reference data because iris.cube.Cube.aggregate_by, used to
+        # compute the reference, produces very small chunks.
+        # https://github.com/SciTools/iris/issues/5455
+        ref_chunks = tuple(
+            -1 if i == axis else chunk
+            for i, chunk in enumerate(cube.lazy_data().chunks)
+        )
+        ref_data = ref_data.rechunk(ref_chunks)
+    with dask.config.set({"array.slicing.split_large_chunks": True}):
+        ref_data_broadcast = da.take(ref_data, indices=indices, axis=axis)
+    data = cube.core_data() - ref_data_broadcast
     cube = cube.copy(data)
     cube.remove_coord(cube_coord)
     return cube


Test patch:
diff --git a/tests/integration/preprocessor/_io/test_load.py b/tests/integration/preprocessor/_io/test_load.py
index 99247bdc5..c0ad0bba8 100644
--- a/tests/integration/preprocessor/_io/test_load.py
+++ b/tests/integration/preprocessor/_io/test_load.py
@@ -10,7 +10,7 @@ import numpy as np
 from iris.coords import DimCoord
 from iris.cube import Cube, CubeList
 
-from esmvalcore.preprocessor._io import concatenate_callback, load
+from esmvalcore.preprocessor._io import load
 
 
 def _create_sample_cube():
@@ -60,7 +60,7 @@ class TestLoad(unittest.TestCase):
                 cube.attributes[attr] = attr
             self._save_cube(cube)
         for temp_file in self.temp_files:
-            cubes = load(temp_file, callback='default')
+            cubes = load(temp_file)
             cube = cubes[0]
             self.assertEqual(1, len(cubes))
             self.assertTrue((cube.data == np.array([1, 2])).all())
@@ -79,7 +79,7 @@ class TestLoad(unittest.TestCase):
                     coord.attributes[attr] = attr
             self._save_cube(cube)
         for temp_file in self.temp_files:
-            cubes = load(temp_file, callback='default')
+            cubes = load(temp_file)
             cube = cubes[0]
             self.assertEqual(1, len(cubes))
             self.assertTrue((cube.data == np.array([1, 2])).all())
@@ -94,7 +94,7 @@ class TestLoad(unittest.TestCase):
         cube = _create_sample_cube()
         temp_file = self._save_cube(cube)
 
-        cubes = load(temp_file, callback=concatenate_callback)
+        cubes = load(temp_file)
         cube = cubes[0]
         self.assertEqual(1, len(cubes))
         self.assertTrue((cube.data == np.array([1, 2])).all())
diff --git a/tests/integration/preprocessor/test_preprocessing_task.py b/tests/integration/preprocessor/test_preprocessing_task.py
index 81f980c9e..3b33e8f44 100644
--- a/tests/integration/preprocessor/test_preprocessing_task.py
+++ b/tests/integration/preprocessor/test_preprocessing_task.py
@@ -16,7 +16,7 @@ def test_load_save_task(tmp_path):
     iris.save(cube, in_file)
     dataset = Dataset(short_name='tas')
     dataset.files = [in_file]
-    dataset._load_with_callback = lambda _: cube.copy()
+    dataset.load = lambda: cube.copy()
 
     # Create task
     task = PreprocessingTask([
@@ -57,11 +57,11 @@ def test_load_save_and_other_task(tmp_path, monkeypatch):
 
     dataset1 = Dataset(short_name='tas', dataset='dataset1')
     dataset1.files = [file1]
-    dataset1._load_with_callback = lambda _: in_cube.copy()
+    dataset1.load = lambda: in_cube.copy()
 
     dataset2 = Dataset(short_name='tas', dataset='dataset1')
     dataset2.files = [file2]
-    dataset2._load_with_callback = lambda _: in_cube.copy()
+    dataset2.load = lambda: in_cube.copy()
 
     # Create some mock preprocessor functions and patch
     # `esmvalcore.preprocessor` so it uses them.
diff --git a/tests/integration/recipe/test_recipe.py b/tests/integration/recipe/test_recipe.py
index ba53388dd..6090e61c1 100644
--- a/tests/integration/recipe/test_recipe.py
+++ b/tests/integration/recipe/test_recipe.py
@@ -82,7 +82,6 @@ MANDATORY_SCRIPT_SETTINGS_KEYS = (
 )
 
 DEFAULT_PREPROCESSOR_STEPS = (
-    'load',
     'remove_supplementary_variables',
     'save',
 )
@@ -106,9 +105,6 @@ def create_test_file(filename, tracking_id=None):
 def _get_default_settings_for_chl(save_filename):
     """Get default preprocessor settings for chl."""
     defaults = {
-        'load': {
-            'callback': 'default'
-        },
         'remove_supplementary_variables': {},
         'save': {
             'compress': False,
@@ -557,9 +553,6 @@ def test_default_fx_preprocessor(tmp_path, patched_datafinder, session):
     assert preproc_dir.startswith(str(tmp_path))
 
     defaults = {
-        'load': {
-            'callback': 'default'
-        },
         'remove_supplementary_variables': {},
         'save': {
             'compress': False,
diff --git a/tests/unit/preprocessor/_time/test_time.py b/tests/unit/preprocessor/_time/test_time.py
index b847a2972..26b8ebf35 100644
--- a/tests/unit/preprocessor/_time/test_time.py
+++ b/tests/unit/preprocessor/_time/test_time.py
@@ -6,6 +6,7 @@ import unittest
 from datetime import datetime
 from typing import List, Tuple
 
+import dask.array as da
 import iris
 import iris.coord_categorisation
 import iris.coords
@@ -1658,8 +1659,9 @@ def make_map_data(number_years=2):
         standard_name='longitude',
     )
     data = np.array([[0, 1], [1, 0]]) * times[:, None, None]
+    chunks = (int(data.shape[0] / 2), 1, 2)
     cube = iris.cube.Cube(
-        data,
+        da.asarray(data, chunks=chunks),
         dim_coords_and_dims=[(time, 0), (lat, 1), (lon, 2)],
     )
     return cube
diff --git a/tests/unit/recipe/test_recipe.py b/tests/unit/recipe/test_recipe.py
index a437221a0..498dd606d 100644
--- a/tests/unit/recipe/test_recipe.py
+++ b/tests/unit/recipe/test_recipe.py
@@ -541,7 +541,6 @@ def test_get_default_settings(mocker):
 
     settings = _recipe._get_default_settings(dataset)
     assert settings == {
-        'load': {'callback': 'default'},
         'remove_supplementary_variables': {},
         'save': {'compress': False, 'alias': 'sic'},
     }
diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py
index 56baaa60d..1a0114e8c 100644
--- a/tests/unit/test_dataset.py
+++ b/tests/unit/test_dataset.py
@@ -1684,7 +1684,6 @@ def test_load(mocker, session):
 
     load_args = {
         'load': {
-            'callback': 'default',
             'ignore_warnings': None,
         },
         'fix_file': {