diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py index dc5992859..476a9651a 100644 --- a/tests/python-gpu/test_gpu_with_dask.py +++ b/tests/python-gpu/test_gpu_with_dask.py @@ -6,7 +6,7 @@ import numpy as np import asyncio import xgboost import subprocess -from hypothesis import given, strategies, settings, note +from hypothesis import given, strategies, settings, note, HealthCheck from hypothesis._settings import duration from test_gpu_updaters import parameter_strategy @@ -18,6 +18,7 @@ from test_with_dask import run_empty_dmatrix_reg # noqa from test_with_dask import run_empty_dmatrix_cls # noqa from test_with_dask import _get_client_workers # noqa from test_with_dask import generate_array # noqa +from test_with_dask import suppress import testing as tm # noqa @@ -171,25 +172,30 @@ class TestDistributedGPU: run_with_dask_dataframe(dxgb.DaskDMatrix, client) run_with_dask_dataframe(dxgb.DaskDeviceQuantileDMatrix, client) - @given(params=parameter_strategy, num_rounds=strategies.integers(1, 20), - dataset=tm.dataset_strategy) - @settings(deadline=duration(seconds=120)) + @given( + params=parameter_strategy, + num_rounds=strategies.integers(1, 20), + dataset=tm.dataset_strategy, + ) + @settings(deadline=duration(seconds=120), suppress_health_check=suppress) @pytest.mark.skipif(**tm.no_dask()) @pytest.mark.skipif(**tm.no_dask_cuda()) - @pytest.mark.parametrize('local_cuda_cluster', [{'n_workers': 2}], indirect=['local_cuda_cluster']) + @pytest.mark.parametrize( + "local_cuda_cluster", [{"n_workers": 2}], indirect=["local_cuda_cluster"] + ) @pytest.mark.mgpu def test_gpu_hist( self, params: Dict, num_rounds: int, dataset: tm.TestDataset, - local_cuda_cluster: LocalCUDACluster + local_cuda_cluster: LocalCUDACluster, ) -> None: with Client(local_cuda_cluster) as client: - run_gpu_hist(params, num_rounds, dataset, dxgb.DaskDMatrix, - client) - run_gpu_hist(params, num_rounds, dataset, - dxgb.DaskDeviceQuantileDMatrix, client) + run_gpu_hist(params, num_rounds, dataset, dxgb.DaskDMatrix, client) + run_gpu_hist( + params, num_rounds, dataset, dxgb.DaskDeviceQuantileDMatrix, client + ) @pytest.mark.skipif(**tm.no_cupy()) @pytest.mark.skipif(**tm.no_dask()) diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py index 0ab17010b..6f696b8d2 100644 --- a/tests/python/test_with_dask.py +++ b/tests/python/test_with_dask.py @@ -14,7 +14,8 @@ from sklearn.datasets import make_classification import sklearn import os import subprocess -from hypothesis import given, settings, note +import hypothesis +from hypothesis import given, settings, note, HealthCheck from test_updaters import hist_parameter_strategy, exact_parameter_strategy from test_with_sklearn import run_feature_weights @@ -23,14 +24,19 @@ if sys.platform.startswith("win"): if tm.no_dask()['condition']: pytest.skip(msg=tm.no_dask()['reason'], allow_module_level=True) - -from distributed import LocalCluster, Client, get_client +from distributed import LocalCluster, Client from distributed.utils_test import client, loop, cluster_fixture import dask.dataframe as dd import dask.array as da from xgboost.dask import DaskDMatrix +if hasattr(HealthCheck, 'function_scoped_fixture'): + suppress = [HealthCheck.function_scoped_fixture] +else: + suppress = hypothesis.utils.conventions.not_set + + kRows = 1000 kCols = 10 kWorkers = 5 @@ -803,7 +809,7 @@ class TestWithDask: @given(params=hist_parameter_strategy, dataset=tm.dataset_strategy) - @settings(deadline=None) + @settings(deadline=None, suppress_health_check=suppress) def test_hist( self, params: Dict, dataset: tm.TestDataset, client: "Client" ) -> None: @@ -812,7 +818,7 @@ class TestWithDask: @given(params=exact_parameter_strategy, dataset=tm.dataset_strategy) - @settings(deadline=None) + @settings(deadline=None, suppress_health_check=suppress) def test_approx( self, client: "Client", params: Dict, dataset: tm.TestDataset ) -> None: