Remove legacy code (#31)

Yard1 · web-flow · commit 91fa2ad435a4 · 2022-05-12T18:04:25.000+02:00
* Remove legacy code

* Fix CI

* Fix
diff --git a/lightgbm_ray/main.py b/lightgbm_ray/main.py
@@ -57,7 +57,7 @@
     ENV, RayActorError, pickle, _PrepareActorTask, RayParams as RayXGBParams,
     _TrainingState, _is_client_connected, is_session_enabled,
     force_on_current_node, _assert_ray_support, _maybe_print_legacy_warning,
-    _Checkpoint, _create_communication_processes, TUNE_USING_PG, RayTaskError,
+    _Checkpoint, _create_communication_processes, RayTaskError,
     RayXGBoostActorAvailable, RayXGBoostTrainingError, _create_placement_group,
     _shutdown, PlacementGroup, ActorHandle, combine_data, _trigger_data_load,
     DEFAULT_PG, _autodetect_resources as _autodetect_resources_base)
@@ -1202,12 +1202,9 @@ def _wrapped(*args, **kwargs):
     placement_strategy = None
     if not ray_params.elastic_training:
         if added_tune_callback:
-            if TUNE_USING_PG:
-                # If Tune is using placement groups, then strategy has already
-                # been set. Don't create an additional placement_group here.
-                placement_strategy = None
-            else:
-                placement_strategy = "PACK"
+            # If Tune is using placement groups, then strategy has already
+            # been set. Don't create an additional placement_group here.
+            placement_strategy = None
         elif ENV.USE_SPREAD_STRATEGY:
             placement_strategy = "SPREAD"
 
diff --git a/lightgbm_ray/tune.py b/lightgbm_ray/tune.py
@@ -1,6 +1,5 @@
 # Tune imports.
-import os
-from typing import Dict, Union, List
+from typing import Dict
 
 import ray
 import logging
@@ -11,12 +10,15 @@
 from lightgbm.callback import CallbackEnv
 
 from xgboost_ray.session import put_queue
-from xgboost_ray.util import Unavailable, force_on_current_node
+from xgboost_ray.util import force_on_current_node
 
 try:
     from ray import tune
     from ray.tune import is_session_enabled
-    from ray.tune.utils import flatten_dict
+    from ray.tune.integration.lightgbm import (
+        TuneReportCallback as OrigTuneReportCallback, _TuneCheckpointCallback
+        as _OrigTuneCheckpointCallback, TuneReportCheckpointCallback as
+        OrigTuneReportCheckpointCallback)
 
     TUNE_INSTALLED = True
 except ImportError:
@@ -25,33 +27,8 @@
     def is_session_enabled():
         return False
 
-    flatten_dict = is_session_enabled
     TUNE_INSTALLED = False
 
