5385 Add support to override experiment tracking configs

monai/bundle/__init__.py

@@ -25,4 +25,4 @@
     verify_metadata,
     verify_net_in_out,
 )
-from .utils import EXPR_KEY, ID_REF_KEY, ID_SEP_KEY, MACRO_KEY, load_bundle_config
+from .utils import DEFAULT_EXP_MGMT_SETTINGS, EXPR_KEY, ID_REF_KEY, ID_SEP_KEY, MACRO_KEY, load_bundle_config

monai/bundle/scripts.py

@@ -27,7 +27,7 @@
 from monai.apps.utils import _basename, download_url, extractall, get_logger
 from monai.bundle.config_item import ConfigComponent
 from monai.bundle.config_parser import ConfigParser
-from monai.bundle.utils import DEFAULT_INFERENCE, DEFAULT_METADATA, DEFAULT_MLFLOW_SETTINGS
+from monai.bundle.utils import DEFAULT_EXP_MGMT_SETTINGS, DEFAULT_INFERENCE, DEFAULT_METADATA
 from monai.config import IgniteInfo, PathLike
 from monai.data import load_net_with_metadata, save_net_with_metadata
 from monai.networks import convert_to_torchscript, copy_model_state, get_state_dict, save_state
@@ -478,16 +478,17 @@ def patch_bundle_tracking(parser: ConfigParser, settings: dict):
         settings: settings for the experiment tracking, should follow the pattern of default settings.
 
     """
-    for k, v in settings["handlers_id"].items():
-        engine = parser.get(v["id"])
-        if engine is not None:
-            handlers = parser.get(v["handlers"])
-            handler_config = settings["configs"].get(k, None)
-            if handler_config is not None:
+    for k, v in settings["configs"].items():
+        if k in settings["handlers_id"]:
+            engine = parser.get(settings["handlers_id"][k]["id"])
+            if engine is not None:
+                handlers = parser.get(settings["handlers_id"][k]["handlers"])
                 if handlers is None:
-                    engine["train_handlers" if k == "trainer" else "val_handlers"] = [handler_config]
+                    engine["train_handlers" if k == "trainer" else "val_handlers"] = [v]
                 else:
-                    handlers.append(handler_config)
+                    handlers.append(v)
+        elif k not in parser:
+            parser[k] = v
 
 
 def run(
@@ -533,6 +534,7 @@ def run(
             if "mlflow", will add `MLFlowHandler` to the parsed bundle with default loggging settings,
             if other string, treat it as file path to load the logging settings, if `dict`,
             treat it as logging settings, otherwise, use all the default settings.
+            will patch the target config content with `tracking handlers` and the top-level items of `configs`.
             example of customized settings:
 
             .. code-block:: python
@@ -544,16 +546,21 @@ def run(
                         "evaluator": {"id": "evaluator", "handlers": "handlers"},
                     },
                     "configs": {
+                        "tracking_uri": "<path>",
                         "trainer": {
                             "_target_": "MLFlowHandler",
-                            "tracking_uri": "<path>",
+                            "tracking_uri": "@tracking_uri",
                             "iteration_log": True,
                             "output_transform": "$monai.handlers.from_engine(['loss'], first=True)",
                         },
-                        "validator": {"_target_": "MLFlowHandler", "tracking_uri": "<path>", "iteration_log": False},
-                        "evaluator": {"_target_": "MLFlowHandler", "tracking_uri": "<path>", "iteration_log": False},
-                    }
-                }
+                        "validator": {
+                            "_target_": "MLFlowHandler", "tracking_uri": "@tracking_uri", "iteration_log": False,
+                        },
+                        "evaluator": {
+                            "_target_": "MLFlowHandler", "tracking_uri": "@tracking_uri", "iteration_log": False,
+                        },
+                    },
+                },
 
         args_file: a JSON or YAML file to provide default values for `runner_id`, `meta_file`,
             `config_file`, `logging`, and override pairs. so that the command line inputs can be simplified.
@@ -591,10 +598,10 @@ def run(
     # the rest key-values in the _args are to override config content
     parser.update(pairs=_args)
 
-    # patch the bundle with mlflow handler
+    # set tracking configs for experiment management
     if tracking_ is not None:
-        if isinstance(tracking_, str) and tracking_ == "mlflow":
-            settings_ = DEFAULT_MLFLOW_SETTINGS
+        if isinstance(tracking_, str) and tracking_ in DEFAULT_EXP_MGMT_SETTINGS:
+            settings_ = DEFAULT_EXP_MGMT_SETTINGS[tracking_]
         else:
             settings_ = ConfigParser.load_config_files(tracking_)
         patch_bundle_tracking(parser=parser, settings=settings_)

monai/bundle/utils.py

@@ -104,22 +104,25 @@
 DEFAULT_MLFLOW_SETTINGS = {
     "handlers_id": DEFAULT_HANDLERS_ID,
     "configs": {
+        "tracking_uri": "$@output_dir + '/mlruns'",
         # MLFlowHandler config for the trainer
         "trainer": {
             "_target_": "MLFlowHandler",
-            "tracking_uri": "$@output_dir + '/mlflow'",
+            "tracking_uri": "@tracking_uri",
             "iteration_log": True,
             "epoch_log": True,
             "tag_name": "train_loss",
             "output_transform": "$monai.handlers.from_engine(['loss'], first=True)",
         },
         # MLFlowHandler config for the validator
-        "validator": {"_target_": "MLFlowHandler", "tracking_uri": "$@output_dir + '/mlflow'", "iteration_log": False},
+        "validator": {"_target_": "MLFlowHandler", "tracking_uri": "@tracking_uri", "iteration_log": False},
         # MLFlowHandler config for the evaluator
-        "evaluator": {"_target_": "MLFlowHandler", "tracking_uri": "$@output_dir + '/mlflow'", "iteration_log": False},
+        "evaluator": {"_target_": "MLFlowHandler", "tracking_uri": "@tracking_uri", "iteration_log": False},
     },
 }
 
+DEFAULT_EXP_MGMT_SETTINGS = {"mlflow": DEFAULT_MLFLOW_SETTINGS}  # default expriment management settings
+
 
 def load_bundle_config(bundle_path: str, *config_names, **load_kw_args) -> Any:
     """

monai/fl/client/monai_algo.py

@@ -20,8 +20,7 @@
 import monai
 from monai.apps.auto3dseg.data_analyzer import DataAnalyzer
 from monai.auto3dseg import SegSummarizer
-from monai.bundle import ConfigParser
-from monai.bundle.config_item import ConfigComponent, ConfigItem
+from monai.bundle import DEFAULT_EXP_MGMT_SETTINGS, ConfigComponent, ConfigItem, ConfigParser, patch_bundle_tracking
 from monai.fl.client import ClientAlgo, ClientAlgoStats
 from monai.fl.utils.constants import (
     BundleKeys,
@@ -93,7 +92,7 @@ def disable_ckpt_loaders(parser):
 
 class MonaiAlgoStats(ClientAlgoStats):
     """
-    Implementation of ``ClientAlgo`` to allow federated learning with MONAI bundle configurations.
+    Implementation of ``ClientAlgoStats`` to allow federated learning with MONAI bundle configurations.
 
     Args:
         bundle_root: path of bundle.
@@ -343,6 +342,38 @@ class MonaiAlgo(ClientAlgo, MonaiAlgoStats):
         multi_gpu: whether to run MonaiAlgo in a multi-GPU setting; defaults to `False`.
         backend: backend to use for torch.distributed; defaults to "nccl".
         init_method: init_method for torch.distributed; defaults to "env://".
+        tracking: enable the experiment tracking feature at runtime with optionally configurable and extensible.
+            if "mlflow", will add `MLFlowHandler` to the parsed bundle with default loggging settings,
+            if other string, treat it as file path to load the logging settings, if `dict`,
+            treat it as logging settings, otherwise, use all the default settings.
+            will patch the target config content with `tracking handlers` and the top-level items of `configs`.
+            example of customized settings:
+
+            .. code-block:: python
+
+                tracking = {
+                    "handlers_id": {
+                        "trainer": {"id": "train#trainer", "handlers": "train#handlers"},
+                        "validator": {"id": "evaluate#evaluator", "handlers": "evaluate#handlers"},
+                        "evaluator": {"id": "evaluator", "handlers": "handlers"},
+                    },
+                    "configs": {
+                        "tracking_uri": "<path>",
+                        "trainer": {
+                            "_target_": "MLFlowHandler",
+                            "tracking_uri": "@tracking_uri",
+                            "iteration_log": True,
+                            "output_transform": "$monai.handlers.from_engine(['loss'], first=True)",
+                        },
+                        "validator": {
+                            "_target_": "MLFlowHandler", "tracking_uri": "@tracking_uri", "iteration_log": False,
+                        },
+                        "evaluator": {
+                            "_target_": "MLFlowHandler", "tracking_uri": "@tracking_uri", "iteration_log": False,
+                        },
+                    },
+                },
+
     """
 
     def __init__(
@@ -365,6 +396,7 @@ def __init__(
         train_data_key: Optional[str] = BundleKeys.TRAIN_DATA,
         eval_data_key: Optional[str] = BundleKeys.VALID_DATA,
         data_stats_transform_list: Optional[list] = None,
+        tracking: Optional[Union[str, dict]] = None,
     ):
         self.logger = logging.getLogger(self.__class__.__name__)
         if config_evaluate_filename == "default":
@@ -387,6 +419,7 @@ def __init__(
         self.train_data_key = train_data_key
         self.eval_data_key = eval_data_key
         self.data_stats_transform_list = data_stats_transform_list
+        self.tracking = tracking
 
         self.app_root = None
         self.train_parser = None
@@ -467,6 +500,15 @@ def initialize(self, extra=None):
             disable_ckpt_loaders(self.train_parser)
             disable_ckpt_loaders(self.eval_parser)
 
+        # set tracking configs for experiment management
+        if self.tracking is not None:
+            if isinstance(self.tracking, str) and self.tracking in DEFAULT_EXP_MGMT_SETTINGS:
+                settings_ = DEFAULT_EXP_MGMT_SETTINGS[self.tracking]
+            else:
+                settings_ = ConfigParser.load_config_files(self.tracking)
+            patch_bundle_tracking(parser=self.train_parser, settings=settings_)
+            patch_bundle_tracking(parser=self.eval_parser, settings=settings_)
+
         # Get trainer, evaluator
         self.trainer = self.train_parser.get_parsed_content(
             BundleKeys.TRAINER, default=ConfigItem(None, BundleKeys.TRAINER)

tests/test_fl_monai_algo.py

@@ -10,11 +10,15 @@
 # limitations under the License.
 
 import os
+import shutil
+import tempfile
 import unittest
+from pathlib import Path
 
 from parameterized import parameterized
 
 from monai.bundle import ConfigParser
+from monai.bundle.utils import DEFAULT_HANDLERS_ID
 from monai.fl.client.monai_algo import MonaiAlgo
 from monai.fl.utils.constants import ExtraItems
 from monai.fl.utils.exchange_object import ExchangeObject
@@ -130,6 +134,7 @@
 
 
 @SkipIfNoModule("ignite")
+@SkipIfNoModule("mlflow")
 class TestFLMonaiAlgo(unittest.TestCase):
     @parameterized.expand([TEST_TRAIN_1, TEST_TRAIN_2, TEST_TRAIN_3])
     def test_train(self, input_params):
@@ -141,6 +146,19 @@ def test_train(self, input_params):
         else:
             config_train_filename = os.path.join(input_params["bundle_root"], input_params["config_train_filename"])
 
+        data_dir = tempfile.mkdtemp()
+        # test experiment management
+        input_params["tracking"] = {
+            "handlers_id": DEFAULT_HANDLERS_ID,
+            "configs": {
+                "trainer": {
+                    "_target_": "MLFlowHandler",
+                    "tracking_uri": Path(data_dir).as_uri() + "/mlflow_override",
+                    "output_transform": "$monai.handlers.from_engine(['loss'], first=True)",
+                }
+            },
+        }
+
         # initialize algo
         algo = MonaiAlgo(**input_params)
         algo.initialize(extra={ExtraItems.CLIENT_NAME: "test_fl"})
@@ -157,6 +175,10 @@ def test_train(self, input_params):
         # test train
         algo.train(data=data, extra={})
         algo.finalize()
+        # must close it as we are changing different temp dir for cases here
+        algo.train_parser.get_parsed_content("train#handlers")[-1].close()
+        self.assertTrue(os.path.exists(f"{data_dir}/mlflow_override"))
+        shutil.rmtree(data_dir)
 
     @parameterized.expand([TEST_EVALUATE_1, TEST_EVALUATE_2, TEST_EVALUATE_3])
     def test_evaluate(self, input_params):

tests/test_integration_bundle_run.py

@@ -82,11 +82,12 @@ def test_shape(self, config_file, expected_shape):
         settings = {
             "handlers_id": DEFAULT_HANDLERS_ID,
             "configs": {
+                "no_epoch": True,  # test override config in the settings file
                 "evaluator": {
                     "_target_": "MLFlowHandler",
-                    "tracking_uri": "$@output_dir + '/mlflow_override'",
-                    "iteration_log": False,
-                }
+                    "tracking_uri": "$@output_dir + '/mlflow_override1'",
+                    "iteration_log": "@no_epoch",
+                },
             },
         }
         settings_file = os.path.join(tempdir, "mlflow.json")
@@ -111,22 +112,23 @@ def test_shape(self, config_file, expected_shape):
             override = f"--network %{overridefile1}#move_net --dataset#_target_ %{overridefile2}"
         # test with `monai.bundle` as CLI entry directly
         cmd = "-m monai.bundle run evaluating --postprocessing#transforms#2#output_postfix seg"
-        cmd += f" {override} --save_dir {tempdir} --output_dir {outdir}"
+        cmd += f" {override} --no_epoch False --save_dir {tempdir} --output_dir {outdir}"
         la = ["coverage", "run"] + cmd.split(" ") + ["--meta_file", meta_file] + ["--config_file", config_file]
         test_env = os.environ.copy()
         print(f"CUDA_VISIBLE_DEVICES in {__file__}", test_env.get("CUDA_VISIBLE_DEVICES"))
         command_line_tests(la + ["--args_file", def_args_file] + ["--tracking", settings_file])
         loader = LoadImage(image_only=True)
         self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_seg.nii.gz")).shape, expected_shape)
-        self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override"))
+        self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override1"))
 
+        tracking_uri = outdir + "/mlflow_override2"  # test override experiment management configs
         # here test the script with `google fire` tool as CLI
-        cmd = "-m fire monai.bundle.scripts run --runner_id evaluating --tracking mlflow"
-        cmd += f" --evaluator#amp False {override} --save_dir {tempdir} --output_dir {outdir}"
+        cmd = "-m fire monai.bundle.scripts run --runner_id evaluating --tracking mlflow --evaluator#amp False"
+        cmd += f" --tracking_uri {tracking_uri} {override} --save_dir {tempdir} --output_dir {outdir}"
         la = ["coverage", "run"] + cmd.split(" ") + ["--meta_file", meta_file] + ["--config_file", config_file]
         command_line_tests(la)
         self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_trans.nii.gz")).shape, expected_shape)
-        self.assertTrue(os.path.exists(f"{tempdir}/mlflow"))
+        self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override2"))
 
 
 if __name__ == "__main__":