import contextlib import os import random import sys import threading import time from concurrent import futures from typing import Any, Callable, Optional from unittest.mock import Mock, patch import grpc import numpy as np import pytest import ray import ray.core.generated.ray_client_pb2 as ray_client_pb2 import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc from ray._common.utils import get_or_create_event_loop from ray.tests.conftest import call_ray_start_context from ray.util.client.common import CLIENT_SERVER_MAX_THREADS, GRPC_OPTIONS @pytest.fixture(scope="module") def call_ray_start_shared(request): # Starts Ray with a ray client server listening on port 50051 request = Mock() request.param = ( "ray start --head --min-worker-port=0 --max-worker-port=0 --port 0 " "--ray-client-server-port=50051" ) with call_ray_start_context(request) as address: yield address # At a high level, these tests rely on an extra RPC server sitting # between the client and the real Ray server to inject errors, drop responses # and drop requests, i.e. at a high level: # Ray Client <-> Middleman Server <-> Proxy Server # Type for middleman hooks used to inject errors Hook = Callable[[Any], None] class MiddlemanDataServicer(ray_client_pb2_grpc.RayletDataStreamerServicer): """ Forwards all requests to the real data servicer. Useful for injecting errors between a client and server pair. """ def __init__( self, on_response: Optional[Hook] = None, on_request: Optional[Hook] = None ): """Initialize the middleman data servicer. Args: on_response: Optional hook to inject errors before sending back a response. on_request: Optional hook to inject errors before forwarding a request. """ self.stub = None self.on_response = on_response self.on_request = on_request def set_channel(self, channel: grpc.Channel) -> None: self.stub = ray_client_pb2_grpc.RayletDataStreamerStub(channel) def _requests(self, request_iterator): for req in request_iterator: if self.on_request: self.on_request(req) yield req def Datapath(self, request_iterator, context): try: for response in self.stub.Datapath( self._requests(request_iterator), metadata=context.invocation_metadata() ): if self.on_response: self.on_response(response) yield response except grpc.RpcError as e: context.set_code(e.code()) context.set_details(e.details()) class MiddlemanLogServicer(ray_client_pb2_grpc.RayletLogStreamerServicer): """ Forwards all requests to the real log servicer. Useful for injecting errors between a client and server pair. """ def __init__(self, on_response: Optional[Hook] = None): """Initialize the middleman log servicer. Args: on_response: Optional hook to inject errors before sending back a response. """ self.stub = None self.on_response = on_response def set_channel(self, channel: grpc.Channel) -> None: self.stub = ray_client_pb2_grpc.RayletLogStreamerStub(channel) def Logstream(self, request_iterator, context): try: for response in self.stub.Logstream( request_iterator, metadata=context.invocation_metadata() ): if self.on_response: self.on_response(response) yield response except grpc.RpcError as e: context.set_code(e.code()) context.set_details(e.details()) class MiddlemanRayletServicer(ray_client_pb2_grpc.RayletDriverServicer): """ Forwards all requests to the raylet driver servicer. Useful for injecting errors between a client and server pair. """ def __init__( self, on_request: Optional[Hook] = None, on_response: Optional[Hook] = None ): """Initialize the middleman raylet servicer. Args: on_request: Optional hook to inject errors before forwarding a request. on_response: Optional hook to inject errors before sending back a response. """ self.stub = None self.on_request = on_request self.on_response = on_response def set_channel(self, channel: grpc.Channel) -> None: self.stub = ray_client_pb2_grpc.RayletDriverStub(channel) def _call_inner_function( self, request: Any, context, method: str ) -> Optional[ray_client_pb2_grpc.RayletDriverStub]: if self.on_request: self.on_request(request) try: response = getattr(self.stub, method)( request, metadata=context.invocation_metadata() ) except grpc.RpcError as e: context.set_code(e.code()) context.set_details(e.details()) raise if self.on_response and method != "GetObject": # GetObject streams response, handle on_response separately self.on_response(response) return response def Init(self, request, context=None) -> ray_client_pb2.InitResponse: return self._call_inner_function(request, context, "Init") def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse: return self._call_inner_function(request, context, "KVPut") def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse: return self._call_inner_function(request, context, "KVGet") def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse: return self._call_inner_function(request, context, "KVDel") def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse: return self._call_inner_function(request, context, "KVList") def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse: return self._call_inner_function(request, context, "KVExists") def ListNamedActors( self, request, context=None ) -> ray_client_pb2.ClientListNamedActorsResponse: return self._call_inner_function(request, context, "ListNamedActors") def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse: # Cluster info is currently used for health checks and isn't retried, so # don't inject errors. # TODO(ckw): update ClusterInfo so that retries are only skipped for PING try: return self.stub.ClusterInfo( request, metadata=context.invocation_metadata() ) except grpc.RpcError as e: context.set_code(e.code()) context.set_details(e.details()) raise def Terminate(self, req, context=None): return self._call_inner_function(req, context, "Terminate") def GetObject(self, request, context=None): for response in self._call_inner_function(request, context, "GetObject"): if self.on_response: self.on_response(response) yield response def PutObject( self, request: ray_client_pb2.PutRequest, context=None ) -> ray_client_pb2.PutResponse: return self._call_inner_function(request, context, "PutObject") def WaitObject( self, request: ray_client_pb2.WaitRequest, context=None ) -> ray_client_pb2.WaitResponse: return self._call_inner_function(request, context, "WaitObject") def Schedule( self, task: ray_client_pb2.ClientTask, context=None ) -> ray_client_pb2.ClientTaskTicket: return self._call_inner_function(task, context, "Schedule") class MiddlemanServer: """ Helper class that wraps the RPC server that middlemans the connection between the client and the real ray server. Useful for injecting errors between a client and server pair. """ def __init__( self, listen_addr: str, real_addr: str, on_log_response: Optional[Hook] = None, on_data_request: Optional[Hook] = None, on_data_response: Optional[Hook] = None, on_task_request: Optional[Hook] = None, on_task_response: Optional[Hook] = None, ): """Initialize the middleman server. Args: listen_addr: The address the middleman server will listen on. real_addr: The address of the real ray server. on_log_response: Optional hook to inject errors before sending back a log response. on_data_request: Optional hook to inject errors before forwarding a data request. on_data_response: Optional hook to inject errors before sending back a data response. on_task_request: Optional hook to inject errors before forwarding a raylet driver request. on_task_response: Optional hook to inject errors before sending back a raylet driver response. """ self.listen_addr = listen_addr self.real_addr = real_addr self.server = grpc.server( futures.ThreadPoolExecutor(max_workers=CLIENT_SERVER_MAX_THREADS), options=GRPC_OPTIONS, ) self.task_servicer = MiddlemanRayletServicer( on_response=on_task_response, on_request=on_task_request ) self.data_servicer = MiddlemanDataServicer( on_response=on_data_response, on_request=on_data_request ) self.logs_servicer = MiddlemanLogServicer(on_response=on_log_response) ray_client_pb2_grpc.add_RayletDriverServicer_to_server( self.task_servicer, self.server ) ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server( self.data_servicer, self.server ) ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server( self.logs_servicer, self.server ) self.server.add_insecure_port(self.listen_addr) self.channel = None self.reset_channel() def reset_channel(self) -> None: """ Manually close and reopen the channel to the real ray server. This simulates a disconnection between the client and the server. """ if self.channel: self.channel.close() self.channel = grpc.insecure_channel(self.real_addr, options=GRPC_OPTIONS) grpc.channel_ready_future(self.channel) self.task_servicer.set_channel(self.channel) self.data_servicer.set_channel(self.channel) self.logs_servicer.set_channel(self.channel) def start(self) -> None: self.server.start() def stop(self, grace: int) -> None: self.server.stop(grace) @contextlib.contextmanager def start_middleman_server( on_log_response=None, on_data_request=None, on_data_response=None, on_task_request=None, on_task_response=None, ): """ Helper context that starts a middleman server listening on port 10011, and a ray client server on port 50051. """ ray._inside_client_test = True middleman = None try: middleman = MiddlemanServer( listen_addr="localhost:10011", real_addr="localhost:50051", on_log_response=on_log_response, on_data_request=on_data_request, on_data_response=on_data_response, on_task_request=on_task_request, on_task_response=on_task_response, ) middleman.start() ray.init("ray://localhost:10011") yield middleman finally: ray._inside_client_test = False ray.util.disconnect() if middleman: middleman.stop(0) def test_disconnect_during_get(call_ray_start_shared): """ Disconnect the proxy and the client in the middle of a long running get """ @ray.remote def slow_result(): time.sleep(20) return 12345 def disconnect(middleman): time.sleep(3) middleman.reset_channel() with start_middleman_server() as middleman: disconnect_thread = threading.Thread(target=disconnect, args=(middleman,)) disconnect_thread.start() result = ray.get(slow_result.remote()) assert result == 12345 disconnect_thread.join() def test_disconnects_during_large_get(call_ray_start_shared): """ Disconnect repeatedly during a large (multi-chunk) get. """ i = 0 started = False def fail_every_three(_): # Inject an error every third time this method is called nonlocal i, started if not started: return i += 1 if i % 3 == 0: raise RuntimeError @ray.remote def large_result(): # 1024x1024x6 float64 matrix (96 MiB). With 5MiB chunk size, # it will take at least 16 chunks to transfer this object. Since # the failure is injected every 3 chunks, this transfer can only # work if the chunked get request retries at the last received chunk # (instead of starting from the beginning each retry) return np.random.random((1024, 1024, 6)) with start_middleman_server(on_task_response=fail_every_three): started = True result = ray.get(large_result.remote()) assert result.shape == (1024, 1024, 6) def test_disconnects_during_large_async_get(call_ray_start_shared): """ Disconnect repeatedly during a large (multi-chunk) async get. """ i = 0 started = False def fail_every_three(_): # Inject an error every third time this method is called nonlocal i, started if not started: return i += 1 if i % 3 == 0: raise RuntimeError @ray.remote def large_result(): # 1024x1024x6 float64 matrix (96 MiB). With 5MiB chunk size, # it will take at least 16 chunks to transfer this object. Since # the failure is injected every 3 chunks, this transfer can only # work if the chunked get request retries at the last received chunk # (instead of starting from the beginning each retry) return np.random.random((1024, 1024, 6)) with start_middleman_server(on_data_response=fail_every_three): started = True async def get_large_result(): return await large_result.remote() result = get_or_create_event_loop().run_until_complete(get_large_result()) assert result.shape == (1024, 1024, 6) def test_disconnect_during_large_put(call_ray_start_shared): """ Disconnect during a large (multi-chunk) put. """ i = 0 started = False def fail_halfway(_): # Inject an error halfway through the object transfer nonlocal i, started if not started: return i += 1 if i == 8: raise RuntimeError with start_middleman_server(on_data_request=fail_halfway): started = True objref = ray.put(np.random.random((1024, 1024, 6))) assert i > 8 # Check that the failure was injected result = ray.get(objref) assert result.shape == (1024, 1024, 6) def test_disconnect_during_large_schedule(call_ray_start_shared): """ Disconnect during a remote call with a large (multi-chunk) argument. """ i = 0 started = False def fail_halfway(_): # Inject an error halfway through the object transfer nonlocal i, started if not started: return i += 1 if i == 8: raise RuntimeError @ray.remote def f(a): return a.shape with start_middleman_server(on_data_request=fail_halfway): started = True a = np.random.random((1024, 1024, 6)) result = ray.get(f.remote(a)) assert i > 8 # Check that the failure was injected assert result == (1024, 1024, 6) def test_valid_actor_state(call_ray_start_shared): """ Repeatedly inject errors in the middle of mutating actor calls. Check at the end that the final state of the actor is consistent with what we would expect had the disconnects not occurred. """ @ray.remote class IncrActor: def __init__(self): self.val = 0 def incr(self): self.val += 1 return self.val i = 0 # This is to prevent erroring in the initial connection logic. started = False def fail_every_seven(_): # Inject an error every seventh time this method is called nonlocal i, started i += 1 if i % 7 == 0 and started: raise RuntimeError with start_middleman_server( on_data_response=fail_every_seven, on_task_request=fail_every_seven, on_task_response=fail_every_seven, ): started = True actor = IncrActor.remote() for _ in range(100): ref = actor.incr.remote() assert ray.get(ref) == 100 def test_valid_actor_state_2(call_ray_start_shared): """ Do a full disconnect (cancel channel) every 11 requests. Failure happens: - before request sent: request never reaches server - before response received: response never reaches server - while get's are being processed """ @ray.remote class IncrActor: def __init__(self): self.val = 0 def incr(self): self.val += 1 return self.val i = 0 with start_middleman_server() as middleman: def fail_every_eleven(_): nonlocal i i += 1 if i % 11 == 0: middleman.reset_channel() middleman.data_servicer.on_response = fail_every_eleven middleman.task_servicer.on_request = fail_every_eleven middleman.task_servicer.on_response = fail_every_eleven actor = IncrActor.remote() for _ in range(100): ref = actor.incr.remote() assert ray.get(ref) == 100 def test_noisy_puts(call_ray_start_shared): """ Randomly kills the data channel with 10% chance when receiving response (requests made it to server, responses dropped) and checks that final result is still consistent """ random.seed(12345) with start_middleman_server() as middleman: def fail_randomly(response: ray_client_pb2.DataResponse): if random.random() < 0.1: raise RuntimeError middleman.data_servicer.on_response = fail_randomly refs = [ray.put(i * 123) for i in range(500)] results = ray.get(refs) for i, result in enumerate(results): assert result == i * 123 def test_client_reconnect_grace_period(call_ray_start_shared): """ Tests that the client gives up attempting to reconnect the channel after the grace period expires. """ # Lower grace period to 5 seconds to save time with patch.dict( os.environ, {"RAY_CLIENT_RECONNECT_GRACE_PERIOD": "5"} ), start_middleman_server() as middleman: assert ray.get(ray.put(42)) == 42 # Close channel middleman.channel.close() start_time = time.time() with pytest.raises(ConnectionError): ray.get(ray.put(42)) # Connection error should have been raised within a reasonable # amount of time. Set to significantly higher than 5 seconds # to account for reconnect backoff timing assert time.time() - start_time < 20 if __name__ == "__main__": sys.exit(pytest.main(["-sv", __file__]))