Dask 2025.4.0 scheduler info compatibility (dmlc#11462) (dmlc#11476)

Co-authored-by: James Bourbeau <[email protected]>

Author: hcho3

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
 
 
 def check_init_estimation_clf(
@@ -177,7 +178,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())