Try fix ray

.github/workflows/platform-ci.yml

@@ -105,7 +105,7 @@ jobs:
             fi
             if [ -n "$WITH_RAY" ] || [ -n "$WITH_RAY_DAG" ] || [ -n "$WITH_RAY_DEPLOY" ]; then
               pip install "ray>=1.8.0,<2.4.0"
-              pip install "xgboost_ray<0.1.14" "protobuf<4"
+              pip install xgboost_ray "protobuf<4"
               # Ray Datasets need pyarrow>=6.0.1
               pip install "pyarrow>=6.0.1"
               pip install lightgbm

mars/dataframe/contrib/raydataset/dataset.py

@@ -11,10 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import operator
 from functools import reduce
 
-from ....utils import lazy_import
+from ....utils import lazy_import, lazy_import_on_load
 from .mldataset import _rechunk_if_needed
 from typing import Dict, List
 
@@ -24,6 +25,17 @@
 ray_dataset = lazy_import("ray.data", rename="ray_dataset")
 
 
+@lazy_import_on_load(ray_dataset, before=True)
+def _patch_xgboost():
+    # _treat_X_doc in xgboost_ray might cause error
+    try:
+        from xgboost.sklearn import XGBRanker
+
+        XGBRanker.predict.__doc__ = XGBRanker.predict.__doc__ or ""
+    except (ImportError, AttributeError):
+        pass
+
+
 def to_ray_dataset(df, num_shards: int = None):
     """Create a Ray Dataset from Mars DataFrame
 

mars/dataframe/contrib/raydataset/tests/test_mldataset.py

@@ -28,11 +28,7 @@
 
 ray = lazy_import("ray")
 ml_dataset = lazy_import("ray.util.data", rename="ml_dataset")
-
-try:
-    import xgboost_ray
-except ImportError:  # pragma: no cover
-    xgboost_ray = None
+xgboost_ray = lazy_import("xgboost_ray")
 try:
     import sklearn
 except ImportError:  # pragma: no cover

mars/storage/vineyard.py

@@ -144,11 +144,11 @@ async def setup(cls, **kwargs) -> Tuple[Dict, Dict]:
             vineyard_store = None
         else:
             vineyard_store = vineyard.deploy.local.start_vineyardd(
-                etcd_endpoints,
-                etcd_prefix,
-                vineyardd_path,
-                vineyard_size,
-                vineyard_socket,
+                etcd_endpoints=etcd_endpoints,
+                etcd_prefix=etcd_prefix,
+                vineyardd_path=vineyardd_path,
+                size=vineyard_size,
+                socket=vineyard_socket,
                 rpc=False,
             )
             vineyard_socket = (

mars/utils.py

@@ -368,26 +368,35 @@ def lazy_import(
 
     class LazyModule(object):
         def __init__(self):
+            self._before_loads = []
             self._on_loads = []
 
         def __getattr__(self, item):
             if item.startswith("_pytest") or item in ("__bases__", "__test__"):
                 raise AttributeError(item)
 
+            for before_load_func in self._before_loads:
+                before_load_func()
+
             real_mod = importlib.import_module(name, package=package)
             if rename in globals:
                 globals[rename] = real_mod
             elif locals is not None:
                 locals[rename] = real_mod
             ret = getattr(real_mod, item)
+
             for on_load_func in self._on_loads:
                 on_load_func()
+
             # make sure on_load hooks only executed once
             self._on_loads = []
             return ret
 
-        def add_load_handler(self, func: Callable):
-            self._on_loads.append(func)
+        def add_load_handler(self, func: Callable, before: bool = False):
+            if before:
+                self._before_loads.append(func)
+            else:
+                self._on_loads.append(func)
             return func
 
     if pkgutil.find_loader(prefix_name) is not None:
@@ -398,7 +407,7 @@ def add_load_handler(self, func: Callable):
         return None
 
 
-def lazy_import_on_load(lazy_mod):
+def lazy_import_on_load(lazy_mod, before: bool = False):
     def wrapper(fun):
         if lazy_mod is not None and hasattr(lazy_mod, "add_load_handler"):
             lazy_mod.add_load_handler(fun)