Ensure that Dask Scheduler is set for Ray-on-Dask (#150)

* wrapper

* bump wrapt

* lint

Author: ijrsvt

setup.py

@@ -10,6 +10,7 @@
     "distributed computing framework Ray.",
     url="https://github.com/ray-project/xgboost_ray",
     install_requires=[
-        "xgboost>=0.90", "ray", "numpy>=1.16,<1.20", "pandas", "pyarrow<5.0.0"
+        "xgboost>=0.90", "ray", "numpy>=1.16,<1.20", "pandas", "pyarrow<5.0.0",
+        "wrapt>=1.12.1"
     ])
-    # pyarrow<5.0.0 pinned until petastorm is updated
+# pyarrow<5.0.0 pinned until petastorm is updated

xgboost_ray/data_sources/dask.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
-from typing import Any, Optional, Sequence, Dict, Union, Tuple
+from typing import Any, List, Optional, Sequence, Dict, Union, Tuple
+import wrapt
 
 import pandas as pd
 
@@ -29,6 +30,14 @@ def _assert_dask_installed():
             "the code should not have been reached.")
 
 
+@wrapt.decorator
+def ensure_ray_dask_initialized(func: Any, instance: Any, args: List[Any],
+                                kwargs: Any) -> Any:
+    _assert_dask_installed()
+    dask.config.set(scheduler=ray_dask_get)
+    return func(*args, **kwargs)
+
+
 class Dask(DataSource):
     """Read from distributed Dask dataframe.
 
@@ -41,6 +50,7 @@ class Dask(DataSource):
     supports_central_loading = True
     supports_distributed_loading = True
 
+    @ensure_ray_dask_initialized
     @staticmethod
     def is_data_type(data: Any,
                      filetype: Optional[RayFileType] = None) -> bool:
@@ -51,6 +61,7 @@ def is_data_type(data: Any,
 
         return isinstance(data, (DaskDataFrame, DaskSeries))
 
+    @ensure_ray_dask_initialized
     @staticmethod
     def load_data(
             data: Any,  # dask.pandas.DataFrame
@@ -79,6 +90,7 @@ def load_data(
 
         return local_df
 
+    @ensure_ray_dask_initialized
     @staticmethod
     def convert_to_series(data: Any) -> pd.Series:
         _assert_dask_installed()
@@ -95,6 +107,7 @@ def convert_to_series(data: Any) -> pd.Series:
 
         return DataSource.convert_to_series(data)
 
+    @ensure_ray_dask_initialized
     @staticmethod
     def get_actor_shards(
             data: Any,  # dask.dataframe.DataFrame
@@ -109,6 +122,7 @@ def get_actor_shards(
 
         return data, assign_partitions_to_actors(ip_to_parts, actor_rank_ips)
 
+    @ensure_ray_dask_initialized
     @staticmethod
     def get_n(data: Any):
         """