Merge pull request #200 from dataiku/feature/sc-77250-reload-mlflow-artifact-registry

Give MLflowHandle the ability to reload artifact_repository_registry

Author: Valentin Thorey

dataikuapi/dss/project.py

@@ -1,4 +1,7 @@
 import time, warnings, sys, os.path as osp
+
+from ..dss_plugin_mlflow import MLflowHandle
+
 from .dataset import DSSDataset, DSSDatasetListItem, DSSManagedDatasetCreationHelper
 from .modelcomparison import DSSModelComparison
 from .jupyternotebook import DSSJupyterNotebook, DSSJupyterNotebookListItem
@@ -1581,7 +1584,7 @@ def list_hive_tables(self, hive_database):
         """
         connection_name = "@virtual(hive-jdbc):" + hive_database
         ret = self.client._perform_json("GET", "/projects/%s/datasets/tables-import/actions/list-tables" % (self.project_key),
-                params = {"connectionName": connection_name} )
+                params={"connectionName": connection_name} )
 
         def to_schema_table_pair(x):
             return {"schema":x.get("databaseName", None), "table":x["table"]}
@@ -1590,14 +1593,21 @@ def to_schema_table_pair(x):
     ########################################################
     # App designer
     ########################################################
-
     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 experiment tracking
     ########################################################
+    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)
+
     def clean_experiment_tracking_db(self):
         """
         Cleans the experiments, runs, params, metrics, tags, etc. for this project

dataikuapi/dss_plugin_mlflow/utils.py

@@ -37,6 +37,12 @@ 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
+        if sys.version_info > (3, 4):
+            from importlib import reload
+        import mlflow
+        reload(mlflow.store.artifact.artifact_repository_registry)
+
         # Setup authentication
         if client._session.auth is not None:
             self.mlflow_env.update({

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):