Dbaas 3809 (#51)

* support for saving/loading sklearn and keras models

* wrong parameter order

* missing read mode for keras file

* get run ids by name

* extra quote

Author: Epstein

splicemachine/mlflow_support/constants.py

@@ -23,3 +23,24 @@ class SparkModelType(Enum):
     REGRESSION = 1
     CLUSTERING_WITH_PROB = 2
     CLUSTERING_WO_PROB = 3
+
+class FileExtensions():
+    """
+    Class containing names for
+    valid File Extensions
+    """
+    spark: str = "spark"
+    keras: str = "h5"
+    h2o: str = "h2o"
+    sklearn: str = "pkl"
+
+    @staticmethod
+    def get_valid() -> tuple:
+        """
+        Return a tuple of the valid file extensions
+        in Database
+        :return: (tuple) valid statuses
+        """
+        return (
+            FileExtensions.spark, FileExtensions.keras, FileExtensions.h2o, FileExtensions.sklearn
+        )

splicemachine/mlflow_support/mlflow_support.py

@@ -1,19 +1,21 @@
 import time
 from collections import defaultdict
 from contextlib import contextmanager
-from io import BytesIO
-from os import path, remove
-from shutil import rmtree
-from zipfile import ZipFile
+from os import path
 from sys import version as py_version
 
 import gorilla
 import mlflow
 import requests
 from requests.auth import HTTPBasicAuth
 from mleap.pyspark import spark_support
-import h2o
 import pyspark
+from pyspark.ml.base import Estimator as SparkModel
+import sklearn
+from sklearn.base import BaseEstimator as ScikitModel
+from tensorflow import __version__ as tf_version
+from tensorflow.keras import __version__ as keras_version
+from tensorflow.keras import Model as KerasModel
 
 from splicemachine.mlflow_support.constants import *
 from splicemachine.mlflow_support.utilities import *
@@ -24,6 +26,7 @@
 _TRACKING_URL = get_pod_uri("mlflow", "5001", _TESTING)
 
 _CLIENT = mlflow.tracking.MlflowClient(tracking_uri=_TRACKING_URL)
+mlflow.client = _CLIENT
 
 _GORILLA_SETTINGS = gorilla.Settings(allow_hit=True, store_hit=True)
 _PYTHON_VERSION = py_version.split('|')[0].strip()
@@ -50,6 +53,23 @@ def _get_current_run_data():
     return _CLIENT.get_run(mlflow.active_run().info.run_id).data
 
 
+@_mlflow_patch('get_run_ids_by_name')
+def _get_run_ids_by_name(run_name, experiment_id=None):
+    """
+    Gets a run id from the run name. If there are multiple runs with the same name, all run IDs are returned
+    :param run_name: The name of the run
+    :param experiment_id: The experiment to search in. If None, all experiments are searched
+    :return: List of run ids
+    """
+    exps = [experiment_id] if experiment_id else _CLIENT.list_experiments()
+    run_ids = []
+    for exp in exps:
+        for run in _CLIENT.search_runs(exp.experiment_id):
+            if run_name == run.data.tags['mlflow.runName']:
+                run_ids.append(run.data.tags['Run ID'])
+    return run_ids
+
+
 @_mlflow_patch('register_splice_context')
 def _register_splice_context(splice_context):
     """
@@ -70,7 +90,7 @@ def _check_for_splice_ctx():
 
     if not hasattr(mlflow, '_splice_context'):
         raise SpliceMachineException(
-            "You must run `mlflow.register_splice_context(py_splice_context) before "
+            "You must run `mlflow.register_splice_context(pysplice_context) before "
             "you can run this mlflow operation!"
         )
 
@@ -135,17 +155,26 @@ def _log_model(model, name='model'):
     run_id = mlflow.active_run().info.run_uuid
     if 'h2o' in model_class.lower():
         mlflow.set_tag('splice.h2o_version', h2o.__version__)
-        model_path = h2o.save_model(model=model, path='/tmp/model', force=True)
-        with open(model_path, 'rb') as artifact:
-            byte_stream = bytearray(bytes(artifact.read()))
-        insert_artifact(mlflow._splice_context, name, byte_stream, run_id, file_ext='h2omodel')
-        rmtree('/tmp/model')
+        H2OUtils.log_h2o_model(mlflow._splice_context, model, name, run_id)
 
-    elif 'spark' in model_class.lower():
+    elif isinstance(model, SparkModel):
         mlflow.set_tag('splice.spark_version', pyspark.__version__)
-        SparkUtils.log_spark_model(mlflow._splice_context, model, name, run_id=run_id)
+        SparkUtils.log_spark_model(mlflow._splice_context, model, name, run_id)
+
+    elif isinstance(model, ScikitModel):
+        mlflow.set_tag('splice.sklearn_version', sklearn.__version__)
+        SKUtils.log_sklearn_model(mlflow._splice_context, model, name, run_id)
+
+    elif isinstance(model, KerasModel): # We can't handle keras models with a different backend
+        mlflow.set_tag('splice.keras_version', keras_version)
+        mlflow.set_tag('splice.tf_version', tf_version)
+        KerasUtils.log_keras_model(mlflow._splice_context, model, name, run_id)
+
+
     else:
-        raise SpliceMachineException('Currently we only support logging Spark and H2O models.')
+        raise SpliceMachineException('Model type not supported for logging.'
+                                     'Currently we support logging Spark, H2O, SKLearn and Keras (TF backend) models.'
+                                     'You can save your model to disk, zip it and run mlflow.log_artifact to save.')
 
 @_mlflow_patch('start_run')
 def _start_run(run_id=None, tags=None, experiment_id=None, run_name=None, nested=False):
@@ -313,14 +342,18 @@ def _load_model(run_id=None, name='model'):
     run_id = run_id or mlflow.active_run().info.run_uuid
     model_blob, file_ext = SparkUtils.retrieve_artifact_stream(mlflow._splice_context, run_id, name)
 
-    if file_ext == 'sparkmodel':
+    if file_ext == FileExtensions.spark:
         model = SparkUtils.load_spark_model(mlflow._splice_context, model_blob)
+    elif file_ext == FileExtensions.h2o:
+        model = H2OUtils.load_h2o_model(model_blob)
+    elif file_ext == FileExtensions.sklearn:
+        model = SKUtils.load_sklearn_model(model_blob)
+    elif file_ext == FileExtensions.keras:
+        model = KerasUtils.load_keras_model(model_blob)
+    else:
+        raise SpliceMachineException(f'Model extension {file_ext} was not a supported model type. '
+                                     f'Supported model extensions are {FileExtensions.get_valid()}')
 
-    elif file_ext == 'h2omodel':
-        with open('/tmp/model', 'wb') as file:
-            file.write(model_blob)
-        model = h2o.load_model('/tmp/model')
-        remove('/tmp/model')
     return model
 
 
@@ -585,7 +618,8 @@ def apply_patches():
     """
     targets = [_register_splice_context, _lp, _lm, _timer, _log_artifact, _log_feature_transformations,
                _log_model_params, _log_pipeline_stages, _log_model, _load_model, _download_artifact,
-               _start_run, _current_run_id, _current_exp_id, _deploy_aws, _deploy_azure, _deploy_db, _login_director]
+               _start_run, _current_run_id, _current_exp_id, _deploy_aws, _deploy_azure, _deploy_db, _login_director,
+               _get_run_ids_by_name]
 
     for target in targets:
         gorilla.apply(gorilla.Patch(mlflow, target.__name__.lstrip('_'), target, settings=_GORILLA_SETTINGS))

