Replace pickle with JSON in Auto3DSeg algo serialization

[aymuos15]

Summary

• Replace algo_to_pickle() with algo_to_json() using compact JSON + MONAI's existing _target_ ConfigParser pattern for truly pickle-free algo metadata serialization
• Add _make_json_serializable() helper to handle numpy arrays, tensors, Path objects
• Backward compatible: algo_from_json() can still load legacy .pkl files (with deprecation warning)
• Keep algo_to_pickle() / algo_from_pickle() as deprecated aliases

Note: Model weights still use torch.save separately—this PR focuses on the algo object serialization.

Test plan

• Linting passes (./runtests.sh --codeformat, ./runtests.sh --ruff)
• Unit tests for _make_json_serializable() and _add_path_with_parent() helpers
• Integration test with tests/apps/test_auto3dseg_bundlegen.py (requires optional deps)

Fixes #8586

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The PR replaces pickle-based Algo serialization in Auto3DSeg with JSON-based serialization by adding algo_to_json and algo_from_json plus helpers (_make_json_serializable, template path resolution, and a legacy pickle loader). Deprecated pickle wrappers delegate to the JSON functions. import/export of bundle algorithm history now prefers algo_object.json with fallback to algo_object.pkl. BundleAlgo and Algo gained state_dict/load_state_dict hooks. Tests for JSON helpers and a license header in tests were added. A minor simplification to history filtering logic was applied.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.71% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title accurately captures the main change: replacing pickle with JSON for Auto3DSeg algo serialization.
Description check	✅ Passed	Description covers objectives, test status, and scope clearly. Most template sections addressed; integration test noted as pending.
Linked Issues check	✅ Passed	Changes directly address #8586: replaces pickle with JSON-based serialization [#8586], maintains backward compatibility with deprecation warnings [#8586], and uses MONAI's ConfigParser pattern [#8586].
Out of Scope Changes check	✅ Passed	All changes align with scope: JSON serialization for algo metadata, backward-compatible pkl loading, helpers for JSON conversion, and state_dict/load_state_dict methods for serialization support. No unrelated modifications detected.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In @monai/auto3dseg/utils.py:
- Around line 474-494: The loop over paths_to_try manipulates sys.path
non-atomically (sys.path.insert/remove) around ConfigParser/
parser.get_parsed_content which can race if concurrent; fix by introducing a
module-level lock (e.g., SYS_PATH_LOCK) and acquire it before any sys.path
modification and release it after cleanup, wrapping the path insert, parser
instantiation (ConfigParser and parser.get_parsed_content) and the removal in a
try/finally so removal always happens and the lock is released; reference the
paths_to_try loop, sys.path.insert/sys.path.remove, ConfigParser,
parser.get_parsed_content, and used_template_path when applying the lock.

🧹 Nitpick comments (10)

monai/apps/auto3dseg/utils.py (1)

67-70: Minor readability nit: condition can be simplified.

The or not only_trained is a double negative. Consider inverting for clarity, though current logic is correct.

monai/auto3dseg/utils.py (9)

281-304: LGTM with minor observation.

Good coverage of common types. The fallback to str(value) on line 304 silently converts unknown types—consider logging a warning for unexpected types to aid debugging.

---

307-312: Missing docstring parameter/return documentation.

Per coding guidelines, docstrings should describe parameters and return values.

Suggested docstring

 def _add_path_with_parent(paths: list[str], path: str | None) -> None:
-    """Add a path and its parent directory to the list if the path is a valid directory."""
+    """
+    Add a path and its parent directory to the list if the path is a valid directory.
+
+    Args:
+        paths: List to append paths to (modified in-place).
+        path: Directory path to add; skipped if None or not a valid directory.
+    """

---

327-333: Hardcoded attribute list may become stale.

Consider defining SERIALIZABLE_ATTRS as a module-level constant or documenting why these specific attributes are serialized. Easier to maintain and extend.

---

346-348: Missing encoding parameter for file write.

Explicit encoding="utf-8" is recommended for cross-platform consistency.

Fix

-    with open(json_filename, "w") as f:
+    with open(json_filename, "w", encoding="utf-8") as f:

---

414-414: Unused kwargs parameter.

Static analysis (ARG001) notes kwargs is unused. Docstring says "reserved for future use"—acceptable, but consider using _ prefix convention or **_kwargs to suppress linter warnings.

---

446-447: Missing encoding parameter for file read.

Same as write—explicit encoding="utf-8" recommended.

Fix

-    with open(filename) as f:
+    with open(filename, encoding="utf-8") as f:

---

449-450: Consider TypeError for type validation.

Static analysis (TRY004) suggests TypeError over ValueError for invalid types. Minor pedantic issue.

---

479-481: Simplify dictionary access.

Static analysis (RUF019) notes unnecessary key check. Use state.get("template_path") instead.

