[dask] Use DMLC_TASK_ID. (#5415)

2020-03-15 16:47:03 +08:00 · 2020-03-15 16:47:03 +08:00 · 668e432e2d
commit 668e432e2d
parent fc88105620
1 changed files with 3 additions and 0 deletions
--- 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)