diff --git a/dan/mlflow.py b/dan/mlflow.py index 366a119af425f233dc58c42c286e81f6cd5fd856..180406f33a0dec286f292847338bd022a2ae4b8f 100644 --- a/dan/mlflow.py +++ b/dan/mlflow.py @@ -3,9 +3,8 @@ import os from contextlib import contextmanager import mlflow -from mlflow.environment_variables import MLFLOW_HTTP_REQUEST_MAX_RETRIES - import requests +from mlflow.environment_variables import MLFLOW_HTTP_REQUEST_MAX_RETRIES from dan import logger @@ -20,6 +19,7 @@ def make_mlflow_request(mlflow_method, *args, **kwargs): logger.error(f"Call to `{str(mlflow_method)}` failed with error: {str(e)}") raise e + def setup_environment(config: dict): """ Get the necessary variables from the config file and put them in the environment variables @@ -39,7 +39,9 @@ def setup_environment(config: dict): # Check max retry setting max_retries = MLFLOW_HTTP_REQUEST_MAX_RETRIES.get() if max_retries and int(max_retries) <= 1: - logger.warning(f"The maximum number of retries for MLflow HTTP requests is set to {max_retries}, which is low. Consider using a higher value.") + logger.warning( + f"The maximum number of retries for MLflow HTTP requests is set to {max_retries}, which is low. Consider using a higher value." + ) def logging_metrics(