Merge pull request #133 from whylabs/dev/andy/mlflow

[bugfix] Fix mlfow tracking with predict API

Author: lalmei

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.1.13-dev3
+current_version = 0.1.13-dev8
 commit = True
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\-(?P<release>[a-z]+)(?P<build>\d+))?

docs/conf.py

@@ -119,7 +119,7 @@ def setup(app):
 # built documents.
 #
 # The short X.Y version.
-version = " 0.1.13-dev3"
+version = "0.1.13-dev8"
 # The full version, including alpha/beta/rc tags.
 release = ""  # Is set by calling `setup.py docs`
 

setup.cfg

@@ -4,7 +4,7 @@
 
 [metadata]
 name = whylogs
-version = 0.1.13-dev3
+version = 0.1.13-dev8
 description = Profile and monitor your ML data pipeline end-to-end
 author = WhyLabs.ai
 author-email = [email protected]

src/whylogs/_version.py

@@ -1,3 +1,3 @@
 """WhyLabs version number."""
 
-__version__ = "0.1.13-dev3"
+__version__ = "0.1.13-dev8"

src/whylogs/mlflow/model_wrapper.py

@@ -1,25 +1,48 @@
 import atexit
+from typing import Union
+
+import numpy as np
+import pandas as pd
 
 import whylogs
 import datetime
 
+from logging import getLogger
 from whylogs.app.config import WHYLOGS_YML
 
+logger = getLogger(__name__)
+
+PyFuncOutput = Union[pd.DataFrame, pd.Series, np.ndarray, list]
+
 
 class ModelWrapper(object):
     def __init__(self, model):
         self.model = model
         self.session = whylogs.get_or_create_session("/opt/ml/model/" + WHYLOGS_YML)
-        self.logger = self.create_logger()
+        self.ylog = self.create_logger()
         self.last_upload_time = datetime.datetime.utcnow()
-        atexit.register(self.logger.close)
+        atexit.register(self.ylog.close)
 
     def create_logger(self):
         # TODO: support different rotation mode and support custom name here
         return self.session.logger('live', with_rotation_time='m')
 
-    def predict(self, df):
-        self.logger.log_dataframe(df)
-        output = self.model.predict(df)
-        self.logger.log_dataframe(df)
+    def predict(self, data: pd.DataFrame) -> PyFuncOutput:
+        """
+        Wrapper around https://www.mlflow.org/docs/latest/_modules/mlflow/pyfunc.html#PyFuncModel.predict
+        This allows us to capture input and predictions into whylogs
+        """
+        self.ylog.log_dataframe(data)
+        output = self.model.predict(data)
+
+        if isinstance(output, np.ndarray) or isinstance(output, pd.Series):
+            data = pd.DataFrame(data=output, columns=['prediction'])
+            self.ylog.log_dataframe(data)
+        elif isinstance(output, pd.DataFrame):
+            self.ylog.log_dataframe(output)
+        elif isinstance(output, list):
+            for e in output:
+                self.ylog.log(feature_name='prediction', value=e)
+        else:
+            logger.warning('Unsupported output  type: %s', type(output))
         return output

src/whylogs/mlflow/patcher.py

@@ -7,6 +7,7 @@
 import pandas as pd
 
 from whylogs.app.logger import Logger
+from whylogs import __version__ as whylogs_version
 
 logger = logging.getLogger(__name__)
 
@@ -153,8 +154,7 @@ def _new_mlflow_conda_env(
 ):
     global _original_mlflow_conda_env
     pip_deps = additional_pip_deps or []
-    if "whylogs" not in pip_deps:
-        pip_deps.append("whylogs")
+    pip_deps.append(f"whylogs=={whylogs_version}")
     return _original_mlflow_conda_env(
         path, additional_conda_deps, pip_deps, additional_conda_channels, install_mlflow
     )