Load esmvalcore.dataset.Dataset objects in parallel using Dask

esmvalcore/dataset.py

@@ -6,12 +6,15 @@
 import re
 import textwrap
 import uuid
+from collections.abc import Iterable
 from copy import deepcopy
 from fnmatch import fnmatchcase
 from itertools import groupby
 from pathlib import Path
-from typing import Any, Iterator, Sequence, Union
+from typing import Any, Iterator, Sequence, TypeVar, Union
 
+import dask
+from dask.delayed import Delayed
 from iris.cube import Cube
 
 from esmvalcore import esgf, local
@@ -79,6 +82,14 @@ def _ismatch(facet_value: FacetValue, pattern: FacetValue) -> bool:
             and fnmatchcase(facet_value, pattern))
 
 
+T = TypeVar('T')
+
+
+def _first(elems: Iterable[T]) -> T:
+    """Return the first element."""
+    return next(iter(elems))
+
+
 class Dataset:
     """Define datasets, find the related files, and load them.
 
@@ -664,9 +675,19 @@ def files(self) -> Sequence[File]:
     def files(self, value):
         self._files = value
 
-    def load(self) -> Cube:
+    def load(self, compute=True) -> Cube | Delayed:
         """Load dataset.
 
+        Parameters
+        ----------
+        compute:
+            If :obj:`True`, return the :class:`~iris.cube.Cube` immediately.
+            If :obj:`False`, return a :class:`~dask.delayed.Delayed` object
+            that can be used to load the cube by calling its
+            :meth:`~dask.delayed.Delayed.compute` method. Multiple datasets
+            can be loaded in parallel by passing a list of such delayeds
+            to :func:`dask.compute`.
+
         Raises
         ------
         InputFilesNotFound
@@ -689,7 +710,7 @@ def load(self) -> Cube:
             supplementary_cubes.append(supplementary_cube)
 
         output_file = _get_output_file(self.facets, self.session.preproc_dir)
-        cubes = preprocess(
+        cubes = dask.delayed(preprocess)(
             [cube],
             'add_supplementary_variables',
             input_files=input_files,
@@ -698,7 +719,10 @@ def load(self) -> Cube:
             supplementary_cubes=supplementary_cubes,
         )
 
-        return cubes[0]
+        cube = dask.delayed(_first)(cubes)
+        if compute:
+            return cube.compute()
+        return cube
 
     def _load(self) -> Cube:
         """Load self.files into an iris cube and return it."""
@@ -713,7 +737,14 @@ def _load(self) -> Cube:
             msg = "\n".join(lines)
             raise InputFilesNotFound(msg)
 
+        input_files = [
+            file.local_file(self.session['download_dir']) if isinstance(
+                file, esgf.ESGFFile) else file for file in self.files
+        ]
         output_file = _get_output_file(self.facets, self.session.preproc_dir)
+        debug = self.session['save_intermediary_cubes']
+
+        # Load all input files and concatenate them.
         fix_dir_prefix = Path(
             self.session._fixed_file_dir,
             self._get_joined_summary_facets('_', join_lists=True) + '_',
@@ -739,6 +770,51 @@ def _load(self) -> Cube:
         settings['concatenate'] = {
             'check_level': self.session['check_level']
         }
+
+        result = []
+        for input_file in input_files:
+            files = dask.delayed(preprocess)(
+                [input_file],
+                'fix_file',
+                input_files=[input_file],
+                output_file=output_file,
+                debug=debug,
+                **settings['fix_file'],
+            )
+            # Multiple cubes may be present in a file.
+            cubes = dask.delayed(preprocess)(
+                files,
+                'load',
+                input_files=[input_file],
+                output_file=output_file,
+                debug=debug,
+                **settings['load'],
+            )
+            # Combine the cubes into a single cube per file.
+            cubes = dask.delayed(preprocess)(
+                cubes,
+                'fix_metadata',
+                input_files=[input_file],
+                output_file=output_file,
+                debug=debug,
+                **settings['fix_metadata'],
+            )
+            cube = dask.delayed(_first)(cubes)
+            result.append(cube)
+
+        # Concatenate the cubes from all files.
+        result = dask.delayed(preprocess)(
+            result,
+            'concatenate',
+            input_files=input_files,
+            output_file=output_file,
+            debug=debug,
+            **settings['concatenate'],
+        )
+
+        # At this point `result` is a list containing a single cube. Apply the
+        # remaining preprocessor functions to this cube.
+        settings.clear()
         settings['cmor_check_metadata'] = {
             'check_level': self.session['check_level'],
             'cmor_table': self.facets['project'],
@@ -762,22 +838,17 @@ def _load(self) -> Cube:
             'frequency': self.facets['frequency'],
             'short_name': self.facets['short_name'],
         }
-
-        result = [
-            file.local_file(self.session['download_dir']) if isinstance(
-                file, esgf.ESGFFile) else file for file in self.files
-        ]
         for step, kwargs in settings.items():
-            result = preprocess(
+            result = dask.delayed(preprocess)(
                 result,
                 step,
-                input_files=self.files,
+                input_files=input_files,
                 output_file=output_file,
-                debug=self.session['save_intermediary_cubes'],
+                debug=debug,
                 **kwargs,
             )
 
-        cube = result[0]
+        cube = dask.delayed(_first)(result)
         return cube
 
     def from_ranges(self) -> list['Dataset']:

tests/integration/dataset/test_dataset.py

@@ -3,6 +3,7 @@
 import iris.coords
 import iris.cube
 import pytest
+from dask.delayed import Delayed
 
 from esmvalcore.config import CFG
 from esmvalcore.dataset import Dataset
@@ -34,7 +35,8 @@ def example_data(tmp_path, monkeypatch):
     monkeypatch.setitem(CFG, 'output_dir', tmp_path / 'output_dir')
 
 
-def test_load(example_data):
+@pytest.mark.parametrize('lazy', [True, False])
+def test_load(example_data, lazy):
     tas = Dataset(
         short_name='tas',
         mip='Amon',
@@ -51,7 +53,11 @@ def test_load(example_data):
     tas.find_files()
     print(tas.files)
 
-    cube = tas.load()
-
+    if lazy:
+        result = tas.load(compute=False)
+        assert isinstance(result, Delayed)
+        cube = result.compute()
+    else:
+        cube = tas.load()
     assert isinstance(cube, iris.cube.Cube)
     assert cube.cell_measures()