Remove calls to deprecated Ray core APIs (#293)

Author: krfricke

requirements/test-requirements.txt

@@ -10,3 +10,4 @@ dask
 #workaround for now
 protobuf<4.0.0
 tensorboardX==2.2
+aiohttp

xgboost_ray/data_sources/dask.py

@@ -144,7 +144,7 @@ def get_ip_to_parts(data: Any) -> Dict[int, Sequence[Any]]:
     # (though Ray tries to keep locality). We need to use that until
     # ray.state.objects() or something like it is available again.
     partition_locations_df = persisted.map_partitions(
-        lambda df: pd.DataFrame([ray.get_runtime_context().node_id.hex()])
+        lambda df: pd.DataFrame([ray.get_runtime_context().get_node_id()])
     ).compute()
     partition_locations = [
         partition_locations_df[0].iloc[i] for i in range(partition_locations_df.size)

xgboost_ray/tests/conftest.py

@@ -14,9 +14,9 @@
 def get_default_fixure_system_config():
     system_config = {
         "object_timeout_milliseconds": 200,
-        "num_heartbeats_timeout": 10,
-        "object_store_full_max_retries": 3,
-        "object_store_full_initial_delay_ms": 100,
+        "health_check_initial_delay_ms": 0,
+        "health_check_failure_threshold": 10,
+        "object_store_full_delay_ms": 100,
     }
     return system_config
 
@@ -26,6 +26,8 @@ def get_default_fixture_ray_kwargs():
     ray_kwargs = {
         "num_cpus": 1,
         "object_store_memory": 150 * 1024 * 1024,
+        "dashboard_port": None,
+        "namespace": "default_test_namespace",
         "_system_config": system_config,
     }
     return ray_kwargs

xgboost_ray/tests/test_colocation.py

@@ -16,12 +16,12 @@
 
 class _MockQueueActor(_QueueActor):
     def get_node_id(self):
-        return ray.state.current_node_id()
+        return ray.get_runtime_context().get_node_id()
 
 
 class _MockEventActor(_EventActor):
     def get_node_id(self):
-        return ray.state.current_node_id()
+        return ray.get_runtime_context().get_node_id()
 
 
 @pytest.mark.usefixtures("ray_start_cluster")
@@ -71,20 +71,20 @@ def test_communication_colocation(self):
             cluster.wait_for_nodes()
             ray.init(address=cluster.address)
 
-            local_node = ray.state.current_node_id()
+            local_node = ray.get_runtime_context().get_node_id()
 
             # Note that these will have the same IP in the test cluster
-            assert len(ray.state.node_ids()) == 2
-            assert local_node in ray.state.node_ids()
+            assert len(ray.nodes()) == 2
+            assert local_node in [node["NodeID"] for node in ray.nodes()]
 
             def _mock_train(*args, _training_state, **kwargs):
                 assert (
                     ray.get(_training_state.queue.actor.get_node_id.remote())
-                    == ray.state.current_node_id()
+                    == ray.get_runtime_context().get_node_id()
                 )
                 assert (
                     ray.get(_training_state.stop_event.actor.get_node_id.remote())
-                    == ray.state.current_node_id()
+                    == ray.get_runtime_context().get_node_id()
                 )
                 return _train(*args, _training_state=_training_state, **kwargs)
 

xgboost_ray/util.py

@@ -85,7 +85,7 @@ def get_current_node_resource_key() -> str:
     If using Ray Client, this will return the resource key for the node that
     is running the client server.
     """
-    current_node_id = ray.get_runtime_context().node_id.hex()
+    current_node_id = ray.get_runtime_context().get_node_id()
     for node in ray.nodes():
         if node["NodeID"] == current_node_id:
             # Found the node.