Selectable code env for MLFlow models

dataikuapi/dss/savedmodel.py

@@ -117,15 +117,17 @@ 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
@@ -135,7 +137,7 @@ def import_mlflow_version_from_path(self, version_id, path):
         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)})
 
         return self.get_mlflow_version_handler(version_id)
@@ -232,13 +234,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,