Merge branch 'dev' into 8587-test-erros-on-pytorch-release-2508-on-series-50

Author: ericspod

monai/__init__.py

@@ -131,7 +131,7 @@ def filter(self, record):
     # workaround related to https://github.com/Project-MONAI/MONAI/issues/7575
     if hasattr(torch.cuda.device_count, "cache_clear"):
         torch.cuda.device_count.cache_clear()
-except BaseException:
+except Exception:
     from .utils.misc import MONAIEnvVars
 
     if MONAIEnvVars.debug():

monai/apps/auto3dseg/auto_runner.py

@@ -26,7 +26,7 @@
 from monai.apps.auto3dseg.hpo_gen import NNIGen
 from monai.apps.auto3dseg.utils import export_bundle_algo_history, import_bundle_algo_history
 from monai.apps.utils import get_logger
-from monai.auto3dseg.utils import algo_to_pickle
+from monai.auto3dseg.utils import algo_to_json
 from monai.bundle import ConfigParser
 from monai.transforms import SaveImage
 from monai.utils import AlgoKeys, has_option, look_up_option, optional_import
@@ -740,7 +740,7 @@ def _train_algo_in_sequence(self, history: list[dict[str, Any]]) -> None:
             acc = algo.get_score()
 
             algo_meta_data = {str(AlgoKeys.SCORE): acc}
-            algo_to_pickle(algo, template_path=algo.template_path, **algo_meta_data)
+            algo_to_json(algo, template_path=algo.template_path, **algo_meta_data)
 
     def _train_algo_in_nni(self, history: list[dict[str, Any]]) -> None:
         """

monai/apps/auto3dseg/bundle_gen.py

@@ -36,7 +36,7 @@
     _prepare_cmd_torchrun,
     _run_cmd_bcprun,
     _run_cmd_torchrun,
-    algo_to_pickle,
+    algo_to_json,
 )
 from monai.bundle.config_parser import ConfigParser
 from monai.config import PathLike
@@ -366,6 +366,39 @@ def get_output_path(self):
         """Returns the algo output paths to find the algo scripts and configs."""
         return self.output_path
 
+    def state_dict(self) -> dict:
+        """
+        Return state for serialization.
+
+        Returns:
+            A dictionary containing the BundleAlgo state to serialize.
+
+        Note:
+            template_path is excluded as it is determined dynamically at load time
+            based on which path successfully imports the Algo class.
+        """
+        return {
+            "data_stats_files": self.data_stats_files,
+            "data_list_file": self.data_list_file,
+            "mlflow_tracking_uri": self.mlflow_tracking_uri,
+            "mlflow_experiment_name": self.mlflow_experiment_name,
+            "output_path": self.output_path,
+            "name": self.name,
+            "best_metric": self.best_metric,
+            "fill_records": self.fill_records,
+            "device_setting": self.device_setting,
+        }
+
+    def load_state_dict(self, state: dict) -> None:
+        """
+        Restore state from a dictionary.
+
+        Args:
+            state: A dictionary containing the state to restore.
+        """
+        for key, value in state.items():
+            setattr(self, key, value)
+
 
 # path to download the algo_templates
 default_algo_zip = (
@@ -658,7 +691,7 @@ def generate(
                 else:
                     gen_algo.export_to_disk(output_folder, name, fold=f_id)
 
-                algo_to_pickle(gen_algo, template_path=algo.template_path)
+                algo_to_json(gen_algo, template_path=algo.template_path)
                 self.history.append(
                     {AlgoKeys.ID: name, AlgoKeys.ALGO: gen_algo}
                 )  # track the previous, may create a persistent history

monai/apps/auto3dseg/data_analyzer.py

@@ -341,7 +341,7 @@ def _get_all_case_stats(
                     _label_argmax = True  # track if label is argmaxed
                     batch_data[self.label_key] = label.to(device)
                 d = summarizer(batch_data)
-            except BaseException as err:
+            except Exception as err:
                 if "image_meta_dict" in batch_data.keys():
                     filename = batch_data["image_meta_dict"][ImageMetaKey.FILENAME_OR_OBJ]
                 else:
@@ -357,7 +357,7 @@ def _get_all_case_stats(
                                 label = torch.argmax(label, dim=0) if label.shape[0] > 1 else label[0]
                             batch_data[self.label_key] = label.to("cpu")
                         d = summarizer(batch_data)
-                    except BaseException as err:
+                    except Exception as err:
                         logger.info(f"Unable to process data {filename} on {device}. {err}")
                         continue
                 else:

monai/apps/auto3dseg/ensemble_builder.py

@@ -216,7 +216,7 @@ def __call__(self, pred_param: dict | None = None) -> list:
             if "image_save_func" in param:
                 try:
                     ensemble_preds = self.ensemble_pred(preds, sigmoid=sigmoid)
-                except BaseException:
+                except Exception:
                     ensemble_preds = self.ensemble_pred([_.to("cpu") for _ in preds], sigmoid=sigmoid)
                 res = img_saver(ensemble_preds)
                 # res is the path to the saved results

monai/apps/auto3dseg/hpo_gen.py

@@ -19,7 +19,7 @@
 
 from monai.apps.auto3dseg.bundle_gen import BundleAlgo
 from monai.apps.utils import get_logger
-from monai.auto3dseg import Algo, AlgoGen, algo_from_pickle, algo_to_pickle
+from monai.auto3dseg import Algo, AlgoGen, algo_from_json, algo_to_json
 from monai.bundle.config_parser import ConfigParser
 from monai.config import PathLike
 from monai.utils import optional_import
@@ -36,7 +36,7 @@ class HPOGen(AlgoGen):
     """
     The base class for hyperparameter optimization (HPO) interfaces to generate algos in the Auto3Dseg pipeline.
     The auto-generated algos are saved at their ``output_path`` on the disk. The files in the ``output_path``
