Dask-cuda: Shuffle benchmark fails while getting transfer logs w/ ucx 1.11

[ayushdg]

The shuffle benchmark seems to be failing with UCX1.11 with the NVLink + no IB configuration. From the looks of the error message the error originates while getting the worker transfer logs rather than the benchmark itself.

Error stacktrace:

Traceback (most recent call last):
  File "local_cudf_shuffle.py", line 247, in <module>
    main(parse_args())
  File "local_cudf_shuffle.py", line 125, in main
    incoming_logs = client.run(lambda dask_worker: dask_worker.incoming_transfer_log)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/client.py", line 2543, in run
    return self.sync(self._run, function, *args, **kwargs)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/client.py", line 859, in sync
    return sync(
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/utils.py", line 326, in sync
    raise exc.with_traceback(tb)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/utils.py", line 309, in f
    result[0] = yield future
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/tornado/gen.py", line 762, in run
    value = future.result()
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/client.py", line 2463, in _run
    responses = await self.scheduler.broadcast(
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/core.py", line 866, in send_recv_from_rpc
    result = await send_recv(comm=comm, op=key, **kwargs)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/core.py", line 667, in send_recv
    raise exc.with_traceback(tb)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/core.py", line 502, in handle_comm
    result = await result
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/scheduler.py", line 5587, in broadcast
    results = await All(
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/utils.py", line 210, in All
    result = await tasks.next()
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/scheduler.py", line 5582, in send_message
    resp = await send_recv(comm, close=True, serializers=serializers, **msg)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/core.py", line 651, in send_recv
    response = await comm.read(deserializers=deserializers)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/distributed/comm/ucx.py", line 300, in read
    await self.ep.recv(each_frame)
  File "/opt/conda/envs/cugraph_bench/lib/python3.8/site-packages/ucp/core.py", line 704, in recv
    ret = await comm.tag_recv(self._ep, buffer, nbytes, tag, name=log)
ucp.exceptions.UCXError: <[Recv #077] ep: 0x7f69b0d1bb40, tag: 0x34fc41eda976c804, nbytes: 84165, type: <class 'numpy.ndarray'>>: Connection reset by remote peer

Command to run the benchmark:

 UCX_MAX_RNDV_RAILS=1 python local_cudf_shuffle.py -t gpu -p ucx --enable-tcp-over-ucx --enable-nvlink --disable-infiniband --disable-rdmacm --runs 20 --in-parts 16 --partition-size 1GB --disable-rmm-pool --no-silence-logs --scheduler-address ucx://...

Note: This doesn't happen in the UCX+IB configuration case

[pentschev]

Thanks for reporting this Ayush. This could be the same as reported in https://github.com/openucx/ucx/issues/6922 , that occurs when a worker sends some data and immediately closes the endpoint, so it's important to assess the chain of events in your case. Could you elaborate at what stage does that happen? You mentioned it happens when a worker transfer logs, I'm not sure what exactly that means.

[ayushdg]

Based on gpu activity it seems like this happens right after all the runs from the shuffle call are completed. Based on the error logs it seems to always happen at this function: incoming_logs = client.run(lambda dask_worker: dask_worker.incoming_transfer_log)

[pentschev]

Is that deterministic? How many workers is your cluster running/what's the minimum number of workers you've seen this error with?

[ayushdg]

Is that deterministic? How many workers is your cluster running/what's the minimum number of workers you've seen this error with?

The error was deterministic (tested over 3 runs), cluster consisted of 16 workers (this is on a single node dgx2 instance)

[pentschev]

Could you also let me know about your RMM and spilling policies? I've been trying to run the benchmark with the same arguments as you but it always OOMs. I've been using a LocalCUDACluster instead of an external cluster that the benchmark connects to with --scheduler-address.

[ayushdg]

Yeah here's the current configuration used to startup the scheduler and workers

export WORKER_RMM_POOL_SIZE=22G
export DASK_UCX__CUDA_COPY=True
export DASK_UCX__TCP=True
export DASK_UCX__NVLINK=True
export DASK_UCX__INFINIBAND=False
export DASK_UCX__RDMACM=False
export DASK_RMM__POOL_SIZE=0.5GB
export UCX_MAX_RNDV_RAILS=1
export UCX_TCP_CM_REUSEADDR=y

SCHEDULER_ARGS="--protocol ucx
                --scheduler-file $SCHEDULER_FILE
               "

WORKER_ARGS="--enable-tcp-over-ucx
             --enable-nvlink 
             --disable-infiniband
             --disable-rdmacm
             --rmm-pool-size=$WORKER_RMM_POOL_SIZE
             --scheduler-file $SCHEDULER_FILE
            "

The command used to run the dask-cuda benchmark explicitly sets --disable-rmm-pool to prevent overriding the rmm pool setup defined here

[pentschev]

Thanks @ayushdg for the details, the reason I was getting OOM was because I forgot to specify I wanted the benchmark to use all the devices. I can reproduce it consistently and without need to separately spin up the cluster with the following command:

$ UCX_MAX_RNDV_RAILS=1 python local_cudf_shuffle.py -t gpu -p ucx --enable-tcp-over-ucx --enable-nvlink --disable-infiniband --disable-rdmacm --runs 20 --in-parts 16 --partition-size 1GB --rmm-pool-size 22G --no-silence-logs -d 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15

I'm reasonably sure the cause for this is indeed https://github.com/openucx/ucx/issues/6922, I was able to write a functional workaround, it simply adds a await asyncio.sleep(0.1) just before the ep.abort() call from UCX-Py. Would you be able to test a branch with those changes in https://github.com/pentschev/distributed/tree/workaround-ucx-connection-reset-by-peer ? I can't guarantee this is a fully-functionaly workaround, as I image sleep(0.1) may not be enough if the other endpoint is still receiving large buffers when UCX.close() is called.

[ayushdg]

Happy to test this workaround at larger scales and report the outcome!

[ayushdg]

Can confirm that using the workaround implemented in https://github.com/pentschev/distributed/tree/workaround-ucx-connection-reset-by-peer seems to fix the error, both on single node as well as at scale.

[pentschev]

@ayushdg thanks for confirming. I've been able to test a proper UCX fix for that issue and it seems to resolve it. It's now being worked to get into upstream UCX and then backported to 1.11, I'll then let you know so you can verify it on your end as well, I hope it will be ready early next week and would be great to get you testing that still next week to ensure we have no other issues remaining before the 1.11 release scheduler for August.