Intermediate Xarray+Dask tutorial

[ncclementi]

Hi folks, bringing here a conversation that started over in the dask slack channel.

@dcherian shared

I hear many requests for “intermediate” dask+xarray tutorial material. I’m sure you must also. Do you have ideas or examples of what this material might cover?
In general, it seems like there’s nice introductory material, and some more complex demos showing awesome scaling and all that, but it is quite hard for users to bridge the gap. I’m wondering if you all have ideas on what kinds of material needs to be developed

Feedback I’ve received on this topic so far:

• debugging of commonish error messages
• tips for more efficient chunking debugging. techniques:
  “what’s the chunk size? Okay, what’s the chunk size after the next operation on the data?” with a sprinkling of “whoa, your dask graph has more than a million tasks? that’s too many!”

cc'ing @jacobtomlinson that express interest in this. Quoting him

I think it would be useful to show some real-world data loading. Some common small-scale operations.

I think it would be great to collaborate on this, (xarray+dask) and maybe create a working group to tackle it. Ideas?
@TomNicholas @jbusecke @paigem would any of you be interested in this?

[paigem]

This sounds like a great idea and something that would be useful to the climate research community - thanks for the ping @ncclementi!

Can you clarify what the proposed output will be? E.g. an organized workshop, a self-guided tutorial on the Dask (or xarray) website, a blog post, etc?

[dcherian]

a self-guided tutorial on the Dask (or xarray) website

This is the most important outcome IMO.

What we have on the xarray side is:

• basic intro https://tutorial.xarray.dev/intermediate/xarray_and_dask.html
• and then apply_ufunc:
  • https://tutorial.xarray.dev/advanced/apply_ufunc/simple_dask_apply_ufunc.html
  • https://tutorial.xarray.dev/advanced/apply_ufunc/vectorize_1d.html
• and xr.map_blocks: https://tutorial.xarray.dev/advanced/map_blocks/map_blocks.html

Once you finish the first tutorial, what do we point to? Everything else is quite complicated!

[ncclementi]

I agree! I think the most important outcome should be a self-guided material. Eventually, I can see this material could be used to organize workshops and tutorials, similar to what we do with the dask and xarray base tutorials.

[jbusecke]

Certainly interested. However I am still a bit unclear what the end product will cover. But absolutely happy to keep chatting.

[dcherian]

I am still a bit unclear what the end product will cover.

This is the fundamental question to answer!

[dcherian]

I thought about this a little and now think the question to answer is "What are intermediate-level concepts a dask.array user must understand?"

The beginner-level concepts are:

• xarray: adds metadata to bare arrays; express computations using metadata
• dask: cuts up arrays into blocks; scheduler decides order to run blocks; blocking allows parallelism and low-memory streaming computations.

The intermediate-level concepts are:

• xarray: abstract computation patterns for common analysis tasks; extensiblity (accessors, wrap other array types)
• dask: ?????

I think one important dask concept might come from the idea of "chunk management". This would cover things like

