Inter-worker communication debugging

[alexander-held]

We have occasionally seen what might be inter-worker communication issues of unclear origin. Here's a small setup to possibly reproduce those. The sleep is purely there for the tasks to be long enough for auto-scaling to work and can be removed otherwise.

import dask
import dask.bag
import numpy as np
import time

NUM_TASKS = 512

x = np.linspace(1, NUM_TASKS, NUM_TASKS)
b = dask.bag.from_sequence(x, npartitions=len(x))
 
def blackbox(num):
    time.sleep(4)
    return num

futures = b.map(blackbox)
task = futures.fold(lambda x,y: x+y, split_every=8)

res, *_ = dask.compute(task, scheduler=client)

res, sum(x)  # result should be sum(x)