From daef6f90168db4cb2e2f59b4885e9bd9274e7cb6 Mon Sep 17 00:00:00 2001
From: fabiovincenzi <93596376+fabiovincenzi@users.noreply.github.com>
Date: Fri, 5 Apr 2024 10:54:59 +0200
Subject: [PATCH] [release/0.5] Update creds import in python client (#1084)

---
 python/fasttrackml/_tracking_service/client.py | 9 +++++++--
 python/pyproject.toml                          | 3 +++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/python/fasttrackml/_tracking_service/client.py b/python/fasttrackml/_tracking_service/client.py
index d94a70dc8..ea0ad04f8 100644
--- a/python/fasttrackml/_tracking_service/client.py
+++ b/python/fasttrackml/_tracking_service/client.py
@@ -3,12 +3,17 @@
 
 from mlflow.entities import Param, RunTag
 from mlflow.store.tracking import GET_METRIC_HISTORY_MAX_RESULTS
-from mlflow.tracking._tracking_service import utils
 from mlflow.tracking._tracking_service.client import TrackingServiceClient
 from mlflow.tracking.metric_value_conversion_utils import (
     convert_metric_value_to_float_if_possible,
 )
 from mlflow.utils import chunk_list
+
+try:
+    from mlflow.utils.credentials import get_default_host_creds
+except ImportError:
+    from mlflow.tracking._tracking_service.utils import _get_default_host_creds as get_default_host_creds
+
 from mlflow.utils.time import get_current_time_millis
 from mlflow.utils.validation import MAX_METRICS_PER_BATCH
 
@@ -20,7 +25,7 @@ class FasttrackmlTrackingServiceClient(TrackingServiceClient):
 
     def __init__(self, tracking_uri):
         super().__init__(tracking_uri)
-        self.custom_store = CustomRestStore(partial(utils._get_default_host_creds, self.tracking_uri))
+        self.custom_store = CustomRestStore(partial(get_default_host_creds, self.tracking_uri))
 
     def log_metric(
         self,
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 16d8d410b..5e5efe011 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -46,6 +46,9 @@ classifiers = [
     "Topic :: Scientific/Engineering",
     "Topic :: Utilities",
 ]
+dependencies = [
+  "mlflow >= 2.8.0, < 3.0",
+]
 
 [project.urls]
 Homepage = "https://fasttrackml.io/"