• make sure chunk sizes are appropriate really early on; what is chunks="auto"
• talk about aligning chunks with files stored on disk (Intermediate Xarray+Dask tutorial #6945 (reply in thread))
• importance of specifying chunksize at data read
• relation between chunksizes and number of tasks.
• rechunking is not cheap in most cases
• while output blocks may be small, intermediate blocks need not be. So keep monitoring chunksizes (and tasks) throughout the pipeline.
• some operations need not work nicely with blocks
  • (like reshape)
• How to easily apply functions blockwise using map_blocks or apply_gufunc when it makes sense.

[ncclementi]

For the "how to chunk" aspect we could re-use most of the array notebook from the new dask tutorial. Where we concentrate more on understanding chunking.

For the map_blocks we could take inspiration from the "Hacking-dask tutorial" that goes over more advanced techniques.

I think having a nice application where we can show this using xarray, real data and a real workflow would be the best. But I'm clueless in this aspect as I'm not familiar with applications. I think here is where it comes the input from the geoscientist.

[ThomasMGeo]

Some interesting ideas/things to maybe cover from @djhoese 's AMS talk: https://ams.confex.com/ams/102ANNUAL/meetingapp.cgi/Paper/398825

Thanks for the share @maxrjones

[djhoese]

And I can always make a much longer listed if needed. Just have to think about all the gotchas I've run into over the years.

[djhoese]

The most recent one which kind of goes along with @jukent's comment below is that the dask chunk sizes used should be the same size or smallerlarger than the internal chunking of the file being read. For example, I had some data where the NetCDF file was chunked at 4096x4096 elements, but my chosen dask chunking was 1024x1024. This made the processing essentially come to a hault because each dask worker had to load the large file chunks and then remove the parts it wasn't using. Apparently this was very wasteful.

Edit: Fix typo. Dask chunks should not be smaller than file internal chunking.

[paigem]

I agree with @djhoese - I have seen a lot of confusion (both myself and others) between the internal chunking on disk of a dataset and the Dask chunking for computation. A detailed explanation of how these chunking schemes work together might be a good fit for an intermediate tutorial.

[jukent]

What I hear a lot from at ESDS office hours is that Dask works great in theory with the curated datasets used in tutorials, but that people run into issues when applying their data - so perhaps we could bridge this gap with a workshop or use enrollee-submitted datasets? A live Dask debugging session?

[ThomasMGeo]

What a cool idea, love the thought of having less 'sanitized' datasets, or datasets where things are weird (which is most datasets!).

[paigem]

Yes!! I very much agree with what @jukent says here! A live Dask debugging session that is then recorded, with detailed explanations of the debugging process (e.g. looking at the task graph, reading the dashboard) and outlining the first steps to take to figure out the root cause of a problem would be a great resource.

[jbusecke]

I think all of these suggestions for 'how to approach problems' are phenomenal (even though hard to realize). To me personally this is exactly what makes users go from beginner to intermediate/advanced: the ability to figure out what is wrong.
I have to admit that I am still at times quite clueless as where to start looking for problems, so a detailed guide/strategy would be extremely useful.

[djhoese]

even though hard to realize

Yeah, this is probably the hardest part about preparing something for this. In the best case it would be a real world case, but that probably means the code is entwined in some larger project/code base or that to keep the code simple the example can't be doing too many things. But if the example isn't doing too many things then it likely won't hit the "good parts" of the debugging process. So then we're talking a fake example that maybe does things wrong on purpose (to be corrected as part of the tutorial) or does things that maybe don't have an obvious purpose as a result. It may also require doing some debugging that doesn't reveal a problem but demonstrates why it is a good thing to check (like changing chunk sizes and seeing how the dask graph and processing times change).

[jukent]

If possible Project Pythia would love to turn this into a Cookbook in our gallery once we have some content in place. I'm keeping my eyes and ears open during NCAR office hours for debugging use cases (just had one where only 1 worker was being used).

[dcherian]

@ncclementi :

real data and a real workflow would be the best.

I'm not sure actually. We have some of these but they don't really teach anything other than "well this computation is possible with that dataset, but not mine for some reason, perhaps I should dramatically rechunk my dataset".

---

After thinking for a while, I think one solution is to demonstrate the same calculation with multiple chunking schemes. This will fail for some, work OK for some, and work spectacularly well for a few. Explaining why and point out clues to detect why this is happening would help teach some intuition about chunking. I now remember I tried this to some extent in flox but it does need improvement.

@djhoese :

the dask chunk sizes used should be the same size or smaller than the internal chunking of the file being read. For example, I had some data where the NetCDF file was chunked at 4096x4096 elements, but my chosen dask chunking was 1024x1024.

This is a great example! We can easily demonstrate good and bad performance with a fake file created with different internal chunks.

The other gotcha here is threads vs processes when reading netcdf4/hdf5 fies.

---

@jbusecke:

think all of these suggestions for 'how to approach problems' are phenomenal (even though hard to realize).

@paigem:

A live Dask debugging session that is then recorded, with detailed explanations of the debugging process (e.g. looking at the task graph, reading the dashboard) and outlining the first steps to take to figure out the root cause of a problem would be a great resource.

Writing up the "debugging" is hard (here's one attempt ), I agree that a video could work well here.

[deeplycloudy]

@dcherian:

I’m wondering if you all have ideas on what kinds of material needs to be developed

@djhoese:

So then we're talking a fake example that maybe does things wrong on purpose (to be corrected as part of the tutorial) or does things that maybe don't have an obvious purpose as a result.

It's possible, by using habits you've developed while working in numpy, to accidentally stop using Dask! Illustrating when/how this can happen could be a good tutorial topic.

An example arose this week in my lab. We wrote a function to QC a long time series of 3D spatial grids, followed by finding the maximum along the vertical dimension.

def qc_reflectivity2(dataset, rhv, ref=None):
    for i in range(len(dataset["time"])):
        refl = np.array(dataset["reflectivity"][i, :, :, :])
        refl[np.array(dataset["cross_correlation_ratio"][i, :, :, :]) < rhv] = -999  # 0
        if ref:
            refl[refl < ref] = np.nan  # -999 #0
        dataset["reflectivity"][i, :, :, :] = refl
        max_refl = dataset["reflectivity"].max(axis=1, skipna=True)
    return max_refl

ds=xr.open_mfdataset('NEXRAD_processing/KHGX_20210602/grid/KHGX20210602_*grid.nc', engine='netcdf4')
max_reflectivity = qc_reflectivity2(ds, 0.9, ref=10)

The example above, which is pretty well-vectorized compared to other approaches one could imagine, nevertheless had the side effect of loading everything into memory, and in the process stripping out all the Dask arrays.

This code accomplishes the same thing, but runs instantly by building a deferred Dask computation. Note the need to avoid using Numpy in a way that is comfortingly familiar to many. That is a change in mindset that probably needs to be taught explicitly.

bad_rhv = ds["cross_correlation_ratio"] < 0.9
bad_ref = ds["reflectivity"] < 10
bad=bad_rhv & bad_refl
max_refl = ds["reflectivity"].where(~bad, np.nan).max(axis=1)