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

Code before:

    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:
        if not self.cur_run:
            raise ValueError("Current Run is not Active to log metrics")

        run_id = self.cur_run.info.run_id
        timestamp = int(time.time() * 1000)
        metrics_arr = [mlflow.entities.Metric(key, value, timestamp, step or 0) for key, value in metrics.items()]
        self.client.log_batch(run_id=run_id, metrics=metrics_arr, params=[], tags=[])

    def _parse_artifacts(self):
        """
        Log artifacts to mlflow. Given a path, all files in the path will be logged recursively.
        Given a file, it will be logged to mlflow.
        """
        artifact_list = []
        for path_name in self.artifacts:
            # in case the input is (None,) by default
            if not path_name:
                continue

Review comment (suggestion):
I am concerned about whether we need to support this new feature for now. Seems this feature is still experimental and may change or be removed in a future release without warning, then we will be easily influenced by their release. Or maybe we should fix the version of mlflow to 2.4.0?
https://mlflow.org/docs/latest/python_api/mlflow.client.html?highlight=log_inputs#mlflow.client.MlflowClient.log_inputs

Code after:
    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()]
        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:
        if not self.cur_run:
            raise ValueError("Current Run is not Active to log metrics")

        run_id = self.cur_run.info.run_id
        timestamp = int(time.time() * 1000)
        metrics_arr = [mlflow.entities.Metric(key, value, timestamp, step or 0) for key, value in metrics.items()]
        self.client.log_batch(run_id=run_id, metrics=metrics_arr, params=[], tags=[])

    def _parse_artifacts(self):
        """
        Log artifacts to mlflow. Given a path, all files in the path will be logged recursively.
        Given a file, it will be logged to mlflow.
        """
        artifact_list = []
        for path_name in self.artifacts: