wehub-resource-sync
226 lines
8.5 KiB
Python
226 lines
8.5 KiB
Python
import inspect
|
|
import warnings
|
|
from functools import lru_cache, wraps
|
|
from typing import Any, NamedTuple
|
|
|
|
import pydantic
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models.signature import (
|
|
_extract_type_hints,
|
|
_is_context_in_predict_function_signature,
|
|
)
|
|
from mlflow.types.type_hints import (
|
|
InvalidTypeHintException,
|
|
_convert_data_to_type_hint,
|
|
_infer_schema_from_list_type_hint,
|
|
_is_type_hint_from_example,
|
|
_signature_cannot_be_inferred_from_type_hint,
|
|
_validate_data_against_type_hint,
|
|
model_validate,
|
|
)
|
|
from mlflow.utils.annotations import filter_user_warnings_once
|
|
from mlflow.utils.warnings_utils import color_warning
|
|
|
|
_INVALID_SIGNATURE_ERROR_MSG = (
|
|
"Model's `{func_name}` method contains invalid parameters: {invalid_params}. "
|
|
"Only the following parameter names are allowed: context, model_input, and params. "
|
|
"Note that invalid parameters will no longer be permitted in future versions."
|
|
)
|
|
|
|
|
|
class FuncInfo(NamedTuple):
|
|
input_type_hint: type[Any] | None
|
|
input_param_name: str
|
|
|
|
|
|
def pyfunc(func):
|
|
"""
|
|
A decorator that forces data validation against type hint of the input data
|
|
in the wrapped method. It is no-op if the type hint is not supported by MLflow.
|
|
|
|
.. note::
|
|
The function that applies this decorator must be a valid `predict` function
|
|
of `mlflow.pyfunc.PythonModel`, or a callable that takes a single input.
|
|
"""
|
|
|
|
func_info = _get_func_info_if_type_hint_supported(func)
|
|
return _wrap_predict_with_pyfunc(func, func_info)
|
|
|
|
|
|
def _wrap_predict_with_pyfunc(func, func_info: FuncInfo | None):
|
|
if func_info is not None:
|
|
model_input_index = _model_input_index_in_function_signature(func)
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
args, kwargs = _validate_model_input(
|
|
args,
|
|
kwargs,
|
|
model_input_index,
|
|
func_info.input_type_hint,
|
|
func_info.input_param_name,
|
|
)
|
|
except Exception as e:
|
|
if isinstance(e, MlflowException):
|
|
raise e
|
|
raise MlflowException(
|
|
"Failed to validate the input data against the type hint "
|
|
f"`{func_info.input_type_hint}`. Error: {e}"
|
|
)
|
|
return func(*args, **kwargs)
|
|
|
|
else:
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
return func(*args, **kwargs)
|
|
|
|
wrapper._is_pyfunc = True
|
|
return wrapper
|
|
|
|
|
|
def wrap_non_list_predict_pydantic(func, input_pydantic_model, validation_error_msg, unpack=False):
|
|
"""
|
|
Used by MLflow defined subclasses of PythonModel that have non-list a pydantic model as input.
|
|
Takes in a dict input, validates it against `input_pydantic_model`, and then creates
|
|
the pydantic model.
|
|
|
|
If `unpack` is True, the validated dict is parsed into the function arguments.
|
|
Otherwise, the whole pydantic object is passed to the function.
|
|
|
|
Args:
|
|
func: The predict/predict_stream method of the PythonModel subclass.
|
|
input_pydantic_model: The pydantic model that the input should be validated against.
|
|
validation_error_msg: The error message to raise if the dict input fails to validate.
|
|
unpack: Whether to unpack the validated dict into the function arguments. Defaults to False.
|
|
|
|
Raises:
|
|
MlflowException: If the input fails to validate against the pydantic model.
|
|
|
|
Returns:
|
|
A function that can take either a dict input or a pydantic object as input.
|
|
"""
|
|
|
|
@wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
if len(args) == 1 and isinstance(args[0], dict):
|
|
try:
|
|
model_validate(input_pydantic_model, args[0])
|
|
pydantic_obj = input_pydantic_model(**args[0])
|
|
except pydantic.ValidationError as e:
|
|
raise MlflowException(
|
|
f"{validation_error_msg} Pydantic validation error: {e}"
|
|
) from e
|
|
else:
|
|
if unpack:
|
|
param_names = inspect.signature(func).parameters.keys() - {"self"}
|
|
kwargs = {k: getattr(pydantic_obj, k) for k in param_names}
|
|
return func(self, **kwargs)
|
|
else:
|
|
return func(self, pydantic_obj)
|
|
else:
|
|
# Before logging, this is equivalent to the behavior from the raw predict method
|
|
# After logging, signature enforcement happens in the _convert_input method
|
|
# of the wrapper class
|
|
return func(self, *args, **kwargs)
|
|
|
|
wrapper._is_pyfunc = True
|
|
return wrapper
|
|
|
|
|
|
def _check_func_signature(func, func_name) -> list[str]:
|
|
parameters = inspect.signature(func).parameters
|
|
param_names = [name for name in parameters.keys() if name != "self"]
|
|
if invalid_params := set(param_names) - {"self", "context", "model_input", "params"}:
|
|
warnings.warn(
|
|
_INVALID_SIGNATURE_ERROR_MSG.format(func_name=func_name, invalid_params=invalid_params),
|
|
FutureWarning,
|
|
stacklevel=2,
|
|
)
|
|
return param_names
|
|
|
|
|
|
@lru_cache
|
|
@filter_user_warnings_once
|
|
def _get_func_info_if_type_hint_supported(func) -> FuncInfo | None:
|
|
"""
|
|
Internal method to check if the predict function has type hints and if they are supported
|
|
by MLflow.
|
|
For PythonModel, the signature must be one of below:
|
|
- predict(self, context, model_input, params=None)
|
|
- predict(self, model_input, params=None)
|
|
For callables, the function must contain only one input argument.
|
|
"""
|
|
param_names = _check_func_signature(func, "predict")
|
|
input_arg_index = 1 if _is_context_in_predict_function_signature(func=func) else 0
|
|
type_hint = _extract_type_hints(func, input_arg_index=input_arg_index).input
|
|
input_param_name = param_names[input_arg_index]
|
|
if type_hint is not None:
|
|
if _signature_cannot_be_inferred_from_type_hint(type_hint) or _is_type_hint_from_example(
|
|
type_hint
|
|
):
|
|
return
|
|
try:
|
|
_infer_schema_from_list_type_hint(type_hint)
|
|
except InvalidTypeHintException as e:
|
|
raise MlflowException(
|
|
f"{e.message} To disable data validation, remove the type hint from the "
|
|
"predict function. Otherwise, fix the type hint."
|
|
)
|
|
# catch other exceptions to avoid breaking model usage
|
|
except Exception as e:
|
|
color_warning(
|
|
message="Type hint used in the model's predict function is not supported "
|
|
f"for MLflow's schema validation. {e} "
|
|
"Remove the type hint to disable this warning. "
|
|
"To enable validation for the input data, specify input example "
|
|
"or model signature when logging the model. ",
|
|
category=UserWarning,
|
|
stacklevel=3,
|
|
color="red",
|
|
)
|
|
else:
|
|
return FuncInfo(input_type_hint=type_hint, input_param_name=input_param_name)
|
|
else:
|
|
color_warning(
|
|
"Add type hints to the `predict` method to enable data validation "
|
|
"and automatic signature inference during model logging. "
|
|
"Check https://mlflow.org/docs/latest/model/python_model.html#type-hint-usage-in-pythonmodel"
|
|
" for more details.",
|
|
stacklevel=1,
|
|
color="yellow",
|
|
category=UserWarning,
|
|
)
|
|
|
|
|
|
def _model_input_index_in_function_signature(func):
|
|
parameters = inspect.signature(func).parameters
|
|
# we need to exclude the first argument if "self" is in the parameters
|
|
index = 1 if "self" in parameters else 0
|
|
if _is_context_in_predict_function_signature(parameters=parameters):
|
|
index += 1
|
|
return index
|
|
|
|
|
|
def _validate_model_input(
|
|
args, kwargs, model_input_index_in_sig, type_hint, model_input_param_name
|
|
):
|
|
model_input = None
|
|
input_pos = None
|
|
if model_input_param_name in kwargs:
|
|
model_input = kwargs[model_input_param_name]
|
|
input_pos = "kwargs"
|
|
elif len(args) >= model_input_index_in_sig + 1:
|
|
model_input = args[model_input_index_in_sig]
|
|
input_pos = model_input_index_in_sig
|
|
if input_pos is not None:
|
|
data = _convert_data_to_type_hint(model_input, type_hint)
|
|
data = _validate_data_against_type_hint(data, type_hint)
|
|
if input_pos == "kwargs":
|
|
kwargs[model_input_param_name] = data
|
|
else:
|
|
args = args[:input_pos] + (data,) + args[input_pos + 1 :]
|
|
return args, kwargs
|