Make xgboost an extra dependency (#151)

Author: Yard1

.github/workflows/test.yaml

@@ -55,7 +55,7 @@ jobs:
         if [ -f requirements-test.txt ]; then python -m pip install -r requirements-test.txt; fi
     - name: Install package
       run: |
-        python -m pip install -e .
+        python -m pip install -e ".[default]"
     - name: Print environment info
       run: |
         ./xgboost_ray/tests/env_info.sh
@@ -86,7 +86,7 @@ jobs:
         if [ -f requirements-test.txt ]; then python -m pip install  -r requirements-test.txt; fi
     - name: Install package
       run: |
-        python -m pip install -e .
+        python -m pip install -e ".[default]"
     - name: Print environment info
       run: |
         ./xgboost_ray/tests/env_info.sh
@@ -124,7 +124,7 @@ jobs:
         python -m pip uninstall -y tabulate
     - name: Install package
       run: |
-        python -m pip install -e .
+        python -m pip install -e ".[default]"
     - name: Print environment info
       run: |
         ./xgboost_ray/tests/env_info.sh
@@ -166,7 +166,7 @@ jobs:
         sudo apt-get install -y --no-install-recommends ninja-build
     - name: Install package
       run: |
-        python -m pip install -e .
+        python -m pip install -e ".[default]"
     - name: Clone XGBoost repo
       uses: actions/checkout@v2
       with:
@@ -212,7 +212,7 @@ jobs:
         if [ -f requirements-test.txt ]; then python -m pip install  -r requirements-test.txt; fi
     - name: Install package
       run: |
-        python -m pip install -e .
+        python -m pip install -e ".[default]"
     - name: Install legacy XGBoost
       run: |
         python -m pip install xgboost==0.90

README.md

@@ -22,15 +22,16 @@ Installation
 You can install the latest XGBoost-Ray release from PIP:
 
 ```bash
-pip install xgboost_ray
+pip install "xgboost_ray[default]"
 ```
 
 If you'd like to install the latest master, use this command instead:
 
 ```bash
-pip install git+https://github.com/ray-project/xgboost_ray.git#xgboost_ray
+pip install "git+https://github.com/ray-project/xgboost_ray.git#xgboost_ray[default]"
 ```
 
+Omitting `[default]` will cause `xgboost` to not be installed as a dependency.
 
 Usage
 -----

setup.py

@@ -10,7 +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",
-        "wrapt>=1.12.1"
-    ])
+        "ray", "numpy>=1.16,<1.20", "pandas", "pyarrow<5.0.0", "wrapt>=1.12.1"
+    ],
+    extras_require={"default": ["xgboost>=0.90"]})
 # pyarrow<5.0.0 pinned until petastorm is updated

xgboost_ray/compat/__init__.py

@@ -1,4 +1,7 @@
-import xgboost as xgb
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from xgboost_ray.xgb import xgboost as xgb
 
 try:
     from xgboost.callback import TrainingCallback
@@ -15,7 +18,7 @@ def __init__(self):
                 self._before_iteration = getattr(self, "before_iteration")
                 self.__dict__["before_iteration"] = True
 
