diff --git a/jdluc/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py new file mode 100644 index 0000000..6c42189 --- /dev/null +++ b/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] diff --git a/jdluc/storage.py b/jdluc/storage.py index 1159d1e..4a5f361 100644 --- a/jdluc/storage.py +++ b/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 diff --git a/tools/mark-zarrs-complete.py b/tools/mark-zarrs-complete.py new file mode 100644 index 0000000..bd9d77c --- /dev/null +++ b/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()