Perhaps we should make it possible to avoid unnecessary logic for "simple" cases like zeros/full/ones?

This block_info dict is modeled after that in Array.map_blocks (though it's not identical, since that can include dtypes). My thinking was that, since they are similar operations, the API should be similar, and similar data should be provided to the tasks.

That being said, it's true that in simpler cases there is some unnecessary logic here when constructing the block info. We could revert these changes to keep CreateArrayDeps simple, and then make a new blockwise layer for things like from_array. So to me it seems the choice is between having a larger number of more specific layer classes, or a smaller number of more general ones like this one.

Or __init__ could take something like block_info=True to indicate that it should provide more than just the chunk shape. I generally like to have more predictable outputs for functions, but can be talked out of it :)

I generally like to have more predictable outputs for functions, but can be talked out of it :)

Makes sense - I do prefer consistency (as you already have it), so I'm not really motivated to talk you out of this unless we discover a need for performance optimizations.

It may be easiest to separately serialize each seed, so that we can still do block_info["seed"] = self.seeds[n] in the getitem definition without needed any __dask_distributed_unpack__ logic in this class. However, we would then need to wrap the underlying creation function to handle deserialization of block_info["seed"].

If we do use a function wrapper, it probably makes sense to fully move the random_state_data operation from the client to the worker (and avoid the need to ship the seeds through the scheduler altogether). Do we know any RNG experts? Can we simply shift the default/random seed by some factor related to the chunk index?

Yeah, I haven't fully thought through the serialization story here. But at least, I think that serialize does drill down into the list and serialize the seeds separately, c.f. https://github.com/dask/distributed/blob/bef0308962345303123aba7ec6730757a61a4dfc/distributed/protocol/serialize.py#L263-L284 , and it shouldn't be necessary to deserialize them on the scheduler.

It would definitely be nice to be able to move a smaller representation of the random state to the workers, but at least this implementations doesn't ship more seed data than the current implementation.

But at least, I think that serialize does drill down into the list and serialize the seeds separately, c.f. https://github.com/dask/distributed/blob/bef0308962345303123aba7ec6730757a61a4dfc/distributed/protocol/serialize.py#L263-L284 , and it shouldn't be necessary to deserialize them on the scheduler.

Ah - You are correct that serialize will internally handle the elements of seeds seperately. I hadn't considered the possibility of simply splitting the header/frames on the scheduler, but that may actually be a good way to go.

Regarding the question of moving random state generation to the workers: I like the idea, but I also don't think you should worry about it in this PR.

I was just experimenting with this. It seems that we will need to add an explicit iterate_collection= argument to serialize. Otherwise, we cannot guarentee that the scheduler will be able to access individual elements of a list/dict without needing to deserialize. I can/will submit a PR to distributed to add this feature, but this means a proper __dask_distributed_unpack__ soution here will not work with the latest release of distributed.

I can/will submit a PR to distributed to add this feature, but this means a proper dask_distributed_unpack soution here will not work with the latest release of distributed.

On second thought, I don't think it makes sense to do this. In general, it seems like a much cleaner solution to just call serialize seperately on each element on the client, rather than adding an extra step of splitting the header/frames on the scheduler.

On second thought, I don't think it makes sense to do this. In general, it seems like a much cleaner solution to just call serialize seperately on each element on the client,

Yep, after taking another look at it, I agree with you

rather than adding an extra step of splitting the header/frames on the scheduler.

Can you elaborate more on what you mean here?

In general, I'm having a tough time reasoning about how deserialization should look on the worker. What would the idiomatic way of doing that look like? i.e.,

1. How does one identify whether one should deserialize? If it is coming back as distributed.protocol.Serialized, what would be the best way to identify that without relying on distributed being installed?
2. I ran into a couple of problems around making sure the numpy deserializer is registered on the worker. What am I doing wrong?

(I'm not demanding answers to the above, just what I'm thinking about :) )

rather than adding an extra step of splitting the header/frames on the scheduler.

Can you elaborate more on what you mean here?

If you call serialize on an entire collection, you will get a (header, frames) tuple, and so the scheduler will need to process/split the header and then convert this single tuple to a list of (header, frame) tuples. The logic is not that complicated, but it is messier than simply serializing the individual elements to begin with.

In general, I'm having a tough time reasoning about how deserialization should look on the worker. What would the idiomatic way of doing that look like? i.e.,

I have experimented with ways to do this, but I haven't established anything that I would call idiomatic yet...

How does one identify whether one should deserialize? If it is coming back as distributed.protocol.Serialized, what would be the best way to identify that without relying on distributed being installed?
I ran into a couple of problems around making sure the numpy deserializer is registered on the worker. What am I doing wrong?

1. How does one identify whether one should deserialize? If it is coming back as distributed.protocol.Serialized, what would be the best way to identify that without relying on distributed being installed?

You have certainly isolated the crux of the problem here. It is easy to decide when to serialize (within an __dask_distributed_pack__ definition). However, it is not obvious when the worker will need to deserialize one or all of its arguments. In my first pass, I was relying on pickle rather than serialize/deserialize, so the function simply needed to check if the argument(s) in question was a bytes object. In 7415, I just moved over to serialize/deserialize, and I am relying on an explicit "serialized" label. I don't "love" that solution, but it seems to work.

2. I ran into a couple of problems around making sure the numpy deserializer is registered on the worker. What am I doing wrong?

Not sure - but I'm hoping we wont need to do anything like this to check if the data is serialized.

Yeah, I am probably doing something wrong, but I am finding that the first time I need to deserialize a numpy array on a worker it fails because that protocol hasn't been registered yet. Subsequent deserializations work fine.

I'm currently forcing the registration by relying on the side effect of this import, but I don't view that as a real solution. Perhaps @jrbourbeau has some insight as to what might be going wrong here?

@rjzamora wanted to flag this as where we'll have to think through some of the array inlining logic.

Got it. So, the original code allows you to include the array in the graph once (with a dedicated key), or to inline it in every task. So far, this PR is effectively inlining the array in every task by including it in the IO function (in graph_from_arraylike). Is that correct?

It seems like you could support inline_array=False, but you may need to further expand CreateArrayDeps to include optional args/kwargs that should be the same for all keys (and make it possible to specify these arguments in BlockwiseCreateArray). I guess if you can do this, you could avoid using partial to embed the array in io funciton altogether?

Got it. So, the original code allows you to include the array in the graph once (with a dedicated key), or to inline it in every task. So far, this PR is effectively inlining the array in every task by including it in the IO function (in graph_from_arraylike). Is that correct?

It seems like you could support inline_array=False, but you may need to further expand CreateArrayDeps to include optional args/kwargs that should be the same for all keys (and make it possible to specify these arguments in BlockwiseCreateArray). I guess if you can do this, you could avoid using partial to embed the array in io funciton altogether?

Yes, you have this exactly right -- I'm planning to re-introduce non-inlining, but haven't tackled it yet. For the case of lazy arrays like Zarr, or small ones, this should already do what we want.