Repository: Project-MONAI/MONAI
File: monai/handlers/mlflow_handler.py (Python)

Code before:
            `engine.state` and `output_transform` inherit from the ignite concept:
            https://pytorch-ignite.ai/concepts/03-state/, explanation and usage example are in the tutorial:
            https://github.com/Project-MONAI/tutorials/blob/master/modules/batch_output_transform.ipynb.
        global_epoch_transform: a callable that is used to customize global epoch number.
            For example, in evaluation, the evaluator engine might want to track synced epoch number
            with the trainer engine.
        state_attributes: expected attributes from `engine.state`, if provided, will extract them
            when epoch completed.
        tag_name: when iteration output is a scalar, `tag_name` is used to track, defaults to `'Loss'`.
        experiment_name: a name for an experiment, defaults to `monai_experiment`.
        run_name: a name for a run in an experiment.
        experiment_param: a dict recording parameters which will not change through the whole workflow,
            like torch version, cuda version and so on.
        artifacts: paths to images that need to be recorded after running the workflow.
        optimizer_param_names: parameter names in the optimizer that need to be recorded during running,
            default 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.

    """

    # parameters that are logged at the start of training
    default_tracking_params = ["max_epochs", "epoch_length"]
    # the finish status of a mlflow run
    run_finish_status = mlflow.entities.RunStatus.to_string(mlflow.entities.RunStatus.FINISHED)

    def __init__(
        self,
        tracking_uri: str | None = None,
        iteration_log: bool | Callable[[Engine, int], bool] = True,
        epoch_log: bool | Callable[[Engine, int], bool] = True,
        epoch_logger: Callable[[Engine], Any] | None = None,
        iteration_logger: Callable[[Engine], Any] | None = None,
        output_transform: Callable = lambda x: x[0],
        global_epoch_transform: Callable = lambda x: x,
        state_attributes: Sequence[str] | None = None,
        tag_name: str = DEFAULT_TAG,
        experiment_name: str = "monai_experiment",
        run_name: str | None = None,
        experiment_param: dict | None = None,
        artifacts: str | Sequence[Path] | None = None,
        optimizer_param_names: str | Sequence[str] = "lr",
        close_on_complete: bool = False,
    ) -> None:
        self.iteration_log = iteration_log
        self.epoch_log = epoch_log
        self.epoch_logger = epoch_logger
        self.iteration_logger = iteration_logger
        self.output_transform = output_transform
        self.global_epoch_transform = global_epoch_transform

Review comment (bug):
`mlflow` is a optional import. This will force to load the mlflow when python is checking the class defs before instantiation, which could be the cause of failures in the min tests (which does not install dev packages). 

Code after:
            The default behavior is to track loss from output[0] as output is a decollated list
            and we replicated loss value for every item of the decollated list.
            `engine.state` and `output_transform` inherit from the ignite concept:
            https://pytorch-ignite.ai/concepts/03-state/, explanation and usage example are in the tutorial:
            https://github.com/Project-MONAI/tutorials/blob/master/modules/batch_output_transform.ipynb.
        global_epoch_transform: a callable that is used to customize global epoch number.
            For example, in evaluation, the evaluator engine might want to track synced epoch number
            with the trainer engine.
        state_attributes: expected attributes from `engine.state`, if provided, will extract them
            when epoch completed.
        tag_name: when iteration output is a scalar, `tag_name` is used to track, defaults to `'Loss'`.
        experiment_name: the experiment name of MLflow, default to `'monai_experiment'`. An experiment can be
            used to record several runs.
        run_name: the run name in an experiment. A run can be used to record information about a workflow,
            like the loss, metrics and so on.
        experiment_param: a dict recording parameters which will not change through the whole workflow,
            like torch version, cuda version and so on.
        artifacts: paths to images that need to be recorded after running the workflow.
        optimizer_param_names: parameter names in the optimizer that need to be recorded during running the
            workflow, default 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.

    """

    # parameters that are logged at the start of training
    default_tracking_params = ["max_epochs", "epoch_length"]

    def __init__(
        self,
        tracking_uri: str | None = None,
        iteration_log: bool | Callable[[Engine, int], bool] = True,
        epoch_log: bool | Callable[[Engine, int], bool] = True,
        epoch_logger: Callable[[Engine], Any] | None = None,
        iteration_logger: Callable[[Engine], Any] | None = None,
        output_transform: Callable = lambda x: x[0],
        global_epoch_transform: Callable = lambda x: x,
        state_attributes: Sequence[str] | None = None,
        tag_name: str = DEFAULT_TAG,
        experiment_name: str = "monai_experiment",
        run_name: str | None = None,
        experiment_param: dict | None = None,
        artifacts: str | Sequence[Path] | None = None,
        optimizer_param_names: str | Sequence[str] = "lr",
        close_on_complete: bool = False,
    ) -> None:
        self.iteration_log = iteration_log
        self.epoch_log = epoch_log
        self.epoch_logger = epoch_logger
        self.iteration_logger = iteration_logger