Fix

             algo_config: dict[str, Any] = {"_target_": target}
-            if "template_path" in state and state["template_path"]:
-                algo_config["template_path"] = state["template_path"]
+            if state.get("template_path"):
+                algo_config["template_path"] = state["template_path"]

---

499-502: State restoration silently skips unknown attributes.

If the JSON contains an attribute not present on the algo object, it's silently ignored. This is likely intentional for forward compatibility, but logging a debug message would help troubleshoot version mismatches.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

📥 Commits

Reviewing files that changed from the base of the PR and between 57fdd59 and 23433a7.

📒 Files selected for processing (3)

• monai/apps/auto3dseg/utils.py
• monai/auto3dseg/__init__.py
• monai/auto3dseg/utils.py

🧰 Additional context used

📓 Path-based instructions (1)

**/*.py

⚙️ CodeRabbit configuration file

Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.

Files:

• monai/apps/auto3dseg/utils.py
• monai/auto3dseg/__init__.py
• monai/auto3dseg/utils.py

🧬 Code graph analysis (2)

monai/apps/auto3dseg/utils.py (2)

monai/auto3dseg/utils.py (2)

• algo_from_json (414-510)
• algo_to_json (315-350)

monai/utils/enums.py (1)

• AlgoKeys (687-699)

monai/auto3dseg/__init__.py (1)

monai/auto3dseg/utils.py (3)

• algo_from_json (414-510)
• algo_from_pickle (669-670)
• algo_to_json (315-350)

🪛 Ruff (0.14.10)

monai/auto3dseg/utils.py

384-384: pickle and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue

(S301)

---

391-391: pickle and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue

(S301)

---

414-414: Unused function argument: kwargs

(ARG001)

---

450-450: Prefer TypeError exception for invalid type

(TRY004)

---

450-450: Avoid specifying long messages outside the exception class

(TRY003)

---

455-455: Avoid specifying long messages outside the exception class

(TRY003)

---

480-480: Unnecessary key check before dictionary access

Replace with dict.get

(RUF019)

---

497-497: Avoid specifying long messages outside the exception class

(TRY003)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (19)

• GitHub Check: min-dep-py3 (3.12)
• GitHub Check: min-dep-os (macOS-latest)
• GitHub Check: min-dep-pytorch (2.7.1)
• GitHub Check: min-dep-py3 (3.11)
• GitHub Check: min-dep-pytorch (2.8.0)
• GitHub Check: min-dep-pytorch (2.6.0)
• GitHub Check: min-dep-pytorch (2.5.1)
• GitHub Check: min-dep-py3 (3.10)
• GitHub Check: min-dep-os (ubuntu-latest)
• GitHub Check: min-dep-py3 (3.9)
• GitHub Check: min-dep-os (windows-latest)
• GitHub Check: flake8-py3 (mypy)
• GitHub Check: quick-py3 (windows-latest)
• GitHub Check: quick-py3 (macOS-latest)
• GitHub Check: quick-py3 (ubuntu-latest)
• GitHub Check: build-docs
• GitHub Check: packaging
• GitHub Check: flake8-py3 (pytype)
• GitHub Check: flake8-py3 (codeformat)

🔇 Additional comments (7)

monai/apps/auto3dseg/utils.py (3)

17-17: LGTM!

Import updated to use new JSON-based functions.

---

45-56: LGTM!

Clean preference logic: JSON first, pickle fallback. Good migration strategy.

---

75-84: LGTM!

Docstring and function body updated to use JSON serialization.

monai/auto3dseg/utils.py (3)

353-411: Legacy pickle loader retained for backward compatibility.

Static analysis flags pickle security (S301). Acceptable here since it's deprecated legacy support with clear warnings. No action required beyond eventual removal.

---

504-504: used_template_path overwrites state-restored template_path.

Line 501-502 restores template_path from state, but line 504 immediately overwrites it with used_template_path (which could be different). Verify this is the intended behavior.

---

662-669: LGTM!

Clean deprecated aliases delegating to new JSON functions.

monai/auto3dseg/__init__.py (1)

27-39: LGTM!

New JSON functions properly exported alongside deprecated pickle aliases.

[coderabbitai]

Actionable comments posted: 4

🤖 Fix all issues with AI agents

In `@monai/auto3dseg/utils.py`:
- Around line 482-501: The current loop that tries paths_to_try inserts a path
into sys.path but only removes it on ImportError; change the logic in the for
loop around the insertion and parser.get_parsed_content() to guarantee cleanup
by using a try/finally: set a local inserted flag when you do sys.path.insert(0,
path), perform parser = ConfigParser(algo_config) and algo =
parser.get_parsed_content() in the try block, and in the finally block remove
the path from sys.path if inserted (and set used_template_path only after
success); ensure exceptions are either re-raised or handled the same way as the
previous ImportError branch so sys.path is never left polluted.
- Around line 395-407: The loop over template_paths_candidates currently inserts
p into sys.path but only removes it on ModuleNotFoundError or after successful
load; to guarantee cleanup wrap the insertion and pickle.loads(algo_bytes) in a
try/finally: set an inserted flag when you add p to sys.path, call
pickle.loads(algo_bytes) in the try block, and in the finally block remove p
from sys.path if inserted; keep the existing ModuleNotFoundError handling logic
but avoid leaving sys.path polluted for any other exception and only break the
loop after a successful load (i.e., after pickle.loads completes).
- Around line 315-358: The algo_meta_data passed into algo_to_json may contain
non-JSON-native types (e.g., numpy types) and must be converted before dumping;
update algo_to_json to apply _make_json_serializable to each value in
algo_meta_data (similar to how state is built) and use the sanitized values when
constructing the data dict (keep the existing keys: "_target_", "_state_",
"template_path", plus the serialized algo_meta_data). Ensure you do not change
the keys, only transform the values via _make_json_serializable before creating
json_filename and calling json.dump.
- Around line 281-313: The dict branch in _make_json_serializable fails when
keys are non-strings; change the dict comprehension to coerce keys to strings
(e.g., {str(k): _make_json_serializable(v) for k, v in value.items()}) so
json.dump() won't error on tuple/Path keys, and update both
_make_json_serializable and _add_path_with_parent docstrings to full
Google-style sections (Args, Returns, and Raises if applicable) describing
parameter types, return values, and any exceptions or side effects; keep the
existing behavior for arrays, numpy types, torch.Tensor, Path-like objects, and
the fallback to str(value).

🧹 Nitpick comments (1)

monai/auto3dseg/utils.py (1)

670-678: Add docstrings for deprecated alias functions.

These aliases lack docstrings. Add short Google-style docstrings for API completeness. As per coding guidelines, please add docstrings with Args/Returns.

💡 Proposed fix

 `@deprecated`(since="1.6", msg_suffix="Use algo_to_json instead.")
 def algo_to_pickle(algo: Algo, template_path: PathLike | None = None, **algo_meta_data: Any) -> str:
+    """Deprecated alias for algo_to_json."""
     return algo_to_json(algo, template_path, **algo_meta_data)

@@
 `@deprecated`(since="1.6", msg_suffix="Use algo_from_json instead.")
 def algo_from_pickle(filename: str, template_path: PathLike | None = None, **kwargs: Any) -> Any:
+    """Deprecated alias for algo_from_json."""
     return algo_from_json(filename, template_path, **kwargs)

[ericspod]

Hi @aymuos15 thanks for the work on this, I think it's going in the right direction. I would like others who have worked on Auto3DSeg to have a look as well, and I have a number of comments below. We're working on a CI issue still but will come back to this one soon.

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@monai/auto3dseg/utils.py`:
- Around line 315-357: The function algo_to_json contains an indentation error
where the line assigning algo_meta_data = {str(k): _make_json_serializable(v)
for k,v in algo_meta_data.items()} is unindented and breaks the function; fix it
by indenting that comprehension so it is inside algo_to_json (aligned with the
preceding statements that build state/target), then ensure the subsequent data
dict uses that local algo_meta_data (so data: dict[str, Any] = {...,
**algo_meta_data}) and the function returns json_filename as before; look for
symbols algo_to_json, algo_meta_data, state, target, and data to make the
correction.
- Around line 421-519: Add unit tests for algo serialization: create tests that
(1) serialize a simple Algo instance using algo_to_json, then read it with
algo_from_json and assert the returned algo object and metadata match
expectations and that algo.template_path and algo.output_path get set/overridden
correctly, (2) exercise the legacy .pkl path by writing a pickle file and
calling algo_from_json on a .pkl filename while asserting a FutureWarning is
raised and that _load_legacy_pickle is used to return the legacy object, and (3)
cover cases where template_path comes from file data and when
algorithm_templates sibling folder is used; use temporary directories/files
(tmp_path), pytest.warns for the deprecation, and reference the functions
algo_to_json, algo_from_json, and _load_legacy_pickle in the tests to locate the
code under test.

🧹 Nitpick comments (1)

monai/auto3dseg/utils.py (1)

360-418: Add full Google‑style docstring to _load_legacy_pickle.

Include Args/Returns/Raises to satisfy docstring requirements.

As per coding guidelines: "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."

[coderabbitai]

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

monai/apps/auto3dseg/bundle_gen.py (1)

39-40: ⚠️ Potential issue | 🟡 Minor

Use algo_to_json instead of deprecated algo_to_pickle.

This file still imports and uses the deprecated function at line 695.

-from monai.auto3dseg.utils import (
-    ...
-    algo_to_pickle,
-)
+from monai.auto3dseg.utils import (
+    ...
+    algo_to_json,
+)

And at line 695:

-                algo_to_pickle(gen_algo, template_path=algo.template_path)
+                algo_to_json(gen_algo, template_path=algo.template_path)

🧹 Nitpick comments (3)

monai/auto3dseg/utils.py (2)

339-341: Specify explicit UTF-8 encoding for cross-platform consistency.

-    with open(json_filename, "w") as f:
+    with open(json_filename, "w", encoding="utf-8") as f:

---

439-440: Specify explicit UTF-8 encoding.

-    with open(filename) as f:
+    with open(filename, encoding="utf-8") as f:

tests/auto3dseg/test_json_serialization.py (1)

24-53: Add test case for Path objects.

The _make_json_serializable docstring mentions Path handling, but there's no test for it.

Suggested addition

from pathlib import Path

def test_path(self) -> None:
    p = Path("/some/path")
    assert _make_json_serializable(p) == "/some/path"

[ericspod]

Hi @daguangxu this is one proposed solution for removing pickle from Auto3DSeg.

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (1)

monai/auto3dseg/utils.py (1)

467-491: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Don’t hide the real import failure here.

Catching broad ImportError and then ending with a generic ValueError drops the actual reason ConfigParser failed to import or construct the class. Only the path-probing case should be retried; otherwise preserve the last exception as the cause.

Suggested fix

     paths_to_try: list[str | None] = list(template_paths) if template_paths else [None]
     algo = None
     used_template_path: str | None = None
+    last_error: Exception | None = None
     for path in paths_to_try:
         path_added = False
         try:
             if path and path not in sys.path:
                 sys.path.insert(0, path)
@@
             parser = ConfigParser(algo_config)
             algo = parser.get_parsed_content()
             used_template_path = path
             break
-        except ImportError as e:
+        except ModuleNotFoundError as e:
+            last_error = e
             logging.debug(f"Failed to instantiate {target} with path {path}: {e}")
             continue
         finally:
             if path_added and path in sys.path:
                 sys.path.remove(path)
 
     if algo is None:
-        raise ValueError(f"Failed to instantiate Algo from target '{target}' with paths {template_paths}")
+        raise ValueError(f"Failed to instantiate Algo from target '{target}' with paths {template_paths}") from last_error

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/auto3dseg/utils.py` around lines 467 - 491, The current loop swallows
the real ImportError; change it to record the last caught ImportError (e.g.,
last_exception) inside the except ImportError block when probing paths
(paths_to_try) and continue retrying, and after the loop, if algo is still None
raise the ValueError using "raise ValueError(... ) from last_exception" so the
original ImportError is preserved as the cause; keep the existing path
insertion/removal logic, references: paths_to_try, path_added, ConfigParser,
parser.get_parsed_content(), algo, and template_paths.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@monai/auto3dseg/utils.py`:
- Around line 331-336: The writer currently stores "template_path" as None when
the template_path argument is omitted, which loses algo.template_path and breaks
algo_from_json's sys.path seeding; change the value written for "template_path"
in the data dict to use the provided template_path if present, otherwise fall
back to the Algo object's template_path (e.g., use getattr(algo,
"template_path", None) converted to str when present) so that algo.template_path
is persisted; update the construction of data (the "template_path" entry)
accordingly where algo_meta_data, target, state, template_path and algo are
used.

---

Duplicate comments:
In `@monai/auto3dseg/utils.py`:
- Around line 467-491: The current loop swallows the real ImportError; change it
to record the last caught ImportError (e.g., last_exception) inside the except
ImportError block when probing paths (paths_to_try) and continue retrying, and
after the loop, if algo is still None raise the ValueError using "raise
ValueError(... ) from last_exception" so the original ImportError is preserved
as the cause; keep the existing path insertion/removal logic, references:
paths_to_try, path_added, ConfigParser, parser.get_parsed_content(), algo, and
template_paths.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: aa51d91f-f83f-4656-afab-3d1594418a03

📥 Commits

Reviewing files that changed from the base of the PR and between 3bc5e2d and e4dbc6d.

📒 Files selected for processing (3)

• monai/apps/auto3dseg/bundle_gen.py
• monai/auto3dseg/algo_gen.py
• monai/auto3dseg/utils.py

✅ Files skipped from review due to trivial changes (1)

• monai/apps/auto3dseg/bundle_gen.py

[aymuos15]

@ericspod

I've changed it so people can choose to keep using pickle through the env var MONAI_ALLOW_PICKLE=1, but it's also behind a deprecation (FutureWarning). Should we remove the deprecation and keep pickle as a permanent opt-in instead?

Ive also addressed all the previous comments.

[ericspod]

I had some minor changes again.

[ericspod]

Looks good to me, this addresses the security concern but we've left in the ability to use pickle for those that need it. We'll leave the deprecation warnings too as they are appropriate.