From 2e5e444302d8378a4b46ada34c2e45f4c416db58 Mon Sep 17 00:00:00 2001 From: Zeeshan Modi <92383127+Zeesejo@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:38:52 +0200 Subject: [PATCH] fix(ZarrAvgMerger): retain independent temporary stores Keep separate temporary directory references for the default Zarr v3 value and count stores so neither backing directory is collected early. Restore readable codec conditionals and treat empty Zarr v2 codec sequences as no compression. Add regression coverage for distinct live temporary directories, store read/write and finalization, and empty list/tuple codec arguments. Fixes #8476 Signed-off-by: Zeeshan Modi <92383127+Zeesejo@users.noreply.github.com> --- monai/inferers/merger.py | 41 +++++++++++---------- tests/inferers/test_zarr_avg_merger.py | 51 ++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 19 deletions(-) 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)