From 580b017b431ef71ab56164464afcf7d042785a90 Mon Sep 17 00:00:00 2001
From: Edward Oakes <ed.nmi.oakes@gmail.com>
Date: Mon, 2 Mar 2020 21:03:01 -0600
Subject: [PATCH] Fix flaky global GC tests (#7407)

---
 python/ray/tests/test_reference_counting.py | 27 +++++++++++++++------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/python/ray/tests/test_reference_counting.py b/python/ray/tests/test_reference_counting.py
index f69cdfbf4..64d7c3dec 100644
--- a/python/ray/tests/test_reference_counting.py
+++ b/python/ray/tests/test_reference_counting.py
@@ -12,7 +12,7 @@ import pytest
 
 import ray
 import ray.cluster_utils
-from ray.test_utils import SignalActor
+from ray.test_utils import SignalActor, wait_for_condition
 from ray.internal.internal_api import global_gc
 
 logger = logging.getLogger(__name__)
@@ -106,9 +106,12 @@ def test_global_gc(shutdown_only):
 
         # GC should be triggered for all workers, including the local driver.
         global_gc()
-        time.sleep(1)
-        assert local_ref() is None
-        assert not any(ray.get([a.has_garbage.remote() for a in actors]))
+
+        def check_refs_gced():
+            return (local_ref() is None and
+                    not any(ray.get([a.has_garbage.remote() for a in actors])))
+
+        wait_for_condition(check_refs_gced, timeout_ms=10000)
     finally:
         gc.enable()
 
@@ -155,16 +158,24 @@ def test_global_gc_when_full(shutdown_only):
         # object store. This should cause the captured ObjectIDs' numpy arrays
         # to be evicted.
         ray.put(np.zeros(80 * 1024 * 1024, dtype=np.uint8))
-        assert local_ref() is None
-        assert not any(ray.get([a.has_garbage.remote() for a in actors]))
+
+        def check_refs_gced():
+            return (local_ref() is None and
+                    not any(ray.get([a.has_garbage.remote() for a in actors])))
+
+        wait_for_condition(check_refs_gced, timeout_ms=10000)
 
         # Local driver.
         local_ref = weakref.ref(LargeObjectWithCyclicRef())
 
         # Remote workers.
         actors = [GarbageHolder.remote() for _ in range(2)]
-        assert local_ref() is not None
-        assert all(ray.get([a.has_garbage.remote() for a in actors]))
+
+        def check_refs_gced():
+            return (local_ref() is None and
+                    not any(ray.get([a.has_garbage.remote() for a in actors])))
+
+        wait_for_condition(check_refs_gced, timeout_ms=10000)
 
         # GC should be triggered for all workers, including the local driver,
         # when a remote task tries to put a return value that doesn't fit in