calculating climatologies efficiently

[rabernat]

An extremely common task in climate science it to calculate a climatology (seasonal average) of some statistics from a long spatio-temporal dataset. This is exactly what I am trying to do in my pangeo use case. To make this concrete, here is an example reproducible on pangeo.pydata.org

import xarray as xr
import gcsfs
ds = xr.open_zarr(gcsfs.GCSMap('pangeo-data/cm2.6/control/temp_salt_u_v-5day_avg')

the dataset looks like this

<xarray.Dataset>
Dimensions:         (nv: 2, st_edges_ocean: 51, st_ocean: 50, time: 1460, xt_ocean: 3600, xu_ocean: 3600, yt_ocean: 2700, yu_ocean: 2700)
Coordinates:
  * nv              (nv) float64 1.0 2.0
  * st_edges_ocean  (st_edges_ocean) float64 0.0 10.07 20.16 30.29 40.47 ...
  * st_ocean        (st_ocean) float64 5.034 15.1 25.22 35.36 45.58 55.85 ...
  * time            (time) object 0181-01-03 12:00:00 0181-01-08 12:00:00 ...
  * xt_ocean        (xt_ocean) float64 -279.9 -279.8 -279.7 -279.6 -279.5 ...
  * xu_ocean        (xu_ocean) float64 -279.9 -279.8 -279.7 -279.6 -279.5 ...
  * yt_ocean        (yt_ocean) float64 -81.11 -81.07 -81.02 -80.98 -80.94 ...
  * yu_ocean        (yu_ocean) float64 -81.09 -81.05 -81.0 -80.96 -80.92 ...
Data variables:
    salt            (time, st_ocean, yt_ocean, xt_ocean) float32 dask.array<shape=(1460, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>
    temp            (time, st_ocean, yt_ocean, xt_ocean) float32 dask.array<shape=(1460, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>
    u               (time, st_ocean, yu_ocean, xu_ocean) float32 dask.array<shape=(1460, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>
    v               (time, st_ocean, yu_ocean, xu_ocean) float32 dask.array<shape=(1460, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>

Its size is close to 12 TB uncompressed.

Calculating the climatology is trivial

ds_mm_clim = ds.groupby('time.month').mean(dim='time')

giving

<xarray.Dataset>
Dimensions:         (month: 12, nv: 2, st_edges_ocean: 51, st_ocean: 50, xt_ocean: 3600, xu_ocean: 3600, yt_ocean: 2700, yu_ocean: 2700)
Coordinates:
  * nv              (nv) float64 1.0 2.0
  * st_edges_ocean  (st_edges_ocean) float64 0.0 10.07 20.16 30.29 40.47 ...
  * st_ocean        (st_ocean) float64 5.034 15.1 25.22 35.36 45.58 55.85 ...
  * xt_ocean        (xt_ocean) float64 -279.9 -279.8 -279.7 -279.6 -279.5 ...
  * xu_ocean        (xu_ocean) float64 -279.9 -279.8 -279.7 -279.6 -279.5 ...
  * yt_ocean        (yt_ocean) float64 -81.11 -81.07 -81.02 -80.98 -80.94 ...
  * yu_ocean        (yu_ocean) float64 -81.09 -81.05 -81.0 -80.96 -80.92 ...
  * month           (month) int64 1 2 3 4 5 6 7 8 9 10 11 12
Data variables:
    salt            (month, st_ocean, yt_ocean, xt_ocean) float32 dask.array<shape=(12, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>
    temp            (month, st_ocean, yt_ocean, xt_ocean) float32 dask.array<shape=(12, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>
    u               (month, st_ocean, yu_ocean, xu_ocean) float32 dask.array<shape=(12, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>
    v               (month, st_ocean, yu_ocean, xu_ocean) float32 dask.array<shape=(12, 50, 2700, 3600), chunksize=(1, 1, 2700, 3600)>

Now I want to either persist this or, even better, save it as a new dataset

# hack to get writable gcsfs target
fs_w_permissions = gcsfs.GCSFileSystem(token='browser')
token = fs_w_permissions.session.credentials
fs_w_token = gcsfs.GCSFileSystem(token=token)
output_path = data_path[:-1] + '-monthly-climatology'
gcsmap_output = gcsfs.mapping.GCSMap(output_path, gcs=fs_w_token)
# save to zarr
save_future = ds_mm_clim.to_zarr(gcsmap_output, compute=False)
save_future.compute(retries=10)

In my head, this should be a pretty "streamable" operation: load aggregate all values of variable salt for month==1 (January), st_ocean==0 (the vertical level); store; and move on to the next variable / month / level.

However, these computations do not run very well on the current stack. The rate of reading data outpaces the rate of writing data, leading to huge memory consumption. This is evident in the dashboard screenshot below:

Here I have a cluster of 80 high-memory workers: 22GB RAM each, a total of 1.76 TB. Yet the cluster has nearly 7TB in memory. The workers are spilling lots of data to disk (I didn't even realize that the workers had significant hard drive space). This seems inefficient, although perhaps there is some logic to it that I don't grasp.

This has been discussed in numerous previous issues.

• xarray groupby monthly mean fail case #99
• groupby on dask objects doesn't handle chunks well pydata/xarray#1832
• Scheduler fail case: centering data with dask.array dask/dask#874

In many ways, this is a duplicate of those issues. However, those issues are also muddled up with problems with worker / scheduler config settings that we have mostly overcome on pangeo.pydata.org. Here workers are not dying--they are just operating in what appears to be a sub-optimal way.

Perhaps now is the time to tackle how this operation is scheduled at the dask level? Or perhaps there is not actually a problem at all? The calculation is slowly ticking forward in an evidently stable way.

[jgerardsimcock]

Worker nodes are allocated 100gb of persistent disk space by default. From what I understand, for gce, you can only adjust this on node-pool creation.

[mrocklin]

I've been working today on a similar problem for a different group of people. It has resulted in a number of small fixes that are slowly adding up to significantly improved performance both for their use case, and also more generally (hopefully some of these fixes carry over to our problem here as well).

What I find very helpful when going through this process is a reproducible example that I can scale down to run easily and quickly on my laptop in less than a minute, and that can also scale down to a managable number of tasks, say 1000, so that I can inspect the graph while it runs.

Creating a reproducible example of the full scale problem that scales down to a laptop can be a challenge, especially as we get out further into particular scientific analyses. It is however, also extremely helpful.

[mrocklin]

As an example of what a minimal example often looks like, here is the problem that I've been working on recently: https://gist.github.com/mrocklin/0eaf8ae2d636915b3de72814068492c2

I've found that the simplicity of the example problem corresponds very strongly with the speed at which progress is made on that problem.

[rabernat]

I apologize for not making my reproducible example minimal enough.

[mrocklin]

Ah, sorry, didn't mean to sound negative. Just wanted to set this up as a possible path for people to be able to push forward on this topic without necessarily understanding the depths of dynamic task scheduling. Arguably most of the work here is distilling the problem down to something on which rapid iteration can occur. I've probably run the notebook above a few hundred times in the last few days (to say nothing of smaller versions in single-threaded situations where I can inspect all of the workers directly).

[rabernat]

I understand and appreciate the value of the approach you described. I didn't post this issue on dask or xarray--I posted it on pangeo, the place where we discuss the actual real-world use cases. This example is already heavily, heavily simplified from the type of analysis that one has to do to achieve actual publishable scientific results.

Given the complexity of this stack, I think such "intermediate" examples have value. This issue represents a canonical operation common to all of climate science.

I will continue to work to simplify it further by eliminating the "real data" part. Do you think that the 'store' aspect is important here? Once I have achieved a more minimal example, would you like it posted here or as a new dask issue?

[rabernat]

Here is a dask-only version of this problem. I am not sure to what extent it reproduces the pathologies of my real use case, but it does appear to have the basic problem of the data generation outpacing the store operation, resulting in a buildup of in-memory data.

import dask.array as dsa
import numpy as np

# create random test data--adjust size at will
ntime = 512
shape = (ntime, 50, 100, 200)
chunkshape = (1, 1) + shape[2:]
data = dsa.random.random(shape, chunks=chunkshape)

# simulate the effects of an xarray groupby operation
n_seasons = 4
sampling = 8
season_index = np.round((np.arange(ntime)/sampling) % n_seasons).astype('int')

# calculate a "climatology" 
arrays = [data[season_index==n].mean(axis=0)[None] for n in range(n_seasons)]
result = dsa.concatenate(arrays, axis=0)

# store it in a fake store
class mock_target(object):
    def __init__(self):
        pass
    def __setitem__(self, *args, **kwargs):
        pass
target = mock_target()
result.store(target, lock=False)

[mrocklin]

Thanks @rabernat ! Here are some initial results:

I reduced the dimensions to (512, 2, 100, 200) to reduce the number of tasks from 50k to around 2k. i then ran visualize on it in the following way:

out = result.store(target, lock=False, compute=False)  # get a lazy delayed value back
out.visualize('dask.pdf', color='order', node_attr={'penwidth': '6'})

This gives you an image that color-codes nodes by their priority. I looked at the image and saw that color seemed to be smooth across the image, which shows that task prioritization was doing a good job. Also being able to just look at the graph gave me good confidence that this is definitely a problem that should be able to run in low memory.

I then ran this with the distributed scheduler and it ran too fast to really see what was going on, so I increased the latter two dimensions. I hope that this stays true to the problem at hand. I watched the status page and saw reasonable-ish looking progress. I watched the graph page and again saw that things looked pretty good. The computation proceeded in a linear fashion from bottom to top when possible.

I happened to have some diagnostic code up that showed the difference between how much memory Dask thinks it should be using, vs how much it actually is using, and saw that there was a significant difference, something like 1.5GB expected vs 3-4GB real. This is consistent with dask/dask#3530 which is a deep issue facing the whole stack that we will need to think hard about in the future.

Now that I looked at a smaller problem that I was able to inspect directly (a lot of the diagnostics fail pretty hard above a few thousand tasks) I switched back to the original array shape and ran the computation again. The tasks here are so small and numerous that the dashboard pretty much choked (the scheduler is probably running at 100% trying to keep up with execution). I increased the size to (512, 50, 1000, 2000) just to give the tasks a little it more weight. It's running now. It seems to be maxing out at around 7GB memory real, 4GB expected (expected becomes real if the memory sharding problem in dask/dask#3530 goes away). I haven't yet dug in to see if this is what we should expect.

Probably the next thing for me to do, if there is no response on the above, is cripple my local network to see if it causes run-away data generation (see also dask/distributed#1989 )

@rabernat is this consistent with what you see? Any thoughts on the above?

[tjcrone]

I'll just chime in with a side note to say that is is very helpful to learn about one of your general procedures for exploring this sort of issue. Thank you @mrocklin!

[rabernat]

@mrocklin - Thanks a lot for this helpful analysis.

I tried adding the following step to the computation to try to represent the load-from-disk step of my original example:

from time import sleep
def slow_down(a, overhead=0.04):
    sleep(overhead)
    return a
data = data.map_blocks(slow_down)

This makes the computation easier to watch from the dashboard (without blowing up memory).

Basically, it doesn't look like this example reproduces the fundamental problem from the original whereby the loading of data outpaces the storing. 😐 This suggests that the problem is more subtle.

One interesting thing to note is that, for the original example on pangeo.pydata.org, things improve if I give it more workers. I re-tried that calculation with 160 workers instead of 80, and it actually ended up keeping much less data in memory (or spilled to disk) than the example with fewer workers. This must be telling us something, but I'm not sure what.

[mrocklin]

It was still a helpful problem to go through. I suspect that if I'm able to simulate a slower network then things will become worse. I'll try to give this a shot some time in the next week (also, disclaimer, I'll be mostly focused on other topics over the next couple weeks)

…

On Fri, May 25, 2018 at 7:36 AM, Ryan Abernathey ***@***.***> wrote: @mrocklin <https://github.com/mrocklin> - Thanks a lot for this helpful analysis. I add the following step to the computation to try to represent the load-from-disk step of my original example: from time import sleepdef slow_down(a, overhead=0.04): sleep(overhead) return a data = data.map_blocks(slow_down) This makes the computation easier to watch from the dashboard (without blowing up memory). Basically, it doesn't look like this example reproduces the fundamental problem from the original whereby the loading of data outpaces the storing. 😐 This suggests that the problem is more subtle. One interesting thing to note is that, for the original example on pangeo.pydata.org, things improve if I give it more workers. I re-tried that calculation with 160 workers instead of 80, and it actually ended up keeping *much less* data in memory (or spilled to disk) than the example with fewer workers. This must be telling us something, but I'm not sure what. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#271 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszJOtygT50KYpAyr6SZbDsegyGNWoks5t1-y4gaJpZM4UMz_Y> .

[rabernat]

I guess it's important to emphasize that, despite what I perceived as an inefficiency, the calculation did work. It basically filled up all of the scratch space, but it didn't crash, and it eventually got through the computation and produced the desired result.

[mrocklin]

I guess that's good to hear :) Also, when creating minimal examples I'm also quite happy to work with xarray if that's easier. I wonder if it might be worth investigating in a generator for random datasets within XArray.

…

On Fri, May 25, 2018 at 11:54 AM, Ryan Abernathey ***@***.***> wrote: I guess it's important to emphasize that, despite what I perceived as an inefficiency, the calculation *did work*. It basically filled up all of the scratch space, but it didn't crash, and it eventually got through the computation and produced the desired result. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#271 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszEj5g9H-shRRVYH5aEvtkhMlaVynks5t2CkmgaJpZM4UMz_Y> .

[mrocklin]

@rabernat can this be closed?

[rabernat]

I would say it can be closed one a version of dask with dask/dask#3648 in it has been deployed to pangeo.pydata.org. Until then, users will continue to run out of memory.

…

On Tue, Jun 26, 2018 at 4:17 PM Matthew Rocklin ***@***.***> wrote: @rabernat <https://github.com/rabernat> can this be closed? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#271 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/ABJFJpgWapnsEoujdBbtzG3wOUBtHbQJks5uApa_gaJpZM4UMz_Y> .

[mrocklin]

cc @martindurant let me know if you want to walk through doing this and have someone nearby just in case On Tue, Jun 26, 2018 at 4:31 PM, Ryan Abernathey <notifications@github.com> wrote:

…

I would say it can be closed one a version of dask with dask/dask#3648 in it has been deployed to pangeo.pydata.org. Until then, users will continue to run out of memory. On Tue, Jun 26, 2018 at 4:17 PM Matthew Rocklin ***@***.***> wrote: > @rabernat <https://github.com/rabernat> can this be closed? > > — > You are receiving this because you were mentioned. > Reply to this email directly, view it on GitHub > <#271 (comment) >, > or mute the thread > <https://github.com/notifications/unsubscribe-auth/ ABJFJpgWapnsEoujdBbtzG3wOUBtHbQJks5uApa_gaJpZM4UMz_Y> > . > — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#271 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszPkP7cYyQRfQk478KzkPnNX-gcSYks5uApoJgaJpZM4UMz_Y> .

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[stale]

This issue has been automatically closed because it had not seen recent activity. The issue can always be reopened at a later date.

[rabernat]

I am trying a new approach to this problem using xarray's new map_blocks function. See this example: https://nbviewer.jupyter.org/gist/rabernat/30e7b747f0e3583b5b776e4093266114

[dcherian]

That's exciting!

And it's mostly a result of the Pangeo Seattle meeting!

We should add some version of this notebook to the xarray docs