Attempt at repro for multi-array return

multi_arr_repro.py

import sys

import dask
import dask.array
from dask.delayed import delayed 

import distributed

try:
    from dask_cuda import LocalCUDACluster as DaskCluster
    client_kwargs = dict(protocol="ucx")
except ImportError:
    from distributed import LocalCluster as DaskCluster
    client_kwargs = dict(protocol="tcp")


try:
    import cupy as xnumpy
except ImportError:
    import numpy as xnumpy


@delayed
def double_halve(x):
    return 2 * x, x / 2


def get_client():
    cluster = DaskCluster()
    client = distributed.Client(cluster, **client_kwargs)
    return client


def main(*argv):
    client = get_client()

    rs = dask.array.random.RandomState(RandomState=xnumpy.random.RandomState)
    a = rs.random((20,), chunks=(5,))

    arrs = [double_halve(e) for e in a.blocks]

    futures = client.compute(arrs)
    results = client.gather(futures)

    print("")
    print(results)
    print("")

    return 0


if __name__ == "__main__":
    sys.exit(main(*sys.argv))