Updates to MLFlow import API

Author: cstenac

dataikuapi/dss/project.py

@@ -699,18 +699,21 @@ def get_saved_model(self, sm_id):
         """
         return DSSSavedModel(self.client, self.project_key, sm_id)
 
-    def create_mlflow_pyfunc_model(self, id):
-        """Starts creation of a new external saved model for storing 3rd party models
+    def create_mlflow_pyfunc_model(self, id, prediction_type = None):
+        """
+        Creates a new external saved model for storing and managing MLFlow models
 
-        :param string id: 
+        :param string id: Identifier for the new saved model in the flow
+        :param string prediction_type: Optional (but needed for most operations). One of BINARY_CLASSIFICATION, MULTICLASS or REGRESSION
         """
         if len(id) != 8:
             raise ValueError("model id must be 8 characters long")
         model = {
-            "projectKey" : self.project_key,
             "id": id,
-            "savedModelType" : "MLFLOW_PYFUNC"
+            "savedModelType" : "MLFLOW_PYFUNC",
+            "predictionType" : prediction_type
         }
+
         self.client._perform_empty("POST", "/projects/%s/savedmodels/" % self.project_key, body = model)
         return self.get_saved_model(id)
 

dataikuapi/dss/savedmodel.py

@@ -117,14 +117,35 @@ def get_origin_ml_task(self):
         if fmi is not None:
             return DSSMLTask.from_full_model_id(self.client, fmi, project_key=self.project_key)
 
-    def import_version_from_folder(self, version_id, folder):
+    def import_mlflow_version_from_path(self, version_id, path):
+        """
+        Create a new version for this saved model from a path containing a MLFlow model.
+
+        Requires the saved model to have been created using :meth:`dataikuapi.dss.project.DSSProject.create_mlflow_pyfunc_model`.
+
+        :param str version_id: Identifier of the version to create
+        :param str path: An absolute path on the local filesystem. Must be a folder, and must contain a MLFlow model
+
+        :return a :class:MLFlowVersionHandler in order to interact with the new MLFlow model version
+        """
+        # TODO: Add a check that it's indeed a MLFlow model folder
+        # TODO: Put it in a proper temp folder
+        # TODO: cleanup the archive
         import shutil
-        shutil.make_archive("tmpmodel", "zip", folder) #[, root_dir[, base_dir[, verbose[, dry_run[, owner[, group[, logger]]]]]]])
+        shutil.make_archive("tmpmodel", "zip", path) #[, root_dir[, base_dir[, verbose[, dry_run[, owner[, group[, logger]]]]]]])
 
         with open("tmpmodel.zip", "rb") as fp:
             self.client._perform_empty("POST", "/projects/%s/savedmodels/%s/versions/%s" % (self.project_key, self.sm_id, version_id),
                 files={"file":("tmpmodel.zip", fp)})
 
+        return self.get_mlflow_version_handler(version_id)
+
+    def get_mlflow_version_handler(self, version_id):
+        """
+        Returns a :class:MLFlowVersionHandler to interact with a MLFlow model version
+        """
+        return MLFlowVersionHandler(self, version_id)
+
     ########################################################
     # Metrics
     ########################################################
@@ -211,6 +232,69 @@ def delete(self):
         """
         return self.client._perform_empty("DELETE", "/projects/%s/savedmodels/%s" % (self.project_key, self.sm_id))
 
+class MLFlowVersionHandler:
+    """Handler to interact with an imported MLFlow model version"""
+    def __init__(self, saved_model, version_id):
+        """Do not call this, use :meth:`DSSSavedModel.get_mlflow_version_handler`"""
+        self.saved_model = saved_model
+        self.version_id = version_id
+
+    def set_core_metadata(self,
+        target_column_name, class_labels = None,
+        get_features_from_dataset=None, features_list = None,
+        output_style="AUTO_DETECT"):
+        """
+        Sets metadata for this MLFlow model version
+
+        In addition to target_column_name, one of get_features_from_dataset or features_list must be passed in order
+        to be able to evaluate performance
+
+        :param str target_column_name: name of the target column. Mandatory in order to be able to evaluate performance
+        :param list class_labels: List of strings, ordered class labels. Mandatory in order to be able to evaluate performance on classification models
+        :param str get_features_from_dataset: Name of a dataset to get feature names from
+        :param list features_list: List of {"name": "feature_name", "type": "feature_type"}
+        """
+
+        metadata = self.saved_model.client._perform_json("GET", "/projects/%s/savedmodels/%s/versions/%s/external-ml/metadata" % (self.saved_model.project_key, self.saved_model.sm_id, self.version_id))
+
+        if target_column_name is not None:
+            metadata["targetColumnName"] = target_column_name
+
+        if class_labels is not None:
+            metadata["classLabels"] = [{"label": l} for l in class_labels]
+
+        if get_features_from_dataset is not None:
+            metadata["gatherFeaturesFromDataset"] = get_features_from_dataset
+
+        # TODO: add support for get_features_from_signature=False,
+        #if get_features_from_signature:
+        #    raise Exception("Get features from signature is not yet implemented")
+
+        # TODO: Add support for features_list, with validation
+
+        self.saved_model.client._perform_empty("PUT", 
+            "/projects/%s/savedmodels/%s/versions/%s/external-ml/metadata" % (self.saved_model.project_key, self.saved_model.sm_id, self.version_id),
+            body=metadata)
+
+    def evaluate(self, dataset_ref):
+        """
+        Evaluates the performance of this model version on a particular dataset.
+        After calling this, the "result screens" of the MLFlow model version will be available
+        (confusion matrix, error distribution, performance metrics, ...)
+        and more information will be available when calling :meth:`DSSSavedModel.get_version_details`
+
+        :meth:`set_core_metadata` must be called before you can evaluate a dataset
+
+        :param str dataset_ref: Name of the evaluation dataset to use (either a dataset name or "PROJECT.datasetName")
+        """
+        # TODO Add support for handling a DSSDataset or dataiku.Dataset as dataset_ref
+        req = {
+            "datasetRef" : dataset_ref
+        }
+        self.saved_model.client._perform_empty("POST",
+            "/projects/%s/savedmodels/%s/versions/%s/external-ml/actions/evaluate" % (self.saved_model.project_key, self.saved_model.sm_id, self.version_id),
+            body=req)
+
 
 class DSSSavedModelSettings:
     """
@@ -225,3 +309,13 @@ def __init__(self, saved_model, settings):
 
     def get_raw(self):
         return self.settings
+
+    @property
+    def prediction_metrics_settings(self):
+        """The settings of evaluation metrics for a prediction saved model"""
+        return self.settings["miniTask"]["modeling"]["metrics"]
+
+    def save(self):
+        """Saves the settings of this saved model"""
+        self.saved_model.client._perform_empty("PUT", "/projects/%s/savedmodels/%s" % (self.saved_model.project_key, self.saved_model.sm_id),
+                    body=self.settings)