Method _fill_cache_start_reader

monai/data/dataset.py:634–695

Check the LMDB cache and write the cache if needed. py-lmdb doesn't have a good support for concurrent write. This method can be used with multiple processes, but it may have a negative impact on the performance. Args: show_progress: whether to show the progress

(self, show_progress=True)

        return torch.load(BytesIO(val), map_location="cpu", weights_only=True)

    def _fill_cache_start_reader(self, show_progress=True):
        """
        Check the LMDB cache and write the cache if needed. py-lmdb doesn't have a good support for concurrent write.
        This method can be used with multiple processes, but it may have a negative impact on the performance.

        Args:
            show_progress: whether to show the progress bar if possible.
        """
        # Close any open read environment before attempting write-mode access
        # to prevent "environment already open" errors when multiple LMDBDataset
        # instances target the same db file
        if self._read_env is not None:
            self._read_env.close()
            self._read_env = None
        # create cache
        self.lmdb_kwargs["readonly"] = False
        env = lmdb.open(path=f"{self.db_file}", subdir=False, **self.lmdb_kwargs)
        if show_progress and not has_tqdm:
            warnings.warn("LMDBDataset: tqdm is not installed. not displaying the caching progress.")
        with env.begin(write=False) as search_txn:
            for item in tqdm(self.data) if has_tqdm and show_progress else self.data:
                key = self.hash_func(item)
                done, retry, val = False, 5, None
                while not done and retry > 0:
                    try:
                        with search_txn.cursor() as cursor:
                            done = cursor.set_key(key)
                        if done:
                            continue
                        if val is None:
                            val = self._pre_transform(deepcopy(item))  # keep the original hashed
                            # val = pickle.dumps(val, protocol=self.pickle_protocol)
                            val = self._safe_serialize(val)
                        with env.begin(write=True) as txn:
                            txn.put(key, val)
                        done = True
                    except lmdb.MapFullError:
                        done, retry = False, retry - 1
                        size = env.info()["map_size"]
                        new_size = size * 2
                        warnings.warn(
                            f"Resizing the cache database from {int(size) >> 20}MB to {int(new_size) >> 20}MB."
                        )
                        env.set_mapsize(new_size)
                    except lmdb.MapResizedError:
                        # the mapsize is increased by another process
                        # set_mapsize with a size of 0 to adopt the new size
                        env.set_mapsize(0)
                if not done:  # still has the map full error
                    size = env.info()["map_size"]
                    env.close()
                    raise ValueError(f"LMDB map size reached, increase size above current size of {size}.")
        size = env.info()["map_size"]
        env.close()
        # read-only database env
        self.lmdb_kwargs["readonly"] = True
        self.lmdb_kwargs["map_size"] = size
        if self.lmdb_kwargs.get("lock", None) is None: