MLFlow next steps

dataikuapi/__init__.py

@@ -1,8 +1,9 @@
 from .dssclient import DSSClient
+from .fmclient import FMClientAWS, FMClientAzure
 
 from .apinode_client import APINodeClient
 from .apinode_admin_client import APINodeAdminClient
 
 from .dss.recipe import GroupingRecipeCreator, JoinRecipeCreator, StackRecipeCreator, WindowRecipeCreator, SyncRecipeCreator, SamplingRecipeCreator, SQLQueryRecipeCreator, CodeRecipeCreator, SplitRecipeCreator, SortRecipeCreator, TopNRecipeCreator, DistinctRecipeCreator, DownloadRecipeCreator, PredictionScoringRecipeCreator, ClusteringScoringRecipeCreator
 
-from .dss.admin import DSSUserImpersonationRule, DSSGroupImpersonationRule
+from .dss.admin import DSSUserImpersonationRule, DSSGroupImpersonationRule

dataikuapi/dss/ml.py

@@ -1867,7 +1867,7 @@ def get_diagnostics(self):
         :returns: list of diagnostics
         :rtype: list of type `dataikuapi.dss.ml.DSSMLDiagnostic`
         """
-        diagnostics = self.details.get("trainDiagnostics", {})
+        diagnostics = self.details.get("mlDiagnostics", {})
         return [DSSMLDiagnostic(d) for d in diagnostics.get("diagnostics", [])]
 
     def generate_documentation(self, folder_id=None, path=None):

dataikuapi/dss/modelevaluationstore.py

@@ -124,18 +124,18 @@ def list_model_evaluations(self):
         :returns: The list of the model evaluations
         :rtype: list of :class:`dataikuapi.dss.modelevaluationstore.DSSModelEvaluation`
         """
-        items = self.client._perform_json("GET", "/projects/%s/modelevaluationstores/%s/runs/" % (self.project_key, self.mes_id))
-        return [DSSModelEvaluation(self, item["ref"]["runId"]) for item in items]
+        items = self.client._perform_json("GET", "/projects/%s/modelevaluationstores/%s/evaluations/" % (self.project_key, self.mes_id))
+        return [DSSModelEvaluation(self, item["ref"]["evaluationId"]) for item in items]
 
-    def get_model_evaluation(self, run_id):
+    def get_model_evaluation(self, evaluation_id):
         """
         Get a handle to interact with a specific model evaluation
 
-        :param string run_id: the id of the desired model evaluation
+        :param string evaluation_id: the id of the desired model evaluation
 
         :returns: A :class:`dataikuapi.dss.modelevaluationstore.DSSModelEvaluation` model evaluation handle
         """
-        return DSSModelEvaluation(self, run_id)
+        return DSSModelEvaluation(self, evaluation_id)
 
     def get_latest_model_evaluation(self):
         """
@@ -146,11 +146,11 @@ def get_latest_model_evaluation(self):
             if the store is not empty, else None
         """
 
-        latest_run_id = self.client._perform_text(
-            "GET", "/projects/%s/modelevaluationstores/%s/latestRunId" % (self.project_key, self.mes_id))
-        if not latest_run_id:
+        latest_evaluation_id = self.client._perform_text(
+            "GET", "/projects/%s/modelevaluationstores/%s/latestEvaluationId" % (self.project_key, self.mes_id))
+        if not latest_evaluation_id:
             return None
-        return DSSModelEvaluation(self, latest_run_id)
+        return DSSModelEvaluation(self, latest_evaluation_id)
 
     def delete_model_evaluations(self, evaluations):
         """
@@ -159,13 +159,13 @@ def delete_model_evaluations(self, evaluations):
         obj = []
         for evaluation in evaluations:
             if isinstance(evaluation, DSSModelEvaluation):
-                obj.append(evaluation.run_id)
+                obj.append(evaluation.evaluation_id)
             elif isinstance(evaluation, dict):
-                obj.append(evaluation['run_id'])
+                obj.append(evaluation['evaluation_id'])
             else:
                 obj.append(evaluation)
         self.client._perform_json(
-                "DELETE", "/projects/%s/modelevaluationstores/%s/runs/" % (self.project_key, self.mes_id, self.run_id), body=obj)
+                "DELETE", "/projects/%s/modelevaluationstores/%s/evaluations/" % (self.project_key, self.mes_id), body=obj)
 
     def build(self, job_type="NON_RECURSIVE_FORCED_BUILD", wait=True, no_fail=False):
         """
