more logs

dan/manager/training.py

@@ -31,13 +31,6 @@ try:
 except ImportError:
     pass
 
-try:
-    import mlflow
-
-    from dan.mlflow import logging_metrics, logging_tags_metrics
-except ImportError:
-    pass
-
 
 class GenericTrainingManager:
     def __init__(self, params):

dan/mlflow.py

@@ -91,7 +91,9 @@ def logging_tags_metrics(
 @contextmanager
 def start_mlflow_run(config: dict):
     """
-    Create an MLflow execution context with the parameters contained in the config file
+    Create an MLflow execution context with the parameters contained in the config file.
+
+    Yields the active MLflow run, as well as a boolean saying whether a new one was created.
 
     :param config: dict, the config of the model
     """
@@ -116,5 +118,5 @@ def start_mlflow_run(config: dict):
     # Start run
     yield mlflow.start_run(
         run_id=run_id, run_name=run_name, experiment_id=experiment_id
-    )
+    ), run_id is None
     mlflow.end_run()

dan/ocr/document/train.py

@@ -304,11 +304,14 @@ def run():
         labels_path = (
             Path(config["dataset_params"]["datasets"][dataset_name]) / "labels.json"
         )
-        with start_mlflow_run(config["mlflow"]) as run:
-            logger.info(f"Started MLflow run with ID ({run.info.run_id})")            
+        with start_mlflow_run(config["mlflow"]) as (run, created):
+            if created:
+                logger.info(f"Started MLflow run with ID ({run.info.run_id})")
+            else:
+                logger.info(f"Resumed MLflow run with ID ({run.info.run_id})")
+
             make_mlflow_request(
-                mlflow_method=mlflow.set_tags,
-                tags={"Dataset": dataset_name}
+                mlflow_method=mlflow.set_tags, tags={"Dataset": dataset_name}
             )
             # Get the labels json file
             with open(labels_path) as json_file: