genesis-flow 1.0.0__py3-none-any.whl

mlflow/pyfunc/scoring_server/client.py

@@ -0,0 +1,146 @@
+import json
+import logging
+import tempfile
+import time
+import uuid
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Optional
+
+import requests
+
+from mlflow.deployments import PredictionsResponse
+from mlflow.environment_variables import MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT
+from mlflow.exceptions import MlflowException
+from mlflow.pyfunc import scoring_server
+from mlflow.utils.proto_json_utils import dump_input_data
+
+_logger = logging.getLogger(__name__)
+
+
+class BaseScoringServerClient(ABC):
+    @abstractmethod
+    def wait_server_ready(self, timeout=30, scoring_server_proc=None):
+        """
+        Wait until the scoring server is ready to accept requests.
+        """
+
+    @abstractmethod
+    def invoke(self, data, params: Optional[dict[str, Any]] = None):
+        """
+        Invoke inference on input data. The input data must be pandas dataframe or numpy array or
+        a dict of numpy arrays.
+
+        Args:
+            data: Model input data.
+            params: Additional parameters to pass to the model for inference.
+
+        Returns:
+            Prediction result.
+        """
+
+
+class ScoringServerClient(BaseScoringServerClient):
+    def __init__(self, host, port):
+        self.url_prefix = f"http://{host}:{port}"
+
+    def ping(self):
+        ping_status = requests.get(url=self.url_prefix + "/ping")
+        if ping_status.status_code != 200:
+            raise Exception(f"ping failed (error code {ping_status.status_code})")
+
+    def get_version(self):
+        resp_status = requests.get(url=self.url_prefix + "/version")
+        if resp_status.status_code != 200:
+            raise Exception(f"version failed (error code {resp_status.status_code})")
+        return resp_status.text
+
+    def wait_server_ready(self, timeout=30, scoring_server_proc=None):
+        begin_time = time.time()
+
+        while True:
+            time.sleep(0.3)
+            try:
+                self.ping()
+                return
+            except Exception:
+                pass
+            if time.time() - begin_time > timeout:
+                break
+            if scoring_server_proc is not None:
+                return_code = scoring_server_proc.poll()
+                if return_code is not None:
+                    raise RuntimeError(f"Server process already exit with returncode {return_code}")
+        raise RuntimeError("Wait scoring server ready timeout.")
+
+    def invoke(self, data, params: Optional[dict[str, Any]] = None):
+        """
+        Args:
+            data: Model input data.
+            params: Additional parameters to pass to the model for inference.
+
+        Returns:
+            :py:class:`PredictionsResponse <mlflow.deployments.PredictionsResponse>` result.
+        """
+        response = requests.post(
+            url=self.url_prefix + "/invocations",
+            data=dump_input_data(data, params=params),
+            headers={"Content-Type": scoring_server.CONTENT_TYPE_JSON},
+        )
+        if response.status_code != 200:
+            raise Exception(
+                f"Invocation failed (error code {response.status_code}, response: {response.text})"
+            )
+        return PredictionsResponse.from_json(response.text)
+
+
+class StdinScoringServerClient(BaseScoringServerClient):
+    def __init__(self, process):
+        super().__init__()
+        self.process = process
+        self.tmpdir = Path(tempfile.mkdtemp())
+        self.output_json = self.tmpdir.joinpath("output.json")
+
+    def wait_server_ready(self, timeout=30, scoring_server_proc=None):
+        return_code = self.process.poll()
+        if return_code is not None:
+            raise RuntimeError(f"Server process already exit with returncode {return_code}")
+
+    def invoke(self, data, params: Optional[dict[str, Any]] = None):
+        """
+        Invoke inference on input data. The input data must be pandas dataframe or numpy array or
+        a dict of numpy arrays.
+
+        Args:
+            data: Model input data.
+            params: Additional parameters to pass to the model for inference.
+
+        Returns:
+            :py:class:`PredictionsResponse <mlflow.deployments.PredictionsResponse>` result.
+        """
+        if not self.output_json.exists():
+            self.output_json.touch()
+
+        request_id = str(uuid.uuid4())
+        request = {
+            "id": request_id,
+            "data": dump_input_data(data, params=params),
+            "output_file": str(self.output_json),
+        }
+        self.process.stdin.write(json.dumps(request) + "\n")
+        self.process.stdin.flush()
+
+        begin_time = time.time()
+        while True:
+            _logger.info("Waiting for scoring to complete...")
+            try:
+                with self.output_json.open(mode="r+") as f:
+                    resp = PredictionsResponse.from_json(f.read())
+                    if resp.get("id") == request_id:
+                        f.truncate(0)
+                        return resp
+            except Exception as e:
+                _logger.debug("Exception while waiting for scoring to complete: %s", e)
+            if time.time() - begin_time > MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.get():
+                raise MlflowException("Scoring timeout")
+            time.sleep(1)

