chore: import upstream snapshot with attribution

python/ray/util/client/server/__init__.py

@@ -0,0 +1 @@
+from ray.util.client.server.server import serve  # noqa

python/ray/util/client/server/__main__.py

@@ -0,0 +1,4 @@
+if __name__ == "__main__":
+    from ray.util.client.server.server import main
+
+    main()

python/ray/util/client/server/dataservicer.py

@@ -0,0 +1,415 @@
+import logging
+import sys
+import time
+from collections import defaultdict
+from queue import Queue
+from threading import Event, Lock, Thread
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Union
+
+import grpc
+
+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._private.client_mode_hook import disable_client_hook
+from ray.util.client.common import (
+    CLIENT_SERVER_MAX_THREADS,
+    OrderedResponseCache,
+    _propagate_error_in_context,
+)
+from ray.util.client.server.server_pickler import loads_from_client
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.util.client.server.server import RayletServicer
+
+logger = logging.getLogger(__name__)
+
+QUEUE_JOIN_SECONDS = 10
+
+
+def _get_reconnecting_from_context(context: Any) -> bool:
+    """
+    Get `reconnecting` from gRPC metadata, or False if missing.
+    """
+    metadata = dict(context.invocation_metadata())
+    val = metadata.get("reconnecting")
+    if val is None or val not in ("True", "False"):
+        logger.error(
+            f'Client connecting with invalid value for "reconnecting": {val}, '
+            "This may be because you have a mismatched client and server "
+            "version."
+        )
+        return False
+    return val == "True"
+
+
+def _should_cache(req: ray_client_pb2.DataRequest) -> bool:
+    """
+    Returns True if the response should to the given request should be cached,
+    false otherwise. At the moment the only requests we do not cache are:
+        - asynchronous gets: These arrive out of order. Skipping caching here
+            is fine, since repeating an async get is idempotent
+        - acks: Repeating acks is idempotent
+        - clean up requests: Also idempotent, and client has likely already
+             wrapped up the data connection by this point.
+        - puts: We should only cache when we receive the final chunk, since
+             any earlier chunks won't generate a response
+        - tasks: We should only cache when we receive the final chunk,
+             since any earlier chunks won't generate a response
+    """
+    req_type = req.WhichOneof("type")
+    if req_type == "get" and req.get.asynchronous:
+        return False
+    if req_type == "put":
+        return req.put.chunk_id == req.put.total_chunks - 1
+    if req_type == "task":
+        return req.task.chunk_id == req.task.total_chunks - 1
+    return req_type not in ("acknowledge", "connection_cleanup")
+
+
+def fill_queue(
+    grpc_input_generator: Iterator[ray_client_pb2.DataRequest],
+    output_queue: "Queue[Union[ray_client_pb2.DataRequest, ray_client_pb2.DataResponse]]",  # noqa: E501
+) -> None:
+    """
+    Pushes incoming requests to a shared output_queue.
+    """
+    try:
+        for req in grpc_input_generator:
+            output_queue.put(req)
+    except grpc.RpcError as e:
+        logger.debug(
+            "closing dataservicer reader thread "
+            f"grpc error reading request_iterator: {e}"
+        )
+    finally:
+        # Set the sentinel value for the output_queue
+        output_queue.put(None)
+
+
+class ChunkCollector:
+    """
+    Helper class for collecting chunks from PutObject or ClientTask messages
+    """
+
+    def __init__(self):
+        self.curr_req_id = None
+        self.last_seen_chunk_id = -1
+        self.data = bytearray()
+
+    def add_chunk(
+        self,
+        req: ray_client_pb2.DataRequest,
+        chunk: Union[ray_client_pb2.PutRequest, ray_client_pb2.ClientTask],
+    ):
+        if self.curr_req_id is not None and self.curr_req_id != req.req_id:
+            raise RuntimeError(
+                "Expected to receive a chunk from request with id "
+                f"{self.curr_req_id}, but found {req.req_id} instead."
+            )
+        self.curr_req_id = req.req_id
+        next_chunk = self.last_seen_chunk_id + 1
+        if chunk.chunk_id < next_chunk:
+            # Repeated chunk, ignore
+            return
+        if chunk.chunk_id > next_chunk:
+            raise RuntimeError(
+                f"A chunk {chunk.chunk_id} of request {req.req_id} was "
+                "received out of order."
+            )
+        elif chunk.chunk_id == self.last_seen_chunk_id + 1:
+            self.data.extend(chunk.data)
+            self.last_seen_chunk_id = chunk.chunk_id
+        return chunk.chunk_id + 1 == chunk.total_chunks
+
+    def reset(self):
+        self.curr_req_id = None
+        self.last_seen_chunk_id = -1
+        self.data = bytearray()
+
+
+class DataServicer(ray_client_pb2_grpc.RayletDataStreamerServicer):
+    def __init__(self, basic_service: "RayletServicer"):
+        self.basic_service = basic_service
+        self.clients_lock = Lock()
+        self.num_clients = 0  # guarded by self.clients_lock
+        # dictionary mapping client_id's to the last time they connected
+        self.client_last_seen: Dict[str, float] = {}
+        # dictionary mapping client_id's to their reconnect grace periods
+        self.reconnect_grace_periods: Dict[str, float] = {}
+        # dictionary mapping client_id's to their response cache
+        self.response_caches: Dict[str, OrderedResponseCache] = defaultdict(
+            OrderedResponseCache
+        )
+        # stopped event, useful for signals that the server is shut down
+        self.stopped = Event()
+        # Helper for collecting chunks from PutObject calls. Assumes that
+        # that put requests from different objects aren't interleaved.
+        self.put_request_chunk_collector = ChunkCollector()
+        # Helper for collecting chunks from ClientTask calls. Assumes that
+        # schedule requests from different remote calls aren't interleaved.
+        self.client_task_chunk_collector = ChunkCollector()
+
+    def Datapath(self, request_iterator, context):
+        start_time = time.time()
+        # set to True if client shuts down gracefully
+        cleanup_requested = False
+        metadata = dict(context.invocation_metadata())
+        client_id = metadata.get("client_id")
+        if client_id is None:
+            logger.error("Client connecting with no client_id")
+            return
+        logger.debug(f"New data connection from client {client_id}: ")
+        accepted_connection = self._init(client_id, context, start_time)
+        response_cache = self.response_caches[client_id]
+        # Set to False if client requests a reconnect grace period of 0
+        reconnect_enabled = True
+        if not accepted_connection:
+            return
+        try:
+            request_queue = Queue()
+            queue_filler_thread = Thread(
+                target=fill_queue, daemon=True, args=(request_iterator, request_queue)
+            )
+            queue_filler_thread.start()
+            """For non `async get` requests, this loop yields immediately
+            For `async get` requests, this loop:
+                 1) does not yield, it just continues
+                 2) When the result is ready, it yields
+            """
+            for req in iter(request_queue.get, None):
+                if isinstance(req, ray_client_pb2.DataResponse):
+                    # Early shortcut if this is the result of an async get.
+                    yield req
+                    continue
+
+                assert isinstance(req, ray_client_pb2.DataRequest)
+                if _should_cache(req) and reconnect_enabled:
+                    cached_resp = response_cache.check_cache(req.req_id)
+                    if isinstance(cached_resp, Exception):
+                        # Cache state is invalid, raise exception
+                        raise cached_resp
+                    if cached_resp is not None:
+                        yield cached_resp
+                        continue
+
+                resp = None
+                req_type = req.WhichOneof("type")
+                if req_type == "init":
+                    resp_init = self.basic_service.Init(req.init)
+                    resp = ray_client_pb2.DataResponse(
+                        init=resp_init,
+                    )
+                    with self.clients_lock:
+                        self.reconnect_grace_periods[
+                            client_id
+                        ] = req.init.reconnect_grace_period
+                        if req.init.reconnect_grace_period == 0:
+                            reconnect_enabled = False
+
+                elif req_type == "get":
+                    if req.get.asynchronous:
+                        get_resp = self.basic_service._async_get_object(
+                            req.get, client_id, req.req_id, request_queue
+                        )
+                        if get_resp is None:
+                            # Skip sending a response for this request and
+                            # continue to the next requst. The response for
+                            # this request will be sent when the object is
+                            # ready.
+                            continue
+                    else:
+                        get_resp = self.basic_service._get_object(req.get, client_id)
+                    resp = ray_client_pb2.DataResponse(get=get_resp)
+                elif req_type == "put":
+                    if not self.put_request_chunk_collector.add_chunk(req, req.put):
+                        # Put request still in progress
+                        continue
+                    put_resp = self.basic_service._put_object(
+                        self.put_request_chunk_collector.data,
+                        req.put.client_ref_id,
+                        client_id,
+                    )
+                    self.put_request_chunk_collector.reset()
+                    resp = ray_client_pb2.DataResponse(put=put_resp)
+                elif req_type == "release":
+                    released = []
+                    for rel_id in req.release.ids:
+                        rel = self.basic_service.release(client_id, rel_id)
+                        released.append(rel)
+                    resp = ray_client_pb2.DataResponse(
+                        release=ray_client_pb2.ReleaseResponse(ok=released)
+                    )
+                elif req_type == "connection_info":
+                    resp = ray_client_pb2.DataResponse(
+                        connection_info=self._build_connection_response()
+                    )
+                elif req_type == "prep_runtime_env":
+                    with self.clients_lock:
+                        resp_prep = self.basic_service.PrepRuntimeEnv(
+                            req.prep_runtime_env
+                        )
+                        resp = ray_client_pb2.DataResponse(prep_runtime_env=resp_prep)
+                elif req_type == "connection_cleanup":
+                    cleanup_requested = True
+                    cleanup_resp = ray_client_pb2.ConnectionCleanupResponse()
+                    resp = ray_client_pb2.DataResponse(connection_cleanup=cleanup_resp)
+                elif req_type == "acknowledge":
+                    # Clean up acknowledged cache entries
+                    response_cache.cleanup(req.acknowledge.req_id)
+                    continue
+                elif req_type == "task":
+                    with self.clients_lock:
+                        task = req.task
+                        if not self.client_task_chunk_collector.add_chunk(req, task):
+                            # Not all serialized arguments have arrived
+                            continue
+                        arglist, kwargs = loads_from_client(
+                            self.client_task_chunk_collector.data, self.basic_service
+                        )
+                        self.client_task_chunk_collector.reset()
+                        resp_ticket = self.basic_service.Schedule(
+                            req.task, arglist, kwargs, context
+                        )
+                        resp = ray_client_pb2.DataResponse(task_ticket=resp_ticket)
+                        del arglist
+                        del kwargs
+                elif req_type == "terminate":
+                    with self.clients_lock:
+                        response = self.basic_service.Terminate(req.terminate, context)
+                        resp = ray_client_pb2.DataResponse(terminate=response)
+                elif req_type == "list_named_actors":
+                    with self.clients_lock:
+                        response = self.basic_service.ListNamedActors(
+                            req.list_named_actors
+                        )
+                        resp = ray_client_pb2.DataResponse(list_named_actors=response)
+                else:
+                    raise Exception(
+                        f"Unreachable code: Request type "
+                        f"{req_type} not handled in Datapath"
+                    )
+                resp.req_id = req.req_id
+                if _should_cache(req) and reconnect_enabled:
+                    response_cache.update_cache(req.req_id, resp)
+                yield resp
+        except Exception as e:
+            logger.exception("Error in data channel:")
+            recoverable = _propagate_error_in_context(e, context)
+            invalid_cache = response_cache.invalidate(e)
+            if not recoverable or invalid_cache:
+                context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+                # Connection isn't recoverable, skip cleanup
+                cleanup_requested = True
+        finally:
+            logger.debug(f"Stream is broken with client {client_id}")
+            queue_filler_thread.join(QUEUE_JOIN_SECONDS)
+            if queue_filler_thread.is_alive():
+                logger.error(
+                    "Queue filler thread failed to join before timeout: {}".format(
+                        QUEUE_JOIN_SECONDS
+                    )
+                )
+            cleanup_delay = self.reconnect_grace_periods.get(client_id)
+            if not cleanup_requested and cleanup_delay is not None:
+                logger.debug(
+                    "Cleanup wasn't requested, delaying cleanup by"
+                    f"{cleanup_delay} seconds."
+                )
+                # Delay cleanup, since client may attempt a reconnect
+                # Wait on the "stopped" event in case the grpc server is
+                # stopped and we can clean up earlier.
+                self.stopped.wait(timeout=cleanup_delay)
+            else:
+                logger.debug("Cleanup was requested, cleaning up immediately.")
+            with self.clients_lock:
+                if client_id not in self.client_last_seen:
+                    logger.debug("Connection already cleaned up.")
+                    # Some other connection has already cleaned up this
+                    # this client's session. This can happen if the client
+                    # reconnects and then gracefully shut's down immediately.
+                    return
+                last_seen = self.client_last_seen[client_id]
+                if last_seen > start_time:
+                    # The client successfully reconnected and updated
+                    # last seen some time during the grace period
+                    logger.debug("Client reconnected, skipping cleanup")
+                    return
+                # Either the client shut down gracefully, or the client
+                # failed to reconnect within the grace period. Clean up
+                # the connection.
+                self.basic_service.release_all(client_id)
+                del self.client_last_seen[client_id]
+                if client_id in self.reconnect_grace_periods:
+                    del self.reconnect_grace_periods[client_id]
+                if client_id in self.response_caches:
+                    del self.response_caches[client_id]
+                self.num_clients -= 1
+                logger.debug(
+                    f"Removed client {client_id}, " f"remaining={self.num_clients}"
+                )
+
+                # It's important to keep the Ray shutdown
+                # within this locked context or else Ray could hang.
+                # NOTE: it is strange to start ray in server.py but shut it
+                # down here. Consider consolidating ray lifetime management.
+                with disable_client_hook():
+                    if self.num_clients == 0:
+                        logger.debug("Shutting down ray.")
+                        ray.shutdown()
+
+    def _init(self, client_id: str, context: Any, start_time: float):
+        """
+        Checks if resources allow for another client.
+        Returns a boolean indicating if initialization was successful.
+        """
+        with self.clients_lock:
+            reconnecting = _get_reconnecting_from_context(context)
+            threshold = int(CLIENT_SERVER_MAX_THREADS / 2)
+            if self.num_clients >= threshold:
+                logger.warning(
+                    f"[Data Servicer]: Num clients {self.num_clients} "
+                    f"has reached the threshold {threshold}. "
+                    f"Rejecting client: {client_id}. "
+                )
+                if log_once("client_threshold"):
+                    logger.warning(
+                        "You can configure the client connection "
+                        "threshold by setting the "
+                        "RAY_CLIENT_SERVER_MAX_THREADS env var "
+                        f"(currently set to {CLIENT_SERVER_MAX_THREADS})."
+                    )
+                context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED)
+                return False
+            if reconnecting and client_id not in self.client_last_seen:
+                # Client took too long to reconnect, session has been
+                # cleaned up.
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+                context.set_details(
+                    "Attempted to reconnect to a session that has already "
+                    "been cleaned up."
+                )
+                return False
+            if client_id in self.client_last_seen:
+                logger.debug(f"Client {client_id} has reconnected.")
+            else:
+                self.num_clients += 1
+                logger.debug(
+                    f"Accepted data connection from {client_id}. "
+                    f"Total clients: {self.num_clients}"
+                )
+            self.client_last_seen[client_id] = start_time
+            return True
+
+    def _build_connection_response(self):
+        with self.clients_lock:
+            cur_num_clients = self.num_clients
+        return ray_client_pb2.ConnectionInfoResponse(
+            num_clients=cur_num_clients,
+            python_version="{}.{}.{}".format(
+                sys.version_info[0], sys.version_info[1], sys.version_info[2]
+            ),
+            ray_version=ray.__version__,
+            ray_commit=ray.__commit__,
+        )

