Add support for mlflow experiment name in auto3dseg (Project-MONAI#7442)

Fixes Project-MONAI#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
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [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>
Signed-off-by: Juan Pablo de la Cruz Gutiérrez <juampatronics@gmail.com>

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:

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:

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