Add state_dict/load_state_dict to Algo classes and fix CI issues

- Add state_dict() and load_state_dict() methods to Algo base class
- Override in BundleAlgo with serializable attributes
- Update algo_to_json/algo_from_json to use state_dict pattern
- Fix Black formatting in utils.py
- Fix Windows path test for cross-platform compatibility

Addresses reviewer feedback from @ericspod to follow PyTorch conventions.

Author: aymuos15

monai/apps/auto3dseg/bundle_gen.py

@@ -367,6 +367,37 @@ 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.
+        """
+        return {
+            "template_path": self.template_path,
+            "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():
+            if hasattr(self, key):
+                setattr(self, key, value)
+
 
 # path to download the algo_templates
 default_algo_zip = (

monai/auto3dseg/algo_gen.py

@@ -43,6 +43,31 @@ def get_output_path(self, *args, **kwargs):
         """Returns the algo output paths for scripts location"""
         pass
 
+    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.
+        """
+        pass
+
 
 class AlgoGen(Randomizable):
     """

monai/auto3dseg/utils.py

@@ -322,19 +322,7 @@ def algo_to_json(algo: Algo, template_path: PathLike | None = None, **algo_meta_
     Returns:
         Filename of the saved Algo object (algo_object.json).
     """
-    attrs = [
-        "template_path",
-        "data_stats_files",
-        "data_list_file",
-        "mlflow_tracking_uri",
-        "mlflow_experiment_name",
-        "output_path",
-        "name",
-        "best_metric",
-        "fill_records",
-        "device_setting"
-    ]
-    state = {a : _make_json_serializable(getattr(algo, a)) for a in attrs if hasattr(algo, a)}
+    state = {k: _make_json_serializable(v) for k, v in algo.state_dict().items()}
 
     # Build target string for dynamic class instantiation
     cls = algo.__class__
@@ -503,9 +491,8 @@ def algo_from_json(filename: str, template_path: PathLike | None = None, **kwarg
         raise ValueError(f"Failed to instantiate Algo from target '{target}' with paths {template_paths}")
 
     # Restore the state (skip template_path as it's set to the working import path below)
-    for attr, value in state.items():
-        if attr != "template_path" and hasattr(algo, attr):
-            setattr(algo, attr, value)
+    state_to_load = {k: v for k, v in state.items() if k != "template_path"}
+    algo.load_state_dict(state_to_load)
 
     # Use the path that successfully imported the class, not the original saved path
     # (the original path may no longer exist if the workdir was moved)

tests/auto3dseg/test_json_serialization.py

@@ -48,7 +48,8 @@ def test_torch_tensor(self) -> None:
 
     def test_path(self) -> None:
         p = Path("/some/path")
-        assert _make_json_serializable(p) == "/some/path"
+        # Use str(p) since path separators differ on Windows vs Unix
+        assert _make_json_serializable(p) == str(p)
 
     def test_fallback(self) -> None:
         class Custom: