Improve deletion in cache dict (now needs context manager)

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## [Unreleased]
 
+[breaking]
+
+- `CacheDict`: deletions require a `write()` context (like writes). [#326]
+- `DumpContext.shared_file`: content suffixes must start with `.`. [#326]
+
+[other]
+
 - `DiscriminatedModel`: optimized look-up. [#313]
 - `steps`: fixed nested infra claim deadlock. [#323]
 

docs/infra/explanation.md

@@ -137,7 +137,7 @@ cache = cachedict.CacheDict(folder=tmp_path, keep_in_ram=True)
 # the dictionary is empty:
 assert not cache
 
-# writes require a context manager for efficiency with multiple writes
+# writes and deletions require a context manager for efficiency with multiple writes
 x = np.random.rand(2, 12)
 with cache.write():
     cache["blublu"] = x

exca/cachedict/core.py

@@ -161,6 +161,8 @@ def __reduce__(self) -> tp.Any:
     def clear(self) -> None:
         self._ram_data.clear()
         self._key_info.clear()
+        self._jsonl_readers.clear()
+        self._folder_modified = -1.0
         if self.folder is None or not self.folder.exists():
             return
         # let's remove content but not the folder to keep same permissions
@@ -185,24 +187,31 @@ def keys(self) -> tp.Iterator[str]:
         keys = set(self._ram_data) | set(self._key_info)
         return iter(keys)
 
-    def _read_info_files(self, max_workers: int = 4) -> None:
+    def _read_info_files(self, max_workers: int = 4, force: bool = False) -> None:
         """Load current info files.
 
         Each writer appends to its own JSONL file, so concurrent writes
         of the same key produce duplicate entries across files.  For
         duplicates, whichever file comes last in iterdir() order wins
         (non-deterministic); duplicates are kept so explicit deletion
-        clears every known copy."""
+        clears every known copy.
+
+        Parameters
+        ----------
+        max_workers:
+            Maximum number of threads reading the info files.
+        force:
+            Read even if the folder is frozen or looks unmodified.
+        """
         if self.folder is None or not self.folder.exists():
             return
         readings = max((r.readings for r in self._jsonl_readers.values()), default=0)
-        if self._jsonl_reading_allowance <= readings:
-            # bypass reloading info files
-            return
+        if not force and self._jsonl_reading_allowance <= readings:
+            return  # bypass reloading info files
         modified = self.folder.lstat().st_mtime
         nothing_new = self._folder_modified == modified
         self._folder_modified = modified
-        if nothing_new:
+        if nothing_new and not force:
             logger.debug("Nothing new to read from info files")
             return  # nothing new!
         cpus = os.cpu_count()
@@ -248,7 +257,7 @@ def _cleanup_orphaned_jsonl_files(self) -> None:
             except FileNotFoundError:
                 self._jsonl_readers.pop(name, None)
                 continue
-            logger.warning("Cleaning up orphaned files for %s", name)
+            logger.debug("Cleaning up orphaned files for %s", name)
             prefix = name.removesuffix("-info.jsonl")
             paths = [*self.folder.glob(f"{prefix}.*"), reader._fp]
             data_dir = self.folder / DumpContext.DATA_DIR
@@ -297,11 +306,12 @@ def _write_ctx(self, value: DumpContext | None) -> None:
 
     @contextlib.contextmanager
     def write(self) -> tp.Iterator["CacheDict[X]"]:
-        """Context manager for writing items to the cache."""
+        """Context manager for writing to (and deleting from) the cache."""
         if self._write_ctx is not None:
             raise RuntimeError("Cannot re-open an already open writer")
         if self.folder is not None:
             self._write_ctx = DumpContext(self.folder, permissions=self.permissions)
+        self._local.deleted_in_scope = False
         try:
             if self._write_ctx is not None:
                 with self._write_ctx:
@@ -312,6 +322,11 @@ def write(self) -> tp.Iterator["CacheDict[X]"]:
             self._write_ctx = None
             if self.folder is not None:
                 utils.best_effort_utime(self.folder)
+                if self._local.deleted_in_scope:
+                    try:
+                        self._read_info_files(force=True)  # sweep emptied jsonl pairs
+                    except Exception as e:  # must not mask the body's exception
+                        logger.warning("Failed to sweep %s: %s", self.folder, e)
 
     @contextlib.contextmanager
     def writer(self) -> tp.Iterator["CacheDict[X]"]:
@@ -328,7 +343,7 @@ def __setitem__(self, key: str, value: X) -> None:
         if not isinstance(key, str):
             raise TypeError(f"Non-string keys are not allowed (got {key!r})")
         if self.folder is not None and self._write_ctx is None:
-            raise RuntimeError("Cannot write outside of a writer context")
+            raise RuntimeError("Cannot write outside of a write() context")
         if self._folder_modified <= 0:
             _ = self.keys()
         if key in self._ram_data or key in self._key_info:
@@ -359,6 +374,9 @@ def __delitem__(self, key: str) -> None:
         if self._dumper is None:
             del self._ram_data[key]
             return
+        if self._write_ctx is None:
+            raise RuntimeError("Cannot delete outside of a write() context")
+        self._local.deleted_in_scope = True
         if key not in self._key_info:
             _ = key in self  # populate _key_info from disk
         self._ram_data.pop(key, None)

exca/cachedict/dumpcontext.py

@@ -232,14 +232,15 @@ def shared_file(self, suffix: str) -> tuple[tp.IO[bytes], str]:
         """Open a shared file for appending. Returns (handle, relative_name).
         Content files go under DATA_DIR/; info files (-info.jsonl)
         stay in the root folder. Reused across calls with the same suffix."""
-        if "." not in suffix:
-            raise ValueError(f"suffix must contain '.', got {suffix!r}")
+        is_info = suffix == self.INFO_SUFFIX
+        if not is_info and not suffix.startswith("."):
+            msg = f"suffix must start with '.' to be reclaimable, got {suffix!r}"
+            raise ValueError(msg)
         if self._stack is None:
             raise RuntimeError("DumpContext must be used as a context manager for writes")
         if threading.get_native_id() != self._thread_id:
             raise RuntimeError("DumpContext must not be shared across threads")
         basename = f"{self._prefix}{suffix}"
-        is_info = suffix == self.INFO_SUFFIX
         name = basename if is_info else f"{self.DATA_DIR}/{basename}"
         if name not in self._files:
             path = self.folder / name

exca/cachedict/handlers.py

@@ -472,7 +472,7 @@ def __dump_info__(cls, ctx: DumpContext, value: tp.Any) -> dict[str, tp.Any]:
             ) from e
         if len(raw) <= cls.MAX_INLINE_SIZE:
             return {"content": value}
-        f, name = ctx.shared_file("-data.jsonl")
+        f, name = ctx.shared_file(".data.jsonl")
         offset = f.tell()
         f.write(raw + b"\n")
         return {"filename": name, "offset": offset, "length": len(raw)}

exca/cachedict/test_cachedict.py

@@ -60,7 +60,10 @@ def test_array_cache(tmp_path: Path, in_ram: bool) -> None:
     cache2 = cd.CacheDict(folder=folder)
     assert isinstance(cache2["blublu"], np.ndarray)
     # del
-    del cache2["blublu"]
+    with pytest.raises(RuntimeError, match=r"write\(\) context"):
+        del cache2["blublu"]
+    with cache2.write():
+        del cache2["blublu"]
     assert set(cache2.keys()) == {"blabla"}
     # clear
     cache2.clear()
@@ -115,6 +118,10 @@ def test_specialized_dump(
         cache_type = cache_type[:-2]
         memmap_cache_size = 0
     proc = psutil.Process()
+    try:
+        proc.open_files()
+    except (psutil.AccessDenied, PermissionError) as e:
+        pytest.skip(f"psutil cannot list open files: {e}")
     cache: cd.CacheDict[tp.Any] = cd.CacheDict(
         folder=tmp_path,
         keep_in_ram=keep_in_ram,
@@ -200,7 +207,8 @@ def test_info_jsonl_deletion(tmp_path: Path) -> None:
             assert out.startswith(b"{") and out.endswith(b"}\n")
     # remove one
     chosen = np.random.choice(keys)
-    del cache[chosen]
+    with cache.write():
+        del cache[chosen]
     assert len(cache) == 2
     cache = cd.CacheDict(folder=tmp_path, keep_in_ram=False)
     assert len(cache) == 2
@@ -215,7 +223,8 @@ def test_info_jsonl_deletion_removes_duplicate_entries(tmp_path: Path) -> None:
 
     cache = cd.CacheDict(folder=tmp_path, keep_in_ram=False)
     assert cache["x"] == 12
-    del cache["x"]
+    with cache.write():
+        del cache["x"]
 
     cache = cd.CacheDict(folder=tmp_path, keep_in_ram=False)
     assert "x" not in cache
@@ -299,10 +308,17 @@ def test_clone_is_view_only(tmp_path: Path) -> None:
         assert revived["k"] == 7
 
 
-@pytest.mark.parametrize("cache_type", ["MemmapArrayFile", "String"])
-def test_orphaned_data_file_cleanup(tmp_path: Path, cache_type: str) -> None:
+@pytest.mark.parametrize("read_before_delete", [False, True])
+@pytest.mark.parametrize("cache_type", ["MemmapArrayFile", "String", "Json"])
+def test_orphaned_data_file_cleanup(
+    tmp_path: Path, cache_type: str, read_before_delete: bool
+) -> None:
     """Test that orphaned data files are cleaned up when all items are deleted."""
-    data: tp.Any = np.random.rand(3, 12) if cache_type == "MemmapArrayFile" else "hello"
+    data: tp.Any = {
+        "MemmapArrayFile": np.random.rand(3, 12),
+        "String": "hello",
+        "Json": {"blob": "x" * 50_000},  # above MAX_INLINE_SIZE -> shared data file
+    }[cache_type]
     cache: cd.CacheDict[tp.Any] = cd.CacheDict(
         folder=tmp_path, keep_in_ram=False, cache_type=cache_type
     )
@@ -311,13 +327,20 @@ def test_orphaned_data_file_cleanup(tmp_path: Path, cache_type: str) -> None:
         for c in "abc":
             ex.submit(_write_items, cache, [f"{c}1", f"{c}2"], data)
     assert len(list(tmp_path.glob("*-info.jsonl"))) == 3
-    # Delete all items from one writer, files still exist (cleanup is lazy)
-    for key in ["a1", "a2", "c1", "b2"]:
-        del cache[key]
-    assert len(list(tmp_path.glob("*-info.jsonl"))) == 3
-    # Trigger cleanup via keys() - orphaned pair should be deleted
+    if read_before_delete:
+        assert len(set(cache.keys())) == 6
+    with cache.write():
+        for key in ["a1", "a2", "c1", "b2"]:
+            del cache[key]
+    remaining = list(tmp_path.glob("*-info.jsonl"))
+    assert len(remaining) == 2, (
+        f"leaving write() should drop the emptied pair {remaining}"
+    )
+    live = {p.name.removesuffix("-info.jsonl") for p in remaining}
+    data_files = (tmp_path / "data").glob("*")
+    stale = [p.name for p in data_files if p.name.split(".")[0] not in live]
+    assert not stale, f"data files outliving their info file {stale}"
     assert set(cache.keys()) == {"b1", "c2"}
-    assert len(list(tmp_path.glob("*-info.jsonl"))) == 2
 
 
 @pytest.mark.parametrize(
@@ -367,9 +390,8 @@ def test_jsonl_edge_cases(tmp_path: Path, content: str, should_delete: bool) ->
     # Write and delete an item to trigger reader initialization for our test file
     with cache.write():
         cache["x"] = np.array([1])
-    del cache["x"]
-    # Trigger cleanup
-    _ = list(cache.keys())
+    with cache.write():
+        del cache["x"]
     # Check result
     for fp in [jsonl, data_file]:
         if should_delete:

exca/cachedict/test_dumpcontext.py

@@ -110,6 +110,8 @@ def test_shared_file_lifecycle(tmp_path: Path) -> None:
     ctx = DumpContext(tmp_path)
     with pytest.raises(RuntimeError, match="context manager"):
         ctx.shared_file(".data")
+    with pytest.raises(ValueError, match="must start with"):
+        ctx.shared_file("-data.jsonl")
     with ctx:
         f1, name1 = ctx.shared_file(".data")
         f2, name2 = ctx.shared_file(".data")

exca/cachedict/test_legacy_format.py

@@ -141,7 +141,8 @@ def test_legacy_external_static_roundtrip(tmp_path: Path) -> None:
     assert set(cache.keys()) == {"key1", "key2"}
     assert cache["key1"] == {"a": 1, "b": [2, 3]}
     assert cache["key2"] == "hello"
-    del cache["key1"]
+    with cache.write():
+        del cache["key1"]
     assert set(cache.keys()) == {"key2"}
     cache2: cd.CacheDict[tp.Any] = cd.CacheDict(folder=tmp_path)
     assert set(cache2.keys()) == {"key2"}
@@ -189,7 +190,8 @@ def test_mixed_old_and_new_format(tmp_path: Path) -> None:
     assert cache2["multiline"] == "line1\nline2\nline3"
     assert cache2["extra"] == "new value"
     # Delete an old item, verify new items survive
-    del cache2["hello"]
+    with cache2.write():
+        del cache2["hello"]
     cache3: cd.CacheDict[tp.Any] = cd.CacheDict(folder=dst)
     assert set(cache3.keys()) == {"multiline", "extra"}
 

exca/map.py

@@ -270,11 +270,12 @@ def _find_missing(self, items: dict[str, tp.Any]) -> dict[str, tp.Any]:
             self._check_configs(write=True)
         if self.mode == "force":
             to_remove = set(items) - set(missing) - state.recomputed
-            if to_remove:
+            if to_remove and isinstance(cache, CacheDict):
                 msg = "Clearing %s items for %s (infra.mode=%s)"
                 logger.warning(msg, len(to_remove), self.uid(), self.mode)
-                for uid in to_remove:
-                    del cache[uid]
+                with cache.write():
+                    for uid in to_remove:
+                        del cache[uid]
             missing = {x: y for x, y in items.items() if x not in state.recomputed}
             if isinstance(cache, CacheDict):
                 # dont record computed items if no cache

exca/steps/backends.py

@@ -367,10 +367,10 @@ def run_and_cache(self) -> None:
                     "Clearing partial results after invalid _run_batch output: %s",
                     self.paths.step_uid,
                 )
-            with self.cache_dict.frozen_cache_folder():
-                for uid in written_uids:
-                    if uid in self.cache_dict:
-                        del self.cache_dict[uid]
+                with self.cache_dict.write(), self.cache_dict.frozen_cache_folder():
+                    for uid in written_uids:
+                        if uid in self.cache_dict:
+                            del self.cache_dict[uid]
             if folder is not None:
                 e.add_note(f"  -> cache may be invalid: {folder}")
             raise
@@ -592,7 +592,7 @@ def _clear_caches(
                 logger.warning("Failed to cancel %s%s: %s", paths.step_uid, uids, e)
         # Success first → a mid-clear crash leaves a recoverable cached
         # error rather than a stale success (fail closed).
-        with cd.frozen_cache_folder():
+        with cd.write(), cd.frozen_cache_folder():
             for uid in uids:
                 if uid in cd:
                     del cd[uid]