mlflow/pyfunc/spark_model_cache.py

@@ -0,0 +1,48 @@
+from mlflow.utils._spark_utils import _SparkDirectoryDistributor
+
+
+class SparkModelCache:
+    """Caches models in memory on Spark Executors, to avoid continually reloading from disk.
+
+    This class has to be part of a different module than the one that _uses_ it. This is
+    because Spark will pickle classes that are defined in the local scope, but relies on
+    Python's module loading behavior for classes in different modules. In this case, we
+    are relying on the fact that Python will load a module at-most-once, and can therefore
+    store per-process state in a static map.
+    """
+
+    # Map from unique name --> (loaded model, local_model_path).
+    _models = {}
+
+    # Number of cache hits we've had, for testing purposes.
+    _cache_hits = 0
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def add_local_model(spark, model_path):
+        """Given a SparkSession and a model_path which refers to a pyfunc directory locally,
+        we will zip the directory up, enable it to be distributed to executors, and return
+        the "archive_path", which should be used as the path in get_or_load().
+        """
+        return _SparkDirectoryDistributor.add_dir(spark, model_path)
+
+    @staticmethod
+    def get_or_load(archive_path):
+        """Given a path returned by add_local_model(), this method will return a tuple of
+        (loaded_model, local_model_path).
+        If this Python process ever loaded the model before, we will reuse that copy.
+        """
+        if archive_path in SparkModelCache._models:
+            SparkModelCache._cache_hits += 1
+            return SparkModelCache._models[archive_path]
+
+        local_model_dir = _SparkDirectoryDistributor.get_or_extract(archive_path)
+
+        # We must rely on a supposed cyclic import here because we want this behavior
+        # on the Spark Executors (i.e., don't try to pickle the load_model function).
+        from mlflow.pyfunc import load_model
+
+        SparkModelCache._models[archive_path] = (load_model(local_model_dir), local_model_dir)
+        return SparkModelCache._models[archive_path]

mlflow/pyfunc/stdin_server.py

@@ -0,0 +1,44 @@
+import argparse
+import inspect
+import json
+import logging
+import sys
+
+from mlflow.pyfunc import scoring_server
+from mlflow.pyfunc.model import _log_warning_if_params_not_in_predict_signature
+
+_logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model-uri")
+args = parser.parse_args()
+
+_logger.info("Loading model from %s", args.model_uri)
+
+model = scoring_server.load_model_with_mlflow_config(args.model_uri)
+input_schema = model.metadata.get_input_schema()
+_logger.info("Loaded model")
+
+_logger.info("Waiting for request")
+for line in sys.stdin:
+    _logger.info("Received request")
+    request = json.loads(line)
+
+    _logger.info("Parsing input data")
+    data = request["data"]
+    data, params = scoring_server._split_data_and_params(data)
+    data = scoring_server.infer_and_parse_data(data, input_schema)
+
+    _logger.info("Making predictions")
+    if "params" in inspect.signature(model.predict).parameters:
+        preds = model.predict(data, params=params)
+    else:
+        _log_warning_if_params_not_in_predict_signature(_logger, params)
+        preds = model.predict(data)
+
+    _logger.info("Writing predictions")
+    with open(request["output_file"], "a") as f:
+        scoring_server.predictions_to_json(preds, f, {"id": request["id"]})
+
+    _logger.info("Done")

mlflow/pyfunc/utils/__init__.py

@@ -0,0 +1,3 @@
+from mlflow.pyfunc.utils.data_validation import pyfunc
+
+__all__ = ["pyfunc"]

