Remove legacy code (#217)

krfricke · web-flow · commit 929b3e3a4318 · 2022-05-12T12:03:50.000+01:00
There are several places where we still have code paths for legacy Ray versions. With the next release we can require newer Ray versions to be used.
diff --git a/setup.py b/setup.py
@@ -10,6 +10,5 @@
     "distributed computing framework Ray.",
     url="https://github.com/ray-project/xgboost_ray",
     install_requires=[
-        "ray>=1.6", "numpy>=1.16", "pandas", "wrapt>=1.12.1",
-        "xgboost>=0.90"
+        "ray>=1.10", "numpy>=1.16", "pandas", "wrapt>=1.12.1", "xgboost>=0.90"
     ])
diff --git a/xgboost_ray/main.py b/xgboost_ray/main.py
@@ -38,15 +38,11 @@ class EarlyStopException(XGBoostError):
     from ray.util.annotations import PublicAPI, DeveloperAPI
     from ray.util.placement_group import PlacementGroup, \
         remove_placement_group, get_current_placement_group
+    from ray.util.queue import Queue
 
-    from xgboost_ray.util import Event, Queue, MultiActorTask, \
-        force_on_current_node
+    from xgboost_ray.util import Event, MultiActorTask, force_on_current_node
 
-    if LooseVersion(ray.__version__) >= LooseVersion("1.5.0"):
-        # https://github.com/ray-project/ray/pull/16437
-        DEFAULT_PG = "default"
-    else:
-        DEFAULT_PG = None
+    DEFAULT_PG = "default"
 
     RAY_INSTALLED = True
 except ImportError:
@@ -63,7 +59,7 @@ def inner_f(*args, **kwargs):
     RAY_INSTALLED = False
 
 from xgboost_ray.tune import _try_add_tune_callback, _get_tune_resources, \
-    TUNE_USING_PG, is_session_enabled
+    is_session_enabled
 
 from xgboost_ray.matrix import RayDMatrix, combine_data, \
     RayDeviceQuantileDMatrix, RayDataIter, concat_dataframes, \
@@ -849,7 +845,7 @@ def _create_communication_processes(added_tune_callback: bool = False):
     node_ip = get_node_ip_address()
     # Have to explicitly set num_cpus to 0.
     placement_option = {"num_cpus": 0}
-    if added_tune_callback and TUNE_USING_PG:
+    if added_tune_callback:
         # If Tune is using placement groups, then we force Queue and
         # StopEvent onto same bundle as the Trainable.
         # This forces all 3 to be on the same node.
@@ -1388,12 +1384,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"
+            # Tune is using placement groups, so the strategy has already
+            # been set. Don't create an additional placement_group here.
+            placement_strategy = None
         elif bool(ENV.USE_SPREAD_STRATEGY):
             placement_strategy = "SPREAD"
 
diff --git a/xgboost_ray/tests/test_colocation.py b/xgboost_ray/tests/test_colocation.py
@@ -8,9 +8,10 @@
 import numpy as np
 
 import ray
+from ray.util.queue import _QueueActor
 from xgboost_ray import train, RayDMatrix, RayParams
 from xgboost_ray.main import _train
-from xgboost_ray.util import _EventActor, _QueueActor
+from xgboost_ray.util import _EventActor
 
 
 class _MockQueueActor(_QueueActor):
@@ -57,7 +58,7 @@ def tearDown(self) -> None:
             shutil.rmtree(self.tmpdir)
         ray.shutdown()
 
-    @patch("xgboost_ray.util._QueueActor", _MockQueueActor)
+    @patch("ray.util.queue._QueueActor", _MockQueueActor)
     @patch("xgboost_ray.util._EventActor", _MockEventActor)
     def test_communication_colocation(self):
         """Checks that Queue and Event actors are colocated with the driver."""
diff --git a/xgboost_ray/tune.py b/xgboost_ray/tune.py
@@ -1,21 +1,14 @@
 # Tune imports.
-import os
-from typing import Dict, Union, List, Optional
+from typing import Dict, Optional
 
 import ray
 
