Update Scheduler.rebalance return value when data is missing

[jrbourbeau]

This PR improves error handling in Scheduler.rebalance when keys which are being rebalanced are found to be missing on workers. Currently when we attempt to return the missing data keys here

distributed/distributed/scheduler.py

Line 3106 in 4f11509

"keys": sum([r["keys"] for r in result if "keys" in r], []),

we get the following error:

TypeError: can only concatenate list (not "dict") to list

because r["keys"] is a dictionary. We update things here to return a tuple of keys.

Example to trigger missing data in rebalance:

import time
import random
from itertools import count

import dask.array as da
from distributed import Client

with Client() as client:
    counter = count()
    while next(counter) < 10:
        x = da.random.normal(size=100, chunks=1)
        x = x.persist()
        
        time.sleep(random.random() * 0.1)
        x = None
        client.rebalance()

Note: there were also several trailing commas that black wanted to add

[mrocklin]

Thanks @jrbourbeau . Should we add some variant of the example as a test?

[jrbourbeau]

I added a simplified version of the example as a test. It relies on the cluster deleting Futures while rebalancing is taking place, which seems less reliable than I usually like for unit tests. That said, I haven't been able to think of a more direct test and it's been failing consistently on my laptop. If CI passes we should be good to merge this and I'll keep an eye out for if this test ends up failing on PRs

[mrocklin]

Thanks @jrbourbeau