Fix progress logging with Dask v2026

esmvalcore/_recipe/recipe.py

@@ -1340,7 +1340,7 @@ def run(self) -> None:
             file.prepare()
         logger.info("Successfully downloaded missing data")
 
-        self.tasks.run(max_parallel_tasks=self.session["max_parallel_tasks"])
+        self.tasks.run(self.session)
         logger.info(
             "Wrote recipe with version numbers and wildcards to:\nfile://%s",
             filled_recipe,

esmvalcore/_task.py

@@ -19,16 +19,21 @@
 from copy import deepcopy
 from pathlib import Path
 from shutil import which
+from typing import TYPE_CHECKING
 
 import dask
+import dask.config
 import psutil
 import yaml
 from distributed import Client
 
-from ._citation import _write_citation_files
-from ._provenance import TrackedFile, get_task_provenance
-from .config._dask import get_distributed_client
-from .config._diagnostics import DIAGNOSTICS, TAGS
+from esmvalcore._citation import _write_citation_files
+from esmvalcore._provenance import TrackedFile, get_task_provenance
+from esmvalcore.config._dask import get_distributed_client
+from esmvalcore.config._diagnostics import DIAGNOSTICS, TAGS
+
+if TYPE_CHECKING:
+    from esmvalcore.config import Session
 
 logger = logging.getLogger(__name__)
 
@@ -270,20 +275,24 @@ def flatten(self):
         tasks.add(self)
         return tasks
 
-    def run(self, input_files: list[str] | None = None) -> None:
+    def run(
+        self,
+        session: Session,
+        input_files: list[str] | None = None,
+    ) -> list[str]:
         """Run task."""
         if not self.output_files:
             if input_files is None:
                 input_files = []
             for task in self.ancestors:
-                input_files.extend(task.run())
+                input_files.extend(task.run(session))
             logger.info(
                 "Starting task %s in process [%s]",
                 self.name,
                 os.getpid(),
             )
             start = datetime.datetime.now()
-            self.output_files = self._run(input_files)
+            self.output_files = self._run(input_files, session=session)
             runtime = datetime.datetime.now() - start
             logger.info(
                 "Successfully completed task %s (priority %s) in %s",
@@ -295,7 +304,7 @@ def run(self, input_files: list[str] | None = None) -> None:
         return self.output_files
 
     @abc.abstractmethod
-    def _run(self, input_files: list[str]) -> list[str]:
+    def _run(self, input_files: list[str], session: Session) -> list[str]:
         """Run task."""
 
     def get_product_attributes(self) -> dict:
@@ -349,7 +358,11 @@ def __init__(self, prev_preproc_dir, preproc_dir, name):
 
         super().__init__(ancestors=None, name=name, products=products)
 
-    def _run(self, _: list[str]) -> list[str]:
+    def _run(
+        self,
+        input_files: list[str],  # noqa: ARG002
+        session: Session,  # noqa: ARG002
+    ) -> list[str]:
         """Return the result of a previous run."""
         metadata = self.get_product_attributes()
 
@@ -588,7 +601,11 @@ def _start_diagnostic_script(self, cmd, env):
             env=complete_env,
         )
 
-    def _run(self, input_files):
+    def _run(
+        self,
+        input_files: list[str],
+        session: Session,  # noqa: ARG002
+    ) -> list[str]:
         """Run the diagnostic script."""
         if self.script is None:  # Run only preprocessor
             return []
@@ -803,17 +820,17 @@ def get_independent(self) -> TaskSet:
                 independent_tasks.add(task)
         return independent_tasks
 
-    def run(self, max_parallel_tasks: int | None = None) -> None:
+    def run(self, session: Session) -> None:
         """Run tasks.
 
         Parameters
         ----------
-        max_parallel_tasks : int
-            Number of processes to run. If `1`, run the tasks sequentially.
+        session:
+            The current session.
         """
-        with get_distributed_client() as client:
+        with get_distributed_client(session) as client:
             if client is None:
-                address = None
+                address: str | None = None
             else:
                 address = client.scheduler.address
                 for task in self.flatten():
@@ -825,19 +842,19 @@ def run(self, max_parallel_tasks: int | None = None) -> None:
                         # Python script.
                         task.settings["scheduler_address"] = address
 
-            if max_parallel_tasks == 1:
-                self._run_sequential()
+            if session["max_parallel_tasks"] == 1:
+                self._run_sequential(session)
             else:
-                self._run_parallel(address, max_parallel_tasks)
+                self._run_parallel(address, session)
 
-    def _run_sequential(self) -> None:
+    def _run_sequential(self, session: Session) -> None:
         """Run tasks sequentially."""
         n_tasks = len(self.flatten())
         logger.info("Running %s tasks sequentially", n_tasks)
 
         tasks = self.get_independent()
         for task in sorted(tasks, key=lambda t: t.priority):
-            task.run()
+            task.run(session)
 
     def _get_dask_config(self, max_parallel_tasks: int) -> dict:
         """Configure the threaded Dask scheduler.
@@ -884,14 +901,19 @@ def _get_dask_config(self, max_parallel_tasks: int) -> dict:
         )
         return {"num_workers": n_workers}
 
-    def _run_parallel(self, scheduler_address, max_parallel_tasks):
+    def _run_parallel(
+        self,
+        scheduler_address: str | None,
+        session: Session,
+    ) -> None:
         """Run tasks in parallel."""
         scheduled = self.flatten()
-        running = {}
+        running: dict[BaseTask, multiprocessing.pool.ApplyResult] = {}
 
         n_tasks = n_scheduled = len(scheduled)
         n_running = 0
 
+        max_parallel_tasks = session["max_parallel_tasks"]
         if max_parallel_tasks is None:
             max_parallel_tasks = available_cpu_count()
         max_parallel_tasks = min(max_parallel_tasks, n_tasks)
@@ -930,7 +952,12 @@ def done(task):
                         if all(done(t) for t in task.ancestors):
                             future = pool.apply_async(
                                 _run_task,
-                                [task, scheduler_address, scheduler_lock],
+                                [
+                                    task,
+                                    session,
+                                    scheduler_address,
+                                    scheduler_lock,
+                                ],
                             )
                             running[task] = future
                             scheduled.remove(task)
@@ -971,7 +998,12 @@ def _copy_results(task, future):
     task.output_files, task.products = future.get()
 
 
-def _run_task(task, scheduler_address, scheduler_lock):
+def _run_task(
+    task: BaseTask,
+    session: Session,
+    scheduler_address: str | None,
+    scheduler_lock: threading.Lock | None,
+) -> tuple[list[str], set[TrackedFile]]:
     """Run task and return the result."""
     if scheduler_address is None:
         client = contextlib.nullcontext()
@@ -980,6 +1012,6 @@ def _run_task(task, scheduler_address, scheduler_lock):
 
     with client:
         task.scheduler_lock = scheduler_lock
-        output_files = task.run()
+        output_files = task.run(session)
 
     return output_files, task.products

esmvalcore/config/_dask.py

@@ -12,16 +12,15 @@
 import dask.config
 from distributed import Client
 
-from esmvalcore.config import CFG
-from esmvalcore.exceptions import (
-    InvalidConfigParameter,
-)
+from esmvalcore.exceptions import InvalidConfigParameter
 
 if TYPE_CHECKING:
     from collections.abc import Generator
 
     from distributed.deploy import Cluster
 
+    from esmvalcore.config import Session
+
 logger = logging.getLogger(__name__)
 
 
@@ -76,9 +75,9 @@ def validate_dask_config(dask_config: Mapping) -> None:
 
 
 @contextlib.contextmanager
-def get_distributed_client() -> Generator[None | Client]:
+def get_distributed_client(session: Session) -> Generator[None | Client]:
     """Get a Dask distributed client."""
-    dask_config = deepcopy(CFG["dask"])
+    dask_config = deepcopy(session["dask"])
     validate_dask_config(dask_config)
 
     # Set up cluster and client according to the selected profile

esmvalcore/config/_logging.py

@@ -78,7 +78,7 @@ def _purge_file_handlers(cfg: dict) -> None:
 def _get_log_files(
     cfg: dict,
     output_dir: os.PathLike | str | None = None,
-) -> list:
+) -> list[str]:
     """Initialize log files for the file handlers."""
     log_files = []
 
@@ -114,21 +114,21 @@ def configure_logging(
     cfg_file: os.PathLike | str | None = None,
     output_dir: os.PathLike | str | None = None,
     console_log_level: str | None = None,
-) -> list:
+) -> list[str]:
     """Configure logging.
 
     Parameters
     ----------
-    cfg_file : str, optional
+    cfg_file:
         Logging config file. If `None`, defaults to `configure-logging.yml`
-    output_dir : str, optional
+    output_dir:
         Output directory for the log files. If `None`, log only to the console.
-    console_log_level : str, optional
+    console_log_level:
         If `None`, use the default (INFO).
 
     Returns
     -------
-    log_files : list
+    :
         Filenames that will be logged to.
     """
     if cfg_file is None:

esmvalcore/preprocessor/__init__.py

@@ -117,6 +117,7 @@
     from dask.delayed import Delayed
     from iris.cube import CubeList
 
+    from esmvalcore.config import Session
     from esmvalcore.dataset import Dataset
     from esmvalcore.typing import FacetValue
 
@@ -855,7 +856,11 @@ def __init__(
         self.debug = debug
         self.write_ncl_interface = write_ncl_interface
 
-    def _run(self, _: list[str]) -> list[str]:  # noqa: C901,PLR0912
+    def _run(  # noqa: C901,PLR0912
+        self,
+        input_files: list[str],  # noqa: ARG002
+        session: Session,
+    ) -> list[str]:
         """Run the preprocessor."""
         for product in self.products:
             product.activity = self.activity
@@ -906,7 +911,11 @@ def _run(self, _: list[str]) -> list[str]:  # noqa: C901,PLR0912
                 "Computing and saving data for preprocessing task %s",
                 self.name,
             )
-            _compute_with_progress(delayeds, description=self.name)
+            _compute_with_progress(
+                delayeds,
+                session=session,
+                description=self.name,
+            )
         finally:
             if self.scheduler_lock is not None:
                 self.scheduler_lock.release()