ray-project--ray/python/ray/tests/test_client_reconnect.py

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__]))