Support Dask 2.0 (#4617)

2019-06-27 20:42:35 -07:00
parent 923e6c86ba
commit a30176907f
3 changed files with 7 additions and 3 deletions
--- 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)