from dask.distributed import get_client, Client

from prefect import Flow, task
from prefect.executors import DaskExecutor

# If n_maps >= n_workers the flow will hang
n_workers = 3
n_maps = 3

client = Client(n_workers=n_workers)
executor = DaskExecutor(client.scheduler_info()['address'])

def add(a, b):
    return a + b

@task
def many_computations(n_many):
    client = get_client()

    futures = [client.submit(add, i, i + 1) for i in range(n_many)]
    values = [f.result() for f in futures]
    values = client.gather(futures)
    total = sum(values)
    
    return total

@task
def sum_task(numbers):
    return sum(numbers)    

map_over = [i for i in range(100, 100+n_maps)]

with Flow("test_flow") as flow:
    mapped_results = many_computations.map(map_over)
    total = sum_task(mapped_results)

status = flow.run(executor=executor)