encapsulate requests in try except

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}")