Pangeo / earth science workflows

[jrbourbeau]

There is a large portion of dask.array users who use Dask in geospatial / climatology / earth science / etc. use cases. We should make sure these use cases are represented here in our benchmarks.

@TomNicholas @dcherian we're trying to establish a set of benchmarks for common workloads faced by Dask users. Ideally these are more realistic and less toy examples. These get used day-to-day by engineers as they run AB tests on various changes that they want to make. Workloads in this suite will naturally be optimized over time. What are 1-3 workflows you'd most like to see run regularly as part of Dask's benchmark suite?

xref #725

[dcherian]

How would you like the code to be written/contributed?

1. Idealized dataset?
2. Presumably some kind of small/big version of the problem, controlled by a small number of parameters?
3. Can xarray be a dependency?

I think some combo of groupby/groupby-anomaly (exercise shuffling, reductions); then rolling/rolling anomaly (to exercise map_overlap) is probably what we want.

@aulemahal / @huard do you have a prototype xclim workflow that would be good to test?

[jrbourbeau]

How would you like the code to be written/contributed?

Ideally there's some Jupyter notebook somewhere that someone could point us to. Maybe you already have a notebook or two in mind?

Ultimately we want to add the workflow in this repo as a pytest-style test (here's an example of an existing embarrassingly parallel workflow). These tests are run regularly and used in AB tests as we're developing dask and distributed. Eventually these workflows will probably graduate into an example repo with notebooks that people could checkout, play with, and adapt to their own use cases.

Idealized dataset?
Presumably some kind of small/big version of the problem, controlled by a small number of parameters?

We're able to use large-scale clusters on the cloud here, so large, realistic datasets are welcome. For example, that embarrassingly parallel workflow processes ~3 TB of data on AWS.

In terms of complexity we'd love something in between 20-200 lines of code. We're trying to strike a balance between these workflows being representative of what large Dask user groups actual do, while also making them approachable (i.e. we don't want people's eyes to glaze over when they see 1000 lines of code).

Can xarray be a dependency?

Yeah, definitely. Workflows using dask or xarray (backed by dask) are totally welcome.

[TomNicholas]

Hi @jrbourbeau - this sounds like a good idea.

We're able to use large-scale clusters on the cloud here, so large, realistic datasets are welcome.

What about requester-pays data?

Can xarray be a dependency?

Yeah, definitely. Workflows using dask or xarray (backed by dask) are totally welcome.

How about xGCM or xhistogram? They both ultimately just wrap xarray / dask, but produce task graphs with structures that conventional xarray usage doesn't. (If not that's okay, it would just make a neat example trickier.)

@jbusecke might also have some suggestions

[TomNicholas]

Hi @jrbourbeau - this sounds like a good idea.

We're able to use large-scale clusters on the cloud here, so large, realistic datasets are welcome.

What about requester-pays data?

Can xarray be a dependency?

Yeah, definitely. Workflows using dask or xarray (backed by dask) are totally welcome.

How about xGCM or xhistogram? They both ultimately just wrap xarray / dask, but produce task graphs with structures that conventional xarray usage doesn't. (If not that's okay, it would just make a neat example trickier.)

@jbusecke might also have some suggestions

[jrbourbeau]

What about requester-pays data?

Yeah, requester-pays is fine. We can spin up a cluster and pay for data access anywhere on AWS (and GCP soon). For example, that embarassingly parallel example I referenced earlier is processing ~3TB with requester_pays=True.

How about xGCM or xhistogram?

Good question. I'm inclined to stick with just dask and xarray dependencies at the moment as they're more broadly used

[aulemahal]

@dcherian I don't think we have a xclim workflow that would be straightforward and simple enough for what I understand is the goal of this issue.

On the other hand, it would be easy to come up with one. Sticking with xarray and dask, I guess one could pull data from the CMIP6 AWS bucket and then reduce it with rolling and resample calls. In my experience, these two would be the main climate workflow bottlenecks ? The workflow would be scalable by parametrizing the size of the spatial subset and model selection.

