Skip to content

Commit

Permalink
Add support for mlflow experiment name in auto3dseg (Project-MONAI#7442)
Browse files Browse the repository at this point in the history
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: Yu0610 <612410030@alum.ccu.edu.tw>
  • Loading branch information
drbeh authored and Yu0610 committed Apr 11, 2024
1 parent 1e328c5 commit ed70499
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 4 deletions.
12 changes: 11 additions & 1 deletion monai/apps/auto3dseg/auto_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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")):
Expand Down Expand Up @@ -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
Expand All @@ -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"]

Expand Down Expand Up @@ -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:
Expand Down
32 changes: 30 additions & 2 deletions monai/apps/auto3dseg/bundle_gen.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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"
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion monai/apps/auto3dseg/ensemble_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit ed70499

Please sign in to comment.