From a30176907f5628020f8238464c2f58907ee724d1 Mon Sep 17 00:00:00 2001 From: Philip Hyunsu Cho Date: Thu, 27 Jun 2019 20:42:35 -0700 Subject: [PATCH] Support Dask 2.0 (#4617) --- python-package/xgboost/dask.py | 6 +++++- tests/ci_build/Dockerfile.cpu | 2 +- tests/ci_build/Dockerfile.gpu | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py index b0be2dbcc..e63c5c124 100644 --- a/python-package/xgboost/dask.py +++ b/python-package/xgboost/dask.py @@ -91,7 +91,11 @@ def create_worker_dmatrix(*args, **kwargs): def _run_with_rabit(rabit_args, func, *args): - os.environ["OMP_NUM_THREADS"] = str(distributed_get_worker().ncores) + worker = distributed_get_worker() + try: + os.environ["OMP_NUM_THREADS"] = str(worker.ncores) + except AttributeError: + os.environ["OMP_NUM_THREADS"] = str(worker.nthreads) try: rabit.init(rabit_args) result = func(*args) diff --git a/tests/ci_build/Dockerfile.cpu b/tests/ci_build/Dockerfile.cpu index 4e7fcfc59..3101e0c88 100644 --- a/tests/ci_build/Dockerfile.cpu +++ b/tests/ci_build/Dockerfile.cpu @@ -24,7 +24,7 @@ RUN \ pip install pyyaml cpplint pylint astroid sphinx numpy scipy pandas matplotlib sh recommonmark guzzle_sphinx_theme mock \ breathe matplotlib graphviz pytest scikit-learn wheel kubernetes urllib3 && \ pip install https://h2o-release.s3.amazonaws.com/datatable/stable/datatable-0.7.0/datatable-0.7.0-cp37-cp37m-linux_x86_64.whl && \ - conda install dask + conda install dask=2.0.0 # Install lightweight sudo (not bound to TTY) RUN set -ex; \ diff --git a/tests/ci_build/Dockerfile.gpu b/tests/ci_build/Dockerfile.gpu index c0a21f517..577ee62e9 100644 --- a/tests/ci_build/Dockerfile.gpu +++ b/tests/ci_build/Dockerfile.gpu @@ -17,7 +17,7 @@ ENV PATH=/opt/python/bin:$PATH # Install Python packages RUN \ pip install numpy pytest scipy scikit-learn pandas matplotlib wheel kubernetes urllib3 graphviz && \ - conda install dask + conda install dask=2.0.0 ENV GOSU_VERSION 1.10