I am imagining something like this https://medium.com/pangeo/easy-ipcc-part-1-multi-model-datatree-469b87cf9114, but stripped down to use only xarray and with a more complex function instead of temperature anomalies, like "yearly number of heat wave days", where "heat wave" is spells of N days above a given threshold. I take this example because it involves both a rolling and a resample.

[jbusecke]

I support @TomNicholas suggestions from above.
I think including some sort of vector calculus operator like div/curl would be incredibly valuable in the long run. This sort of operation usually works fine with toy examples but can turn out to be really tricky when running under real-world conditions due to:

• combining several data_arrays of different size (e.g. full 4D array of velocities and a 2D array of cell area)
• not reducing output dimensionality via aggregation (e.g. output data-arrays are the same size as input)
• array padding, which increases the amount of comms between workers

[TomNicholas]

Yeah, requester-pays is fine. We can spin up a cluster and pay for data access anywhere on AWS (and GCP soon). For example, that embarassingly parallel example I referenced earlier is processing ~3TB with requester_pays=True.

Okay that's great.

@jrbourbeau can we specify the resources we want the workload to be tested on? i.e. specify that the total RAM available is larger than a chunk but smaller than the full calculation.

Good question. I'm inclined to stick with just dask and xarray dependencies at the moment as they're more broadly used

I think including some sort of vector calculus operator like div/curl would be incredibly valuable in the long run.

@jbusecke do you think we could make a simplified version of the first step of the hero calculation using just xarray & dask? I'm thinking we use LLC4320, but call xarray.pad explicitly then broadcast metrics across the larger time-dependent arrays. That should create a task graph that has the same gotchas but without importing xGCM. The histogram workload is maybe less important to capture because that literally just calls dask.blockwise.

[jrbourbeau]

can we specify the resources we want the workload to be tested on?

Yeah, definitely. If you look here at the coiled.Cluster API docs, we can specify a specific AWS instance type or the amount of memory / CPU we want workers to have and Coiled will pick an appropriate instance type.

these two would be the main climate workflow bottlenecks
That should create a task graph that has the same gotchas

Just to clarify, for these workflows we're primarily interested in common operations that Pangeo folks want to do. In particular, I'm hopeful that these are representative in general so:

1. If we want to make a change to dask or distributed, we'll be able to have a sense for what broad impact it might have on Pangeo folks
2. If some geospatial / oceanographer / etc. person is curious about Dask and asks themself "Can Dask help me?", they'll be able to looks at these workflows and see something that looks familiar to what they'd want to do

If these common geospatial operations also happen to stress Dask and fail in some cases, then that's also really useful information. I just want to make sure that we include operations that are representative of what geospatial user groups often want to perform, not necessarily cases where Dask falls over.

[jrbourbeau]

Looking through some of the existing pangeo example, I'm wondering if something like https://gallery.pangeo.io/repos/pangeo-gallery/physical-oceanography/01_sea-surface-height.html or https://gallery.pangeo.io/repos/pangeo-gallery/cmip6/global_mean_surface_temp.html would be considered representative. @dcherian @TomNicholas @jbusecke do those notebooks look like common operations the pangeo community cares about?

[dcherian]

Those seem too "easy" at a quick glance.

Here's one pangeo-data/distributed-array-examples#2 . Gabe and I discussed this a bit in
pydata/xarray#6709

This point

I think there is a small but important increase in complexity here because we do ds.u.mean(), ds.v.mean(), (ds.u*ds.v).mean() all together so each chunk of ds.u and ds.v is used for two different outputs.

is definitely a common "feature" of Pangeo workloads.

EDIT: We could add some pad and diff operations to make it look like what Tom's talking about I think

[dcherian]

This one was fixed by worker saturation for a single year of data. Not sure how well it does for all 60 years of data.

The example there uses 20 years, it's taking quite a while for the computation to start.

EDIT: This example takes a loooooooong time to show up on the dashboard. There are lots of tasks. But once it starts, it feels snappy, though I do see some white space in the task stream.

[dcherian]

Finally a pangeo classic, anomalies with respect to a group mean: pangeo-data/distributed-array-examples#4 .

I don't have time to try it now, but excited to see what you coilies figure out :P