Reduce time for some multi-gpu tests (#8288)

* Faster dask tests

* Reuse AllReducer objects in tests.

* Faster boost from prediction tests.

* Use rmm dask fixture.

* Speed up dask demo.

* mypy

* Format with black.

* mypy

* Clang-tidy

Co-authored-by: Hyunsu Philip Cho <chohyu01@cs.washington.edu>

tests/python-gpu/conftest.py

@@ -22,8 +22,8 @@ def setup_rmm_pool(request, pytestconfig):
         rmm.reinitialize(pool_allocator=True, initial_pool_size=1024*1024*1024,
                          devices=list(range(get_n_gpus())))
 
-@pytest.fixture(scope='function')
-def local_cuda_cluster(request, pytestconfig):
+@pytest.fixture(scope='class')
+def local_cuda_client(request, pytestconfig):
     kwargs = {}
     if hasattr(request, 'param'):
         kwargs.update(request.param)
@@ -31,13 +31,12 @@ def local_cuda_cluster(request, pytestconfig):
         if not has_rmm():
             raise ImportError('The --use-rmm-pool option requires the RMM package')
         import rmm
-        from dask_cuda.utils import get_n_gpus
         kwargs['rmm_pool_size'] = '2GB'
     if tm.no_dask_cuda()['condition']:
         raise ImportError('The local_cuda_cluster fixture requires dask_cuda package')
     from dask_cuda import LocalCUDACluster
-    with LocalCUDACluster(**kwargs) as cluster:
-        yield cluster
+    from dask.distributed import Client
+    yield Client(LocalCUDACluster(**kwargs))
 
 def pytest_addoption(parser):
     parser.addoption('--use-rmm-pool', action='store_true', default=False, help='Use RMM pool')