wehub-resource-sync
113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
import sys
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from ray.exceptions import (
|
|
GetTimeoutError,
|
|
ObjectLostError,
|
|
RayTaskError,
|
|
TaskCancelledError,
|
|
WorkerCrashedError,
|
|
)
|
|
from ray.tests.client_test_utils import create_remote_signal_actor
|
|
from ray.util.client.ray_client_helpers import ray_start_client_server
|
|
|
|
|
|
def valid_exceptions(use_force):
|
|
if use_force:
|
|
return (RayTaskError, TaskCancelledError, WorkerCrashedError, ObjectLostError)
|
|
else:
|
|
return (RayTaskError, TaskCancelledError)
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
|
|
@pytest.mark.parametrize("use_force", [True, False])
|
|
def test_cancel_chain(ray_start_regular, use_force):
|
|
with ray_start_client_server() as ray:
|
|
SignalActor = create_remote_signal_actor(ray)
|
|
signaler = SignalActor.remote()
|
|
|
|
@ray.remote
|
|
def wait_for(t):
|
|
return ray.get(t[0])
|
|
|
|
obj1 = wait_for.remote([signaler.wait.remote()])
|
|
obj2 = wait_for.remote([obj1])
|
|
obj3 = wait_for.remote([obj2])
|
|
obj4 = wait_for.remote([obj3])
|
|
|
|
assert len(ray.wait([obj1], timeout=0.1)[0]) == 0
|
|
ray.cancel(obj1, force=use_force)
|
|
for ob in [obj1, obj2, obj3, obj4]:
|
|
with pytest.raises(valid_exceptions(use_force)):
|
|
ray.get(ob)
|
|
|
|
signaler2 = SignalActor.remote()
|
|
obj1 = wait_for.remote([signaler2.wait.remote()])
|
|
obj2 = wait_for.remote([obj1])
|
|
obj3 = wait_for.remote([obj2])
|
|
obj4 = wait_for.remote([obj3])
|
|
|
|
assert len(ray.wait([obj3], timeout=0.1)[0]) == 0
|
|
ray.cancel(obj3, force=use_force)
|
|
for ob in [obj3, obj4]:
|
|
with pytest.raises(valid_exceptions(use_force)):
|
|
ray.get(ob)
|
|
|
|
with pytest.raises(GetTimeoutError):
|
|
ray.get(obj1, timeout=0.1)
|
|
|
|
with pytest.raises(GetTimeoutError):
|
|
ray.get(obj2, timeout=0.1)
|
|
|
|
signaler2.send.remote()
|
|
ray.get(obj1)
|
|
|
|
|
|
def test_kill_cancel_metadata(ray_start_regular):
|
|
"""
|
|
Verifies that client worker's terminate_actor and terminate_task methods
|
|
pass worker's metadata attribute server to the grpc stub's Terminate
|
|
method.
|
|
|
|
This is done by mocking the grpc stub's Terminate method to raise an
|
|
exception with argument equal to the key of the metadata. We then verify
|
|
that the exception is raised when calling ray.kill and ray.cancel.
|
|
"""
|
|
with ray_start_client_server(metadata=[("key", "value")]) as ray:
|
|
|
|
@ray.remote
|
|
class A:
|
|
pass
|
|
|
|
@ray.remote
|
|
def f():
|
|
time.sleep(1000)
|
|
|
|
class MetadataIsCorrectlyPassedException(Exception):
|
|
pass
|
|
|
|
def mock_terminate(self, term):
|
|
raise MetadataIsCorrectlyPassedException(self._metadata[1][0])
|
|
|
|
# Mock stub's Terminate method to raise an exception.
|
|
stub = ray.get_context().api.worker.data_client
|
|
stub.Terminate = mock_terminate.__get__(stub)
|
|
|
|
# Verify the expected exception is raised with ray.kill.
|
|
# Check that argument of the exception matches "key" from the
|
|
# metadata above.
|
|
actor = A.remote()
|
|
with pytest.raises(MetadataIsCorrectlyPassedException, match="key"):
|
|
ray.kill(actor)
|
|
|
|
# Verify the expected exception is raised with ray.cancel.
|
|
task_ref = f.remote()
|
|
with pytest.raises(MetadataIsCorrectlyPassedException, match="key"):
|
|
ray.cancel(task_ref)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|