wehub-resource-sync
152 lines
4.3 KiB
Python
152 lines
4.3 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
|
|
|
|
class MyError(Exception):
|
|
pass
|
|
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self) -> None:
|
|
self.count = 0
|
|
|
|
def increment(self) -> int:
|
|
c = self.count
|
|
self.count += 1
|
|
return c
|
|
|
|
def get_count(self) -> int:
|
|
return self.count
|
|
|
|
|
|
# TODO: also do work for async and threaded actors
|
|
@ray.remote(max_task_retries=3)
|
|
class TroubleMaker:
|
|
@ray.method(max_task_retries=5, retry_exceptions=[MyError])
|
|
def may_raise_n_times(self, counter, n):
|
|
"""
|
|
Raises if there were n calls before this call.
|
|
Returns the number of calls before this call, if it's > n.
|
|
"""
|
|
c = ray.get(counter.increment.remote())
|
|
print(f"may_raise_n_times, n = {n}, count = {c}")
|
|
if c < n:
|
|
print(f"method raises in {c} th call, want {n} times")
|
|
raise MyError()
|
|
return c
|
|
|
|
@ray.method(retry_exceptions=[MyError])
|
|
def raise_or_exit(self, counter, actions):
|
|
"""
|
|
Increments the counter and performs an action based on the param actions[count].
|
|
If count >= len(actions), return the count.
|
|
|
|
Note: This method doesn't set `max_task_retries`. Ray expects it to inherit
|
|
max_task_retries = 3.
|
|
|
|
@param actions: List["raise" | "exit"]
|
|
"""
|
|
|
|
c = ray.get(counter.increment.remote())
|
|
action = "return" if c >= len(actions) else actions[c]
|
|
print(f"raise_or_exit, action = {action}, count = {c}")
|
|
|
|
if action == "raise":
|
|
raise MyError()
|
|
elif action == "exit":
|
|
sys.exit(1)
|
|
else:
|
|
return c
|
|
|
|
|
|
def test_method_raise_5_times(shutdown_only):
|
|
counter = Counter.remote()
|
|
trouble_maker = TroubleMaker.remote()
|
|
assert ray.get(trouble_maker.may_raise_n_times.remote(counter, 5)) == 5
|
|
assert ray.get(counter.get_count.remote()) == 6
|
|
|
|
|
|
def test_method_raise_no_over_retry(shutdown_only):
|
|
counter = Counter.remote()
|
|
trouble_maker = TroubleMaker.remote()
|
|
with pytest.raises(MyError):
|
|
ray.get(trouble_maker.may_raise_n_times.remote(counter, 6))
|
|
assert ray.get(counter.get_count.remote()) == 6
|
|
|
|
|
|
def test_method_no_retry_without_retry_exceptions(shutdown_only):
|
|
counter = Counter.remote()
|
|
trouble_maker = TroubleMaker.remote()
|
|
with pytest.raises(MyError):
|
|
ray.get(
|
|
trouble_maker.may_raise_n_times.options(retry_exceptions=False).remote(
|
|
counter, 5
|
|
)
|
|
)
|
|
assert ray.get(counter.get_count.remote()) == 1
|
|
|
|
|
|
def test_options_takes_precedence(shutdown_only):
|
|
counter = Counter.remote()
|
|
trouble_maker = TroubleMaker.remote()
|
|
assert (
|
|
ray.get(
|
|
trouble_maker.may_raise_n_times.options(max_task_retries=10).remote(
|
|
counter, 10
|
|
)
|
|
)
|
|
== 10
|
|
)
|
|
assert ray.get(counter.get_count.remote()) == 11
|
|
|
|
|
|
def test_options_takes_precedence_no_over_retry(shutdown_only):
|
|
counter = Counter.remote()
|
|
trouble_maker = TroubleMaker.remote()
|
|
|
|
with pytest.raises(MyError):
|
|
ray.get(
|
|
trouble_maker.may_raise_n_times.options(max_task_retries=10).remote(
|
|
counter, 11
|
|
)
|
|
)
|
|
assert ray.get(counter.get_count.remote()) == 11
|
|
|
|
|
|
def test_method_exit_no_over_retry_max_restarts(shutdown_only):
|
|
"""
|
|
Even if we have enough max_task_retries, we may still raise due to max_restarts.
|
|
"""
|
|
counter = Counter.remote()
|
|
trouble_maker = TroubleMaker.options(max_restarts=1).remote()
|
|
with pytest.raises(ray.exceptions.RayActorError):
|
|
assert ray.get(
|
|
trouble_maker.raise_or_exit.options(max_task_retries=4).remote(
|
|
counter, ["raise", "exit", "exit"]
|
|
)
|
|
)
|
|
# 2 calls: 1 initial + 1 exception-retry + 1 exit-retry (then no more)
|
|
assert ray.get(counter.get_count.remote()) == 3
|
|
|
|
|
|
def test_exit_only_no_over_retry(shutdown_only):
|
|
"""
|
|
Sanity testing: only do exit-retry works
|
|
"""
|
|
counter = Counter.remote()
|
|
trouble_maker = TroubleMaker.options(max_restarts=2).remote()
|
|
ray.get(
|
|
trouble_maker.raise_or_exit.options(max_task_retries=2).remote(
|
|
counter, ["exit", "exit"]
|
|
)
|
|
) == 2
|
|
assert ray.get(counter.get_count.remote()) == 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|