From bd81ff09af2054a0d1fa723e84f3dc32c5045137 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 3 Jul 2020 00:36:11 +0800
Subject: [PATCH] Cleanup.

---
 python-package/xgboost/dask.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py
index f7c40ff28519..65cfa85896d3 100644
--- a/python-package/xgboost/dask.py
+++ b/python-package/xgboost/dask.py
@@ -34,6 +34,11 @@
 from .sklearn import XGBModel, XGBRegressorBase, XGBClassifierBase
 from .sklearn import xgboost_model_doc
 
+try:
+    from distributed import Client
+except ImportError:
+    Client = None
+
 # Current status is considered as initial support, many features are
 # not properly supported yet.
 #
@@ -360,7 +365,7 @@ def get_worker_data_shape(self, worker):
             cols = c
         return (rows, cols)
 
-from distributed import Client
+
 def _get_rabit_args(worker_map, client: Client, host_ip=None, port=None):
     '''Get rabit context arguments from data distribution in DaskDMatrix.'''
     msg = 'Please provide both IP and port'