Merge branch 'further_cicd_cleanup' of github.com:ericspod/MONAI into further_cicd_cleanup

Author: ericspod

monai/data/dataset.py

@@ -609,24 +609,26 @@ def __init__(
         # the cache is created without multi-threading
         self._read_env: Any | None = None
         # this runs on the primary thread/process
-        self._fill_cache_start_reader(show_progress=self.progress)
-        print(f"Accessing lmdb file: {self.db_file.absolute()}.")
+        read_env = self._fill_cache_start_reader(show_progress=self.progress)
+        read_env.close()
 
     def set_data(self, data: Sequence):
         """
         Set the input data and delete all the out-dated cache content.
-
         """
+        self.close()
         super().set_data(data=data)
         self._read_env = self._fill_cache_start_reader(show_progress=self.progress)
 
     def _safe_serialize(self, val):
+        """Serialize the tensor/array `val` using the pickle protocol, and return its bytes object."""
         out = BytesIO()
         torch.save(convert_to_tensor(val), out, pickle_protocol=self.pickle_protocol)
         out.seek(0)
         return out.read()
 
     def _safe_deserialize(self, val):
+        """Load the object from the given bytes data, this must be loadable as weights only using `torch.load`."""
         return torch.load(BytesIO(val), map_location="cpu", weights_only=True)
 
     def _fill_cache_start_reader(self, show_progress=True):
@@ -687,10 +689,7 @@ def _fill_cache_start_reader(self, show_progress=True):
         return lmdb.open(path=f"{self.db_file}", subdir=False, **self.lmdb_kwargs)
 
     def _cachecheck(self, item_transformed):
-        """
-        if the item is not found in the lmdb file, resolves to the persistent cache default behaviour.
-
-        """
+        """If the item is not found in the lmdb file, resolves to the persistent cache default behaviour."""
         if self._read_env is None:
             # this runs on multiple processes, each one should have its own env.
             self._read_env = self._fill_cache_start_reader(show_progress=False)
@@ -715,7 +714,14 @@ def info(self):
         out = dict(self._read_env.info())
         out["size"] = len(self.data)
         out["filename"] = f"{self.db_file.absolute()}"
+        self.close()
         return out
+    
+    def close(self):
+        """Close the read environment and set it to None, if it exists."""
+        if self._read_env:
+            self._read_env.close()
+            self._read_env = None
 
 
 class CacheDataset(Dataset):

tests/data/test_lmdbdataset.py

@@ -19,7 +19,7 @@
 import numpy as np
 from parameterized import parameterized
 
-from monai.data import LMDBDataset, json_hashing
+from monai.data import LMDBDataset, json_hashing, pickle_hashing
 from monai.transforms import Compose, LoadImaged, SimulateDelayd, Transform
 from tests.test_utils import skip_if_windows
 
@@ -76,7 +76,7 @@
         SimulateDelayd(keys=["image", "label", "extra"], delay_time=[1e-7, 1e-6, 1e-5]),
     ],
     (128, 128, 128),
-    {"db_name": "testdb", "lmdb_kwargs": {"map_size": 2 * 1024**2}},
+    {"db_name": "testdb", "lmdb_kwargs": {"map_size": 50 * 1024**2}},
 ]
 
 
@@ -91,36 +91,16 @@ def __call__(self, data):
 
 @skip_if_windows
 class TestLMDBDataset(unittest.TestCase):
-    def test_cache(self):
-        """testing no inplace change to the hashed item"""
+    @parameterized.expand([(pickle_hashing,), (json_hashing,)])
+    def test_cache(self, hash_func):
+        """Testing no inplace change to the hashed item."""
         items = [[list(range(i))] for i in range(5)]
 
         with tempfile.TemporaryDirectory() as tempdir:
