Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Load esmvalcore.dataset.Dataset objects in parallel using Dask #2517

Open
wants to merge 9 commits into
base: main
Choose a base branch
from

Conversation

bouweandela
Copy link
Member

@bouweandela bouweandela commented Sep 9, 2024

Description

Load the individual files in a dataset in parallel using Dask and add the option to get a dask.delayed.Delayed back from esmvalcore.dataset.Dataset that can be fed to dask.compute to get an iris.cube.Cube. This can considerably speed up loading datasets that consist of many files or, when used with the delayed option, speed up loading multiple datasets.

Related to #2300 and #2316

Link to documentation: https://esmvaltool--2517.org.readthedocs.build/projects/ESMValCore/en/2517/api/esmvalcore.dataset.html#esmvalcore.dataset.Dataset.load


Before you get started

Checklist

It is the responsibility of the author to make sure the pull request is ready to review. The icons indicate whether the item will be subject to the 🛠 Technical or 🧪 Scientific review.


To help with the number pull requests:

@bouweandela bouweandela added the dask related to improvements using Dask label Sep 9, 2024
Copy link

codecov bot commented Sep 9, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.77%. Comparing base (e9fccbd) to head (81b41a0).

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2517   +/-   ##
=======================================
  Coverage   94.77%   94.77%           
=======================================
  Files         251      251           
  Lines       14266    14286   +20     
=======================================
+ Hits        13520    13540   +20     
  Misses        746      746           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Copy link
Contributor

@valeriupredoi valeriupredoi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is brilliant, bud! I've been meaning to get delayed in places in Core for some time. Got one possible nagging comment through - from https://docs.dask.org/en/stable/delayed-best-practices.html they say "Every delayed task has an overhead of a few hundred microseconds. Usually this is ok, but it can become a problem if you apply dask.delayed too finely. In this case, it’s often best to break up your many tasks into batches or use one of the Dask collections to help you." - I am guessing this applies to O(millions) (at least) but can we maybe run a test with one of those mega recipes that loads hundreds of datasets?

@valeriupredoi
Copy link
Contributor

oh and maybe a line or two in the documentation perhaps? Bit of an advanced topic, so maybe a very short reference

@bouweandela
Copy link
Member Author

bouweandela commented Oct 23, 2024

This may need a bit more testing. The recipe below fails with the distributed scheduler and the iris main branch.

# ESMValTool
# recipe_python.yml
#
# See https://docs.esmvaltool.org/en/latest/recipes/recipe_examples.html
# for a description of this recipe.
#
# See https://docs.esmvaltool.org/projects/esmvalcore/en/latest/recipe/overview.html
# for a description of the recipe format.
---
documentation:
  description: |
    Example recipe that plots a map and timeseries of temperature.

  title: Recipe that runs an example diagnostic written in Python.

  authors:
    - andela_bouwe
    - righi_mattia

  maintainer:
    - schlund_manuel

  references:
    - acknow_project

  projects:
    - esmval
    - c3s-magic

datasets:
  - {dataset: FGOALS-f3-L, ensemble: 'r1i1p1f1', grid: gn}

preprocessors:
  # See https://docs.esmvaltool.org/projects/esmvalcore/en/latest/recipe/preprocessor.html
  # for a description of the preprocessor functions.

  annual_mean_global:
    area_statistics:
      operator: mean
    annual_statistics:
      operator: mean
    convert_units:
      units: degrees_C

diagnostics:

  timeseries:
    description: Annual mean temperature in Amsterdam and global mean since 1850.
    themes:
      - phys
    realms:
      - atmos
    variables:
      tos_global:
        short_name: tos
        mip: Omon
        project: CMIP6
        exp: [historical, ssp585]
        preprocessor: annual_mean_global
        timerange: 1850/2100
        caption: Annual global mean {long_name} according to {dataset}.
    scripts:
      script1:
        script: examples/diagnostic.py
        quickplot:
          plot_type: plot

@bouweandela
Copy link
Member Author

This issue mentioned above is fixed by SciTools/iris#6187.

@@ -765,6 +798,51 @@ def _load(self) -> Cube:
**self.facets,
}
settings["concatenate"] = {"check_level": self.session["check_level"]}

result = []
for input_file in input_files:
Copy link
Contributor

@schlunma schlunma Oct 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This changes how data is passed through the different preprocessor functions, doesn't it?

Right now, for example, fix_metadata will get ALL cubes from ALL files as input. With this change here, it will only get the cubes from one file, right?

I know that fix_metadata itself groups by file, but this is already very problematic (see #1806 and #2551).

I also fear that this might have other undesired side effects. Why do you need to treat these first preprocessor functions differently in the new code?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do you need to treat these first preprocessor functions differently in the new code?

To improve parallelism. Like this, each input file can be loaded and preprocessed up to the concatenate step in parallel.

This changes how data is passed through the different preprocessor functions, doesn't it?

No, it just takes the grouping out of fix_metadata and implements it in the function calling fix_metadata to enable additional parallelism. If this pull request is merged, #2551 would need to be updated to do the grouping here instead of inside fix_metadata.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, I think I misunderstood the code in the first place. The function preprocess is not at all straightforward when it comes to handling of input and output types...I agree that the behavior has not changed.

I will test this with a couple of recipes once Levante is running again next week. In the meantime, would it make sense to remove the grouping of files in fix_metadata? It would be confusing to have this in two places of the code. I know that this wouldn't be strictly backwards-compatible, but the grouping was only enabled if the cubes have a source_file attribute (which is probably only the case when used within ESMValTool). I highly doubt that this function would be very useful outside of ESMValTool anyway.

Copy link
Contributor

@schlunma schlunma left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I fear that this will break existing recipes due to changes to the preprocessing pipeline. Will remove this block once resolved.

@bouweandela
Copy link
Member Author

Sorry, I fear that this will break existing recipes due to changes to the preprocessing pipeline. Will remove this block once resolved.

Thanks for reviewing @schlunma! As far as I can see it does not change the preprocessing pipeline, but maybe you can find a case where it does? Maybe you could run a recipe that you think could potentially be broken as a test and report back the result?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
dask related to improvements using Dask
Development

Successfully merging this pull request may close these issues.

3 participants