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

[dask] Fix nthread config with dask sklearn wrapper. (#7633)

Browse Source
This commit is contained in:
Jiaming Yuan
2022-02-08 06:38:32 +08:00
committed by GitHub
parent d152c59a9c
commit 3e693e4f97
2 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/python/test_with_dask.py
View File

@@ -563,10 +563,13 @@ def run_dask_classifier(
assert list(history.keys())[0] == "validation_0"
assert list(history["validation_0"].keys())[0] == metric
assert len(list(history["validation_0"])) == 1
config = json.loads(classifier.get_booster().save_config())
n_threads = int(config["learner"]["generic_param"]["nthread"])
assert n_threads != 0 and n_threads != os.cpu_count()
forest = int(
json.loads(classifier.get_booster().save_config())["learner"][
"gradient_booster"
]["gbtree_train_param"]["num_parallel_tree"]
config["learner"]["gradient_booster"]["gbtree_train_param"]["num_parallel_tree"]
)
if model == "boosting":
assert len(history["validation_0"][metric]) == 2
@@ -591,6 +594,7 @@ def run_dask_classifier(
np.testing.assert_allclose(single_node_proba, probas)
else:
import cupy
cupy.testing.assert_allclose(single_node_proba, probas)
# Test with dataframe, not shared with GPU as cupy doesn't work well with da.unique.