Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix flaky data initialization test. (#6318)

Browse Source
This commit is contained in:
Jiaming Yuan 2020-10-30 03:11:22 +08:00 committed by GitHub
parent 608bda7052
commit c80657b542
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/python/test_with_dask.py
View File

@ -768,7 +768,9 @@ class TestDaskCallbacks:
def worker_fn(worker_addr, data_ref): def worker_fn(worker_addr, data_ref):
with xgb.dask.RabitContext(rabit_args): with xgb.dask.RabitContext(rabit_args):
local_dtrain = xgb.dask._dmatrix_from_worker_map(**data_ref) 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( futures = client.map(
worker_fn, workers, [m.create_fn_args()] * len(workers), worker_fn, workers, [m.create_fn_args()] * len(workers),