Use BlockwiseDep for map_blocks with block_id or block_info

[bmerry]

When the map_blocks function has block_id or block_info arguments, it is populated with information about the current block. Previously this was done by synthesizing all the information at the time map_blocks was called and stored in a da.Array of objects. Now that Blockwise supports I/O-deps, it's feasible to defer this work until the Blockwise is materialised.

I've made this a draft PR because I have no idea what I'm doing with the serialization and need some advice. I've made a distributed test pass, but after seeing some comments in other PRs I'm aware that there are some subtleties about what may be deserialized on the scheduler (in this case I think an np.ndarray is getting deserialized there).

Maybe a solution would be to have _BlockInfo.__getitem__ generate a task to generate the block_info, rather than generating it itself? That would also have the advantage that constructing the block info would be left to the workers.

I've extended da.blockwise to accept BlockwiseDep arguments (and documented it), so that map_blocks can pass them through (since map_blocks calls da.blockwise rather than dask.blockwise). I probably still need to add some direct tests, rather than relying on the coverage from map_blocks.

The performance looks good. Here's some benchmark code:

#!/usr/bin/env python3

from time import monotonic

import dask.array as da
import numpy as np


def func(a, b, block_info=None):
    return a + b


a = da.from_array(np.arange(200)[np.newaxis, :], chunks=1)
b = da.from_array(np.arange(200)[:, np.newaxis], chunks=1)
t0 = monotonic()
c = da.map_blocks(func, a, b)
t1 = monotonic()
c.compute()
t2 = monotonic()
print(f'Prepare: {t1 - t0:.3f}')
print(f'Compute: {t2 - t1:.3f}')

Before:

Prepare: 0.446
Compute: 7.457

After:

Prepare: 0.001
Compute: 6.715

• Tests added / passed
• Passes black dask / flake8 dask / isort dask

[bmerry]

cc @rjzamora @dask/array

[bmerry]

@madsbk any tips on the serialization? If I understand correctly (which I'm not at all sure I do), to_serialize says "serialize how you like on the client, but don't deserialize on the scheduler, pass it to the worker to deserialize." However, I don't know how to get the worker to do the deserialization: I end up with a Serialize (or Serialized, I forget) object in the graph instead of the original object. This might be because it's nested inside a task.

[madsbk]

to_serialize says "serialize how you like on the client, but don't deserialize on the scheduler, pass it to the worker to deserialize."

Yes this is correct. The worker should deserialize the data automatically but the serialization logic in Distributed is very messy so it is hard to know exactly why it doesn't happen :/
One way to make the worker deserialize a to_serialize object is to give the object as an argument to the materialized tasks.

FYI: I am working on redesigning the serialization logic in Distributed, which will handle issues like this. E.g. the to_serialize function will be an optional flag to avoid unnecessary deserializations, it will not be required for correctness.

[bmerry]

FYI: I am working on redesigning the serialization logic in Distributed, which will handle issues like this.

Sounds good. What's the timeframe, and is there a ticket I can watch so that I'll know when it's in place? If it's happening soon I might just leave this until then.

One way to make the worker deserialize a to_serialize object is to give the object as an argument to the materialized tasks.

To check that I understand you correctly: if X is a to_serialize object, and the layer materialization produces a task like this:

(func, X, 123)

it should work, whereas a task

(dict, [["key1", X], ["key2", 123]])

might not? (the latter is approximately what I'm currently doing, although there are two layers of dict involved).

[bmerry]

I've updated the PR with a different approach - see commit message of bc2b235 for details. @rjzamora I feel like this new approach might not be in the "spirit" of BlockwiseDep, in that I'm just using an empty BlockwiseDepDict to access the block index, and all the logic for processing that block index into an input to the task is handled by a wrapper function and some constant (indices is None) arguments to the Blockwise. Let me know if you think I should revert bc2b235.

[bmerry]

I've added some extra tests, so I think this is ready for review. I'm still not sure if I'm on the right path with serialization - it seems a bit different to the related code.

[bmerry]

The failing tests seem to be flaky and not related to this PR.