Create dedicated class for MLflow custom endpoints (#203)

* Create dedicated class for MLflow custom endpoint

* Remove unused parameter

Author: Valentin Thorey

dataikuapi/dss/mlflow.py

@@ -0,0 +1,18 @@
+class DSSMLflowExtension(object):
+    """
+    A handle to interact with specific endpoints of the DSS MLflow integration.
+
+    Do not create this directly, use :meth:`dataikuapi.dss.DSSProject.get_mlflow_extension`
+    """
+
+    def __init__(self, client, project_key):
+        self.client = client
+        self.project = client.get_project(project_key)
+        self.project_key = project_key
+
+    def list_models(self, run_id):
+        response = self.client._perform_http(
+            "GET", "/api/2.0/mlflow/extension/models/{}".format(run_id),
+            headers={"x-dku-mlflow-project-key": self.project_key}
+        )
+        return response.json()

dataikuapi/dss/project.py

@@ -12,6 +12,7 @@
 from .managedfolder import DSSManagedFolder
 from .savedmodel import DSSSavedModel
 from .modelevaluationstore import DSSModelEvaluationStore
+from .mlflow import DSSMLflowExtension
 from .job import DSSJob, DSSJobWaiter
 from .scenario import DSSScenario, DSSScenarioListItem
 from .continuousactivity import DSSContinuousActivity
@@ -33,8 +34,8 @@ class DSSProject(object):
     Do not create this class directly, instead use :meth:`dataikuapi.DSSClient.get_project`
     """
     def __init__(self, client, project_key):
-       self.client = client
-       self.project_key = project_key
+        self.client = client
+        self.project_key = project_key
 
     def get_summary(self):
         """
@@ -1608,6 +1609,15 @@ def setup_mlflow(self, managed_folder="mlflow_artifacts", host=None):
         """
         return MLflowHandle(client=self.client, project_key=self.project_key, managed_folder=managed_folder, host=host)
 
+    def get_mlflow_extension(self):
+        """
+        Get a handle to interact with the extension of MLflow provided by DSS
+
+        :returns: A :class:`dataikuapi.dss.mlflow.DSSMLflowExtension` Mlflow Extension handle
+
+        """
+        return DSSMLflowExtension(client=self.client, project_key=self.project_key)
+
     def clean_experiment_tracking_db(self):
         """
         Cleans the experiments, runs, params, metrics, tags, etc. for this project

dataikuapi/dssclient.py

@@ -1028,7 +1028,7 @@ def set_license(self, license):
     # Internal Request handling
     ########################################################
 
-    def _perform_http(self, method, path, params=None, body=None, stream=False, files=None, raw_body=None):
+    def _perform_http(self, method, path, params=None, body=None, stream=False, files=None, raw_body=None, headers=None):
         if body is not None:
             body = json.dumps(body)
         if raw_body is not None:
@@ -1038,8 +1038,9 @@ def _perform_http(self, method, path, params=None, body=None, stream=False, file
             http_res = self._session.request(
                     method, "%s/dip/publicapi%s" % (self.host, path),
                     params=params, data=body,
-                    files = files,
-                    stream = stream)
+                    files=files,
+                    stream=stream,
+                    headers=headers)
             http_res.raise_for_status()
             return http_res
         except exceptions.HTTPError: