Stealing sensitive tests

[fjetter]

We'd like to have a couple of tests that are related to stealing sensitive workloads.

1. Very simple example, currently causing problems (Root-ish tasks all schedule onto one worker dask/distributed#6573)

import time
import dask
import distributed

# dask.config.set({"distributed.scheduler.work-stealing": False})
client = distributed.Client(n_workers=4, threads_per_worker=1)
root = dask.delayed(lambda n: "x" * n)(dask.utils.parse_bytes("1MiB"), dask_key_name="root")
results = [dask.delayed(lambda *args: None)(root, i) for i in range(10000)]
dask.compute(results)

2. Upscaling situation (example from Poor work scheduling when cluster adapts size dask/distributed#4471)

import time
import distributed
import dask.array as da
import dask


client = distributed.Client(n_workers=1, threads_per_worker=1, memory_limit='1G')
print(client.dashboard_link)

# Slow task.
def func1(chunk):
    if sum(chunk.shape) != 0: # Make initialization fast
        time.sleep(5)
    return chunk

def func2(chunk):
    return chunk

data = da.zeros((30, 30, 30), chunks=5)
result = data.map_overlap(func1, depth=1, dtype=data.dtype)
result = result.map_overlap(func2, depth=1, dtype=data.dtype)
future = client.compute(result)

print('started computation')

time.sleep(11)
# print('scaling to 4 workers')
# client.cluster.scale(4)

time.sleep(5)
print('scaling to 20 workers')
client.cluster.scale(20)

_ = future.result()

3. An example that shows a very inhomogeneous compute. This should mimic a long tail in a computation graph where work stealing should reduce wall time, e.g. by using an asymmetric distribution of sleep values