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

[dask] speed up tests (#7020)

Browse Source
This commit is contained in:
jmoralez 2021-06-10 22:43:01 -05:00 committed by GitHub
parent f79cc4a7a4
commit 25514e104a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tests/python/test_with_dask.py
View File

@ -28,7 +28,6 @@ if tm.no_dask()['condition']:
pytest.skip(msg=tm.no_dask()['reason'], allow_module_level=True) pytest.skip(msg=tm.no_dask()['reason'], allow_module_level=True)
from distributed import LocalCluster, Client from distributed import LocalCluster, Client
from distributed.utils_test import client, loop, cluster_fixture
import dask.dataframe as dd import dask.dataframe as dd
import dask.array as da import dask.array as da
from xgboost.dask import DaskDMatrix from xgboost.dask import DaskDMatrix
@ -40,6 +39,20 @@ else:
suppress = hypothesis.utils.conventions.not_set # type:ignore suppress = hypothesis.utils.conventions.not_set # type:ignore
@pytest.fixture(scope='module')
def cluster():
with LocalCluster(
n_workers=2, threads_per_worker=2, dashboard_address=None
) as dask_cluster:
yield dask_cluster
@pytest.fixture
def client(cluster):
with Client(cluster) as dask_client:
yield dask_client
kRows = 1000 kRows = 1000
kCols = 10 kCols = 10
kWorkers = 5 kWorkers = 5