@@ -263,11 +263,11 @@ class DSSModelEvaluation:
     Do not create this class directly, instead use :meth:`dataikuapi.dss.DSSModelEvaluationStore.get_model_evaluation`
     """
 
-    def __init__(self, model_evaluation_store, run_id):
+    def __init__(self, model_evaluation_store, evaluation_id):
         self.model_evaluation_store = model_evaluation_store
         self.client = model_evaluation_store.client
         # unpack some fields
-        self.run_id = run_id
+        self.evaluation_id = evaluation_id
         self.project_key = model_evaluation_store.project_key
         self.mes_id = model_evaluation_store.mes_id
 
@@ -276,23 +276,23 @@ def get_full_info(self):
         Retrieve the model evaluation with its performance data
         """
         data = self.client._perform_json(
-            "GET", "/projects/%s/modelevaluationstores/%s/runs/%s" % (self.project_key, self.mes_id, self.run_id))
+            "GET", "/projects/%s/modelevaluationstores/%s/evaluations/%s" % (self.project_key, self.mes_id, self.evaluation_id))
         return DSSModelEvaluationFullInfo(self, data)
 
     def get_full_id(self):
-        return "ME-{}-{}-{}".format(self.project_key, self.mes_id, self.run_id)
+        return "ME-{}-{}-{}".format(self.project_key, self.mes_id, self.evaluation_id)
 
     def delete(self):
         """
         Remove this model evaluation
         """
-        obj = [self.run_id]
+        obj = [self.evaluation_id]
         self.client._perform_json(
-                "DELETE", "/projects/%s/modelevaluationstores/%s/runs/" % (self.project_key, self.mes_id), body=obj)
+                "DELETE", "/projects/%s/modelevaluationstores/%s/evaluations/" % (self.project_key, self.mes_id), body=obj)
 
     @property
     def full_id(self):
-        return "ME-%s-%s-%s"%(self.project_key, self.mes_id, self.run_id)
+        return "ME-%s-%s-%s"%(self.project_key, self.mes_id, self.evaluation_id)
 
     def compute_data_drift(self, reference=None, data_drift_params=None, wait=True):
         """
@@ -310,7 +310,7 @@ def compute_data_drift(self, reference=None, data_drift_params=None, wait=True):
             reference = reference.full_id
 
         future_response = self.client._perform_json(
-            "POST", "/projects/%s/modelevaluationstores/%s/runs/%s/computeDataDrift" % (self.project_key, self.mes_id, self.run_id),
+            "POST", "/projects/%s/modelevaluationstores/%s/evaluations/%s/computeDataDrift" % (self.project_key, self.mes_id, self.evaluation_id),
             body={
                 "referenceId": reference,
                 "dataDriftParams": data_drift_params
@@ -325,7 +325,7 @@ def get_metrics(self):
         :return: the metrics, as a JSON object
         """
         return self.client._perform_json(
-            "GET", "/projects/%s/modelevaluationstores/%s/runs/%s/metrics" % (self.project_key, self.mes_id, self.run_id))
+            "GET", "/projects/%s/modelevaluationstores/%s/evaluations/%s/metrics" % (self.project_key, self.mes_id, self.evaluation_id))
 
     def get_sample_df(self):
         """
@@ -337,12 +337,12 @@ def get_sample_df(self):
         buf = BytesIO()
         with self.client._perform_raw(
                 "GET",
-                "/projects/%s/modelevaluationstores/%s/runs/%s/sample" % (self.project_key, self.mes_id, self.run_id)
+                "/projects/%s/modelevaluationstores/%s/evaluations/%s/sample" % (self.project_key, self.mes_id, self.evaluation_id)
         ).raw as f:
             buf.write(f.read())
         schema_txt = self.client._perform_raw(
             "GET",
-            "/projects/%s/modelevaluationstores/%s/runs/%s/schema" % (self.project_key, self.mes_id, self.run_id)
+            "/projects/%s/modelevaluationstores/%s/evaluations/%s/schema" % (self.project_key, self.mes_id, self.evaluation_id)
         ).text
         schema = json.loads(schema_txt)
         import pandas as pd

dataikuapi/dss/project.py

@@ -709,22 +709,22 @@ def get_saved_model(self, sm_id):
         """
         return DSSSavedModel(self.client, self.project_key, sm_id)
 
-    def create_mlflow_pyfunc_model(self, id, prediction_type = None):
+    def create_mlflow_pyfunc_model(self, name, prediction_type = None):
         """
         Creates a new external saved model for storing and managing MLFlow models
 
-        :param string id: Identifier for the new saved model in the flow
+        :param string name: Human readable name 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")
+        if not name:
+            raise ValueError("name can not be empty")
         model = {
-            "id": id,
             "savedModelType" : "MLFLOW_PYFUNC",
-            "predictionType" : prediction_type
+            "predictionType" : prediction_type,
+            "name": name
         }
 
-        self.client._perform_empty("POST", "/projects/%s/savedmodels/" % self.project_key, body = model)
+        id = self.client._perform_json("POST", "/projects/%s/savedmodels/" % self.project_key, body = model)["id"]
         return self.get_saved_model(id)
 
     ########################################################

dataikuapi/dss/savedmodel.py

@@ -117,26 +117,30 @@ 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_mlflow_version_from_path(self, version_id, path):
+    def import_mlflow_version_from_path(self, version_id, path, code_env_name = "INHERIT"):
         """
         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
