Is there any way of having `.map_blocks` be even more opaque to dask?

[max-sixty]

Is your feature request related to a problem?

Currently I have a workload which does something a bit like:

ds = open_zarr(source)
(
    ds.assign(
        x=ds.foo * ds.bar
        y=ds.foo + ds.bar
    ).to_zarr(dest)
)

(the actual calc is a bit more complicated! And while I don't have a MVCE of the full calc, I pasted a task graph below)

Dask — while very impressive in many ways — handles this extremely badly, because it attempts to load the whole of ds into memory before writing out any chunks. There are lots of issues on this in the dask repo; it seems like an intractable problem for dask.

Describe the solution you'd like

I was hoping to make the internals of this task opaque to dask, so it became a much dumber task runner — just map over the blocks, running the function and writing the result, block by block. I thought I had some success with .map_blocks last week — the internals of the calc are now opaque at least. But the dask cluster is falling over again, I think because the write is seen as a separate task.

Is there any way to make the write more opaque too?

Describe alternatives you've considered

I've built a homegrown thing which is really hacky which does this on a custom scheduler — just runs the functions and writes with region. I'd much prefer to use & contribute to the broader ecosystem...

Additional context

(It's also possible I'm making some basic error — and I do remember it working much better last week — so please feel free to direct me / ask me for more examples, if this doesn't ring true)

[dcherian]

dask.delayed maybe?

You have a shuffle like thing going from Tower 3 to Tower 4. That's probably it. What's happening there?

Also you can write from within map_blocks and return 0. Does that fit your workload

[spencerkclark]

@max-sixty I have run into similar issues, and wrote a package called xpartition to facilitate this kind of crude / pragmatic writing approach. I have been hesitant to advertise it widely as I have many of the same reservations as you, but despite that I find myself using it regularly for important things.

[alexamici]

@max-sixty I was doing a similar thing last week, hit the same issue and after trying a thousand strategies I ended up working around it with a custom scheduler inside the map_blocks as you did 😢

[TomNicholas]

FYI cubed.map_blocks exists, so it might be interesting to try using that on this problem. However xr.map_blocks doesn't (yet) have the ability to defer to cubed instead of dask, so you would have to write your computation using the array-level function cubed.map_blocks.

EDIT: Also doing this

I was hoping to make the internals of this task opaque to dask, so it became a much dumber task runner — just map over the blocks, running the function and writing the result, block by block.

is essentially what I would like to do with xr.map_blocks to make it support using cubed.map_blocks. The reason I haven't made this change yet is because I don't know enough about the internals of xr.map_blocks to know why it doesn't just work by calling dask.map_blocks (I assume there is a good reason!).

[max-sixty]

Eeeek, sounds like this is a common issue then... Without wanting to blast @s anytime anything is a problem, I'll politely CC @martindurant & @mrocklin for info, as I count four alternatives that folks have built, within a few hours of me posting the issue...

(Goes without saying that I am very happy to be wrong, not trying to throw stones at the hard work your team has done)

---

Also you can write from within map_blocks and return 0

That would work well, but it needs the indexes, and that's effectively prevented by #8409, because we get >1GB task graphs.

---

You have a shuffle like thing going from Tower 3 to Tower 4. That's probably it. What's happening there?

Yes, good observation.

Speculatively — I think this is splitting the source read by data variable, then shuffling to do the calc, then writing each destination data variable.

If that's correct, then another approach could be to try and keep the tasks for each block together, rather than organizing by each data variable.

I'll try and get to checking whether that assertion is correct.

---

FYI cubed.map_blocks exists, so it might be interesting to try using that on this problem.

I'm very up for checking this out, and some of the team pointed me towards it (hence my DM last week @TomNicholas !). I am way over on trying new things vs getting work done, and I would need to set up an env for running cubed.

I'd also like to try beam!

[mrocklin]

I'll defer to @fjetter, who has been playing a lot with Dask graph ordering recently and having decent success. He may have a new version for you to try.

If you're able to provide an mcve that recreates your currently unhappy graph I'd encourage you to do so. There's pretty good activity in this space right now.

[dcherian]

I think this is splitting the source read by data variable, then shuffling to do the calc, then writing each destination data variable.
If that's correct, then another approach could be to try and keep the tasks for each block together, rather than organizing by each data variable.

IDK you have one dask array per variable anyway. But maybe the viz is misleading. This kind of shuffle type workload is where distributed falls apart. Also why are you splitting by variable? map_blocks can handle datasets.

In general, I'm also confused as to how a shuffle type workload gets expressed as a blockwise workload without a major rechunking/rearrangement (i.e. a shuffle) in the middle.

that's effectively prevented by #8409, because we get >1GB task graphs.

Could try with the current state of #8412

[max-sixty]

IDK you have one dask array per variable anyway. But maybe the viz is misleading. This kind of shuffle type workload is where distributed falls apart. Also why are you splitting by variable? map_blocks can handle datasets.

It wasn't me splitting by variable! 😄 dask seems to be doing it.

I'll try and get together an MCVE...

[max-sixty]

OK, here's an MCVE. It requires more specificity than I expected — I think my read that things were working last week was correct actually, because the calcs were slightly different. It relies on the nested calcs, such that return d doesn't trigger the issue.

(though that was the beauty of .map_blocks for me — that it made the graph opaque to dask, and small changes in the inputs couldn't cause catastrophic runtime impact)

ds = xr.Dataset(
    data_vars=dict(
        a=(("x", "y"), np.arange(80).reshape(8, 10)),
        b=(("y"), np.arange(10)),
    )
).chunk(x=1)

def f(ds):
    d = ds.a * ds.b
    return (d * ds.b).sum("y")

result = ds.map_blocks(f).compute()

Here's the graph — you can see the lack of basic parallelism that @dcherian pointed out above — though obv less severe than the big example above:

[TomNicholas]

Excellent @max-sixty!

@tomwhite we should add this to the distributed arrays repo examples, and try it with cubed

[dcherian]

This is basically the same as pangeo-data/distributed-array-examples#2 which @fjetter tried to fix with dask/dask#10535

EDIT: And yeah the viz was misleading. It's not a shuffle, it's blockwise ops on different arrays, but only purely blockwise if you're unchunked on y

[TomNicholas]

@dcherian right yeah, I see that now. We should still try it out!