Merge pull request #198 from dataiku/feature/sc-77060-refactor-setup-mlflow-context-manager

Refactor setup_mlflow context manager to be used either as a context manager or using the provided MLflowHandle

Author: Valentin Thorey

dataikuapi/dss_plugin_mlflow/__init__.py

@@ -1 +1 @@
-from .utils import load_dss_mlflow_plugin
+from .utils import MLflowHandle

dataikuapi/dss_plugin_mlflow/utils.py

@@ -1,26 +1,76 @@
 import os
 import sys
 import tempfile
+import shutil
+from base64 import b64encode
 
 
-def load_dss_mlflow_plugin():
-    """ Function to dynamically add entrypoints for MLflow
-
-    MLflow uses entrypoints==0.3 to load entrypoints from plugins at import time.
-    This function adds dss-mlflow-plugin entrypoints dynamically by adding them in sys.path
-    at call time.
-    """
-    tempdir = tempfile.mkdtemp()
-    plugin_dir = os.path.join(tempdir, "dss-plugin-mlflow.egg-info")
-    if not os.path.isdir(plugin_dir):
-        os.makedirs(plugin_dir)
-    with open(os.path.join(plugin_dir, "entry_points.txt"), "w") as f:
-        f.write(
-            "[mlflow.request_header_provider]\n"
-            "unused=dataikuapi.dss_plugin_mlflow.header_provider:PluginDSSHeaderProvider\n"
-            "[mlflow.artifact_repository]\n"
-            "dss-managed-folder=dataikuapi.dss_plugin_mlflow.artifact_repository:PluginDSSManagedFolderArtifactRepository\n"
-        )
-    # Load plugin
-    sys.path.insert(0, tempdir)
-    return tempdir
+class MLflowHandle:
+    def __init__(self, client, project_key, managed_folder="mlflow_artifacts", host=None):
+        """ Add the MLflow-plugin parts of dataikuapi to MLflow local setup.
+
+        This method deals with
+        1. importing dynamically the DSS MLflow plugin:
+        MLflow uses entrypoints==0.3 to load entrypoints from plugins at import time.
+        We add dss-mlflow-plugin entrypoints dynamically by adding them in sys.path
+        at call time.
+
+        2. Setup the correct environment variables to setup the plugin to work
+        with DSS backend as the tracking backend and enable using DSS managed folder
+        as artifact location.
+        """
+        self.project_key = project_key
+        self.client = client
+        self.mlflow_env = {}
+
+        # Load DSS as the plugin
+        self.tempdir = tempfile.mkdtemp()
+        plugin_dir = os.path.join(self.tempdir, "dss-plugin-mlflow.egg-info")
+        if not os.path.isdir(plugin_dir):
+            os.makedirs(plugin_dir)
+        with open(os.path.join(plugin_dir, "entry_points.txt"), "w") as f:
+            f.write(
+                "[mlflow.request_header_provider]\n"
+                "unused=dataikuapi.dss_plugin_mlflow.header_provider:PluginDSSHeaderProvider\n"
+                "[mlflow.artifact_repository]\n"
+                "dss-managed-folder=dataikuapi.dss_plugin_mlflow.artifact_repository:PluginDSSManagedFolderArtifactRepository\n"
+            )
+        sys.path.insert(0, self.tempdir)
+
+        # Setup authentication
+        if client._session.auth is not None:
+            self.mlflow_env.update({
+                "DSS_MLFLOW_HEADER": "Authorization",
+                "DSS_MLFLOW_TOKEN": "Basic {}".format(
+                    b64encode("{}:".format(self.client._session.auth.username).encode("utf-8")).decode("utf-8")),
+                "DSS_MLFLOW_APIKEY": self.client.api_key
+            })
+        elif client.internal_ticket:
+            self.mlflow_env.update({
+                "DSS_MLFLOW_HEADER": "X-DKU-APITicket",
+                "DSS_MLFLOW_TOKEN": self.client.internal_ticket,
+                "DSS_MLFLOW_INTERNAL_TICKET": self.client.internal_ticket
+            })
+        # Set host, tracking URI and project key
+        self.mlflow_env.update({
+            "DSS_MLFLOW_PROJECTKEY": project_key,
+            "MLFLOW_TRACKING_URI": self.client.host + "/dip/publicapi" if host is None else host,
+            "DSS_MLFLOW_HOST": self.client.host,
+            "DSS_MLFLOW_MANAGED_FOLDER": managed_folder
+        })
+        os.environ.update(self.mlflow_env)
+
+    def clear(self):
+        shutil.rmtree(self.tempdir)
+        for variable in self.mlflow_env:
+            os.environ.pop(variable, None)
+
+    def import_mlflow(self):
+        import mlflow
+        return mlflow
+
+    def __enter__(self):
+        return self.import_mlflow()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.clear()

dataikuapi/dssclient.py

@@ -1,14 +1,11 @@
 import json
-import os
-import shutil
 
 from requests import Session
 from requests import exceptions
 from requests.auth import HTTPBasicAuth
-from contextlib import contextmanager
 
 from dataikuapi.dss.notebook import DSSNotebook
-from .dss_plugin_mlflow import load_dss_mlflow_plugin
+from .dss_plugin_mlflow import MLflowHandle
 from .dss.future import DSSFuture
 from .dss.projectfolder import DSSProjectFolder
 from .dss.project import DSSProject
@@ -21,7 +18,6 @@
 from .dss.apideployer import DSSAPIDeployer
 from .dss.projectdeployer import DSSProjectDeployer
 import os.path as osp
-from base64 import b64encode
 from .utils import DataikuException, dku_basestring_type
 
 class DSSClient(object):
@@ -1095,7 +1091,6 @@ def get_object_discussions(self, project_key, object_type, object_id):
     ########################################################
     # MLflow
     ########################################################
-    @contextmanager
     def setup_mlflow(self, project_key, managed_folder="mlflow_artifacts", host=None):
         """
         Setup the dss-plugin for MLflow
@@ -1104,38 +1099,7 @@ def setup_mlflow(self, project_key, managed_folder="mlflow_artifacts", host=None
         :param str managed_folder: managed folder where artifacts are stored
         :param str host: setup a custom host if the backend used is not DSS
         """
-        tempdir = load_dss_mlflow_plugin()
-        mlflow_env = {}
-        if self._session.auth is not None:
-            mlflow_env.update({
-                "DSS_MLFLOW_HEADER": "Authorization",
-                "DSS_MLFLOW_TOKEN": "Basic {}".format(
-                    b64encode("{}:".format(self._session.auth.username).encode("utf-8")).decode("utf-8")),
-                "DSS_MLFLOW_APIKEY": self.api_key
-            })
-        elif self.internal_ticket:
-            mlflow_env.update({
-                "DSS_MLFLOW_HEADER": "X-DKU-APITicket",
-                "DSS_MLFLOW_TOKEN": self.internal_ticket,
-                "DSS_MLFLOW_INTERNAL_TICKET": self.internal_ticket
-            })
-        mlflow_env.update({
-            "DSS_MLFLOW_PROJECTKEY": project_key,
-            "MLFLOW_TRACKING_URI": self.host + "/dip/publicapi" if host is None else host,
-            "DSS_MLFLOW_HOST": self.host,
-            "DSS_MLFLOW_MANAGED_FOLDER": managed_folder,
-        })
-        os.environ.update(mlflow_env)
-
-        try:
-            import mlflow
-            yield mlflow
-        except Exception as e:
-            raise e
-        finally:
-            shutil.rmtree(tempdir)
-            for variable in mlflow_env:
-                os.environ.pop(variable, None)
+        return MLflowHandle(client=self, project_key=project_key, managed_folder=managed_folder, host=host)
 
 
 class TemporaryImportHandle(object):