5649 auto close mlflow run

monai/bundle/scripts.py

@@ -561,6 +561,7 @@ def run(
                             "run_name": "@run_name",
                             "iteration_log": True,
                             "output_transform": "$monai.handlers.from_engine(['loss'], first=True)",
+                            "close_on_complete": True,
                         },
                         "validator": {
                             "_target_": "MLFlowHandler",
@@ -577,6 +578,7 @@ def run(
                             "experiment_name": "@experiment_name",
                             "run_name": "@run_name",
                             "iteration_log": False,
+                            "close_on_complete": True,
                         },
                     },
                 },

monai/bundle/utils.py

@@ -122,6 +122,7 @@
             "epoch_log": True,
             "tag_name": "train_loss",
             "output_transform": "$monai.handlers.from_engine(['loss'], first=True)",
+            "close_on_complete": True,
         },
         # MLFlowHandler config for the validator
         "validator": {
@@ -140,6 +141,7 @@
             "experiment_name": "@experiment_name",
             "run_name": "@run_name",
             "iteration_log": False,
+            "close_on_complete": True,
         },
     },
 }

monai/handlers/mlflow_handler.py

@@ -82,6 +82,7 @@ class MLFlowHandler:
         artifacts: paths to images that need to be recorded after a whole run.
         optimizer_param_names: parameters' name in optimizer that need to be record during runing,
             defaults to "lr".
+        close_on_complete: whether to close the mlflow run in `complete` phase in workflow, default to False.
 
     For more details of MLFlow usage, please refer to: https://mlflow.org/docs/latest/index.html.
 
@@ -101,11 +102,12 @@ def __init__(
         global_epoch_transform: Callable = lambda x: x,
         state_attributes: Optional[Sequence[str]] = None,
         tag_name: str = DEFAULT_TAG,
-        experiment_name: str = "default_experiment",
+        experiment_name: str = "monai_experiment",
         run_name: Optional[str] = None,
         experiment_param: Optional[Dict] = None,
         artifacts: Optional[Union[str, Sequence[Path]]] = None,
         optimizer_param_names: Union[str, Sequence[str]] = "lr",
+        close_on_complete: bool = False,
     ) -> None:
         if tracking_uri is not None:
             mlflow.set_tracking_uri(tracking_uri)
@@ -124,6 +126,7 @@ def __init__(
         self.artifacts = ensure_tuple(artifacts)
         self.optimizer_param_names = ensure_tuple(optimizer_param_names)
         self.client = mlflow.MlflowClient()
+        self.close_on_complete = close_on_complete
 
     def _delete_exist_param_in_dict(self, param_dict: Dict) -> None:
         """
@@ -156,6 +159,8 @@ def attach(self, engine: Engine) -> None:
             engine.add_event_handler(Events.EPOCH_COMPLETED, self.epoch_completed)
         if not engine.has_event_handler(self.complete, Events.COMPLETED):
             engine.add_event_handler(Events.COMPLETED, self.complete)
+        if self.close_on_complete and (not engine.has_event_handler(self.close, Events.COMPLETED)):
+            engine.add_event_handler(Events.COMPLETED, self.close)
 
     def start(self, engine: Engine) -> None:
         """

tests/test_fl_monai_algo.py

@@ -155,6 +155,7 @@ def test_train(self, input_params):
                     "_target_": "MLFlowHandler",
                     "tracking_uri": Path(data_dir).as_uri() + "/mlflow_override",
                     "output_transform": "$monai.handlers.from_engine(['loss'], first=True)",
+                    "close_on_complete": True,
                 }
             },
         }
@@ -175,8 +176,6 @@ def test_train(self, input_params):
         # test train
         algo.train(data=data, extra={})
         algo.finalize()
-        # must close it as we are changing different temp dir for cases here
-        algo.train_parser.get_parsed_content("train#handlers")[-1].close()
         self.assertTrue(os.path.exists(f"{data_dir}/mlflow_override"))
         shutil.rmtree(data_dir)
 

tests/test_handler_mlflow.py

@@ -53,6 +53,7 @@ def _update_metric(engine):
                 state_attributes=["test"],
                 experiment_param=experiment_param,
                 artifacts=[artifact_path],
+                close_on_complete=True,
             )
             handler.attach(engine)
             engine.run(range(3), max_epochs=2)