diff --git a/dan/mlflow.py b/dan/mlflow.py index 063c2b91d96e81f62914b0d14fa14596473b481c..c1d6cd3f9d8cfd5e7b05deb6a967ca3a0a472d3c 100644 --- a/dan/mlflow.py +++ b/dan/mlflow.py @@ -3,11 +3,19 @@ import os from contextlib import contextmanager import mlflow +import requests from mlflow.exceptions import MlflowException from dan import logger +def make_mlflow_request(mlflow_method, *args, **kwargs): + try: + mlflow_method(*args, **kwargs) + except requests.exceptions.ConnectionError as e: + logger.error(f"Call to `{str(mlflow_method)}` failed with error: {str(e)}") + + def setup_environment(config: dict): """ Get the necessary variables from the config file and put them in the environment variables @@ -45,7 +53,9 @@ def logging_metrics( mlflow_values = { f"{step}_{name}": value for name, value in display_values.items() } - mlflow.log_metrics(mlflow_values, epoch) + make_mlflow_request( + mlflow_method=mlflow.log_metrics, metrics=mlflow_values, step=epoch + ) def logging_tags_metrics( @@ -66,7 +76,7 @@ def logging_tags_metrics( mlflow_values = { f"{step}_{name}": value for name, value in display_values.items() } - mlflow.set_tags(mlflow_values) + make_mlflow_request(mlflow_method=mlflow.set_tags, tags=mlflow_values) @contextmanager @@ -84,7 +94,9 @@ def start_mlflow_run(config: dict): experiment_id = config.get("experiment_id") assert experiment_id, "Missing MLflow experiment ID in the configuration" try: - mlflow.set_experiment(experiment_id=experiment_id) + make_mlflow_request( + mlflow_method=mlflow.set_experiment, experiment_id=experiment_id + ) logger.info(f"Run Experiment ID : {experiment_id} on MLFlow") except MlflowException as e: logger.error(f"Couldn't set Mlflow experiment with ID: {experiment_id}")