Support passing the managed folder as an id or a Folder in setup_mlflow

Author: alarive

dataikuapi/dss/project.py

@@ -1612,8 +1612,10 @@ def setup_mlflow(self, managed_folder, host=None):
         """
         Setup the dss-plugin for MLflow
 
-        :param object managed_folder: a :class:`dataikuapi.dss.DSSManagedFolder` where MLflow artifacts should be stored.
-        :param str host: setup a custom host if the backend used is not DSS
+        :param object managed_folder: the managed folder where MLflow artifacts should be stored.
+                                      Can be either a managed folder id as a string,
+                                      a :class:`dataikuapi.dss.DSSManagedFolder`, or a :class:`dataiku.Folder`
+        :param str host: setup a custom host if the backend used is not DSS.
         """
         return MLflowHandle(client=self.client, project=self, managed_folder=managed_folder, host=host)
 

dataikuapi/dss_plugin_mlflow/utils.py

@@ -62,30 +62,39 @@ def __init__(self, client, project, managed_folder, host=None):
                 "DSS_MLFLOW_INTERNAL_TICKET": self.client.internal_ticket
             })
 
-        if not isinstance(managed_folder, DSSManagedFolder):
-            raise TypeError('managed_folder must a DSSManagedFolder.')
-
         if not client._session.verify:
             self.mlflow_env.update({"MLFLOW_TRACKING_INSECURE_TLS": "true"})
         elif isinstance(client._session.verify, str):
             self.mlflow_env.update({"MLFLOW_TRACKING_SERVER_CERT_PATH": client._session.verify})
 
+        mf_full_id = managed_folder
+        if isinstance(managed_folder, DSSManagedFolder):
+            mf_full_id = managed_folder.project.project_key + "." + managed_folder.id
+        elif hasattr(managed_folder, 'name'): # True if dataiku.Folder
+            mf_full_id = managed_folder.name
+
+        if not isinstance(mf_full_id, str):
+            raise TypeError('Type of managed_folder must be "str", "DSSManagedFolder" or "dataiku.Folder".')
+
+        if not "." in mf_full_id:
+            mf_full_id = self.project_key + "." + mf_full_id
+
+        mf_project = mf_full_id.split(".")[0]
+        mf_id = mf_full_id.split(".")[1]
 
-        mf_project = managed_folder.project.project_key
-        mf_id = managed_folder.id
         try:
             client.get_project(mf_project).get_managed_folder(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))
+                logging.error('The managed folder "%s" does not exist, please create it in your project flow before running this command.' % (mf_full_id))
             raise
 
         # Set host, tracking URI, project key and managed_folder_id
         self.mlflow_env.update({
             "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_project + "." + mf_id
+            "DSS_MLFLOW_MANAGED_FOLDER_ID": mf_full_id
         })
 
         os.environ.update(self.mlflow_env)