Fix RayDeviceQuantileDMatrix (#69)

This fixes the RayDeviceQuantileDMatrix usage for the new data source interface. Also gets rid of Ray core calls to private APIs.

Author: krfricke

xgboost_ray/main.py

@@ -209,19 +209,29 @@ def _set_omp_num_threads():
 
 def _get_dmatrix(data: RayDMatrix, param: Dict) -> xgb.DMatrix:
     if isinstance(data, RayDeviceQuantileDMatrix):
-        if isinstance(param["data"], list):
-            dm_param = {
-                "feature_names": data.feature_names,
-                "feature_types": data.feature_types,
-                "missing": data.missing,
-            }
-            if not isinstance(data, xgb.DeviceQuantileDMatrix):
-                pass
-            param.update(dm_param)
-            it = RayDataIter(**param)
-            matrix = xgb.DeviceQuantileDMatrix(it, **dm_param)
-        else:
-            matrix = xgb.DeviceQuantileDMatrix(**param)
+        # If we only got a single data shard, create a list so we can
+        # iterate over it
+        if not isinstance(param["data"], list):
+            param["data"] = [param["data"]]
+
+            if not isinstance(param["label"], list):
+                param["label"] = [param["label"]]
+            if not isinstance(param["weight"], list):
+                param["weight"] = [param["weight"]]
+            if not isinstance(param["data"], list):
+                param["base_margin"] = [param["base_margin"]]
+
+        param["label_lower_bound"] = [None]
+        param["label_upper_bound"] = [None]
+
+        dm_param = {
+            "feature_names": data.feature_names,
+            "feature_types": data.feature_types,
+            "missing": data.missing,
+        }
+        param.update(dm_param)
+        it = RayDataIter(**param)
+        matrix = xgb.DeviceQuantileDMatrix(it, **dm_param)
     else:
         if isinstance(param["data"], list):
             dm_param = {
@@ -436,7 +446,8 @@ def train(self, rabit_args: List[str], params: Dict[str, Any],
             if num_threads > 0:
                 local_params["num_threads"] = num_threads
             else:
-                local_params["nthread"] = ray.utils.get_num_cpus()
+                local_params["nthread"] = sum(
+                    num for _, num in ray.get_resource_ids().get("CPU", []))
 
         if dtrain not in self._data:
             self.load_data(dtrain)

xgboost_ray/matrix.py

@@ -578,9 +578,9 @@ class RayDMatrix:
     def __init__(self,
                  data: Data,
                  label: Optional[Data] = None,
-                 missing: Optional[float] = None,
                  weight: Optional[Data] = None,
                  base_margin: Optional[Data] = None,
+                 missing: Optional[float] = None,
                  label_lower_bound: Optional[Data] = None,
                  label_upper_bound: Optional[Data] = None,
                  feature_names: Optional[List[str]] = None,
@@ -730,12 +730,50 @@ def __eq__(self, other):
 class RayDeviceQuantileDMatrix(RayDMatrix):
     """Currently just a thin wrapper for type detection"""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self,
+                 data: Data,
+                 label: Optional[Data] = None,
+                 weight: Optional[Data] = None,
+                 base_margin: Optional[Data] = None,
+                 missing: Optional[float] = None,
+                 label_lower_bound: Optional[Data] = None,
+                 label_upper_bound: Optional[Data] = None,
+                 feature_names: Optional[List[str]] = None,
+                 feature_types: Optional[List[np.dtype]] = None,
+                 *args,
+                 **kwargs):
         if cp is None:
             raise RuntimeError(
                 "RayDeviceQuantileDMatrix requires cupy to be installed."
-                "\nFIX THIS by installing cupy: `pip install cupy`")
-        super(RayDeviceQuantileDMatrix, self).__init__(*args, **kwargs)
+                "\nFIX THIS by installing cupy: `pip install cupy-cudaXYZ` "
+                "where XYZ is your local CUDA version.")
+        if label_lower_bound or label_upper_bound:
+            raise RuntimeError(
+                "RayDeviceQuantileDMatrix does not support "
+                "`label_lower_bound` and `label_upper_bound` (just as the "
+                "xgboost.DeviceQuantileDMatrix). Please pass None instead.")
+        super(RayDeviceQuantileDMatrix, self).__init__(
+            data=data,
+            label=label,
+            weight=weight,
+            base_margin=base_margin,
+            missing=missing,
+            label_lower_bound=None,
+            label_upper_bound=None,
+            feature_names=feature_names,
+            feature_types=feature_types,
+            *args,
+            **kwargs)
+
+    def get_data(
+            self, rank: int, num_actors: Optional[int] = None
+    ) -> Dict[str, Union[None, pd.DataFrame, List[Optional[pd.DataFrame]]]]:
+        data_dict = super(RayDeviceQuantileDMatrix, self).get_data(
+            rank=rank, num_actors=num_actors)
+        # Remove some dict keys here that are generated automatically
+        data_dict.pop("label_lower_bound", None)
+        data_dict.pop("label_upper_bound", None)
+        return data_dict
 
 
 def _can_load_distributed(source: Data) -> bool: