/python/ray/dashboard/state_api_utils.py

import dataclasses
from dataclasses import asdict, fields
from typing import Awaitable, Callable, List, Tuple

import aiohttp.web

from ray.dashboard.optional_utils import rest_response
from ray.dashboard.utils import HTTPStatusCode
from ray.util.state.common import (
    DEFAULT_LIMIT,
    DEFAULT_RPC_TIMEOUT,
    RAY_MAX_LIMIT_FROM_API_SERVER,
    ListApiOptions,
    ListApiResponse,
    PredicateType,
    StateSchema,
    SummaryApiOptions,
    SummaryApiResponse,
    SupportedFilterType,
    filter_fields,
)
from ray.util.state.exception import DataSourceUnavailable
from ray.util.state.util import convert_string_to_type


def do_reply(
    status_code: HTTPStatusCode, error_message: str, result: ListApiResponse, **kwargs
):
    return rest_response(
        status_code=status_code,
        message=error_message,
        result=result,
        convert_google_style=False,
        **kwargs,
    )


async def handle_list_api(
    list_api_fn: Callable[[ListApiOptions], Awaitable[ListApiResponse]],
    req: aiohttp.web.Request,
):
    try:
        result = await list_api_fn(option=options_from_req(req))
        return do_reply(
            status_code=HTTPStatusCode.OK,
            error_message="",
            result=asdict(result),
        )
    except ValueError as e:
        return do_reply(
            status_code=HTTPStatusCode.BAD_REQUEST,
            error_message=str(e),
            result=None,
        )
    except DataSourceUnavailable as e:
        return do_reply(
            status_code=HTTPStatusCode.INTERNAL_ERROR,
            error_message=str(e),
            result=None,
        )


def _get_filters_from_req(
    req: aiohttp.web.Request,
) -> List[Tuple[str, PredicateType, SupportedFilterType]]:
    filter_keys = req.query.getall("filter_keys", [])
    filter_predicates = req.query.getall("filter_predicates", [])
    filter_values = req.query.getall("filter_values", [])
    assert len(filter_keys) == len(filter_values)
    filters = []
    for key, predicate, val in zip(filter_keys, filter_predicates, filter_values):
        filters.append((key, predicate, val))
    return filters


def options_from_req(req: aiohttp.web.Request) -> ListApiOptions:
    """Obtain `ListApiOptions` from the aiohttp request."""
    limit = int(
        req.query.get("limit") if req.query.get("limit") is not None else DEFAULT_LIMIT
    )

    if limit > RAY_MAX_LIMIT_FROM_API_SERVER:
        raise ValueError(
            f"Given limit {limit} exceeds the supported "
            f"limit {RAY_MAX_LIMIT_FROM_API_SERVER}. Use a lower limit, or set the "
            f"`RAY_MAX_LIMIT_FROM_API_SERVER` environment variable to a larger value."
        )

    timeout = int(req.query.get("timeout", 30))
    filters = _get_filters_from_req(req)
    detail = convert_string_to_type(req.query.get("detail", False), bool)
    exclude_driver = convert_string_to_type(req.query.get("exclude_driver", True), bool)

    return ListApiOptions(
        limit=limit,
        timeout=timeout,
        filters=filters,
        detail=detail,
        exclude_driver=exclude_driver,
    )


def summary_options_from_req(req: aiohttp.web.Request) -> SummaryApiOptions:
    timeout = int(req.query.get("timeout", DEFAULT_RPC_TIMEOUT))
    filters = _get_filters_from_req(req)
    summary_by = req.query.get("summary_by", None)
    return SummaryApiOptions(timeout=timeout, filters=filters, summary_by=summary_by)


async def handle_summary_api(
    summary_fn: Callable[[SummaryApiOptions], SummaryApiResponse],
    req: aiohttp.web.Request,
):
    result = await summary_fn(option=summary_options_from_req(req))
    return do_reply(
        status_code=HTTPStatusCode.OK,
        error_message="",
        result=asdict(result),
    )


