Hang on dask.array Program with 2 Nodes

I have a very simple implementation of linear regression (attached, see below) using the dask.array interface that was working a few months ago with Dask 1.14 but I've since upgraded and now it is hanging with Dask 2.6 when I run it on two nodes. If I run on just a single node it works fine. If I go back to Dask 1.14 it also works. I always see the program run at least a few iterations (usually 4-6) before hanging. I flush all stdout calls so I know precisely how far I've made it before things hang.

I'm running on two NVIDIA DGX boxes each of which are connected by both ethernet and Infiniband (it's unclear to me which network is being used for UDP/TCP communication). Both nodes are running Ubuntu 18.04 LTS and have a shared IBM GPFS file system. I'm using python 3.6 from anaconda. I setup my Dask cluster from inside my slurm job as follows:

scontrol show hostname > hostfile.txt
dash-ssh --hostfile hostfile.txt --nthreads 8 --nprocs 1 &

where hostfile.txt contains the network names of the two nodes. I run my program as follows:

python hang_linreg.py -s 1 -i 100 -f 32 -n 100 --cluster $SLURMD_NODENAME:8786 --profile linreg_hang.html

Given that this still works with an older version of Dask, but hangs with the most recent version I'm inclined to believe that the issue is with Dask itself. Please let me know if there is additional information that I can provide to help with debugging.

[hang_linreg.txt](https://github.com/dask/distributed/files/3747644/hang_linreg.txt)