ray-project--ray/python/ray/serve/context.py

"""
This file stores global state for a Serve application. Deployment replicas
can use this state to access metadata or the Serve controller.
"""

import asyncio
import contextvars
import logging
import os
from collections import defaultdict
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional

import ray
from ray.serve._private.client import ServeControllerClient
from ray.serve._private.common import DeploymentID, ReplicaID
from ray.serve._private.config import DeploymentConfig
from ray.serve._private.constants import (
    RAY_SERVE_INTERNAL_DEPLOYMENT_ACTOR_NAME_ENV_VAR,
    RAY_SERVE_INTERNAL_DEPLOYMENT_APP_NAME_ENV_VAR,
    RAY_SERVE_INTERNAL_DEPLOYMENT_CODE_VERSION_ENV_VAR,
    RAY_SERVE_INTERNAL_DEPLOYMENT_NAME_ENV_VAR,
    SERVE_CONTROLLER_NAME,
    SERVE_LOGGER_NAME,
    SERVE_NAMESPACE,
)
from ray.serve._private.replica_result import ReplicaResult
from ray.serve._private.utils import get_deployment_actor_name
from ray.serve.exceptions import RayServeException
from ray.serve.gang import GangContext
from ray.serve.grpc_util import RayServegRPCContext
from ray.serve.schema import ReplicaRank
from ray.util.annotations import DeveloperAPI

logger = logging.getLogger(SERVE_LOGGER_NAME)

_INTERNAL_REPLICA_CONTEXT: "ReplicaContext" = None
_INTERNAL_DEPLOYMENT_ACTOR_CONTEXT: "DeploymentActorContext" = None
_global_client: ServeControllerClient = None
# Ray job id (driver session) that created the cached client above. ray.shutdown()
# leaves this module global intact, so after a reconnect (new job id) the cached
# handle is unusable and must not be reused. See #64647.
_global_client_job_id: Optional[str] = None


@DeveloperAPI
@dataclass
class ReplicaContext:
    """Stores runtime context info for replicas.

    Fields:
        - app_name: name of the application the replica is a part of.
        - deployment: name of the deployment the replica is a part of.
        - replica_tag: unique ID for the replica.
        - servable_object: instance of the user class/function this replica is running.
        - rank: the rank of the replica.
        - world_size: the number of replicas in the deployment.
        - gang_context: context information for the gang the replica is part of.
        - code_version: code version of the deployment (for get_deployment_actor).
    """

    replica_id: ReplicaID
    servable_object: Callable
    _deployment_config: DeploymentConfig
    rank: ReplicaRank
    world_size: int
    _handle_registration_callback: Optional[Callable[[DeploymentID], None]] = None
    gang_context: Optional[GangContext] = None
    code_version: Optional[str] = None

    @property
    def app_name(self) -> str:
        return self.replica_id.deployment_id.app_name

    @property
    def deployment(self) -> str:
        return self.replica_id.deployment_id.name

    @property
    def replica_tag(self) -> str:
        return self.replica_id.unique_id


@DeveloperAPI
@dataclass
class DeploymentActorContext:
    """Stores runtime context info for deployment-scoped actors."""

    deployment_id: DeploymentID
    actor_name: str
    code_version: Optional[str] = None

    @property
    def app_name(self) -> str:
        return self.deployment_id.app_name

    @property
    def deployment(self) -> str:
        return self.deployment_id.name


def _get_global_client(
    raise_if_no_controller_running: bool = True,
) -> Optional[ServeControllerClient]:
    """Gets the global client, which stores the controller's handle.

    Args:
        raise_if_no_controller_running: Whether to raise an exception if
            there is no currently running Serve controller.

    Returns:
        ServeControllerClient to the running Serve controller. If there
        is no running controller and raise_if_no_controller_running is
        set to False, returns None.

    Raises:
        RayServeException: If there is no running Serve controller actor
            and raise_if_no_controller_running is set to True.
    """

    if _cached_client_from_current_session():
        return _global_client

    # No usable cache (never set, or left over from a previous driver session).
    # Drop any stale handle and rediscover the controller by name.
    _disconnect()
    return _connect(raise_if_no_controller_running)


def _check_cached_client_alive() -> tuple:
    """Health-check the cached controller client.

    Returns:
        (client, had_cached) tuple.
        - ``(client, True)``: cached client from the current session is alive.
        - ``(None, True)``: cached client from the current session existed but
          is unreachable; the cache has been cleared. Callers should **not**
          attempt to reconnect via ``_connect()`` because GCS is likely dead and
          ``ray.get_actor()`` would hang until the 60-second C++ GCS
          reconnection timeout kills the process.
        - ``(None, False)``: no usable cached client (never set, or left over
          from a previous driver session). Callers may safely call
          ``_get_global_client()`` to discover a running controller.
    """

    if not _cached_client_from_current_session():
        # Either nothing cached, or a handle left behind by a previous driver
        # session. In the latter case GCS is alive (we are connected now), so it
        # is safe for callers to reconnect. Drop any stale handle.
        _disconnect()
        return None, False

    try:
        ray.get(_global_client._controller.check_alive.remote(), timeout=5)
        return _global_client, True
    except Exception as e:
        logger.info(f"The cached controller has died or is unreachable: {e}.")
        _disconnect()
        return None, True


def _set_global_client(client):
    global _global_client, _global_client_job_id
    _global_client = client
    _global_client_job_id = (
        ray.get_runtime_context().get_job_id() if client is not None else None
    )


def _disconnect():
    """Forget the cached controller client for this driver session.

    Mirrors ``_connect()``. This does **not** shut Serve down on the cluster; it
    only drops the local cached handle so the next call rediscovers the
    controller by name.
    """
    _set_global_client(None)


def _cached_client_from_current_session() -> bool:
    """Whether a usable cached client exists for the current Ray session.

    ``_global_client`` is a module global that survives ``ray.shutdown()``, so a
    handle cached by a previous driver session cannot be used after the driver
    reconnects (which yields a new job id). Comparing job ids detects that case.
    """
    if _global_client is None or not ray.is_initialized():
        return False
    return _global_client_job_id == ray.get_runtime_context().get_job_id()


def _get_internal_replica_context():
    return _INTERNAL_REPLICA_CONTEXT


def _get_internal_deployment_actor_context():
    global _INTERNAL_DEPLOYMENT_ACTOR_CONTEXT

    if _INTERNAL_DEPLOYMENT_ACTOR_CONTEXT is not None:
        return _INTERNAL_DEPLOYMENT_ACTOR_CONTEXT

    app_name = os.environ.get(RAY_SERVE_INTERNAL_DEPLOYMENT_APP_NAME_ENV_VAR)
    deployment_name = os.environ.get(RAY_SERVE_INTERNAL_DEPLOYMENT_NAME_ENV_VAR)
    actor_name = os.environ.get(RAY_SERVE_INTERNAL_DEPLOYMENT_ACTOR_NAME_ENV_VAR)

    if app_name is None or deployment_name is None or actor_name is None:
        return None

    _INTERNAL_DEPLOYMENT_ACTOR_CONTEXT = DeploymentActorContext(
        deployment_id=DeploymentID(name=deployment_name, app_name=app_name),
        actor_name=actor_name,
        code_version=os.environ.get(RAY_SERVE_INTERNAL_DEPLOYMENT_CODE_VERSION_ENV_VAR),
    )
    return _INTERNAL_DEPLOYMENT_ACTOR_CONTEXT


def _get_deployment_actor(actor_name: str):
    """Get a handle to a deployment-scoped actor by name.

    Thin wrapper around ``ray.get_actor`` with the Serve deployment-actor naming
    convention. See ``serve.get_deployment_actor`` docstring for behavior,
    ``ValueError``/``RayActorError`` expectations, and refresh patterns.
    """
    internal_context = _get_internal_replica_context()
    if internal_context is None:
        raise RayServeException(
            "`serve.get_deployment_actor()` may only be called from within "
            "a Ray Serve deployment replica."
        )
    deployment_id = internal_context.replica_id.deployment_id
    return ray.get_actor(
        get_deployment_actor_name(
            deployment_id,
            actor_name,
            code_version=internal_context.code_version,
        ),
        namespace=SERVE_NAMESPACE,
    )


def _set_internal_replica_context(
    *,
    replica_id: ReplicaID,
    servable_object: Callable,
    _deployment_config: DeploymentConfig,
    rank: ReplicaRank,
    world_size: int,
    handle_registration_callback: Optional[Callable[[str, str], None]] = None,
    gang_context: Optional[GangContext] = None,
    code_version: Optional[str] = None,
):
    global _INTERNAL_REPLICA_CONTEXT
    _INTERNAL_REPLICA_CONTEXT = ReplicaContext(
        replica_id=replica_id,
        servable_object=servable_object,
        _deployment_config=_deployment_config,
        rank=rank,
        world_size=world_size,
        _handle_registration_callback=handle_registration_callback,
        gang_context=gang_context,
        code_version=code_version,
    )


def _connect(raise_if_no_controller_running: bool = True) -> ServeControllerClient:
    """Connect to an existing Serve application on this Ray cluster.

    If called from within a replica, this will connect to the same Serve
    app that the replica is running in.

    Args:
        raise_if_no_controller_running: If ``True``, raise when no Serve
            controller actor is found. If ``False``, return ``None`` instead.

    Returns:
        ServeControllerClient that encapsulates a Ray actor handle to the
        existing Serve application's Serve Controller. None if there is
        no running Serve controller actor and raise_if_no_controller_running
        is set to False.

    Raises:
        RayServeException: If there is no running Serve controller actor
            and raise_if_no_controller_running is set to True.
    """

    # Initialize ray if needed.
    ray._private.worker.global_worker._filter_logs_by_job = False
    if not ray.is_initialized():
        ray.init(namespace=SERVE_NAMESPACE)

    # Try to get serve controller if it exists
    try:
        controller = ray.get_actor(SERVE_CONTROLLER_NAME, namespace=SERVE_NAMESPACE)
    except ValueError:
        if raise_if_no_controller_running:
            raise RayServeException(
                "There is no Serve instance running on this Ray cluster."
            )
        return

    client = ServeControllerClient(
        controller,
    )
    _set_global_client(client)
    return client


# Serve request context var which is used for storing the internal
# request context information.
# route_prefix: http url route path, e.g. http://127.0.0.1:/app
#     the route is "/app". When you send requests by handle,
#     the route is empty.
# request_id: the request id is generated from http proxy, the value
#     shouldn't be changed when the variable is set.
#     This can be from the client and is used for logging.
# _internal_request_id: the request id is generated from the proxy. Used to track the
#     request objects in the system.
# note:
#   The request context is readonly to avoid potential
#       async task conflicts when using it concurrently.


@dataclass(frozen=True)
class _RequestContext:
    route: str = ""
    request_id: str = ""
    _internal_request_id: str = ""
    app_name: str = ""
    multiplexed_model_id: str = ""
    session_id: str = ""
    grpc_context: Optional[RayServegRPCContext] = None
    is_http_request: bool = False
    cancel_on_parent_request_cancel: bool = False
    # The client address in "host:port" format, if available.
    _client: str = ""
    # Ray tracing context for this request (if tracing is enabled)
    # This is extracted from _ray_trace_ctx kwarg at the replica entry point
    # Advanced users can access this to propagate tracing to external systems
    _ray_trace_ctx: Optional[dict] = None


_serve_request_context = contextvars.ContextVar(
    "Serve internal request context variable", default=None
)

_serve_batch_request_context = contextvars.ContextVar(
    "Serve internal batching request context variable", default=None
)


def _get_serve_request_context():
    """Get the current request context.

    Returns:
        The current request context
    """

    if _serve_request_context.get() is None:
        _serve_request_context.set(_RequestContext())
    return _serve_request_context.get()


def _get_serve_batch_request_context():
    """Get the list of request contexts for the current batch."""
    if _serve_batch_request_context.get() is None:
        _serve_batch_request_context.set([])
    return _serve_batch_request_context.get()


def _set_request_context(
    route: str = "",
    request_id: str = "",
    _internal_request_id: str = "",
    app_name: str = "",
    multiplexed_model_id: str = "",
):
    """Set the request context. If the value is not set,
    the current context value will be used."""

    current_request_context = _get_serve_request_context()

    _serve_request_context.set(
        _RequestContext(
            route=route or current_request_context.route,
            request_id=request_id or current_request_context.request_id,
            _internal_request_id=_internal_request_id
            or current_request_context._internal_request_id,
            app_name=app_name or current_request_context.app_name,
            multiplexed_model_id=multiplexed_model_id
            or current_request_context.multiplexed_model_id,
        )
    )


def _unset_request_context():
    """Unset the request context."""
    _serve_request_context.set(_RequestContext())


def _set_batch_request_context(request_contexts: List[_RequestContext]):
    """Add the request context to the batch request context."""
    _serve_batch_request_context.set(request_contexts)


# `_requests_pending_assignment` is a map from request ID to a
# dictionary of asyncio tasks.
# The request ID points to an ongoing request that is executing on the
# current replica, and the asyncio tasks are ongoing tasks started on
# the router to assign child requests to downstream replicas.

# A dictionary is used over a set to track the asyncio tasks for more
# efficient addition and deletion time complexity. A uniquely generated
# `response_id` is used to identify each task.

_requests_pending_assignment: Dict[str, Dict[str, asyncio.Task]] = defaultdict(dict)


# Note that the functions below that manipulate
# `_requests_pending_assignment` are NOT thread-safe. They are only
# expected to be called from the same thread/asyncio event-loop.


def _get_requests_pending_assignment(parent_request_id: str) -> Dict[str, asyncio.Task]:
    if parent_request_id in _requests_pending_assignment:
        return _requests_pending_assignment[parent_request_id]

    return {}


def _add_request_pending_assignment(parent_request_id: str, response_id: str, task):
    # NOTE: `parent_request_id` is the `internal_request_id` corresponding
    # to an ongoing Serve request, so it is always non-empty.
    _requests_pending_assignment[parent_request_id][response_id] = task


def _remove_request_pending_assignment(parent_request_id: str, response_id: str):
    if response_id in _requests_pending_assignment[parent_request_id]:
        del _requests_pending_assignment[parent_request_id][response_id]

    if len(_requests_pending_assignment[parent_request_id]) == 0:
        del _requests_pending_assignment[parent_request_id]


# `_in_flight_requests` is a map from request ID to a dictionary of replica results.
# The request ID points to an ongoing Serve request, and the replica results are
# in-flight child requests that have been assigned to a downstream replica.

# A dictionary is used over a set to track the replica results for more
# efficient addition and deletion time complexity. A uniquely generated
# `response_id` is used to identify each replica result.

_in_flight_requests: Dict[str, Dict[str, ReplicaResult]] = defaultdict(dict)

# Note that the functions below that manipulate `_in_flight_requests`
# are NOT thread-safe. They are only expected to be called from the
# same thread/asyncio event-loop.


def _get_in_flight_requests(parent_request_id):
    if parent_request_id in _in_flight_requests:
        return _in_flight_requests[parent_request_id]

    return {}


def _add_in_flight_request(parent_request_id, response_id, replica_result):
    _in_flight_requests[parent_request_id][response_id] = replica_result


def _remove_in_flight_request(parent_request_id, response_id):
    if response_id in _in_flight_requests[parent_request_id]:
        del _in_flight_requests[parent_request_id][response_id]

    if len(_in_flight_requests[parent_request_id]) == 0:
        del _in_flight_requests[parent_request_id]