-    may contain scripts that define the algo, configuration files, and pickle files that save the internal states
+    may contain scripts that define the algo, configuration files, and JSON files that save the internal states
     of the algo before/after the training. Compared to the BundleGen class, HPOGen generates Algo on-the-fly, so
     training and algo generation may be executed alternatively and take a long time to finish the generation process.
 
@@ -72,7 +72,7 @@ class NNIGen(HPOGen):
 
     Args:
         algo: an Algo object (e.g. BundleAlgo) with defined methods: ``get_output_path`` and train
-            and supports saving to and loading from pickle files via ``algo_from_pickle`` and ``algo_to_pickle``.
+            and supports saving to and loading via ``algo_from_json`` and ``algo_to_json``.
         params: a set of parameter to override the algo if override is supported by Algo subclass.
 
     Examples::
@@ -81,16 +81,16 @@ class NNIGen(HPOGen):
         ├── algorithm_templates
         │   └── unet
         ├── unet_0
-        │   ├── algo_object.pkl
+        │   ├── algo_object.json
         │   ├── configs
         │   └── scripts
         ├── unet_0_learning_rate_0.01
-        │   ├── algo_object.pkl
+        │   ├── algo_object.json
         │   ├── configs
         │   ├── model_fold0
         │   └── scripts
         └── unet_0_learning_rate_0.1
-            ├── algo_object.pkl
+            ├── algo_object.json
             ├── configs
             ├── model_fold0
             └── scripts
@@ -129,10 +129,10 @@ def __init__(self, algo: Algo | None = None, params: dict | None = None):
             else:
                 self.algo = algo
 
-            self.obj_filename = algo_to_pickle(self.algo, template_path=self.algo.template_path)
+            self.obj_filename = algo_to_json(self.algo, template_path=self.algo.template_path)
 
     def get_obj_filename(self):
-        """Return the filename of the dumped pickle algo object."""
+        """Return the filename of the dumped algo object."""
         return self.obj_filename
 
     def print_bundle_algo_instruction(self):
@@ -190,7 +190,7 @@ def generate(self, output_folder: str = ".") -> None:
         task_id = self.get_task_id()
         task_prefix = os.path.basename(self.algo.get_output_path())
         write_path = os.path.join(output_folder, task_prefix + task_id)
