Replace pickle with JSON in Auto3DSeg algo serialization

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
@@ -367,6 +367,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 = (
@@ -659,7 +692,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/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,11 +42,18 @@ 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:
@@ -57,7 +64,7 @@ def import_bundle_algo_history(
 
         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/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,

monai/auto3dseg/algo_gen.py

@@ -38,6 +38,30 @@ def get_score(self, *args, **kwargs):
     def get_output_path(self, *args, **kwargs):
         """Returns the algo output paths for scripts location"""
 
+    def state_dict(self) -> dict:
+        """
+        Return state for serialization.
+
+        Subclasses should override this method to return a dictionary of
+        attributes that need to be serialized. This follows the PyTorch
+        convention for state management.
+
+        Returns:
+            A dictionary containing the state to serialize.
+        """
+        return {}
+
+    def load_state_dict(self, state: dict) -> None:
+        """
+        Restore state from a dictionary.
+
+        Subclasses should override this method to restore their state
+        from the dictionary returned by state_dict().
+
+        Args:
+            state: A dictionary containing the state to restore.
+        """
+
 
 class AlgoGen(Randomizable):
     """