ImportError: cannot import name 'TOTAL_MEMORY' from [BUG] 'distributed.worker' (/opt/anaconda3/lib/python3.7/site-packages/distributed/worker.py) ​

[Dendiiiii]

Describe the bug:
Creating a LocalCUDACluster results in an ImportError.

Environment:
GCP Notebook with Image: RAPIDS XGboost (with Intel® MKL-DNN/MKL and CUDA 10.0) [EXPERIMENTAL]

Code and Error Messages

from dask_cuda import LocalCUDACluster

~/miniconda/envs/rapidsai/lib/python3.7/site-packages/dask_cuda/__init__.py in <module>
----> 1 from .local_cuda_cluster import LocalCUDACluster
      2 from .dgx import DGX
      3 
      4 from ._version import get_versions
      5 __version__ = get_versions()['version']

~/miniconda/envs/rapidsai/lib/python3.7/site-packages/dask_cuda/local_cuda_cluster.py in <module>
      5 from dask.distributed import LocalCluster
      6 from distributed.utils import parse_bytes
----> 7 from distributed.worker import TOTAL_MEMORY

I saw some people solve this issue by upgrading some dependencies but this GCP environment already gives the update-to-date package versions. I tried to use conda/pip install dask-cuda, but it did not work.

I solved this issue by installing the dask-cuda from the GitHub source using:
pip install git+https://github.com/rapidsai/dask-cuda.git

Doing git clone and python setup.py install does not work. Each time you restart the kernel, the issue will pop up again.

Hope this will help.