[dask] Workarounds for different Dask versions. (#11436)

---------

Co-authored-by: TomAugspurger <toaugspurger@nvidia.com>

Author: trivialfis

python-package/xgboost/dask/__init__.py

@@ -55,7 +55,7 @@
 import logging
 from collections import defaultdict
 from contextlib import contextmanager
-from functools import partial, update_wrapper
+from functools import cache, partial, update_wrapper
 from threading import Thread
 from typing import (
     Any,
@@ -85,6 +85,8 @@
 from dask import dataframe as dd
 from dask.delayed import Delayed
 from distributed import Future
+from packaging.version import Version
+from packaging.version import parse as parse_version
 
 from .. import collective, config
 from .._typing import FeatureNames, FeatureTypes, IterationRange
@@ -171,6 +173,21 @@
 LOGGER = logging.getLogger("[xgboost.dask]")
 
 
+@cache
+def _DASK_VERSION() -> Version:
+    return parse_version(dask.__version__)
+
+
+@cache
+def _DASK_2024_12_1() -> bool:
+    return _DASK_VERSION() >= parse_version("2024.12.1")
+
+
+@cache
+def _DASK_2025_3_0() -> bool:
+    return _DASK_VERSION() >= parse_version("2025.3.0")
+
+
 def _try_start_tracker(
     n_workers: int,
     addrs: List[Union[Optional[str], Optional[Tuple[str, int]]]],
@@ -1491,6 +1508,18 @@ async def _predict_async(
         )
         if isinstance(predts, dd.DataFrame):
             predts = predts.to_dask_array()
+            # Make sure the booster is part of the task graph implicitly
+            # only needed for certain versions of dask.
+            if _DASK_2024_12_1() and not _DASK_2025_3_0():
+                # Fixes this issue for dask>=2024.1.1,<2025.3.0
+                # Dask==2025.3.0 fails with:
+                #     RuntimeError: Attempting to use an asynchronous
+                #     Client in a synchronous context of `dask.compute`
+                #
+                # Dask==2025.4.0 fails with:
+                #     TypeError: Value type is not supported for data
+                #     iterator:<class 'distributed.client.Future'>
+                predts = predts.persist()
         return predts
 
     @_deprecate_positional_args