python/ray/util/client/server/logservicer.py

@@ -0,0 +1,125 @@
+"""This file responds to log stream requests and forwards logs
+with its handler.
+"""
+import io
+import logging
+import queue
+import threading
+import uuid
+
+import grpc
+
+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._private.ray_logging import global_worker_stdstream_dispatcher
+from ray._private.worker import print_worker_logs
+from ray.util.client.common import CLIENT_SERVER_MAX_THREADS
+
+logger = logging.getLogger(__name__)
+
+
+class LogstreamHandler(logging.Handler):
+    def __init__(self, queue, level):
+        super().__init__()
+        self.queue = queue
+        self.level = level
+
+    def emit(self, record: logging.LogRecord):
+        logdata = ray_client_pb2.LogData()
+        logdata.msg = record.getMessage()
+        logdata.level = record.levelno
+        logdata.name = record.name
+        self.queue.put(logdata)
+
+
+class StdStreamHandler:
+    def __init__(self, queue):
+        self.queue = queue
+        self.id = str(uuid.uuid4())
+
+    def handle(self, data):
+        logdata = ray_client_pb2.LogData()
+        logdata.level = -2 if data["is_err"] else -1
+        logdata.name = "stderr" if data["is_err"] else "stdout"
+        with io.StringIO() as file:
+            print_worker_logs(data, file)
+            logdata.msg = file.getvalue()
+        self.queue.put(logdata)
+
+    def register_global(self):
+        global_worker_stdstream_dispatcher.add_handler(self.id, self.handle)
+
+    def unregister_global(self):
+        global_worker_stdstream_dispatcher.remove_handler(self.id)
+
+
+def log_status_change_thread(log_queue, request_iterator):
+    std_handler = StdStreamHandler(log_queue)
+    current_handler = None
+    root_logger = logging.getLogger("ray")
+    default_level = root_logger.getEffectiveLevel()
+    try:
+        for req in request_iterator:
+            if current_handler is not None:
+                root_logger.setLevel(default_level)
+                root_logger.removeHandler(current_handler)
+                std_handler.unregister_global()
+            if not req.enabled:
+                current_handler = None
+                continue
+            current_handler = LogstreamHandler(log_queue, req.loglevel)
+            std_handler.register_global()
+            root_logger.addHandler(current_handler)
+            root_logger.setLevel(req.loglevel)
+    except grpc.RpcError as e:
+        logger.debug(f"closing log thread " f"grpc error reading request_iterator: {e}")
+    finally:
+        if current_handler is not None:
+            root_logger.setLevel(default_level)
+            root_logger.removeHandler(current_handler)
+            std_handler.unregister_global()
+        log_queue.put(None)
+
+
+class LogstreamServicer(ray_client_pb2_grpc.RayletLogStreamerServicer):
+    def __init__(self):
+        super().__init__()
+        self.num_clients = 0
+        self.client_lock = threading.Lock()
+
+    def Logstream(self, request_iterator, context):
+        initialized = False
+        with self.client_lock:
+            threshold = CLIENT_SERVER_MAX_THREADS / 2
+            if self.num_clients + 1 >= threshold:
+                context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED)
+                logger.warning(
+                    f"Logstream: Num clients {self.num_clients} has reached "
+                    f"the threshold {threshold}. Rejecting new connection."
+                )
+                return
+            self.num_clients += 1
+            initialized = True
+            logger.info(
+                "New logs connection established. " f"Total clients: {self.num_clients}"
+            )
+        log_queue = queue.Queue()
+        thread = threading.Thread(
+            target=log_status_change_thread,
+            args=(log_queue, request_iterator),
+            daemon=True,
+        )
+        thread.start()
+        try:
+            queue_iter = iter(log_queue.get, None)
+            for record in queue_iter:
+                if record is None:
+                    break
+                yield record
+        except grpc.RpcError as e:
+            logger.debug(f"Closing log channel: {e}")
+        finally:
+            thread.join()
+            with self.client_lock:
+                if initialized:
+                    self.num_clients -= 1

python/ray/util/client/server/proxier.py

