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
@@ -167,7 +168,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 +371,33 @@ 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.
+        #
+        # The write goes to a temporary file that is then moved into place, rather than
+        # directly to the target path. A process interrupted while writing -- a worker
+        # hitting a wall-clock limit or being preempted, which is routine when the cache is
+        # filled from a batch scheduler -- would otherwise leave a truncated file behind
+        # that later runs treat as a valid cache entry. The temporary name includes host and
+        # process id so that several workers sharing a cache directory, possibly on a
+        # network filesystem, cannot overwrite each other's partial writes.
+        #
+        # An existing entry is not normally rewritten, but two workers can pass that check
+        # at the same time and both proceed, so the move has to tolerate an occupied
+        # destination; Path.replace does, whereas Path.rename raises on Windows in that case.
+        #
+        # Compression is passed explicitly because pandas would otherwise infer it from the
+        # file name, and the temporary name does not carry the suffix the destination has.
+        # Deriving it from the destination keeps the stored format exactly as before.
         result_to_cache = {k: v for k, v in result.items() if k != "assemblies"}
-        pd.to_pickle(result_to_cache, cache_file_path)
+        compression = "gzip" if cache_file_path.suffix == ".gz" else "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)
+            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,18 @@
+import os
+import socket
 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
+
+# A small structure that ships with the test data, so the cache tests below do not depend
+# on a local PDB mirror.
+STRUCTURE = TEST_DATA_IO / "2hhb.cif.gz"
 
 TEST_CASES = [
     "4NDZ",  # 29K atoms, large enough to test caching without too much variance
@@ -100,5 +108,92 @@ 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_cached_entry_is_gzip_compressed(tmp_path: Path) -> None:
+    """The stored format is unchanged: entries are gzip compressed, as their name says.
+
+    Cache files are named `.pkl.gz` and pandas infers the compression from that name. Writing
+    through a temporary file would lose the inference, since the temporary name does not carry
+    the suffix, so the compression has to be passed explicitly. This test pins the resulting
+    format down.
+    """
+    parse(STRUCTURE, cache_dir=tmp_path, save_to_cache=True)
+    (entry,) = _cache_files(tmp_path)
+    assert entry.name.endswith(".pkl.gz")
+    gzip_magic = bytes.fromhex("1f8b")
+    assert entry.read_bytes()[:2] == gzip_magic, "cache entry is not gzip compressed"
+
+    # ...and it is still readable, i.e. the format matches what the reader expects.
+    result = parse(STRUCTURE, cache_dir=tmp_path, load_from_cache=True)
+    assert result["asym_unit"].array_length() > 0
+
+
+def test_cache_write_is_atomic(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """An interrupted write leaves no cache entry behind.
+
+    Without an atomic write, a process killed while serialising would leave a truncated file
+    that later runs would treat as a valid cache entry. The write is made to fail part way
+    through; afterwards the cache directory must contain neither an entry nor a leftover
+    temporary file.
+    """
+    real_to_pickle = pd.to_pickle
+
+    def failing_to_pickle(obj, path, *args, **kwargs):
+        # Write a partial file first, so the test would fail if the target path were written
+        # to directly instead of via a temporary file.
+        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"
+
+
+def test_cache_write_tolerates_a_destination_created_concurrently(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Moving the finished file into place works even if the entry appeared meanwhile.
+
+    An existing entry is normally not rewritten, but two workers can pass that check at the
+    same time and both go on to write, so the move of the second one finds its destination
+    occupied. `Path.replace` overwrites it; `Path.rename` would raise `FileExistsError` on
+    Windows and leave the worker's temporary file behind. The race is reproduced here by
+    creating the destination while the temporary file is being written.
+
+    Note that the distinction between the two only shows on Windows: POSIX `rename` replaces
+    an existing destination silently, so on Linux and macOS this test passes either way and
+    covers only that the entry ends up complete and no temporary file is stranded.
+    """
+    real_to_pickle = pd.to_pickle
+    # Rebuild the suffix the implementation appends, so the destination can be derived from
+    # the temporary path without assuming anything about the host name.
+    suffix = f".{socket.gethostname().replace(os.sep, '_')}.{os.getpid()}.tmp"
+
+    def to_pickle_and_simulate_other_worker(obj, path, *args, **kwargs):
+        real_to_pickle(obj, path, *args, **kwargs)
+        tmp = Path(path)
+        assert tmp.name.endswith(suffix), "cache write no longer uses the expected temporary name"
+        tmp.with_name(tmp.name[: -len(suffix)]).write_bytes(b"written by another worker")
+
+    monkeypatch.setattr(pd, "to_pickle", to_pickle_and_simulate_other_worker)
+    parse(STRUCTURE, cache_dir=tmp_path, save_to_cache=True)
+    monkeypatch.setattr(pd, "to_pickle", real_to_pickle)
+
+    assert len(_cache_files(tmp_path)) == 1, "the concurrent write left more than one entry"
+    assert not list(tmp_path.rglob("*.tmp")), "the move left a temporary file behind"
+
+    # The entry is the one this worker wrote, not the placeholder, and it is readable.
+    result = parse(STRUCTURE, cache_dir=tmp_path, load_from_cache=True)
+    assert result["asym_unit"].array_length() > 0
+
+
 if __name__ == "__main__":
     pytest.main([__file__])