Move setup_mlflow to DSSProject

Author: Valentin Thorey

dataikuapi/dss/project.py

@@ -21,6 +21,7 @@
 from .analysis import DSSAnalysis
 from .flow import DSSProjectFlow
 from .app import DSSAppManifest
+from .dss_plugin_mlflow import MLflowHandle
 
 
 class DSSProject(object):
@@ -1595,6 +1596,18 @@ def get_app_manifest(self):
         raw_data = self.client._perform_json("GET", "/projects/%s/app-manifest" % self.project_key)
         return DSSAppManifest(self.client, raw_data, self.project_key)
 
+    ########################################################
+    # MLflow
+    ########################################################
+    def setup_mlflow(self, managed_folder="mlflow_artifacts", host=None):
+        """
+        Setup the dss-plugin for MLflow
+
+        :param str managed_folder: managed folder where artifacts are stored
+        :param str host: setup a custom host if the backend used is not DSS
+        """
+        return MLflowHandle(client=self.client, project_key=self.project_key, managed_folder=managed_folder, host=host)
+
 
 class TablesImportDefinition(object):
     """

dataikuapi/dss_plugin_mlflow/utils.py

@@ -38,11 +38,9 @@ def __init__(self, client, project_key, managed_folder="mlflow_artifacts", host=
         sys.path.insert(0, self.tempdir)
 
         # Reload the artifact_repository_registry in case MLflow was imported beforehand
-        import mlflow
-        try:  # python 3.4+
+        if sys.version_info > (3, 4):
             from importlib import reload
-        except ImportError:
-            pass
+        import mlflow
         reload(mlflow.store.artifact.artifact_repository_registry)
 
         # Setup authentication

dataikuapi/dssclient.py

@@ -5,7 +5,6 @@
 from requests.auth import HTTPBasicAuth
 
 from dataikuapi.dss.notebook import DSSNotebook
-from .dss_plugin_mlflow import MLflowHandle
 from .dss.future import DSSFuture
 from .dss.projectfolder import DSSProjectFolder
 from .dss.project import DSSProject
@@ -1088,19 +1087,6 @@ def get_object_discussions(self, project_key, object_type, object_id):
         """
         return DSSObjectDiscussions(self, project_key, object_type, object_id)
 
-    ########################################################
-    # MLflow
-    ########################################################
-    def setup_mlflow(self, project_key, managed_folder="mlflow_artifacts", host=None):
-        """
-        Setup the dss-plugin for MLflow
-
-        :param str project_key: identifier of the project to access
-        :param str managed_folder: managed folder where artifacts are stored
-        :param str host: setup a custom host if the backend used is not DSS
-        """
-        return MLflowHandle(client=self, project_key=project_key, managed_folder=managed_folder, host=host)
-
 
 class TemporaryImportHandle(object):
     def __init__(self, client, import_id):