From ca0f73f3ac9c1a20f6a508bbf82930aba64350c6 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Thu, 15 Oct 2026 19:38:57 +0200 Subject: [PATCH 1/4] storage: a zarr cache entry counts only once its write has finished For zarr files, set up a .complete marker file to check when data is fully written. This replaces relying on zarr.json -- that doesn't work because zarr.json is the first to be written. This might happen from a failed or aborted calculation, or because multiple processes are trying to write the same file (e.g. running 2 countries which share a tile) A partially-written zarr might be not cause any warning signs later, since it might be techncially valid in spite of missing data. Now an empty .complete marker is written beside the zarr after its last chunk, and the entry counts as cached only once the marker exists. A zarr found without its marker is a write that was killed part-way or is still running: serialize raises FileExistsError and leaves it for someone to resolve by hand, usually by deleting the zarr, rather than writing over it and risking a race with a live writer. Entries written before this change have no marker, so each raises once until deleted. --- jdluc/__tests__/storage_test.py | 58 +++++++++++++++++++++++++++++++++ jdluc/storage.py | 22 ++++++++++--- 2 files changed, 75 insertions(+), 5 deletions(-) create mode 100644 jdluc/__tests__/storage_test.py diff --git a/jdluc/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py new file mode 100644 index 0000000..2e5e48e --- /dev/null +++ b/jdluc/__tests__/storage_test.py @@ -0,0 +1,58 @@ +import pathlib + +import numpy +import pytest +import xarray + +from jdluc import storage +from jdluc.config import Config + + +@pytest.fixture(autouse=True) +def config_without_dot_env(monkeypatch: pytest.MonkeyPatch) -> None: + # Writing a zarr reads number_of_dask_workers from .env, which CI does not have + config = Config( + harvard_dataverse_guestbook_json="", + ingest_root="", + number_of_dask_workers=2, + scratch_root="", + usda_nass_api_key="", + ) + 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 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) + + cacher.serialize(get_dset()) + + assert storage.path_exists(uri=cacher.exists_uri) + assert_reads_back(cacher) + + +def test_zarr_cacher_refuses_to_write_over_an_unfinished_zarr( + tmp_path: pathlib.Path, +) -> None: + cacher = storage.ZarrCacher(hash_key="abc", root=str(tmp_path)) + # A write killed part-way leaves zarr.json but no marker + storage.write_dask_dataset_to_zarr(dset=get_dset(), path_to_zarr=cacher.uri) + assert storage.path_exists(uri=f"{cacher.uri:s}/zarr.json") + assert not storage.path_exists(uri=cacher.exists_uri) + + with pytest.raises(FileExistsError): + cacher.serialize(get_dset()) + + assert not storage.path_exists(uri=cacher.exists_uri) diff --git a/jdluc/storage.py b/jdluc/storage.py index 1159d1e..684fda5 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -7,6 +7,7 @@ import threading import typing +import fsspec # type: ignore[import-untyped] import pandas import xarray @@ -16,15 +17,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) @@ -150,6 +147,13 @@ def serialize(self, df: pandas.DataFrame) -> None: @dataclasses.dataclass class ZarrCacher: + """A zarr at `uri`, complete once the `.complete` marker beside it exists. + + The marker is written after the last chunk, so a zarr without one is a write that was + killed part-way or is still running. Writing over it could race a live writer, so it is + an error left for someone to resolve by hand, usually by deleting the zarr. + """ + hash_key: str root: str @@ -159,13 +163,21 @@ def uri(self) -> str: @property def exists_uri(self) -> str: - return f"{self.uri:s}/zarr.json" + 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.uri): + raise FileExistsError( + f"{self.uri=:s} exists without {self.exists_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.exists_uri) + fs.touch(path) class CacherDecoratorProtocol(typing.Protocol[R]): From 8bab50011e37b4f40af9961d38b11f3838f765d2 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Thu, 15 Oct 2026 20:50:29 +0200 Subject: [PATCH 2/4] tests: unit tests run without a .env An autouse fixture in jdluc/__tests__/conftest.py stubs Config.from_dot_env for every test not marked integration. Each field gets its type's empty value, with two dask workers, so a unit test that reaches for a storage root fails instead of touching a bucket. Integration tests still read the real .env. This replaces the fixture storage_test carried for itself, so the next unit test that writes a zarr needs no copy of it, and a new Config field needs no edit to the tests. --- jdluc/__tests__/conftest.py | 27 +++++++++++++++++++++++++++ jdluc/__tests__/storage_test.py | 14 -------------- 2 files changed, 27 insertions(+), 14 deletions(-) create mode 100644 jdluc/__tests__/conftest.py diff --git a/jdluc/__tests__/conftest.py b/jdluc/__tests__/conftest.py new file mode 100644 index 0000000..f5bcde1 --- /dev/null +++ b/jdluc/__tests__/conftest.py @@ -0,0 +1,27 @@ +import dataclasses +import typing + +import pytest + +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 unit 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( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch +) -> None: + """Unit tests need no .env; integration tests still read the real one.""" + if request.node.get_closest_marker("integration"): + return + config = get_unit_test_config() + monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config)) diff --git a/jdluc/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py index 2e5e48e..34a3953 100644 --- a/jdluc/__tests__/storage_test.py +++ b/jdluc/__tests__/storage_test.py @@ -5,20 +5,6 @@ import xarray from jdluc import storage -from jdluc.config import Config - - -@pytest.fixture(autouse=True) -def config_without_dot_env(monkeypatch: pytest.MonkeyPatch) -> None: - # Writing a zarr reads number_of_dask_workers from .env, which CI does not have - config = Config( - harvard_dataverse_guestbook_json="", - ingest_root="", - number_of_dask_workers=2, - scratch_root="", - usda_nass_api_key="", - ) - monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config)) def get_dset() -> xarray.Dataset: From f9daffd5c78877f39be3756952aee14ada4fac05 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Wed, 11 Nov 2026 15:56:25 +0100 Subject: [PATCH 3/4] storage: distinguish incomplete from complete zarrs; marker inside the zarr Review feedback on the complete marker: - The marker moves inside the zarr, to `{uri}/complete`, so deleting the zarr deletes it too. - exists_uri is zarr.json again, and a new complete_uri sits alongside it (the same for parquet, whose write is atomic). The cache decorator checks complete_uri for a hit. - serialize raises ZarrExistsAndIncompleteException or ZarrExistsAndCompleteException (both FileExistsError), so callers can choose how to handle each state. - mark_existing_zarrs_complete backfills the marker into zarrs written before it existed, so they are not recomputed. - Zarr's warning about the unrecognized `complete` object is filtered when opening a store. - The no-.env test setup moves from conftest.py into storage_test.py. --- jdluc/__tests__/conftest.py | 27 ------------ jdluc/__tests__/storage_test.py | 70 +++++++++++++++++++++++++++---- jdluc/storage.py | 74 +++++++++++++++++++++++++++------ 3 files changed, 125 insertions(+), 46 deletions(-) delete mode 100644 jdluc/__tests__/conftest.py diff --git a/jdluc/__tests__/conftest.py b/jdluc/__tests__/conftest.py deleted file mode 100644 index f5bcde1..0000000 --- a/jdluc/__tests__/conftest.py +++ /dev/null @@ -1,27 +0,0 @@ -import dataclasses -import typing - -import pytest - -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 unit 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( - request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch -) -> None: - """Unit tests need no .env; integration tests still read the real one.""" - if request.node.get_closest_marker("integration"): - return - config = get_unit_test_config() - monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config)) diff --git a/jdluc/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py index 34a3953..b0146c8 100644 --- a/jdluc/__tests__/storage_test.py +++ b/jdluc/__tests__/storage_test.py @@ -1,10 +1,30 @@ +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: @@ -19,26 +39,62 @@ def assert_reads_back(cacher: storage.ZarrCacher) -> None: ) +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_unfinished_zarr( +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)) - # A write killed part-way leaves zarr.json but no marker - storage.write_dask_dataset_to_zarr(dset=get_dset(), path_to_zarr=cacher.uri) - assert storage.path_exists(uri=f"{cacher.uri:s}/zarr.json") - assert not storage.path_exists(uri=cacher.exists_uri) + write_without_marker(cacher) - with pytest.raises(FileExistsError): + with pytest.raises(storage.ZarrExistsAndIncompleteException): cacher.serialize(get_dset()) - assert not storage.path_exists(uri=cacher.exists_uri) + 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) + + +def test_mark_existing_zarrs_complete(tmp_path: pathlib.Path) -> None: + unmarked = storage.ZarrCacher(hash_key="unmarked", root=str(tmp_path)) + write_without_marker(unmarked) + complete = storage.ZarrCacher(hash_key="complete", root=str(tmp_path)) + complete.serialize(get_dset()) + + assert storage.mark_existing_zarrs_complete(root=str(tmp_path)) == [unmarked.uri] + assert not storage.path_exists(uri=unmarked.complete_uri) + + assert storage.mark_existing_zarrs_complete(root=str(tmp_path), dry_run=False) == [ + unmarked.uri + ] + assert storage.path_exists(uri=unmarked.complete_uri) + assert_reads_back(unmarked) + + assert storage.mark_existing_zarrs_complete(root=str(tmp_path), dry_run=False) == [] diff --git a/jdluc/storage.py b/jdluc/storage.py index 684fda5..94bcb81 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -6,6 +6,7 @@ import os import threading import typing +import warnings import fsspec # type: ignore[import-untyped] import pandas @@ -104,7 +105,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 @@ -121,6 +127,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: ... @@ -134,7 +142,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}") @@ -145,13 +155,24 @@ 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 beside it exists. - - The marker is written after the last chunk, so a zarr without one is a write that was - killed part-way or is still running. Writing over it could race a live writer, so it is - an error left for someone to resolve by hand, usually by deleting the zarr. + """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 @@ -163,23 +184,52 @@ def uri(self) -> str: @property def exists_uri(self) -> str: - return f"{self.uri:s}.complete" + 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.uri): - raise FileExistsError( - f"{self.uri=:s} exists without {self.exists_uri=:s}: a write was killed " + 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.exists_uri) + fs, path = fsspec.url_to_fs(url=self.complete_uri) fs.touch(path) +def mark_existing_zarrs_complete(root: str, dry_run: bool = True) -> list[str]: + """Touch the `complete` marker in every zarr under `root` that has `zarr.json` but no marker. + + For zarrs written before the marker existed, which are trusted as complete. Returns the + zarrs marked (or that would be, under `dry_run`). Idempotent. + """ + fs, path = fsspec.url_to_fs(url=root) + protocol = root.split("://")[0] + "://" if "://" in root else "" + marked = [] + for zarr_json in sorted(fs.glob(f"{path.rstrip('/'):s}/*.zarr/zarr.json")): + zarr_path = zarr_json.removesuffix("/zarr.json") + if fs.exists(f"{zarr_path:s}/complete"): + continue + uri = protocol + zarr_path + logger.info(f"{'Would mark' if dry_run else 'Marking'} {uri=:s} complete") + if not dry_run: + fs.touch(f"{zarr_path:s}/complete") + marked.append(uri) + return marked + + class CacherDecoratorProtocol(typing.Protocol[R]): def __call__(self, func: typing.Callable[P, R]) -> typing.Callable[P, R]: ... @@ -229,7 +279,7 @@ 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): + if not path_exists(uri=cacher.complete_uri): 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 From 21506bf3aad360093d58168a2b81ea055aade2fc Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Wed, 11 Nov 2026 16:14:51 +0100 Subject: [PATCH 4/4] storage: tolerate another process finishing a zarr first; move marker migration to tools/ - The cache decorator suppresses ZarrExistsAndCompleteException: when another process finishes the same entry while this one computes, its result is used instead of crashing the run. ZarrExistsAndIncompleteException still raises. - A test drives the decorator through each state: absent, incomplete, completed by another process, and a hit. - mark_existing_zarrs_complete leaves the library for tools/mark-zarrs-complete.py, matching the other one-off scripts: dry run by default, --apply to write, root defaulting to SCRATCH_ROOT. It builds paths through ZarrCacher, so the marker name lives in one place. Its docstring covers what it trusts, that only zarrs directly under the root are checked, and that it must run while nothing is writing to the root. --- jdluc/__tests__/storage_test.py | 51 +++++++++++++++++++--------- jdluc/storage.py | 26 +++----------- tools/mark-zarrs-complete.py | 60 +++++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 38 deletions(-) create mode 100644 tools/mark-zarrs-complete.py diff --git a/jdluc/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py index b0146c8..6c42189 100644 --- a/jdluc/__tests__/storage_test.py +++ b/jdluc/__tests__/storage_test.py @@ -82,19 +82,38 @@ def test_zarr_cacher_refuses_to_write_over_a_complete_zarr( assert_reads_back(cacher) -def test_mark_existing_zarrs_complete(tmp_path: pathlib.Path) -> None: - unmarked = storage.ZarrCacher(hash_key="unmarked", root=str(tmp_path)) - write_without_marker(unmarked) - complete = storage.ZarrCacher(hash_key="complete", root=str(tmp_path)) - complete.serialize(get_dset()) - - assert storage.mark_existing_zarrs_complete(root=str(tmp_path)) == [unmarked.uri] - assert not storage.path_exists(uri=unmarked.complete_uri) - - assert storage.mark_existing_zarrs_complete(root=str(tmp_path), dry_run=False) == [ - unmarked.uri - ] - assert storage.path_exists(uri=unmarked.complete_uri) - assert_reads_back(unmarked) - - assert storage.mark_existing_zarrs_complete(root=str(tmp_path), dry_run=False) == [] +@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 94bcb81..4a5f361 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -1,3 +1,4 @@ +import contextlib import dataclasses import functools import hashlib @@ -209,27 +210,6 @@ def serialize(self, dset: xarray.Dataset) -> None: fs.touch(path) -def mark_existing_zarrs_complete(root: str, dry_run: bool = True) -> list[str]: - """Touch the `complete` marker in every zarr under `root` that has `zarr.json` but no marker. - - For zarrs written before the marker existed, which are trusted as complete. Returns the - zarrs marked (or that would be, under `dry_run`). Idempotent. - """ - fs, path = fsspec.url_to_fs(url=root) - protocol = root.split("://")[0] + "://" if "://" in root else "" - marked = [] - for zarr_json in sorted(fs.glob(f"{path.rstrip('/'):s}/*.zarr/zarr.json")): - zarr_path = zarr_json.removesuffix("/zarr.json") - if fs.exists(f"{zarr_path:s}/complete"): - continue - uri = protocol + zarr_path - logger.info(f"{'Would mark' if dry_run else 'Marking'} {uri=:s} complete") - if not dry_run: - fs.touch(f"{zarr_path:s}/complete") - marked.append(uri) - return marked - - class CacherDecoratorProtocol(typing.Protocol[R]): def __call__(self, func: typing.Callable[P, R]) -> typing.Callable[P, R]: ... @@ -280,7 +260,9 @@ def inner(*args: P.args, **kwargs: P.kwargs) -> R: cacher_cls(hash_key=hash_key, root=config.scratch_root), ) if not path_exists(uri=cacher.complete_uri): - cacher.serialize(func(*args, **kwargs)) + # 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()