Store the managed folder id in the MLflow artifact URI

Author: alarive

dataikuapi/dss/project.py

@@ -1608,14 +1608,14 @@ def get_app_manifest(self):
 
     # MLflow experiment tracking
     ########################################################
-    def setup_mlflow(self, managed_folder="mlflow_artifacts", host=None):
+    def setup_mlflow(self, managed_folder_name="mlflow_artifacts", host=None):
         """
         Setup the dss-plugin for MLflow
 
-        :param str managed_folder: managed folder where artifacts are stored
+        :param str managed_folder_name: name of the managed folder where artifacts should be 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)
+        return MLflowHandle(client=self.client, project_key=self.project_key, managed_folder_name=managed_folder_name, host=host)
 
     def get_mlflow_extension(self):
         """

dataikuapi/dss_plugin_mlflow/artifact_repository.py

@@ -6,17 +6,16 @@
 
 
 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)
-
+    if not parsed.netloc or parsed.netloc == '.':
+        raise Exception("Could not find a managed folder id in URI: %s" % uri)
+    return parsed
 
 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"),
@@ -28,14 +27,19 @@ def __init__(self, artifact_uri):
                 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])
+        parsed_uri = parse_dss_managed_folder_uri(artifact_uri)
+        self.managed_folder = self.__get_managed_folder(parsed_uri.netloc)
+        self.base_artifact_path = os.path.normpath(parsed_uri.path)
+
+    def __get_managed_folder(self, managed_folder_smart_id):
+        chunks = managed_folder_smart_id.split('.')
+        if len(chunks) == 1:
+            return self.project.get_managed_folder(chunks[0])
+        elif len(chunks) == 2:
+            project = self.client.get_project(chunks[0])
+            return project.get_managed_folder(chunks[1])
         else:
-            self.managed_folder = self.project.create_managed_folder(os.environ.get("DSS_MLFLOW_MANAGED_FOLDER"))
+            raise Exception("Invalid managed folder id: %s" % managed_folder_smart_id)
 
     def log_artifact(self, local_file, artifact_path=None):
         """

dataikuapi/dss_plugin_mlflow/header_provider.py

@@ -11,5 +11,6 @@ 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"),
+            "x-dku-mlflow-managed-folder-id": os.environ.get("DSS_MLFLOW_MANAGED_FOLDER_ID"),
         }
         return headers

dataikuapi/dss_plugin_mlflow/utils.py

@@ -6,7 +6,7 @@
 
 
 class MLflowHandle:
-    def __init__(self, client, project_key, managed_folder="mlflow_artifacts", host=None):
+    def __init__(self, client, project_key, managed_folder_name, host=None):
         """ Add the MLflow-plugin parts of dataikuapi to MLflow local setup.
 
         This method deals with
@@ -61,9 +61,30 @@ def __init__(self, client, project_key, managed_folder="mlflow_artifacts", host=
         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
+            "DSS_MLFLOW_HOST": self.client.host
         })
+
+        # Get or create managed folder with name 'managed_folder_name'
+        project = self.client.get_project(project_key)
+        managed_folders = [
+            x["id"] for x in project.list_managed_folders()
+            if x["name"] == managed_folder_name
+        ]
+        managed_folder = None
+        if len(managed_folders) > 0:
+            managed_folder = project.get_managed_folder(managed_folders[0])
+        else:
+            managed_folder = project.create_managed_folder(managed_folder_name)
+
+
+        # TODO: don't allow to pass a managed folder name, everything is already
+        # wired in artifact_repository and in the backend for DSS_MLFLOW_MANAGED_FOLDER_ID to contain
+        # a smart ref. So, instead of managed_folder_name, we should take a DSSManagedFolder
+        # or a managed folder id/smart ref and set DSS_MLFLOW_MANAGED_FOLDER_ID to that (smart) id.
+        self.mlflow_env.update({
+            "DSS_MLFLOW_MANAGED_FOLDER_ID": managed_folder.id
+        })
+
         os.environ.update(self.mlflow_env)
 
     def clear(self):