genesis-flow 1.0.0__py3-none-any.whl

mlflow/types/type_hints.py

@@ -0,0 +1,647 @@
+import base64
+import logging
+from datetime import datetime
+from functools import lru_cache
+from typing import Any, NamedTuple, Optional, TypeVar, Union, get_args, get_origin
+
+import pydantic
+import pydantic.fields
+
+from mlflow.environment_variables import _MLFLOW_IS_IN_SERVING_ENVIRONMENT
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.types.schema import (
+    COLSPEC_TYPES,
+    AnyType,
+    Array,
+    ColSpec,
+    DataType,
+    Map,
+    Object,
+    Property,
+    Schema,
+)
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER, model_dump_compat
+from mlflow.utils.warnings_utils import color_warning
+
+FIELD_TYPE = pydantic.fields.FieldInfo if IS_PYDANTIC_V2_OR_NEWER else pydantic.fields.ModelField
+_logger = logging.getLogger(__name__)
+NONE_TYPE = type(None)
+UNION_TYPES = (Union,)
+try:
+    # this import is only available in Python 3.10+
+    from types import UnionType
+
+    UNION_TYPES += (UnionType,)
+except ImportError:
+    pass
+
+# special type hint that can be used to convert data to
+# the input example type after data validation
+TypeFromExample = TypeVar("TypeFromExample")
+OPTIONAL_INPUT_MSG = (
+    "Input cannot be Optional type. Fix this by removing the "
+    "Optional wrapper from the type hint. To use optional fields, "
+    "use a Pydantic-based type hint definition. See "
+    "https://docs.pydantic.dev/latest/api/base_model/ for pydantic "
+    "BaseModel examples. Check https://mlflow.org/docs/latest/model/python_model.html#supported-type-hints"
+    " for more details."
+)
+
+
+# numpy types are not supported
+TYPE_HINTS_TO_DATATYPE_MAPPING = {
+    int: DataType.long,
+    str: DataType.string,
+    bool: DataType.boolean,
+    float: DataType.double,
+    bytes: DataType.binary,
+    datetime: DataType.datetime,
+}
+
+SUPPORTED_TYPE_HINT_MSG = (
+    "Type hints must be a list[...] where collection element type is one of these types: "
+    f"{list(TYPE_HINTS_TO_DATATYPE_MAPPING.keys())}, pydantic BaseModel subclasses, "
+    "lists and dictionaries of primitive types, or typing.Any. Check "
+    "https://mlflow.org/docs/latest/model/python_model.html#supported-type-hints for more details."
+)
+
+
+def _try_import_numpy():
+    try:
+        import numpy
+
+        return numpy
+    except ImportError:
+        return
+
+
+@lru_cache(maxsize=1)
+def type_hints_no_signature_inference():
+    """
+    This function returns a tuple of types that can be used
+    as type hints, but no schema can be inferred from them.
+
+    ..note::
+        These types can not be used as nested types in other type hints.
+    """
+    type_hints = ()
+    try:
+        import pandas as pd
+
+        type_hints += (
+            pd.DataFrame,
+            pd.Series,
+        )
+    except ImportError:
+        pass
+
+    try:
+        import numpy as np
+
+        type_hints += (np.ndarray,)
+    except ImportError:
+        pass
+
+    try:
+        from scipy.sparse import csc_matrix, csr_matrix
+
+        type_hints += (csc_matrix, csr_matrix)
+    except ImportError:
+        pass
+
+    return type_hints
+
+
+class ColSpecType(NamedTuple):
+    dtype: COLSPEC_TYPES
+    required: bool
+
+
+class UnsupportedTypeHintException(MlflowException):
+    def __init__(self, type_hint):
+        super().__init__(
+            f"Unsupported type hint `{_type_hint_repr(type_hint)}`. {SUPPORTED_TYPE_HINT_MSG}",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+
+class InvalidTypeHintException(MlflowException):
+    def __init__(self, *, message):
+        super().__init__(message, error_code=INVALID_PARAMETER_VALUE)
+
+
+def _signature_cannot_be_inferred_from_type_hint(type_hint: type[Any]) -> bool:
+    return type_hint in type_hints_no_signature_inference()
+
+
+def _is_type_hint_from_example(type_hint: type[Any]) -> bool:
+    return type_hint == TypeFromExample
+
+
+def _is_example_valid_for_type_from_example(example: Any) -> bool:
+    allowed_types = (list,)
+    try:
+        import pandas as pd
+
+        allowed_types += (pd.DataFrame, pd.Series)
+    except ImportError:
+        pass
+    return isinstance(example, allowed_types)
+
+
+def _convert_dataframe_to_example_format(data: Any, input_example: Any) -> Any:
+    import numpy as np
+    import pandas as pd
+
+    if isinstance(data, pd.DataFrame):
+        if isinstance(input_example, pd.DataFrame):
+            return data
+        if isinstance(input_example, pd.Series):
+            data = data.iloc[:, 0]
+            data.name = input_example.name
+            return data
+        if np.isscalar(input_example):
+            return data.iloc[0, 0]
+        if isinstance(input_example, dict):
+            if len(data) == 1:
+                return data.to_dict(orient="records")[0]
+            else:
+                # This case shouldn't happen
+                _logger.warning("Cannot convert DataFrame to a single dictionary.")
+                return data
+        if isinstance(input_example, list):
+            # list[scalar]
+            if len(data.columns) == 1 and all(np.isscalar(x) for x in input_example):
+                return data.iloc[:, 0].tolist()
+            else:
+                # NB: there are some cases that this doesn't work well, but it's the best we can do
+                # e.g. list of dictionaries with different keys
+                # [{"a": 1}, {"b": 2}] -> pd.DataFrame(...) during schema enforcement
+                # here -> [{'a': 1.0, 'b': nan}, {'a': nan, 'b': 2.0}]
+                return data.to_dict(orient="records")
+
+    return data
+
+
+def _infer_colspec_type_from_type_hint(type_hint: type[Any]) -> ColSpecType:
+    """
+    Infer the ColSpec type from a type hint.
+    The inferred dtype should be one of the supported data types in COLSPEC_TYPES.
+    """
+    if type_hint == Any:
+        color_warning(
+            message="Any type hint is inferred as AnyType, and MLflow doesn't validate the data "
+            "for this type. Please use a more specific type hint to enable data validation.",
+            stacklevel=2,
+            color="yellow_bold",
+        )
+        return ColSpecType(dtype=AnyType(), required=True)
+    if datatype := TYPE_HINTS_TO_DATATYPE_MAPPING.get(type_hint):
+        return ColSpecType(dtype=datatype, required=True)
+    elif _is_pydantic_type_hint(type_hint):
+        dtype = _infer_type_from_pydantic_model(type_hint)
+        return ColSpecType(dtype=dtype, required=True)
+    elif origin_type := get_origin(type_hint):
+        args = get_args(type_hint)
+        if origin_type is list:
+            internal_type = _get_element_type_of_list_type_hint(type_hint)
+            return ColSpecType(
+                dtype=Array(_infer_colspec_type_from_type_hint(type_hint=internal_type).dtype),
+                required=True,
+            )
+        if origin_type is dict:
+            if len(args) == 2:
+                if args[0] != str:
+                    raise InvalidTypeHintException(
+                        message=f"Dictionary key type must be str, got {args[0]} in type hint "
+                        f"{_type_hint_repr(type_hint)}"
+                    )
+                return ColSpecType(
+                    dtype=Map(_infer_colspec_type_from_type_hint(type_hint=args[1]).dtype),
+                    required=True,
+                )
+            raise InvalidTypeHintException(
+                message="Dictionary type hint must contain two element types, got "
+                f"{_type_hint_repr(type_hint)}"
+            )
+        if origin_type in UNION_TYPES:
+            if NONE_TYPE in args:
+                # This case shouldn't happen, but added for completeness
+                if len(args) < 2:
+                    raise InvalidTypeHintException(
+                        message=f"Union type hint must contain at least one non-None type, "
+                        f"got {_type_hint_repr(type_hint)}"
+                    )
+                # Optional type
+                elif len(args) == 2:
+                    effective_type = next((arg for arg in args if arg is not NONE_TYPE), None)
+                    return ColSpecType(
+                        dtype=_infer_colspec_type_from_type_hint(effective_type).dtype,
+                        required=False,
+                    )
+                # Optional Union type
+                else:
+                    _logger.warning(
+                        "Union type hint with multiple non-None types is inferred as AnyType, "
+                        "and MLflow doesn't validate the data against its element types."
+                    )
+                    return ColSpecType(dtype=AnyType(), required=False)
+            # Union type with all valid types is matched as AnyType
+            else:
+                _logger.warning(
+                    "Union type hint is inferred as AnyType, and MLflow doesn't validate the data "
+                    "against its element types."
+                )
+                return ColSpecType(dtype=AnyType(), required=True)
+    _raise_type_hint_error(type_hint)
+
+
+def _raise_type_hint_error(type_hint: type[Any]) -> None:
+    if (
+        type_hint
+        in (
+            list,
+            dict,
+            Optional,
+        )
+        + UNION_TYPES
+    ):
+        raise InvalidTypeHintException(
+            message=f"Invalid type hint `{_type_hint_repr(type_hint)}`, it must include "
+            f"a valid element type. {SUPPORTED_TYPE_HINT_MSG}"
+        )
+    raise UnsupportedTypeHintException(type_hint=type_hint)
+
+
+def _infer_type_from_pydantic_model(model: pydantic.BaseModel) -> Object:
+    """
+    Infer the object schema from a pydantic model.
+    """
+    if _is_pydantic_type_hint(model):
+        fields = model_fields(model)
+    else:
+        raise TypeError(f"model must be a Pydantic model class, but got {type(model)}")
+
+    properties = []
+    invalid_fields = []
+    for field_name, field_info in fields.items():
+        annotation = field_info.annotation
+        # this shouldn't happen since pydantic has checks for missing annotations
+        # but added here to avoid potential edge cases
+        if annotation is None:
+            invalid_fields.append(field_name)
+            continue
+        colspec_type = _infer_colspec_type_from_type_hint(annotation)
+        if colspec_type.required is False and field_required(field_info):
+            raise InvalidTypeHintException(
+                message=f"Optional field `{field_name}` in Pydantic model `{model.__name__}` "
+                "doesn't have a default value. Please set default value to None for this field."
+            )
+        properties.append(
+            Property(
+                name=field_name,
+                dtype=colspec_type.dtype,
+                required=colspec_type.required,
+            )
+        )
+    if invalid_fields:
+        raise InvalidTypeHintException(
+            message="The following fields in the Pydantic model do not have type annotations: "
+            f"{invalid_fields}. Please add type annotations to these fields."
+        )
+
+    return Object(properties=properties)
+
+
+def _is_pydantic_type_hint(type_hint: type[Any]) -> bool:
+    try:
+        return issubclass(type_hint, pydantic.BaseModel)
+    # inspect.isclass(dict[str, int]) is True, but issubclass raises a TypeError
+    except TypeError:
+        return False
+
+
+def model_fields(
+    model: pydantic.BaseModel,
+) -> dict[str, type[FIELD_TYPE]]:
+    if IS_PYDANTIC_V2_OR_NEWER:
+        return model.model_fields
+    return model.__fields__
+
+
+def model_validate(model: pydantic.BaseModel, values: Any) -> None:
+    if IS_PYDANTIC_V2_OR_NEWER:
+        # use strict mode to avoid any data conversion here
+        # e.g. "123" will not be converted to 123 if the type is int
+        model.model_validate(values, strict=True)
+    else:
+        model.validate(values)
+
+
+def field_required(field: type[FIELD_TYPE]) -> bool:
+    if IS_PYDANTIC_V2_OR_NEWER:
+        return field.is_required()
+    return field.required
+
+
+def _get_element_type_of_list_type_hint(type_hint: type[list[Any]]) -> Any:
+    """
+    Get the element type of list[...] type hint
+    """
+    args = get_args(type_hint)
+    # Optional[list[...]]
+    if type(None) in args:
+        raise MlflowException.invalid_parameter_value(OPTIONAL_INPUT_MSG)
+    # a valid list[...] type hint must only contain one argument
+    if len(args) == 0:
+        raise InvalidTypeHintException(
+            message=f"Type hint `{_type_hint_repr(type_hint)}` doesn't contain a collection "
+            "element type. Fix by adding an element type to the collection type definition, "
+            "e.g. `list[str]` instead of `list`."
+        )
+    if len(args) > 1:
+        raise InvalidTypeHintException(
+            message=f"Type hint `{_type_hint_repr(type_hint)}` contains {len(args)} element types. "
+            "Collections must have only a single type definition e.g. `list[int]` is valid; "
+            "`list[str, int]` is invalid."
+        )
+    return args[0]
+
+
+def _is_list_type_hint(type_hint: type[Any]) -> bool:
+    origin_type = _get_origin_type(type_hint)
+    return type_hint == list or origin_type is list
+
+
+def _infer_schema_from_list_type_hint(type_hint: type[list[Any]]) -> Schema:
+    """
+    Infer schema from a list type hint.
+    The type hint must be list[...], and the inferred schema contains a
+    single ColSpec, where the type is based on the element type of the list type hint,
+    since ColSpec represents a column's data type of the dataset.
+    e.g. list[int] -> Schema([ColSpec(type=DataType.long, required=True)])
+    A valid `predict` function of a pyfunc model must use list type hint for the input.
+    """
+    if not _is_list_type_hint(type_hint):
+        # This should be invalid, but to keep backwards compatibility of ChatCompletionRequest
+        # type hint used in some rag models, we raise UnsupportedTypeHintException here
+        # so that the model with such type hint can still be logged
+        raise MlflowException.invalid_parameter_value(
+            message="Type hints must be wrapped in list[...] because MLflow assumes the "
+            "predict method to take multiple input instances. Specify your type hint as "
+            f"`list[{_type_hint_repr(type_hint)}]` for a valid signature."
+        )
+    internal_type = _get_element_type_of_list_type_hint(type_hint)
+    return _infer_schema_from_type_hint(internal_type)
+
+
+def _infer_schema_from_type_hint(type_hint: type[Any]) -> Schema:
+    col_spec_type = _infer_colspec_type_from_type_hint(type_hint)
+    # Creating Schema with unnamed optional inputs is not supported
+    if col_spec_type.required is False:
+        raise InvalidTypeHintException(message=OPTIONAL_INPUT_MSG)
+    return Schema([ColSpec(type=col_spec_type.dtype, required=col_spec_type.required)])
+
+
+def _validate_data_against_type_hint(data: Any, type_hint: type[Any]) -> Any:
+    """
+    Validate the data against provided type hint.
+    The allowed conversions are:
+        dictionary data with Pydantic model type hint -> Pydantic model instance
+
+    Args:
+        data: The data to validate
+        type_hint: The type hint to validate against
+    """
+    if _is_pydantic_type_hint(type_hint):
+        # if data is a pydantic model instance, convert it to a dictionary for validation
+        if isinstance(data, pydantic.BaseModel):
+            data_dict = model_dump_compat(data)
+        elif isinstance(data, dict):
+            data_dict = data
+        else:
+            raise MlflowException.invalid_parameter_value(
+                "Expecting example to be a dictionary or pydantic model instance for "
+                f"Pydantic type hint, got {type(data)}"
+            )
+        try:
+            model_validate(type_hint, data_dict)
+        except pydantic.ValidationError as e:
+            raise MlflowException.invalid_parameter_value(
+                message=f"Data doesn't match type hint, error: {e}. Expected fields in the "
+                f"type hint: {model_fields(type_hint)}; passed data: {data_dict}. Check "
+                "https://mlflow.org/docs/latest/model/python_model.html#pydantic-model-type-hints-data-conversion"
+                " for more details.",
+            ) from e
+        else:
+            return type_hint(**data_dict) if isinstance(data, dict) else data
+    elif type_hint == Any:
+        return data
+    elif type_hint in TYPE_HINTS_TO_DATATYPE_MAPPING:
+        if _MLFLOW_IS_IN_SERVING_ENVIRONMENT.get():
+            data = _parse_data_for_datatype_hint(data=data, type_hint=type_hint)
+        if isinstance(data, type_hint):
+            return data
+        raise MlflowException.invalid_parameter_value(
+            f"Expected type {_type_hint_repr(type_hint)}, but got {type(data).__name__}"
+        )
+    elif origin_type := get_origin(type_hint):
+        args = get_args(type_hint)
+        if origin_type is list:
+            return _validate_list_elements(element_type=args[0], data=data)
+        elif origin_type is dict:
+            return _validate_dict_elements(element_type=args[1], data=data)
+        elif origin_type in UNION_TYPES:
+            # Optional type
+            if NONE_TYPE in args:
+                if data is None:
+                    return data
+                if len(args) == 2:
+                    effective_type = next((arg for arg in args if arg is not NONE_TYPE), None)
+                    return _validate_data_against_type_hint(data=data, type_hint=effective_type)
+            # Union type with all valid types is matched as AnyType
+            # no validation needed for AnyType
+            return data
+    _raise_type_hint_error(type_hint)
+
+
+def _parse_data_for_datatype_hint(data: Any, type_hint: type[Any]) -> Any:
+    """
+    Parse the data based on the type hint.
+    This should only be used in MLflow serving environment to convert
+    json data to the expected format.
+    Allowed conversions:
+        - string data with datetime type hint -> datetime object
+        - string data with bytes type hint -> bytes object
+    """
+    if type_hint == bytes and isinstance(data, str):
+        # The assumption is that the data is base64 encoded, and
+        # scoring server accepts base64 encoded string for bytes fields.
+        # MLflow uses the same method for saving input example
+        # via base64.encodebytes(x).decode("ascii")
+        return base64.decodebytes(bytes(data, "utf8"))
+    if type_hint == datetime and isinstance(data, str):
+        # The assumption is that the data is in ISO format
+        return datetime.fromisoformat(data)
+    return data
+
+
+class ValidationResult(NamedTuple):
+    value: Optional[Any] = None
+    error_message: Optional[str] = None
+
+
+def _get_data_validation_result(data: Any, type_hint: type[Any]) -> ValidationResult:
+    try:
+        value = _validate_data_against_type_hint(data=data, type_hint=type_hint)
+        return ValidationResult(value=value)
+    except MlflowException as e:
+        return ValidationResult(error_message=e.message)
+
+
+def _type_hint_repr(type_hint: type[Any]) -> str:
+    return (
+        type_hint.__name__
+        if _is_pydantic_type_hint(type_hint) or type(type_hint) == type
+        else str(type_hint)
+    )
+
+
+def _validate_list_elements(element_type: type[Any], data: Any) -> list[Any]:
+    if not isinstance(data, list):
+        raise MlflowException.invalid_parameter_value(
+            f"Expected list, but got {type(data).__name__}"
+        )
+    invalid_elems = []
+    result = []
+    for elem in data:
+        validation_result = _get_data_validation_result(data=elem, type_hint=element_type)
+        if validation_result.error_message:
+            invalid_elems.append((str(elem), validation_result.error_message))
+        else:
+            result.append(validation_result.value)
+    if invalid_elems:
+        invalid_elems_msg = (
+            f"{invalid_elems[:5]} ... (truncated)" if len(invalid_elems) > 5 else invalid_elems
+        )
+        raise MlflowException.invalid_parameter_value(
+            f"Failed to validate data against type hint `list[{_type_hint_repr(element_type)}]`, "
+            f"invalid elements: {invalid_elems_msg}"
+        )
+    return result
+
+
+def _validate_dict_elements(element_type: type[Any], data: Any) -> dict[str, Any]:
+    if not isinstance(data, dict):
+        raise MlflowException.invalid_parameter_value(
+            f"Expected dict, but got {type(data).__name__}"
+        )
+    invalid_elems = {}
+    result = {}
+    for key, value in data.items():
+        if not isinstance(key, str):
+            invalid_elems[str(key)] = f"Key must be a string, got {type(key).__name__}"
+            continue
+        validation_result = _get_data_validation_result(data=value, type_hint=element_type)
+        if validation_result.error_message:
+            invalid_elems[key] = validation_result.error_message
+        else:
+            result[key] = validation_result.value
+    if invalid_elems:
+        raise MlflowException.invalid_parameter_value(
+            f"Failed to validate data against type hint "
+            f"`dict[str, {_type_hint_repr(element_type)}]`, "
+            f"invalid elements: {invalid_elems}"
+        )
+    return result
+
+
+def _get_origin_type(type_hint: type[Any]) -> Any:
+    """
+    Get the origin type of a type hint.
+    If the type hint is Union type, return the origin type of the effective type.
+    If the type hint is Union type with multiple effective types, return Any.
+    """
+    origin_type = get_origin(type_hint)
+    if origin_type in UNION_TYPES:
+        args = get_args(type_hint)
+        if NONE_TYPE in args and len(args) == 2:
+            effective_type = next((arg for arg in args if arg is not NONE_TYPE), None)
+            return _get_origin_type(effective_type)
+        else:
+            # Union types match Any
+            return Any
+    return origin_type
+
+
+def _convert_data_to_type_hint(data: Any, type_hint: type[Any]) -> Any:
+    """
+    Convert data to the expected format based on the type hint.
+    This function is used in data validation of @pyfunc to support compatibility with
+    functions such as mlflow.evaluate and spark_udf since they accept pandas DF as input.
+    NB: the input pandas DataFrame must contain a single column with the same type as the type hint.
+    Supported conversions:
+        - pandas DataFrame with a single column + list[...] type hint -> list
+        - pandas DataFrame with multiple columns + list[dict[...]] type hint -> list[dict[...]]
+    """
+    import pandas as pd
+
+    result = data
+    if isinstance(data, pd.DataFrame) and type_hint != pd.DataFrame:
+        origin_type = _get_origin_type(type_hint)
+        if origin_type is not list:
+            raise MlflowException(
+                "Only `list[...]` type hint supports pandas DataFrame input "
+                f"with a single column. But got {_type_hint_repr(type_hint)}."
+            )
+        element_type = _get_element_type_of_list_type_hint(type_hint)
+        # This is needed for list[dict] or list[pydantic.BaseModel] type hints
+        # since the data can be converted to pandas DataFrame with multiple columns
+        # inside spark_udf
+        if element_type is dict or _is_pydantic_type_hint(element_type):
+            # if the column is 0, then each row is a dictionary
+            if list(data.columns) == [0]:
+                result = data.iloc[:, 0].tolist()
+            else:
+                result = data.to_dict(orient="records")
+        else:
+            if len(data.columns) != 1:
+                # TODO: remove the warning and raise Exception once the bug about evaluate
+                # DF containing multiple columns is fixed
+                _logger.warning(
+                    "`predict` function with list[...] type hints of non-dictionary collection "
+                    "type only supports pandas DataFrame with a single column. But got "
+                    f"{len(data.columns)} columns. The data will be converted to a list "
+                    "of the first column."
+                )
+            result = data.iloc[:, 0].tolist()
+        # only sanitize the data when it's converted from pandas DataFrame
+        # since spark_udf implicitly converts lists into numpy arrays
+        return _sanitize_data(result)
+
+    return data
+
+
+def _sanitize_data(data: Any) -> Any:
+    """
+    Sanitize the data by converting any numpy lists to Python lists.
+    This is needed because spark_udf (pandas_udf) implicitly converts lists into numpy arrays.
+
+    For example, below udf demonstrates the behavior:
+        df = spark.createDataFrame(pd.DataFrame({"input": [["a", "b"], ["c", "d"]]}))
+        @pandas_udf(ArrayType(StringType()))
+        def my_udf(input_series: pd.Series) -> pd.Series:
+            print(type(input_series.iloc[0]))
+        df.withColumn("output", my_udf("input")).show()
+    """
+    if np := _try_import_numpy():
+        if isinstance(data, np.ndarray):
+            data = data.tolist()
+        if isinstance(data, list):
+            data = [_sanitize_data(elem) for elem in data]
+        if isinstance(data, dict):
+            data = {key: _sanitize_data(value) for key, value in data.items()}
+        if isinstance(data, float) and np.isnan(data):
+            data = None
+    return data