storage: a zarr cache entry counts only once its write has finished

jdluc/__tests__/storage_test.py

@@ -0,0 +1,119 @@
+import dataclasses
+import pathlib
+import typing
+
+import numpy
+import pytest
+import xarray
+
+from jdluc import storage
+from jdluc.config import Config
+
+# The fields a unit test needs a real value for
+TEST_ENV = {"number_of_dask_workers": 2}
+
+
+def get_unit_test_config() -> Config:
+    """TEST_ENV over every other field left empty, so a test that reaches for a root fails rather than touching real storage."""
+    types = typing.get_type_hints(Config)
+    empty = {field.name: types[field.name]() for field in dataclasses.fields(Config)}
+    return Config(**{**empty, **TEST_ENV})
+
+
+@pytest.fixture(autouse=True)
+def config_without_dot_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Writing a zarr reads the config; these tests need no .env."""
+    config = get_unit_test_config()
+    monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config))
+
+
+def get_dset() -> xarray.Dataset:
+    return xarray.Dataset(
+        {"band": (("y", "x"), numpy.arange(16, dtype="float32").reshape(4, 4))}
+    ).chunk({"x": 2, "y": 2})
+
+
+def assert_reads_back(cacher: storage.ZarrCacher) -> None:
+    xarray.testing.assert_equal(
+        cacher.deserialize()["band"].drop_vars("spatial_ref"), get_dset()["band"]
+    )
+
+
+def write_without_marker(cacher: storage.ZarrCacher) -> None:
+    """A write killed part-way: zarr.json, but no marker."""
+    storage.write_dask_dataset_to_zarr(dset=get_dset(), path_to_zarr=cacher.uri)
+    assert storage.path_exists(uri=cacher.exists_uri)
+    assert not storage.path_exists(uri=cacher.complete_uri)
+
+
+def test_zarr_cacher_marks_a_finished_write_complete(tmp_path: pathlib.Path) -> None:
+    cacher = storage.ZarrCacher(hash_key="abc", root=str(tmp_path))
+    assert not storage.path_exists(uri=cacher.exists_uri)
+    assert not storage.path_exists(uri=cacher.complete_uri)
+
+    cacher.serialize(get_dset())
+
+    assert storage.path_exists(uri=cacher.exists_uri)
+    assert storage.path_exists(uri=cacher.complete_uri)
+    assert_reads_back(cacher)
+
+
+def test_zarr_cacher_refuses_to_write_over_an_incomplete_zarr(
+    tmp_path: pathlib.Path,
+) -> None:
+    cacher = storage.ZarrCacher(hash_key="abc", root=str(tmp_path))
+    write_without_marker(cacher)
+
+    with pytest.raises(storage.ZarrExistsAndIncompleteException):
+        cacher.serialize(get_dset())
+
+    assert not storage.path_exists(uri=cacher.complete_uri)
+
+
+def test_zarr_cacher_refuses_to_write_over_a_complete_zarr(
+    tmp_path: pathlib.Path,
+) -> None:
+    cacher = storage.ZarrCacher(hash_key="abc", root=str(tmp_path))
+    cacher.serialize(get_dset())
+
+    with pytest.raises(storage.ZarrExistsAndCompleteException):
+        cacher.serialize(get_dset())
+
+    assert_reads_back(cacher)
+
+
+@pytest.fixture
+def scratch_root(
+    tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+) -> pathlib.Path:
+    config = dataclasses.replace(get_unit_test_config(), scratch_root=str(tmp_path))
+    monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config))
+    return tmp_path
+
+
+def test_cache_to_zarr_handles_each_state(scratch_root: pathlib.Path) -> None:
+    calls = []
+
+    # `attempt` is left out of the hash, so each call misses the in-memory cache but
+    # shares one zarr
+    @storage.cache_to_zarr(ignored_args=["attempt"], version=0)
+    def compute(attempt: int) -> xarray.Dataset:
+        calls.append(attempt)
+        if attempt == 3:  # another process finishes the entry while this one computes
+            marker.touch()
+        return get_dset()
+
+    compute(attempt=1)  # absent: written and marked complete
+    (zarr,) = scratch_root.glob("*.zarr")
+    marker = zarr / "complete"
+    assert marker.exists()
+
+    marker.unlink()
+    with pytest.raises(storage.ZarrExistsAndIncompleteException):
+        compute(attempt=2)
+
+    xarray.testing.assert_equal(
+        compute(attempt=3)["band"].drop_vars("spatial_ref"), get_dset()["band"]
+    )
+    compute(attempt=4)  # complete: a hit, so compute does not run
+    assert calls == [1, 2, 3]

jdluc/storage.py

@@ -1,3 +1,4 @@
+import contextlib
 import dataclasses
 import functools
 import hashlib
@@ -6,7 +7,9 @@
 import os
 import threading
 import typing
+import warnings
 
+import fsspec  # type: ignore[import-untyped]
 import pandas
 import xarray
 
@@ -16,15 +19,11 @@
 
 
 def path_exists(uri: str) -> bool:
-    import fsspec  # type: ignore[import-untyped]
-
     fs, path = fsspec.url_to_fs(url=uri)
     return fs.exists(path)
 
 
 def put_file(local_path: str, uri: str) -> None:
-    import fsspec  # type: ignore[import-untyped]
-
     logger.info(f"Writing file {local_path=:s} to {uri=:s}")
     fs, path = fsspec.url_to_fs(url=uri)
     fs.makedirs(os.path.dirname(p=path), exist_ok=True)
@@ -107,7 +106,12 @@ def open_zarr_to_dask_dataset(path_to_zarr: str) -> xarray.Dataset:
 
     path_to_zarr = path_to_zarr.rstrip("/")
     logger.info(f"Loading {path_to_zarr=:s}")
-    dset = xarray.open_zarr(consolidated=False, store=path_to_zarr)
+    with warnings.catch_warnings():
+        # ZarrCacher's `complete` marker sits in the store root, which zarr warns about
+        warnings.filterwarnings(
+            "ignore", message="Object at complete is not recognized"
+        )
+        dset = xarray.open_zarr(consolidated=False, store=path_to_zarr)
     assert dset.chunks is not None
     assert isinstance(dset, xarray.Dataset)
     assert dset.encoding["source"] == path_to_zarr
@@ -124,6 +128,8 @@ def __init__(self, *, hash_key: str, root: str) -> None: ...
     def uri(self) -> str: ...
     @property
     def exists_uri(self) -> str: ...
+    @property
+    def complete_uri(self) -> str: ...
     def deserialize(self) -> R: ...
     def serialize(self, value: R) -> None: ...
 
@@ -137,7 +143,9 @@ class ParquetCacher:
     def uri(self) -> str:
         return join_uri(prefix=self.hash_key + ".parquet", root=self.root)
 
+    # A parquet write is atomic, so existing is being complete
     exists_uri = uri
+    complete_uri = uri
 
     def deserialize(self) -> pandas.DataFrame:
         logger.info(f"Loading from {self.uri=:s}")
@@ -148,8 +156,26 @@ def serialize(self, df: pandas.DataFrame) -> None:
         df.to_parquet(path=self.uri)
 
 
+class ZarrExistsAndCompleteException(FileExistsError):
+    """A zarr cache entry is already fully written."""
+
+
+class ZarrExistsAndIncompleteException(FileExistsError):
+    """A zarr cache entry exists, but its write was killed part-way or is still running."""
+
+
 @dataclasses.dataclass
 class ZarrCacher:
+    """A zarr at `uri`, complete once the `complete` marker inside it exists.
+
+    `zarr.json` is the first object written, so it only shows that a write began. The
+    marker is written after the last chunk, which gives three states: absent, exists but
+    incomplete (a write killed part-way, or still running), and exists and complete.
+    `serialize` raises on either existing state rather than writing over it, so callers
+    can choose how to handle it (poll, skip, raise). An incomplete zarr is usually
+    resolved by hand, by deleting it once no process is writing it.
+    """
+
     hash_key: str
     root: str
 
@@ -161,11 +187,27 @@ def uri(self) -> str:
     def exists_uri(self) -> str:
         return f"{self.uri:s}/zarr.json"
 
+    @property
+    def complete_uri(self) -> str:
+        return f"{self.uri:s}/complete"
+
     def deserialize(self) -> xarray.Dataset:
         return open_zarr_to_dask_dataset(path_to_zarr=self.uri)
 
     def serialize(self, dset: xarray.Dataset) -> None:
+        if path_exists(uri=self.complete_uri):
+            raise ZarrExistsAndCompleteException(
+                f"{self.uri=:s} is already complete: another process finished writing it."
+            )
+        if path_exists(uri=self.exists_uri):
+            raise ZarrExistsAndIncompleteException(
+                f"{self.uri=:s} exists without {self.complete_uri=:s}: a write was killed "
+                "part-way or is still running. Once no process is writing it, delete the "
+                "zarr to recompute it."
+            )
         write_dask_dataset_to_zarr(dset=dset, path_to_zarr=self.uri)
+        fs, path = fsspec.url_to_fs(url=self.complete_uri)
+        fs.touch(path)
 
 
 class CacherDecoratorProtocol(typing.Protocol[R]):
@@ -217,8 +259,10 @@ def inner(*args: P.args, **kwargs: P.kwargs) -> R:
                 CacherProtocol[R],
                 cacher_cls(hash_key=hash_key, root=config.scratch_root),
             )
-            if not path_exists(uri=cacher.exists_uri):
-                cacher.serialize(func(*args, **kwargs))
+            if not path_exists(uri=cacher.complete_uri):
+                # Another process may finish this entry first; its result is as good as ours
+                with contextlib.suppress(ZarrExistsAndCompleteException):
+                    cacher.serialize(func(*args, **kwargs))
             # Always return the deserialized cache so cache hits and misses
             # look identical and so downstream operations start with a fresh
             # compute graph

tools/mark-zarrs-complete.py

@@ -0,0 +1,60 @@
+"""Mark zarrs written before the `complete` marker existed as complete, so they are not recomputed.
+
+`ZarrCacher` counts a zarr as cached only once its `complete` marker exists, and zarrs written
+before the marker have none. This trusts every one of them, including any that are actually
+partial: delete a zarr you suspect instead. Only zarrs directly under the root are checked, which
+is where `ZarrCacher` writes them.
+
+Run it before the marker change merges, and once more after deploy to catch zarrs the old code
+wrote in between. Run it only while nothing is writing to the root, since a zarr mid-write would be
+marked complete. It lists what it would mark unless given `--apply`, and is idempotent.
+
+  uv run python tools/mark-zarrs-complete.py
+  uv run python tools/mark-zarrs-complete.py --apply
+"""
+
+import argparse
+import logging
+import os
+
+import fsspec  # type: ignore[import-untyped]
+
+from jdluc import config, storage
+
+logger = logging.getLogger(__name__)
+
+
+def main() -> None:
+    logging.basicConfig(format="%(levelname)s - %(message)s", level=logging.INFO)
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--root", help="defaults to SCRATCH_ROOT from .env")
+    parser.add_argument(
+        "--apply", action="store_true", help="write the markers, not just list them"
+    )
+    args = parser.parse_args()
+    root = args.root or config.Config.from_dot_env().scratch_root
+
+    fs, path = fsspec.url_to_fs(url=root)
+    cachers = [
+        storage.ZarrCacher(
+            hash_key=os.path.basename(os.path.dirname(zarr_json)).removesuffix(".zarr"),
+            root=root,
+        )
+        for zarr_json in sorted(fs.glob(f"{path.rstrip('/'):s}/*.zarr/zarr.json"))
+    ]
+    unmarked = [c for c in cachers if not storage.path_exists(uri=c.complete_uri)]
+    for cacher in unmarked:
+        logger.info(
+            f"{'Marking' if args.apply else 'Would mark'} {cacher.uri:s} complete"
+        )
+        if args.apply:
+            fs.touch(fsspec.url_to_fs(url=cacher.complete_uri)[1])
+    logger.info(
+        f"{len(unmarked):d} of {len(cachers):d} zarrs {'marked' if args.apply else 'to mark'}"
+    )
+
+
+if __name__ == "__main__":
+    main()