Add DSS plugin for MLflow (#191)

Valentin Thorey · lpenet · web-flow · commit 67cf94b30b4c · 2021-12-09T09:08:16.000+01:00
* Add DSS plugin for MLflow

* Add doc to mlflow plugin

* Fix inconsistent wording

* Prevent creating multiple tempdir for mlflow plugin

* Set host while setting up mlflow plugin

* Rename MLflow plugin for consistency

* Add artifact part of the mlflow plugin

* Add log_artifact to mlflow plugin

* Add log_artifacts and list_artifacts

* Adapt wording to backend

* Add deleta and download to artifact plugin

* Add context manager for dss plugin

* Revert formatting

* Fix typo

Co-authored-by: Ludovic Pénet &lt;ludovic.penet@dataiku.com&gt;

Co-authored-by: Ludovic Pénet &lt;ludovic.penet@dataiku.com&gt;
diff --git a/dataikuapi/dss_plugin_mlflow/__init__.py b/dataikuapi/dss_plugin_mlflow/__init__.py
@@ -0,0 +1 @@
+from .utils import load_dss_mlflow_plugin
diff --git a/dataikuapi/dss_plugin_mlflow/artifact_repository.py b/dataikuapi/dss_plugin_mlflow/artifact_repository.py
@@ -0,0 +1,244 @@
+import os
+import posixpath
+import tempfile
+import urllib
+from dataikuapi import DSSClient
+
+
+def parse_dss_managed_folder_uri(uri):
+    """Parse an S3 URI, returning (bucket, path)"""
+    parsed = urllib.parse.urlparse(uri)
+    if parsed.scheme != "dss-managed-folder":
+        raise Exception("Not a DSS Managed Folder URI: %s" % uri)
+    return os.path.normpath(parsed.path)
+
+
+class PluginDSSManagedFolderArtifactRepository:
+
+    def __init__(self, artifact_uri):
+        self.base_artifact_path = parse_dss_managed_folder_uri(artifact_uri)
+        if os.environ.get("DSS_MLFLOW_APIKEY") is not None:
+            self.client = DSSClient(
+                os.environ.get("DSS_MLFLOW_HOST"),
+                api_key=os.environ.get("DSS_MLFLOW_APIKEY")
+            )
+        else:
+            self.client = DSSClient(
+                os.environ.get("DSS_MLFLOW_HOST"),
+                internal_ticket=os.environ.get("DSS_MLFLOW_INTERNAL_TICKET")
+            )
+        self.project = self.client.get_project(os.environ.get("DSS_MLFLOW_PROJECTKEY"))
+        managed_folders = [
+            x["id"] for x in self.project.list_managed_folders()
+            if x["name"] == os.environ.get("DSS_MLFLOW_MANAGED_FOLDER")
+        ]
+        if len(managed_folders) > 0:
+            self.managed_folder = self.project.get_managed_folder(managed_folders[0])
+        else:
+            self.managed_folder = self.project.create_managed_folder(os.environ.get("DSS_MLFLOW_MANAGED_FOLDER"))
+
+    def log_artifact(self, local_file, artifact_path=None):
+        """
+        Log a local file as an artifact, optionally taking an ``artifact_path`` to place it in
+        within the run's artifacts. Run artifacts can be organized into directories, so you can
+        place the artifact in a directory this way.
+
+        :param local_file: Path to artifact to log
+        :param artifact_path: Directory within the run's artifact directory in which to log the
+                              artifact.
+        """
+        path = (
+            os.path.join(self.base_artifact_path, artifact_path) if artifact_path else self.base_artifact_path
+        )
+        self.managed_folder.put_file(os.path.join(path, os.path.basename(local_file)), open(local_file))
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        """
+        Log the files in the specified local directory as artifacts, optionally taking
+        an ``artifact_path`` to place them in within the run's artifacts.
+
+        :param local_dir: Directory of local artifacts to log
+        :param artifact_path: Directory within the run's artifact directory in which to log the
+                              artifacts
+        """
+        path = (
+            os.path.join(self.base_artifact_path, artifact_path) if artifact_path else self.base_artifact_path
+        )
+        self.managed_folder.upload_folder(path, local_dir)
+
+    def list_artifacts(self, path):
+        """
+        Return all the artifacts for this run_id directly under path. If path is a file, returns
+        an empty list. Will error if path is neither a file nor directory.
+
+        :param path: Relative source path that contains desired artifacts
+
+        :return: List of artifacts as FileInfo listed directly under path.
+        """
+        from pathlib import Path
+        path = Path(os.path.join(self.base_artifact_path, path))
+        files = [x["path"] for x in self.managed_folder.list_contents().get("items", []) if path in Path(x["path"]).parents]
+        return files
+
+    def _is_directory(self, artifact_path):
+        listing = self.list_artifacts(artifact_path)
+        return len(listing) > 0
+
+    def _create_download_destination(self, src_artifact_path, dst_local_dir_path=None):
+        """
+        Creates a local filesystem location to be used as a destination for downloading the artifact
+        specified by `src_artifact_path`. The destination location is a subdirectory of the
+        specified `dst_local_dir_path`, which is determined according to the structure of
+        `src_artifact_path`. For example, if `src_artifact_path` is `dir1/file1.txt`, then the
+        resulting destination path is `<dst_local_dir_path>/dir1/file1.txt`. Local directories are
+        created for the resulting destination location if they do not exist.
+
+        :param src_artifact_path: A relative, POSIX-style path referring to an artifact stored
+                                  within the repository's artifact root location.
+                                  `src_artifact_path` should be specified relative to the
+                                  repository's artifact root location.
+        :param dst_local_dir_path: The absolute path to a local filesystem directory in which the
+                                   local destination path will be contained. The local destination
+                                   path may be contained in a subdirectory of `dst_root_dir` if
+                                   `src_artifact_path` contains subdirectories.
+        :return: The absolute path to a local filesystem location to be used as a destination
+                 for downloading the artifact specified by `src_artifact_path`.
+        """
+        src_artifact_path = src_artifact_path.rstrip("/")  # Ensure correct dirname for trailing '/'
+        dirpath = posixpath.dirname(src_artifact_path)
+        local_dir_path = os.path.join(dst_local_dir_path, dirpath)
+        local_file_path = os.path.join(dst_local_dir_path, src_artifact_path)
+        if not os.path.exists(local_dir_path):
+            os.makedirs(local_dir_path, exist_ok=True)
+        return local_file_path
+
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it.
+        The caller is responsible for managing the lifecycle of the downloaded artifacts.
+
+        :param artifact_path: Relative source path to the desired artifacts.
+        :param dst_path: Absolute path of the local filesystem destination directory to which to
+                         download the specified artifacts. This directory must already exist.
+                         If unspecified, the artifacts will either be downloaded to a new
+                         uniquely-named directory on the local filesystem or will be returned
+                         directly in the case of the LocalArtifactRepository.
+
+        :return: Absolute path of the local filesystem location containing the desired artifacts.
+        """
+        from mlflow.exceptions import MlflowException
+        from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, RESOURCE_DOES_NOT_EXIST
+
+        # TODO: Probably need to add a more efficient method to stream just a single artifact
+        #       without downloading it, or to get a pre-signed URL for cloud storage.
+        def download_artifact(src_artifact_path, dst_local_dir_path):
+            """
+            Download the file artifact specified by `src_artifact_path` to the local filesystem
+            directory specified by `dst_local_dir_path`.
+
+            :param src_artifact_path: A relative, POSIX-style path referring to a file artifact
+                                      stored within the repository's artifact root location.
+                                      `src_artifact_path` should be specified relative to the
+                                      repository's artifact root location.
+            :param dst_local_dir_path: Absolute path of the local filesystem destination directory
+                                       to which to download the specified artifact. The downloaded
+                                       artifact may be written to a subdirectory of
+                                       `dst_local_dir_path` if `src_artifact_path` contains
+                                       subdirectories.
+            :return: A local filesystem path referring to the downloaded file.
+            """
+            local_destination_file_path = self._create_download_destination(
+                src_artifact_path=src_artifact_path, dst_local_dir_path=dst_local_dir_path
+            )
+            self._download_file(
+                remote_file_path=src_artifact_path, local_path=local_destination_file_path
+            )
+            return local_destination_file_path
+
+        def download_artifact_dir(src_artifact_dir_path, dst_local_dir_path):
+            local_dir = os.path.join(dst_local_dir_path, src_artifact_dir_path)
+            dir_content = [  # prevent infinite loop, sometimes the dir is recursively included
+                file_info
+                for file_info in self.list_artifacts(src_artifact_dir_path)
+                if file_info.path != "." and file_info.path != src_artifact_dir_path
+            ]
+            if not dir_content:  # empty dir
+                if not os.path.exists(local_dir):
+                    os.makedirs(local_dir, exist_ok=True)
+            else:
+                for file_info in dir_content:
+                    if file_info.is_dir:
+                        download_artifact_dir(
+                            src_artifact_dir_path=file_info.path,
+                            dst_local_dir_path=dst_local_dir_path,
+                        )
+                    else:
+                        download_artifact(
+                            src_artifact_path=file_info.path, dst_local_dir_path=dst_local_dir_path
+                        )
+            return local_dir
+
+        if dst_path is None:
+            dst_path = tempfile.mkdtemp()
+        dst_path = os.path.abspath(dst_path)
+
+        if not os.path.exists(dst_path):
+            raise MlflowException(
+                message=(
+                    "The destination path for downloaded artifacts does not"
+                    " exist! Destination path: {dst_path}".format(dst_path=dst_path)
+                ),
+                error_code=RESOURCE_DOES_NOT_EXIST,
+            )
+        elif not os.path.isdir(dst_path):
+            raise MlflowException(
+                message=(
+                    "The destination path for downloaded artifacts must be a directory!"
+                    " Destination path: {dst_path}".format(dst_path=dst_path)
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        # Check if the artifacts points to a directory
+        if self._is_directory(artifact_path):
+            return download_artifact_dir(
+                src_artifact_dir_path=artifact_path, dst_local_dir_path=dst_path
+            )
+        else:
+            return download_artifact(src_artifact_path=artifact_path, dst_local_dir_path=dst_path)
+
+    def _download_file(self, remote_file_path, local_path):
+        """
+        Download the file at the specified relative remote path and saves
+        it at the specified local path.
+
+        :param remote_file_path: Source path to the remote file, relative to the root
+                                 directory of the artifact repository.
+        :param local_path: The path to which to save the downloaded file.
+        """
+        with self.managed_folder.get_file(remote_file_path) as remote_file:
+            with open(local_path, "wb") as local_file:
+                for line in remote_file:
+                    local_file.write(line)
+
+    def delete_artifacts(self, artifact_path=None):
+        """
+        Delete the artifacts at the specified location.
+        Supports the deletion of a single file or of a directory. Deletion of a directory
+        is recursive.
+        :param artifact_path: Path of the artifact to delete
+        """
+        path = (
+            os.path.join(self.base_artifact_path, artifact_path) if artifact_path else self.base_artifact_path
+        )
+        self.managed_folder.delete_file(path)
+
+
+def verify_artifact_path(artifact_path):
+    from mlflow.exceptions import MlflowException
+    from mlflow.utils.validation import path_not_unique, bad_path_message
+    if artifact_path and path_not_unique(artifact_path):
+        raise MlflowException(
+            "Invalid artifact path: '%s'. %s" % (artifact_path, bad_path_message(artifact_path))
+        )
diff --git a/dataikuapi/dss_plugin_mlflow/header_provider.py b/dataikuapi/dss_plugin_mlflow/header_provider.py
@@ -0,0 +1,15 @@
+import os
+
+class PluginDSSHeaderProvider:
+
+    def in_context(self):
+        env_variables = {"DSS_MLFLOW_HEADER", "DSS_MLFLOW_TOKEN", "DSS_MLFLOW_PROJECTKEY"}
+        if len(env_variables.difference(set(os.environ))) == 0:
+            return True
+
+    def request_headers(self):
+        headers = {
+            os.environ.get("DSS_MLFLOW_HEADER"): os.environ.get("DSS_MLFLOW_TOKEN"),
+            "x-dku-mlflow-project-key": os.environ.get("DSS_MLFLOW_PROJECTKEY"),
+        }
+        return headers
diff --git a/dataikuapi/dss_plugin_mlflow/utils.py b/dataikuapi/dss_plugin_mlflow/utils.py
@@ -0,0 +1,26 @@
+import os
+import sys
+import tempfile
+
+
+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
diff --git a/dataikuapi/dssclient.py b/dataikuapi/dssclient.py
@@ -1,10 +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 load_dss_mlflow_plugin
 from .dss.future import DSSFuture
 from .dss.projectfolder import DSSProjectFolder
 from .dss.project import DSSProject
@@ -17,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):
@@ -1087,6 +1092,52 @@ def get_object_discussions(self, project_key, object_type, object_id):
         """
         return DSSObjectDiscussions(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
+
+        :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
+        """
+        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):
     def __init__(self, client, import_id):
         self.client = client

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+from .utils import load_dss_mlflow_plugin`