Skip to content

fix(ZarrAvgMerger): use separate tmpdir attrs to prevent value_store tmpdir GC on zarr v3 #8816

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Zeesejo wants to merge 1 commit into
base: dev
Choose a base branch
from
+16 −37
Open

fix(ZarrAvgMerger): use separate tmpdir attrs to prevent value_store tmpdir GC on zarr v3 #8816

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 16 additions & 37 deletions monai/inferers/merger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,32 +286,33 @@ def __init__(
self.value_dtype = value_dtype
self.count_dtype = count_dtype
self.store = store
self.tmpdir: TemporaryDirectory | None
# Use separate tmpdir attributes so both TemporaryDirectory objects remain alive
# for the full lifetime of this instance. Previously a single `self.tmpdir` was
# overwritten when both value_store and count_store were None, causing the first
# TemporaryDirectory to be garbage-collected immediately.
self.value_tmpdir: TemporaryDirectory | None = None
self.count_tmpdir: TemporaryDirectory | None = None

# Handle zarr v3 vs older versions
is_zarr_v3 = version_geq(get_package_version("zarr"), "3.0.0")

if is_zarr_v3:
if value_store is None:
self.tmpdir = TemporaryDirectory()
self.value_store = zarr.storage.LocalStore(self.tmpdir.name) # type: ignore
self.value_tmpdir = TemporaryDirectory()
self.value_store = zarr.storage.LocalStore(self.value_tmpdir.name) # type: ignore
else:
self.value_store = value_store # type: ignore
if count_store is None:
self.tmpdir = TemporaryDirectory()
self.count_store = zarr.storage.LocalStore(self.tmpdir.name) # type: ignore
self.count_tmpdir = TemporaryDirectory()
self.count_store = zarr.storage.LocalStore(self.count_tmpdir.name) # type: ignore
else:
self.count_store = count_store # type: ignore
else:
self.tmpdir = None
self.value_store = zarr.storage.TempStore() if value_store is None else value_store # type: ignore
self.count_store = zarr.storage.TempStore() if count_store is None else count_store # type: ignore

self.chunks = chunks

# Handle compressor/codecs based on zarr version
is_zarr_v3 = version_geq(get_package_version("zarr"), "3.0.0")

# Initialize codecs/compressor attributes with proper types
self.codecs: list | None = None
self.value_codecs: list | None = None
Expand All @@ -322,50 +323,28 @@ def __init__(
if codecs is not None:
self.codecs = codecs
elif compressor is not None:
# Convert compressor to codec format
if isinstance(compressor, (list, tuple)):
self.codecs = compressor
else:
self.codecs = [compressor]
self.codecs = compressor if isinstance(compressor, (list, tuple)) else [compressor]
else:
self.codecs = None

if value_codecs is not None:
self.value_codecs = value_codecs
elif value_compressor is not None:
if isinstance(value_compressor, (list, tuple)):
self.value_codecs = value_compressor
else:
self.value_codecs = [value_compressor]
self.value_codecs = value_compressor if isinstance(value_compressor, (list, tuple)) else [value_compressor]
else:
self.value_codecs = None

if count_codecs is not None:
self.count_codecs = count_codecs
elif count_compressor is not None:
if isinstance(count_compressor, (list, tuple)):
self.count_codecs = count_compressor
else:
self.count_codecs = [count_compressor]
self.count_codecs = count_compressor if isinstance(count_compressor, (list, tuple)) else [count_compressor]
else:
self.count_codecs = None
else:
# For zarr v2, use compressors
if codecs is not None:
# If codecs are specified in v2, use the first codec as compressor
self.codecs = codecs[0] if isinstance(codecs, (list, tuple)) else codecs
else:
self.codecs = compressor # type: ignore[assignment]

if value_codecs is not None:
self.value_codecs = value_codecs[0] if isinstance(value_codecs, (list, tuple)) else value_codecs
else:
self.value_codecs = value_compressor # type: ignore[assignment]

if count_codecs is not None:
self.count_codecs = count_codecs[0] if isinstance(count_codecs, (list, tuple)) else count_codecs
else:
self.count_codecs = count_compressor # type: ignore[assignment]
self.codecs = codecs[0] if isinstance(codecs, (list, tuple)) else codecs if codecs is not None else compressor # type: ignore[assignment]
self.value_codecs = value_codecs[0] if isinstance(value_codecs, (list, tuple)) else value_codecs if value_codecs is not None else value_compressor # type: ignore[assignment]
self.count_codecs = count_codecs[0] if isinstance(count_codecs, (list, tuple)) else count_codecs if count_codecs is not None else count_compressor # type: ignore[assignment]
Comment on lines +345 to +347
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Guard empty codec lists in the v2 fallback.

codecs[0], value_codecs[0], and count_codecs[0] will raise IndexError for []. Please treat an empty sequence as None or raise a targeted ValueError, and add a regression test for that path. As per coding guidelines, "Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Ensure new or modified definitions will be covered by existing or new unit tests."

Suggested fix 
+        def _unwrap_v2_codec_arg(arg):
+            if isinstance(arg, (list, tuple)):
+                return arg[0] if arg else None
+            return arg
+
         else:
             # For zarr v2, use compressors
-            self.codecs = codecs[0] if isinstance(codecs, (list, tuple)) else codecs if codecs is not None else compressor  # type: ignore[assignment]
-            self.value_codecs = value_codecs[0] if isinstance(value_codecs, (list, tuple)) else value_codecs if value_codecs is not None else value_compressor  # type: ignore[assignment]
-            self.count_codecs = count_codecs[0] if isinstance(count_codecs, (list, tuple)) else count_codecs if count_codecs is not None else count_compressor  # type: ignore[assignment]
+            self.codecs = _unwrap_v2_codec_arg(codecs) if codecs is not None else compressor  # type: ignore[assignment]
+            self.value_codecs = _unwrap_v2_codec_arg(value_codecs) if value_codecs is not None else value_compressor  # type: ignore[assignment]
+            self.count_codecs = _unwrap_v2_codec_arg(count_codecs) if count_codecs is not None else count_compressor  # type: ignore[assignment]
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@monai/inferers/merger.py` around lines 345 - 347, The current v2 fallback
uses codecs[0], value_codecs[0], and count_codecs[0] which will IndexError on
empty sequences; update the logic in the Merger initializer (the assignment to
self.codecs, self.value_codecs, self.count_codecs) to treat an empty list/tuple
as None (i.e., if isinstance(x, (list,tuple)) and len(x) > 0 then use x[0], elif
x is not None use x, else use the corresponding
compressor/value_compressor/count_compressor fallback), or alternatively raise a
clear ValueError when an explicit empty sequence is provided; also add a
regression unit test that passes empty lists for
codecs/value_codecs/count_codecs and asserts the correct fallback behavior or
the expected ValueError.


# Create zarr arrays with appropriate parameters based on version
if is_zarr_v3:
Expand Down
Loading