fix test_worker_stats (ray-project#7655)

* fix test_worker_stats * fix lint error * fix lint error Co-authored-by: senlin.zsl <[email protected]>
LecJackS · Mar 20, 2020 · 7d08b41 · 7d08b41
1 parent e69664b
commit 7d08b41
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/python/ray/test_utils.py b/python/ray/test_utils.py
@@ -6,6 +6,7 @@
 import sys
 import tempfile
 import time
+import socket
 
 import ray
 
@@ -234,3 +235,26 @@ def put_object(obj, use_ray_put):
         return ray.put(obj)
     else:
         return _put.remote(obj)
+
+
+def wait_until_server_available(address,
+                                timeout_ms=5000,
+                                retry_interval_ms=100):
+    ip_port = address.split(":")
+    ip = ip_port[0]
+    port = int(ip_port[1])
+    time_elapsed = 0
+    start = time.time()
+    while time_elapsed <= timeout_ms:
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.settimeout(1)
+        try:
+            s.connect((ip, port))
+        except Exception:
+            time_elapsed = (time.time() - start) * 1000
+            time.sleep(retry_interval_ms / 1000.0)
+            s.close()
+            continue
+        s.close()
+        return True
+    return False
diff --git a/python/ray/tests/test_metrics.py b/python/ray/tests/test_metrics.py
@@ -11,7 +11,8 @@
 from ray.core.generated import reporter_pb2
 from ray.core.generated import reporter_pb2_grpc
 from ray.test_utils import (RayTestTimeoutException,
-                            wait_until_succeeded_without_exception)
+                            wait_until_succeeded_without_exception,
+                            wait_until_server_available)
 
 import psutil  # We must import psutil after ray because we bundle it with ray.
 
@@ -125,6 +126,8 @@ def actor_killed(PID):
         else:
             return False
 
+    assert (wait_until_server_available(addresses["webui_url"]) is True)
+
     webui_url = addresses["webui_url"]
     webui_url = webui_url.replace("localhost", "http://127.0.0.1")
     for worker in reply.workers_stats:
@@ -185,6 +188,8 @@ def getpid(self):
     c.local_store.remote()
     c.remote_store.remote()
 
+    assert (wait_until_server_available(addresses["webui_url"]) is True)
+
     start_time = time.time()
     while True:
         time.sleep(1)
@@ -267,6 +272,7 @@ def __init__(self):
     ActorRequiringGPU.remote()
 
     def test_infeasible_actor(ray_addresses):
+        assert (wait_until_server_available(addresses["webui_url"]) is True)
         webui_url = ray_addresses["webui_url"].replace("localhost",
                                                        "http://127.0.0.1")
         raylet_info = requests.get(webui_url + "/api/raylet_info").json()
@@ -306,6 +312,7 @@ def __init__(self):
     assert parent_actor is not None
 
     def test_pending_actor(ray_addresses):
+        assert (wait_until_server_available(addresses["webui_url"]) is True)
         webui_url = ray_addresses["webui_url"].replace("localhost",
                                                        "http://127.0.0.1")
         raylet_info = requests.get(webui_url + "/api/raylet_info").json()