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

Code before:
            self._log_params(self.experiment_param)

        attrs = {attr: getattr(engine.state, attr, None) for attr in self.default_tracking_params}
        self._delete_exist_param_in_dict(attrs)
        self._log_params(attrs)

        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:
            for _ in range(3):
                try:
                    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)
                    break
                except MlflowException as e:
                    if "RESOURCE_ALREADY_EXISTS" in str(e):
                        logger.warning("Experiment already exists; delaying before retrying.")
                        time.sleep(1)
                        continue
                    else:
                        raise e

        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}"

Review comment (suggestion):
I don't think this line is needed. Furthermore, if the exception is raised three times in a row the loop will silently do nothing. In place of this line we should check if this is the last time through the loop and raise the exception if so instead of sleeping. The check after this loop isn't sufficient since it raises a different exception which won't have the  information that the MlflowException would have.

Code after:
            self._log_params(self.experiment_param)

        attrs = {attr: getattr(engine.state, attr, None) for attr in self.default_tracking_params}
        self._delete_exist_param_in_dict(attrs)
        self._log_params(attrs)

        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:
            for _retry_time in range(3):
                try:
                    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)
                    break
                except MlflowException as e:
                    if "RESOURCE_ALREADY_EXISTS" in str(e):
                        logger.warning("Experiment already exists; delaying before retrying.")
                        time.sleep(1)
                        if _retry_time == 2:
                            raise e
                    else:
                        raise e

        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))