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

[CI] Fix Dask Pytest fixture (#6024)

Browse Source
This commit is contained in:
Philip Hyunsu Cho 2020-08-17 16:45:22 -07:00 committed by GitHub
parent d240463b38
commit 14d5ce712c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/python-gpu/conftest.py
View File

@ -32,14 +32,12 @@ def local_cuda_cluster(request, pytestconfig):
raise ImportError('The --use-rmm-pool option requires the RMM package') raise ImportError('The --use-rmm-pool option requires the RMM package')
import rmm import rmm
from dask_cuda.utils import get_n_gpus from dask_cuda.utils import get_n_gpus
rmm.reinitialize()
kwargs['rmm_pool_size'] = '2GB' kwargs['rmm_pool_size'] = '2GB'
if tm.no_dask_cuda()['condition']: if tm.no_dask_cuda()['condition']:
raise ImportError('The local_cuda_cluster fixture requires dask_cuda package') raise ImportError('The local_cuda_cluster fixture requires dask_cuda package')
from dask_cuda import LocalCUDACluster from dask_cuda import LocalCUDACluster
cluster = LocalCUDACluster(**kwargs) with LocalCUDACluster(**kwargs) as cluster:
yield cluster yield cluster
cluster.close()
def pytest_addoption(parser): def pytest_addoption(parser):
parser.addoption('--use-rmm-pool', action='store_true', default=False, help='Use RMM pool') parser.addoption('--use-rmm-pool', action='store_true', default=False, help='Use RMM pool')