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

monai/inferers/merger.py

@@ -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
@@ -322,7 +323,6 @@ 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:
@@ -351,19 +351,22 @@ def __init__(
                 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
+            if isinstance(codecs, (list, tuple)):
+                self.codecs = codecs[0] if codecs else None
+            elif codecs is not None:
+                self.codecs = 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
+            if isinstance(value_codecs, (list, tuple)):
+                self.value_codecs = value_codecs[0] if value_codecs else None
+            elif value_codecs is not None:
+                self.value_codecs = 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
+            if isinstance(count_codecs, (list, tuple)):
+                self.count_codecs = count_codecs[0] if count_codecs else None
+            elif count_codecs is not None:
+                self.count_codecs = count_codecs
             else:
                 self.count_codecs = count_compressor  # type: ignore[assignment]
 

tests/inferers/test_zarr_avg_merger.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import gc
 import os
 import unittest
 from tempfile import TemporaryDirectory
@@ -439,6 +440,56 @@ def test_zarr_avg_merger_patches(self, arguments, patch_locations, expected):
         # check if the result is matching the expectation
         assert_allclose(output[:], expected.numpy())
 
+    @unittest.skipUnless(version_geq(get_package_version("zarr"), "3.0.0"), "Requires zarr v3")
+    def test_temporary_stores_remain_alive(self):
+        """Default stores must retain distinct backing directories after garbage collection."""
+        merger = ZarrAvgMerger(merged_shape=(1, 1, 2, 2), store=self.merged_name)
+        self.assertIsNotNone(merger.value_tmpdir)
+        self.assertIsNotNone(merger.count_tmpdir)
+        self.addCleanup(merger.value_tmpdir.cleanup)
+        self.addCleanup(merger.count_tmpdir.cleanup)
+        gc.collect()
+        self.assertNotEqual(merger.value_tmpdir.name, merger.count_tmpdir.name)
+        self.assertTrue(os.path.isdir(merger.value_tmpdir.name))
+        self.assertTrue(os.path.isdir(merger.count_tmpdir.name))
+
+    @unittest.skipUnless(version_geq(get_package_version("zarr"), "3.0.0"), "Requires zarr v3")
+    def test_temporary_stores_read_write(self):
+        """Both default stores must remain usable through reopening and finalization."""
+        merger = ZarrAvgMerger(merged_shape=(1, 1, 2, 2), store=self.merged_name)
+        for name in ("value_tmpdir", "count_tmpdir"):
+            tmpdir = getattr(merger, name, None)
+            if tmpdir is not None:
+                self.addCleanup(tmpdir.cleanup)
+        gc.collect()
+        values = zarr.open_array(store=merger.value_store, mode="r+")
+        counts = zarr.open_array(store=merger.count_store, mode="r+")
+        values[:] = 6
+        counts[:] = 2
+        assert_allclose(merger.get_values()[:], np.full((1, 1, 2, 2), 6))
+        assert_allclose(merger.get_counts()[:], np.full((1, 1, 2, 2), 2))
+        assert_allclose(merger.finalize()[:], np.full((1, 1, 2, 2), 3))
+
+    @parameterized.expand([("list", []), ("tuple", ())])
+    @unittest.skipIf(version_geq(get_package_version("zarr"), "3.0.0"), "Requires zarr v2")
+    def test_empty_codec_sequences(self, _, empty_codecs):
+        """Explicit empty codec sequences disable compression without indexing errors."""
+        merger = ZarrAvgMerger(
+            merged_shape=(1, 1, 2, 2),
+            store=self.merged_name,
+            value_store=zarr.storage.MemoryStore(),
+            count_store=zarr.storage.MemoryStore(),
+            codecs=empty_codecs,
+            value_codecs=empty_codecs,
+            count_codecs=empty_codecs,
+        )
+        self.assertIsNone(merger.output.compressor)
+        self.assertIsNone(merger.values.compressor)
+        self.assertIsNone(merger.counts.compressor)
+        patch = torch.ones((1, 1, 2, 2))
+        merger.aggregate(patch, (0, 0))
+        assert_allclose(merger.finalize()[:], patch.numpy())
+
     def test_zarr_avg_merger_finalized_error(self):
         with self.assertRaises(ValueError):
             merger = ZarrAvgMerger(merged_shape=(1, 3, 2, 3), store=self.merged_name)