@@ -0,0 +1,938 @@
+import atexit
+import json
+import logging
+import socket
+import sys
+import time
+import traceback
+import urllib
+from concurrent import futures
+from dataclasses import dataclass
+from itertools import chain
+from threading import Event, Lock, RLock, Thread
+from typing import Callable, Dict, List, Optional, Tuple
+from urllib.parse import urlparse, urlunparse
+
+import grpc
+
+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
+import ray.core.generated.runtime_env_agent_pb2 as runtime_env_agent_pb2
+from ray._common.network_utils import (
+    build_address,
+    get_localhost_ip,
+    is_ipv6,
+    is_localhost,
+)
+from ray._common.tls_utils import add_port_to_grpc_server
+from ray._common.utils import env_integer
+from ray._private.authentication.http_token_authentication import (
+    format_authentication_http_error,
+    get_auth_headers_if_auth_enabled,
+)
+from ray._private.client_mode_hook import disable_client_hook
+from ray._private.grpc_utils import init_grpc_channel
+from ray._private.parameter import RayParams
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.services import (
+    ProcessInfo,
+    get_node_with_retry,
+    start_ray_client_server,
+)
+from ray._private.utils import detect_fate_sharing_support
+from ray._raylet import GcsClient
+from ray.cloudpickle.compat import pickle
+from ray.exceptions import AuthenticationError
+from ray.job_config import JobConfig
+from ray.util.client.common import (
+    CLIENT_SERVER_MAX_THREADS,
+    GRPC_OPTIONS,
+    ClientServerHandle,
+    _get_client_id_from_context,
+    _propagate_error_in_context,
+)
+from ray.util.client.server.dataservicer import _get_reconnecting_from_context
+
+# Import psutil after ray so the packaged version is used.
+import psutil
+
+logger = logging.getLogger(__name__)
+
+CHECK_PROCESS_INTERVAL_S = 30
+
+MIN_SPECIFIC_SERVER_PORT = 23000
+MAX_SPECIFIC_SERVER_PORT = 24000
+
+CHECK_CHANNEL_TIMEOUT_S = env_integer("RAY_CLIENT_SERVER_CHECK_CHANNEL_TIMEOUT_S", 30)
+
+LOGSTREAM_RETRIES = 5
+LOGSTREAM_RETRY_INTERVAL_SEC = 2
+
+
+@dataclass
+class SpecificServer:
+    port: int
+    process_handle_future: futures.Future
+    channel: "grpc._channel.Channel"
+
+    def is_ready(self) -> bool:
+        """Check if the server is ready or not (doesn't block)."""
+        return self.process_handle_future.done()
+
+    def wait_ready(self, timeout: Optional[float] = None) -> None:
+        """
+        Wait for the server to actually start up.
+        """
+        res = self.process_handle_future.result(timeout=timeout)
+        if res is None:
+            # This is only set to none when server creation specifically fails.
+            raise RuntimeError("Server startup failed.")
+
+    def poll(self) -> Optional[int]:
+        """Check if the process has exited."""
+        try:
+            proc = self.process_handle_future.result(timeout=0.1)
+            if proc is not None:
+                return proc.process.poll()
+        except futures.TimeoutError:
+            return
+
+    def kill(self) -> None:
+        """Try to send a KILL signal to the process."""
+        try:
+            proc = self.process_handle_future.result(timeout=0.1)
+            if proc is not None:
+                proc.process.kill()
+        except futures.TimeoutError:
+            # Server has not been started yet.
+            pass
+
+    def set_result(self, proc: Optional[ProcessInfo]) -> None:
+        """Set the result of the internal future if it is currently unset."""
+        if not self.is_ready():
+            self.process_handle_future.set_result(proc)
+
+
+def _match_running_client_server(command: List[str]) -> bool:
+    """
+    Detects if the main process in the given command is the RayClient Server.
+    This works by ensuring that the command is of the form:
+        <py_executable> -m ray.util.client.server <args>
+    """
+    flattened = " ".join(command)
+    return "-m ray.util.client.server" in flattened
+
+
+class ProxyManager:
+    def __init__(
+        self,
+        address: Optional[str],
+        runtime_env_agent_address: str,
+        *,
+        session_dir: Optional[str] = None,
+        redis_username: Optional[str] = None,
+        redis_password: Optional[str] = None,
+        node_id: Optional[str] = None,
+    ):
+        self.servers: Dict[str, SpecificServer] = dict()
+        self.server_lock = RLock()
+        self._address = address
+        self._redis_username = redis_username
+        self._redis_password = redis_password
+        self._free_ports: List[int] = list(
+            range(MIN_SPECIFIC_SERVER_PORT, MAX_SPECIFIC_SERVER_PORT)
+        )
+
+        if runtime_env_agent_address:
+            parsed = urlparse(runtime_env_agent_address)
+            # runtime env agent self-assigns a free port, fetch it from GCS
+            if parsed.port is None or parsed.port == 0:
+                if node_id is None:
+                    raise ValueError(
+                        "node_id is required when runtime_env_agent_address "
+                        "has no port specified"
+                    )
+                node_info = get_node_with_retry(address, node_id)
+                runtime_env_agent_address = urlunparse(
+                    parsed._replace(
+                        netloc=f"{parsed.hostname}:{node_info['runtime_env_agent_port']}"
+                    )
+                )
+
+        self._runtime_env_agent_address = runtime_env_agent_address
+
+        self._check_thread = Thread(target=self._check_processes, daemon=True)
+        self._check_thread.start()
+
+        self.fate_share = bool(detect_fate_sharing_support())
+        self._node: Optional[ray._private.node.Node] = None
+        atexit.register(self._cleanup)
+
+    def _get_unused_port(self, family: int = socket.AF_INET) -> int:
+        """
+        Search for a port in _free_ports that is unused.
+        """
+        with self.server_lock:
+            num_ports = len(self._free_ports)
+            for _ in range(num_ports):
+                port = self._free_ports.pop(0)
+                s = socket.socket(family, socket.SOCK_STREAM)
+                try:
+                    s.bind(("", port))
+                except OSError:
+                    self._free_ports.append(port)
+                    continue
+                finally:
+                    s.close()
+                return port
+        raise RuntimeError("Unable to succeed in selecting a random port.")
+
+    @property
+    def address(self) -> str:
+        """
+        Returns the provided Ray bootstrap address, or creates a new cluster.
+        """
+        if self._address:
+            return self._address
+        # Start a new, locally scoped cluster.
+        connection_tuple = ray.init()
+        self._address = connection_tuple["address"]
+        self._session_dir = connection_tuple["session_dir"]
+        return self._address
+
+    @property
+    def node(self) -> ray._private.node.Node:
+        """Gets a 'ray.Node' object for this node (the head node).
+        If it does not already exist, one is created using the bootstrap
+        address.
+        """
+        if self._node:
+            return self._node
+        ray_params = RayParams(gcs_address=self.address)
+
+        self._node = ray._private.node.Node(
+            ray_params,
+            head=False,
+            shutdown_at_exit=False,
+            spawn_reaper=False,
+            connect_only=True,
+        )
+
+        return self._node
+
+    def create_specific_server(self, client_id: str) -> SpecificServer:
+        """
+        Create, but not start a SpecificServer for a given client. This
+        method must be called once per client.
+        """
+        with self.server_lock:
+            assert (
+                self.servers.get(client_id) is None
+            ), f"Server already created for Client: {client_id}"
+
+            host = get_localhost_ip()
+            port = self._get_unused_port(
+                socket.AF_INET6 if is_ipv6(host) else socket.AF_INET
+            )
+
+            server = SpecificServer(
+                port=port,
+                process_handle_future=futures.Future(),
+                channel=init_grpc_channel(
+                    build_address(host, port), options=GRPC_OPTIONS
+                ),
+            )
+            self.servers[client_id] = server
+            return server
+
+    def _create_runtime_env(
+        self,
+        serialized_runtime_env: str,
+        runtime_env_config: str,
+        specific_server: SpecificServer,
+    ):
+        """Increase the runtime_env reference by sending an RPC to the agent.
+
+        Includes retry logic to handle the case when the agent is
+        temporarily unreachable (e.g., hasn't been started up yet).
+        """
+        logger.info(
+            f"Increasing runtime env reference for "
+            f"ray_client_server_{specific_server.port}."
+            f"Serialized runtime env is {serialized_runtime_env}."
+        )
+
+        assert (
+            len(self._runtime_env_agent_address) > 0
+        ), "runtime_env_agent_address not set"
+
+        create_env_request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest(
+            serialized_runtime_env=serialized_runtime_env,
+            runtime_env_config=runtime_env_config,
+            job_id=f"ray_client_server_{specific_server.port}".encode("utf-8"),
+            source_process="client_server",
+        )
+
+        retries = 0
+        max_retries = 5
+        wait_time_s = 0.5
+        last_exception = None
+        while retries <= max_retries:
+            try:
+                url = urllib.parse.urljoin(
+                    self._runtime_env_agent_address, "/get_or_create_runtime_env"
+                )
+                data = create_env_request.SerializeToString()
+                headers = {"Content-Type": "application/octet-stream"}
+                headers.update(**get_auth_headers_if_auth_enabled(headers))
+                req = urllib.request.Request(
+                    url, data=data, method="POST", headers=headers
+                )
+                response = urllib.request.urlopen(req, timeout=None)
+                response_data = response.read()
+                r = runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply()
+                r.ParseFromString(response_data)
+
+                if r.status == runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_OK:
+                    return r.serialized_runtime_env_context
+                elif (
+                    r.status
+                    == runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_FAILED
+                ):
+                    raise RuntimeError(
+                        "Failed to create runtime_env for Ray client "
+                        f"server, it is caused by:\n{r.error_message}"
+                    )
+                else:
+                    assert False, f"Unknown status: {r.status}."
+            except urllib.error.HTTPError as e:
+                body = ""
+                try:
+                    body = e.read().decode("utf-8", "ignore")
+                except Exception:
+                    body = e.reason if hasattr(e, "reason") else str(e)
+
+                formatted_error = format_authentication_http_error(e.code, body or "")
+                if formatted_error:
+                    raise AuthenticationError(formatted_error) from e
+
+                # Treat non-auth HTTP errors like URLError (retry with backoff)
+                last_exception = e
+                logger.warning(
+                    f"GetOrCreateRuntimeEnv request failed with HTTP {e.code}: {body or e}. "
+                    f"Retrying after {wait_time_s}s. "
+                    f"{max_retries-retries} retries remaining."
+                )
+
+            except urllib.error.URLError as e:
+                last_exception = e
+                logger.warning(
+                    f"GetOrCreateRuntimeEnv request failed: {e}. "
+                    f"Retrying after {wait_time_s}s. "
+                    f"{max_retries-retries} retries remaining."
+                )
+
+            # Exponential backoff.
+            time.sleep(wait_time_s)
+            retries += 1
+            wait_time_s *= 2
+
+        raise TimeoutError(
+            f"GetOrCreateRuntimeEnv request failed after {max_retries} attempts."
+            f" Last exception: {last_exception}"
+        )
+
+    def start_specific_server(self, client_id: str, job_config: JobConfig) -> bool:
+        """
+        Start up a RayClient Server for an incoming client to
+        communicate with. Returns whether creation was successful.
+        """
+        specific_server = self._get_server_for_client(client_id)
+        assert specific_server, f"Server has not been created for: {client_id}"
+
+        output, error = self.node.get_log_file_handles(
+            f"ray_client_server_{specific_server.port}", unique=True
+        )
+
+        serialized_runtime_env = job_config._get_serialized_runtime_env()
+        runtime_env_config = job_config._get_proto_runtime_env_config()
+        if not serialized_runtime_env or serialized_runtime_env == "{}":
+            # TODO(edoakes): can we just remove this case and always send it
+            # to the agent?
+            serialized_runtime_env_context = RuntimeEnvContext().serialize()
+        else:
+            serialized_runtime_env_context = self._create_runtime_env(
+                serialized_runtime_env=serialized_runtime_env,
+                runtime_env_config=runtime_env_config,
+                specific_server=specific_server,
+            )
+
+        proc = start_ray_client_server(
+            self.address,
+            get_localhost_ip(),
+            specific_server.port,
+            stdout_file=output,
+            stderr_file=error,
+            fate_share=self.fate_share,
+            server_type="specific-server",
+            serialized_runtime_env_context=serialized_runtime_env_context,
+            redis_username=self._redis_username,
+            redis_password=self._redis_password,
+        )
+
+        # Wait for the process being run transitions from the shim process
+        # to the actual RayClient Server.
+        pid = proc.process.pid
+        if sys.platform != "win32":
+            psutil_proc = psutil.Process(pid)
+        else:
+            psutil_proc = None
+        # Don't use `psutil` on Win32
+        while psutil_proc is not None:
+            if proc.process.poll() is not None:
+                logger.error(f"SpecificServer startup failed for client: {client_id}")
+                break
+            cmd = psutil_proc.cmdline()
+            if _match_running_client_server(cmd):
+                break
+            logger.debug("Waiting for Process to reach the actual client server.")
+            time.sleep(0.5)
+        specific_server.set_result(proc)
+        logger.info(
+            f"SpecificServer started on port: {specific_server.port} "
+            f"with PID: {pid} for client: {client_id}"
+        )
+        return proc.process.poll() is None
+
+    def _get_server_for_client(self, client_id: str) -> Optional[SpecificServer]:
+        with self.server_lock:
+            client = self.servers.get(client_id)
+            if client is None:
+                logger.error(f"Unable to find channel for client: {client_id}")
+            return client
+
+    def has_channel(self, client_id: str) -> bool:
+        server = self._get_server_for_client(client_id)
+        if server is None:
+            return False
+
+        return server.is_ready()
+
+    def get_channel(
+        self,
+        client_id: str,
+    ) -> Optional["grpc._channel.Channel"]:
+        """
+        Find the gRPC Channel for the given client_id. This will block until
+        the server process has started.
+        """
+        server = self._get_server_for_client(client_id)
+        if server is None:
+            return None
+        # Wait for the SpecificServer to become ready.
+        server.wait_ready()
+        try:
+            grpc.channel_ready_future(server.channel).result(
+                timeout=CHECK_CHANNEL_TIMEOUT_S
+            )
+            return server.channel
+        except grpc.FutureTimeoutError:
+            logger.exception(f"Timeout waiting for channel for {client_id}")
+            return None
+
+    def _check_processes(self):
+        """
+        Keeps the internal servers dictionary up-to-date with running servers.
+        """
+        while True:
+            with self.server_lock:
+                for client_id, specific_server in list(self.servers.items()):
+                    if specific_server.poll() is not None:
+                        logger.info(
+                            f"Specific server {client_id} is no longer running"
+                            f", freeing its port {specific_server.port}"
+                        )
+                        del self.servers[client_id]
+                        # Port is available to use again.
+                        self._free_ports.append(specific_server.port)
+
+            time.sleep(CHECK_PROCESS_INTERVAL_S)
+
+    def _cleanup(self) -> None:
+        """
+        Forcibly kill all spawned RayClient Servers. This ensures cleanup
+        for platforms where fate sharing is not supported.
+        """
+        for server in self.servers.values():
+            server.kill()
+
+
+class RayletServicerProxy(ray_client_pb2_grpc.RayletDriverServicer):
+    def __init__(self, ray_connect_handler: Callable, proxy_manager: ProxyManager):
+        self.proxy_manager = proxy_manager
+        self.ray_connect_handler = ray_connect_handler
+
+    def _call_inner_function(
+        self, request, context, method: str
+    ) -> Optional[ray_client_pb2_grpc.RayletDriverStub]:
+        client_id = _get_client_id_from_context(context)
+        chan = self.proxy_manager.get_channel(client_id)
+        if not chan:
+            logger.error(f"Channel for Client: {client_id} not found!")
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return None
+
+        stub = ray_client_pb2_grpc.RayletDriverStub(chan)
+        try:
+            metadata = [("client_id", client_id)]
+            if context:
+                metadata = context.invocation_metadata()
+            return getattr(stub, method)(request, metadata=metadata)
+        except Exception as e:
+            # Error while proxying -- propagate the error's context to user
+            logger.exception(f"Proxying call to {method} failed!")
+            _propagate_error_in_context(e, context)
+
+    def _has_channel_for_request(self, context):
+        client_id = _get_client_id_from_context(context)
+        return self.proxy_manager.has_channel(client_id)
+
+    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:
+        """Proxies internal_kv.put.
+
+        This is used by the working_dir code to upload to the GCS before
+        ray.init is called. In that case (if we don't have a server yet)
+        we directly make the internal KV call from the proxier.
+
+        Otherwise, we proxy the call to the downstream server as usual.
+        """
+        if self._has_channel_for_request(context):
+            return self._call_inner_function(request, context, "KVPut")
+
+        with disable_client_hook():
+            already_exists = ray.experimental.internal_kv._internal_kv_put(
+                request.key, request.value, overwrite=request.overwrite
+            )
+        return ray_client_pb2.KVPutResponse(already_exists=already_exists)
+
+    def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse:
+        """Proxies internal_kv.get.
+
+        This is used by the working_dir code to upload to the GCS before
+        ray.init is called. In that case (if we don't have a server yet)
+        we directly make the internal KV call from the proxier.
+
+        Otherwise, we proxy the call to the downstream server as usual.
+        """
+        if self._has_channel_for_request(context):
+            return self._call_inner_function(request, context, "KVGet")
+
+        with disable_client_hook():
+            value = ray.experimental.internal_kv._internal_kv_get(request.key)
+        return ray_client_pb2.KVGetResponse(value=value)
+
+    def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse:
+        """Proxies internal_kv.delete.
+
+        This is used by the working_dir code to upload to the GCS before
+        ray.init is called. In that case (if we don't have a server yet)
+        we directly make the internal KV call from the proxier.
+
+        Otherwise, we proxy the call to the downstream server as usual.
+        """
+        if self._has_channel_for_request(context):
+            return self._call_inner_function(request, context, "KVDel")
+
+        with disable_client_hook():
+            ray.experimental.internal_kv._internal_kv_del(request.key)
+        return ray_client_pb2.KVDelResponse()
+
+    def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse:
+        """Proxies internal_kv.list.
+
+        This is used by the working_dir code to upload to the GCS before
+        ray.init is called. In that case (if we don't have a server yet)
+        we directly make the internal KV call from the proxier.
+
+        Otherwise, we proxy the call to the downstream server as usual.
+        """
+        if self._has_channel_for_request(context):
+            return self._call_inner_function(request, context, "KVList")
+
+        with disable_client_hook():
+            keys = ray.experimental.internal_kv._internal_kv_list(request.prefix)
+        return ray_client_pb2.KVListResponse(keys=keys)
+
+    def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse:
+        """Proxies internal_kv.exists.
+
+        This is used by the working_dir code to upload to the GCS before
+        ray.init is called. In that case (if we don't have a server yet)
+        we directly make the internal KV call from the proxier.
+
+        Otherwise, we proxy the call to the downstream server as usual.
+        """
+        if self._has_channel_for_request(context):
+            return self._call_inner_function(request, context, "KVExists")
+
+        with disable_client_hook():
+            exists = ray.experimental.internal_kv._internal_kv_exists(request.key)
+        return ray_client_pb2.KVExistsResponse(exists=exists)
+
+    def PinRuntimeEnvURI(
+        self, request, context=None
+    ) -> ray_client_pb2.ClientPinRuntimeEnvURIResponse:
+        """Proxies internal_kv.pin_runtime_env_uri.
+
+        This is used by the working_dir code to upload to the GCS before
+        ray.init is called. In that case (if we don't have a server yet)
+        we directly make the internal KV call from the proxier.
+
+        Otherwise, we proxy the call to the downstream server as usual.
+        """
+        if self._has_channel_for_request(context):
+            return self._call_inner_function(request, context, "PinRuntimeEnvURI")
+
+        with disable_client_hook():
+            ray.experimental.internal_kv._pin_runtime_env_uri(
+                request.uri, expiration_s=request.expiration_s
+            )
+        return ray_client_pb2.ClientPinRuntimeEnvURIResponse()
+
+    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:
+
+        # NOTE: We need to respond to the PING request here to allow the client
+        # to continue with connecting.
+        if request.type == ray_client_pb2.ClusterInfoType.PING:
+            resp = ray_client_pb2.ClusterInfoResponse(json=json.dumps({}))
+            return resp
+        return self._call_inner_function(request, context, "ClusterInfo")
+
+    def Terminate(self, req, context=None):
+        return self._call_inner_function(req, context, "Terminate")
+
+    def GetObject(self, request, context=None):
+        try:
+            yield from self._call_inner_function(request, context, "GetObject")
+        except Exception as e:
+            # Error while iterating over response from GetObject stream
+            logger.exception("Proxying call to GetObject failed!")
+            _propagate_error_in_context(e, context)
+
+    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, context=None) -> ray_client_pb2.WaitResponse:
+        return self._call_inner_function(request, context, "WaitObject")
+
+    def Schedule(self, task, context=None) -> ray_client_pb2.ClientTaskTicket:
+        return self._call_inner_function(task, context, "Schedule")
+
+
+def ray_client_server_env_prep(job_config: JobConfig) -> JobConfig:
+    return job_config
+
+
+def prepare_runtime_init_req(
+    init_request: ray_client_pb2.DataRequest,
+) -> Tuple[ray_client_pb2.DataRequest, JobConfig]:
+    """
+    Extract JobConfig and possibly mutate InitRequest before it is passed to
+    the specific RayClient Server.
+    """
+    init_type = init_request.WhichOneof("type")
+    assert init_type == "init", (
+        "Received initial message of type " f"{init_type}, not 'init'."
+    )
+    req = init_request.init
+    job_config = JobConfig()
+    if req.job_config:
+        job_config = pickle.loads(req.job_config)
+    new_job_config = ray_client_server_env_prep(job_config)
+    modified_init_req = ray_client_pb2.InitRequest(
+        job_config=pickle.dumps(new_job_config),
+        ray_init_kwargs=init_request.init.ray_init_kwargs,
+        reconnect_grace_period=init_request.init.reconnect_grace_period,
+    )
+
+    init_request.init.CopyFrom(modified_init_req)
+    return (init_request, new_job_config)
+
+
+class RequestIteratorProxy:
+    def __init__(self, request_iterator):
+        self.request_iterator = request_iterator
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            return next(self.request_iterator)
+        except grpc.RpcError as e:
+            # To stop proxying already CANCLLED request stream gracefully,
+            # we only translate the exact grpc.RpcError to StopIteration,
+            # not its subsclasses. ex: grpc._Rendezvous
+            # https://github.com/grpc/grpc/blob/v1.43.0/src/python/grpcio/grpc/_server.py#L353-L354
+            # This fixes the https://github.com/ray-project/ray/issues/23865
+            if type(e) is not grpc.RpcError:
+                raise e  # re-raise other grpc exceptions
+            logger.exception(
+                "Stop iterating cancelled request stream with the following exception:"
+            )
+            raise StopIteration
+
+
+class DataServicerProxy(ray_client_pb2_grpc.RayletDataStreamerServicer):
+    def __init__(self, proxy_manager: ProxyManager):
+        self.num_clients = 0
+        # dictionary mapping client_id's to the last time they connected
+        self.clients_last_seen: Dict[str, float] = {}
+        self.reconnect_grace_periods: Dict[str, float] = {}
+        self.clients_lock = Lock()
+        self.proxy_manager = proxy_manager
+        self.stopped = Event()
+
+    def modify_connection_info_resp(
+        self, init_resp: ray_client_pb2.DataResponse
+    ) -> ray_client_pb2.DataResponse:
+        """
+        Modify the `num_clients` returned the ConnectionInfoResponse because
+        individual SpecificServers only have **one** client.
+        """
+        init_type = init_resp.WhichOneof("type")
+        if init_type != "connection_info":
+            return init_resp
+        modified_resp = ray_client_pb2.DataResponse()
+        modified_resp.CopyFrom(init_resp)
+        with self.clients_lock:
+            modified_resp.connection_info.num_clients = self.num_clients
+        return modified_resp
+
+    def Datapath(self, request_iterator, context):
+        request_iterator = RequestIteratorProxy(request_iterator)
+        cleanup_requested = False
+        start_time = time.time()
+        client_id = _get_client_id_from_context(context)
+        if client_id == "":
+            return
+        reconnecting = _get_reconnecting_from_context(context)
+
+        if reconnecting:
+            with self.clients_lock:
+                if client_id not in self.clients_last_seen:
+                    # Client took too long to reconnect, session has already
+                    # been cleaned up
+                    context.set_code(grpc.StatusCode.NOT_FOUND)
+                    context.set_details(
+                        "Attempted to reconnect a session that has already "
+                        "been cleaned up"
+                    )
+                    return
+                self.clients_last_seen[client_id] = start_time
+            server = self.proxy_manager._get_server_for_client(client_id)
+            channel = self.proxy_manager.get_channel(client_id)
+            # iterator doesn't need modification on reconnect
+            new_iter = request_iterator
+        else:
+            # Create Placeholder *before* reading the first request.
+            server = self.proxy_manager.create_specific_server(client_id)
+            with self.clients_lock:
+                self.clients_last_seen[client_id] = start_time
+                self.num_clients += 1
+
+        try:
+            if not reconnecting:
+                logger.info(f"New data connection from client {client_id}: ")
+                init_req = next(request_iterator)
+                with self.clients_lock:
+                    self.reconnect_grace_periods[
+                        client_id
+                    ] = init_req.init.reconnect_grace_period
+                try:
+                    modified_init_req, job_config = prepare_runtime_init_req(init_req)
+                    if not self.proxy_manager.start_specific_server(
+                        client_id, job_config
+                    ):
+                        logger.error(
+                            f"Server startup failed for client: {client_id}, "
+                            f"using JobConfig: {job_config}!"
+                        )
+                        raise RuntimeError(
+                            "Starting Ray client server failed. See "
+                            f"ray_client_server_{server.port}.err for "
+                            "detailed logs."
+                        )
+                    channel = self.proxy_manager.get_channel(client_id)
+                    if channel is None:
+                        logger.error(f"Channel not found for {client_id}")
+                        raise RuntimeError(
+                            "Proxy failed to Connect to backend! Check "
+                            "`ray_client_server.err` and "
+                            f"`ray_client_server_{server.port}.err` on the "
+                            "head node of the cluster for the relevant logs. "
+                            "By default these are located at "
+                            "/tmp/ray/session_latest/logs."
+                        )
+                except Exception:
+                    init_resp = ray_client_pb2.DataResponse(
+                        init=ray_client_pb2.InitResponse(
+                            ok=False, msg=traceback.format_exc()
+                        )
+                    )
+                    init_resp.req_id = init_req.req_id
+                    yield init_resp
+                    return None
+
+                new_iter = chain([modified_init_req], request_iterator)
+
+            stub = ray_client_pb2_grpc.RayletDataStreamerStub(channel)
+            metadata = [("client_id", client_id), ("reconnecting", str(reconnecting))]
+            resp_stream = stub.Datapath(new_iter, metadata=metadata)
+            for resp in resp_stream:
+                resp_type = resp.WhichOneof("type")
+                if resp_type == "connection_cleanup":
+                    # Specific server is skipping cleanup, proxier should too
+                    cleanup_requested = True
+                yield self.modify_connection_info_resp(resp)
+        except Exception as e:
+            logger.exception("Proxying Datapath failed!")
+            # Propogate error through context
+            recoverable = _propagate_error_in_context(e, context)
+            if not recoverable:
+                # Client shouldn't attempt to recover, clean up connection
+                cleanup_requested = True
+        finally:
+            cleanup_delay = self.reconnect_grace_periods.get(client_id)
+            if not cleanup_requested and cleanup_delay is not None:
+                # Delay cleanup, since client may attempt a reconnect
+                # Wait on stopped event in case the server closes and we
+                # can clean up earlier
+                self.stopped.wait(timeout=cleanup_delay)
+            with self.clients_lock:
+                if client_id not in self.clients_last_seen:
+                    logger.info(f"{client_id} not found. Skipping clean up.")
+                    # Connection has already been cleaned up
+                    return
+                last_seen = self.clients_last_seen[client_id]
+                logger.info(
+                    f"{client_id} last started stream at {last_seen}. Current "
+                    f"stream started at {start_time}."
+                )
+                if last_seen > start_time:
+                    logger.info("Client reconnected. Skipping cleanup.")
+                    # Client has reconnected, don't clean up
+                    return
+                logger.debug(f"Client detached: {client_id}")
+                self.num_clients -= 1
+                del self.clients_last_seen[client_id]
+                if client_id in self.reconnect_grace_periods:
+                    del self.reconnect_grace_periods[client_id]
+                server.set_result(None)
+
+
+class LogstreamServicerProxy(ray_client_pb2_grpc.RayletLogStreamerServicer):
+    def __init__(self, proxy_manager: ProxyManager):
+        super().__init__()
+        self.proxy_manager = proxy_manager
+
+    def Logstream(self, request_iterator, context):
+        request_iterator = RequestIteratorProxy(request_iterator)
+        client_id = _get_client_id_from_context(context)
+        if client_id == "":
+            return
+        logger.debug(f"New logstream connection from client {client_id}: ")
+
+        channel = None
+        # We need to retry a few times because the LogClient *may* connect
+        # Before the DataClient has finished connecting.
+        for i in range(LOGSTREAM_RETRIES):
+            channel = self.proxy_manager.get_channel(client_id)
+
+            if channel is not None:
+                break
+            logger.warning(f"Retrying Logstream connection. {i+1} attempts failed.")
+            time.sleep(LOGSTREAM_RETRY_INTERVAL_SEC)
+
+        if channel is None:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details(
+                "Logstream proxy failed to connect. Channel for client "
+                f"{client_id} not found."
+            )
+            return None
+
+        stub = ray_client_pb2_grpc.RayletLogStreamerStub(channel)
+
+        resp_stream = stub.Logstream(
+            request_iterator, metadata=[("client_id", client_id)]
+        )
+        try:
+            for resp in resp_stream:
+                yield resp
+        except Exception:
+            logger.exception("Proxying Logstream failed!")
+
+
+def serve_proxier(
+    host: str,
+    port: int,
+    gcs_address: Optional[str],
+    *,
+    redis_username: Optional[str] = None,
+    redis_password: Optional[str] = None,
+    session_dir: Optional[str] = None,
+    runtime_env_agent_address: Optional[str] = None,
+    node_id: Optional[str] = None,
+):
+    # Initialize internal KV to be used to upload and download working_dir
+    # before calling ray.init within the RayletServicers.
+    # NOTE(edoakes): redis_address and redis_password should only be None in
+    # tests.
+    if gcs_address is not None:
+        gcs_cli = GcsClient(address=gcs_address)
+        ray.experimental.internal_kv._initialize_internal_kv(gcs_cli)
+
+    from ray._private.grpc_utils import create_grpc_server_with_interceptors
+
+    server = create_grpc_server_with_interceptors(
+        max_workers=CLIENT_SERVER_MAX_THREADS,
+        thread_name_prefix="ray_client_proxier",
+        options=GRPC_OPTIONS,
+        asynchronous=False,
+    )
+    proxy_manager = ProxyManager(
+        gcs_address,
+        session_dir=session_dir,
+        redis_username=redis_username,
+        redis_password=redis_password,
+        runtime_env_agent_address=runtime_env_agent_address,
+        node_id=node_id,
+    )
+    task_servicer = RayletServicerProxy(None, proxy_manager)
+    data_servicer = DataServicerProxy(proxy_manager)
+    logs_servicer = LogstreamServicerProxy(proxy_manager)
+    ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server)
+    ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server)
+    ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server)
+    if not is_localhost(host):
+        add_port_to_grpc_server(server, build_address(get_localhost_ip(), port))
+    add_port_to_grpc_server(server, build_address(host, port))
+    server.start()
+    return ClientServerHandle(
+        task_servicer=task_servicer,
+        data_servicer=data_servicer,
+        logs_servicer=logs_servicer,
+        grpc_server=server,
+    )

python/ray/util/client/server/server.py

@@ -0,0 +1,968 @@
+import base64
+import functools
+import gc
+import inspect
+import json
+import logging
+import math
+import os
+import pickle
+import queue
+import threading
+import time
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Optional, Set, Union
+
+import grpc
+
+import ray
+import ray._private.state
+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 import cloudpickle
+from ray._common.network_utils import (
+    build_address,
+    get_all_interfaces_ip,
+    get_localhost_ip,
+    is_localhost,
+)
+from ray._common.tls_utils import add_port_to_grpc_server
+from ray._private import ray_constants
+from ray._private.client_mode_hook import disable_client_hook
+from ray._private.ray_constants import env_integer
+from ray._private.ray_logging import setup_logger
+from ray._private.ray_logging.logging_config import LoggingConfig
+from ray._private.services import canonicalize_bootstrap_address_or_die
+from ray._raylet import GcsClient
+from ray.job_config import JobConfig
+from ray.util.client.common import (
+    CLIENT_SERVER_MAX_THREADS,
+    GRPC_OPTIONS,
+    OBJECT_TRANSFER_CHUNK_SIZE,
+    ClientServerHandle,
+    ResponseCache,
+)
+from ray.util.client.server.dataservicer import DataServicer
+from ray.util.client.server.logservicer import LogstreamServicer
+from ray.util.client.server.proxier import serve_proxier
+from ray.util.client.server.server_pickler import dumps_from_server, loads_from_client
+from ray.util.client.server.server_stubs import current_server
+
+logger = logging.getLogger(__name__)
+
+TIMEOUT_FOR_SPECIFIC_SERVER_S = env_integer("TIMEOUT_FOR_SPECIFIC_SERVER_S", 30)
+
+
+def _use_response_cache(func):
+    """
+    Decorator for gRPC stubs. Before calling the real stubs, checks if there's
+    an existing entry in the caches. If there is, then return the cached
+    entry. Otherwise, call the real function and use the real cache
+    """
+
+    @functools.wraps(func)
+    def wrapper(self, request, context):
+        metadata = dict(context.invocation_metadata())
+        expected_ids = ("client_id", "thread_id", "req_id")
+        if any(i not in metadata for i in expected_ids):
+            # Missing IDs, skip caching and call underlying stub directly
+            return func(self, request, context)
+
+        # Get relevant IDs to check cache
+        client_id = metadata["client_id"]
+        thread_id = metadata["thread_id"]
+        req_id = int(metadata["req_id"])
+
+        # Check if response already cached
+        response_cache = self.response_caches[client_id]
+        cached_entry = response_cache.check_cache(thread_id, req_id)
+        if cached_entry is not None:
+            if isinstance(cached_entry, Exception):
+                # Original call errored, propagate error
+                context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+                context.set_details(str(cached_entry))
+                raise cached_entry
+            return cached_entry
+
+        try:
+            # Response wasn't cached, call underlying stub and cache result
+            resp = func(self, request, context)
+        except Exception as e:
+            # Unexpected error in underlying stub -- update cache and
+            # propagate to user through context
+            response_cache.update_cache(thread_id, req_id, e)
+            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+            context.set_details(str(e))
+            raise
+        response_cache.update_cache(thread_id, req_id, resp)
+        return resp
+
+    return wrapper
+
+
+class RayletServicer(ray_client_pb2_grpc.RayletDriverServicer):
+    def __init__(self, ray_connect_handler: Callable):
+        """Construct a raylet service
+
+        Args:
+           ray_connect_handler: Function to connect to ray cluster
+        """
+        # Stores client_id -> (ref_id -> ObjectRef)
+        self.object_refs: Dict[str, Dict[bytes, ray.ObjectRef]] = defaultdict(dict)
+        # Stores client_id -> (client_ref_id -> ref_id (in self.object_refs))
+        self.client_side_ref_map: Dict[str, Dict[bytes, bytes]] = defaultdict(dict)
+        self.function_refs = {}
+        self.actor_refs: Dict[bytes, ray.ActorHandle] = {}
+        self.actor_owners: Dict[str, Set[bytes]] = defaultdict(set)
+        self.registered_actor_classes = {}
+        self.named_actors = set()
+        self.state_lock = threading.Lock()
+        self.ray_connect_handler = ray_connect_handler
+        self.response_caches: Dict[str, ResponseCache] = defaultdict(ResponseCache)
+
+    def Init(
+        self, request: ray_client_pb2.InitRequest, context=None
+    ) -> ray_client_pb2.InitResponse:
+        if request.job_config:
+            job_config = pickle.loads(request.job_config)
+            job_config._client_job = True
+        else:
+            job_config = None
+        current_job_config = None
+        with disable_client_hook():
+            if ray.is_initialized():
+                worker = ray._private.worker.global_worker
+                current_job_config = worker.core_worker.get_job_config()
+            else:
+                extra_kwargs = json.loads(request.ray_init_kwargs or "{}")
+                # Reconstruct LoggingConfig from dict after InitRequest ray_init_kwargs is parsed from JSON on the server.
+                if "logging_config" in extra_kwargs and isinstance(
+                    extra_kwargs["logging_config"], dict
+                ):
+                    extra_kwargs["logging_config"] = LoggingConfig.from_dict(
+                        extra_kwargs["logging_config"]
+                    )
+                try:
+                    self.ray_connect_handler(job_config, **extra_kwargs)
+                except Exception as e:
+                    logger.exception("Running Ray Init failed:")
+                    return ray_client_pb2.InitResponse(
+                        ok=False,
+                        msg=f"Call to `ray.init()` on the server failed with: {e}",
+                    )
+        if job_config is None:
+            return ray_client_pb2.InitResponse(ok=True)
+
+        # NOTE(edoakes): this code should not be necessary anymore because we
+        # only allow a single client/job per server. There is an existing test
+        # that tests the behavior of multiple clients with the same job config
+        # connecting to one server (test_client_init.py::test_num_clients),
+        # so I'm leaving it here for now.
+        job_config = job_config._get_proto_job_config()
+        # If the server has been initialized, we need to compare whether the
+        # runtime env is compatible.
+        if current_job_config:
+            job_uris = set(job_config.runtime_env_info.uris.working_dir_uri)
+            job_uris.update(job_config.runtime_env_info.uris.py_modules_uris)
+            current_job_uris = set(
+                current_job_config.runtime_env_info.uris.working_dir_uri
+            )
+            current_job_uris.update(
+                current_job_config.runtime_env_info.uris.py_modules_uris
+            )
+            if job_uris != current_job_uris and len(job_uris) > 0:
+                return ray_client_pb2.InitResponse(
+                    ok=False,
+                    msg="Runtime environment doesn't match "
+                    f"request one {job_config.runtime_env_info.uris} "
+                    f"current one {current_job_config.runtime_env_info.uris}",
+                )
+        return ray_client_pb2.InitResponse(ok=True)
+
+    @_use_response_cache
+    def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse:
+        try:
+            with disable_client_hook():
+                already_exists = ray.experimental.internal_kv._internal_kv_put(
+                    request.key,
+                    request.value,
+                    overwrite=request.overwrite,
+                    namespace=request.namespace,
+                )
+        except Exception as e:
+            return_exception_in_context(e, context)
+            already_exists = False
+        return ray_client_pb2.KVPutResponse(already_exists=already_exists)
+
+    def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse:
+        try:
+            with disable_client_hook():
+                value = ray.experimental.internal_kv._internal_kv_get(
+                    request.key, namespace=request.namespace
+                )
+        except Exception as e:
+            return_exception_in_context(e, context)
+            value = b""
+        return ray_client_pb2.KVGetResponse(value=value)
+
+    @_use_response_cache
+    def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse:
+        try:
+            with disable_client_hook():
+                deleted_num = ray.experimental.internal_kv._internal_kv_del(
+                    request.key,
+                    del_by_prefix=request.del_by_prefix,
+                    namespace=request.namespace,
+                )
+        except Exception as e:
+            return_exception_in_context(e, context)
+            deleted_num = 0
+        return ray_client_pb2.KVDelResponse(deleted_num=deleted_num)
+
+    def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse:
+        try:
+            with disable_client_hook():
+                keys = ray.experimental.internal_kv._internal_kv_list(
+                    request.prefix, namespace=request.namespace
+                )
+        except Exception as e:
+            return_exception_in_context(e, context)
+            keys = []
+        return ray_client_pb2.KVListResponse(keys=keys)
+
+    def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse:
+        try:
+            with disable_client_hook():
+                exists = ray.experimental.internal_kv._internal_kv_exists(
+                    request.key, namespace=request.namespace
+                )
+        except Exception as e:
+            return_exception_in_context(e, context)
+            exists = False
+        return ray_client_pb2.KVExistsResponse(exists=exists)
+
+    def ListNamedActors(
+        self, request, context=None
+    ) -> ray_client_pb2.ClientListNamedActorsResponse:
+        with disable_client_hook():
+            actors = ray.util.list_named_actors(all_namespaces=request.all_namespaces)
+
+        return ray_client_pb2.ClientListNamedActorsResponse(
+            actors_json=json.dumps(actors)
+        )
+
+    def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse:
+        resp = ray_client_pb2.ClusterInfoResponse()
+        resp.type = request.type
+        if request.type == ray_client_pb2.ClusterInfoType.CLUSTER_RESOURCES:
+            with disable_client_hook():
+                resources = ray.cluster_resources()
+            # Normalize resources into floats
+            # (the function may return values that are ints)
+            float_resources = {k: float(v) for k, v in resources.items()}
+            resp.resource_table.CopyFrom(
+                ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources)
+            )
+        elif request.type == ray_client_pb2.ClusterInfoType.AVAILABLE_RESOURCES:
+            with disable_client_hook():
+                resources = ray.available_resources()
+            # Normalize resources into floats
+            # (the function may return values that are ints)
+            float_resources = {k: float(v) for k, v in resources.items()}
+            resp.resource_table.CopyFrom(
+                ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources)
+            )
+        elif request.type == ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT:
+            ctx = ray_client_pb2.ClusterInfoResponse.RuntimeContext()
+            with disable_client_hook():
+                rtc = ray.get_runtime_context()
+                ctx.job_id = ray._common.utils.hex_to_binary(rtc.get_job_id())
+                ctx.node_id = ray._common.utils.hex_to_binary(rtc.get_node_id())
+                ctx.worker_id = ray._common.utils.hex_to_binary(rtc.get_worker_id())
+                ctx.namespace = rtc.namespace
+                ctx.capture_client_tasks = (
+                    rtc.should_capture_child_tasks_in_placement_group
+                )
+                ctx.gcs_address = rtc.gcs_address
+                ctx.runtime_env = rtc.get_runtime_env_string()
+                ctx.session_name = rtc.get_session_name()
+            resp.runtime_context.CopyFrom(ctx)
+        else:
+            with disable_client_hook():
+                resp.json = self._return_debug_cluster_info(request, context)
+        return resp
+
+    def _return_debug_cluster_info(self, request, context=None) -> str:
+        """Handle ClusterInfo requests that only return a json blob."""
+        data = None
+        if request.type == ray_client_pb2.ClusterInfoType.NODES:
+            data = ray.nodes()
+        elif request.type == ray_client_pb2.ClusterInfoType.IS_INITIALIZED:
+            data = ray.is_initialized()
+        elif request.type == ray_client_pb2.ClusterInfoType.TIMELINE:
+            data = ray.timeline()
+        elif request.type == ray_client_pb2.ClusterInfoType.PING:
+            data = {}
+        elif request.type == ray_client_pb2.ClusterInfoType.DASHBOARD_URL:
+            data = {"dashboard_url": ray._private.worker.get_dashboard_url()}
+        else:
+            raise TypeError("Unsupported cluster info type")
+        return json.dumps(data)
+
+    def release(self, client_id: str, id: bytes) -> bool:
+        with self.state_lock:
+            if client_id in self.object_refs:
+                if id in self.object_refs[client_id]:
+                    logger.debug(f"Releasing object {id.hex()} for {client_id}")
+                    del self.object_refs[client_id][id]
+                    return True
+
+            if client_id in self.actor_owners:
+                if id in self.actor_owners[client_id]:
+                    logger.debug(f"Releasing actor {id.hex()} for {client_id}")
+                    self.actor_owners[client_id].remove(id)
+                    if self._can_remove_actor_ref(id):
+                        logger.debug(f"Deleting reference to actor {id.hex()}")
+                        del self.actor_refs[id]
+                    return True
+
+            return False
+
+    def release_all(self, client_id):
+        with self.state_lock:
+            self._release_objects(client_id)
+            self._release_actors(client_id)
+        # NOTE: Try to actually dereference the object and actor refs.
+        # Otherwise dereferencing will happen later, which may run concurrently
+        # with ray.shutdown() and will crash the process. The crash is a bug
+        # that should be fixed eventually.
+        gc.collect()
+
+    def _can_remove_actor_ref(self, actor_id_bytes):
+        no_owner = not any(
+            actor_id_bytes in actor_list for actor_list in self.actor_owners.values()
+        )
+        return no_owner and actor_id_bytes not in self.named_actors
+
+    def _release_objects(self, client_id):
+        if client_id not in self.object_refs:
+            logger.debug(f"Releasing client with no references: {client_id}")
+            return
+        count = len(self.object_refs[client_id])
+        del self.object_refs[client_id]
+        if client_id in self.client_side_ref_map:
+            del self.client_side_ref_map[client_id]
+        if client_id in self.response_caches:
+            del self.response_caches[client_id]
+        logger.debug(f"Released all {count} objects for client {client_id}")
+
+    def _release_actors(self, client_id):
+        if client_id not in self.actor_owners:
+            logger.debug(f"Releasing client with no actors: {client_id}")
+            return
+
+        count = 0
+        actors_to_remove = self.actor_owners.pop(client_id)
+        for id_bytes in actors_to_remove:
+            count += 1
+            if self._can_remove_actor_ref(id_bytes):
+                logger.debug(f"Deleting reference to actor {id_bytes.hex()}")
+                del self.actor_refs[id_bytes]
+
+        logger.debug(f"Released all {count} actors for client: {client_id}")
+
+    @_use_response_cache
+    def Terminate(self, req, context=None):
+        if req.WhichOneof("terminate_type") == "task_object":
+            try:
+                object_ref = self.object_refs[req.client_id][req.task_object.id]
+                with disable_client_hook():
+                    ray.cancel(
+                        object_ref,
+                        force=req.task_object.force,
+                        recursive=req.task_object.recursive,
+                    )
+            except Exception as e:
+                return_exception_in_context(e, context)
+        elif req.WhichOneof("terminate_type") == "actor":
+            try:
+                actor_ref = self.actor_refs[req.actor.id]
+                with disable_client_hook():
+                    ray.kill(actor_ref, no_restart=req.actor.no_restart)
+            except Exception as e:
+                return_exception_in_context(e, context)
+        else:
+            raise RuntimeError(
+                "Client requested termination without providing a valid terminate_type"
+            )
+        return ray_client_pb2.TerminateResponse(ok=True)
+
+    def _async_get_object(
+        self,
+        request: ray_client_pb2.GetRequest,
+        client_id: str,
+        req_id: int,
+        result_queue: queue.Queue,
+        context=None,
+    ) -> Optional[ray_client_pb2.GetResponse]:
+        """Attempts to schedule a callback to push the GetResponse to the
+        main loop when the desired object is ready. If there is some failure
+        in scheduling, a GetResponse will be immediately returned.
+        """
+        if len(request.ids) != 1:
+            raise ValueError(
+                f"Async get() must have exactly 1 Object ID. Actual: {request}"
+            )
+        rid = request.ids[0]
+        ref = self.object_refs[client_id].get(rid, None)
+        if not ref:
+            return ray_client_pb2.GetResponse(
+                valid=False,
+                error=cloudpickle.dumps(
+                    ValueError(
+                        f"ClientObjectRef with id {rid} not found for "
+                        f"client {client_id}"
+                    )
+                ),
+            )
+        try:
+            logger.debug("async get: %s" % ref)
+            with disable_client_hook():
+
+                def send_get_response(result: Any) -> None:
+                    """Pushes GetResponses to the main DataPath loop to send
+                    to the client. This is called when the object is ready
+                    on the server side."""
+                    try:
+                        serialized = dumps_from_server(result, client_id, self)
+                        total_size = len(serialized)
+                        assert total_size > 0, "Serialized object cannot be zero bytes"
+                        total_chunks = math.ceil(
+                            total_size / OBJECT_TRANSFER_CHUNK_SIZE
+                        )
+                        for chunk_id in range(request.start_chunk_id, total_chunks):
+                            start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
+                            end = min(
+                                total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE
+                            )
+                            get_resp = ray_client_pb2.GetResponse(
+                                valid=True,
+                                data=serialized[start:end],
+                                chunk_id=chunk_id,
+                                total_chunks=total_chunks,
+                                total_size=total_size,
+                            )
+                            chunk_resp = ray_client_pb2.DataResponse(
+                                get=get_resp, req_id=req_id
+                            )
+                            result_queue.put(chunk_resp)
+                    except Exception as exc:
+                        get_resp = ray_client_pb2.GetResponse(
+                            valid=False, error=cloudpickle.dumps(exc)
+                        )
+                        resp = ray_client_pb2.DataResponse(get=get_resp, req_id=req_id)
+                        result_queue.put(resp)
+
+                ref._on_completed(send_get_response)
+                return None
+
+        except Exception as e:
+            return ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
+
+    def GetObject(self, request: ray_client_pb2.GetRequest, context):
+        metadata = dict(context.invocation_metadata())
+        client_id = metadata.get("client_id")
+        if client_id is None:
+            yield ray_client_pb2.GetResponse(
+                valid=False,
+                error=cloudpickle.dumps(
+                    ValueError("client_id is not specified in request metadata")
+                ),
+            )
+        else:
+            yield from self._get_object(request, client_id)
+
+    def _get_object(self, request: ray_client_pb2.GetRequest, client_id: str):
+        objectrefs = []
+        for rid in request.ids:
+            ref = self.object_refs[client_id].get(rid, None)
+            if ref:
+                objectrefs.append(ref)
+            else:
+                yield ray_client_pb2.GetResponse(
+                    valid=False,
+                    error=cloudpickle.dumps(
+                        ValueError(
+                            f"ClientObjectRef {rid} is not found for client {client_id}"
+                        )
+                    ),
+                )
+                return
+        try:
+            logger.debug("get: %s" % objectrefs)
+            with disable_client_hook():
+                items = ray.get(objectrefs, timeout=request.timeout)
+        except Exception as e:
+            yield ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e))
+            return
+        serialized = dumps_from_server(items, client_id, self)
+        total_size = len(serialized)
+        assert total_size > 0, "Serialized object cannot be zero bytes"
+        total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE)
+        for chunk_id in range(request.start_chunk_id, total_chunks):
+            start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE
+            end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE)
+            yield ray_client_pb2.GetResponse(
+                valid=True,
+                data=serialized[start:end],
+                chunk_id=chunk_id,
+                total_chunks=total_chunks,
+                total_size=total_size,
+            )
+
+    def PutObject(
+        self, request: ray_client_pb2.PutRequest, context=None
+    ) -> ray_client_pb2.PutResponse:
+        """gRPC entrypoint for unary PutObject"""
+        return self._put_object(request.data, request.client_ref_id, "", context)
+
+    def _put_object(
+        self,
+        data: Union[bytes, bytearray],
+        client_ref_id: bytes,
+        client_id: str,
+        context: Optional[grpc.ServicerContext] = None,
+    ) -> ray_client_pb2.PutResponse:
+        """Put an object in the cluster with ray.put() via gRPC.
+
+        Args:
+            data: Pickled data. Can either be bytearray if this is called
+              from the dataservicer, or bytes if called from PutObject.
+            client_ref_id: The id associated with this object on the client.
+            client_id: The client who owns this data, for tracking when to
+              delete this reference.
+            context: gRPC context.
+
+        Returns:
+            A ``PutResponse`` containing the resulting object ref id, or an
+            error payload if the put failed.
+        """
+        try:
+            obj = loads_from_client(data, self)
+            with disable_client_hook():
+                objectref = ray.put(obj)
+        except Exception as e:
+            logger.exception("Put failed:")
+            return ray_client_pb2.PutResponse(
+                id=b"", valid=False, error=cloudpickle.dumps(e)
+            )
+
+        self.object_refs[client_id][objectref.binary()] = objectref
+        if len(client_ref_id) > 0:
+            self.client_side_ref_map[client_id][client_ref_id] = objectref.binary()
+        logger.debug("put: %s" % objectref)
+        return ray_client_pb2.PutResponse(id=objectref.binary(), valid=True)
+
+    def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse:
+        object_refs = []
+        for rid in request.object_ids:
+            if rid not in self.object_refs[request.client_id]:
+                raise Exception(
+                    "Asking for a ref not associated with this client: %s" % str(rid)
+                )
+            object_refs.append(self.object_refs[request.client_id][rid])
+        num_returns = request.num_returns
+        timeout = request.timeout
+        try:
+            with disable_client_hook():
+                ready_object_refs, remaining_object_refs = ray.wait(
+                    object_refs,
+                    num_returns=num_returns,
+                    timeout=timeout if timeout != -1 else None,
+                )
+        except Exception as e:
+            # TODO(ameer): improve exception messages.
+            logger.error(f"Exception {e}")
+            return ray_client_pb2.WaitResponse(valid=False)
+        logger.debug(
+            "wait: %s %s" % (str(ready_object_refs), str(remaining_object_refs))
+        )
+        ready_object_ids = [
+            ready_object_ref.binary() for ready_object_ref in ready_object_refs
+        ]
+        remaining_object_ids = [
+            remaining_object_ref.binary()
+            for remaining_object_ref in remaining_object_refs
+        ]
+        return ray_client_pb2.WaitResponse(
+            valid=True,
+            ready_object_ids=ready_object_ids,
+            remaining_object_ids=remaining_object_ids,
+        )
+
+    def Schedule(
+        self,
+        task: ray_client_pb2.ClientTask,
+        arglist: List[Any],
+        kwargs: Dict[str, Any],
+        context=None,
+    ) -> ray_client_pb2.ClientTaskTicket:
+        logger.debug(
+            "schedule: %s %s"
+            % (task.name, ray_client_pb2.ClientTask.RemoteExecType.Name(task.type))
+        )
+        try:
+            with disable_client_hook():
+                if task.type == ray_client_pb2.ClientTask.FUNCTION:
+                    result = self._schedule_function(task, arglist, kwargs, context)
+                elif task.type == ray_client_pb2.ClientTask.ACTOR:
+                    result = self._schedule_actor(task, arglist, kwargs, context)
+                elif task.type == ray_client_pb2.ClientTask.METHOD:
+                    result = self._schedule_method(task, arglist, kwargs, context)
+                elif task.type == ray_client_pb2.ClientTask.NAMED_ACTOR:
+                    result = self._schedule_named_actor(task, context)
+                else:
+                    raise NotImplementedError(
+                        "Unimplemented Schedule task type: %s"
+                        % ray_client_pb2.ClientTask.RemoteExecType.Name(task.type)
+                    )
+                result.valid = True
+                return result
+        except Exception as e:
+            logger.debug("Caught schedule exception", exc_info=True)
+            return ray_client_pb2.ClientTaskTicket(
+                valid=False, error=cloudpickle.dumps(e)
+            )
+
+    def _schedule_method(
+        self,
+        task: ray_client_pb2.ClientTask,
+        arglist: List[Any],
+        kwargs: Dict[str, Any],
+        context=None,
+    ) -> ray_client_pb2.ClientTaskTicket:
+        actor_handle = self.actor_refs.get(task.payload_id)
+        if actor_handle is None:
+            raise Exception("Can't run an actor the server doesn't have a handle for")
+        method = getattr(actor_handle, task.name)
+        opts = decode_options(task.options)
+        if opts is not None:
+            method = method.options(**opts)
+        output = method.remote(*arglist, **kwargs)
+        ids = self.unify_and_track_outputs(output, task.client_id)
+        return ray_client_pb2.ClientTaskTicket(return_ids=ids)
+
+    def _schedule_actor(
+        self,
+        task: ray_client_pb2.ClientTask,
+        arglist: List[Any],
+        kwargs: Dict[str, Any],
+        context=None,
+    ) -> ray_client_pb2.ClientTaskTicket:
+        remote_class = self.lookup_or_register_actor(
+            task.payload_id, task.client_id, decode_options(task.baseline_options)
+        )
+        opts = decode_options(task.options)
+        if opts is not None:
+            remote_class = remote_class.options(**opts)
+        with current_server(self):
+            actor = remote_class.remote(*arglist, **kwargs)
+        self.actor_refs[actor._actor_id.binary()] = actor
+        self.actor_owners[task.client_id].add(actor._actor_id.binary())
+        return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()])
+
+    def _schedule_function(
+        self,
+        task: ray_client_pb2.ClientTask,
+        arglist: List[Any],
+        kwargs: Dict[str, Any],
+        context=None,
+    ) -> ray_client_pb2.ClientTaskTicket:
+        remote_func = self.lookup_or_register_func(
+            task.payload_id, task.client_id, decode_options(task.baseline_options)
+        )
+        opts = decode_options(task.options)
+        if opts is not None:
+            remote_func = remote_func.options(**opts)
+        with current_server(self):
+            output = remote_func.remote(*arglist, **kwargs)
+        ids = self.unify_and_track_outputs(output, task.client_id)
+        return ray_client_pb2.ClientTaskTicket(return_ids=ids)
+
+    def _schedule_named_actor(
+        self, task: ray_client_pb2.ClientTask, context=None
+    ) -> ray_client_pb2.ClientTaskTicket:
+        assert len(task.payload_id) == 0
+        # Convert empty string back to None.
+        actor = ray.get_actor(task.name, task.namespace or None)
+        bin_actor_id = actor._actor_id.binary()
+        if bin_actor_id not in self.actor_refs:
+            self.actor_refs[bin_actor_id] = actor
+        self.actor_owners[task.client_id].add(bin_actor_id)
+        self.named_actors.add(bin_actor_id)
+        return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()])
+
+    def lookup_or_register_func(
+        self, id: bytes, client_id: str, options: Optional[Dict]
+    ) -> ray.remote_function.RemoteFunction:
+        with disable_client_hook():
+            if id not in self.function_refs:
+                funcref = self.object_refs[client_id][id]
+                func = ray.get(funcref)
+                if not inspect.isfunction(func):
+                    raise Exception(
+                        "Attempting to register function that isn't a function."
+                    )
+                if options is None or len(options) == 0:
+                    self.function_refs[id] = ray.remote(func)
+                else:
+                    self.function_refs[id] = ray.remote(**options)(func)
+        return self.function_refs[id]
+
+    def lookup_or_register_actor(
+        self, id: bytes, client_id: str, options: Optional[Dict]
+    ):
+        with disable_client_hook():
+            if id not in self.registered_actor_classes:
+                actor_class_ref = self.object_refs[client_id][id]
+                actor_class = ray.get(actor_class_ref)
+                if not inspect.isclass(actor_class):
+                    raise Exception("Attempting to schedule actor that isn't a class.")
+                if options is None or len(options) == 0:
+                    reg_class = ray.remote(actor_class)
+                else:
+                    reg_class = ray.remote(**options)(actor_class)
+                self.registered_actor_classes[id] = reg_class
+
+        return self.registered_actor_classes[id]
+
+    def unify_and_track_outputs(self, output, client_id):
+        if output is None:
+            outputs = []
+        elif isinstance(output, list):
+            outputs = output
+        else:
+            outputs = [output]
+        for out in outputs:
+            if out.binary() in self.object_refs[client_id]:
+                logger.warning(f"Already saw object_ref {out}")
+            self.object_refs[client_id][out.binary()] = out
+        return [out.binary() for out in outputs]
+
+
+def return_exception_in_context(err, context):
+    if context is not None:
+        context.set_details(encode_exception(err))
+        # Note: https://grpc.github.io/grpc/core/md_doc_statuscodes.html
+        # ABORTED used here since it should never be generated by the
+        # grpc lib -- this way we know the error was generated by ray logic
+        context.set_code(grpc.StatusCode.ABORTED)
+
+
+def encode_exception(exception) -> str:
+    data = cloudpickle.dumps(exception)
+    return base64.standard_b64encode(data).decode()
+
+
+def decode_options(options: ray_client_pb2.TaskOptions) -> Optional[Dict[str, Any]]:
+    if not options.pickled_options:
+        return None
+    opts = pickle.loads(options.pickled_options)
+    assert isinstance(opts, dict)
+
+    return opts
+
+
+def serve(host: str, port: int, ray_connect_handler=None):
+    def default_connect_handler(
+        job_config: JobConfig = None, **ray_init_kwargs: Dict[str, Any]
+    ):
+        with disable_client_hook():
+            if not ray.is_initialized():
+                return ray.init(job_config=job_config, **ray_init_kwargs)
+
+    from ray._private.grpc_utils import create_grpc_server_with_interceptors
+
+    ray_connect_handler = ray_connect_handler or default_connect_handler
+    server = create_grpc_server_with_interceptors(
+        max_workers=CLIENT_SERVER_MAX_THREADS,
+        thread_name_prefix="ray_client_server",
+        options=GRPC_OPTIONS,
+        asynchronous=False,
+    )
+    task_servicer = RayletServicer(ray_connect_handler)
+    data_servicer = DataServicer(task_servicer)
+    logs_servicer = LogstreamServicer()
+    ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server)
+    ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server)
+    ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server)
+    if not is_localhost(host):
+        add_port_to_grpc_server(server, build_address(get_localhost_ip(), port))
+    add_port_to_grpc_server(server, build_address(host, port))
+    current_handle = ClientServerHandle(
+        task_servicer=task_servicer,
+        data_servicer=data_servicer,
+        logs_servicer=logs_servicer,
+        grpc_server=server,
+    )
+    server.start()
+    return current_handle
+
+
+def init_and_serve(host: str, port: int, *args, **kwargs):
+    with disable_client_hook():
+        # Disable client mode inside the worker's environment
+        info = ray.init(*args, **kwargs)
+
+    def ray_connect_handler(job_config=None, **ray_init_kwargs):
+        # Ray client will disconnect from ray when
+        # num_clients == 0.
+        if ray.is_initialized():
+            return info
+        else:
+            return ray.init(job_config=job_config, *args, **kwargs)
+
+    server_handle = serve(host, port, ray_connect_handler=ray_connect_handler)
+    return (server_handle, info)
+
+
+def shutdown_with_server(server, _exiting_interpreter=False):
+    server.stop(1)
+    with disable_client_hook():
+        ray.shutdown(_exiting_interpreter=_exiting_interpreter)
+
+
+def create_ray_handler(address, redis_password, redis_username=None):
+    def ray_connect_handler(job_config: JobConfig = None, **ray_init_kwargs):
+        if address:
+            if redis_password:
+                ray.init(
+                    address=address,
+                    _redis_username=redis_username,
+                    _redis_password=redis_password,
+                    job_config=job_config,
+                    **ray_init_kwargs,
+                )
+            else:
+                ray.init(address=address, job_config=job_config, **ray_init_kwargs)
+        else:
+            ray.init(job_config=job_config, **ray_init_kwargs)
+
+    return ray_connect_handler
+
+
+def try_create_gcs_client(address: Optional[str]) -> Optional[GcsClient]:
+    """
+    Try to create a gcs client based on the command line args or by
+    autodetecting a running Ray cluster.
+    """
+    address = canonicalize_bootstrap_address_or_die(address)
+    return GcsClient(address=address)
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--host",
+        type=str,
+        default=get_all_interfaces_ip(),
+        help="Host IP to bind to. Defaults to all interfaces (0.0.0.0/::).",
+    )
+    parser.add_argument("-p", "--port", type=int, default=10001, help="Port to bind to")
+    parser.add_argument(
+        "--mode",
+        type=str,
+        choices=["proxy", "legacy", "specific-server"],
+        default="proxy",
+    )
+    parser.add_argument(
+        "--address", required=False, type=str, help="Address to use to connect to Ray"
+    )
+    parser.add_argument(
+        "--redis-username",
+        required=False,
+        type=str,
+        help="username for connecting to Redis",
+    )
+    parser.add_argument(
+        "--runtime-env-agent-address",
+        required=False,
+        type=str,
+        default=None,
+        help="The port to use for connecting to the runtime_env_agent.",
+    )
+    parser.add_argument(
+        "--node-id",
+        required=False,
+        type=str,
+        default=None,
+        help="The hex ID of this node.",
+    )
+    args, _ = parser.parse_known_args()
+    redis_password = os.environ.get(ray_constants.RAY_REDIS_PASSWORD_ENV)
+    setup_logger(ray_constants.LOGGER_LEVEL, ray_constants.LOGGER_FORMAT)
+
+    ray_connect_handler = create_ray_handler(
+        args.address, redis_password, args.redis_username
+    )
+
+    hostport = build_address(args.host, args.port)
+    args_str = str(args)
+    logger.info(f"Starting Ray Client server on {hostport}, args {args_str}")
+    if args.mode == "proxy":
+        server = serve_proxier(
+            args.host,
+            args.port,
+            args.address,
+            redis_username=args.redis_username,
+            redis_password=redis_password,
+            runtime_env_agent_address=args.runtime_env_agent_address,
+            node_id=args.node_id,
+        )
+    else:
+        server = serve(args.host, args.port, ray_connect_handler)
+
+    try:
+        idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S
+        while True:
+            health_report = {
+                "time": time.time(),
+            }
+
+            try:
+                if not ray.experimental.internal_kv._internal_kv_initialized():
+                    gcs_client = try_create_gcs_client(args.address)
+                    ray.experimental.internal_kv._initialize_internal_kv(gcs_client)
+                ray.experimental.internal_kv._internal_kv_put(
+                    "ray_client_server",
+                    json.dumps(health_report),
+                    namespace=ray_constants.KV_NAMESPACE_HEALTHCHECK,
+                )
+            except Exception as e:
+                logger.error(
+                    f"[{args.mode}] Failed to put health check on {args.address}"
+                )
+                logger.exception(e)
+
+            time.sleep(1)
+            if args.mode == "specific-server":
+                if server.data_servicer.num_clients > 0:
+                    idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S
+                else:
+                    idle_checks_remaining -= 1
+                if idle_checks_remaining == 0:
+                    raise KeyboardInterrupt()
+                if (
+                    idle_checks_remaining % 5 == 0
+                    and idle_checks_remaining != TIMEOUT_FOR_SPECIFIC_SERVER_S
+                ):
+                    logger.info(f"{idle_checks_remaining} idle checks before shutdown.")
+
+    except KeyboardInterrupt:
+        server.stop(0)
+
+
+if __name__ == "__main__":
+    main()

python/ray/util/client/server/server_pickler.py

@@ -0,0 +1,124 @@
+"""Implements the client side of the client/server pickling protocol.
+
+These picklers are aware of the server internals and can find the
+references held for the client within the server.
+
+More discussion about the client/server pickling protocol can be found in:
+
+  ray/util/client/client_pickler.py
+
+ServerPickler dumps ray objects from the server into the appropriate stubs.
+ClientUnpickler loads stubs from the client and finds their associated handle
+in the server instance.
+"""
+import io
+from typing import TYPE_CHECKING, Any
+
+import ray
+import ray.cloudpickle as cloudpickle
+from ray._private.client_mode_hook import disable_client_hook
+from ray.util.client.client_pickler import PickleStub
+from ray.util.client.server.server_stubs import (
+    ClientReferenceActor,
+    ClientReferenceFunction,
+)
+
+if TYPE_CHECKING:
+    from ray.util.client.server.server import RayletServicer
+
+import pickle  # noqa: F401
+
+
+class ServerPickler(cloudpickle.CloudPickler):
+    def __init__(self, client_id: str, server: "RayletServicer", *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client_id = client_id
+        self.server = server
+
+    def persistent_id(self, obj):
+        if isinstance(obj, ray.ObjectRef):
+            obj_id = obj.binary()
+            if obj_id not in self.server.object_refs[self.client_id]:
+                # We're passing back a reference, probably inside a reference.
+                # Let's hold onto it.
+                self.server.object_refs[self.client_id][obj_id] = obj
+            return PickleStub(
+                type="Object",
+                client_id=self.client_id,
+                ref_id=obj_id,
+                name=None,
+                baseline_options=None,
+            )
+        elif isinstance(obj, ray.actor.ActorHandle):
+            actor_id = obj._actor_id.binary()
+            if actor_id not in self.server.actor_refs:
+                # We're passing back a handle, probably inside a reference.
+                self.server.actor_refs[actor_id] = obj
+            if actor_id not in self.server.actor_owners[self.client_id]:
+                self.server.actor_owners[self.client_id].add(actor_id)
+            return PickleStub(
+                type="Actor",
+                client_id=self.client_id,
+                ref_id=obj._actor_id.binary(),
+                name=None,
+                baseline_options=None,
+            )
+        return None
+
+
+class ClientUnpickler(pickle.Unpickler):
+    def __init__(self, server, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.server = server
+
+    def persistent_load(self, pid):
+        assert isinstance(pid, PickleStub)
+        if pid.type == "Ray":
+            return ray
+        elif pid.type == "Object":
+            return self.server.object_refs[pid.client_id][pid.ref_id]
+        elif pid.type == "Actor":
+            return self.server.actor_refs[pid.ref_id]
+        elif pid.type == "RemoteFuncSelfReference":
+            return ClientReferenceFunction(pid.client_id, pid.ref_id)
+        elif pid.type == "RemoteFunc":
+            return self.server.lookup_or_register_func(
+                pid.ref_id, pid.client_id, pid.baseline_options
+            )
+        elif pid.type == "RemoteActorSelfReference":
+            return ClientReferenceActor(pid.client_id, pid.ref_id)
+        elif pid.type == "RemoteActor":
+            return self.server.lookup_or_register_actor(
+                pid.ref_id, pid.client_id, pid.baseline_options
+            )
+        elif pid.type == "RemoteMethod":
+            actor = self.server.actor_refs[pid.ref_id]
+            return getattr(actor, pid.name)
+        else:
+            raise NotImplementedError("Uncovered client data type")
+
+
+def dumps_from_server(
+    obj: Any, client_id: str, server_instance: "RayletServicer", protocol=None
+) -> bytes:
+    with io.BytesIO() as file:
+        sp = ServerPickler(client_id, server_instance, file, protocol=protocol)
+        sp.dump(obj)
+        return file.getvalue()
+
+
+def loads_from_client(
+    data: bytes,
+    server_instance: "RayletServicer",
+    *,
+    fix_imports=True,
+    encoding="ASCII",
+    errors="strict"
+) -> Any:
+    with disable_client_hook():
+        if isinstance(data, str):
+            raise TypeError("Can't load pickle from unicode string")
+        file = io.BytesIO(data)
+        return ClientUnpickler(
+            server_instance, file, fix_imports=fix_imports, encoding=encoding
+        ).load()

python/ray/util/client/server/server_stubs.py

@@ -0,0 +1,66 @@
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+
+_current_server = None
+
+
+@contextmanager
+def current_server(r):
+    global _current_server
+    remote = _current_server
+    _current_server = r
+    try:
+        yield
+    finally:
+        _current_server = remote
+
+
+class ClientReferenceSentinel(ABC):
+    def __init__(self, client_id, id):
+        self.client_id = client_id
+        self.id = id
+
+    def __reduce__(self):
+        remote_obj = self.get_remote_obj()
+        if remote_obj is None:
+            return (self.__class__, (self.client_id, self.id))
+        return (identity, (remote_obj,))
+
+    @abstractmethod
+    def get_remote_obj(self):
+        pass
+
+    def get_real_ref_from_server(self):
+        global _current_server
+        if _current_server is None:
+            return None
+        client_map = _current_server.client_side_ref_map.get(self.client_id, None)
+        if client_map is None:
+            return None
+        return client_map.get(self.id, None)
+
+
+class ClientReferenceActor(ClientReferenceSentinel):
+    def get_remote_obj(self):
+        global _current_server
+        real_ref_id = self.get_real_ref_from_server()
+        if real_ref_id is None:
+            return None
+        return _current_server.lookup_or_register_actor(
+            real_ref_id, self.client_id, None
+        )
+
+
+class ClientReferenceFunction(ClientReferenceSentinel):
+    def get_remote_obj(self):
+        global _current_server
+        real_ref_id = self.get_real_ref_from_server()
+        if real_ref_id is None:
+            return None
+        return _current_server.lookup_or_register_func(
+            real_ref_id, self.client_id, None
+        )
+
+
+def identity(x):
+    return x