Refactor.

python-package/xgboost/dask.py

@@ -351,34 +351,31 @@ def _get_worker_parts(worker_map, meta_names, worker):
     assert list_of_parts, 'data in ' + worker.address + ' was moved.'
     assert isinstance(list_of_parts, list)
 
-    # `_get_worker_parts` is launched inside worker.  In dask side
-    # this should be equal to `worker._get_client`.
-    client = get_client()
-    list_of_parts = client.gather(list_of_parts)
-    data = None
-    labels = None
-    weights = None
-    base_margin = None
-    label_lower_bound = None
-    label_upper_bound = None
-
-    local_data = list(zip(*list_of_parts))
-    data = local_data[0]
-
-    for i, part in enumerate(local_data[1:]):
-        if meta_names[i] == 'labels':
-            labels = part
-        if meta_names[i] == 'weights':
-            weights = part
-        if meta_names[i] == 'base_margin':
-            base_margin = part
-        if meta_names[i] == 'label_lower_bound':
-            label_lower_bound = part
-        if meta_names[i] == 'label_upper_bound':
-            label_upper_bound = part
-
-    return (data, labels, weights, base_margin, label_lower_bound,
-            label_upper_bound)
+    with worker_client() as client:
+        list_of_parts = client.gather(list_of_parts)
+        local_data = list(zip(*list_of_parts))
+
+        data = local_data[0]
+        labels = None
+        weights = None
+        base_margin = None
+        label_lower_bound = None
+        label_upper_bound = None
+
+        for i, part in enumerate(local_data[1:]):
+            if meta_names[i] == 'labels':
+                labels = part
+            if meta_names[i] == 'weights':
+                weights = part
+            if meta_names[i] == 'base_margin':
+                base_margin = part
+            if meta_names[i] == 'label_lower_bound':
+                label_lower_bound = part
+            if meta_names[i] == 'label_upper_bound':
+                label_upper_bound = part
+
+        return (data, labels, weights, base_margin, label_lower_bound,
+                label_upper_bound)
 
 
 class DaskPartitionIter(DataIter):  # pylint: disable=R0902