From 3dca228f9df6949a9d6fa44d522eb06be47f80f4 Mon Sep 17 00:00:00 2001
From: Yoann Schneider <yschneider@teklia.com>
Date: Fri, 10 Feb 2023 12:15:01 +0100
Subject: [PATCH] encapsulate requests in try except
---
dan/mlflow.py | 18 +++++++++++++++---
1 file changed, 15 insertions(+), 3 deletions(-)
diff --git a/dan/mlflow.py b/dan/mlflow.py
index 063c2b91..c1d6cd3f 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}")
--
GitLab