Skip to content

Commit 24fe5bd

Browse files
committed
Add list_experiments to MLflow extensions. Document a bit list_models.
1 parent d51ba05 commit 24fe5bd

File tree

1 file changed

+19
-0
lines changed

1 file changed

+19
-0
lines changed

dataikuapi/dss/mlflow.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,27 @@ def __init__(self, client, project_key):
1111
self.project_key = project_key
1212

1313
def list_models(self, run_id):
14+
"""
15+
Returns the list of models of given run
16+
17+
:param run_id: run_id for which to return a list of models
18+
:type run_id: str
19+
"""
1420
response = self.client._perform_http(
1521
"GET", "/api/2.0/mlflow/extension/models/{}".format(run_id),
1622
headers={"x-dku-mlflow-project-key": self.project_key}
1723
)
1824
return response.json()
25+
26+
def list_experiments(self):
27+
"""
28+
Returns the list of experiments in the DSS project for which MLflow integration
29+
is setup
30+
31+
:rtype: dict
32+
"""
33+
response = self.client._perform_http(
34+
"GET", "/api/2.0/mlflow/experiments/list",
35+
headers={"x-dku-mlflow-project-key": self.project_key}
36+
)
37+
return response.json()

0 commit comments

Comments
 (0)