-try:
-    from typing import OrderedDict
-except ImportError:
-    from collections import OrderedDict
-
 import logging
 
 from ray.util.annotations import PublicAPI
 
 from xgboost_ray.xgb import xgboost as xgb
 
-from xgboost_ray.compat import TrainingCallback
 from xgboost_ray.session import put_queue, get_rabit_rank
 from xgboost_ray.util import Unavailable, force_on_current_node
 
@@ -42,90 +35,7 @@ def is_session_enabled():
     flatten_dict = is_session_enabled
     TUNE_INSTALLED = False
 
-# Todo(krfricke): Remove after next ray core release
-if not hasattr(OrigTuneReportCallback, "_get_report_dict") or not issubclass(
-        OrigTuneReportCallback, TrainingCallback):
-    TUNE_LEGACY = True
-else:
-    TUNE_LEGACY = False
-
-# Todo(amogkam): Remove after Ray 1.3 release.
-try:
-    from ray.tune import PlacementGroupFactory
-
-    TUNE_USING_PG = True
-except ImportError:
-    TUNE_USING_PG = False
-    PlacementGroupFactory = Unavailable
-
-if TUNE_LEGACY and TUNE_INSTALLED:
-    # Until the next release, keep compatible callbacks here.
-    class TuneReportCallback(OrigTuneReportCallback, TrainingCallback):
-        def _get_report_dict(self, evals_log):
-            if isinstance(evals_log, OrderedDict):
-                # xgboost>=1.3
-                result_dict = flatten_dict(evals_log, delimiter="-")
-                for k in list(result_dict):
-                    result_dict[k] = result_dict[k][0]
-            else:
-                # xgboost<1.3
-                result_dict = dict(evals_log)
-            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 after_iteration(self, model, epoch: int, evals_log: Dict):
-            if get_rabit_rank() == 0:
-                report_dict = self._get_report_dict(evals_log)
-                put_queue(lambda: tune.report(**report_dict))
-
-    class _TuneCheckpointCallback(_OrigTuneCheckpointCallback,
-                                  TrainingCallback):
-        def __init__(self, filename: str, frequency: int):
-            super(_TuneCheckpointCallback, self).__init__(filename)
-            self._frequency = frequency
-
-        @staticmethod
-        def _create_checkpoint(model, 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 after_iteration(self, model, epoch: int, evals_log: Dict):
-            if get_rabit_rank() == 0:
-                put_queue(lambda: self._create_checkpoint(
-                    model, epoch, self._filename, self._frequency))
-
-    class TuneReportCheckpointCallback(OrigTuneReportCheckpointCallback,
-                                       TrainingCallback):
-        _checkpoint_callback_cls = _TuneCheckpointCallback
-        _report_callbacks_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_callbacks_cls(metrics)
-
-        def after_iteration(self, model, epoch: int, evals_log: Dict):
-            if get_rabit_rank() == 0:
-                self._checkpoint.after_iteration(model, epoch, evals_log)
-                self._report.after_iteration(model, epoch, evals_log)
-
-elif TUNE_INSTALLED:
+if TUNE_INSTALLED:
     # New style callbacks.
     class TuneReportCallback(OrigTuneReportCallback):
         def after_iteration(self, model, epoch: int, evals_log: Dict):
@@ -168,15 +78,10 @@ def _try_add_tune_callback(kwargs: Dict):
                         target="xgboost_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(
@@ -203,35 +108,21 @@ def _get_tune_resources(num_actors: int, cpus_per_actor: int,
                         resources_per_actor: Optional[Dict]):
     """Returns object to use for ``resources_per_trial`` with Ray Tune."""
     if TUNE_INSTALLED:
-        if not TUNE_USING_PG:
-            resources_per_actor = {} if not resources_per_actor \
-                else resources_per_actor
-            extra_custom_resources = {
-                k: v * num_actors
-                for k, v in resources_per_actor.items()
-            }
-            return dict(
-                cpu=1,
-                extra_cpu=cpus_per_actor * num_actors,
-                extra_gpu=gpus_per_actor * num_actors,
-                extra_custom_resources=extra_custom_resources,
-            )
-        else:
-            from ray.tune import PlacementGroupFactory
-
-            head_bundle = {"CPU": 1}
-            child_bundle = {"CPU": cpus_per_actor, "GPU": gpus_per_actor}
-            child_bundle_extra = {} if resources_per_actor is None else \
-                resources_per_actor
-            child_bundles = [{
-                **child_bundle,
-                **child_bundle_extra
-            } for _ in range(num_actors)]
-            bundles = [head_bundle] + child_bundles
-            placement_group_factory = PlacementGroupFactory(
-                bundles, strategy="PACK")
-
-            return placement_group_factory
+        from ray.tune import PlacementGroupFactory
+
+        head_bundle = {"CPU": 1}
+        child_bundle = {"CPU": cpus_per_actor, "GPU": gpus_per_actor}
+        child_bundle_extra = {} if resources_per_actor is None else \
+            resources_per_actor
+        child_bundles = [{
+            **child_bundle,
+            **child_bundle_extra
+        } for _ in range(num_actors)]
+        bundles = [head_bundle] + child_bundles
+        placement_group_factory = PlacementGroupFactory(
+            bundles, strategy="PACK")
+
+        return placement_group_factory
     else:
         raise RuntimeError("Tune is not installed, so `get_tune_resources` is "
                            "not supported. You can install Ray Tune via `pip "
diff --git a/xgboost_ray/util.py b/xgboost_ray/util.py
@@ -4,7 +4,6 @@
 
 import ray
 from ray.util.annotations import DeveloperAPI
-from ray.util.queue import Queue as RayQueue, Empty, Full
 
 
 @DeveloperAPI
@@ -50,78 +49,6 @@ def shutdown(self):
         self.actor = None
 
 
-# Remove after Ray 1.2 release.
-if getattr(RayQueue, "shutdown", None) is not None:
-    from ray.util.queue import _QueueActor
-else:
-    # Have to copy the class here so that we can subclass this for mocking.
-    # If we have the @ray.remote decorator, then we can't subclass it.
-    class _QueueActor:
-        def __init__(self, maxsize):
-            self.maxsize = maxsize
-            self.queue = asyncio.Queue(self.maxsize)
-
-        def qsize(self):
-            return self.queue.qsize()
-
-        def empty(self):
-            return self.queue.empty()
-
-        def full(self):
-            return self.queue.full()
-
-        async def put(self, item, timeout=None):
-            try:
-                await asyncio.wait_for(self.queue.put(item), timeout)
-            except asyncio.TimeoutError:
-                raise Full
-
-        async def get(self, timeout=None):
-            try:
-                return await asyncio.wait_for(self.queue.get(), timeout)
-            except asyncio.TimeoutError:
-                raise Empty
-
-        def put_nowait(self, item):
-            self.queue.put_nowait(item)
-
-        def put_nowait_batch(self, items):
-            # If maxsize is 0, queue is unbounded, so no need to check size.
-            if self.maxsize > 0 and len(items) + self.qsize() > self.maxsize:
-                raise Full(f"Cannot add {len(items)} items to queue of size "
-                           f"{self.qsize()} and maxsize {self.maxsize}.")
-            for item in items:
-                self.queue.put_nowait(item)
-
-        def get_nowait(self):
-            return self.queue.get_nowait()
-
-        def get_nowait_batch(self, num_items):
-            if num_items > self.qsize():
-                raise Empty(f"Cannot get {num_items} items from queue of size "
-                            f"{self.qsize()}.")
-            return [self.queue.get_nowait() for _ in range(num_items)]
-
-
-# Remove after Ray 1.2 release.
-@DeveloperAPI
-class Queue(RayQueue):
-    def __init__(self, maxsize: int = 0,
-                 actor_options: Optional[Dict] = None) -> None:
-        actor_options = actor_options or {}
-        self.maxsize = maxsize
-        self.actor = ray.remote(_QueueActor).options(**actor_options).remote(
-            self.maxsize)
-
-    def shutdown(self):
-        if getattr(RayQueue, "shutdown", None) is not None:
-            super(Queue, self).shutdown()
-        else:
-            if self.actor:
-                ray.kill(self.actor)
-            self.actor = None
-
-
 @DeveloperAPI
 class MultiActorTask:
     """Utility class to hold multiple futures.