-try:
-    from ray.tune.integration.lightgbm import \
-        TuneReportCallback as OrigTuneReportCallback, \
-        _TuneCheckpointCallback as _OrigTuneCheckpointCallback, \
-        TuneReportCheckpointCallback as OrigTuneReportCheckpointCallback
-except ImportError:
-    TuneReportCallback = _TuneCheckpointCallback = \
-        TuneReportCheckpointCallback = Unavailable
-    OrigTuneReportCallback = _OrigTuneCheckpointCallback = \
-        OrigTuneReportCheckpointCallback = object
-
-if not hasattr(OrigTuneReportCallback, "_get_report_dict"):
-    TUNE_LEGACY = True
-else:
-    TUNE_LEGACY = False
-
-try:
-    from ray.tune import PlacementGroupFactory
-
-    TUNE_USING_PG = True
-except ImportError:
-    TUNE_USING_PG = False
-    PlacementGroupFactory = Unavailable
-
 
 class _TuneLGBMRank0Mixin:
     """Mixin to allow for dynamic setting of rank so that only
@@ -69,115 +46,8 @@ def is_rank_0(self, val: bool):
         self._is_rank_0 = val
 
 
-if TUNE_LEGACY and TUNE_INSTALLED:
-
-    class TuneReportCallback(_TuneLGBMRank0Mixin, OrigTuneReportCallback):
-        """Create a callback that reports metrics to Ray Tune."""
-        order = 20
-
-        def __init__(
-                self,
-                metrics: Union[None, str, List[str], Dict[str, str]] = None):
-            if isinstance(metrics, str):
-                metrics = [metrics]
-            self._metrics = metrics
-
-        def _get_report_dict(self,
-                             evals_log: Dict[str, Dict[str, list]]) -> dict:
-            result_dict = flatten_dict(evals_log, delimiter="-")
-            if not self._metrics:
-                report_dict = result_dict
-            else:
-                report_dict = {}
-                for key in self._metrics:
-                    if isinstance(self._metrics, dict):
-                        metric = self._metrics[key]
-                    else:
-                        metric = key
-                    report_dict[key] = result_dict[metric]
-            return report_dict
-
-        def _get_eval_result(self, env: CallbackEnv) -> dict:
-            eval_result = {}
-            for data_name, eval_name, result, _ in env.evaluation_result_list:
-                if data_name not in eval_result:
-                    eval_result[data_name] = {}
-                eval_result[data_name][eval_name] = result
-            return eval_result
-
-        def __call__(self, env: CallbackEnv) -> None:
-            if not self.is_rank_0:
-                return
-            eval_result = self._get_eval_result(env)
-            report_dict = self._get_report_dict(eval_result)
-            put_queue(lambda: tune.report(**report_dict))
-
-    class _TuneCheckpointCallback(_TuneLGBMRank0Mixin,
-                                  _OrigTuneCheckpointCallback):
-        """LightGBM checkpoint callback"""
-        order = 19
-
-        def __init__(self,
-                     filename: str = "checkpoint",
-                     frequency: int = 5,
-                     *,
-                     is_rank_0: bool = False):
-            self._filename = filename
-            self._frequency = frequency
-            self.is_rank_0 = is_rank_0
-
-        @staticmethod
-        def _create_checkpoint(model: Booster, epoch: int, filename: str,
-                               frequency: int):
-            if epoch % frequency > 0:
-                return
-            with tune.checkpoint_dir(step=epoch) as checkpoint_dir:
-                model.save_model(os.path.join(checkpoint_dir, filename))
-
-        def __call__(self, env: CallbackEnv) -> None:
-            if not self.is_rank_0:
-                return
-            put_queue(lambda: self._create_checkpoint(
-                env.model, env.iteration, self._filename, self._frequency))
-
-    class TuneReportCheckpointCallback(_TuneLGBMRank0Mixin,
-                                       OrigTuneReportCheckpointCallback):
-        """Creates a callback that reports metrics and checkpoints model."""
-        order = 21
-
-        _checkpoint_callback_cls = _TuneCheckpointCallback
-        _report_callback_cls = TuneReportCallback
-
-        def __init__(
-                self,
-                metrics: Union[None, str, List[str], Dict[str, str]] = None,
-                filename: str = "checkpoint",
-                frequency: int = 5):
-            self._checkpoint = self._checkpoint_callback_cls(
-                filename, frequency)
-            self._report = self._report_callback_cls(metrics)
-
-        @property
-        def is_rank_0(self) -> bool:
-            try:
-                return self._is_rank_0
-            except AttributeError:
-                return False
-
-        @is_rank_0.setter
-        def is_rank_0(self, val: bool):
-            self._is_rank_0 = val
-            if hasattr(self, "_checkpoint"):
-                self._checkpoint.is_rank_0 = val
-            if hasattr(self, "_report"):
-                self._report.is_rank_0 = val
-
-        def __call__(self, env: CallbackEnv) -> None:
-            self._checkpoint(env)
-            self._report(env)
+if TUNE_INSTALLED:
 
-elif TUNE_INSTALLED:
-    # New style callbacks.
     class TuneReportCallback(_TuneLGBMRank0Mixin, OrigTuneReportCallback):
         def __call__(self, env: CallbackEnv) -> None:
             if not self.is_rank_0:
@@ -241,15 +111,10 @@ def _try_add_tune_callback(kwargs: Dict):
                         target="lightgbm_ray.tune.TuneReportCallback"))
                 has_tune_callback = True
             elif isinstance(cb, OrigTuneReportCheckpointCallback):
-                if TUNE_LEGACY:
-                    replace_cb = TuneReportCheckpointCallback(
-                        metrics=cb._report._metrics,
-                        filename=cb._checkpoint._filename)
-                else:
-                    replace_cb = TuneReportCheckpointCallback(
-                        metrics=cb._report._metrics,
-                        filename=cb._checkpoint._filename,
-                        frequency=cb._checkpoint._frequency)
+                replace_cb = TuneReportCheckpointCallback(
+                    metrics=cb._report._metrics,
+                    filename=cb._checkpoint._filename,
+                    frequency=cb._checkpoint._frequency)
                 new_callbacks.append(replace_cb)
                 logging.warning(
                     REPLACE_MSG.format(
diff --git a/setup.py b/setup.py
@@ -9,4 +9,4 @@
     long_description="A distributed backend for LightGBM built on top of "
     "distributed computing framework Ray.",
     url="https://github.com/ray-project/lightgbm_ray",
-    install_requires=["lightgbm>=3.2.1", "xgboost_ray>=0.1.8"])
+    install_requires=["lightgbm>=3.2.1", "xgboost_ray>=0.1.9"])