-        def __call__(self, callback_env: xgb.core.CallbackEnv):
+        def __call__(self, callback_env: "xgb.core.CallbackEnv"):
             if hasattr(self, "_before_iteration"):
                 self._before_iteration(
                     model=callback_env.model,

xgboost_ray/data_sources/data_source.py

@@ -1,12 +1,14 @@
-from typing import Any, Optional, Sequence, Tuple, Dict, List
+from typing import Any, Optional, Sequence, Tuple, Dict, List, TYPE_CHECKING
 
 from enum import Enum
 
 import pandas as pd
-import xgboost as xgb
 
 from ray.actor import ActorHandle
 
+if TYPE_CHECKING:
+    from xgboost_ray.xgb import xgboost as xgb
+
 
 class RayFileType(Enum):
     """Enum for different file types (used for overrides)."""
@@ -86,7 +88,7 @@ def load_data(data: Any,
         raise NotImplementedError
 
     @staticmethod
-    def update_feature_names(matrix: xgb.DMatrix,
+    def update_feature_names(matrix: "xgb.DMatrix",
                              feature_names: Optional[List[str]]):
         """Optionally update feature names before training/prediction
 

xgboost_ray/data_sources/numpy.py

@@ -1,12 +1,14 @@
-from typing import Any, Optional, Sequence, List
+from typing import Any, Optional, Sequence, List, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
-import xgboost as xgb
 
 from xgboost_ray.data_sources.data_source import DataSource, RayFileType
 from xgboost_ray.data_sources.pandas import Pandas
 
+if TYPE_CHECKING:
+    from xgboost_ray.xgb import xgboost as xgb
+
 
 class Numpy(DataSource):
     """Read from numpy arrays."""
@@ -17,7 +19,7 @@ def is_data_type(data: Any,
         return isinstance(data, np.ndarray)
 
     @staticmethod
-    def update_feature_names(matrix: xgb.DMatrix,
+    def update_feature_names(matrix: "xgb.DMatrix",
                              feature_names: Optional[List[str]]):
         # Potentially unset feature names
         matrix.feature_names = feature_names

xgboost_ray/main.py

@@ -14,7 +14,7 @@
 import numpy as np
 import pandas as pd
 
-import xgboost as xgb
+from xgboost_ray.xgb import xgboost as xgb
 from xgboost.core import XGBoostError, EarlyStopException
 
 from xgboost_ray.callback import DistributedCallback, \
@@ -76,16 +76,18 @@
 ELASTIC_RESTART_GRACE_PERIOD_S = int(
     os.getenv("RXGB_ELASTIC_RESTART_GRACE_PERIOD_S", 10))
 
+xgboost_version = xgb.__version__ if xgb else "0.0.0"
+
 LEGACY_WARNING = (
     f"You are using `xgboost_ray` with a legacy XGBoost version "
-    f"(version {xgb.__version__}). While we try to support "
+    f"(version {xgboost_version}). While we try to support "
     f"older XGBoost versions, please note that this library is only "
     f"fully tested and supported for XGBoost >= 1.4. Please consider "
     f"upgrading your XGBoost version (`pip install -U xgboost`).")
 
 # XGBoost version as an int tuple for comparisions
 XGBOOST_VERSION_TUPLE = tuple(
-    int(x) for x in re.sub(r"[^\.0-9]", "", xgb.__version__).split("."))
+    int(x) for x in re.sub(r"[^\.0-9]", "", xgboost_version).split("."))
 
 
 class RayXGBoostTrainingError(RuntimeError):
@@ -1144,6 +1146,11 @@ def train(
     """
     os.environ.setdefault("RAY_IGNORE_UNHANDLED_ERRORS", "1")
 
+    if xgb is None:
+        raise ImportError(
+            "xgboost package is not installed. XGBoost-Ray WILL NOT WORK. "
+            "FIX THIS by running `pip install \"xgboost-ray[default]\"`.")
+
     if _remote is None:
         _remote = _is_client_connected() and \
                   not is_session_enabled()
@@ -1528,6 +1535,11 @@ def predict(model: xgb.Booster,
     """
     os.environ.setdefault("RAY_IGNORE_UNHANDLED_ERRORS", "1")
 
+    if xgb is None:
+        raise ImportError(
+            "xgboost package is not installed. XGBoost-Ray WILL NOT WORK. "
+            "FIX THIS by running `pip install \"xgboost-ray[default]\"`.")
+
     if _remote is None:
         _remote = _is_client_connected() and \
                   not is_session_enabled()

xgboost_ray/matrix.py

@@ -2,7 +2,7 @@
 import uuid
 from enum import Enum
 from typing import Union, Optional, Tuple, Iterable, List, Dict, Sequence, \
-    Callable, Type
+    Callable, Type, TYPE_CHECKING
 
 from ray.actor import ActorHandle
 
@@ -13,7 +13,6 @@
 
 import numpy as np
 import pandas as pd
-import xgboost as xgb
 
 import os
 
@@ -43,6 +42,9 @@ class RayDataset:
     DataIter = object
     LEGACY_MATRIX = True
 
+if TYPE_CHECKING:
+    from xgboost_ray.xgb import xgboost as xgb
+
 Data = Union[str, List[str], np.ndarray, pd.DataFrame, pd.Series, MLDataset]
 
 
@@ -194,7 +196,7 @@ def assert_enough_shards_for_actors(self, num_actors: int):
         # Pass per default
         pass
 
-    def update_matrix_properties(self, matrix: xgb.DMatrix):
+    def update_matrix_properties(self, matrix: "xgb.DMatrix"):
         data_source = self.get_data_source()
         data_source.update_feature_names(matrix, self.feature_names)
 
@@ -780,7 +782,7 @@ def unload_data(self):
                 del self.refs[rank][name]
         self.loaded = False
 
-    def update_matrix_properties(self, matrix: xgb.DMatrix):
+    def update_matrix_properties(self, matrix: "xgb.DMatrix"):
         self.loader.update_matrix_properties(matrix)
 
     def __hash__(self):

xgboost_ray/tests/release/cluster_cpu.yaml

@@ -37,4 +37,4 @@ file_mounts_sync_continuously: false
 setup_commands:
     - pip install -U {{env["RAY_WHEEL"] | default("ray")}}
     - pip install dask pytest
-    - pip install -U {{env["XGBOOST_RAY_PACKAGE"] | default("xgboost_ray")}}
+    - pip install -U {{env["XGBOOST_RAY_PACKAGE"] | default("xgboost_ray")}}[default]

xgboost_ray/tests/release/setup_xgboost.sh

@@ -5,7 +5,7 @@ pip install pytest
 pip uninstall -y xgboost_ray || true
 
 # Install xgboost package
-pip install -U ${XGBOOST_RAY_PACKAGE:-xgboost_ray}
+pip install -U ${XGBOOST_RAY_PACKAGE:-xgboost_ray}[default]
 
 # Create test dataset
 sudo mkdir -p /data || true