Separate out _RayXGBoostActor (#120)

Author: Yard1

xgboost_ray/main.py

@@ -358,7 +358,6 @@ def _validate_ray_params(ray_params: Union[None, RayParams, dict]) \
     return ray_params
 
 
-@ray.remote
 class RayXGBoostActor:
     """Remote Ray XGBoost actor class.
 
@@ -613,6 +612,11 @@ def predict(self, model: xgb.Booster, data: RayDMatrix, **kwargs):
         return predictions
 
 
+@ray.remote
+class _RemoteRayXGBoostActor(RayXGBoostActor):
+    pass
+
+
 class _PrepareActorTask(MultiActorTask):
     def __init__(self, actor: ActorHandle, queue: Queue, stop_event: Event,
                  load_data: List[RayDMatrix]):
@@ -663,7 +667,7 @@ def _create_actor(
     # Send DEFAULT_PG here, which changed in Ray > 1.4.0
     # If we send `None`, this will ignore the parent placement group and
     # lead to errors e.g. when used within Ray Tune
-    return RayXGBoostActor.options(
+    return _RemoteRayXGBoostActor.options(
         num_cpus=num_cpus_per_actor,
         num_gpus=num_gpus_per_actor,
         resources=resources_per_actor,

xgboost_ray/tests/test_data_source.py

@@ -9,7 +9,7 @@
 from ray import ObjectRef
 
 from xgboost_ray.data_sources import Modin, Dask
-from xgboost_ray.main import RayXGBoostActor
+from xgboost_ray.main import _RemoteRayXGBoostActor
 
 from xgboost_ray.data_sources.modin import MODIN_INSTALLED
 from xgboost_ray.data_sources.dask import DASK_INSTALLED
@@ -279,7 +279,7 @@ def create_remote_df(arr):
 
         # Create ray actors
         actors = [
-            RayXGBoostActor.options(resources={
+            _RemoteRayXGBoostActor.options(resources={
                 f"node:{nip}": 0.1
             }).remote(rank=rank, num_actors=len(actor_nodes))
             for rank, nip in enumerate(actor_node_ips)
@@ -414,7 +414,7 @@ def create_remote_df(arr):
 
         # Create ray actors
         actors = [
-            RayXGBoostActor.options(resources={
+            _RemoteRayXGBoostActor.options(resources={
                 f"node:{nip}": 0.1
             }).remote(rank=rank, num_actors=len(actor_nodes))
             for rank, nip in enumerate(actor_node_ips)

xgboost_ray/tests/test_fault_tolerance.py

@@ -418,7 +418,7 @@ def testSameResultWithAndWithoutError(self):
 
     @patch("xgboost_ray.main._PrepareActorTask", _FakeTask)
     @patch("xgboost_ray.elastic._PrepareActorTask", _FakeTask)
-    @patch("xgboost_ray.main.RayXGBoostActor", MagicMock)
+    @patch("xgboost_ray.main._RemoteRayXGBoostActor", MagicMock)
     @patch("xgboost_ray.main.ELASTIC_RESTART_GRACE_PERIOD_S", 30)
     @patch("xgboost_ray.elastic.ELASTIC_RESTART_GRACE_PERIOD_S", 30)
     def testMaybeScheduleNewActors(self):