Repository: mlflow/mlflow

Issue:
[FR] Add possibility to unwrap python_model from PyFuncModel class
### Willingness to contribute

Yes. I can contribute this feature independently.

### Proposal Summary

FR could be solved by adding `unwrap(self)` method to `PyFuncModel` class. `unwrap(self)` would return `self._model_impl`, thus providing direct access to the wrapped python model. It could also return `self._model_impl` if user wants that.

### Motivation

> #### What is the use case for this feature?
Users would be able to load their custom python objects by calling `unwrap()` method when using `mlflow.pyfunc.load_model`.

> #### Why is this use case valuable to support for MLflow users in general?
This is valuable because methods such as `save_model, log_model` and `load_model` could still be used, but much more flexibility with custom models would be achieved.
> #### Why is this use case valuable to support for your project(s) or organization?
This is valuable to my projects because we often use custom models and need access to it's custom methods.
> #### Why is it currently difficult to achieve this use case?
The current implementation of `PyFuncModel` class implements only one useful method for interaction with loaded model (`predict(data) -> predictions`). When using a custom python object as a model (while inheriting `mlflow.pyfunc.PythonModel`) it's very limiting to have only one method exposed.


### Details

_No response_

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/pyfunc/__init__.py b/mlflow/pyfunc/__init__.py
--- a/mlflow/pyfunc/__init__.py
+++ b/mlflow/pyfunc/__init__.py
@@ -207,45 +207,55 @@
   artifacts.
 """
 
+import collections
 import importlib
-import tempfile
+import logging
+import os
 import signal
+import subprocess
 import sys
+import tempfile
+import threading
+from copy import deepcopy
+from typing import Any, Union, Iterator, Tuple
 
 import numpy as np
-import os
 import pandas
 import yaml
-from copy import deepcopy
-import logging
-import threading
-import collections
-import subprocess
 
-from typing import Any, Union, Iterator, Tuple
 import mlflow
 import mlflow.pyfunc.model
+from mlflow.exceptions import MlflowException
 from mlflow.models import Model, ModelSignature, ModelInputExample
 from mlflow.models.model import MLMODEL_FILE_NAME
-from mlflow.models.utils import PyFuncInput, PyFuncOutput, _enforce_schema, _save_example
+from mlflow.models.utils import (
+    PyFuncInput,
+    PyFuncOutput,
+    _enforce_schema,
+    _save_example,
+)
+from mlflow.protos.databricks_pb2 import (
+    INVALID_PARAMETER_VALUE,
+    RESOURCE_DOES_NOT_EXIST,
+)
 from mlflow.pyfunc.model import (  # pylint: disable=unused-import
     PythonModel,
     PythonModelContext,
     get_default_conda_env,
 )
 from mlflow.pyfunc.model import get_default_pip_requirements
+from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
-from mlflow.utils import PYTHON_VERSION, get_major_minor_py_version, _is_in_ipython_notebook
-from mlflow.utils.annotations import deprecated
-from mlflow.utils.file_utils import _copy_file_or_tree, write_to
-from mlflow.utils.model_utils import (
-    _get_flavor_configuration,
-    _validate_and_copy_code_paths,
-    _add_code_from_conf_to_system_path,
-    _get_flavor_configuration_from_uri,
-    _validate_and_prepare_target_save_path,
+from mlflow.utils import (
+    PYTHON_VERSION,
+    get_major_minor_py_version,
+    _is_in_ipython_notebook,
 )
-from mlflow.utils.uri import append_to_uri_path
+from mlflow.utils import env_manager as _EnvManager
+from mlflow.utils import find_free_port
+from mlflow.utils.annotations import deprecated
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+from mlflow.utils.docstring_utils import format_docstring, LOG_MODEL_PARAM_DOCS
 from mlflow.utils.environment import (
     _validate_env_arguments,
     _process_pip_requirements,
@@ -256,24 +266,22 @@
     _PYTHON_ENV_FILE_NAME,
     _PythonEnv,
 )
-from mlflow.utils import env_manager as _EnvManager
-from mlflow.utils.docstring_utils import format_docstring, LOG_MODEL_PARAM_DOCS
-from mlflow.utils.databricks_utils import is_in_databricks_runtime
+from mlflow.utils.file_utils import _copy_file_or_tree, write_to
 from mlflow.utils.file_utils import get_or_create_tmp_dir, get_or_create_nfs_tmp_dir
-from mlflow.utils.process import cache_return_value_per_process
-from mlflow.exceptions import MlflowException
-from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
-from mlflow.protos.databricks_pb2 import (
-    INVALID_PARAMETER_VALUE,
-    RESOURCE_DOES_NOT_EXIST,
+from mlflow.utils.model_utils import (
+    _get_flavor_configuration,
+    _validate_and_copy_code_paths,
+    _add_code_from_conf_to_system_path,
+    _get_flavor_configuration_from_uri,
+    _validate_and_prepare_target_save_path,
 )
-
+from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir
+from mlflow.utils.process import cache_return_value_per_process
 from mlflow.utils.requirements_utils import (
     _check_requirement_satisfied,
     _parse_requirements,
 )
-from mlflow.utils import find_free_port
-from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir
+from mlflow.utils.uri import append_to_uri_path
 
 FLAVOR_NAME = "python_function"
 MAIN = "loader_module"
@@ -372,6 +380,18 @@ def predict(self, data: PyFuncInput) -> PyFuncOutput:
             data = _enforce_schema(data, input_schema)
         return self._predict_fn(data)
 
+    def unwrap_python_model(self):
+        """
+        Unwrap the underlying Python model object.
+        """
+        try:
+            python_model = self._model_impl.python_model
+            if python_model is None:
+                raise AttributeError("Expected python_model attribute not to be None.")
+        except AttributeError as e:
+            raise MlflowException("Unable to retrieve base model object from pyfunc.") from e
+        return python_model
+
     def __eq__(self, other):
         if not isinstance(other, PyFuncModel):
             return False
@@ -750,7 +770,13 @@ def spark_udf(spark, model_uri, result_type="double", env_manager="local"):
         DataType as SparkDataType,
         StructType as SparkStructType,
     )
-    from pyspark.sql.types import DoubleType, IntegerType, FloatType, LongType, StringType
+    from pyspark.sql.types import (
+        DoubleType,
+        IntegerType,
+        FloatType,
+        LongType,
+        StringType,
+    )
     from mlflow.models.cli import _get_flavor_backend
 
     _EnvManager.validate(env_manager)
@@ -781,7 +807,8 @@ def spark_udf(spark, model_uri, result_type="double", env_manager="local"):
         )
 
     local_model_path = _download_artifact_from_uri(
-        artifact_uri=model_uri, output_path=_create_model_downloading_tmp_dir(should_use_nfs)
+        artifact_uri=model_uri,
+        output_path=_create_model_downloading_tmp_dir(should_use_nfs),
     )
 
     if env_manager == _EnvManager.LOCAL:
@@ -972,7 +999,8 @@ def server_redirect_log_thread_func(child_stdout):
                     sys.stdout.write("[model server] " + decoded)
 
             server_redirect_log_thread = threading.Thread(
-                target=server_redirect_log_thread_func, args=(scoring_server_proc.stdout,)
+                target=server_redirect_log_thread_func,
+                args=(scoring_server_proc.stdout,),
             )
             server_redirect_log_thread.setDaemon(True)
             server_redirect_log_thread.start()


Test patch:
diff --git a/tests/pyfunc/test_pyfunc_class_methods.py b/tests/pyfunc/test_pyfunc_class_methods.py
new file mode 100644
--- /dev/null
+++ b/tests/pyfunc/test_pyfunc_class_methods.py
@@ -0,0 +1,25 @@
+import mlflow
+from mlflow.pyfunc import PythonModel, log_model, load_model
+
+
+def test_unwrap_python_model_from_pyfunc_class():
+    class MyModel(PythonModel):
+        def __init__(self, param_1: str, param_2: int):
+            self.param_1 = param_1
+            self.param_2 = param_2
+
+        def predict(self, context, model_input):
+            return model_input + self.param_2
+
+        def upper_param_1(self):
+            return self.param_1.upper()
+
+    with mlflow.start_run():
+        model = MyModel("this is test message", 2)
+        model_uri = log_model(python_model=model, artifact_path="mlruns").model_uri
+        loaded_model = load_model(model_uri).unwrap_python_model()
+        assert isinstance(loaded_model, MyModel)
+        assert loaded_model.param_1 == "this is test message"
+        assert loaded_model.param_2 == 2
+        assert loaded_model.predict(None, 1) == 3
+        assert loaded_model.upper_param_1() == "THIS IS TEST MESSAGE"