From 332dab94eb51020f05244ee3fff63fdb345bbe43 Mon Sep 17 00:00:00 2001 From: Herbert Wendler Date: Thu, 13 Aug 2026 14:44:20 +0200 Subject: [PATCH 1/4] fix(io): write parse cache atomically Cache entries were written straight to their target path. A process that died while serialising - a worker hitting a wall-clock limit or being preempted, which is routine when the cache is populated from a batch scheduler - left a truncated file behind that subsequent runs accepted as a valid cache entry, so the failure surfaced later as an unrelated parse error. Entries are now written to a temporary file that is moved into place once complete. The temporary name carries 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 now passed explicitly. pandas infers it from the file name, and the temporary name does not carry the suffix that the destination has, so without this the entries would silently be stored uncompressed. It is derived from the destination, which keeps the stored format unchanged. Tests cover an interrupted write leaving neither a cache entry nor a temporary file, a destination created concurrently, and the stored format still being gzip compressed. --- src/atomworks/io/parser.py | 34 +++++++++- tests/io/components/test_caching.py | 97 ++++++++++++++++++++++++++++- 2 files changed, 127 insertions(+), 4 deletions(-) diff --git a/src/atomworks/io/parser.py b/src/atomworks/io/parser.py index d57e75fe..3213ebce 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 @@ -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 diff --git a/tests/io/components/test_caching.py b/tests/io/components/test_caching.py index a6759220..ba03699c 100644 --- a/tests/io/components/test_caching.py +++ b/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__]) From 1cce0c92c895d682373d7f08a0d61ae3227c321a Mon Sep 17 00:00:00 2001 From: Herbert Wendler Date: Thu, 13 Aug 2026 15:22:00 +0200 Subject: [PATCH 2/4] style(docs): apply make format to conf.py Running `make format` reformats this file: the html_js_files entry uses five spaces of indentation and single quotes, and the file lacks a trailing newline. The change is whitespace and quoting only. Kept as a separate commit so that the accompanying cache fix stays confined to its own scope. --- docs/conf.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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"}), +] From e87e7fad47f21632f5335c5d678ee793bf41423e Mon Sep 17 00:00:00 2001 From: hwendler <104029565+hwendler@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:30:32 +0200 Subject: [PATCH 3/4] Apply suggestions from code review Co-authored-by: Nathaniel Corley --- src/atomworks/io/parser.py | 18 +++--------------- tests/io/components/test_caching.py | 2 -- 2 files changed, 3 insertions(+), 17 deletions(-) diff --git a/src/atomworks/io/parser.py b/src/atomworks/io/parser.py index 3213ebce..6520cf7d 100644 --- a/src/atomworks/io/parser.py +++ b/src/atomworks/io/parser.py @@ -373,21 +373,9 @@ def parse( # 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. + # 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"} compression = "gzip" if cache_file_path.suffix == ".gz" else "infer" node = socket.gethostname().replace(os.sep, "_") diff --git a/tests/io/components/test_caching.py b/tests/io/components/test_caching.py index ba03699c..044d7f17 100644 --- a/tests/io/components/test_caching.py +++ b/tests/io/components/test_caching.py @@ -10,8 +10,6 @@ from atomworks.io.utils.testing import assert_same_atom_array 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 = [ From 1367359d4ac40c56e2f3558c96242e2c5fcdbd16 Mon Sep 17 00:00:00 2001 From: Herbert Wendler Date: Tue, 18 Aug 2026 21:43:10 +0200 Subject: [PATCH 4/4] fix(io): map cache suffix to compression; drop redundant cache tests Handle every suffix utils/compression recognises (.gz, .gzip, .zst) instead of .gz alone, so an entry is never written uncompressed under a compressed name. Note pandas infers .gz and .zst but not .gzip. Drop two tests that add no coverage: the gzip-format test, since read_pickle raises BadGzipFile on an uncompressed .pkl.gz and the end-to-end cache test already round-trips; and the concurrent-destination test, which only discriminates on Windows while CI runs ubuntu-latest. Co-Authored-By: Claude Opus 5 --- src/atomworks/io/parser.py | 8 +++- tests/io/components/test_caching.py | 70 +---------------------------- 2 files changed, 9 insertions(+), 69 deletions(-) diff --git a/src/atomworks/io/parser.py b/src/atomworks/io/parser.py index 6520cf7d..2f791502 100644 --- a/src/atomworks/io/parser.py +++ b/src/atomworks/io/parser.py @@ -86,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.""" @@ -377,11 +381,13 @@ def parse( # 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"} - compression = "gzip" if cache_file_path.suffix == ".gz" else "infer" + # 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) diff --git a/tests/io/components/test_caching.py b/tests/io/components/test_caching.py index 044d7f17..07e3577d 100644 --- a/tests/io/components/test_caching.py +++ b/tests/io/components/test_caching.py @@ -1,5 +1,3 @@ -import os -import socket import time from pathlib import Path @@ -111,38 +109,12 @@ def _cache_files(cache_dir: Path) -> list[Path]: 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. - """ + """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): - # Write a partial file first, so the test would fail if the target path were written - # to directly instead of via a temporary file. + # 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") @@ -155,43 +127,5 @@ def failing_to_pickle(obj, path, *args, **kwargs): 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__])