MLFLowHandler - How to close the trainer/evaluator pipeline when completed?

[whikwon]

From the tutorial code,
I can attach a MLFlowHandlers to each trainer and evaluator like the code below.

train_mlflow_handler = MLFlowHandler(tracking_uri=Path(mlflow_dir).as_uri(), output_transform=lambda x: x)
train_mlflow_handler.attach(trainer)
val_mlflow_handler = MLFlowHandler(
    tracking_uri=Path(mlflow_dir).as_uri(),
    # no need to plot loss value, so disable per iteration output
    output_transform=lambda x: None,
    # fetch global epoch number from trainer
    global_epoch_transform=lambda x: trainer.state.epoch,
)
val_mlflow_handler.attach(evaluator)

Is there a way to stop MLflow after the trainer and evaluator have finished? Adding argument close_on_complete=True in train and evaluator causes a new run to be created each time the trainer or evaluator is closed, which is problematic.

Thanks in advance.

[KumoLiu]

Hi @whikwon, did you sett close_on_complete=True within the train_mlflow_handler instead of within val_mlflow_handler? The reason being that if it's not set within train_mlflow_handler, a new run for training might inadvertently be started after evaluation considering the training is still not finished.

Thanks.

[whikwon]

Hi @KumoLiu,
Yes, I tried setting close_on_complete=True for the trainer, but when the trainer finishes and the evaluator performs the final validation, it creates a new run at that time.

[KumoLiu]

Hi @whikwon,
Please use SupervisedTrainer and ValidationHandler instead create_supervised_trainer which can consider validation as a part of training.
An example here:
https://github.com/Project-MONAI/tutorials/blob/570b19e678e06aba26f53df8bc73f848ae2984ba/pathology/hovernet/training.py#L310

[whikwon]

Oh, I got it. Thank you so much.