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

---------

Co-authored-by: TomAugspurger <[email protected]>

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, wraps
+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]]]],
@@ -1476,35 +1493,33 @@ async def _predict_async(
         iteration_range: Optional[IterationRange],
     ) -> Any:
         iteration_range = self._get_iteration_range(iteration_range)
-        if self._can_use_inplace_predict():
-            predts = await inplace_predict(
-                client=self.client,
-                model=self.get_booster(),
-                data=data,
-                iteration_range=iteration_range,
-                predict_type="margin" if output_margin else "value",
-                missing=self.missing,
-                base_margin=base_margin,
-                validate_features=validate_features,
-            )
-            if isinstance(predts, dd.DataFrame):
-                predts = predts.to_dask_array()
-        else:
-            test_dmatrix: DaskDMatrix = await DaskDMatrix(  # type: ignore
-                self.client,
-                data=data,
-                base_margin=base_margin,
-                missing=self.missing,
-                feature_types=self.feature_types,
-            )
-            predts = await predict(
-                self.client,
-                model=self.get_booster(),
-                data=test_dmatrix,
-                output_margin=output_margin,
-                validate_features=validate_features,
-                iteration_range=iteration_range,
-            )
+        # Dask doesn't support gblinear and accepts only Dask collection types (array
+        # and dataframe). We can perform inplace predict.
+        assert self._can_use_inplace_predict()
+        predts = await inplace_predict(
+            client=self.client,
+            model=self.get_booster(),
+            data=data,
+            iteration_range=iteration_range,
+            predict_type="margin" if output_margin else "value",
+            missing=self.missing,
+            base_margin=base_margin,
+            validate_features=validate_features,
+        )
+        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