-
+        :param str code_env_name: Name of the code env to use for this model version. The code env must contain at least 
+                                  mlflow and the package(s) corresponding to the used MLFlow-compatible frameworks.
+                                  If value is "INHERIT", the default active code env of the project will be used
         :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
+        import os
         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),
+            self.client._perform_empty("POST", "/projects/%s/savedmodels/%s/versions/%s?codeEnvName=%s" % (self.project_key, self.sm_id, version_id, code_env_name),
                 files={"file":("tmpmodel.zip", fp)})
+        os.remove("tmpmodel.zip")
 
         return self.get_mlflow_version_handler(version_id)
 
@@ -232,13 +236,33 @@ def delete(self):
         """
         return self.client._perform_empty("DELETE", "/projects/%s/savedmodels/%s" % (self.project_key, self.sm_id))
 
+class MLFlowVersionSettings:
+    """Handle for the settings of an imported MLFlow model version"""
+
+    def __init__(self, version_handler, data):
+        self.version_handler = version_handler
+        self.data = data
+
+    @property
+    def raw(self):
+        return self.data
+
+    def save(self):
+        self.version_handler.saved_model.client._perform_empty("PUT", 
+            "/projects/%s/savedmodels/%s/versions/%s/external-ml/metadata" % (self.version_handler.saved_model.project_key, self.version_handler.saved_model.sm_id, self.version_handler.version_id),
+            body=self.data)
+
 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 get_settings(self):
+        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))
+        return MLFlowVersionSettings(self, metadata)
+
     def set_core_metadata(self,
         target_column_name, class_labels = None,
         get_features_from_dataset=None, features_list = None,

dataikuapi/fm/future.py

@@ -0,0 +1,101 @@
+import sys, time
+
+
+class FMFuture(object):
+    """
+    A future on the DSS instance
+    """
+
+    def __init__(
+        self, client, job_id, state=None, result_wrapper=lambda result: result
+    ):
+        self.client = client
+        self.job_id = job_id
+        self.state = state
+        self.state_is_peek = True
+        self.result_wrapper = result_wrapper
+
+    @staticmethod
+    def from_resp(client, resp, result_wrapper=lambda result: result):
+        """Creates a DSSFuture from a parsed JSON response"""
+        return FMFuture(
+            client, resp.get("jobId", None), state=resp, result_wrapper=result_wrapper
+        )
+
+    @classmethod
+    def get_result_wait_if_needed(cls, client, ret):
+        if "jobId" in ret:
+            future = FMFuture(client, ret["jobId"], ret)
+            future.wait_for_result()
+            return future.get_result()
+        else:
+            return ret["result"]
+
+    def abort(self):
+        """
+        Abort the future
+        """
+        return self.client._perform_tenant_empty("DELETE", "/futures/%s" % self.job_id)
+
+    def get_state(self):
+        """
+        Get the status of the future, and its result if it's ready
+        """
+        self.state = self.client._perform_tenant_json(
+            "GET", "/futures/%s" % self.job_id, params={"peek": False}
+        )
+        self.state_is_peek = False
+        return self.state
+
+    def peek_state(self):
+        """
+        Get the status of the future, and its result if it's ready
+        """
+        self.state = self.client._perform_tenant_json(
+            "GET", "/futures/%s" % self.job_id, params={"peek": True}
+        )
+        self.state_is_peek = True
+        return self.state
+
+    def get_result(self):
+        """
+        Get the future result if it's ready, raises an Exception otherwise
+        """
+        if (
+            self.state is None
+            or not self.state.get("hasResult", False)
+            or self.state_is_peek
+        ):
+            self.get_state()
+        if self.state.get("hasResult", False):
+            return self.result_wrapper(self.state.get("result", None))
+        else:
+            raise Exception("Result not ready")
+
+    def has_result(self):
+        """
+        Checks whether the future has a result ready
+        """
+        if self.state is None or not self.state.get("hasResult", False):
+            self.get_state()
+        return self.state.get("hasResult", False)
+
+    def wait_for_result(self):
+        """
+        Wait and get the future result
+        """
+        if self.state.get("hasResult", False):
+            return self.result_wrapper(self.state.get("result", None))
+        if (
+            self.state is None
+            or not self.state.get("hasResult", False)
+            or self.state_is_peek
+        ):
+            self.get_state()
+        while not self.state.get("hasResult", False):
+            time.sleep(5)
+            self.get_state()
+        if self.state.get("hasResult", False):
+            return self.result_wrapper(self.state.get("result", None))
+        else:
+            raise Exception("No result")