Revert "Refactor setup_mlflow contextmanager"

This reverts commit 9e5c95b.

Author: Valentin Thorey

dataikuapi/dss_plugin_mlflow/__init__.py

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

dataikuapi/dss_plugin_mlflow/utils.py

@@ -1,75 +1,26 @@
 import os
 import sys
 import tempfile
-import shutil
-from base64 import b64encode
 
 
-class MLflowHandle:
-    def __init__(self, client, project_key, managed_folder, host=None):
-        """ Add the MLflow-plugin parts of dataikuapi to MLflow local setup.
-
-        This functions 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 environemnent 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 self.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,
-        })
-        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()
+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

dataikuapi/dssclient.py

@@ -1,11 +1,14 @@
 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 MLflowHandle
+from .dss_plugin_mlflow import load_dss_mlflow_plugin
 from .dss.future import DSSFuture
 from .dss.projectfolder import DSSProjectFolder
 from .dss.project import DSSProject
@@ -18,6 +21,7 @@
 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):
@@ -1091,6 +1095,7 @@ 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
@@ -1099,7 +1104,38 @@ 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
         """
-        return MLflowHandle(client=self, project_key=project_key, managed_folder=managed_folder, host=host)
+        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)
 
 
 class TemporaryImportHandle(object):