diff --git a/exca/cachedict.py b/exca/cachedict.py index 7b6b31c4..c1d10228 100644 --- a/exca/cachedict.py +++ b/exca/cachedict.py @@ -118,8 +118,7 @@ def __init__( self._key_info: dict[str, DumpInfo] = {} # json info file reading self._folder_modified = -1.0 - self._info_files_last: dict[str, int] = {} - self._jsonl_readings = 0 # for perf + self._jsonl_readers: dict[str, JsonlReader] = {} self._jsonl_reading_allowance = float("inf") # keep loaders live for optimized loading # (instances are reinstantiated for dumping though, to make sure they are unique) @@ -160,10 +159,10 @@ def _read_info_files(self) -> None: """Load current info files""" if self.folder is None: return - if self._jsonl_reading_allowance <= self._jsonl_readings: + readings = max((r.readings for r in self._jsonl_readers.values()), default=0) + if self._jsonl_reading_allowance <= readings: # bypass reloading info files return - self._jsonl_readings += 1 folder = Path(self.folder) # read all existing jsonl files find_cmd = 'find . -type f -name "*-info.jsonl"' @@ -179,52 +178,8 @@ def _read_info_files(self) -> None: out = e.output # stderr contains missing tmp files names = out.decode("utf8").splitlines() for name in names: - fp = folder / name - last = 0 - meta = {} - fail = "" - with fp.open("rb") as f: - for k, line in enumerate(f): - if fail: - msg = f"Failed to read non-last line #{k - 1} in {fp}:\n{fail!r}" - raise RuntimeError(msg) - count = len(line) - last = last + count - line = line.strip() - if not line: - logger.debug("Skipping empty line #%s", k) - continue - strline = line.decode("utf8") - if not k: - if not strline.startswith(METADATA_TAG): - raise RuntimeError(f"metadata missing in info file {fp}") - strline = strline[len(METADATA_TAG) :] - try: - info = json.loads(strline) - except json.JSONDecodeError: - msg = "Failed to read to line #%s in %s in info file %s" - logger.warning(msg, k, name, strline) - # last line could be currently being written? - # (let's be robust to it) - fail = strline - last -= count # move back for next read - continue - if not k: # metadata - meta = info - new_last = self._info_files_last.get(fp.name, last) - if new_last > last: - last = new_last - msg = "Forwarding to byte %s in info file %s" - logger.debug(msg, last, name) - f.seek(last) - continue - key = info.pop("#key") - dinfo = DumpInfo( - jsonl=fp, byte_range=(last - count, last), **meta, content=info - ) - self._key_info[key] = dinfo - - self._info_files_last[fp.name] = last # f.tell() + jreader = self._jsonl_readers.setdefault(name, JsonlReader(folder / name)) + self._key_info.update(jreader.read()) def values(self) -> tp.Iterable[X]: for key in self: @@ -309,7 +264,8 @@ def frozen_cache_folder(self) -> tp.Iterator[None]: This is useful to speed up __contains__ statement with many missing items, which could trigger thousands of file rereads """ - self._jsonl_reading_allowance = self._jsonl_readings + 1 + readings = max((r.readings for r in self._jsonl_readers.values()), default=0) + self._jsonl_reading_allowance = readings + 1 try: yield finally: @@ -375,6 +331,7 @@ def __setitem__(self, key: str, value: X) -> None: if cd.folder is not None: if self._info_filepath is None: raise RuntimeError("Cannot write out of a writer context") + fp = self._info_filepath if self._dumper is None: self._dumper = DumperLoader.CLASSES[cd.cache_type](cd.folder) self._exit_stack.enter_context(self._dumper.open()) @@ -387,7 +344,6 @@ def __setitem__(self, key: str, value: X) -> None: meta = {"cache_type": cd.cache_type} if self._info_handle is None: # create the file only when required to avoid leaving empty files for some time - fp = self._info_filepath self._info_handle = self._exit_stack.enter_context(fp.open("ab")) if not self._info_handle.tell(): meta_str = METADATA_TAG + json.dumps(meta) + "\n" @@ -397,13 +353,14 @@ def __setitem__(self, key: str, value: X) -> None: self._info_handle.write(b + b"\n") info.pop("#key") dinfo = DumpInfo( - jsonl=self._info_filepath, + jsonl=fp, byte_range=(current, current + len(b) + 1), content=info, **meta, ) cd._key_info[key] = dinfo - cd._info_files_last[self._info_filepath.name] = self._info_handle.tell() + last = self._info_handle.tell() + cd._jsonl_readers.setdefault(fp.name, JsonlReader(fp))._last = last # reading will reload to in-memory cache if need be # (since dumping may have loaded the underlying data, let's not keep it) if cd.permissions is not None: @@ -413,3 +370,61 @@ def __setitem__(self, key: str, value: X) -> None: except Exception: # pylint: disable=broad-except pass # avoid issues in case of overlapping processes os.utime(cd.folder) # make sure the modified time is updated + + +class JsonlReader: + def __init__(self, filepath: str | Path) -> None: + self._fp = Path(filepath) + self._last = 0 + self.readings = 0 + + def read(self) -> dict[str, DumpInfo]: + self.readings += 1 + last = 0 + meta = {} + fail = "" + key_info = {} + with self._fp.open("rb") as f: + for k, line in enumerate(f): + if fail: + msg = ( + f"Failed to read non-last line #{k - 1} in {self._fp}:\n{fail!r}" + ) + raise RuntimeError(msg) + count = len(line) + last = last + count + line = line.strip() + if not line: + logger.debug("Skipping empty line #%s", k) + continue + strline = line.decode("utf8") + if not k: + if not strline.startswith(METADATA_TAG): + raise RuntimeError(f"metadata missing in info file {self._fp}") + strline = strline[len(METADATA_TAG) :] + try: + info = json.loads(strline) + except json.JSONDecodeError: + msg = "Failed to read to line #%s in %s in info file %s" + logger.warning(msg, k, self._fp.name, strline) + # last line could be currently being written? + # (let's be robust to it) + fail = strline + last -= count # move back for next read + continue + if not k: # metadata + meta = info + new_last = max(last, self._last) + if new_last > last: + last = new_last + msg = "Forwarding to byte %s in info file %s" + logger.debug(msg, last, self._fp.name) + f.seek(last) + continue + key = info.pop("#key") + dinfo = DumpInfo( + jsonl=self._fp, byte_range=(last - count, last), **meta, content=info + ) + key_info[key] = dinfo + self._last = last + return key_info diff --git a/exca/test_map.py b/exca/test_map.py index 5be155f4..c45f3ee7 100644 --- a/exca/test_map.py +++ b/exca/test_map.py @@ -110,15 +110,15 @@ def test_map_infra_cache_dict_calls(tmp_path: Path) -> None: whatever = Whatever(infra={"folder": tmp_path, "cluster": "local"}) # type: ignore cd = whatever.infra.cache_dict _ = list(whatever.process([1, 2, 3, 4])) - assert cd._jsonl_readings == 3 + assert max(r.readings for r in cd._jsonl_readers.values()) == 1 whatever = Whatever(infra={"folder": tmp_path, "cluster": "local"}) # type: ignore cd = whatever.infra.cache_dict _ = list(whatever.process([1])) - assert cd._jsonl_readings == 1 + assert max(r.readings for r in cd._jsonl_readers.values()) == 1 _ = list(whatever.process([2, 3, 4])) - assert cd._jsonl_readings == 1 + assert max(r.readings for r in cd._jsonl_readers.values()) == 1 _ = list(whatever.process([5])) - assert cd._jsonl_readings == 4 + assert max(r.readings for r in cd._jsonl_readers.values()) == 2 def test_missing_yield() -> None: