Do not create managed folder automatically for mlflow

Author: alarive

dataikuapi/dss/project.py

@@ -1625,14 +1625,15 @@ def get_app_manifest(self):
 
     # MLflow experiment tracking
     ########################################################
-    def setup_mlflow(self, managed_folder_name="mlflow_artifacts", host=None):
+    def setup_mlflow(self, managed_folder, host=None):
         """
         Setup the dss-plugin for MLflow
 
-        :param str managed_folder_name: name of the managed folder where artifacts should be stored
+        :param managed_folder: Managed folder where artifacts should be stored.
+                               Either a managed folder id, or an instance of :class:`dataikuapi.dss.DSSManagedFolder`.
         :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_name=managed_folder_name, host=host)
+        return MLflowHandle(client=self.client, project=self, managed_folder=managed_folder, host=host)
 
     def get_mlflow_extension(self):
         """

dataikuapi/dss_plugin_mlflow/utils.py

@@ -1,3 +1,6 @@
+from ..dss.utils import AnyLoc
+from ..utils import DataikuException
+import logging
 import os
 import sys
 import tempfile
@@ -6,7 +9,7 @@
 
 
 class MLflowHandle:
-    def __init__(self, client, project_key, managed_folder_name, host=None):
+    def __init__(self, client, project, managed_folder, host=None):
         """ Add the MLflow-plugin parts of dataikuapi to MLflow local setup.
 
         This method deals with
@@ -19,7 +22,8 @@ def __init__(self, client, project_key, managed_folder_name, host=None):
         with DSS backend as the tracking backend and enable using DSS managed folder
         as artifact location.
         """
-        self.project_key = project_key
+        self.project = project
+        self.project_key = project.project_key
         self.client = client
         self.mlflow_env = {}
 
@@ -57,32 +61,28 @@ def __init__(self, client, project_key, managed_folder_name, host=None):
                 "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
-        })
 
-        # 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.
+        def full_id(mf):
+            return mf.project.project_key + "." + mf.id
+
+        def get_managed_folder(smart_mf_id):
+            return client.get_project(smart_mf_id.project_key).get_managed_folder(smart_mf_id.object_id)
+
+        mf_id = managed_folder if isinstance(managed_folder, str) else full_id(managed_folder)
+        smart_mf_id = AnyLoc.from_ref(self.project_key, mf_id)
+        try:
+            get_managed_folder(smart_mf_id).get_definition()
+        except DataikuException as e:
+            if "NotFoundException" in str(e):
+                logging.error('The managed folder "%s" does not exist, please create it in your project flow before running this command.' % (mf_id))
+            raise
+
+        # Set host, tracking URI, project key and managed_folder_id
         self.mlflow_env.update({
-            "DSS_MLFLOW_MANAGED_FOLDER_ID": managed_folder.id
+            "DSS_MLFLOW_PROJECTKEY": self.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_ID": mf_id
         })
 
         os.environ.update(self.mlflow_env)