Faster dask tests

Author: RAMitchell

python-package/xgboost/testing/dask.py

@@ -338,13 +338,13 @@ def create_dmatrix(
         return DMatrixT(*args, **kwargs)
 
     def run(DMatrixT: Type[dxgb.DaskDMatrix]) -> None:
-        enc, reenc, y, _, _ = make_recoded(device, n_features=96)
+        enc, reenc, y, _, _ = make_recoded(device, n_features=16)
         to = get_client_workers(client)
 
         denc, dreenc, dy = (
-            dd.from_pandas(enc, npartitions=8).persist(workers=to),
-            dd.from_pandas(reenc, npartitions=8).persist(workers=to),
-            da.from_array(y, chunks=(y.shape[0] // 8,)).persist(workers=to),
+            dd.from_pandas(enc, npartitions=2).persist(workers=to),
+            dd.from_pandas(reenc, npartitions=2).persist(workers=to),
+            da.from_array(y, chunks=(y.shape[0] // 2,)).persist(workers=to),
         )
 
         Xy = create_dmatrix(DMatrixT, client, denc, dy, enable_categorical=True)

tests/test_distributed/test_with_dask/conftest.py

@@ -36,6 +36,17 @@ def client(cluster: LocalCluster) -> Generator[Client, None, None]:
         yield dask_client
 
 
+@pytest.fixture(autouse=True)
+def client_as_current(request: pytest.FixtureRequest) -> Generator[None, None, None]:
+    for name in ("client", "client_one_worker"):
+        if name in request.fixturenames:
+            dask_client = request.getfixturevalue(name)
+            with dask_client.as_current():
+                yield
+            return
+    yield
+
+
 @pytest.fixture(scope="session")
 def client_one_worker() -> Generator[Client, None, None]:
     n_threads = os.cpu_count()