chore: import upstream snapshot with attribution

python/ray/util/state/__init__.py

@@ -0,0 +1,49 @@
+from ray.util.state.api import (
+    StateApiClient,
+    get_actor,
+    get_job,
+    get_log,
+    get_node,
+    get_objects,
+    get_placement_group,
+    get_task,
+    get_worker,
+    list_actors,
+    list_cluster_events,
+    list_jobs,
+    list_logs,
+    list_nodes,
+    list_objects,
+    list_placement_groups,
+    list_runtime_envs,
+    list_tasks,
+    list_workers,
+    summarize_actors,
+    summarize_objects,
+    summarize_tasks,
+)
+
+__all__ = [
+    "get_actor",
+    "get_log",
+    "get_node",
+    "get_objects",
+    "get_placement_group",
+    "get_task",
+    "get_worker",
+    "get_job",
+    "list_actors",
+    "list_jobs",
+    "list_nodes",
+    "list_placement_groups",
+    "list_tasks",
+    "list_workers",
+    "list_objects",
+    "list_runtime_envs",
+    "list_logs",
+    "list_cluster_events",
+    "summarize_actors",
+    "summarize_objects",
+    "summarize_tasks",
+    "StateApiClient",
+]

python/ray/util/state/exception.py

@@ -0,0 +1,18 @@
+"""Internal Error"""
+
+
+class DataSourceUnavailable(Exception):
+    pass
+
+
+"""User-facing Error"""
+
+
+class RayStateApiException(Exception):
+    pass
+
+
+class ServerUnavailable(RayStateApiException):
+    """Thrown when failing to connect to dashboard server"""
+
+    pass

python/ray/util/state/state_manager.py

@@ -0,0 +1,510 @@
+import dataclasses
+import inspect
+import json
+import logging
+from functools import wraps
+from typing import Any, Dict, List, Optional, Tuple
+
+import aiohttp
+import grpc
+from grpc.aio._call import UnaryStreamCall
+
+import ray.dashboard.consts as dashboard_consts
+import ray.dashboard.modules.log.log_consts as log_consts
+from ray._common.network_utils import build_address
+from ray._common.utils import hex_to_binary
+from ray._private import ray_constants
+from ray._private.authentication.http_token_authentication import (
+    get_auth_headers_if_auth_enabled,
+)
+from ray._raylet import ActorID, GcsClient, JobID, NodeID, TaskID
+from ray.core.generated import gcs_service_pb2_grpc
+from ray.core.generated.gcs_pb2 import ActorTableData, GcsNodeInfo
+from ray.core.generated.gcs_service_pb2 import (
+    FilterPredicate,
+    GetAllActorInfoReply,
+    GetAllActorInfoRequest,
+    GetAllNodeInfoRequest,
+    GetAllPlacementGroupReply,
+    GetAllPlacementGroupRequest,
+    GetAllWorkerInfoReply,
+    GetAllWorkerInfoRequest,
+    GetTaskEventsReply,
+    GetTaskEventsRequest,
+)
+from ray.core.generated.node_manager_pb2 import (
+    GetObjectsInfoReply,
+    GetObjectsInfoRequest,
+)
+from ray.core.generated.node_manager_pb2_grpc import NodeManagerServiceStub
+from ray.core.generated.reporter_pb2 import (
+    ListLogsReply,
+    ListLogsRequest,
+    StreamLogRequest,
+)
+from ray.core.generated.reporter_pb2_grpc import LogServiceStub
+from ray.core.generated.runtime_env_agent_pb2 import (
+    GetRuntimeEnvsInfoReply,
+    GetRuntimeEnvsInfoRequest,
+)
+from ray.dashboard.modules.job.common import JobInfoStorageClient
+from ray.dashboard.modules.job.pydantic_models import JobDetails, JobType
+from ray.dashboard.modules.job.utils import get_driver_jobs
+from ray.util.state.common import (
+    RAY_MAX_LIMIT_FROM_DATA_SOURCE,
+    PredicateType,
+    SupportedFilterType,
+)
+from ray.util.state.exception import DataSourceUnavailable
+
+logger = logging.getLogger(__name__)
+
+_STATE_MANAGER_GRPC_OPTIONS = [
+    *ray_constants.GLOBAL_GRPC_OPTIONS,
+    ("grpc.max_send_message_length", ray_constants.GRPC_CPP_MAX_MESSAGE_SIZE),
+    ("grpc.max_receive_message_length", ray_constants.GRPC_CPP_MAX_MESSAGE_SIZE),
+]
+
+
+def handle_grpc_network_errors(func):
+    """Decorator to add a network handling logic.
+
+    It is a helper method for `StateDataSourceClient`.
+    The method can only be used for async methods.
+    """
+    assert inspect.iscoroutinefunction(func)
+
+    @wraps(func)
+    async def api_with_network_error_handler(*args: Any, **kwargs: Any):
+        """Apply the network error handling logic to each APIs,
+        such as retry or exception policies.
+
+        Args:
+            *args: Positional arguments forwarded to the wrapped coroutine.
+            **kwargs: Keyword arguments forwarded to the wrapped coroutine.
+
+        Returns:
+            If RPC succeeds, it returns what the original function returns.
+            If RPC fails, it raises exceptions.
+
+        Raises:
+            DataSourceUnavailable: if the source is unavailable because it is down
+                or there's a slow network issue causing timeout.
+
+            Exception: Otherwise, the raw network exceptions (e.g., gRPC) will be
+                raised.
+        """
+        try:
+            return await func(*args, **kwargs)
+        except grpc.aio.AioRpcError as e:
+            if (
+                e.code() == grpc.StatusCode.DEADLINE_EXCEEDED
+                or e.code() == grpc.StatusCode.UNAVAILABLE
+            ):
+                raise DataSourceUnavailable(
+                    "Failed to query the data source. "
+                    "It is either there's a network issue, or the source is down."
+                ) from e
+            else:
+                logger.exception(e)
+                raise e
+
+    return api_with_network_error_handler
+
+
+class StateDataSourceClient:
+    """The client to query states from various data sources such as Raylet, GCS, Agents.
+
+    Note that it doesn't directly query core workers. They are proxied through raylets.
+
+    The module is not in charge of service discovery. The caller is responsible for
+    finding services and register stubs through `register*` APIs.
+
+    Non `register*` APIs
+    - Return the protobuf directly if it succeeds to query the source.
+    - Raises an exception if there's any network issue.
+    - throw a ValueError if it cannot find the source.
+    """
+
+    def __init__(self, gcs_channel: grpc.aio.Channel, gcs_client: GcsClient):
+        self.register_gcs_client(gcs_channel)
+        self._job_client = JobInfoStorageClient(gcs_client)
+        self._gcs_client = gcs_client
+        self._client_session = aiohttp.ClientSession()
+
+    def register_gcs_client(self, gcs_channel: grpc.aio.Channel):
+        self._gcs_actor_info_stub = gcs_service_pb2_grpc.ActorInfoGcsServiceStub(
+            gcs_channel
+        )
+        self._gcs_pg_info_stub = gcs_service_pb2_grpc.PlacementGroupInfoGcsServiceStub(
+            gcs_channel
+        )
+        self._gcs_worker_info_stub = gcs_service_pb2_grpc.WorkerInfoGcsServiceStub(
+            gcs_channel
+        )
+        self._gcs_task_info_stub = gcs_service_pb2_grpc.TaskInfoGcsServiceStub(
+            gcs_channel
+        )
+
+    def get_raylet_stub(self, ip: str, port: int):
+        from ray._private.grpc_utils import init_grpc_channel
+
+        options = _STATE_MANAGER_GRPC_OPTIONS
+        channel = init_grpc_channel(build_address(ip, port), options, asynchronous=True)
+        return NodeManagerServiceStub(channel)
+
+    async def get_log_service_stub(self, node_id: NodeID) -> LogServiceStub:
+        """Returns None if the agent on the node is not registered in Internal KV."""
+        from ray._private.grpc_utils import init_grpc_channel
+
+        agent_addr = await self._gcs_client.async_internal_kv_get(
+            f"{dashboard_consts.DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX}{node_id.hex()}".encode(),
+            namespace=ray_constants.KV_NAMESPACE_DASHBOARD,
+            timeout=dashboard_consts.GCS_RPC_TIMEOUT_SECONDS,
+        )
+        if not agent_addr:
+            return None
+        ip, http_port, grpc_port = json.loads(agent_addr)
+        options = ray_constants.GLOBAL_GRPC_OPTIONS
+        channel = init_grpc_channel(
+            build_address(ip, grpc_port), options=options, asynchronous=True
+        )
+        return LogServiceStub(channel)
+
+    async def ip_to_node_id(self, ip: Optional[str]) -> Optional[str]:
+        """Return the node id in hex that corresponds to the given ip.
+
+        Args:
+            ip: The ip address.
+
+        Returns:
+            None if the corresponding id doesn't exist.
+            Node id otherwise. If None node_ip is given,
+            it will also return None.
+        """
+        if not ip:
+            return None
+        # Uses the dashboard agent keys to find ip -> id mapping.
+        agent_addr = await self._gcs_client.async_internal_kv_get(
+            f"{dashboard_consts.DASHBOARD_AGENT_ADDR_IP_PREFIX}{ip}".encode(),
+            namespace=ray_constants.KV_NAMESPACE_DASHBOARD,
+            timeout=dashboard_consts.GCS_RPC_TIMEOUT_SECONDS,
+        )
+        if not agent_addr:
+            return None
+        node_id, http_port, grpc_port = json.loads(agent_addr)
+        return node_id
+
+    @handle_grpc_network_errors
+    async def get_all_actor_info(
+        self,
+        timeout: int = None,
+        limit: int = RAY_MAX_LIMIT_FROM_DATA_SOURCE,
+        filters: Optional[List[Tuple[str, PredicateType, SupportedFilterType]]] = None,
+    ) -> Optional[GetAllActorInfoReply]:
+        if filters is None:
+            filters = []
+
+        req_filters = GetAllActorInfoRequest.Filters()
+        for filter in filters:
+            key, predicate, value = filter
+            if predicate != "=":
+                # We only support EQUAL predicate for source side filtering.
+                continue
+            if key == "actor_id":
+                req_filters.actor_id = ActorID(hex_to_binary(value)).binary()
+            elif key == "state":
+                # Convert to uppercase.
+                value = value.upper()
+                if value not in ActorTableData.ActorState.keys():
+                    raise ValueError(f"Invalid actor state for filtering: {value}")
+                req_filters.state = ActorTableData.ActorState.Value(value)
+            elif key == "job_id":
+                req_filters.job_id = JobID(hex_to_binary(value)).binary()
+
+        request = GetAllActorInfoRequest(limit=limit, filters=req_filters)
+        reply = await self._gcs_actor_info_stub.GetAllActorInfo(
+            request, timeout=timeout
+        )
+        return reply
+
+    @handle_grpc_network_errors
+    async def get_all_task_info(
+        self,
+        timeout: int = None,
+        limit: int = RAY_MAX_LIMIT_FROM_DATA_SOURCE,
+        filters: Optional[List[Tuple[str, PredicateType, SupportedFilterType]]] = None,
+        exclude_driver: bool = False,
+    ) -> Optional[GetTaskEventsReply]:
+
+        if filters is None:
+            filters = []
+
+        req_filters = GetTaskEventsRequest.Filters()
+        for filter in filters:
+            key, predicate, value = filter
+            filter_predicate = None
+            if predicate == "=":
+                filter_predicate = FilterPredicate.EQUAL
+            elif predicate == "!=":
+                filter_predicate = FilterPredicate.NOT_EQUAL
+            else:
+                # We only support EQUAL and NOT_EQUAL predicate for source side
+                # filtering. If invalid predicates were specified, it should already be
+                # raised when the filters arguments are parsed
+                assert False, "Invalid predicate: " + predicate
+
+            if key == "actor_id":
+                actor_filter = GetTaskEventsRequest.Filters.ActorIdFilter()
+                actor_filter.actor_id = ActorID(hex_to_binary(value)).binary()
+                actor_filter.predicate = filter_predicate
+                req_filters.actor_filters.append(actor_filter)
+
+            elif key == "job_id":
+                job_filter = GetTaskEventsRequest.Filters.JobIdFilter()
+                job_filter.job_id = JobID(hex_to_binary(value)).binary()
+                job_filter.predicate = filter_predicate
+                req_filters.job_filters.append(job_filter)
+
+            elif key == "task_id":
+                task_filter = GetTaskEventsRequest.Filters.TaskIdFilter()
+                task_filter.task_id = TaskID(hex_to_binary(value)).binary()
+                task_filter.predicate = filter_predicate
+                req_filters.task_filters.append(task_filter)
+
+            elif key == "name":
+                task_name_filter = GetTaskEventsRequest.Filters.TaskNameFilter()
+                task_name_filter.task_name = value
+                task_name_filter.predicate = filter_predicate
+                req_filters.task_name_filters.append(task_name_filter)
+
+            elif key == "state":
+                state_filter = GetTaskEventsRequest.Filters.StateFilter()
+                state_filter.state = value
+                state_filter.predicate = filter_predicate
+                req_filters.state_filters.append(state_filter)
+
+            else:
+                continue
+
+        req_filters.exclude_driver = exclude_driver
+
+        request = GetTaskEventsRequest(limit=limit, filters=req_filters)
+        reply = await self._gcs_task_info_stub.GetTaskEvents(request, timeout=timeout)
+        return reply
+
+    @handle_grpc_network_errors
+    async def get_all_placement_group_info(
+        self, timeout: int = None, limit: int = RAY_MAX_LIMIT_FROM_DATA_SOURCE
+    ) -> Optional[GetAllPlacementGroupReply]:
+
+        request = GetAllPlacementGroupRequest(limit=limit)
+        reply = await self._gcs_pg_info_stub.GetAllPlacementGroup(
+            request, timeout=timeout
+        )
+        return reply
+
+    async def get_all_node_info(
+        self,
+        timeout: int = None,
+        limit: int = RAY_MAX_LIMIT_FROM_DATA_SOURCE,
+        filters: Optional[List[Tuple[str, PredicateType, SupportedFilterType]]] = None,
+    ) -> Optional[Tuple[Dict[NodeID, GcsNodeInfo], int]]:
+        """Returns node info and the number of filtered nodes.
+
+        Args:
+            timeout: Timeout in seconds
+            limit: Maximum number of nodes to return
+            filters: List of (key, predicate, value) tuples. Supported keys:
+                - "node_id": Filter by node ID (hex string)
+                - "state": Filter by node state (string, e.g., "ALIVE")
+                - "node_name": Filter by node name
+
+        Returns:
+            A tuple of (node_infos, num_filtered) where:
+                - node_infos: Dict[NodeID, GcsNodeInfo] mapping node IDs to their info
+                - num_filtered: Number of nodes filtered out by the query
+        """
+        if filters is None:
+            filters = []
+
+        node_selectors = []
+        state_filter = None
+        for filter in filters:
+            key, predicate, value = filter
+            if predicate != "=":
+                # We only support EQUAL predicate for source side filtering.
+                continue
+
+            if key == "node_id":
+                node_selector = GetAllNodeInfoRequest.NodeSelector()
+                node_selector.node_id = NodeID(hex_to_binary(value)).binary()
+                node_selectors.append(node_selector)
+            elif key == "state":
+                value = value.upper()
+                if value not in GcsNodeInfo.GcsNodeState.keys():
+                    raise ValueError(f"Invalid node state for filtering: {value}")
+                state_filter = GcsNodeInfo.GcsNodeState.Value(value)
+            elif key == "node_name":
+                node_selector = GetAllNodeInfoRequest.NodeSelector()
+                node_selector.node_name = value
+                node_selectors.append(node_selector)
+            else:
+                continue
+
+        reply = await self._gcs_client.async_get_all_node_info(
+            timeout=timeout,
+            node_selectors=node_selectors,
+            state_filter=state_filter,
+            limit=limit,
+        )
+        return reply
+
+    @handle_grpc_network_errors
+    async def get_all_worker_info(
+        self,
+        timeout: int = None,
+        limit: int = RAY_MAX_LIMIT_FROM_DATA_SOURCE,
+        filters: Optional[List[Tuple[str, PredicateType, SupportedFilterType]]] = None,
+    ) -> Optional[GetAllWorkerInfoReply]:
+
+        if filters is None:
+            filters = []
+
+        req_filters = GetAllWorkerInfoRequest.Filters()
+        for filter in filters:
+            key, predicate, value = filter
+            # Special treatments for the Ray Debugger.
+            if (
+                key == "num_paused_threads"
+                and predicate in ("!=", ">")
+                and value == "0"
+            ):
+                req_filters.exist_paused_threads = True
+                continue
+            if key == "is_alive" and predicate == "=" and value == "True":
+                req_filters.is_alive = True
+                continue
+            else:
+                continue
+
+        request = GetAllWorkerInfoRequest(limit=limit, filters=req_filters)
+        reply = await self._gcs_worker_info_stub.GetAllWorkerInfo(
+            request, timeout=timeout
+        )
+        return reply
+
+    # TODO(rickyx):
+    # This is currently mirroring dashboard/modules/job/job_head.py::list_jobs
+    # We should eventually unify the logic.
+    async def get_job_info(self, timeout: int = None) -> List[JobDetails]:
+        # Cannot use @handle_grpc_network_errors because async def is not supported yet.
+
+        driver_jobs, submission_job_drivers = await get_driver_jobs(
+            self._gcs_client, timeout=timeout
+        )
+        submission_jobs = await self._job_client.get_all_jobs(timeout=timeout)
+        submission_jobs = [
+            JobDetails(
+                **dataclasses.asdict(job),
+                submission_id=submission_id,
+                job_id=submission_job_drivers.get(submission_id).id
+                if submission_id in submission_job_drivers
+                else None,
+                driver_info=submission_job_drivers.get(submission_id),
+                type=JobType.SUBMISSION,
+            )
+            for submission_id, job in submission_jobs.items()
+        ]
+
+        return list(driver_jobs.values()) + submission_jobs
+
+    @handle_grpc_network_errors
+    async def get_object_info(
+        self,
+        node_manager_ip: str,
+        node_manager_port: int,
+        timeout: int = None,
+        limit: int = RAY_MAX_LIMIT_FROM_DATA_SOURCE,
+    ) -> Optional[GetObjectsInfoReply]:
+        stub = self.get_raylet_stub(node_manager_ip, node_manager_port)
+
+        reply = await stub.GetObjectsInfo(
+            GetObjectsInfoRequest(limit=limit),
+            timeout=timeout,
+        )
+        return reply
+
+    async def get_runtime_envs_info(
+        self,
+        node_ip: str,
+        runtime_env_agent_port: int,
+        timeout: int = None,
+        limit: int = RAY_MAX_LIMIT_FROM_DATA_SOURCE,
+    ) -> Optional[GetRuntimeEnvsInfoReply]:
+        if not node_ip or not runtime_env_agent_port:
+            raise ValueError(
+                f"Expected non empty node ip and runtime env agent port, got {node_ip} and {runtime_env_agent_port}."
+            )
+        timeout = aiohttp.ClientTimeout(total=timeout)
+        url = f"http://{build_address(node_ip, runtime_env_agent_port)}/get_runtime_envs_info"
+        request = GetRuntimeEnvsInfoRequest(limit=limit)
+        data = request.SerializeToString()
+        headers = get_auth_headers_if_auth_enabled({})
+        async with self._client_session.post(
+            url, data=data, timeout=timeout, headers=headers
+        ) as resp:
+            if resp.status >= 200 and resp.status < 300:
+                response_data = await resp.read()
+                reply = GetRuntimeEnvsInfoReply()
+                reply.ParseFromString(response_data)
+                return reply
+            else:
+                raise DataSourceUnavailable(
+                    "Failed to query the runtime env agent for get_runtime_envs_info. "
+                    "Either there's a network issue, or the source is down. "
+                    f"Response is {resp.status}, reason {resp.reason}"
+                )
+
+    @handle_grpc_network_errors
+    async def list_logs(
+        self, node_id: str, glob_filter: str, timeout: int = None
+    ) -> ListLogsReply:
+        stub = await self.get_log_service_stub(NodeID.from_hex(node_id))
+        if not stub:
+            raise ValueError(f"Agent for node id: {node_id} doesn't exist.")
+        return await stub.ListLogs(
+            ListLogsRequest(glob_filter=glob_filter), timeout=timeout
+        )
+
+    @handle_grpc_network_errors
+    async def stream_log(
+        self,
+        node_id: str,
+        log_file_name: str,
+        keep_alive: bool,
+        lines: int,
+        interval: Optional[float],
+        timeout: int,
+        start_offset: Optional[int] = None,
+        end_offset: Optional[int] = None,
+    ) -> UnaryStreamCall:
+        stub = await self.get_log_service_stub(NodeID.from_hex(node_id))
+        if not stub:
+            raise ValueError(f"Agent for node id: {node_id} doesn't exist.")
+
+        stream = stub.StreamLog(
+            StreamLogRequest(
+                keep_alive=keep_alive,
+                log_file_name=log_file_name,
+                lines=lines,
+                interval=interval,
+                start_offset=start_offset,
+                end_offset=end_offset,
+            ),
+            timeout=timeout,
+        )
+        metadata = await stream.initial_metadata()
+        if metadata.get(log_consts.LOG_GRPC_ERROR) is not None:
+            raise ValueError(metadata.get(log_consts.LOG_GRPC_ERROR))
+        return stream

