Dask-cuda: [FEA] All to All Initialization

[quasiben]

Awhile @kkraus14 and I were chatting about doing an all-to-all for ucx to establish various IPC handles as part of a warmup before a workflow. With a bit of help from @jakirkham I think encoding something like the following in a function (dask_cuda.all_to_all(client)) would work (functional example below):

from itertools import combinations

from dask.distributed import Client, wait

def main():
    n_workers = 3
    client = Client(n_workers=n_workers)

    workers = list(client.scheduler_info()["workers"])
    futs = []
    for w in workers:
        bit_of_data = b"0" * 1
        data = client.map(lambda x: bit_of_data, range(1), pure=False, workers=[w])
        futs.append(data[0])

    wait(futs)

    # assert 1 key on each worker
    for w, keys in client.who_has().items():
        assert len(keys) == 1

    def f(x,y):
        pass

    new_futs = []
    for w in workers:
        for fut1, fut2 in combinations(futs, 2):
            data = client.submit(f, fut1, fut2, workers=[w], pure=False)
            new_futs.append(data)

    wait(new_futs)
    # delete noop funcs futures
    client.cancel(new_futs)

    # assert all to all has resulted in all data on every worker
    for key, w in client.has_what().items():
        assert len(w) == n_workers

if __name__ == "__main__":
    main()

Note we would need to replace bit_of_data with GPU data -- cupy.arange(10) or something smallish but larger thank 8192 bytes

[pentschev]

This seems like a good idea, I would suggest this as an option rather than enforcing it on every start. Also note that we would have to check for rendezvous numbers instead of hardcoding 8192 as is the default in UCX-Py, see https://github.com/rapidsai/ucx-py/blob/branch-0.16/ucp/__init__.py#L31 .

[quasiben]

Definitely an option!

[jakirkham]

As to checking the rendezvous threshold, we should be able to just check ucp.get_config()["RNDV_THRESH"], which gets the value that UCX itself is using.