From 9460a5375b4ba6d2201cb9e59f8cde20d2633e2f Mon Sep 17 00:00:00 2001
From: Gagandeep Singh <gdp.1807@gmail.com>
Date: Fri, 29 Oct 2021 22:45:12 +0530
Subject: [PATCH] Added retry logic in test_basic::test_ray_options (#19832)

* Added retry logic in test_ray_options

* Applied linting format

* Made test consistent
---
 python/ray/tests/test_basic.py | 43 ++++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 13 deletions(-)

diff --git a/python/ray/tests/test_basic.py b/python/ray/tests/test_basic.py
index c3b515aa7..d302ff4d7 100644
--- a/python/ray/tests/test_basic.py
+++ b/python/ray/tests/test_basic.py
@@ -303,27 +303,44 @@ def test_ray_options(shutdown_only):
 
     @ray.remote(
         num_cpus=2, num_gpus=3, memory=150 * 2**20, resources={"custom1": 1})
-    def foo():
-        import time
-        # Sleep for a heartbeat period to ensure resources changing reported.
-        time.sleep(0.1)
-        return ray.available_resources()
+    def foo(expected_resources):
+        # Possibly wait until the available resources have been updated
+        # (there might be a delay due to heartbeats)
+        retries = 10
+        keys = ["CPU", "GPU", "custom1"]
+        while retries >= 0:
+            resources = ray.available_resources()
+            do_return = True
+            for key in keys:
+                if resources[key] != expected_resources[key]:
+                    print(key, resources[key], expected_resources[key])
+                    do_return = False
+                    break
+            if do_return:
+                return resources["memory"]
+            time.sleep(0.1)
+            retries -= 1
+        raise RuntimeError("Number of retries exceeded")
 
-    without_options = ray.get(foo.remote())
-    with_options = ray.get(
+    expected_resources_without_options = {
+        "CPU": 8.0,
+        "GPU": 7.0,
+        "custom1": 1.0
+    }
+    memory_available_without_options = ray.get(
+        foo.remote(expected_resources_without_options))
+
+    expected_resources_with_options = {"CPU": 7.0, "GPU": 6.0, "custom1": 1.5}
+    memory_available_with_options = ray.get(
         foo.options(
             num_cpus=3,
             num_gpus=4,
             memory=50 * 2**20,
             resources={
                 "custom1": 0.5
-            }).remote())
+            }).remote(expected_resources_with_options))
 
-    to_check = ["CPU", "GPU", "memory", "custom1"]
-    for key in to_check:
-        print(key, without_options[key], with_options[key])
-        assert without_options[key] != with_options[key], key
-    assert without_options != with_options
+    assert memory_available_without_options < memory_available_with_options
 
 
 @pytest.mark.skipif(client_test_enabled(), reason="internal api")