Client.compute doesn't materialize delayed chunk of persisted array if provided as a dict value

[jmoralez]

Describe the issue:
When persisting a dask array and trying to compute one of its chunks, the return value is the dask key instead of the computed array if the chunk is provided as a dict value.

Minimal Complete Verifiable Example:

import dask.array as da
from dask.distributed import Client


if __name__ == '__main__':
    client = Client()
    X = da.random.random((100, 2), chunks=(100, 2)).persist()
    delayed_chunk = X.to_delayed()[0, 0]
    print(type(client.compute(delayed_chunk).result())) # numpy.ndarray
    print(type(client.compute([delayed_chunk])[0].result())) # numpy.ndarray
    print(type(client.compute({'x': delayed_chunk}).result()['x']))  # tuple (the dask key)

Anything else we need to know?:
On 2024.11.2 these all return an array.

Environment:

• Dask version: 2024.12.0
• Python version: 3.10.16
• Operating System: Ubuntu
• Install method (conda, pip, source): pip install dask[distributed]

[jrbourbeau]

Thanks for the issue and nice minimal example @jmoralez -- I'm able to reproduce. @hendrikmakait maybe this is similar to some of the other client.compute task spec stuff you did a week or two ago?

[IzerOnadimQC]

The breakage is not coming from distributed but dask itself.

I think the responsible commit is dask/dask@5b115c4. I can reproduce the behaviour described above when installing distributed 2024.12.0 along with that particular commit, but installing the previous commit (dask/dask@fe9f453) restores the old behaviour.

@phofl I don't understand these changes well enough to see the issue. Could you please take a look?