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

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

Browse Source 
This is consistent with the thread overriding behavior.
This commit is contained in:
Jiaming Yuan
2021-10-20 15:16:40 +08:00
committed by GitHub
parent b8e8f0fcd9
commit f999897615
3 changed files with 44 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/python-gpu/test_gpu_with_dask.py
View File

@@ -434,7 +434,7 @@ class TestDistributedGPU:
def worker_fn(worker_addr: str, data_ref: Dict) -> None:
with dxgb.RabitContext(rabit_args):
local_dtrain = dxgb._dmatrix_from_list_of_parts(**data_ref)
local_dtrain = dxgb._dmatrix_from_list_of_parts(**data_ref, nthread=7)
fw_rows = local_dtrain.get_float_info("feature_weights").shape[0]
assert fw_rows == local_dtrain.num_col()

4
tests/python/test_with_dask.py
View File

@@ -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