1783 lines
72 KiB
Python
1783 lines
72 KiB
Python
import ast
|
|
import base64
|
|
import json
|
|
import math
|
|
import operator
|
|
import re
|
|
import shlex
|
|
from typing import Any
|
|
|
|
import sqlparse
|
|
from packaging.version import Version
|
|
from sqlparse.sql import (
|
|
Comparison,
|
|
Identifier,
|
|
IdentifierList,
|
|
Parenthesis,
|
|
Statement,
|
|
Token,
|
|
TokenList,
|
|
)
|
|
from sqlparse.tokens import Token as TokenType
|
|
|
|
from mlflow.entities import RunInfo
|
|
from mlflow.entities.model_registry.model_version_stages import STAGE_DELETED_INTERNAL
|
|
from mlflow.entities.model_registry.prompt import IS_PROMPT_TAG_KEY
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.store.db.db_types import MSSQL, MYSQL, POSTGRES, SQLITE
|
|
from mlflow.tracing.constant import TraceMetadataKey, TraceTagKey
|
|
from mlflow.utils.mlflow_tags import (
|
|
MLFLOW_DATASET_CONTEXT,
|
|
)
|
|
|
|
|
|
def _convert_like_pattern_to_regex(pattern, flags=0):
|
|
if not pattern.startswith("%"):
|
|
pattern = "^" + pattern
|
|
if not pattern.endswith("%"):
|
|
pattern = pattern + "$"
|
|
return re.compile(pattern.replace("_", ".").replace("%", ".*"), flags)
|
|
|
|
|
|
def _like(string, pattern):
|
|
return _convert_like_pattern_to_regex(pattern).match(string) is not None
|
|
|
|
|
|
def _ilike(string, pattern):
|
|
return _convert_like_pattern_to_regex(pattern, flags=re.IGNORECASE).match(string) is not None
|
|
|
|
|
|
def _join_in_comparison_tokens(tokens, search_traces=False):
|
|
"""
|
|
Find a sequence of tokens that matches the pattern of an IN comparison or a NOT IN comparison,
|
|
join the tokens into a single Comparison token. Otherwise, return the original list of tokens.
|
|
"""
|
|
if Version(sqlparse.__version__) < Version("0.4.4"):
|
|
# In sqlparse < 0.4.4, IN is treated as a comparison, we don't need to join tokens
|
|
return tokens
|
|
|
|
non_whitespace_tokens = [t for t in tokens if not t.is_whitespace]
|
|
joined_tokens = []
|
|
num_tokens = len(non_whitespace_tokens)
|
|
iterator = enumerate(non_whitespace_tokens)
|
|
while elem := next(iterator, None):
|
|
index, first = elem
|
|
# We need at least 3 tokens to form an IN comparison or a NOT IN comparison
|
|
if num_tokens - index < 3:
|
|
joined_tokens.extend(non_whitespace_tokens[index:])
|
|
break
|
|
|
|
if search_traces:
|
|
# timestamp
|
|
if first.match(ttype=TokenType.Name.Builtin, values=["timestamp", "timestamp_ms"]):
|
|
(_, second) = next(iterator, (None, None))
|
|
(_, third) = next(iterator, (None, None))
|
|
if any(x is None for x in [second, third]):
|
|
raise MlflowException(
|
|
f"Invalid comparison clause with token `{first}, {second}, {third}`, "
|
|
"expected 3 tokens",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if (
|
|
second.match(
|
|
ttype=TokenType.Operator.Comparison,
|
|
values=SearchTraceUtils.VALID_NUMERIC_ATTRIBUTE_COMPARATORS,
|
|
)
|
|
and third.ttype == TokenType.Literal.Number.Integer
|
|
):
|
|
joined_tokens.append(Comparison(TokenList([first, second, third])))
|
|
continue
|
|
else:
|
|
joined_tokens.extend([first, second, third])
|
|
|
|
# Wait until we encounter an identifier token
|
|
if not isinstance(first, Identifier):
|
|
joined_tokens.append(first)
|
|
continue
|
|
|
|
(_, second) = next(iterator)
|
|
(_, third) = next(iterator)
|
|
|
|
# IN
|
|
if (
|
|
isinstance(first, Identifier)
|
|
and second.match(ttype=TokenType.Keyword, values=["IN"])
|
|
and isinstance(third, Parenthesis)
|
|
):
|
|
joined_tokens.append(Comparison(TokenList([first, second, third])))
|
|
continue
|
|
|
|
(_, fourth) = next(iterator, (None, None))
|
|
if fourth is None:
|
|
joined_tokens.extend([first, second, third])
|
|
break
|
|
|
|
# NOT IN
|
|
if (
|
|
isinstance(first, Identifier)
|
|
and second.match(ttype=TokenType.Keyword, values=["NOT"])
|
|
and third.match(ttype=TokenType.Keyword, values=["IN"])
|
|
and isinstance(fourth, Parenthesis)
|
|
):
|
|
joined_tokens.append(
|
|
Comparison(TokenList([first, Token(TokenType.Keyword, "NOT IN"), fourth]))
|
|
)
|
|
continue
|
|
|
|
joined_tokens.extend([first, second, third, fourth])
|
|
|
|
return joined_tokens
|
|
|
|
|
|
class SearchUtils:
|
|
LIKE_OPERATOR = "LIKE"
|
|
ILIKE_OPERATOR = "ILIKE"
|
|
ASC_OPERATOR = "asc"
|
|
DESC_OPERATOR = "desc"
|
|
VALID_ORDER_BY_TAGS = [ASC_OPERATOR, DESC_OPERATOR]
|
|
VALID_METRIC_COMPARATORS = {">", ">=", "!=", "=", "<", "<="}
|
|
VALID_PARAM_COMPARATORS = {"!=", "=", LIKE_OPERATOR, ILIKE_OPERATOR}
|
|
VALID_TAG_COMPARATORS = {"!=", "=", LIKE_OPERATOR, ILIKE_OPERATOR}
|
|
VALID_STRING_ATTRIBUTE_COMPARATORS = {"!=", "=", LIKE_OPERATOR, ILIKE_OPERATOR, "IN", "NOT IN"}
|
|
VALID_NUMERIC_ATTRIBUTE_COMPARATORS = VALID_METRIC_COMPARATORS
|
|
VALID_DATASET_COMPARATORS = {"!=", "=", LIKE_OPERATOR, ILIKE_OPERATOR, "IN", "NOT IN"}
|
|
_BUILTIN_NUMERIC_ATTRIBUTES = {"start_time", "end_time"}
|
|
_ALTERNATE_NUMERIC_ATTRIBUTES = {"created", "Created"}
|
|
_ALTERNATE_STRING_ATTRIBUTES = {"run name", "Run name", "Run Name"}
|
|
NUMERIC_ATTRIBUTES = set(
|
|
list(_BUILTIN_NUMERIC_ATTRIBUTES) + list(_ALTERNATE_NUMERIC_ATTRIBUTES)
|
|
)
|
|
DATASET_ATTRIBUTES = {"name", "digest", "context"}
|
|
VALID_SEARCH_ATTRIBUTE_KEYS = set(
|
|
RunInfo.get_searchable_attributes()
|
|
+ list(_ALTERNATE_NUMERIC_ATTRIBUTES)
|
|
+ list(_ALTERNATE_STRING_ATTRIBUTES)
|
|
)
|
|
VALID_ORDER_BY_ATTRIBUTE_KEYS = set(
|
|
RunInfo.get_orderable_attributes() + list(_ALTERNATE_NUMERIC_ATTRIBUTES)
|
|
)
|
|
_METRIC_IDENTIFIER = "metric"
|
|
_ALTERNATE_METRIC_IDENTIFIERS = {"metrics"}
|
|
_PARAM_IDENTIFIER = "parameter"
|
|
_ALTERNATE_PARAM_IDENTIFIERS = {"parameters", "param", "params"}
|
|
_TAG_IDENTIFIER = "tag"
|
|
_ALTERNATE_TAG_IDENTIFIERS = {"tags"}
|
|
_ATTRIBUTE_IDENTIFIER = "attribute"
|
|
_ALTERNATE_ATTRIBUTE_IDENTIFIERS = {"attr", "attributes", "run"}
|
|
_DATASET_IDENTIFIER = "dataset"
|
|
_ALTERNATE_DATASET_IDENTIFIERS = {"datasets"}
|
|
_IDENTIFIERS = [
|
|
_METRIC_IDENTIFIER,
|
|
_PARAM_IDENTIFIER,
|
|
_TAG_IDENTIFIER,
|
|
_ATTRIBUTE_IDENTIFIER,
|
|
_DATASET_IDENTIFIER,
|
|
]
|
|
_VALID_IDENTIFIERS = set(
|
|
_IDENTIFIERS
|
|
+ list(_ALTERNATE_METRIC_IDENTIFIERS)
|
|
+ list(_ALTERNATE_PARAM_IDENTIFIERS)
|
|
+ list(_ALTERNATE_TAG_IDENTIFIERS)
|
|
+ list(_ALTERNATE_ATTRIBUTE_IDENTIFIERS)
|
|
+ list(_ALTERNATE_DATASET_IDENTIFIERS)
|
|
)
|
|
STRING_VALUE_TYPES = {TokenType.Literal.String.Single}
|
|
DELIMITER_VALUE_TYPES = {TokenType.Punctuation}
|
|
WHITESPACE_VALUE_TYPE = TokenType.Text.Whitespace
|
|
NUMERIC_VALUE_TYPES = {TokenType.Literal.Number.Integer, TokenType.Literal.Number.Float}
|
|
# Registered Models Constants
|
|
ORDER_BY_KEY_TIMESTAMP = "timestamp"
|
|
ORDER_BY_KEY_LAST_UPDATED_TIMESTAMP = "last_updated_timestamp"
|
|
ORDER_BY_KEY_MODEL_NAME = "name"
|
|
VALID_ORDER_BY_KEYS_REGISTERED_MODELS = {
|
|
ORDER_BY_KEY_TIMESTAMP,
|
|
ORDER_BY_KEY_LAST_UPDATED_TIMESTAMP,
|
|
ORDER_BY_KEY_MODEL_NAME,
|
|
}
|
|
VALID_TIMESTAMP_ORDER_BY_KEYS = {ORDER_BY_KEY_TIMESTAMP, ORDER_BY_KEY_LAST_UPDATED_TIMESTAMP}
|
|
# We encourage users to use timestamp for order-by
|
|
RECOMMENDED_ORDER_BY_KEYS_REGISTERED_MODELS = {ORDER_BY_KEY_MODEL_NAME, ORDER_BY_KEY_TIMESTAMP}
|
|
|
|
@staticmethod
|
|
def get_comparison_func(comparator):
|
|
return {
|
|
">": operator.gt,
|
|
">=": operator.ge,
|
|
"=": operator.eq,
|
|
"!=": operator.ne,
|
|
"<=": operator.le,
|
|
"<": operator.lt,
|
|
"LIKE": _like,
|
|
"ILIKE": _ilike,
|
|
"IN": lambda x, y: x in y,
|
|
"NOT IN": lambda x, y: x not in y,
|
|
}[comparator]
|
|
|
|
@staticmethod
|
|
def get_sql_comparison_func(comparator, dialect):
|
|
import sqlalchemy as sa
|
|
|
|
def comparison_func(column, value):
|
|
if comparator == "LIKE":
|
|
return column.like(value)
|
|
elif comparator == "ILIKE":
|
|
return column.ilike(value)
|
|
elif comparator == "IN":
|
|
return column.in_(value)
|
|
elif comparator == "NOT IN":
|
|
return ~column.in_(value)
|
|
return SearchUtils.get_comparison_func(comparator)(column, value)
|
|
|
|
def mssql_comparison_func(column, value):
|
|
if not isinstance(column.type, sa.types.String):
|
|
return comparison_func(column, value)
|
|
|
|
collated = column.collate("Japanese_Bushu_Kakusu_100_CS_AS_KS_WS")
|
|
return comparison_func(collated, value)
|
|
|
|
def mysql_comparison_func(column, value):
|
|
if not isinstance(column.type, sa.types.String):
|
|
return comparison_func(column, value)
|
|
|
|
# MySQL is case insensitive by default, so we need to use the binary operator to
|
|
# perform case sensitive comparisons.
|
|
templates = {
|
|
# Use non-binary ahead of binary comparison for runtime performance
|
|
"=": "({column} = :value AND BINARY {column} = :value)",
|
|
"!=": "({column} != :value OR BINARY {column} != :value)",
|
|
"LIKE": "({column} LIKE :value AND BINARY {column} LIKE :value)",
|
|
}
|
|
if comparator in templates:
|
|
column = f"{column.class_.__tablename__}.{column.key}"
|
|
return sa.text(templates[comparator].format(column=column)).bindparams(
|
|
sa.bindparam("value", value=value, unique=True)
|
|
)
|
|
|
|
return comparison_func(column, value)
|
|
|
|
return {
|
|
POSTGRES: comparison_func,
|
|
SQLITE: comparison_func,
|
|
MSSQL: mssql_comparison_func,
|
|
MYSQL: mysql_comparison_func,
|
|
}[dialect]
|
|
|
|
@staticmethod
|
|
def translate_key_alias(key):
|
|
if key in ["created", "Created"]:
|
|
return "start_time"
|
|
if key in ["run name", "Run name", "Run Name"]:
|
|
return "run_name"
|
|
return key
|
|
|
|
@classmethod
|
|
def _trim_ends(cls, string_value):
|
|
return string_value[1:-1]
|
|
|
|
@classmethod
|
|
def _is_quoted(cls, value, pattern):
|
|
return len(value) >= 2 and value.startswith(pattern) and value.endswith(pattern)
|
|
|
|
@classmethod
|
|
def _trim_backticks(cls, entity_type):
|
|
"""Remove backticks from identifier like `param`, if they exist."""
|
|
if cls._is_quoted(entity_type, "`"):
|
|
return cls._trim_ends(entity_type)
|
|
return entity_type
|
|
|
|
@classmethod
|
|
def _strip_quotes(cls, value, expect_quoted_value=False):
|
|
"""
|
|
Remove quotes for input string.
|
|
Values of type strings are expected to have quotes.
|
|
Keys containing special characters are also expected to be enclose in quotes.
|
|
"""
|
|
if cls._is_quoted(value, "'") or cls._is_quoted(value, '"'):
|
|
return cls._trim_ends(value)
|
|
elif expect_quoted_value:
|
|
raise MlflowException(
|
|
"Parameter value is either not quoted or unidentified quote "
|
|
f"types used for string value {value}. Use either single or double "
|
|
"quotes.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
return value
|
|
|
|
@classmethod
|
|
def _valid_entity_type(cls, entity_type):
|
|
entity_type = cls._trim_backticks(entity_type)
|
|
if entity_type not in cls._VALID_IDENTIFIERS:
|
|
raise MlflowException(
|
|
f"Invalid entity type '{entity_type}'. Valid values are {cls._IDENTIFIERS}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
if entity_type in cls._ALTERNATE_PARAM_IDENTIFIERS:
|
|
return cls._PARAM_IDENTIFIER
|
|
elif entity_type in cls._ALTERNATE_METRIC_IDENTIFIERS:
|
|
return cls._METRIC_IDENTIFIER
|
|
elif entity_type in cls._ALTERNATE_TAG_IDENTIFIERS:
|
|
return cls._TAG_IDENTIFIER
|
|
elif entity_type in cls._ALTERNATE_ATTRIBUTE_IDENTIFIERS:
|
|
return cls._ATTRIBUTE_IDENTIFIER
|
|
elif entity_type in cls._ALTERNATE_DATASET_IDENTIFIERS:
|
|
return cls._DATASET_IDENTIFIER
|
|
else:
|
|
# one of ("metric", "parameter", "tag", or "attribute") since it a valid type
|
|
return entity_type
|
|
|
|
@classmethod
|
|
def _get_identifier(cls, identifier, valid_attributes):
|
|
try:
|
|
tokens = identifier.split(".", 1)
|
|
if len(tokens) == 1:
|
|
key = tokens[0]
|
|
entity_type = cls._ATTRIBUTE_IDENTIFIER
|
|
else:
|
|
entity_type, key = tokens
|
|
except ValueError:
|
|
raise MlflowException(
|
|
f"Invalid identifier {identifier!r}. Columns should be specified as "
|
|
"'attribute.<key>', 'metric.<key>', 'tag.<key>', 'dataset.<key>', or "
|
|
"'param.'.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
identifier = cls._valid_entity_type(entity_type)
|
|
key = cls._trim_backticks(cls._strip_quotes(key))
|
|
if identifier == cls._ATTRIBUTE_IDENTIFIER and key not in valid_attributes:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid attribute key '{key}' specified. Valid keys are '{valid_attributes}'"
|
|
)
|
|
elif identifier == cls._DATASET_IDENTIFIER and key not in cls.DATASET_ATTRIBUTES:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid dataset key '{key}' specified. Valid keys are '{cls.DATASET_ATTRIBUTES}'"
|
|
)
|
|
return {"type": identifier, "key": key}
|
|
|
|
@classmethod
|
|
def _get_value(cls, identifier_type, key, token):
|
|
if identifier_type == cls._METRIC_IDENTIFIER:
|
|
if token.ttype not in cls.NUMERIC_VALUE_TYPES:
|
|
raise MlflowException(
|
|
f"Expected numeric value type for metric. Found {token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return token.value
|
|
elif identifier_type == cls._PARAM_IDENTIFIER or identifier_type == cls._TAG_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
raise MlflowException(
|
|
"Expected a quoted string value for "
|
|
f"{identifier_type} (e.g. 'my-value'). Got value "
|
|
f"{token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif identifier_type == cls._ATTRIBUTE_IDENTIFIER:
|
|
if key in cls.NUMERIC_ATTRIBUTES:
|
|
if token.ttype not in cls.NUMERIC_VALUE_TYPES:
|
|
raise MlflowException(
|
|
f"Expected numeric value type for numeric attribute: {key}. "
|
|
f"Found {token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return token.value
|
|
elif token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
elif isinstance(token, Parenthesis):
|
|
if key != "run_id":
|
|
raise MlflowException(
|
|
"Only the 'run_id' attribute supports comparison with a list of quoted "
|
|
"string values.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls._parse_run_ids(token)
|
|
else:
|
|
raise MlflowException(
|
|
f"Expected a quoted string value for attributes. Got value {token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif identifier_type == cls._DATASET_IDENTIFIER:
|
|
if key in cls.DATASET_ATTRIBUTES and (
|
|
token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier)
|
|
):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
elif isinstance(token, Parenthesis):
|
|
if key not in ("name", "digest", "context"):
|
|
raise MlflowException(
|
|
"Only the dataset 'name' and 'digest' supports comparison with a list of "
|
|
"quoted string values.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls._parse_run_ids(token)
|
|
else:
|
|
raise MlflowException(
|
|
"Expected a quoted string value for dataset attributes. "
|
|
f"Got value {token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
# Expected to be either "param" or "metric".
|
|
raise MlflowException(
|
|
"Invalid identifier type. Expected one of "
|
|
f"{[cls._METRIC_IDENTIFIER, cls._PARAM_IDENTIFIER]}."
|
|
)
|
|
|
|
@classmethod
|
|
def _validate_comparison(cls, tokens, search_traces=False):
|
|
base_error_string = "Invalid comparison clause"
|
|
if len(tokens) != 3:
|
|
raise MlflowException(
|
|
f"{base_error_string}. Expected 3 tokens found {len(tokens)}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if not isinstance(tokens[0], Identifier):
|
|
if not search_traces:
|
|
raise MlflowException(
|
|
f"{base_error_string}. Expected 'Identifier' found '{tokens[0]}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if search_traces and not tokens[0].match(
|
|
ttype=TokenType.Name.Builtin, values=["timestamp", "timestamp_ms"]
|
|
):
|
|
raise MlflowException(
|
|
f"{base_error_string}. Expected 'TokenType.Name.Builtin' found '{tokens[0]}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if not isinstance(tokens[1], Token) and tokens[1].ttype != TokenType.Operator.Comparison:
|
|
raise MlflowException(
|
|
f"{base_error_string}. Expected comparison found '{tokens[1]}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if not isinstance(tokens[2], Token) and (
|
|
tokens[2].ttype not in cls.STRING_VALUE_TYPES.union(cls.NUMERIC_VALUE_TYPES)
|
|
or isinstance(tokens[2], Identifier)
|
|
):
|
|
raise MlflowException(
|
|
f"{base_error_string}. Expected value token found '{tokens[2]}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
@classmethod
|
|
def _get_comparison(cls, comparison):
|
|
stripped_comparison = [token for token in comparison.tokens if not token.is_whitespace]
|
|
cls._validate_comparison(stripped_comparison)
|
|
comp = cls._get_identifier(stripped_comparison[0].value, cls.VALID_SEARCH_ATTRIBUTE_KEYS)
|
|
comp["comparator"] = stripped_comparison[1].value
|
|
comp["value"] = cls._get_value(comp.get("type"), comp.get("key"), stripped_comparison[2])
|
|
return comp
|
|
|
|
@classmethod
|
|
def _invalid_statement_token_search_runs(cls, token):
|
|
if (
|
|
isinstance(token, Comparison)
|
|
or token.is_whitespace
|
|
or token.match(ttype=TokenType.Keyword, values=["AND"])
|
|
):
|
|
return False
|
|
return True
|
|
|
|
@classmethod
|
|
def _process_statement(cls, statement):
|
|
# check validity
|
|
tokens = _join_in_comparison_tokens(statement.tokens)
|
|
invalids = list(filter(cls._invalid_statement_token_search_runs, tokens))
|
|
if len(invalids) > 0:
|
|
invalid_clauses = ", ".join(f"'{token}'" for token in invalids)
|
|
raise MlflowException(
|
|
f"Invalid clause(s) in filter string: {invalid_clauses}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return [cls._get_comparison(si) for si in tokens if isinstance(si, Comparison)]
|
|
|
|
@classmethod
|
|
def parse_search_filter(cls, filter_string):
|
|
if not filter_string:
|
|
return []
|
|
try:
|
|
parsed = sqlparse.parse(filter_string)
|
|
except Exception:
|
|
raise MlflowException(
|
|
f"Error on parsing filter '{filter_string}'", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
if len(parsed) == 0 or not isinstance(parsed[0], Statement):
|
|
raise MlflowException(
|
|
f"Invalid filter '{filter_string}'. Could not be parsed.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif len(parsed) > 1:
|
|
raise MlflowException(
|
|
f"Search filter contained multiple expression {filter_string!r}. "
|
|
"Provide AND-ed expression list.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls._process_statement(parsed[0])
|
|
|
|
@classmethod
|
|
def is_metric(cls, key_type, comparator):
|
|
if key_type == cls._METRIC_IDENTIFIER:
|
|
if comparator not in cls.VALID_METRIC_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' not one of '{cls.VALID_METRIC_COMPARATORS}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def is_param(cls, key_type, comparator):
|
|
if key_type == cls._PARAM_IDENTIFIER:
|
|
if comparator not in cls.VALID_PARAM_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' not one of '{cls.VALID_PARAM_COMPARATORS}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def is_tag(cls, key_type, comparator):
|
|
if key_type == cls._TAG_IDENTIFIER:
|
|
if comparator not in cls.VALID_TAG_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' not one of '{cls.VALID_TAG_COMPARATORS}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def is_attribute(cls, key_type, key_name, comparator):
|
|
return cls.is_string_attribute(key_type, key_name, comparator) or cls.is_numeric_attribute(
|
|
key_type, key_name, comparator
|
|
)
|
|
|
|
@classmethod
|
|
def is_string_attribute(cls, key_type, key_name, comparator):
|
|
if key_type == cls._ATTRIBUTE_IDENTIFIER and key_name not in cls.NUMERIC_ATTRIBUTES:
|
|
if comparator not in cls.VALID_STRING_ATTRIBUTE_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' not one of "
|
|
f"'{cls.VALID_STRING_ATTRIBUTE_COMPARATORS}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def is_numeric_attribute(cls, key_type, key_name, comparator):
|
|
if key_type == cls._ATTRIBUTE_IDENTIFIER and key_name in cls.NUMERIC_ATTRIBUTES:
|
|
if comparator not in cls.VALID_NUMERIC_ATTRIBUTE_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' not one of "
|
|
f"'{cls.VALID_STRING_ATTRIBUTE_COMPARATORS}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def is_dataset(cls, key_type, comparator):
|
|
if key_type == cls._DATASET_IDENTIFIER:
|
|
if comparator not in cls.VALID_DATASET_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' "
|
|
f"not one of '{cls.VALID_DATASET_COMPARATORS}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def _does_run_match_clause(cls, run, sed):
|
|
key_type = sed.get("type")
|
|
key = sed.get("key")
|
|
value = sed.get("value")
|
|
comparator = sed.get("comparator").upper()
|
|
|
|
key = SearchUtils.translate_key_alias(key)
|
|
|
|
if cls.is_metric(key_type, comparator):
|
|
lhs = run.data.metrics.get(key, None)
|
|
value = float(value)
|
|
elif cls.is_param(key_type, comparator):
|
|
lhs = run.data.params.get(key, None)
|
|
elif cls.is_tag(key_type, comparator):
|
|
lhs = run.data.tags.get(key, None)
|
|
elif cls.is_string_attribute(key_type, key, comparator):
|
|
lhs = getattr(run.info, key)
|
|
elif cls.is_numeric_attribute(key_type, key, comparator):
|
|
lhs = getattr(run.info, key)
|
|
value = int(value)
|
|
elif cls.is_dataset(key_type, comparator):
|
|
if key == "context":
|
|
return any(
|
|
SearchUtils.get_comparison_func(comparator)(tag.value if tag else None, value)
|
|
for dataset_input in run.inputs.dataset_inputs
|
|
for tag in dataset_input.tags
|
|
if tag.key == MLFLOW_DATASET_CONTEXT
|
|
)
|
|
else:
|
|
return any(
|
|
SearchUtils.get_comparison_func(comparator)(
|
|
getattr(dataset_input.dataset, key), value
|
|
)
|
|
for dataset_input in run.inputs.dataset_inputs
|
|
)
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid search expression type '{key_type}'", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
if lhs is None:
|
|
return False
|
|
|
|
return SearchUtils.get_comparison_func(comparator)(lhs, value)
|
|
|
|
@classmethod
|
|
def filter(cls, runs, filter_string):
|
|
"""Filters a set of runs based on a search filter string."""
|
|
if not filter_string:
|
|
return runs
|
|
parsed = cls.parse_search_filter(filter_string)
|
|
|
|
def run_matches(run):
|
|
return all(cls._does_run_match_clause(run, s) for s in parsed)
|
|
|
|
return [run for run in runs if run_matches(run)]
|
|
|
|
@classmethod
|
|
def _validate_order_by_and_generate_token(cls, order_by):
|
|
try:
|
|
parsed = sqlparse.parse(order_by)
|
|
except Exception:
|
|
raise MlflowException(
|
|
f"Error on parsing order_by clause '{order_by}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if len(parsed) != 1 or not isinstance(parsed[0], Statement):
|
|
raise MlflowException(
|
|
f"Invalid order_by clause '{order_by}'. Could not be parsed.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
statement = parsed[0]
|
|
ttype_for_timestamp = (
|
|
TokenType.Name.Builtin
|
|
if Version(sqlparse.__version__) >= Version("0.4.3")
|
|
else TokenType.Keyword
|
|
)
|
|
|
|
if len(statement.tokens) == 1 and isinstance(statement[0], Identifier):
|
|
token_value = statement.tokens[0].value
|
|
elif len(statement.tokens) == 1 and statement.tokens[0].match(
|
|
ttype=ttype_for_timestamp, values=[cls.ORDER_BY_KEY_TIMESTAMP]
|
|
):
|
|
token_value = cls.ORDER_BY_KEY_TIMESTAMP
|
|
elif (
|
|
statement.tokens[0].match(
|
|
ttype=ttype_for_timestamp, values=[cls.ORDER_BY_KEY_TIMESTAMP]
|
|
)
|
|
and all(token.is_whitespace for token in statement.tokens[1:-1])
|
|
and statement.tokens[-1].ttype == TokenType.Keyword.Order
|
|
):
|
|
token_value = cls.ORDER_BY_KEY_TIMESTAMP + " " + statement.tokens[-1].value
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid order_by clause '{order_by}'. Could not be parsed.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return token_value
|
|
|
|
@classmethod
|
|
def _parse_order_by_string(cls, order_by):
|
|
token_value = cls._validate_order_by_and_generate_token(order_by)
|
|
is_ascending = True
|
|
tokens = shlex.split(token_value.replace("`", '"'))
|
|
if len(tokens) > 2:
|
|
raise MlflowException(
|
|
f"Invalid order_by clause '{order_by}'. Could not be parsed.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif len(tokens) == 2:
|
|
order_token = tokens[1].lower()
|
|
if order_token not in cls.VALID_ORDER_BY_TAGS:
|
|
raise MlflowException(
|
|
f"Invalid ordering key in order_by clause '{order_by}'.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
is_ascending = order_token == cls.ASC_OPERATOR
|
|
token_value = tokens[0]
|
|
return token_value, is_ascending
|
|
|
|
@classmethod
|
|
def parse_order_by_for_search_runs(cls, order_by):
|
|
token_value, is_ascending = cls._parse_order_by_string(order_by)
|
|
identifier = cls._get_identifier(token_value.strip(), cls.VALID_ORDER_BY_ATTRIBUTE_KEYS)
|
|
return identifier["type"], identifier["key"], is_ascending
|
|
|
|
@classmethod
|
|
def parse_order_by_for_search_registered_models(cls, order_by):
|
|
token_value, is_ascending = cls._parse_order_by_string(order_by)
|
|
token_value = token_value.strip()
|
|
if token_value not in cls.VALID_ORDER_BY_KEYS_REGISTERED_MODELS:
|
|
raise MlflowException(
|
|
f"Invalid order by key '{token_value}' specified. Valid keys "
|
|
f"are '{cls.RECOMMENDED_ORDER_BY_KEYS_REGISTERED_MODELS}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return token_value, is_ascending
|
|
|
|
@classmethod
|
|
def _get_value_for_sort(cls, run, key_type, key, ascending):
|
|
"""Returns a tuple suitable to be used as a sort key for runs."""
|
|
sort_value = None
|
|
key = SearchUtils.translate_key_alias(key)
|
|
if key_type == cls._METRIC_IDENTIFIER:
|
|
sort_value = run.data.metrics.get(key)
|
|
elif key_type == cls._PARAM_IDENTIFIER:
|
|
sort_value = run.data.params.get(key)
|
|
elif key_type == cls._TAG_IDENTIFIER:
|
|
sort_value = run.data.tags.get(key)
|
|
elif key_type == cls._ATTRIBUTE_IDENTIFIER:
|
|
sort_value = getattr(run.info, key)
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid order_by entity type '{key_type}'", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
|
|
# Return a key such that None values are always at the end.
|
|
is_none = sort_value is None
|
|
is_nan = isinstance(sort_value, float) and math.isnan(sort_value)
|
|
fill_value = (1 if ascending else -1) * math.inf
|
|
|
|
if is_none:
|
|
sort_value = fill_value
|
|
elif is_nan:
|
|
sort_value = -fill_value
|
|
|
|
is_none_or_nan = is_none or is_nan
|
|
|
|
return (is_none_or_nan, sort_value) if ascending else (not is_none_or_nan, sort_value)
|
|
|
|
@classmethod
|
|
def sort(cls, runs, order_by_list):
|
|
"""Sorts a set of runs based on their natural ordering and an overriding set of order_bys.
|
|
Runs are naturally ordered first by start time descending, then by run id for tie-breaking.
|
|
"""
|
|
runs = sorted(runs, key=lambda run: (-run.info.start_time, run.info.run_uuid))
|
|
if not order_by_list:
|
|
return runs
|
|
# NB: We rely on the stability of Python's sort function, so that we can apply
|
|
# the ordering conditions in reverse order.
|
|
for order_by_clause in reversed(order_by_list):
|
|
(key_type, key, ascending) = cls.parse_order_by_for_search_runs(order_by_clause)
|
|
|
|
runs = sorted(
|
|
runs,
|
|
key=lambda run: cls._get_value_for_sort(run, key_type, key, ascending),
|
|
reverse=not ascending,
|
|
)
|
|
return runs
|
|
|
|
@classmethod
|
|
def parse_start_offset_from_page_token(cls, page_token):
|
|
# Note: the page_token is expected to be a base64-encoded JSON that looks like
|
|
# { "offset": xxx }. However, this format is not stable, so it should not be
|
|
# relied upon outside of this method.
|
|
if not page_token:
|
|
return 0
|
|
|
|
try:
|
|
decoded_token = base64.b64decode(page_token)
|
|
except TypeError:
|
|
raise MlflowException(
|
|
"Invalid page token, could not base64-decode", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
except base64.binascii.Error:
|
|
raise MlflowException(
|
|
"Invalid page token, could not base64-decode", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
|
|
try:
|
|
parsed_token = json.loads(decoded_token)
|
|
except ValueError:
|
|
raise MlflowException(
|
|
f"Invalid page token, decoded value={decoded_token}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
offset_str = parsed_token.get("offset")
|
|
if not offset_str:
|
|
raise MlflowException(
|
|
f"Invalid page token, parsed value={parsed_token}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
try:
|
|
offset = int(offset_str)
|
|
except ValueError:
|
|
raise MlflowException(
|
|
f"Invalid page token, not stringable {offset_str}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
return offset
|
|
|
|
@classmethod
|
|
def create_page_token(cls, offset):
|
|
return base64.b64encode(json.dumps({"offset": offset}).encode("utf-8"))
|
|
|
|
@classmethod
|
|
def paginate(cls, runs, page_token, max_results):
|
|
"""Paginates a set of runs based on an offset encoded into the page_token and a max
|
|
results limit. Returns a pair containing the set of paginated runs, followed by
|
|
an optional next_page_token if there are further results that need to be returned.
|
|
"""
|
|
start_offset = cls.parse_start_offset_from_page_token(page_token)
|
|
final_offset = start_offset + max_results
|
|
|
|
paginated_runs = runs[start_offset:final_offset]
|
|
next_page_token = None
|
|
if final_offset < len(runs):
|
|
next_page_token = cls.create_page_token(final_offset)
|
|
return (paginated_runs, next_page_token)
|
|
|
|
# Model Registry specific parser
|
|
# TODO: Tech debt. Refactor search code into common utils, tracking server, and model
|
|
# registry specific code.
|
|
|
|
VALID_SEARCH_KEYS_FOR_MODEL_VERSIONS = {"name", "run_id", "source_path"}
|
|
VALID_SEARCH_KEYS_FOR_REGISTERED_MODELS = {"name"}
|
|
|
|
@classmethod
|
|
def _check_valid_identifier_list(cls, value_token):
|
|
if len(value_token._groupable_tokens) == 0:
|
|
raise MlflowException(
|
|
"While parsing a list in the query,"
|
|
" expected a non-empty list of string values, but got empty list",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
# Single element (e.g. `('x')`)
|
|
if (
|
|
len(value_token._groupable_tokens) == 1
|
|
and value_token._groupable_tokens[0].ttype is TokenType.String.Single
|
|
):
|
|
return
|
|
|
|
# Multiple elements (e.g. `('x','y')`)
|
|
if not isinstance(value_token._groupable_tokens[0], IdentifierList):
|
|
raise MlflowException(
|
|
"While parsing a list in the query,"
|
|
" expected a non-empty list of string values, but got ill-formed list.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif not all(
|
|
token.ttype
|
|
in {*cls.STRING_VALUE_TYPES, *cls.DELIMITER_VALUE_TYPES, cls.WHITESPACE_VALUE_TYPE}
|
|
for token in value_token._groupable_tokens[0].tokens
|
|
):
|
|
raise MlflowException(
|
|
"While parsing a list in the query, expected string value, punctuation, "
|
|
f"or whitespace, but got different type in list: {value_token}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
@classmethod
|
|
def _parse_list_from_sql_token(cls, token):
|
|
try:
|
|
str_or_tuple = ast.literal_eval(token.value)
|
|
return [str_or_tuple] if isinstance(str_or_tuple, str) else str_or_tuple
|
|
except SyntaxError:
|
|
raise MlflowException(
|
|
"While parsing a list in the query,"
|
|
" expected a non-empty list of string values, but got ill-formed list.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
@classmethod
|
|
def _parse_run_ids(cls, token):
|
|
cls._check_valid_identifier_list(token)
|
|
run_id_list = cls._parse_list_from_sql_token(token)
|
|
# Because MySQL IN clause is case-insensitive, but all run_ids only contain lower
|
|
# case letters, so that we filter out run_ids containing upper case letters here.
|
|
return [run_id for run_id in run_id_list if run_id.islower()]
|
|
|
|
|
|
class SearchExperimentsUtils(SearchUtils):
|
|
VALID_SEARCH_ATTRIBUTE_KEYS = {"name", "creation_time", "last_update_time"}
|
|
VALID_ORDER_BY_ATTRIBUTE_KEYS = {"name", "experiment_id", "creation_time", "last_update_time"}
|
|
NUMERIC_ATTRIBUTES = {"creation_time", "last_update_time"}
|
|
|
|
@classmethod
|
|
def _invalid_statement_token_search_experiments(cls, token):
|
|
if (
|
|
isinstance(token, Comparison)
|
|
or token.is_whitespace
|
|
or token.match(ttype=TokenType.Keyword, values=["AND"])
|
|
):
|
|
return False
|
|
return True
|
|
|
|
@classmethod
|
|
def _process_statement(cls, statement):
|
|
tokens = _join_in_comparison_tokens(statement.tokens)
|
|
invalids = list(filter(cls._invalid_statement_token_search_experiments, tokens))
|
|
if len(invalids) > 0:
|
|
invalid_clauses = ", ".join(map(str, invalids))
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid clause(s) in filter string: {invalid_clauses}"
|
|
)
|
|
return [cls._get_comparison(t) for t in tokens if isinstance(t, Comparison)]
|
|
|
|
@classmethod
|
|
def _get_identifier(cls, identifier, valid_attributes):
|
|
tokens = identifier.split(".", maxsplit=1)
|
|
if len(tokens) == 1:
|
|
key = tokens[0]
|
|
identifier = cls._ATTRIBUTE_IDENTIFIER
|
|
else:
|
|
entity_type, key = tokens
|
|
valid_entity_types = ("attribute", "tag", "tags")
|
|
if entity_type not in valid_entity_types:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid entity type '{entity_type}'. "
|
|
f"Valid entity types are {valid_entity_types}"
|
|
)
|
|
identifier = cls._valid_entity_type(entity_type)
|
|
|
|
key = cls._trim_backticks(cls._strip_quotes(key))
|
|
if identifier == cls._ATTRIBUTE_IDENTIFIER and key not in valid_attributes:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid attribute key '{key}' specified. Valid keys are '{valid_attributes}'"
|
|
)
|
|
return {"type": identifier, "key": key}
|
|
|
|
@classmethod
|
|
def _get_comparison(cls, comparison):
|
|
stripped_comparison = [token for token in comparison.tokens if not token.is_whitespace]
|
|
cls._validate_comparison(stripped_comparison)
|
|
left, comparator, right = stripped_comparison
|
|
comp = cls._get_identifier(left.value, cls.VALID_SEARCH_ATTRIBUTE_KEYS)
|
|
comp["comparator"] = comparator.value
|
|
comp["value"] = cls._get_value(comp.get("type"), comp.get("key"), right)
|
|
return comp
|
|
|
|
@classmethod
|
|
def parse_order_by_for_search_experiments(cls, order_by):
|
|
token_value, is_ascending = cls._parse_order_by_string(order_by)
|
|
identifier = cls._get_identifier(token_value.strip(), cls.VALID_ORDER_BY_ATTRIBUTE_KEYS)
|
|
return identifier["type"], identifier["key"], is_ascending
|
|
|
|
@classmethod
|
|
def is_attribute(cls, key_type, comparator):
|
|
if key_type == cls._ATTRIBUTE_IDENTIFIER:
|
|
if comparator not in cls.VALID_STRING_ATTRIBUTE_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' not one of "
|
|
f"'{cls.VALID_STRING_ATTRIBUTE_COMPARATORS}'"
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def _does_experiment_match_clause(cls, experiment, sed):
|
|
key_type = sed.get("type")
|
|
key = sed.get("key")
|
|
value = sed.get("value")
|
|
comparator = sed.get("comparator").upper()
|
|
|
|
if cls.is_string_attribute(key_type, key, comparator):
|
|
lhs = getattr(experiment, key)
|
|
elif cls.is_numeric_attribute(key_type, key, comparator):
|
|
lhs = getattr(experiment, key)
|
|
value = float(value)
|
|
elif cls.is_tag(key_type, comparator):
|
|
if key not in experiment.tags:
|
|
return False
|
|
lhs = experiment.tags.get(key, None)
|
|
if lhs is None:
|
|
return experiment
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid search expression type '{key_type}'", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
|
|
return SearchUtils.get_comparison_func(comparator)(lhs, value)
|
|
|
|
@classmethod
|
|
def filter(cls, experiments, filter_string):
|
|
if not filter_string:
|
|
return experiments
|
|
parsed = cls.parse_search_filter(filter_string)
|
|
|
|
def experiment_matches(experiment):
|
|
return all(cls._does_experiment_match_clause(experiment, s) for s in parsed)
|
|
|
|
return list(filter(experiment_matches, experiments))
|
|
|
|
@classmethod
|
|
def _get_sort_key(cls, order_by_list):
|
|
order_by = []
|
|
parsed_order_by = map(cls.parse_order_by_for_search_experiments, order_by_list)
|
|
for type_, key, ascending in parsed_order_by:
|
|
if type_ == "attribute":
|
|
order_by.append((key, ascending))
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(f"Invalid order_by entity: {type_}")
|
|
|
|
# Add a tie-breaker
|
|
if not any(key == "experiment_id" for key, _ in order_by):
|
|
order_by.append(("experiment_id", False))
|
|
|
|
# https://stackoverflow.com/a/56842689
|
|
class _Sorter:
|
|
def __init__(self, obj, ascending):
|
|
self.obj = obj
|
|
self.ascending = ascending
|
|
|
|
# Only need < and == are needed for use as a key parameter in the sorted function
|
|
def __eq__(self, other):
|
|
return other.obj == self.obj
|
|
|
|
def __lt__(self, other):
|
|
if self.obj is None:
|
|
return False
|
|
elif other.obj is None:
|
|
return True
|
|
elif self.ascending:
|
|
return self.obj < other.obj
|
|
else:
|
|
return other.obj < self.obj
|
|
|
|
def _apply_sorter(experiment, key, ascending):
|
|
attr = getattr(experiment, key)
|
|
return _Sorter(attr, ascending)
|
|
|
|
return lambda experiment: tuple(_apply_sorter(experiment, k, asc) for (k, asc) in order_by)
|
|
|
|
@classmethod
|
|
def sort(cls, experiments, order_by_list):
|
|
return sorted(experiments, key=cls._get_sort_key(order_by_list))
|
|
|
|
|
|
# https://stackoverflow.com/a/56842689
|
|
class _Reversor:
|
|
def __init__(self, obj):
|
|
self.obj = obj
|
|
|
|
# Only need < and == are needed for use as a key parameter in the sorted function
|
|
def __eq__(self, other):
|
|
return other.obj == self.obj
|
|
|
|
def __lt__(self, other):
|
|
return other.obj < self.obj
|
|
|
|
|
|
def _apply_reversor(model, key, ascending):
|
|
attr = getattr(model, key)
|
|
return attr if ascending else _Reversor(attr)
|
|
|
|
|
|
class SearchModelUtils(SearchUtils):
|
|
NUMERIC_ATTRIBUTES = {"creation_timestamp", "last_updated_timestamp"}
|
|
VALID_SEARCH_ATTRIBUTE_KEYS = {"name"}
|
|
VALID_ORDER_BY_KEYS_REGISTERED_MODELS = {"name", "creation_timestamp", "last_updated_timestamp"}
|
|
|
|
@classmethod
|
|
def _does_registered_model_match_clauses(cls, model, sed):
|
|
key_type = sed.get("type")
|
|
key = sed.get("key")
|
|
value = sed.get("value")
|
|
comparator = sed.get("comparator").upper()
|
|
|
|
# what comparators do we support here?
|
|
if cls.is_string_attribute(key_type, key, comparator):
|
|
lhs = getattr(model, key)
|
|
elif cls.is_numeric_attribute(key_type, key, comparator):
|
|
lhs = getattr(model, key)
|
|
value = int(value)
|
|
elif cls.is_tag(key_type, comparator):
|
|
# NB: We should use the private attribute `_tags` instead of the `tags` property
|
|
# to consider all tags including reserved ones.
|
|
lhs = model._tags.get(key, None)
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid search expression type '{key_type}'", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
|
|
# NB: Handling the special `mlflow.prompt.is_prompt` tag. This tag is used for
|
|
# distinguishing between prompt models and normal models. For example, we want to
|
|
# search for models only by the following filter string:
|
|
#
|
|
# tags.`mlflow.prompt.is_prompt` != 'true'
|
|
# tags.`mlflow.prompt.is_prompt` = 'false'
|
|
#
|
|
# However, models do not have this tag, so lhs is None in this case. Instead of returning
|
|
# False like normal tag filter, we need to return True here.
|
|
if key == IS_PROMPT_TAG_KEY and lhs is None:
|
|
return (comparator == "=" and value == "false") or (
|
|
comparator == "!=" and value == "true"
|
|
)
|
|
|
|
if lhs is None:
|
|
return False
|
|
|
|
return SearchUtils.get_comparison_func(comparator)(lhs, value)
|
|
|
|
@classmethod
|
|
def filter(cls, registered_models, filter_string):
|
|
"""Filters a set of registered models based on a search filter string."""
|
|
if not filter_string:
|
|
return registered_models
|
|
parsed = cls.parse_search_filter(filter_string)
|
|
|
|
def registered_model_matches(model):
|
|
return all(cls._does_registered_model_match_clauses(model, s) for s in parsed)
|
|
|
|
return [
|
|
registered_model
|
|
for registered_model in registered_models
|
|
if registered_model_matches(registered_model)
|
|
]
|
|
|
|
@classmethod
|
|
def parse_order_by_for_search_registered_models(cls, order_by):
|
|
token_value, is_ascending = cls._parse_order_by_string(order_by)
|
|
identifier = SearchExperimentsUtils._get_identifier(
|
|
token_value.strip(), cls.VALID_ORDER_BY_KEYS_REGISTERED_MODELS
|
|
)
|
|
return identifier["type"], identifier["key"], is_ascending
|
|
|
|
@classmethod
|
|
def _get_sort_key(cls, order_by_list):
|
|
order_by = []
|
|
parsed_order_by = map(cls.parse_order_by_for_search_registered_models, order_by_list or [])
|
|
for type_, key, ascending in parsed_order_by:
|
|
if type_ == "attribute":
|
|
order_by.append((key, ascending))
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(f"Invalid order_by entity: {type_}")
|
|
|
|
# Add a tie-breaker
|
|
if not any(key == "name" for key, _ in order_by):
|
|
order_by.append(("name", True))
|
|
|
|
return lambda model: tuple(_apply_reversor(model, k, asc) for (k, asc) in order_by)
|
|
|
|
@classmethod
|
|
def sort(cls, models, order_by_list):
|
|
return sorted(models, key=cls._get_sort_key(order_by_list))
|
|
|
|
@classmethod
|
|
def _process_statement(cls, statement):
|
|
tokens = _join_in_comparison_tokens(statement.tokens)
|
|
invalids = list(filter(cls._invalid_statement_token_search_model_registry, tokens))
|
|
if len(invalids) > 0:
|
|
invalid_clauses = ", ".join(map(str, invalids))
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid clause(s) in filter string: {invalid_clauses}"
|
|
)
|
|
return [cls._get_comparison(t) for t in tokens if isinstance(t, Comparison)]
|
|
|
|
@classmethod
|
|
def _get_model_search_identifier(cls, identifier, valid_attributes):
|
|
tokens = identifier.split(".", maxsplit=1)
|
|
if len(tokens) == 1:
|
|
key = tokens[0]
|
|
identifier = cls._ATTRIBUTE_IDENTIFIER
|
|
else:
|
|
entity_type, key = tokens
|
|
valid_entity_types = ("attribute", "tag", "tags")
|
|
if entity_type not in valid_entity_types:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid entity type '{entity_type}'. "
|
|
f"Valid entity types are {valid_entity_types}"
|
|
)
|
|
identifier = (
|
|
cls._TAG_IDENTIFIER if entity_type in ("tag", "tags") else cls._ATTRIBUTE_IDENTIFIER
|
|
)
|
|
|
|
if identifier == cls._ATTRIBUTE_IDENTIFIER and key not in valid_attributes:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid attribute key '{key}' specified. Valid keys are '{valid_attributes}'"
|
|
)
|
|
|
|
key = cls._trim_backticks(cls._strip_quotes(key))
|
|
return {"type": identifier, "key": key}
|
|
|
|
@classmethod
|
|
def _get_comparison(cls, comparison):
|
|
stripped_comparison = [token for token in comparison.tokens if not token.is_whitespace]
|
|
cls._validate_comparison(stripped_comparison)
|
|
left, comparator, right = stripped_comparison
|
|
comp = cls._get_model_search_identifier(left.value, cls.VALID_SEARCH_ATTRIBUTE_KEYS)
|
|
comp["comparator"] = comparator.value.upper()
|
|
comp["value"] = cls._get_value(comp.get("type"), comp.get("key"), right)
|
|
return comp
|
|
|
|
@classmethod
|
|
def _get_value(cls, identifier_type, key, token):
|
|
if identifier_type == cls._TAG_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
raise MlflowException(
|
|
"Expected a quoted string value for "
|
|
f"{identifier_type} (e.g. 'my-value'). Got value "
|
|
f"{token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif identifier_type == cls._ATTRIBUTE_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
elif isinstance(token, Parenthesis):
|
|
if key != "run_id":
|
|
raise MlflowException(
|
|
"Only the 'run_id' attribute supports comparison with a list of quoted "
|
|
"string values.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls._parse_run_ids(token)
|
|
else:
|
|
raise MlflowException(
|
|
"Expected a quoted string value or a list of quoted string values for "
|
|
f"attributes. Got value {token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
# Expected to be either "param" or "metric".
|
|
raise MlflowException(
|
|
"Invalid identifier type. Expected one of "
|
|
f"{[cls._ATTRIBUTE_IDENTIFIER, cls._TAG_IDENTIFIER]}.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
@classmethod
|
|
def _invalid_statement_token_search_model_registry(cls, token):
|
|
if (
|
|
isinstance(token, Comparison)
|
|
or token.is_whitespace
|
|
or token.match(ttype=TokenType.Keyword, values=["AND"])
|
|
):
|
|
return False
|
|
return True
|
|
|
|
|
|
class SearchModelVersionUtils(SearchUtils):
|
|
NUMERIC_ATTRIBUTES = {"version_number", "creation_timestamp", "last_updated_timestamp"}
|
|
VALID_SEARCH_ATTRIBUTE_KEYS = {
|
|
"name",
|
|
"version_number",
|
|
"run_id",
|
|
"source_path",
|
|
}
|
|
VALID_ORDER_BY_ATTRIBUTE_KEYS = {
|
|
"name",
|
|
"version_number",
|
|
"creation_timestamp",
|
|
"last_updated_timestamp",
|
|
}
|
|
VALID_STRING_ATTRIBUTE_COMPARATORS = {"!=", "=", "LIKE", "ILIKE", "IN"}
|
|
|
|
@classmethod
|
|
def _does_model_version_match_clauses(cls, mv, sed):
|
|
key_type = sed.get("type")
|
|
key = sed.get("key")
|
|
value = sed.get("value")
|
|
comparator = sed.get("comparator").upper()
|
|
|
|
if cls.is_string_attribute(key_type, key, comparator):
|
|
lhs = getattr(mv, "source" if key == "source_path" else key)
|
|
elif cls.is_numeric_attribute(key_type, key, comparator):
|
|
if key == "version_number":
|
|
key = "version"
|
|
lhs = getattr(mv, key)
|
|
value = int(value)
|
|
elif cls.is_tag(key_type, comparator):
|
|
lhs = mv.tags.get(key, None)
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid search expression type '{key_type}'", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
|
|
# NB: Handling the special `mlflow.prompt.is_prompt` tag. This tag is used for
|
|
# distinguishing between prompt models and normal models. For example, we want to
|
|
# search for models only by the following filter string:
|
|
#
|
|
# tags.`mlflow.prompt.is_prompt` != 'true'
|
|
# tags.`mlflow.prompt.is_prompt` = 'false'
|
|
#
|
|
# However, models do not have this tag, so lhs is None in this case. Instead of returning
|
|
# False like normal tag filter, we need to return True here.
|
|
if key == IS_PROMPT_TAG_KEY and lhs is None:
|
|
return (comparator == "=" and value == "false") or (
|
|
comparator == "!=" and value == "true"
|
|
)
|
|
|
|
if lhs is None:
|
|
return False
|
|
|
|
if comparator == "IN" and isinstance(value, (set, list)):
|
|
return lhs in set(value)
|
|
|
|
return SearchUtils.get_comparison_func(comparator)(lhs, value)
|
|
|
|
@classmethod
|
|
def filter(cls, model_versions, filter_string):
|
|
"""Filters a set of model versions based on a search filter string."""
|
|
model_versions = [mv for mv in model_versions if mv.current_stage != STAGE_DELETED_INTERNAL]
|
|
if not filter_string:
|
|
return model_versions
|
|
parsed = cls.parse_search_filter(filter_string)
|
|
|
|
def model_version_matches(mv):
|
|
return all(cls._does_model_version_match_clauses(mv, s) for s in parsed)
|
|
|
|
return [mv for mv in model_versions if model_version_matches(mv)]
|
|
|
|
@classmethod
|
|
def parse_order_by_for_search_model_versions(cls, order_by):
|
|
token_value, is_ascending = cls._parse_order_by_string(order_by)
|
|
identifier = SearchExperimentsUtils._get_identifier(
|
|
token_value.strip(), cls.VALID_ORDER_BY_ATTRIBUTE_KEYS
|
|
)
|
|
return identifier["type"], identifier["key"], is_ascending
|
|
|
|
@classmethod
|
|
def _get_sort_key(cls, order_by_list):
|
|
order_by = []
|
|
parsed_order_by = map(cls.parse_order_by_for_search_model_versions, order_by_list or [])
|
|
for type_, key, ascending in parsed_order_by:
|
|
if type_ == "attribute":
|
|
# Need to add this mapping because version is a keyword in sql
|
|
if key == "version_number":
|
|
key = "version"
|
|
order_by.append((key, ascending))
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(f"Invalid order_by entity: {type_}")
|
|
|
|
# Add a tie-breaker
|
|
if not any(key == "name" for key, _ in order_by):
|
|
order_by.append(("name", True))
|
|
if not any(key == "version_number" for key, _ in order_by):
|
|
order_by.append(("version", False))
|
|
|
|
return lambda model_version: tuple(
|
|
_apply_reversor(model_version, k, asc) for (k, asc) in order_by
|
|
)
|
|
|
|
@classmethod
|
|
def sort(cls, model_versions, order_by_list):
|
|
return sorted(model_versions, key=cls._get_sort_key(order_by_list))
|
|
|
|
@classmethod
|
|
def _get_model_version_search_identifier(cls, identifier, valid_attributes):
|
|
tokens = identifier.split(".", maxsplit=1)
|
|
if len(tokens) == 1:
|
|
key = tokens[0]
|
|
identifier = cls._ATTRIBUTE_IDENTIFIER
|
|
else:
|
|
entity_type, key = tokens
|
|
valid_entity_types = ("attribute", "tag", "tags")
|
|
if entity_type not in valid_entity_types:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid entity type '{entity_type}'. "
|
|
f"Valid entity types are {valid_entity_types}"
|
|
)
|
|
identifier = (
|
|
cls._TAG_IDENTIFIER if entity_type in ("tag", "tags") else cls._ATTRIBUTE_IDENTIFIER
|
|
)
|
|
|
|
if identifier == cls._ATTRIBUTE_IDENTIFIER and key not in valid_attributes:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid attribute key '{key}' specified. Valid keys are '{valid_attributes}'"
|
|
)
|
|
|
|
key = cls._trim_backticks(cls._strip_quotes(key))
|
|
return {"type": identifier, "key": key}
|
|
|
|
@classmethod
|
|
def _get_comparison(cls, comparison):
|
|
stripped_comparison = [token for token in comparison.tokens if not token.is_whitespace]
|
|
cls._validate_comparison(stripped_comparison)
|
|
left, comparator, right = stripped_comparison
|
|
comp = cls._get_model_version_search_identifier(left.value, cls.VALID_SEARCH_ATTRIBUTE_KEYS)
|
|
comp["comparator"] = comparator.value.upper()
|
|
comp["value"] = cls._get_value(comp.get("type"), comp.get("key"), right)
|
|
return comp
|
|
|
|
@classmethod
|
|
def _get_value(cls, identifier_type, key, token):
|
|
if identifier_type == cls._TAG_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
raise MlflowException(
|
|
"Expected a quoted string value for "
|
|
f"{identifier_type} (e.g. 'my-value'). Got value "
|
|
f"{token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif identifier_type == cls._ATTRIBUTE_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
elif isinstance(token, Parenthesis):
|
|
if key != "run_id":
|
|
raise MlflowException(
|
|
"Only the 'run_id' attribute supports comparison with a list of quoted "
|
|
"string values.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls._parse_run_ids(token)
|
|
elif token.ttype in cls.NUMERIC_VALUE_TYPES:
|
|
if key not in cls.NUMERIC_ATTRIBUTES:
|
|
raise MlflowException(
|
|
f"Only the '{cls.NUMERIC_ATTRIBUTES}' attributes support comparison with "
|
|
"numeric values.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if token.ttype == TokenType.Literal.Number.Integer:
|
|
return int(token.value)
|
|
elif token.ttype == TokenType.Literal.Number.Float:
|
|
return float(token.value)
|
|
else:
|
|
raise MlflowException(
|
|
"Expected a quoted string value or a list of quoted string values for "
|
|
f"attributes. Got value {token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
# Expected to be either "param" or "metric".
|
|
raise MlflowException(
|
|
"Invalid identifier type. Expected one of "
|
|
f"{[cls._ATTRIBUTE_IDENTIFIER, cls._TAG_IDENTIFIER]}.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
@classmethod
|
|
def _process_statement(cls, statement):
|
|
tokens = _join_in_comparison_tokens(statement.tokens)
|
|
invalids = list(filter(cls._invalid_statement_token_search_model_version, tokens))
|
|
if len(invalids) > 0:
|
|
invalid_clauses = ", ".join(map(str, invalids))
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid clause(s) in filter string: {invalid_clauses}"
|
|
)
|
|
return [cls._get_comparison(t) for t in tokens if isinstance(t, Comparison)]
|
|
|
|
@classmethod
|
|
def _invalid_statement_token_search_model_version(cls, token):
|
|
if (
|
|
isinstance(token, Comparison)
|
|
or token.is_whitespace
|
|
or token.match(ttype=TokenType.Keyword, values=["AND"])
|
|
):
|
|
return False
|
|
return True
|
|
|
|
@classmethod
|
|
def parse_search_filter(cls, filter_string):
|
|
if not filter_string:
|
|
return []
|
|
try:
|
|
parsed = sqlparse.parse(filter_string)
|
|
except Exception:
|
|
raise MlflowException(
|
|
f"Error on parsing filter '{filter_string}'", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
if len(parsed) == 0 or not isinstance(parsed[0], Statement):
|
|
raise MlflowException(
|
|
f"Invalid filter '{filter_string}'. Could not be parsed.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif len(parsed) > 1:
|
|
raise MlflowException(
|
|
f"Search filter contained multiple expression {filter_string!r}. "
|
|
"Provide AND-ed expression list.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls._process_statement(parsed[0])
|
|
|
|
|
|
class SearchTraceUtils(SearchUtils):
|
|
"""
|
|
Utility class for searching traces.
|
|
"""
|
|
|
|
VALID_SEARCH_ATTRIBUTE_KEYS = {
|
|
"request_id",
|
|
"timestamp",
|
|
"timestamp_ms",
|
|
"execution_time",
|
|
"execution_time_ms",
|
|
"status",
|
|
# The following keys are mapped to tags or metadata
|
|
"name",
|
|
"run_id",
|
|
}
|
|
VALID_ORDER_BY_ATTRIBUTE_KEYS = {
|
|
"experiment_id",
|
|
"timestamp",
|
|
"timestamp_ms",
|
|
"execution_time",
|
|
"execution_time_ms",
|
|
"status",
|
|
"request_id",
|
|
# The following keys are mapped to tags or metadata
|
|
"name",
|
|
"run_id",
|
|
}
|
|
|
|
NUMERIC_ATTRIBUTES = {
|
|
"timestamp_ms",
|
|
"timestamp",
|
|
"execution_time_ms",
|
|
"execution_time",
|
|
}
|
|
|
|
# For now, don't support LIKE/ILIKE operators for trace search because it may
|
|
# cause performance issues with large attributes and tags. We can revisit this
|
|
# decision if we find a way to support them efficiently.
|
|
VALID_TAG_COMPARATORS = {"!=", "="}
|
|
VALID_STRING_ATTRIBUTE_COMPARATORS = {"!=", "=", "IN", "NOT IN"}
|
|
|
|
_REQUEST_METADATA_IDENTIFIER = "request_metadata"
|
|
_TAG_IDENTIFIER = "tag"
|
|
_ATTRIBUTE_IDENTIFIER = "attribute"
|
|
|
|
# These are aliases for the base identifiers
|
|
# e.g. trace.status is equivalent to attribute.status
|
|
_ALTERNATE_IDENTIFIERS = {
|
|
"tags": _TAG_IDENTIFIER,
|
|
"attributes": _ATTRIBUTE_IDENTIFIER,
|
|
"trace": _ATTRIBUTE_IDENTIFIER,
|
|
"metadata": _REQUEST_METADATA_IDENTIFIER,
|
|
}
|
|
_IDENTIFIERS = {_TAG_IDENTIFIER, _REQUEST_METADATA_IDENTIFIER, _ATTRIBUTE_IDENTIFIER}
|
|
_VALID_IDENTIFIERS = _IDENTIFIERS | set(_ALTERNATE_IDENTIFIERS.keys())
|
|
|
|
SUPPORT_IN_COMPARISON_ATTRIBUTE_KEYS = {"name", "status", "request_id", "run_id"}
|
|
|
|
# Some search keys are defined differently in the DB models.
|
|
# E.g. "name" is mapped to TraceTagKey.TRACE_NAME
|
|
SEARCH_KEY_TO_TAG = {
|
|
"name": TraceTagKey.TRACE_NAME,
|
|
}
|
|
SEARCH_KEY_TO_METADATA = {
|
|
"run_id": TraceMetadataKey.SOURCE_RUN,
|
|
}
|
|
# Alias for attribute keys
|
|
SEARCH_KEY_TO_ATTRIBUTE = {
|
|
"timestamp": "timestamp_ms",
|
|
"execution_time": "execution_time_ms",
|
|
}
|
|
|
|
@classmethod
|
|
def filter(cls, traces, filter_string):
|
|
"""Filters a set of traces based on a search filter string."""
|
|
if not filter_string:
|
|
return traces
|
|
parsed = cls.parse_search_filter_for_search_traces(filter_string)
|
|
|
|
def trace_matches(trace):
|
|
return all(cls._does_trace_match_clause(trace, s) for s in parsed)
|
|
|
|
return list(filter(trace_matches, traces))
|
|
|
|
@classmethod
|
|
def _does_trace_match_clause(cls, trace, sed):
|
|
type_ = sed.get("type")
|
|
key = sed.get("key")
|
|
value = sed.get("value")
|
|
comparator = sed.get("comparator").upper()
|
|
|
|
if cls.is_tag(type_, comparator):
|
|
lhs = trace.tags.get(key)
|
|
elif cls.is_request_metadata(type_, comparator):
|
|
lhs = trace.request_metadata.get(key)
|
|
elif cls.is_attribute(type_, key, comparator):
|
|
lhs = getattr(trace, key)
|
|
elif sed.get("type") == cls._TAG_IDENTIFIER:
|
|
lhs = trace.tags.get(key)
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid search key '{key}', supported are {cls.VALID_SEARCH_ATTRIBUTE_KEYS}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if lhs is None:
|
|
return False
|
|
|
|
return SearchUtils.get_comparison_func(comparator)(lhs, value)
|
|
|
|
@classmethod
|
|
def sort(cls, traces, order_by_list):
|
|
return sorted(traces, key=cls._get_sort_key(order_by_list))
|
|
|
|
@classmethod
|
|
def parse_order_by_for_search_traces(cls, order_by):
|
|
token_value, is_ascending = cls._parse_order_by_string(order_by)
|
|
identifier = cls._get_identifier(token_value.strip(), cls.VALID_ORDER_BY_ATTRIBUTE_KEYS)
|
|
identifier = cls._replace_key_to_tag_or_metadata(identifier)
|
|
return identifier["type"], identifier["key"], is_ascending
|
|
|
|
@classmethod
|
|
def parse_search_filter_for_search_traces(cls, filter_string):
|
|
parsed = cls.parse_search_filter(filter_string)
|
|
return [cls._replace_key_to_tag_or_metadata(p) for p in parsed]
|
|
|
|
@classmethod
|
|
def _replace_key_to_tag_or_metadata(cls, parsed: dict[str, Any]):
|
|
"""
|
|
Replace search key to tag or metadata key if it is in the mapping.
|
|
"""
|
|
key = parsed.get("key").lower()
|
|
if key in cls.SEARCH_KEY_TO_TAG:
|
|
parsed["type"] = cls._TAG_IDENTIFIER
|
|
parsed["key"] = cls.SEARCH_KEY_TO_TAG[key]
|
|
elif key in cls.SEARCH_KEY_TO_METADATA:
|
|
parsed["type"] = cls._REQUEST_METADATA_IDENTIFIER
|
|
parsed["key"] = cls.SEARCH_KEY_TO_METADATA[key]
|
|
elif key in cls.SEARCH_KEY_TO_ATTRIBUTE:
|
|
parsed["key"] = cls.SEARCH_KEY_TO_ATTRIBUTE[key]
|
|
return parsed
|
|
|
|
@classmethod
|
|
def is_request_metadata(cls, key_type, comparator):
|
|
if key_type == cls._REQUEST_METADATA_IDENTIFIER:
|
|
# Request metadata accepts the same set of comparators as tags
|
|
if comparator not in cls.VALID_TAG_COMPARATORS:
|
|
raise MlflowException(
|
|
f"Invalid comparator '{comparator}' not one of '{cls.VALID_TAG_COMPARATORS}'",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
def _valid_entity_type(cls, entity_type):
|
|
entity_type = cls._trim_backticks(entity_type)
|
|
if entity_type not in cls._VALID_IDENTIFIERS:
|
|
raise MlflowException(
|
|
f"Invalid entity type '{entity_type}'. Valid values are {cls._VALID_IDENTIFIERS}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif entity_type in cls._ALTERNATE_IDENTIFIERS:
|
|
return cls._ALTERNATE_IDENTIFIERS[entity_type]
|
|
else:
|
|
return entity_type
|
|
|
|
@classmethod
|
|
def _get_sort_key(cls, order_by_list):
|
|
order_by = []
|
|
parsed_order_by = map(cls.parse_order_by_for_search_traces, order_by_list or [])
|
|
for type_, key, ascending in parsed_order_by:
|
|
if type_ == "attribute":
|
|
order_by.append((key, ascending))
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Invalid order_by entity `{type_}` with key `{key}`"
|
|
)
|
|
|
|
# Add a tie-breaker
|
|
if not any(key == "timestamp_ms" for key, _ in order_by):
|
|
order_by.append(("timestamp_ms", False))
|
|
if not any(key == "request_id" for key, _ in order_by):
|
|
order_by.append(("request_id", True))
|
|
|
|
return lambda trace: tuple(_apply_reversor(trace, k, asc) for (k, asc) in order_by)
|
|
|
|
@classmethod
|
|
def _get_value(cls, identifier_type, key, token):
|
|
if identifier_type == cls._TAG_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
elif isinstance(token, Parenthesis):
|
|
return cls._parse_attribute_lists(token)
|
|
raise MlflowException(
|
|
"Expected a quoted string value for "
|
|
f"{identifier_type} (e.g. 'my-value'). Got value "
|
|
f"{token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif identifier_type == cls._ATTRIBUTE_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
elif isinstance(token, Parenthesis):
|
|
if key not in cls.SUPPORT_IN_COMPARISON_ATTRIBUTE_KEYS:
|
|
raise MlflowException(
|
|
f"Only attributes in {cls.SUPPORT_IN_COMPARISON_ATTRIBUTE_KEYS} "
|
|
"supports comparison with a list of quoted string values.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls._parse_attribute_lists(token)
|
|
elif token.ttype in cls.NUMERIC_VALUE_TYPES:
|
|
if key not in cls.NUMERIC_ATTRIBUTES:
|
|
raise MlflowException(
|
|
f"Only the '{cls.NUMERIC_ATTRIBUTES}' attributes support comparison with "
|
|
"numeric values.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if token.ttype == TokenType.Literal.Number.Integer:
|
|
return int(token.value)
|
|
elif token.ttype == TokenType.Literal.Number.Float:
|
|
return float(token.value)
|
|
else:
|
|
raise MlflowException(
|
|
"Expected a quoted string value or a list of quoted string values for "
|
|
f"attributes. Got value {token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
elif identifier_type == cls._REQUEST_METADATA_IDENTIFIER:
|
|
if token.ttype in cls.STRING_VALUE_TYPES or isinstance(token, Identifier):
|
|
return cls._strip_quotes(token.value, expect_quoted_value=True)
|
|
else:
|
|
raise MlflowException(
|
|
"Expected a quoted string value for "
|
|
f"{identifier_type} (e.g. 'my-value'). Got value "
|
|
f"{token.value}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
# Expected to be either "param" or "metric".
|
|
raise MlflowException(
|
|
f"Invalid identifier type: {identifier_type}. "
|
|
f"Expected one of {cls._VALID_IDENTIFIERS}.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
@classmethod
|
|
def _parse_attribute_lists(cls, token):
|
|
cls._check_valid_identifier_list(token)
|
|
return cls._parse_list_from_sql_token(token)
|
|
|
|
@classmethod
|
|
def _process_statement(cls, statement):
|
|
# check validity
|
|
tokens = _join_in_comparison_tokens(statement.tokens, search_traces=True)
|
|
invalids = list(filter(cls._invalid_statement_token_search_traces, tokens))
|
|
if len(invalids) > 0:
|
|
invalid_clauses = ", ".join(f"'{token}'" for token in invalids)
|
|
raise MlflowException(
|
|
f"Invalid clause(s) in filter string: {invalid_clauses}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return [cls._get_comparison(si) for si in tokens if isinstance(si, Comparison)]
|
|
|
|
@classmethod
|
|
def _invalid_statement_token_search_traces(cls, token):
|
|
if (
|
|
isinstance(token, Comparison)
|
|
or token.is_whitespace
|
|
or token.match(ttype=TokenType.Keyword, values=["AND"])
|
|
):
|
|
return False
|
|
return True
|
|
|
|
@classmethod
|
|
def _get_comparison(cls, comparison):
|
|
stripped_comparison = [token for token in comparison.tokens if not token.is_whitespace]
|
|
cls._validate_comparison(stripped_comparison, search_traces=True)
|
|
comp = cls._get_identifier(stripped_comparison[0].value, cls.VALID_SEARCH_ATTRIBUTE_KEYS)
|
|
comp["comparator"] = stripped_comparison[1].value
|
|
comp["value"] = cls._get_value(comp.get("type"), comp.get("key"), stripped_comparison[2])
|
|
return comp
|