def convert_filters_type(
    filter: List[Tuple[str, PredicateType, SupportedFilterType]],
    schema: StateSchema,
) -> List[Tuple[str, PredicateType, SupportedFilterType]]:
    """Convert the given filter's type to SupportedFilterType.

    This method is necessary because click can only accept a single type
    for its tuple (which is string in this case).

    Args:
        filter: A list of filter which is a tuple of (key, val).
        schema: The state schema. It is used to infer the type of the column for filter.

    Returns:
        A new list of filters with correct types that match the schema.
    """
    new_filter = []
    if dataclasses.is_dataclass(schema):
        schema = {field.name: field.type for field in fields(schema)}
    else:
        schema = schema.schema_dict()

    for col, predicate, val in filter:
        if col in schema:
            column_type = schema[col]
            try:
                isinstance(val, column_type)
            except TypeError:
                # Calling `isinstance` to the Literal type raises a TypeError.
                # Ignore this case.
                pass
            else:
                if isinstance(val, column_type):
                    # Do nothing.
                    pass
                elif column_type is int or column_type == "integer":
                    try:
                        val = convert_string_to_type(val, int)
                    except ValueError:
                        raise ValueError(
                            f"Invalid filter `--filter {col} {val}` for a int type "
                            "column. Please provide an integer filter "
                            f"`--filter {col} [int]`"
                        )
                elif column_type is float or column_type == "number":
                    try:
                        val = convert_string_to_type(
                            val,
                            float,
                        )
                    except ValueError:
                        raise ValueError(
                            f"Invalid filter `--filter {col} {val}` for a float "
                            "type column. Please provide an integer filter "
                            f"`--filter {col} [float]`"
                        )
                elif column_type is bool or column_type == "boolean":
                    try:
                        val = convert_string_to_type(val, bool)
                    except ValueError:
                        raise ValueError(
                            f"Invalid filter `--filter {col} {val}` for a boolean "
                            "type column. Please provide "
                            f"`--filter {col} [True|true|1]` for True or "
                            f"`--filter {col} [False|false|0]` for False."
                        )
        new_filter.append((col, predicate, val))
    return new_filter


def do_filter(
    data: List[dict],
    filters: List[Tuple[str, PredicateType, SupportedFilterType]],
    state_dataclass: StateSchema,
    detail: bool,
) -> List[dict]:
    """Return the filtered data given filters.

    Args:
        data: A list of state data.
        filters: A list of KV tuple to filter data (key, val). The data is filtered
            if data[key] != val.
        state_dataclass: The state schema.
        detail: If True, include detail-only columns; otherwise drop them.

    Returns:
        A list of filtered state data in dictionary. Each state data's
        unnecessary columns are filtered by the given state_dataclass schema.
    """
    filters = convert_filters_type(filters, state_dataclass)
    result = []
    for datum in data:
        match = True
        for filter_column, filter_predicate, filter_value in filters:
            filterable_columns = state_dataclass.filterable_columns()
            filter_column = filter_column.lower()
            if filter_column not in filterable_columns:
                raise ValueError(
                    f"The given filter column {filter_column} is not supported. "
                    "Enter filters with –-filter key=value "
                    "or –-filter key!=value "
                    f"Supported filter columns: {filterable_columns}"
                )

            if filter_column not in datum:
                match = False
            elif filter_predicate == "=":
                if isinstance(filter_value, str) and isinstance(
                    datum[filter_column], str
                ):
                    # Case insensitive match for string filter values.
                    match = datum[filter_column].lower() == filter_value.lower()
                elif isinstance(filter_value, str) and isinstance(
                    datum[filter_column], bool
                ):
                    match = datum[filter_column] == convert_string_to_type(
                        filter_value, bool
                    )
                elif isinstance(filter_value, str) and isinstance(
                    datum[filter_column], int
                ):
                    match = datum[filter_column] == convert_string_to_type(
                        filter_value, int
                    )
                else:
                    match = datum[filter_column] == filter_value
            elif filter_predicate == "!=":
                if isinstance(filter_value, str) and isinstance(
                    datum[filter_column], str
                ):
                    match = datum[filter_column].lower() != filter_value.lower()
                else:
                    match = datum[filter_column] != filter_value
            else:
                raise ValueError(
                    f"Unsupported filter predicate {filter_predicate} is given. "
                    "Available predicates: =, !=."
                )

            if not match:
                break

        if match:
            result.append(filter_fields(datum, state_dataclass, detail))
    return result