python/ray/util/state/util.py

@@ -0,0 +1,62 @@
+from typing import Optional, Union
+
+
+def convert_string_to_type(
+    val: Optional[Union[str, int, float, bool]], convert_type: Union[int, float, bool]
+) -> Union[int, float, bool]:
+    """Convert the given value to a convert type.
+
+    If the given val is None, it will just return None without the conversion.
+
+    It supports,
+        str -> int/float/bool
+        int -> int
+        bool -> bool
+        float -> float
+    """
+    if val is None:
+        return None
+    elif type(val) is convert_type:
+        return val
+    elif convert_type is int:
+        try:
+            val = int(val)
+        except ValueError:
+            raise ValueError(
+                f"Failed to convert a value {val} of type {type(val)} to {convert_type}"
+            )
+    elif convert_type is float:
+        try:
+            val = float(val)
+        except ValueError:
+            raise ValueError(
+                f"Failed to convert a value {val} of type {type(val)} to {convert_type}"
+            )
+    elif convert_type is bool:
+        # Without this, "False" will become True.
+        if val == "False" or val == "false" or val == "0":
+            val = False
+        elif val == "True" or val == "true" or val == "1":
+            val = True
+        else:
+            raise ValueError(
+                f"Failed to convert a value {val} of type {type(val)} to {convert_type}"
+            )
+    else:
+        assert False, f"Unsupported convert type {convert_type}"
+    return val
+
+
+def record_deprecated_state_api_import():
+    import warnings
+
+    from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
+
+    warnings.warn(
+        "Ray state API is no longer experimental. Please import from `ray.util.state`. "
+        "instead. Importing from `ray.experimental` will be deprecated in "
+        "future releases. ",
+        DeprecationWarning,
+    )
+
+    record_extra_usage_tag(TagKey.EXPERIMENTAL_STATE_API_IMPORT, "1")