diff --git a/monai/inferers/merger.py b/monai/inferers/merger.py index 31e9b5d6322..f179ac6a077 100644 --- a/monai/inferers/merger.py +++ b/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] diff --git a/tests/inferers/test_zarr_avg_merger.py b/tests/inferers/test_zarr_avg_merger.py index b4e26d6bebd..d921bc0a7f9 100644 --- a/tests/inferers/test_zarr_avg_merger.py +++ b/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)