Serving input example support (mlflow#12710)

Signed-off-by: Serena Ruan <[email protected]>

mlflow/catboost/__init__.py

@@ -124,18 +124,18 @@ def save_model(
     _validate_and_prepare_target_save_path(path)
     code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
 
-    if signature is None and input_example is not None:
+    if mlflow_model is None:
+        mlflow_model = Model()
+    saved_example = _save_example(mlflow_model, input_example, path)
+
+    if signature is None and saved_example is not None:
         wrapped_model = _CatboostModelWrapper(cb_model)
-        signature = _infer_signature_from_input_example(input_example, wrapped_model)
+        signature = _infer_signature_from_input_example(saved_example, wrapped_model)
     elif signature is False:
         signature = None
 
-    if mlflow_model is None:
-        mlflow_model = Model()
     if signature is not None:
         mlflow_model.signature = signature
-    if input_example is not None:
-        _save_example(mlflow_model, input_example, path)
     if metadata is not None:
         mlflow_model.metadata = metadata
 

mlflow/diviner/__init__.py

@@ -30,6 +30,7 @@
 from mlflow.exceptions import MlflowException
 from mlflow.models import Model, ModelInputExample, ModelSignature
 from mlflow.models.model import MLMODEL_FILE_NAME
+from mlflow.models.signature import _infer_signature_from_input_example
 from mlflow.models.utils import _save_example
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
@@ -157,10 +158,13 @@ def save_model(
 
     if mlflow_model is None:
         mlflow_model = Model()
+    saved_example = _save_example(mlflow_model, input_example, str(path))
+    if signature is None and saved_example is not None:
+        wrapped_model = _DivinerModelWrapper(diviner_model)
+        signature = _infer_signature_from_input_example(saved_example, wrapped_model)
+
     if signature is not None:
         mlflow_model.signature = signature
-    if input_example is not None:
-        _save_example(mlflow_model, input_example, str(path))
     if metadata is not None:
         mlflow_model.metadata = metadata
 

mlflow/fastai/__init__.py

@@ -26,6 +26,7 @@
 from mlflow import pyfunc
 from mlflow.models import Model, ModelInputExample, ModelSignature
 from mlflow.models.model import MLMODEL_FILE_NAME
+from mlflow.models.signature import _infer_signature_from_input_example
 from mlflow.models.utils import _save_example
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
@@ -163,10 +164,13 @@ def save_model(
 
     if mlflow_model is None:
         mlflow_model = Model()
+    saved_example = _save_example(mlflow_model, input_example, path)
+    if signature is None and saved_example is not None:
+        wrapped_model = _FastaiModelWrapper(fastai_learner)
+        signature = _infer_signature_from_input_example(saved_example, wrapped_model)
+
     if signature is not None:
         mlflow_model.signature = signature
-    if input_example is not None:
-        _save_example(mlflow_model, input_example, path)
     if metadata is not None:
         mlflow_model.metadata = metadata
 

mlflow/gluon/__init__.py

@@ -224,18 +224,18 @@ def save_model(
     os.makedirs(data_path)
     code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
 
-    if signature is None and input_example is not None:
+    if mlflow_model is None:
+        mlflow_model = Model()
+    saved_example = _save_example(mlflow_model, input_example, path)
+
+    if signature is None and saved_example is not None:
         wrapped_model = _GluonModelWrapper(gluon_model)
-        signature = _infer_signature_from_input_example(input_example, wrapped_model)
+        signature = _infer_signature_from_input_example(saved_example, wrapped_model)
     elif signature is False:
         signature = None
 
-    if mlflow_model is None:
-        mlflow_model = Model()
     if signature is not None:
         mlflow_model.signature = signature
-    if input_example is not None:
-        _save_example(mlflow_model, input_example, path)
     if metadata is not None:
         mlflow_model.metadata = metadata
 

mlflow/h2o/__init__.py

@@ -110,18 +110,18 @@ def save_model(
     os.makedirs(model_data_path)
     code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
 
-    if signature is None and input_example is not None:
+    if mlflow_model is None:
+        mlflow_model = Model()
+    saved_example = _save_example(mlflow_model, input_example, path)
+
+    if signature is None and saved_example is not None:
         wrapped_model = _H2OModelWrapper(h2o_model)
-        signature = _infer_signature_from_input_example(input_example, wrapped_model)
+        signature = _infer_signature_from_input_example(saved_example, wrapped_model)
     elif signature is False:
         signature = None
 
-    if mlflow_model is None:
-        mlflow_model = Model()
     if signature is not None:
         mlflow_model.signature = signature
-    if input_example is not None:
-        _save_example(mlflow_model, input_example, path)
     if metadata is not None:
         mlflow_model.metadata = metadata
 

mlflow/langchain/__init__.py

@@ -284,12 +284,14 @@ def load_retriever(persist_directory):
 
     code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
 
+    if mlflow_model is None:
+        mlflow_model = Model()
+    saved_example = _save_example(mlflow_model, input_example, path, example_no_conversion)
+
     if signature is None:
-        if input_example is not None:
+        if saved_example is not None:
             wrapped_model = _LangChainModelWrapper(lc_model)
-            signature = _infer_signature_from_input_example(
-                input_example, wrapped_model, no_conversion=example_no_conversion
-            )
+            signature = _infer_signature_from_input_example(saved_example, wrapped_model)
         else:
             if hasattr(lc_model, "input_keys"):
                 input_columns = [
@@ -320,13 +322,8 @@ def load_retriever(persist_directory):
                 else None
             )
 
-    if mlflow_model is None:
-        mlflow_model = Model()
     if signature is not None:
         mlflow_model.signature = signature
-
-    if input_example is not None:
-        _save_example(mlflow_model, input_example, path, example_no_conversion)
     if metadata is not None:
         mlflow_model.metadata = metadata
 

mlflow/lightgbm/__init__.py

@@ -178,18 +178,18 @@ def save_model(
     model_data_path = os.path.join(path, model_data_subpath)
     code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
 
-    if signature is None and input_example is not None:
+    if mlflow_model is None:
+        mlflow_model = Model()
+    saved_example = _save_example(mlflow_model, input_example, path)
+
+    if signature is None and saved_example is not None:
         wrapped_model = _LGBModelWrapper(lgb_model)
-        signature = _infer_signature_from_input_example(input_example, wrapped_model)
+        signature = _infer_signature_from_input_example(saved_example, wrapped_model)
     elif signature is False:
         signature = None
 
-    if mlflow_model is None:
-        mlflow_model = Model()
     if signature is not None:
         mlflow_model.signature = signature
-    if input_example is not None:
-        _save_example(mlflow_model, input_example, path)
     if metadata is not None:
         mlflow_model.metadata = metadata
 

mlflow/models/__init__.py

@@ -73,16 +73,24 @@
 try:
     from mlflow.models.python_api import predict
     from mlflow.models.signature import ModelSignature, infer_signature, set_signature
-    from mlflow.models.utils import ModelInputExample, add_libraries_to_model, validate_schema
+    from mlflow.models.utils import (
+        ModelInputExample,
+        add_libraries_to_model,
+        convert_input_example_to_serving_input,
+        validate_schema,
+        validate_serving_input,
+    )
 
     __all__ += [
         "ModelSignature",
         "ModelInputExample",
         "infer_signature",
         "validate_schema",
         "add_libraries_to_model",
+        "convert_input_example_to_serving_input",
         "set_signature",
         "predict",
+        "validate_serving_input",
     ]
 except ImportError:
     pass

mlflow/models/model.py

@@ -353,6 +353,21 @@ def get_params_schema(self):
         """
         return getattr(self.signature, "params", None)
 
+    def get_serving_input(self, path: str):
+        """
+        Load serving input example from a model directory. Returns None if there is no serving input
+        example.
+
+        Args:
+            path: Path to the model directory.
+
+        Returns:
+            Serving input example or None if the model has no serving input example.
+        """
+        from mlflow.models.utils import _load_serving_input_example
+
+        return _load_serving_input_example(self, path)
+
     def load_input_example(self, path: str):
         """
         Load the input example saved along a model. Returns None if there is no example metadata
@@ -747,9 +762,36 @@ def log(
                     await_registration_for=await_registration_for,
                     local_model_path=local_path,
                 )
-        model_info = mlflow_model.get_model_info()
-        if registered_model is not None:
-            model_info.registered_model_version = registered_model.version
+            model_info = mlflow_model.get_model_info()
+            if registered_model is not None:
+                model_info.registered_model_version = registered_model.version
+
+            # validate input example works for serving when logging the model
+            serving_input = mlflow_model.get_serving_input(local_path)
+            if mlflow_model.signature is None and serving_input is None:
+                _logger.warning(
+                    "Input example should be provided to infer model signature if the model "
+                    "signature is not provided when logging the model."
+                )
+            if serving_input:
+                from mlflow.models import validate_serving_input
+
+                try:
+                    validate_serving_input(model_info.model_uri, serving_input)
+                except Exception as e:
+                    _logger.warning(
+                        f"Failed to validate serving input example {serving_input}. "
+                        "Alternatively, you can avoid passing input example and pass model "
+                        "signature instead when logging the model. To ensure the input example "
+                        "is valid prior to serving, please try calling "
+                        "`mlflow.models.validate_serving_input` on the model uri and serving "
+                        "input example. A serving input example can be generated from model "
+                        "input example using "
+                        "`mlflow.models.convert_input_example_to_serving_input` function.\n"
+                        f"Got error: {e}",
+                        exc_info=_logger.isEnabledFor(logging.DEBUG),
+                    )
+
         return model_info
 
 

mlflow/models/signature.py

@@ -18,7 +18,7 @@
 from mlflow.exceptions import MlflowException
 from mlflow.models import Model
 from mlflow.models.model import MLMODEL_FILE_NAME
-from mlflow.models.utils import ModelInputExample, _contains_params, _Example
+from mlflow.models.utils import _contains_params, _Example
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, RESOURCE_DOES_NOT_EXIST
 from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
 from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
@@ -351,32 +351,38 @@ def _infer_signature_from_type_hints(func, input_arg_index, input_example=None):
 
 
 def _infer_signature_from_input_example(
-    input_example: ModelInputExample, wrapped_model, no_conversion=False
+    input_example: Optional[_Example], wrapped_model
 ) -> Optional[ModelSignature]:
     """
     Infer the signature from an example input and a PyFunc wrapped model. Catches all exceptions.
 
     Args:
-        input_example: An instance representing a typical input to the model.
+        input_example: Saved _Example object that contains input example instance.
         wrapped_model: A PyFunc wrapped model which has a `predict` method.
 
     Returns:
         A `ModelSignature` object containing the inferred schema of both the model's inputs
         based on the `input_example` and the model's outputs based on the prediction from the
         `wrapped_model`.
     """
+    from mlflow.pyfunc import _validate_prediction_input
+
+    if input_example is None:
+        return None
+
     try:
-        if _contains_params(input_example):
-            input_example, params = input_example
-        else:
-            params = None
-        if not no_conversion:
-            example = _Example(input_example)
-            # Copy the input example so that it is not mutated by predict()
-            input_example = deepcopy(example.inference_data)
-        input_schema = _infer_schema(input_example)
+        # Copy the input example so that it is not mutated by predict()
+        input_data = deepcopy(input_example.inference_data)
+        params = input_example.inference_params
+
+        input_schema = _infer_schema(input_data)
         params_schema = _infer_param_schema(params) if params else None
-        prediction = wrapped_model.predict(input_example, params=params)
+        # do the same validation as pyfunc predict to make sure the signature is correctly
+        # applied to the model
+        input_data, params = _validate_prediction_input(
+            input_data, params, input_schema, params_schema
+        )
+        prediction = wrapped_model.predict(input_data, params=params)
         # For column-based inputs, 1D numpy arrays likely signify row-based predictions. Thus, we
         # convert them to a Pandas series for inferring as a single ColSpec Schema.
         if (