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

Code before:
        experiment = self.experiment
        if not experiment:
            experiment = self.client.get_experiment_by_name(self.experiment_name)
            if not experiment:
                experiment_id = self.client.create_experiment(self.experiment_name)
                experiment = self.client.get_experiment(experiment_id)

        if experiment.lifecycle_stage != mlflow.entities.LifecycleStage.ACTIVE:
            raise ValueError(f"Cannot set a deleted experiment '{self.experiment_name}' as the active experiment")
        self.experiment = experiment

    @staticmethod
    def _get_pandas_dataset_info(pandas_dataset):
        dataset_name = pandas_dataset.name
        return {
            f"{dataset_name}_digest": pandas_dataset.digest,
            f"{dataset_name}_samples": pandas_dataset.profile["num_rows"],
        }

    def _log_dataset(self, sample_dict: dict[str, Any], context: str = "nontrain") -> None:
        if not self.cur_run:
            raise ValueError("Current Run is not Active to log the dataset")

        # Need to update the self.cur_run to sync the dataset log, otherwise the `inputs` info will be out-of-date.
        self.cur_run = self.client.get_run(self.cur_run.info.run_id)
        logged_train_set = [x for x in self.cur_run.inputs.dataset_inputs if "train" == x.dataset.name[: len("train")]]
        logged_nontrain_set = [
            x for x in self.cur_run.inputs.dataset_inputs if "nontrain" == x.dataset.name[: len("nontrain")]
        ]
        # In case there are more datasets.
        dataset_cnt = str(len(logged_nontrain_set if context == "nontrain" else logged_train_set))
        dataset_name = f"{context}_dataset_{dataset_cnt}"
        sample_df = pandas.DataFrame(sample_dict)
        dataset = mlflow.data.from_pandas(sample_df, name=dataset_name)
        exist_dataset_list = list(
            filter(lambda x: x.dataset.digest == dataset.digest, self.cur_run.inputs.dataset_inputs)
        )

        if not len(exist_dataset_list):
            datasets = [mlflow.entities.DatasetInput(dataset._to_mlflow_entity())]
            self.client.log_inputs(run_id=self.cur_run.info.run_id, datasets=datasets)
            dataset_info = MLFlowHandler._get_pandas_dataset_info(dataset)
            self._log_params(dataset_info)

    def _log_params(self, params: dict[str, Any]) -> None:
        if not self.cur_run:
            raise ValueError("Current Run is not Active to log params")
        params_arr = [mlflow.entities.Param(key, str(value)) for key, value in params.items()]
        self.client.log_batch(run_id=self.cur_run.info.run_id, metrics=[], params=params_arr, tags=[])

    def _log_metrics(self, metrics: dict[str, Any], step: int | None = None) -> None:

Review comment (suggestion):
`if "train" == x.dataset.name[: len("train")]]` looks very confusing...
Maybe `x.dataset.name.startswith("train")`?

Thanks.


Code after:
        if self.dataset_logger:
            self.dataset_logger(self.dataset_dict)
        else:
            self._default_dataset_log(self.dataset_dict)

    def _set_experiment(self):
        experiment = self.experiment
        if not experiment:
            experiment = self.client.get_experiment_by_name(self.experiment_name)
            if not experiment:
                experiment_id = self.client.create_experiment(self.experiment_name)
                experiment = self.client.get_experiment(experiment_id)

        if experiment.lifecycle_stage != mlflow.entities.LifecycleStage.ACTIVE:
            raise ValueError(f"Cannot set a deleted experiment '{self.experiment_name}' as the active experiment")
        self.experiment = experiment

    @staticmethod
    def _get_pandas_dataset_info(pandas_dataset):
        dataset_name = pandas_dataset.name
        return {
            f"{dataset_name}_digest": pandas_dataset.digest,
            f"{dataset_name}_samples": pandas_dataset.profile["num_rows"],
        }

    def _log_dataset(self, sample_dict: dict[str, Any], context: str = "train") -> None:
        if not self.cur_run:
            raise ValueError("Current Run is not Active to log the dataset")

        # Need to update the self.cur_run to sync the dataset log, otherwise the `inputs` info will be out-of-date.
        self.cur_run = self.client.get_run(self.cur_run.info.run_id)
        logged_set = [x for x in self.cur_run.inputs.dataset_inputs if x.dataset.name.startswith(context)]
        # In case there are datasets with the same name.
        dataset_count = str(len(logged_set))
        dataset_name = f"{context}_dataset_{dataset_count}"
        sample_df = pandas.DataFrame(sample_dict)
        dataset = mlflow.data.from_pandas(sample_df, name=dataset_name)
        exist_dataset_list = list(
            filter(lambda x: x.dataset.digest == dataset.digest, self.cur_run.inputs.dataset_inputs)
        )

        if not len(exist_dataset_list):
            datasets = [mlflow.entities.DatasetInput(dataset._to_mlflow_entity())]
            self.client.log_inputs(run_id=self.cur_run.info.run_id, datasets=datasets)
            dataset_info = MLFlowHandler._get_pandas_dataset_info(dataset)
            self._log_params(dataset_info)

    def _log_params(self, params: dict[str, Any]) -> None:
        if not self.cur_run:
            raise ValueError("Current Run is not Active to log params")
        params_arr = [mlflow.entities.Param(key, str(value)) for key, value in params.items()]