splicemachine/mlflow_support/utilities.py

@@ -1,14 +1,19 @@
 from os import environ as env_vars, popen as rbash, system as bash, remove
 from sys import getsizeof
+from shutil import rmtree
+from pickle import dumps as save_pickle_string, loads as load_pickle_string
+from io import BytesIO
+from h5py import File as h5_file
 import re
 
-from pyspark.ml import Pipeline, PipelineModel
 from pyspark.ml.base import Model as SparkModel
+from tensorflow.keras.models import load_model as load_kr_model
 from py4j.java_gateway import java_import
 
 from splicemachine.spark.constants import SQL_TYPES
 from splicemachine.mlflow_support.constants import *
 from mleap.pyspark.spark_support import SimpleSparkSerializer
+
 import h2o
 
 from pyspark.ml.pipeline import PipelineModel
@@ -102,8 +107,25 @@ def get_h2omojo_model(splice_context, model):
         raw_mojo = jvm.MojoModel.load(model_path)
         java_mojo_c = jvm.EasyPredictModelWrapper.Config().setModel(raw_mojo)
         java_mojo = jvm.EasyPredictModelWrapper(java_mojo_c)
+        remove('/tmp/model.zip')
         return java_mojo, raw_mojo
 
+    @staticmethod
+    def log_h2o_model(splice_context, model, name, run_id):
+        model_path = h2o.save_model(model=model, path='/tmp/model', force=True)
+        with open(model_path, 'rb') as artifact:
+            byte_stream = bytearray(bytes(artifact.read()))
+        insert_artifact(splice_context, name, byte_stream, run_id, file_ext=FileExtensions.h2o)
+        rmtree('/tmp/model')
+
+    @staticmethod
+    def load_h2o_model(model_blob):
+        with open('/tmp/model', 'wb') as file:
+            file.write(model_blob)
+        model = h2o.load_model('/tmp/model')
+        remove('/tmp/model')
+        return model
+
     @staticmethod
     def insert_h2omojo_model(splice_context, run_id, model):
         model_exists = splice_context.df(
@@ -121,7 +143,29 @@ def insert_h2omojo_model(splice_context, run_id, model):
             insert_model(splice_context, run_id, byte_array, 'h2omojo', h2o.__version__)
 
 
+class SKUtils:
+    @staticmethod
+    def log_sklearn_model(splice_context, model, name, run_id):
+        byte_stream = save_pickle_string(model)
+        insert_artifact(splice_context, name, byte_stream, run_id, file_ext=FileExtensions.sklearn)
 
+    @staticmethod
+    def load_sklearn_model(model_blob):
+        return load_pickle_string(model_blob)
+
+class KerasUtils:
+    @staticmethod
+    def log_keras_model(splice_context, model, name, run_id):
+        model.save('/tmp/model.h5')
+        with open('/tmp/model.h5', 'rb') as f:
+            byte_stream = bytearray(bytes(f.read()))
+        insert_artifact(splice_context, name, byte_stream, run_id, file_ext=FileExtensions.keras)
+        remove('/tmp/model.h5')
+
+    @staticmethod
+    def load_keras_model(model_blob):
+        hfile = h5_file(BytesIO(model_blob), 'r')
+        return load_kr_model(hfile)
 
 class SparkUtils:
     @staticmethod
@@ -285,8 +329,8 @@ def get_model_type(pipeline_or_model):
 
         return m_type
     @staticmethod
-    def log_spark_model(splice_ctx, model, name, run_id):
-        jvm = splice_ctx.jvm
+    def log_spark_model(splice_context, model, name, run_id):
+        jvm = splice_context.jvm
         java_import(jvm, "java.io.{BinaryOutputStream, ObjectOutputStream, ByteArrayInputStream}")
 
         if not SparkUtils.is_spark_pipeline(model):
@@ -299,8 +343,8 @@ def log_spark_model(splice_ctx, model, name, run_id):
         oos.writeObject(model._to_java())
         oos.flush()
         oos.close()
-        insert_artifact(splice_ctx, name, baos.toByteArray(), run_id,
-                    file_ext='sparkmodel')  # write the byte stream to the db as a BLOB
+        insert_artifact(splice_context, name, baos.toByteArray(), run_id,
+                    file_ext='spark')  # write the byte stream to the db as a BLOB
 
     @staticmethod
     def load_spark_model(splice_ctx, spark_pipeline_blob):
@@ -456,13 +500,13 @@ def get_mleap_model(splice_context, fittedPipe, df, run_id: str):
         bash('mkdir /tmp')
     # Serialize the Spark model into Mleap format
     if f'{run_id}.zip' in rbash('ls /tmp').read():
-        bash(f'rm /tmp/{run_id}.zip')
+        remove(f'/tmp/{run_id}.zip')
     fittedPipe.serializeToBundle(f"jar:file:///tmp/{run_id}.zip", df)
 
     jvm = splice_context.jvm
     java_import(jvm, "com.splicemachine.mlrunner.FileRetriever")
     obj = jvm.FileRetriever.loadBundle(f'jar:file:///tmp/{run_id}.zip')
-    bash(f'rm /tmp/{run_id}.zip"')
+    remove(f'/tmp/{run_id}.zip')
     return obj
 
 
@@ -522,7 +566,6 @@ def create_data_table(splice_context, schema_table_name, schema_str, primary_key
             f'A model has already been deployed to table {schema_table_name}. We currently only support deploying 1 model per table')
     SQL_TABLE = f'CREATE TABLE {schema_table_name} (\n' + schema_str
 
-    # FIXME: Add the run_id as a column with constant default value to always be the run_id
     pk_cols = ''
     for i in primary_key:
         # If pk is already in the schema_string, don't add another column. PK may be an existing value