Repository: mars-project/mars

Issue:
[BUG] Mars worker not recovered on ray master
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
In ray master, If a actor created with `max_restarts=-1` is restarting, call actor method will raise exception instead of pending in caller, which make ray worker failover not work. 

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version: 3.7.9
2. The version of Mars you use: master
3. Versions of crucial packages, such as ray, numpy, scipy and pandas: ray master
4. Full stack of the error.
![image](https://user-images.githubusercontent.com/12445254/169981083-5098b0de-3ced-4930-b8ff-e374c6676140.png)
5. Minimized code to reproduce the error.
```
import ray
import time

@ray.remote(max_restarts=-1)
class A:
  def __init__(self):
    if ray.get_runtime_context().was_current_actor_reconstructed:
      import os
      os._exit(-1)
    print(ray.get_runtime_context().was_current_actor_reconstructed)
    time.sleep(3)

  def f(self):
    return 1

  def f1():
    time.sleep(30)
a = A.remote()
print(ray.get(a.f.remote()))
r = a.f1.remote()
ray.kill(a, no_restart=False)
try:
  ray.get(r)
except Exception:
  pass
print(ray.get(a.f.remote()))
```
**Expected behavior**
This can be fixed by specifying max_retries=-1 when call actro methods

**Additional context**
Add any other context about the problem here.



Discussion:

Resolving patch:
diff --git a/mars/deploy/oscar/ray.py b/mars/deploy/oscar/ray.py
--- a/mars/deploy/oscar/ray.py
+++ b/mars/deploy/oscar/ray.py
@@ -37,7 +37,7 @@
     AbstractClusterBackend,
 )
 from ...services import NodeRole
-from ...utils import lazy_import
+from ...utils import lazy_import, retry_callable
 from ..utils import (
     load_config,
     get_third_party_modules_from_config,
@@ -274,7 +274,10 @@ async def reconstruct_worker(self, address: str):
         async def _reconstruct_worker():
             logger.info("Reconstruct worker %s", address)
             actor = ray.get_actor(address)
-            state = await actor.state.remote()
+            # ray call will error when actor is restarting
+            state = await retry_callable(
+                actor.state.remote, ex_type=ray.exceptions.RayActorError, sync=False
+            )()
             if state == RayPoolState.SERVICE_READY:
                 logger.info("Worker %s is service ready.")
                 return
diff --git a/mars/oscar/backends/ray/pool.py b/mars/oscar/backends/ray/pool.py
--- a/mars/oscar/backends/ray/pool.py
+++ b/mars/oscar/backends/ray/pool.py
@@ -28,7 +28,7 @@
 
 from ... import ServerClosed
 from ....serialization.ray import register_ray_serializers
-from ....utils import lazy_import, ensure_coverage
+from ....utils import lazy_import, ensure_coverage, retry_callable
 from ..config import ActorPoolConfig
 from ..message import CreateActorMessage
 from ..pool import (
@@ -130,14 +130,27 @@ async def start_sub_pool(
             f"process_index {process_index} is not consistent with index {_process_index} "
             f"in external_address {external_address}"
         )
+        actor_handle = config["kwargs"]["sub_pool_handles"][external_address]
+        state = await retry_callable(
+            actor_handle.state.remote, ex_type=ray.exceptions.RayActorError, sync=False
+        )()
+        if state is RayPoolState.SERVICE_READY:  # pragma: no cover
+            logger.info("Ray sub pool %s is alive, kill it first.", external_address)
+            await kill_and_wait(actor_handle, no_restart=False)
+            # Wait sub pool process restarted.
+            await retry_callable(
+                actor_handle.state.remote,
+                ex_type=ray.exceptions.RayActorError,
+                sync=False,
+            )()
         logger.info("Start to start ray sub pool %s.", external_address)
         create_sub_pool_timeout = 120
-        actor_handle = config["kwargs"]["sub_pool_handles"][external_address]
-        done, _ = await asyncio.wait(
-            [actor_handle.set_actor_pool_config.remote(actor_pool_config)],
-            timeout=create_sub_pool_timeout,
-        )
-        if not done:  # pragma: no cover
+        try:
+            await asyncio.wait_for(
+                actor_handle.set_actor_pool_config.remote(actor_pool_config),
+                timeout=create_sub_pool_timeout,
+            )
+        except asyncio.TimeoutError:  # pragma: no cover
             msg = (
                 f"Can not start ray sub pool {external_address} in {create_sub_pool_timeout} seconds.",
             )
@@ -153,6 +166,10 @@ async def wait_sub_pools_ready(cls, create_pool_tasks: List[asyncio.Task]):
 
     async def recover_sub_pool(self, address: str):
         process = self.sub_processes[address]
+        # ray call will error when actor is restarting
+        await retry_callable(
+            process.state.remote, ex_type=ray.exceptions.RayActorError, sync=False
+        )()
         await process.start.remote()
 
         if self._auto_recover == "actor":
diff --git a/mars/utils.py b/mars/utils.py
--- a/mars/utils.py
+++ b/mars/utils.py
@@ -1698,3 +1698,41 @@ def ensure_coverage():
             pass
         else:
             cleanup_on_sigterm()
+
+
+def retry_callable(
+    callable_,
+    ex_type: type = Exception,
+    wait_interval=1,
+    max_retries=-1,
+    sync: bool = None,
+):
+    if inspect.iscoroutinefunction(callable_) or sync is False:
+
+        @functools.wraps(callable)
+        async def retry_call(*args, **kwargs):
+            num_retried = 0
+            while max_retries < 0 or num_retried < max_retries:
+                num_retried += 1
+                try:
+                    return await callable_(*args, **kwargs)
+                except ex_type:
+                    await asyncio.sleep(wait_interval)
+
+    else:
+
+        @functools.wraps(callable)
+        def retry_call(*args, **kwargs):
+            num_retried = 0
+            ex = None
+            while max_retries < 0 or num_retried < max_retries:
+                num_retried += 1
+                try:
+                    return callable_(*args, **kwargs)
+                except ex_type as e:
+                    ex = e
+                    time.sleep(wait_interval)
+            assert ex is not None
+            raise ex  # pylint: disable-msg=E0702
+
+    return retry_call


Test patch:
diff --git a/mars/deploy/oscar/tests/test_ray.py b/mars/deploy/oscar/tests/test_ray.py
--- a/mars/deploy/oscar/tests/test_ray.py
+++ b/mars/deploy/oscar/tests/test_ray.py
@@ -578,7 +578,7 @@ async def remote(self):
     class FakeActor:
         state = FakeActorMethod()
 
-    def _get_actor(*args):
+    def _get_actor(*args, **kwargs):
         return FakeActor
 
     async def _stop_worker(*args):
@@ -677,7 +677,8 @@ async def test_auto_scale_in(ray_large_cluster):
         assert await autoscaler_ref.get_dynamic_worker_nums() == 2
 
 
-@pytest.mark.timeout(timeout=1000)
+@pytest.mark.skip("Enable it when ray ownership bug is fixed")
+@pytest.mark.timeout(timeout=200)
 @pytest.mark.parametrize("ray_large_cluster", [{"num_nodes": 4}], indirect=True)
 @require_ray
 @pytest.mark.asyncio
diff --git a/mars/tests/test_utils.py b/mars/tests/test_utils.py
--- a/mars/tests/test_utils.py
+++ b/mars/tests/test_utils.py
@@ -616,3 +616,37 @@ def __call__(self, *args, **kwargs):
 def test_gen_random_id(id_length):
     rnd_id = utils.new_random_id(id_length)
     assert len(rnd_id) == id_length
+
+
+@pytest.mark.asyncio
+async def test_retry_callable():
+    assert utils.retry_callable(lambda x: x)(1) == 1
+    assert utils.retry_callable(lambda x: 0)(1) == 0
+
+    class CustomException(BaseException):
+        pass
+
+    def f1(x):
+        nonlocal num_retried
+        num_retried += 1
+        if num_retried == 3:
+            return x
+        raise CustomException
+
+    num_retried = 0
+    with pytest.raises(CustomException):
+        utils.retry_callable(f1)(1)
+    assert utils.retry_callable(f1, ex_type=CustomException)(1) == 1
+    num_retried = 0
+    with pytest.raises(CustomException):
+        utils.retry_callable(f1, max_retries=2, ex_type=CustomException)(1)
+    num_retried = 0
+    assert utils.retry_callable(f1, max_retries=3, ex_type=CustomException)(1) == 1
+
+    async def f2(x):
+        return f1(x)
+
+    num_retried = 0
+    with pytest.raises(CustomException):
+        await utils.retry_callable(f2)(1)
+    assert await utils.retry_callable(f2, ex_type=CustomException)(1) == 1