Fix/atomic parse cache writes

docs/conf.py

@@ -97,5 +97,5 @@
 }
 
 html_js_files = [
-     ('https://scripts.simpleanalyticscdn.com/latest.js', {'async': 'async', 'defer': 'defer'}),
-]
+    ("https://scripts.simpleanalyticscdn.com/latest.js", {"async": "async", "defer": "defer"}),
+]

src/atomworks/io/parser.py

@@ -13,6 +13,7 @@
 import io
 import logging
 import os
+import socket
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Literal
@@ -85,6 +86,10 @@
 _CACHE_SHARDING_DEPTH = 2  # Use 2-level sharding by default (e.g., ab/cd/abcdef123456/)
 _CACHE_SHARDING_CHARS_PER_DIR = 2  # Number of characters per directory level
 
+# Cache-file suffix -> pandas compression, covering what `utils.compression` recognises.
+# Note pandas infers `.gz` and `.zst` but not `.gzip`.
+_CACHE_COMPRESSION = {".gz": "gzip", ".gzip": "gzip", ".zst": "zstd"}
+
 
 def _get_atomworks_version() -> str:
     """Lazy import of atomworks version to avoid circular imports."""
@@ -167,7 +172,10 @@ def parse(
             If not provided, the file type will be inferred automatically.
         load_from_cache (bool, optional): Whether to load pre-compiled results from cache. Defaults to False.
         cache_dir (PathLike, optional): Directory path to save pre-compiled results. Defaults to None.
-        save_to_cache (bool, optional): Whether to save the results to cache when building the structure. Defaults to False.
+        save_to_cache (bool, optional): Whether to save the results to cache when building the structure.
+            Defaults to False. An entry is written to a temporary file and then moved into place, so a
+            process interrupted while writing leaves no partial entry behind, and several processes may
+            fill a shared cache directory concurrently. An entry that already exists is not rewritten.
 
         **Parsing arguments:**
         ccd_mirror_path (str, optional): Path to the local mirror of the Chemical Component Dictionary (recommended).
@@ -367,9 +375,23 @@ def parse(
         # Ensure all parent directories exist
         cache_file_path.parent.mkdir(parents=True, exist_ok=True)
 
-        # Save the result to the cache, excluding the assemblies
+        # Save the result to the cache, excluding the assemblies.
+        #
+        # Write to a temp file (named with host and pid to avoid collisions between
+        # workers sharing the cache) and atomically move it into place, so an interrupted
+        # write can't leave a corrupt cache entry
         result_to_cache = {k: v for k, v in result.items() if k != "assemblies"}
-        pd.to_pickle(result_to_cache, cache_file_path)
+        # Explicit: pandas would infer compression from the temp name, which has no suffix
+        compression = _CACHE_COMPRESSION.get(cache_file_path.suffix, "infer")
+        node = socket.gethostname().replace(os.sep, "_")
+        tmp_path = cache_file_path.with_name(f"{cache_file_path.name}.{node}.{os.getpid()}.tmp")
+        try:
+            pd.to_pickle(result_to_cache, tmp_path, compression=compression)
+            # replace() not rename(): two workers can race here, and rename() raises on Windows
+            tmp_path.replace(cache_file_path)
+        except BaseException:
+            tmp_path.unlink(missing_ok=True)
+            raise
 
     return result
 

tests/io/components/test_caching.py

@@ -1,10 +1,14 @@
 import time
+from pathlib import Path
 
+import pandas as pd
 import pytest
 
 from atomworks.io.parser import parse
 from atomworks.io.utils.testing import assert_same_atom_array
-from tests.io.conftest import get_pdb_path
+from tests.io.conftest import TEST_DATA_IO, get_pdb_path
+
+STRUCTURE = TEST_DATA_IO / "2hhb.cif.gz"
 
 TEST_CASES = [
     "4NDZ",  # 29K atoms, large enough to test caching without too much variance
@@ -100,5 +104,28 @@ def different_args_parse():
     assert abs(different_args_elapsed_time - normal_elapsed_time) < normal_elapsed_time * 0.8
 
 
+def _cache_files(cache_dir: Path) -> list[Path]:
+    """Return the cache entries below `cache_dir`, ignoring temporary write files."""
+    return [p for p in cache_dir.rglob("*") if p.is_file() and not p.name.endswith(".tmp")]
+
+
+def test_cache_write_is_atomic(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """An interrupted write leaves neither a cache entry nor a temporary file behind."""
+    real_to_pickle = pd.to_pickle
+
+    def failing_to_pickle(obj, path, *args, **kwargs):
+        # Partial file first, so the test fails if the target path were written directly.
+        Path(path).write_bytes(b"partial")
+        raise KeyboardInterrupt("interrupted while writing the cache")
+
+    monkeypatch.setattr(pd, "to_pickle", failing_to_pickle)
+    with pytest.raises(KeyboardInterrupt):
+        parse(STRUCTURE, cache_dir=tmp_path, save_to_cache=True)
+    monkeypatch.setattr(pd, "to_pickle", real_to_pickle)
+
+    assert not _cache_files(tmp_path), "an interrupted write left a cache entry behind"
+    assert not list(tmp_path.rglob("*.tmp")), "an interrupted write left a temporary file behind"
+
+
 if __name__ == "__main__":
     pytest.main([__file__])