From c27cc980e078113cec3355f892344a4e6f5a9515 Mon Sep 17 00:00:00 2001 From: "Dr. Behrooz Hashemian" <3968947+drbeh@users.noreply.github.com> Date: Tue, 6 Feb 2024 03:58:15 -0500 Subject: [PATCH] Add support for mlflow experiment name in auto3dseg (#7442) Fixes #7441 This PR enable Auto3DSeg users to manage their runs and experiment more efficiently in MLFlow under arbitrary experiment names, by providing experiment name as an input parameter. ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality). change). - [x] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [x] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [x] In-line docstrings updated. --------- Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com> --- monai/apps/auto3dseg/auto_runner.py | 12 ++++++++- monai/apps/auto3dseg/bundle_gen.py | 32 ++++++++++++++++++++++-- monai/apps/auto3dseg/ensemble_builder.py | 2 +- 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/monai/apps/auto3dseg/auto_runner.py b/monai/apps/auto3dseg/auto_runner.py index e4c2d908b7..52a0824227 100644 --- a/monai/apps/auto3dseg/auto_runner.py +++ b/monai/apps/auto3dseg/auto_runner.py @@ -85,6 +85,7 @@ class AutoRunner: can be skipped based on the analysis on the dataset from Auto3DSeg DataAnalyzer. mlflow_tracking_uri: a tracking URI for MLflow server which could be local directory or address of the remote tracking Server; MLflow runs will be recorded locally in algorithms' model folder if the value is None. + mlflow_experiment_name: the name of the experiment in MLflow server. kwargs: image writing parameters for the ensemble inference. The kwargs format follows the SaveImage transform. For more information, check https://docs.monai.io/en/stable/transforms.html#saveimage. @@ -212,6 +213,7 @@ def __init__( templates_path_or_url: str | None = None, allow_skip: bool = True, mlflow_tracking_uri: str | None = None, + mlflow_experiment_name: str | None = None, **kwargs: Any, ): if input is None and os.path.isfile(os.path.join(os.path.abspath(work_dir), "input.yaml")): @@ -253,6 +255,7 @@ def __init__( self.hpo = hpo and has_nni self.hpo_backend = hpo_backend self.mlflow_tracking_uri = mlflow_tracking_uri + self.mlflow_experiment_name = mlflow_experiment_name self.kwargs = deepcopy(kwargs) # parse input config for AutoRunner param overrides @@ -268,7 +271,13 @@ def __init__( if param in self.data_src_cfg and isinstance(self.data_src_cfg[param], bool): setattr(self, param, self.data_src_cfg[param]) # e.g. self.analyze = self.data_src_cfg["analyze"] - for param in ["algos", "hpo_backend", "templates_path_or_url", "mlflow_tracking_uri"]: # override from config + for param in [ + "algos", + "hpo_backend", + "templates_path_or_url", + "mlflow_tracking_uri", + "mlflow_experiment_name", + ]: # override from config if param in self.data_src_cfg: setattr(self, param, self.data_src_cfg[param]) # e.g. self.algos = self.data_src_cfg["algos"] @@ -813,6 +822,7 @@ def run(self): data_stats_filename=self.datastats_filename, data_src_cfg_name=self.data_src_cfg_name, mlflow_tracking_uri=self.mlflow_tracking_uri, + mlflow_experiment_name=self.mlflow_experiment_name, ) if self.gpu_customization: diff --git a/monai/apps/auto3dseg/bundle_gen.py b/monai/apps/auto3dseg/bundle_gen.py index 03b9c8bbf4..8a54d18be7 100644 --- a/monai/apps/auto3dseg/bundle_gen.py +++ b/monai/apps/auto3dseg/bundle_gen.py @@ -85,7 +85,8 @@ def __init__(self, template_path: PathLike): self.template_path = template_path self.data_stats_files = "" self.data_list_file = "" - self.mlflow_tracking_uri = None + self.mlflow_tracking_uri: str | None = None + self.mlflow_experiment_name: str | None = None self.output_path = "" self.name = "" self.best_metric = None @@ -139,7 +140,16 @@ def set_mlflow_tracking_uri(self, mlflow_tracking_uri: str | None) -> None: the remote tracking Server; MLflow runs will be recorded locally in algorithms' model folder if the value is None. """ - self.mlflow_tracking_uri = mlflow_tracking_uri # type: ignore + self.mlflow_tracking_uri = mlflow_tracking_uri + + def set_mlflow_experiment_name(self, mlflow_experiment_name: str | None) -> None: + """ + Set the experiment name for MLflow server + + Args: + mlflow_experiment_name: a string to specify the experiment name for MLflow server. + """ + self.mlflow_experiment_name = mlflow_experiment_name def fill_template_config(self, data_stats_filename: str, algo_path: str, **kwargs: Any) -> dict: """ @@ -447,6 +457,7 @@ class BundleGen(AlgoGen): mlflow_tracking_uri: a tracking URI for MLflow server which could be local directory or address of the remote tracking Server; MLflow runs will be recorded locally in algorithms' model folder if the value is None. + mlfow_experiment_name: a string to specify the experiment name for MLflow server. .. code-block:: bash python -m monai.apps.auto3dseg BundleGen generate --data_stats_filename="../algorithms/datastats.yaml" @@ -460,6 +471,7 @@ def __init__( data_stats_filename: str | None = None, data_src_cfg_name: str | None = None, mlflow_tracking_uri: str | None = None, + mlflow_experiment_name: str | None = None, ): if algos is None or isinstance(algos, (list, tuple, str)): if templates_path_or_url is None: @@ -513,6 +525,7 @@ def __init__( self.data_stats_filename = data_stats_filename self.data_src_cfg_name = data_src_cfg_name self.mlflow_tracking_uri = mlflow_tracking_uri + self.mlflow_experiment_name = mlflow_experiment_name self.history: list[dict] = [] def set_data_stats(self, data_stats_filename: str) -> None: @@ -552,10 +565,23 @@ def set_mlflow_tracking_uri(self, mlflow_tracking_uri): """ self.mlflow_tracking_uri = mlflow_tracking_uri + def set_mlflow_experiment_name(self, mlflow_experiment_name): + """ + Set the experiment name for MLflow server + + Args: + mlflow_experiment_name: a string to specify the experiment name for MLflow server. + """ + self.mlflow_experiment_name = mlflow_experiment_name + def get_mlflow_tracking_uri(self): """Get the tracking URI for MLflow server""" return self.mlflow_tracking_uri + def get_mlflow_experiment_name(self): + """Get the experiment name for MLflow server""" + return self.mlflow_experiment_name + def get_history(self) -> list: """Get the history of the bundleAlgo object with their names/identifiers""" return self.history @@ -608,10 +634,12 @@ def generate( data_stats = self.get_data_stats() data_src_cfg = self.get_data_src() mlflow_tracking_uri = self.get_mlflow_tracking_uri() + mlflow_experiment_name = self.get_mlflow_experiment_name() gen_algo = deepcopy(algo) gen_algo.set_data_stats(data_stats) gen_algo.set_data_source(data_src_cfg) gen_algo.set_mlflow_tracking_uri(mlflow_tracking_uri) + gen_algo.set_mlflow_experiment_name(mlflow_experiment_name) name = f"{gen_algo.name}_{f_id}" if allow_skip: diff --git a/monai/apps/auto3dseg/ensemble_builder.py b/monai/apps/auto3dseg/ensemble_builder.py index e29745e5cf..b2bea806de 100644 --- a/monai/apps/auto3dseg/ensemble_builder.py +++ b/monai/apps/auto3dseg/ensemble_builder.py @@ -464,7 +464,7 @@ def set_ensemble_method(self, ensemble_method_name: str = "AlgoEnsembleBestByFol ensemble_method_name, supported=["AlgoEnsembleBestN", "AlgoEnsembleBestByFold"] ) if self.ensemble_method_name == "AlgoEnsembleBestN": - n_best = kwargs.pop("n_best", False) or 2 + n_best = kwargs.pop("n_best", 2) self.ensemble_method = AlgoEnsembleBestN(n_best=n_best) elif self.ensemble_method_name == "AlgoEnsembleBestByFold": self.ensemble_method = AlgoEnsembleBestByFold(n_fold=self.num_fold) # type: ignore