From c80657b542ead48bd317a20ea359b13a9b81933d Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Fri, 30 Oct 2020 03:11:22 +0800 Subject: [PATCH] Fix flaky data initialization test. (#6318) --- tests/python/test_with_dask.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py index f759c63bd..3bcc5c865 100644 --- a/tests/python/test_with_dask.py +++ b/tests/python/test_with_dask.py @@ -768,7 +768,9 @@ class TestDaskCallbacks: def worker_fn(worker_addr, data_ref): with xgb.dask.RabitContext(rabit_args): local_dtrain = xgb.dask._dmatrix_from_worker_map(**data_ref) - assert local_dtrain.num_row() == kRows / n_workers + total = np.array([local_dtrain.num_row()]) + total = xgb.rabit.allreduce(total, xgb.rabit.Op.SUM) + assert total[0] == kRows futures = client.map( worker_fn, workers, [m.create_fn_args()] * len(workers),