wehub-resource-sync
267 lines
9.5 KiB
Python
267 lines
9.5 KiB
Python
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
|