diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 0ce5a9b607..69eac61af8 100644 --- a/monai/bundle/scripts.py +++ b/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, }, }, }, diff --git a/monai/bundle/utils.py b/monai/bundle/utils.py index 4b1adeffea..8ffcd11208 100644 --- a/monai/bundle/utils.py +++ b/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, }, }, } diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py index a06f2013aa..961b91f087 100644 --- a/monai/handlers/mlflow_handler.py +++ b/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: """ diff --git a/tests/test_fl_monai_algo.py b/tests/test_fl_monai_algo.py index 6b234fbe96..426bb44cfc 100644 --- a/tests/test_fl_monai_algo.py +++ b/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) diff --git a/tests/test_handler_mlflow.py b/tests/test_handler_mlflow.py index de164cadb8..f79707384f 100644 --- a/tests/test_handler_mlflow.py +++ b/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)