-        self.obj_filename = os.path.join(write_path, "algo_object.pkl")
+        self.obj_filename = os.path.join(write_path, "algo_object.json")
 
         if isinstance(self.algo, BundleAlgo):
             self.algo.export_to_disk(
@@ -214,15 +214,15 @@ def run_algo(self, obj_filename: str, output_folder: str = ".", template_path: P
         The python interface for NNI to run.
 
         Args:
-            obj_filename: the pickle-exported Algo object.
+            obj_filename: the serialized Algo object.
             output_folder: the root path of the algorithms templates.
             template_path: the algorithm_template. It must contain algo.py in the follow path:
                 ``{algorithm_templates_dir}/{network}/scripts/algo.py``
         """
         if not os.path.isfile(obj_filename):
             raise ValueError(f"{obj_filename} is not found")
 
-        self.algo, algo_meta_data = algo_from_pickle(obj_filename, template_path=template_path)
+        self.algo, algo_meta_data = algo_from_json(obj_filename, template_path=template_path)
 
         # step 1 sample hyperparams
         params = self.get_hyperparameters()
@@ -235,7 +235,7 @@ def run_algo(self, obj_filename: str, output_folder: str = ".", template_path: P
         acc = self.algo.get_score()
         algo_meta_data = {str(AlgoKeys.SCORE): acc}
 
-        algo_to_pickle(self.algo, template_path=self.algo.template_path, **algo_meta_data)
+        algo_to_json(self.algo, template_path=self.algo.template_path, **algo_meta_data)
         self.set_score(acc)
 
 
@@ -250,7 +250,7 @@ class OptunaGen(HPOGen):
 
     Args:
         algo: an Algo object (e.g. BundleAlgo). The object must at least define two methods: get_output_path and train
-            and supports saving to and loading from pickle files via ``algo_from_pickle`` and ``algo_to_pickle``.
+            and supports saving to and loading via ``algo_from_json`` and ``algo_to_json``.
         params: a set of parameter to override the algo if override is supported by Algo subclass.
 
     Examples::
@@ -259,16 +259,16 @@ class OptunaGen(HPOGen):
         ├── algorithm_templates
         │   └── unet
         ├── unet_0
-        │   ├── algo_object.pkl
+        │   ├── algo_object.json
         │   ├── configs
         │   └── scripts
         ├── unet_0_learning_rate_0.01
-        │   ├── algo_object.pkl
+        │   ├── algo_object.json
         │   ├── configs
         │   ├── model_fold0
         │   └── scripts
         └── unet_0_learning_rate_0.1
-            ├── algo_object.pkl
+            ├── algo_object.json
             ├── configs
             ├── model_fold0
             └── scripts
@@ -296,10 +296,10 @@ def __init__(self, algo: Algo | None = None, params: dict | None = None) -> None
             else:
                 self.algo = algo
 
-            self.obj_filename = algo_to_pickle(self.algo, template_path=self.algo.template_path)
+            self.obj_filename = algo_to_json(self.algo, template_path=self.algo.template_path)
 
     def get_obj_filename(self):
-        """Return the dumped pickle object of algo."""
+        """Return the dumped object of algo."""
         return self.obj_filename
 
     def get_hyperparameters(self):
@@ -329,7 +329,7 @@ def __call__(
         Callable that Optuna will use to optimize the hyper-parameters
 
         Args:
-            obj_filename: the pickle-exported Algo object.
+            obj_filename: the serialized Algo object.
             output_folder: the root path of the algorithms templates.
             template_path: the algorithm_template. It must contain algo.py in the follow path:
                 ``{algorithm_templates_dir}/{network}/scripts/algo.py``
@@ -364,7 +364,7 @@ def generate(self, output_folder: str = ".") -> None:
         task_id = self.get_task_id()
         task_prefix = os.path.basename(self.algo.get_output_path())
         write_path = os.path.join(output_folder, task_prefix + task_id)
-        self.obj_filename = os.path.join(write_path, "algo_object.pkl")
+        self.obj_filename = os.path.join(write_path, "algo_object.json")
 
         if isinstance(self.algo, BundleAlgo):
             self.algo.export_to_disk(output_folder, task_prefix + task_id, fill_with_datastats=False)
@@ -377,15 +377,15 @@ def run_algo(self, obj_filename: str, output_folder: str = ".", template_path: P
         The python interface for NNI to run.
 
         Args:
-            obj_filename: the pickle-exported Algo object.
+            obj_filename: the serialized Algo object.
             output_folder: the root path of the algorithms templates.
             template_path: the algorithm_template. It must contain algo.py in the follow path:
                 ``{algorithm_templates_dir}/{network}/scripts/algo.py``
         """
         if not os.path.isfile(obj_filename):
             raise ValueError(f"{obj_filename} is not found")
 
-        self.algo, algo_meta_data = algo_from_pickle(obj_filename, template_path=template_path)
+        self.algo, algo_meta_data = algo_from_json(obj_filename, template_path=template_path)
 
         # step 1 sample hyperparams
         params = self.get_hyperparameters()
@@ -397,5 +397,5 @@ def run_algo(self, obj_filename: str, output_folder: str = ".", template_path: P
         # step 4 report validation acc to controller
         acc = self.algo.get_score()
         algo_meta_data = {str(AlgoKeys.SCORE): acc}
-        algo_to_pickle(self.algo, template_path=self.algo.template_path, **algo_meta_data)
+        algo_to_json(self.algo, template_path=self.algo.template_path, **algo_meta_data)
         self.set_score(acc)

monai/apps/auto3dseg/utils.py

@@ -14,7 +14,7 @@
 import os
 
 from monai.apps.auto3dseg.bundle_gen import BundleAlgo
-from monai.auto3dseg import algo_from_pickle, algo_to_pickle
+from monai.auto3dseg import algo_from_json, algo_to_json
 from monai.utils.enums import AlgoKeys
 
 __all__ = ["import_bundle_algo_history", "export_bundle_algo_history", "get_name_from_algo_id"]
@@ -42,22 +42,29 @@ def import_bundle_algo_history(
         if not os.path.isdir(write_path):
             continue
 
-        obj_filename = os.path.join(write_path, "algo_object.pkl")
-        if not os.path.isfile(obj_filename):  # saved mode pkl
+        # Prefer JSON format, fall back to legacy pickle
+        json_filename = os.path.join(write_path, "algo_object.json")
+        pkl_filename = os.path.join(write_path, "algo_object.pkl")
+
+        if os.path.isfile(json_filename):
+            obj_filename = json_filename
+        elif os.path.isfile(pkl_filename):
+            obj_filename = pkl_filename
+        else:
             continue
 
-        algo, algo_meta_data = algo_from_pickle(obj_filename, template_path=template_path)
+        algo, algo_meta_data = algo_from_json(obj_filename, template_path=template_path)
 
         best_metric = algo_meta_data.get(AlgoKeys.SCORE, None)
         if best_metric is None:
             try:
                 best_metric = algo.get_score()
-            except BaseException:
+            except Exception:
                 pass
 
         is_trained = best_metric is not None
 
-        if (only_trained and is_trained) or not only_trained:
+        if is_trained or not only_trained:
             history.append(
                 {AlgoKeys.ID: name, AlgoKeys.ALGO: algo, AlgoKeys.SCORE: best_metric, AlgoKeys.IS_TRAINED: is_trained}
             )
@@ -67,14 +74,14 @@ def import_bundle_algo_history(
 
 def export_bundle_algo_history(history: list[dict[str, BundleAlgo]]) -> None:
     """
-    Save all the BundleAlgo in the history to algo_object.pkl in each individual folder
+    Save all the BundleAlgo in the history to algo_object.json in each individual folder.
 
     Args:
         history: a List of Bundle. Typically, the history can be obtained from BundleGen get_history method
     """
     for algo_dict in history:
         algo = algo_dict[AlgoKeys.ALGO]
-        algo_to_pickle(algo, template_path=algo.template_path)
+        algo_to_json(algo, template_path=algo.template_path)
 
 
 def get_name_from_algo_id(id: str) -> str:

monai/apps/detection/metrics/coco.py

@@ -541,7 +541,7 @@ def _compute_stats_single_threshold(
         for save_idx, array_index in enumerate(inds):
             precision[save_idx] = pr[array_index]
             th_scores[save_idx] = dt_scores_sorted[array_index]
-    except BaseException:
+    except Exception:
         pass
 
     return recall, np.array(precision), np.array(th_scores)

monai/apps/nnunet/nnunetv2_runner.py

@@ -200,7 +200,7 @@ def __init__(
             from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name
 
             self.dataset_name = maybe_convert_to_dataset_name(int(self.dataset_name_or_id))
-        except BaseException:
+        except Exception:
             logger.warning(
                 f"Dataset with name/ID: {self.dataset_name_or_id} cannot be found in the record. "
                 "Please ignore the message above if you are running the pipeline from a fresh start. "
@@ -278,7 +278,7 @@ def convert_dataset(self):
                 num_input_channels=num_input_channels,
                 output_datafolder=raw_data_foldername,
             )
-        except BaseException as err:
+        except Exception as err:
             logger.warning(f"Input config may be incorrect. Detail info: error/exception message is:\n {err}")
             return
 

monai/auto3dseg/__init__.py

@@ -25,7 +25,9 @@
 from .operations import Operations, SampleOperations, SummaryOperations
 from .seg_summarizer import SegSummarizer
 from .utils import (
+    algo_from_json,
     algo_from_pickle,
+    algo_to_json,
     algo_to_pickle,
     concat_multikeys_to_dict,
     concat_val_to_np,