Commit 3d2e7d38 authored by Yuxin Wu's avatar Yuxin Wu

fix size count in LMDBData

parent 9af1c032
......@@ -58,7 +58,7 @@ notifications:
recipients:
- ppwwyyxxc@gmail.com
on_success: never
on_failure: change
on_failure: always
- webhooks:
urls:
- https://webhooks.gitter.im/e/cede9dbbf6630b3704b3
......
......@@ -79,8 +79,9 @@ class LMDBData(RNGDataFlow):
self._shuffle = shuffle
self.open_lmdb()
logger.info("Found {} entries in {}".format(self._size, self._lmdb_path))
self._size = self._txn.stat()['entries']
self._set_keys(keys)
logger.info("Found {} entries in {}".format(self._size, self._lmdb_path))
def _set_keys(self, keys=None):
def find_keys(txn, size):
......@@ -94,12 +95,16 @@ class LMDBData(RNGDataFlow):
pbar.update()
return keys
if self._shuffle:
if keys is None:
# get the list of keys either from __keys__ or by iterating
try:
self.keys = loads(self._txn.get('__keys__'))
except Exception:
except:
self.keys = None
else:
self._size -= 1 # delete this item
if self._shuffle:
if keys is None:
if self.keys is None:
self.keys = find_keys(self._txn, self._size)
else:
# check if key-format like '{:0>8d}' was given
......@@ -114,7 +119,6 @@ class LMDBData(RNGDataFlow):
readonly=True, lock=False, readahead=True,
map_size=1099511627776 * 2, max_readers=100)
self._txn = self._lmdb.begin()
self._size = self._txn.stat()['entries']
def reset_state(self):
self._lmdb.close()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment