diff --git a/docs/conf.py b/docs/conf.py index 358c8db6..855b0b14 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -97,5 +97,5 @@ } html_js_files = [ - ('https://scripts.simpleanalyticscdn.com/latest.js', {'async': 'async', 'defer': 'defer'}), -] \ No newline at end of file + ("https://scripts.simpleanalyticscdn.com/latest.js", {"async": "async", "defer": "defer"}), +] diff --git a/src/atomworks/io/parser.py b/src/atomworks/io/parser.py index d57e75fe..2f791502 100644 --- a/src/atomworks/io/parser.py +++ b/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 diff --git a/tests/io/components/test_caching.py b/tests/io/components/test_caching.py index a6759220..07e3577d 100644 --- a/tests/io/components/test_caching.py +++ b/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__])