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

@@ -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():
+    def foo(expected_resources):
-        import time
+        # Possibly wait until the available resources have been updated
-        # Sleep for a heartbeat period to ensure resources changing reported.
+        # (there might be a delay due to heartbeats)
-        time.sleep(0.1)
+        retries = 10
-        return ray.available_resources()
+        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())
+    expected_resources_without_options = {
-    with_options = ray.get(
+        "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"]
+    assert memory_available_without_options < memory_available_with_options
-    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
 
 
 @pytest.mark.skipif(client_test_enabled(), reason="internal api")