[dask] Use DMLC_TASK_ID. (dmlc#5415)

trivialfis · Mar 24, 2020 · eff45ae · eff45ae
1 parent 01a1f80
commit eff45ae
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py
@@ -69,6 +69,9 @@ class RabitContext:
     '''A context controling rabit initialization and finalization.'''
     def __init__(self, args):
         self.args = args
+        worker = distributed_get_worker()
+        self.args.append(
+            ('DMLC_TASK_ID=[xgboost.dask]:' + str(worker.address)).encode())
 
     def __enter__(self):
         rabit.init(self.args)