Add continue_training parameter to nnUNetExecutor and update job config timeout

Author: SimoneBendazzoli93

monai/nvflare/nnunet_executor.py

@@ -124,6 +124,7 @@ def __init__(
         modality_list=None,
         train_extra_configs=None,
         exclude_vars=None,
+        continue_training=False,
     ):
         super().__init__()
 
@@ -147,6 +148,7 @@ def __init__(
         self.bundle_root = bundle_root
         self.train_extra_configs = train_extra_configs
         self.modality_list = modality_list
+        self.continue_training = continue_training
 
     def handle_event(self, event_type: str, fl_ctx: FLContext):
         if event_type == EventType.START_RUN:
@@ -305,6 +307,7 @@ def train(self):
             dataset_name_or_id=self.nnunet_config["dataset_name_or_id"],
             run_with_bundle=True if self.bundle_root is not None else False,
             bundle_root=self.bundle_root,
+            continue_training=self.continue_training
         )
         outgoing_dxo = DXO(data_kind=DataKind.COLLECTION, data=validation_summary, meta={})
         return outgoing_dxo.to_shareable()

monai/nvflare/nvflare_generate_job_configs.py

@@ -665,7 +665,7 @@ def train_config(clients, experiment, root_dir, script_dir, nvflare_exec):
                     "min_responses_required": 0,
                     "wait_time_after_min_received": 10,
                     "task_name": task_name,
-                    "timeout": 6000,
+                    "timeout": 600000,
                 },
             }
         ],
@@ -699,6 +699,7 @@ def train_config(clients, experiment, root_dir, script_dir, nvflare_exec):
                             },
                             "client_name": clients[client_id]["client_name"],
                             "tracking_uri": experiment["tracking_uri"],
+                            "continue_training": experiment["continue_training"],
                         },
                     },
                 }

monai/nvflare/nvflare_nnunet.py

@@ -46,6 +46,8 @@ def train(
     fold=0,
     bundle_root=None,
     mlflow_token=None,
+    continue_training=False,
+    resume_epoch="latest",
 ):
     """
 
@@ -75,6 +77,10 @@ def train(
         Root directory for MONAI bundle, by default None.
     mlflow_token : str, optional
         Token for MLflow authentication, by default None.
+    continue_training : bool, optional
+        Whether to continue training from a checkpoint, by default False.
+    resume_epoch : int, optional
+        Epoch to resume training from, by default "latest".
 
     Returns
     -------
@@ -89,15 +95,19 @@ def train(
     else:
         os.environ["BUNDLE_ROOT"] = bundle_root
         os.environ["PYTHONPATH"] = os.environ["PYTHONPATH"] + ":" + bundle_root
+        config_files = os.path.join(bundle_root, "configs", "train_resume.yaml")
+        if continue_training:
+            config_files = [os.path.join(bundle_root, "configs", "train.yaml"), os.path.join(bundle_root, "configs", "train_continue.yaml")]
         monai.bundle.run(
-            config_file=Path(bundle_root).joinpath("configs/train.yaml"),
+            config_file=config_files,
             bundle_root=bundle_root,
             nnunet_trainer_class_name=trainer_class_name,
             mlflow_experiment_name=experiment_name,
             mlflow_run_name="run_" + client_name,
             tracking_uri=tracking_uri,
             fold_id=fold,
             nnunet_root_folder=nnunet_root_dir,
+            reload_checkpoint_epoch=resume_epoch
         )
         nnunet_config = {"dataset_name_or_id": dataset_name_or_id, "nnunet_trainer": trainer_class_name}
         convert_monai_bundle_to_nnunet(nnunet_config, bundle_root)
@@ -619,6 +629,7 @@ def prepare_bundle(bundle_config, train_extra_configs=None):
         train_config["mlflow_run_name"] = bundle_config["mlflow_run_name"]
 
         train_config["data_src_cfg"] = "$@nnunet_root_folder+'/data_src_cfg.yaml'"
+        train_config["nnunet_root_folder"] = "."
         train_config["runner"] = {
             "_target_": "nnUNetV2Runner",
             "input_config": "$@data_src_cfg",