diff --git a/monai/bundle/__init__.py b/monai/bundle/__init__.py index 38cbb2c264..fba3aee994 100644 --- a/monai/bundle/__init__.py +++ b/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 diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index fc0d2a09b5..052b5ab763 100644 --- a/monai/bundle/scripts.py +++ b/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": "", "trainer": { "_target_": "MLFlowHandler", - "tracking_uri": "", + "tracking_uri": "@tracking_uri", "iteration_log": True, "output_transform": "$monai.handlers.from_engine(['loss'], first=True)", }, - "validator": {"_target_": "MLFlowHandler", "tracking_uri": "", "iteration_log": False}, - "evaluator": {"_target_": "MLFlowHandler", "tracking_uri": "", "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_) diff --git a/monai/bundle/utils.py b/monai/bundle/utils.py index 42d38bfeea..99afb59b58 100644 --- a/monai/bundle/utils.py +++ b/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: """ diff --git a/monai/fl/client/monai_algo.py b/monai/fl/client/monai_algo.py index f18e7faa63..40819c5612 100644 --- a/monai/fl/client/monai_algo.py +++ b/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": "", + "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) diff --git a/tests/test_fl_monai_algo.py b/tests/test_fl_monai_algo.py index 0627235a18..6b234fbe96 100644 --- a/tests/test_fl_monai_algo.py +++ b/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): diff --git a/tests/test_integration_bundle_run.py b/tests/test_integration_bundle_run.py index 8d2076e235..986dd77076 100644 --- a/tests/test_integration_bundle_run.py +++ b/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__":