Force managed_folder to be a DSSManagedFolder in setup_mlflow

Author: alarive

dataikuapi/dss/project.py

@@ -1,4 +1,4 @@
-import time, warnings, sys, os.path as osp
+import warnings, os.path as osp
 
 from ..dss_plugin_mlflow import MLflowHandle
 
@@ -1629,8 +1629,7 @@ def setup_mlflow(self, managed_folder, host=None):
         """
         Setup the dss-plugin for MLflow
 
-        :param managed_folder: Managed folder where artifacts should be stored.
-                               Either a managed folder id, or an instance of :class:`dataikuapi.dss.DSSManagedFolder`.
+        :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
         """
         return MLflowHandle(client=self.client, project=self, managed_folder=managed_folder, host=host)

dataikuapi/dss_plugin_mlflow/utils.py

@@ -1,5 +1,5 @@
-from ..dss.utils import AnyLoc
 from ..utils import DataikuException
+from ..dss.managedfolder import DSSManagedFolder
 import logging
 import os
 import sys
@@ -62,16 +62,13 @@ def __init__(self, client, project, managed_folder, host=None):
                 "DSS_MLFLOW_INTERNAL_TICKET": self.client.internal_ticket
             })
 
-        def full_id(mf):
-            return mf.project.project_key + "." + mf.id
+        if not isinstance(managed_folder, DSSManagedFolder):
+            raise TypeError('managed_folder must a DSSManagedFolder.')
 
-        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)
+        mf_project = managed_folder.project.project_key
+        mf_id = managed_folder.id
         try:
-            get_managed_folder(smart_mf_id).get_definition()
+            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))
@@ -82,7 +79,7 @@ def get_managed_folder(smart_mf_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
+            "DSS_MLFLOW_MANAGED_FOLDER_ID": mf_project + "." + mf_id
         })
 
         os.environ.update(self.mlflow_env)