genesis-flow 1.0.0__py3-none-any.whl

mlflow/pyfunc/scoring_server/__init__.py

@@ -0,0 +1,604 @@
+"""
+Scoring server for python model format.
+The passed int model is expected to have function:
+   predict(pandas.Dataframe) -> pandas.DataFrame
+
+Input, expected in text/csv or application/json format,
+is parsed into pandas.DataFrame and passed to the model.
+
+Defines four endpoints:
+    /ping used for health check
+    /health (same as /ping)
+    /version used for getting the mlflow version
+    /invocations used for scoring
+"""
+
+import asyncio
+import inspect
+import json
+import logging
+import os
+import shlex
+import sys
+import traceback
+from functools import wraps
+from typing import Any, NamedTuple, Optional
+
+from mlflow.environment_variables import (
+    _MLFLOW_IS_IN_SERVING_ENVIRONMENT,
+    MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT,
+)
+
+# NB: We need to be careful what we import form mlflow here. Scoring server is used from within
+# model's conda environment. The version of mlflow doing the serving (outside) and the version of
+# mlflow in the model's conda environment (inside) can differ. We should therefore keep mlflow
+# dependencies to the minimum here.
+# ALl of the mlflow dependencies below need to be backwards compatible.
+from mlflow.exceptions import MlflowException
+from mlflow.pyfunc.model import _log_warning_if_params_not_in_predict_signature
+from mlflow.types import ParamSchema, Schema
+from mlflow.utils import reraise
+from mlflow.utils.file_utils import path_to_local_file_uri
+from mlflow.utils.proto_json_utils import (
+    MlflowInvalidInputException,
+    NumpyEncoder,
+    _get_jsonable_obj,
+    dataframe_from_parsed_json,
+    parse_tf_serving_input,
+)
+from mlflow.version import VERSION
+
+try:
+    from mlflow.pyfunc import PyFuncModel, load_model
+except ImportError:
+    from mlflow.pyfunc import load_pyfunc as load_model
+from io import StringIO
+
+from mlflow.protos.databricks_pb2 import BAD_REQUEST, INVALID_PARAMETER_VALUE
+from mlflow.pyfunc.utils.serving_data_parser import is_unified_llm_input
+
+_SERVER_MODEL_PATH = "__pyfunc_model_path__"
+SERVING_MODEL_CONFIG = "SERVING_MODEL_CONFIG"
+
+CONTENT_TYPE_CSV = "text/csv"
+CONTENT_TYPE_JSON = "application/json"
+
+CONTENT_TYPES = [
+    CONTENT_TYPE_CSV,
+    CONTENT_TYPE_JSON,
+]
+
+_logger = logging.getLogger(__name__)
+
+DF_RECORDS = "dataframe_records"
+DF_SPLIT = "dataframe_split"
+INSTANCES = "instances"
+INPUTS = "inputs"
+
+SUPPORTED_FORMATS = {DF_RECORDS, DF_SPLIT, INSTANCES, INPUTS}
+SERVING_PARAMS_KEY = "params"
+
+REQUIRED_INPUT_FORMAT = (
+    f"The input must be a JSON dictionary with exactly one of the input fields {SUPPORTED_FORMATS}"
+)
+SCORING_PROTOCOL_CHANGE_INFO = (
+    "IMPORTANT: The MLflow Model scoring protocol has changed in MLflow version 2.0. If you are"
+    " seeing this error, you are likely using an outdated scoring request format. To resolve the"
+    " error, either update your request format or adjust your MLflow Model's requirements file to"
+    " specify an older version of MLflow (for example, change the 'mlflow' requirement specifier"
+    " to 'mlflow==1.30.0'). If you are making a request using the MLflow client"
+    " (e.g. via `mlflow.pyfunc.spark_udf()`), upgrade your MLflow client to a version >= 2.0 in"
+    " order to use the new request format. For more information about the updated MLflow"
+    " Model scoring protocol in MLflow 2.0, see"
+    " https://mlflow.org/docs/latest/models.html#deploy-mlflow-models."
+)
+
+
+def load_model_with_mlflow_config(model_uri):
+    extra_kwargs = {}
+    if model_config_json := os.environ.get(SERVING_MODEL_CONFIG):
+        extra_kwargs["model_config"] = json.loads(model_config_json)
+
+    return load_model(model_uri, **extra_kwargs)
+
+
+# Keep this method to maintain compatibility with MLServer
+# https://github.com/SeldonIO/MLServer/blob/caa173ab099a4ec002a7c252cbcc511646c261a6/runtimes/mlflow/mlserver_mlflow/runtime.py#L13C5-L13C31
+def infer_and_parse_json_input(json_input, schema: Schema = None):
+    """
+    Args:
+        json_input: A JSON-formatted string representation of TF serving input or a Pandas
+                    DataFrame, or a stream containing such a string representation.
+        schema: Optional schema specification to be used during parsing.
+    """
+    if isinstance(json_input, dict):
+        decoded_input = json_input
+    else:
+        try:
+            decoded_input = json.loads(json_input)
+        except json.decoder.JSONDecodeError as ex:
+            raise MlflowException(
+                message=(
+                    "Failed to parse input from JSON. Ensure that input is a valid JSON"
+                    f" formatted string. Error: '{ex}'. Input: \n{json_input}\n"
+                ),
+                error_code=BAD_REQUEST,
+            )
+    if isinstance(decoded_input, dict):
+        format_keys = set(decoded_input.keys()).intersection(SUPPORTED_FORMATS)
+        if len(format_keys) != 1:
+            message = f"Received dictionary with input fields: {list(decoded_input.keys())}"
+            raise MlflowException(
+                message=f"{REQUIRED_INPUT_FORMAT}. {message}. {SCORING_PROTOCOL_CHANGE_INFO}",
+                error_code=BAD_REQUEST,
+            )
+        input_format = format_keys.pop()
+        if input_format in (INSTANCES, INPUTS):
+            return parse_tf_serving_input(decoded_input, schema=schema)
+
+        elif input_format in (DF_SPLIT, DF_RECORDS):
+            # NB: skip the dataframe_ prefix
+            pandas_orient = input_format[10:]
+            return dataframe_from_parsed_json(
+                decoded_input[input_format], pandas_orient=pandas_orient, schema=schema
+            )
+    elif isinstance(decoded_input, list):
+        message = "Received a list"
+        raise MlflowException(
+            message=f"{REQUIRED_INPUT_FORMAT}. {message}. {SCORING_PROTOCOL_CHANGE_INFO}",
+            error_code=BAD_REQUEST,
+        )
+    else:
+        message = f"Received unexpected input type '{type(decoded_input)}'"
+        raise MlflowException(
+            message=f"{REQUIRED_INPUT_FORMAT}. {message}.", error_code=BAD_REQUEST
+        )
+
+
+def _decode_json_input(json_input):
+    """
+    Args:
+        json_input: A JSON-formatted string representation of TF serving input or a Pandas
+                    DataFrame, or a stream containing such a string representation.
+
+    Returns:
+        A dictionary representation of the JSON input.
+    """
+    if isinstance(json_input, dict):
+        return json_input
+
+    try:
+        decoded_input = json.loads(json_input)
+    except json.decoder.JSONDecodeError as ex:
+        raise MlflowInvalidInputException(
+            "Ensure that input is a valid JSON formatted string. "
+            f"Error: '{ex!r}'\nInput: \n{json_input}\n"
+        ) from ex
+
+    if isinstance(decoded_input, dict):
+        return decoded_input
+    if isinstance(decoded_input, list):
+        raise MlflowInvalidInputException(f"{REQUIRED_INPUT_FORMAT}. Received a list.")
+
+    raise MlflowInvalidInputException(
+        f"{REQUIRED_INPUT_FORMAT}. Received unexpected input type '{type(decoded_input)}."
+    )
+
+
+def _split_data_and_params_for_llm_input(json_input, param_schema: Optional[ParamSchema]):
+    data = {}
+    params = {}
+    schema_params = {param.name for param in param_schema.params} if param_schema else {}
+
+    for key, value in json_input.items():
+        # if the model defines a param schema, then we can add
+        # it to the params dict. otherwise, add it to the data
+        # dict to prevent it from being ignored at inference time
+        if key in schema_params:
+            params[key] = value
+        else:
+            data[key] = value
+
+    return data, params
+
+
+def _split_data_and_params(json_input):
+    input_dict = _decode_json_input(json_input)
+    data = {k: v for k, v in input_dict.items() if k in SUPPORTED_FORMATS}
+    params = input_dict.pop(SERVING_PARAMS_KEY, None)
+    return data, params
+
+
+def infer_and_parse_data(data, schema: Schema = None):
+    """
+    Args:
+        data: A dictionary representation of TF serving input or a Pandas
+            DataFrame, or a stream containing such a string representation.
+        schema: Optional schema specification to be used during parsing.
+    """
+
+    format_keys = set(data.keys()).intersection(SUPPORTED_FORMATS)
+    if len(format_keys) != 1:
+        message = f"Received dictionary with input fields: {list(data.keys())}"
+        raise MlflowException(
+            message=f"{REQUIRED_INPUT_FORMAT}. {message}. {SCORING_PROTOCOL_CHANGE_INFO}",
+            error_code=BAD_REQUEST,
+        )
+    input_format = format_keys.pop()
+    if input_format in (INSTANCES, INPUTS):
+        return parse_tf_serving_input(data, schema=schema)
+
+    if input_format in (DF_SPLIT, DF_RECORDS):
+        pandas_orient = input_format[10:]  # skip the dataframe_ prefix
+        return dataframe_from_parsed_json(
+            data[input_format], pandas_orient=pandas_orient, schema=schema
+        )
+
+
+def parse_csv_input(csv_input, schema: Schema = None):
+    """
+    Args:
+        csv_input: A CSV-formatted string representation of a Pandas DataFrame, or a stream
+                   containing such a string representation.
+        schema: Optional schema specification to be used during parsing.
+    """
+    import pandas as pd
+
+    try:
+        if schema is None:
+            return pd.read_csv(csv_input)
+        else:
+            dtypes = dict(zip(schema.input_names(), schema.pandas_types()))
+            return pd.read_csv(csv_input, dtype=dtypes)
+    except Exception as e:
+        _handle_serving_error(
+            error_message=(
+                "Failed to parse input as a Pandas DataFrame. Ensure that the input is"
+                " a valid CSV-formatted Pandas DataFrame produced using the"
+                f" `pandas.DataFrame.to_csv()` method. Error: '{e}'"
+            ),
+            error_code=BAD_REQUEST,
+        )
+
+
+def unwrapped_predictions_to_json(raw_predictions, output):
+    predictions = _get_jsonable_obj(raw_predictions, pandas_orient="records")
+    return json.dump(predictions, output, cls=NumpyEncoder)
+
+
+def predictions_to_json(raw_predictions, output, metadata=None):
+    if metadata and "predictions" in metadata:
+        raise MlflowException(
+            "metadata cannot contain 'predictions' key", error_code=INVALID_PARAMETER_VALUE
+        )
+    predictions = _get_jsonable_obj(raw_predictions, pandas_orient="records")
+    return json.dump({"predictions": predictions, **(metadata or {})}, output, cls=NumpyEncoder)
+
+
+def _handle_serving_error(error_message, error_code, include_traceback=True):
+    """
+    Logs information about an exception thrown by model inference code that is currently being
+    handled and reraises it with the specified error message. The exception stack trace
+    is also included in the reraised error message.
+
+    Args:
+        error_message: A message for the reraised exception.
+        error_code: An appropriate error code for the reraised exception. This should be one of
+            the codes listed in the `mlflow.protos.databricks_pb2` proto.
+        include_traceback: Whether to include the current traceback in the returned error.
+    """
+    if include_traceback:
+        traceback_buf = StringIO()
+        traceback.print_exc(file=traceback_buf)
+        traceback_str = traceback_buf.getvalue()
+        e = MlflowException(message=error_message, error_code=error_code, stack_trace=traceback_str)
+    else:
+        e = MlflowException(message=error_message, error_code=error_code)
+    reraise(MlflowException, e)
+
+
+class InvocationsResponse(NamedTuple):
+    response: str
+    status: int
+    mimetype: str
+
+
+def invocations(data, content_type, model, input_schema):
+    type_parts = list(map(str.strip, content_type.split(";")))
+    mime_type = type_parts[0]
+    parameter_value_pairs = type_parts[1:]
+    parameter_values = {
+        key: value for pair in parameter_value_pairs for key, _, value in [pair.partition("=")]
+    }
+
+    charset = parameter_values.get("charset", "utf-8").lower()
+    if charset != "utf-8":
+        return InvocationsResponse(
+            response="The scoring server only supports UTF-8",
+            status=415,
+            mimetype="text/plain",
+        )
+
+    unexpected_content_parameters = set(parameter_values.keys()).difference({"charset"})
+    if unexpected_content_parameters:
+        return InvocationsResponse(
+            response=(
+                f"Unrecognized content type parameters: "
+                f"{', '.join(unexpected_content_parameters)}. "
+                f"{SCORING_PROTOCOL_CHANGE_INFO}"
+            ),
+            status=415,
+            mimetype="text/plain",
+        )
+
+    # The traditional JSON request/response format, wraps the data with one of the supported keys
+    # like "dataframe_split" and "predictions". For LLM use cases, we also support unwrapped JSON
+    # payload, to provide unified prediction interface.
+    should_parse_as_unified_llm_input = False
+
+    if mime_type == CONTENT_TYPE_CSV:
+        # Convert from CSV to pandas
+        if isinstance(data, bytes):
+            data = data.decode("utf-8")
+        csv_input = StringIO(data)
+        data = parse_csv_input(csv_input=csv_input, schema=input_schema)
+        params = None
+    elif mime_type == CONTENT_TYPE_JSON:
+        parsed_json_input = _parse_json_data(data, model.metadata, input_schema)
+        data = parsed_json_input.data
+        params = parsed_json_input.params
+        should_parse_as_unified_llm_input = parsed_json_input.is_unified_llm_input
+    else:
+        return InvocationsResponse(
+            response=(
+                "This predictor only supports the following content types:"
+                f" Types: {CONTENT_TYPES}."
+                f" Got '{content_type}'."
+            ),
+            status=415,
+            mimetype="text/plain",
+        )
+
+    # Do the prediction
+    # NB: utils._validate_serving_input mimic the scoring process here to validate input_example
+    # work for serving, so any changes here should be reflected there as well
+    try:
+        if "params" in inspect.signature(model.predict).parameters:
+            raw_predictions = model.predict(data, params=params)
+        else:
+            _log_warning_if_params_not_in_predict_signature(_logger, params)
+            raw_predictions = model.predict(data)
+    except MlflowException as e:
+        if "Failed to enforce schema" in e.message:
+            _logger.warning(
+                "If using `instances` as input key, we internally convert "
+                "the data type from `records` (List[Dict]) type to "
+                "`list` (Dict[str, List]) type if the data is a pandas "
+                "dataframe representation. This might cause schema changes. "
+                "Please use `inputs` to avoid this conversion.\n"
+            )
+        e.message = f"Failed to predict data '{data}'. \nError: {e.message}"
+        raise e
+    except Exception:
+        raise MlflowException(
+            message=(
+                "Encountered an unexpected error while evaluating the model. Verify"
+                " that the serialized input Dataframe is compatible with the model for"
+                " inference."
+            ),
+            error_code=BAD_REQUEST,
+            stack_trace=traceback.format_exc(),
+        )
+    result = StringIO()
+
+    # if the data was formatted using the unified LLM format,
+    # then return the data without the "predictions" key
+    if should_parse_as_unified_llm_input:
+        unwrapped_predictions_to_json(raw_predictions, result)
+    else:
+        predictions_to_json(raw_predictions, result)
+
+    return InvocationsResponse(response=result.getvalue(), status=200, mimetype="application/json")
+
+
+class ParsedJsonInput(NamedTuple):
+    data: Any
+    params: Optional[dict[str, Any]]
+    is_unified_llm_input: bool
+
+
+def _parse_json_data(data, metadata, input_schema):
+    json_input = _decode_json_input(data)
+    _is_unified_llm_input = is_unified_llm_input(json_input)
+    # no data parsing for unified LLM input format
+    if _is_unified_llm_input:
+        # Unified LLM input format
+        if hasattr(metadata, "get_params_schema"):
+            params_schema = metadata.get_params_schema()
+        else:
+            params_schema = None
+        data, params = _split_data_and_params_for_llm_input(json_input, params_schema)
+    else:
+        # Traditional json input format
+        data, params = _split_data_and_params(data)
+        # data only needs to be parsed if the model signature is not from type hint
+        # default to True for backwards compatibility
+        should_parse_data = (
+            not metadata._is_signature_from_type_hint()
+            if hasattr(metadata, "_is_signature_from_type_hint")
+            else True
+        )
+        if should_parse_data:
+            data = infer_and_parse_data(data, input_schema)
+        else:
+            if INPUTS not in data:
+                raise MlflowException.invalid_parameter_value(
+                    "Request payload must be a dictionary with 'inputs' key when "
+                    f"the model contains a valid type hint. Found keys in payload: {data.keys()}."
+                )
+            data = data[INPUTS]
+    return ParsedJsonInput(data, params, _is_unified_llm_input)
+
+
+def _async_catch_mlflow_exception(func):
+    from fastapi.responses import Response
+
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except MlflowException as e:
+            return Response(
+                content=e.serialize_as_json(),
+                status_code=e.get_http_status_code(),
+                media_type="application/json",
+            )
+
+    return wrapper
+
+
+def init(model: PyFuncModel):
+    """
+    Initialize the server. Loads pyfunc model from the path.
+    """
+    from fastapi import FastAPI, Request
+    from fastapi.responses import Response
+
+    app = FastAPI()
+    input_schema = model.metadata.get_input_schema()
+    # set the environment variable to indicate that we are in a serving environment
+    os.environ[_MLFLOW_IS_IN_SERVING_ENVIRONMENT.name] = "true"
+    timeout = MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.get()
+
+    @app.middleware("http")
+    async def timeout_middleware(request: Request, call_next):
+        try:
+            return await asyncio.wait_for(call_next(request), timeout=timeout)
+        except (asyncio.TimeoutError, TimeoutError):
+            return Response(
+                content="Request processing time exceeded limit",
+                status_code=504,
+                media_type="application/json",
+            )
+
+    @app.route("/ping", methods=["GET"])
+    @app.route("/health", methods=["GET"])
+    async def ping(request: Request):
+        """
+        Determine if the container is working and healthy.
+        We declare it healthy if we can load the model successfully.
+        """
+        health = model is not None
+        status = 200 if health else 404
+        return Response(content="\n", status_code=status, media_type="application/json")
+
+    @app.route("/version", methods=["GET"])
+    async def version(request: Request):
+        """
+        Returns the current mlflow version.
+        """
+        return Response(content=VERSION, status_code=200, media_type="application/json")
+
+    @app.route("/invocations", methods=["POST"])
+    @_async_catch_mlflow_exception
+    async def transformation(request: Request):
+        """
+        Do an inference on a single batch of data. In this sample server,
+        we take data as CSV or json, convert it to a Pandas DataFrame or Numpy,
+        generate predictions and convert them back to json.
+        """
+
+        data = await request.body()
+        content_type = request.headers.get("content-type")
+        # TODO: convert "invocations" to an async method to make internal logic fully non-blocking.
+        result = await asyncio.to_thread(invocations, data, content_type, model, input_schema)
+
+        return Response(
+            content=result.response, status_code=result.status, media_type=result.mimetype
+        )
+
+    return app
+
+
+def _predict(model_uri, input_path, output_path, content_type):
+    from mlflow.pyfunc.utils.environment import _simulate_serving_environment
+
+    with _simulate_serving_environment():
+        pyfunc_model = load_model(model_uri)
+
+        should_parse_as_unified_llm_input = False
+        if content_type == "json":
+            if input_path is None:
+                input_str = sys.stdin.read()
+            else:
+                with open(input_path) as f:
+                    input_str = f.read()
+            parsed_json_input = _parse_json_data(
+                data=input_str,
+                metadata=pyfunc_model.metadata,
+                input_schema=pyfunc_model.metadata.get_input_schema(),
+            )
+            df = parsed_json_input.data
+            params = parsed_json_input.params
+            should_parse_as_unified_llm_input = parsed_json_input.is_unified_llm_input
+        elif content_type == "csv":
+            df = (
+                parse_csv_input(input_path)
+                if input_path is not None
+                else parse_csv_input(sys.stdin)
+            )
+            params = None
+        else:
+            raise Exception(f"Unknown content type '{content_type}'")
+
+        if "params" in inspect.signature(pyfunc_model.predict).parameters:
+            raw_predictions = pyfunc_model.predict(df, params=params)
+        else:
+            _log_warning_if_params_not_in_predict_signature(_logger, params)
+            raw_predictions = pyfunc_model.predict(df)
+
+        parse_output_func = (
+            unwrapped_predictions_to_json
+            if should_parse_as_unified_llm_input
+            else predictions_to_json
+        )
+
+        if output_path is None:
+            parse_output_func(raw_predictions, sys.stdout)
+        else:
+            with open(output_path, "w") as fout:
+                parse_output_func(raw_predictions, fout)
+
+
+def _serve(model_uri, port, host):
+    pyfunc_model = load_model(model_uri)
+    init(pyfunc_model).run(port=port, host=host)
+
+
+def get_cmd(
+    model_uri: str,
+    port: Optional[int] = None,
+    host: Optional[int] = None,
+    timeout: Optional[int] = None,
+    nworkers: Optional[int] = None,
+) -> tuple[str, dict[str, str]]:
+    local_uri = path_to_local_file_uri(model_uri)
+    timeout = timeout or MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.get()
+
+    args = []
+    if host:
+        args.append(f"--host {shlex.quote(host)}")
+
+    if port:
+        args.append(f"--port {port}")
+
+    if nworkers:
+        args.append(f"--workers {nworkers}")
+
+    command = f"uvicorn {' '.join(args)} mlflow.pyfunc.scoring_server.app:app"
+
+    command_env = os.environ.copy()
+    command_env[_SERVER_MODEL_PATH] = local_uri
+    command_env[MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.name] = str(timeout)
+
+    return command, command_env

mlflow/pyfunc/scoring_server/app.py

@@ -0,0 +1,7 @@
+import os
+
+from mlflow.pyfunc import scoring_server
+
+app = scoring_server.init(
+    scoring_server.load_model_with_mlflow_config(os.environ[scoring_server._SERVER_MODEL_PATH])
+)