Robust mlflow requests

Requires !57 (merged)

Closes #28 (closed)

dan/ocr/document/train.py

@@ -24,7 +24,7 @@ try:
 
     MLFLOW = True
     logger.info("MLflow Logging available.")
-    from dan.mlflow import start_mlflow_run
+    from dan.mlflow import make_mlflow_request, start_mlflow_run
 except ImportError:
     MLFLOW = False
 
@@ -71,16 +71,16 @@ def get_config():
     Retrieve model configuration
     """
     dataset_name = "esposalles"
-    dataset_level = "page"
-    dataset_variant = ""
-    dataset_path = "/home/training_data/ATR_paragraph/Esposalles"
+    dataset_level = "record"
+    dataset_variant = "_debug"
+    dataset_path = "."
     params = {
         "mlflow": {
             "dataset_name": dataset_name,
             "run_name": "Test log DAN",
             "s3_endpoint_url": "",
             "tracking_uri": "",
-            "experiment_id": "9",
+            "experiment_id": "0",
             "aws_access_key_id": "",
             "aws_secret_access_key": "",
         },
@@ -103,6 +103,11 @@ def get_config():
                     (dataset_name, "val"),
                 ],
             },
+            "test": {
+                "{}-test".format(dataset_name): [
+                    (dataset_name, "test"),
+                ],
+            },
             "config": {
                 "load_in_memory": True,  # Load all images in CPU memory
                 "worker_per_gpu": 4,  # Num of parallel processes per gpu for data loading
@@ -287,18 +292,27 @@ def run():
             / "labels.json"
         )
         with start_mlflow_run(config["mlflow"]) as run:
-            logger.info(f"Set tags to MLflow on {config['mlflow']['run_name']}")
-            mlflow.set_tags({"Dataset": config["mlflow"]["dataset_name"]})
+            logger.info(f"Started MLflow run with ID ({run.info.run_id})")
+
+            make_mlflow_request(
+                mlflow_method=mlflow.set_tags, tags={"Dataset": dataset_name}
+            )
 
             # Get the labels json file
             with open(labels_path) as json_file:
                 labels_artifact = json.load(json_file)
 
             # Log MLflow artifacts
-            mlflow.log_dict(config_artifact, "config.json")
-            mlflow.log_dict(labels_artifact, "labels.json")
+            for artifact, filename in [
+                (config_artifact, "config.json"),
+                (labels_artifact, "labels.json"),
+            ]:
+                make_mlflow_request(
+                    mlflow_method=mlflow.log_dict,
+                    dictionary=artifact,
+                    artifact_file=filename,
+                )
 
-            logger.info(f"Started MLflow run with ID ({run.info.run_id})")
             if (
                 config["training_params"]["use_ddp"]
                 and not config["training_params"]["force_cpu"]