[dask] Use nthread in DMatrix construction. (#7337)

This is consistent with the thread overriding behavior.
2021-10-20 15:16:40 +08:00
parent b8e8f0fcd9
commit f999897615
3 changed files with 44 additions and 33 deletions
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -1275,7 +1275,9 @@ class TestWithDask:

                def worker_fn(worker_addr: str, data_ref: Dict) -> None:
                    with xgb.dask.RabitContext(rabit_args):
-                        local_dtrain = xgb.dask._dmatrix_from_list_of_parts(**data_ref)
+                        local_dtrain = xgb.dask._dmatrix_from_list_of_parts(
+                            **data_ref, nthread=7
+                        )
                        total = np.array([local_dtrain.num_row()])
                        total = xgb.rabit.allreduce(total, xgb.rabit.Op.SUM)
                        assert total[0] == kRows