[backport][dask] Fix with the latest Dask. (dmlc#11291) (dmlc#11302)

Author: trivialfis

python-package/xgboost/dask/__init__.py

@@ -83,6 +83,8 @@
 from dask import array as da
 from dask import bag as db
 from dask import dataframe as dd
+from dask.delayed import Delayed
+from distributed import Future
 
 from .. import collective, config
 from .._typing import FeatureNames, FeatureTypes, IterationRange
@@ -336,7 +338,7 @@ def __init__(
 
         self._n_cols = data.shape[1]
         assert isinstance(self._n_cols, int)
-        self.worker_map: Dict[str, List[distributed.Future]] = defaultdict(list)
+        self.worker_map: Dict[str, List[Future]] = defaultdict(list)
         self.is_quantile: bool = False
 
         self._init = client.sync(
@@ -369,7 +371,6 @@ async def _map_local_data(
         label_upper_bound: Optional[_DaskCollection] = None,
     ) -> "DaskDMatrix":
         """Obtain references to local data."""
-        from dask.delayed import Delayed
 
         def inconsistent(
             left: List[Any], left_name: str, right: List[Any], right_name: str
@@ -381,49 +382,39 @@ def inconsistent(
             )
             return msg
 
-        def check_columns(parts: numpy.ndarray) -> None:
-            # x is required to be 2 dim in __init__
-            assert parts.ndim == 1 or parts.shape[1], (
-                "Data should be"
-                " partitioned by row. To avoid this specify the number"
-                " of columns for your dask Array explicitly. e.g."
-                " chunks=(partition_size, X.shape[1])"
-            )
-
-        def to_delayed(d: _DaskCollection) -> List[Delayed]:
-            """Breaking data into partitions, a trick borrowed from
-            dask_xgboost. `to_delayed` downgrades high-level objects into numpy or
-            pandas equivalents.
-
-            """
+        def to_futures(d: _DaskCollection) -> List[Future]:
+            """Breaking data into partitions."""
             d = client.persist(d)
-            delayed_obj = d.to_delayed()
-            if isinstance(delayed_obj, numpy.ndarray):
-                # da.Array returns an array to delayed objects
-                check_columns(delayed_obj)
-                delayed_list: List[Delayed] = delayed_obj.flatten().tolist()
-            else:
-                # dd.DataFrame
-                delayed_list = delayed_obj
-            return delayed_list
+            if (
+                hasattr(d.partitions, "shape")
+                and len(d.partitions.shape) > 1
+                and d.partitions.shape[1] > 1
+            ):
+                raise ValueError(
+                    "Data should be"
+                    " partitioned by row. To avoid this specify the number"
+                    " of columns for your dask Array explicitly. e.g."
+                    " chunks=(partition_size, -1])"
+                )
+            return client.futures_of(d)
 
-        def flatten_meta(meta: Optional[_DaskCollection]) -> Optional[List[Delayed]]:
+        def flatten_meta(meta: Optional[_DaskCollection]) -> Optional[List[Future]]:
             if meta is not None:
-                meta_parts: List[Delayed] = to_delayed(meta)
+                meta_parts: List[Future] = to_futures(meta)
                 return meta_parts
             return None
 
-        X_parts = to_delayed(data)
+        X_parts = to_futures(data)
         y_parts = flatten_meta(label)
         w_parts = flatten_meta(weights)
         margin_parts = flatten_meta(base_margin)
         qid_parts = flatten_meta(qid)
         ll_parts = flatten_meta(label_lower_bound)
         lu_parts = flatten_meta(label_upper_bound)
 
-        parts: Dict[str, List[Delayed]] = {"data": X_parts}
+        parts: Dict[str, List[Future]] = {"data": X_parts}
 
-        def append_meta(m_parts: Optional[List[Delayed]], name: str) -> None:
+        def append_meta(m_parts: Optional[List[Future]], name: str) -> None:
             if m_parts is not None:
                 assert len(X_parts) == len(m_parts), inconsistent(
                     X_parts, "X", m_parts, name
@@ -437,12 +428,12 @@ def append_meta(m_parts: Optional[List[Delayed]], name: str) -> None:
         append_meta(ll_parts, "label_lower_bound")
         append_meta(lu_parts, "label_upper_bound")
         # At this point, `parts` looks like:
-        # [(x0, x1, ..), (y0, y1, ..), ..] in delayed form
+        # [(x0, x1, ..), (y0, y1, ..), ..] in future form
 
         # turn into list of dictionaries.
-        packed_parts: List[Dict[str, Delayed]] = []
+        packed_parts: List[Dict[str, Future]] = []
         for i in range(len(X_parts)):
-            part_dict: Dict[str, Delayed] = {}
+            part_dict: Dict[str, Future] = {}
             for key, value in parts.items():
                 part_dict[key] = value[i]
             packed_parts.append(part_dict)
@@ -451,16 +442,17 @@ def append_meta(m_parts: Optional[List[Delayed]], name: str) -> None:
         # pylint: disable=no-member
         delayed_parts: List[Delayed] = list(map(dask.delayed, packed_parts))
         # At this point, the mental model should look like:
-        # [(x0, y0, ..), (x1, y1, ..), ..] in delayed form
+        # [{"data": x0, "label": y0, ..}, {"data": x1, "label": y1, ..}, ..]
 
-        # convert delayed objects into futures and make sure they are realized
-        fut_parts: List[distributed.Future] = client.compute(delayed_parts)
+        # Convert delayed objects into futures and make sure they are realized
+        #
+        # This also makes partitions to align (co-locate) on workers (X_0, y_0 should be
+        # on the same worker).
+        fut_parts: List[Future] = client.compute(delayed_parts)
         await distributed.wait(fut_parts)  # async wait for parts to be computed
 
-        # maybe we can call dask.align_partitions here to ease the partition alignment?
-
         for part in fut_parts:
-            # Each part is [x0, y0, w0, ...] in future form.
+            # Each part is [{"data": x0, "label": y0, ..}, ...] in future form.
             assert part.status == "finished", part.status
 
         # Preserving the partition order for prediction.
@@ -473,7 +465,7 @@ def append_meta(m_parts: Optional[List[Delayed]], name: str) -> None:
             keys=[part.key for part in fut_parts]
         )
 
-        worker_map: Dict[str, List[distributed.Future]] = defaultdict(list)
+        worker_map: Dict[str, List[Future]] = defaultdict(list)
 
         for key, workers in who_has.items():
             worker_map[next(iter(workers))].append(key_to_partition[key])

src/c_api/c_api.cc

@@ -76,7 +76,7 @@ void XGBBuildInfoDevice(Json *p_info) {
 #endif
 
 XGB_DLL int XGBuildInfo(char const **out) {
-  API_BEGIN();
+  API_BEGIN_UNGUARD()
   xgboost_CHECK_C_ARG_PTR(out);
   Json info{Object{}};
 
@@ -135,14 +135,14 @@ XGB_DLL int XGBuildInfo(char const **out) {
 }
 
 XGB_DLL int XGBRegisterLogCallback(void (*callback)(const char*)) {
-  API_BEGIN_UNGUARD();
+  API_BEGIN_UNGUARD()
   LogCallbackRegistry* registry = LogCallbackRegistryStore::Get();
   registry->Register(callback);
   API_END();
 }
 
 XGB_DLL int XGBSetGlobalConfig(const char* json_str) {
-  API_BEGIN();
+  API_BEGIN_UNGUARD()
 
   xgboost_CHECK_C_ARG_PTR(json_str);
   Json config{Json::Load(StringView{json_str})};
@@ -204,7 +204,7 @@ XGB_DLL int XGBSetGlobalConfig(const char* json_str) {
 }
 
 XGB_DLL int XGBGetGlobalConfig(const char** json_str) {
-  API_BEGIN();
+  API_BEGIN_UNGUARD()
   auto const& global_config = *GlobalConfigThreadLocalStore::Get();
   Json config {ToJson(global_config)};
   auto const* mgr = global_config.__MANAGER__();
@@ -246,6 +246,8 @@ XGB_DLL int XGDMatrixCreateFromFile(const char *fname, int silent, DMatrixHandle
   xgboost_CHECK_C_ARG_PTR(fname);
   xgboost_CHECK_C_ARG_PTR(out);
 
+  LOG(WARNING) << error::DeprecatedFunc(__func__, "2.0.0", "XGDMatrixCreateFromURI");
+
   Json config{Object()};
   config["uri"] = std::string{fname};
   config["silent"] = silent;

tests/python-gpu/test_gpu_prediction.py

@@ -287,7 +287,7 @@ def predict_dense(x):
     def test_inplace_predict_cupy(self):
         self.run_inplace_predict_cupy(0)
 
-    @pytest.mark.skipif(**tm.no_cupy())
+    @pytest.mark.skip
     @pytest.mark.mgpu
     def test_inplace_predict_cupy_specified_device(self):
         import cupy as cp