mlflow/pyfunc/utils/data_validation.py

@@ -0,0 +1,224 @@
+import inspect
+import warnings
+from functools import lru_cache, wraps
+from typing import Any, NamedTuple, Optional
+
+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: Optional[type[Any]]
+    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: Optional[FuncInfo]):
+    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) -> Optional[FuncInfo]:
+    """
+    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

mlflow/pyfunc/utils/environment.py

@@ -0,0 +1,22 @@
+import os
+from contextlib import contextmanager
+
+from mlflow.environment_variables import _MLFLOW_IS_IN_SERVING_ENVIRONMENT
+
+
+@contextmanager
+def _simulate_serving_environment():
+    """
+    Some functions (e.g. validate_serving_input) replicate the data transformation logic
+    that happens in the model serving environment to validate data before model deployment.
+    This context manager can be used to simulate the serving environment for such functions.
+    """
+    original_value = _MLFLOW_IS_IN_SERVING_ENVIRONMENT.get_raw()
+    try:
+        _MLFLOW_IS_IN_SERVING_ENVIRONMENT.set("true")
+        yield
+    finally:
+        if original_value is not None:
+            os.environ[_MLFLOW_IS_IN_SERVING_ENVIRONMENT.name] = original_value
+        else:
+            del os.environ[_MLFLOW_IS_IN_SERVING_ENVIRONMENT.name]

mlflow/pyfunc/utils/input_converter.py

@@ -0,0 +1,47 @@
+from dataclasses import fields, is_dataclass
+from typing import Union, get_args, get_origin
+
+from mlflow.utils.annotations import experimental
+
+
+def _is_optional_dataclass(field_type) -> bool:
+    """
+    Check if the field type is an Optional containing a dataclass.
+    Currently, ... | None (in Python 3.10) is not supported.
+    """
+    if get_origin(field_type) is Union:
+        inner_types = get_args(field_type)
+        # Check if it's a Union[Dataclass, NoneType] (i.e., Optional[Dataclass])
+        if len(inner_types) == 2 and any(t is type(None) for t in inner_types):
+            effective_type = next(t for t in get_args(field_type) if t is not type(None))
+            return is_dataclass(effective_type)
+    return False
+
+
+@experimental(version="2.18.0")
+def _hydrate_dataclass(dataclass_type, data):
+    """Recursively create an instance of the dataclass_type from data."""
+    if not (is_dataclass(dataclass_type) or _is_optional_dataclass(dataclass_type)):
+        raise ValueError(f"{dataclass_type.__name__} is not a dataclass")
+    if data is None:
+        return None
+
+    field_names = {f.name: f.type for f in fields(dataclass_type)}
+    kwargs = {}
+    for key, field_type in field_names.items():
+        if key in data:
+            value = data[key]
+            if is_dataclass(field_type):
+                kwargs[key] = _hydrate_dataclass(field_type, value)
+            elif _is_optional_dataclass(field_type):
+                effective_type = next(t for t in get_args(field_type) if t is not type(None))
+                kwargs[key] = _hydrate_dataclass(effective_type, value)
+            elif get_origin(field_type) == list:
+                item_type = get_args(field_type)[0]
+                if is_dataclass(item_type):
+                    kwargs[key] = [_hydrate_dataclass(item_type, item) for item in value]
+                else:
+                    kwargs[key] = value
+            else:
+                kwargs[key] = value
+    return dataclass_type(**kwargs)

mlflow/pyfunc/utils/serving_data_parser.py

@@ -0,0 +1,11 @@
+from typing import Any
+
+# Support unwrapped JSON with these keys for LLM use cases of Chat, Completions, Embeddings tasks
+LLM_CHAT_KEY = "messages"
+LLM_COMPLETIONS_KEY = "prompt"
+LLM_EMBEDDINGS_KEY = "input"
+SUPPORTED_LLM_FORMATS = {LLM_CHAT_KEY, LLM_COMPLETIONS_KEY, LLM_EMBEDDINGS_KEY}
+
+
+def is_unified_llm_input(json_input: dict[str, Any]):
+    return any(x in json_input for x in SUPPORTED_LLM_FORMATS)