Dask 2025.4.0 scheduler info compatibility (#11462)

Author: jrbourbeau

python-package/xgboost/dask/__init__.py

@@ -254,15 +254,11 @@ def __init__(self, **args: CollArgsVals) -> None:
         super().__init__(**args)
 
         worker = distributed.get_worker()
-        with distributed.worker_client() as client:
-            info = client.scheduler_info()
-            w = info["workers"][worker.address]
-            wid = w["id"]
         # We use task ID for rank assignment which makes the RABIT rank consistent (but
         # not the same as task ID is string and "10" is sorted before "2") with dask
-        # worker ID. This outsources the rank assignment to dask and prevents
+        # worker name. This outsources the rank assignment to dask and prevents
         # non-deterministic issue.
-        self.args["DMLC_TASK_ID"] = f"[xgboost.dask-{wid}]:" + str(worker.address)
+        self.args["DMLC_TASK_ID"] = f"[xgboost.dask-{worker.name}]:{worker.address}"
 
 
 def _get_client(client: Optional["distributed.Client"]) -> "distributed.Client":
@@ -923,12 +919,11 @@ def train(  # pylint: disable=unused-argument
 
     """
     client = _get_client(client)
-    args = locals()
     return client.sync(
         _train_async,
         global_config=config.get_config(),
         dconfig=_get_dask_config(),
-        **args,
+        **locals(),
     )
 
 

python-package/xgboost/testing/dask.py

@@ -7,6 +7,7 @@
 from dask import array as da
 from dask import dataframe as dd
 from distributed import Client, get_worker
+from packaging.version import parse as parse_version
 from sklearn.datasets import make_classification
 
 import xgboost as xgb
@@ -15,7 +16,7 @@
 from xgboost.testing.updater import get_basescore
 
 from .. import dask as dxgb
-from ..dask import _get_rabit_args
+from ..dask import _DASK_VERSION, _get_rabit_args
 from .data import make_batches
 from .data import make_categorical as make_cat_local
 
@@ -179,7 +180,8 @@ def get_rabit_args(client: Client, n_workers: int) -> Any:
 
 def get_client_workers(client: Client) -> List[str]:
     "Get workers from a dask client."
-    workers = client.scheduler_info()["workers"]
+    kwargs = {"n_workers": -1} if _DASK_VERSION() >= parse_version("2025.4.0") else {}
+    workers = client.scheduler_info(**kwargs)["workers"]
     return list(workers.keys())