-            ds = LMDBDataset(items, transform=_InplaceXform(), cache_dir=tempdir, lmdb_kwargs={"map_size": 10 * 1024})
-            self.assertEqual(items, [[[]], [[0]], [[0, 1]], [[0, 1, 2]], [[0, 1, 2, 3]]])
-            ds1 = LMDBDataset(items, transform=_InplaceXform(), cache_dir=tempdir, lmdb_kwargs={"map_size": 10 * 1024})
-            self.assertEqual(list(ds1), list(ds))
-            self.assertEqual(items, [[[]], [[0]], [[0, 1]], [[0, 1, 2]], [[0, 1, 2, 3]]])
-
-            ds = LMDBDataset(
-                items,
-                transform=_InplaceXform(),
-                cache_dir=tempdir,
-                lmdb_kwargs={"map_size": 10 * 1024},
-                hash_func=json_hashing,
-            )
-            self.assertEqual(items, [[[]], [[0]], [[0, 1]], [[0, 1, 2]], [[0, 1, 2, 3]]])
-            ds1 = LMDBDataset(
-                items,
-                transform=_InplaceXform(),
-                cache_dir=tempdir,
-                lmdb_kwargs={"map_size": 10 * 1024},
-                hash_func=json_hashing,
-            )
-            self.assertEqual(list(ds1), list(ds))
-            self.assertEqual(items, [[[]], [[0]], [[0, 1]], [[0, 1, 2]], [[0, 1, 2, 3]]])
-
-        self.assertTrue(isinstance(ds1.info(), dict))
+            ds = LMDBDataset(items, transform=_InplaceXform(), cache_dir=tempdir, hash_func=hash_func, progress=False)
+            self.assertEqual(items, [[[]], [[0]], [[0, 1]], [[0, 1, 2]], [[0, 1, 2, 3]]], "Input dataset mutated.")
+            self.assertNotEqual(items, list(ds), "Output data unmodified by transform.")
+            self.assertIsInstance(ds.info(), dict)
 
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7])
     def test_shape(self, transform, expected_shape, kwargs=None):
@@ -146,16 +126,30 @@ def test_shape(self, transform, expected_shape, kwargs=None):
                 },
             ]
 
+            # update the data to cache
+            test_data_new = [
+                {
+                    "image": os.path.join(tempdir, "test_image1_new.nii.gz"),
+                    "label": os.path.join(tempdir, "test_label1_new.nii.gz"),
+                    "extra": os.path.join(tempdir, "test_extra1_new.nii.gz"),
+                },
+                {
+                    "image": os.path.join(tempdir, "test_image2_new.nii.gz"),
+                    "label": os.path.join(tempdir, "test_label2_new.nii.gz"),
+                    "extra": os.path.join(tempdir, "test_extra2_new.nii.gz"),
+                },
+            ]
+
             cache_dir = os.path.join(os.path.join(tempdir, "cache"), "data")
-            dataset_precached = LMDBDataset(
-                data=test_data, transform=transform, progress=False, cache_dir=cache_dir, **kwargs
-            )
+            ds_args = dict(data=test_data, transform=transform, progress=False, cache_dir=cache_dir, **kwargs)
+
+            dataset_precached = LMDBDataset(**ds_args)
             data1_precached = dataset_precached[0]
             data2_precached = dataset_precached[1]
 
-            dataset_postcached = LMDBDataset(
-                data=test_data, transform=transform, progress=False, cache_dir=cache_dir, **kwargs
-            )
+            dataset_precached.close()
+
+            dataset_postcached = LMDBDataset(**ds_args)
             data1_postcached = dataset_postcached[0]
             data2_postcached = dataset_postcached[1]
 
@@ -179,19 +173,6 @@ def test_shape(self, transform, expected_shape, kwargs=None):
                 self.assertTupleEqual(data2_postcached["label"].shape, expected_shape)
                 self.assertTupleEqual(data2_postcached["extra"].shape, expected_shape)
 
-            # update the data to cache
-            test_data_new = [
-                {
-                    "image": os.path.join(tempdir, "test_image1_new.nii.gz"),
-                    "label": os.path.join(tempdir, "test_label1_new.nii.gz"),
-                    "extra": os.path.join(tempdir, "test_extra1_new.nii.gz"),
-                },
-                {
-                    "image": os.path.join(tempdir, "test_image2_new.nii.gz"),
-                    "label": os.path.join(tempdir, "test_label2_new.nii.gz"),
-                    "extra": os.path.join(tempdir, "test_extra2_new.nii.gz"),
-                },
-            ]
             # test new exchanged cache content
             if transform is None:
                 dataset_postcached.set_data(data=test_data_new)