Improve cluster.wait_for_nodes() API. (#3712)

* Separate out functionality for querying client table and improve cluster.wait_for_nodes() API.

* Linting

* Add back logging statements.

* info -> debug

test/failure_test.py

@@ -639,7 +639,7 @@ def ray_start_two_nodes():
 # the monitor to detect enough missed heartbeats.
 def test_warning_for_dead_node(ray_start_two_nodes):
     cluster = ray_start_two_nodes
-    cluster.wait_for_nodes(2)
+    cluster.wait_for_nodes()
 
     client_ids = {item["ClientID"] for item in ray.global_state.client_table()}
 

test/multi_node_test_2.py

@@ -5,6 +5,7 @@ from __future__ import print_function
 import json
 import logging
 import pytest
+import time
 
 import ray
 import ray.services as services
@@ -82,10 +83,10 @@ def test_internal_config(start_connected_longer_cluster):
     cluster.wait_for_nodes()
 
     cluster.remove_node(worker)
-    cluster.wait_for_nodes(retries=10)
+    time.sleep(1)
     assert ray.global_state.cluster_resources()["CPU"] == 2
 
-    cluster.wait_for_nodes(retries=20)
+    time.sleep(2)
     assert ray.global_state.cluster_resources()["CPU"] == 1
 
 

test/runtest.py

@@ -1772,7 +1772,7 @@ def test_multiple_local_schedulers(ray_start_cluster):
     cluster.add_node(num_cpus=5, num_gpus=5)
     cluster.add_node(num_cpus=10, num_gpus=1)
     ray.init(redis_address=cluster.redis_address)
-    cluster.wait_for_nodes(3)
+    cluster.wait_for_nodes()
 
     # Define a bunch of remote functions that all return the socket name of
     # the plasma store. Since there is a one-to-one correspondence between