Retry and exception for hang on memory store full (#5143)

python/ray/tests/test_failure.py

@@ -4,6 +4,7 @@ from __future__ import print_function
 
 import json
 import os
+import pyarrow.plasma as plasma
 import pytest
 import sys
 import tempfile
@@ -719,3 +720,48 @@ def test_connect_with_disconnected_node(shutdown_only):
     # There is no connection error to a dead node.
     info = relevant_errors(ray_constants.RAYLET_CONNECTION_ERROR)
     assert len(info) == 0
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster_head", [{
+        "num_cpus": 5,
+        "object_store_memory": 10**7
+    }],
+    indirect=True)
+@pytest.mark.parametrize("num_actors", [1, 2, 5])
+def test_parallel_actor_fill_plasma_retry(ray_start_cluster_head, num_actors):
+    @ray.remote
+    class LargeMemoryActor(object):
+        def some_expensive_task(self):
+            return np.zeros(10**7 // 2, dtype=np.uint8)
+
+    actors = [LargeMemoryActor.remote() for _ in range(num_actors)]
+    for _ in range(10):
+        pending = [a.some_expensive_task.remote() for a in actors]
+        while pending:
+            [done], pending = ray.wait(pending, num_returns=1)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster_head", [{
+        "num_cpus": 2,
+        "object_store_memory": 10**7
+    }],
+    indirect=True)
+def test_fill_plasma_exception(ray_start_cluster_head):
+    @ray.remote
+    class LargeMemoryActor(object):
+        def some_expensive_task(self):
+            return np.zeros(10**7 + 2, dtype=np.uint8)
+
+        def test(self):
+            return 1
+
+    actor = LargeMemoryActor.remote()
+    with pytest.raises(ray.exceptions.RayTaskError):
+        ray.get(actor.some_expensive_task.remote())
+    # Make sure actor does not die
+    ray.get(actor.test.remote())
+
+    with pytest.raises(plasma.PlasmaStoreFull):
+        ray.put(np.zeros(10**7 + 2, dtype=np.uint8))