storage: a zarr cache entry counts only once its write has finished by danohu · Pull Request #16 · cornerstone-data/luc

danohu · 2026-10-15T17:52:56Z

For zarr cache entries, write a complete marker once the data is fully written, and count an entry as cached only when the marker exists.

Relying on zarr.json doesn't work, because it is the first object written. A failed or aborted run, or two processes writing the same entry (e.g. two countries that share a tile), can leave a zarr that looks cached, and it may be technically valid while missing data.

What changes

ZarrCacher writes an empty {uri}/complete after the last chunk. It sits inside the zarr, so deleting the zarr deletes it too.
Cachers get a complete_uri alongside exists_uri. For parquet, whose write is atomic, they are the same. The cache decorator checks complete_uri for a hit.
serialize distinguishes three states and raises on the two existing ones, so callers can choose how to handle them (poll, skip, raise):
- absent → write, then touch the marker
- exists but incomplete → ZarrExistsAndIncompleteException (killed part-way, or still running; resolve by hand, usually by deleting the zarr once nothing is writing it)
- exists and complete → ZarrExistsAndCompleteException (another process finished first)
Both subclass FileExistsError.
The cache decorator suppresses ZarrExistsAndCompleteException: if another process finishes the same entry while this one computes, the decorator returns that result rather than crashing. ZarrExistsAndIncompleteException still propagates, so a leftover partial zarr stops the run until someone deletes it.
tools/mark-zarrs-complete.py adds the marker to existing zarrs (see the push plan below).
The storage unit tests run without a .env (fixture local to storage_test.py), and cover each serialize state and the decorator through each state.

Push plan

Existing zarrs have no marker, so after this merges each one would raise ZarrExistsAndIncompleteException instead of being reused. To keep them, mark them complete with tools/mark-zarrs-complete.py. It runs from the repo root and works on SCRATCH_ROOT from .env unless given --root. Without --apply it only lists what it would mark.

Run it only while nothing is writing to the root. It cannot tell an old zarr from one being written right now, and would mark a zarr mid-write as complete.

Before merge, with no pipeline running, list what would be marked, then mark it. This is harmless to the current code, which ignores the marker:
```
uv run python tools/mark-zarrs-complete.py
uv run python tools/mark-zarrs-complete.py --apply
```
Merge and deploy.
Once no job started before the deploy is still running, and before starting new ones, run --apply again to catch zarrs the old code wrote in between. It is idempotent: entries already marked are skipped.

Caveats

The migration trusts every pre-existing zarr, including any that are actually partial. Anyone who suspects a specific entry should delete it.
It only checks zarrs directly under the root, which is where ZarrCacher writes them.
Check-then-write is still not atomic: two processes that both see "absent" will both write. The marker stops a half-written zarr from being read as complete; it isn't a lock.
Zarr warns about the unrecognized complete object when opening a store; open_zarr_to_dask_dataset filters that one warning.

jeffseif

From the IRL discussion:

we would need a push plan where all existing zarr's have the complete file touched.

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.

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.

…e 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.

… 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.

danohu · 2026-11-11T17:27:33Z

@jeffseif all the changes are now in-place, re-review please:

backfill command (assumes everything existing is valid, and nothing else is touching zarr files at that time)
limited the test config setup to these tests
moved the .complete files into the zarr directories (causes some warnings which need to be ignored)
switched to the 3-state model (absent, incomplete, complete)

jeffseif

Logic and testing look solid -- this is a welcome improvement!

I left some comments and gave you a proactive approve so that you can address them and then merge or merge immediately if you disagree with them.

jeffseif · 2026-11-11T20:28:44Z

+
+
+def assert_reads_back(cacher: storage.ZarrCacher) -> None:
+    xarray.testing.assert_equal(


TIL, I'll try to use it in future tests!

jeffseif · 2026-11-11T20:35:18Z

+
+
+def test_cache_to_zarr_handles_each_state(scratch_root: pathlib.Path) -> None:
+    calls = []


nit: we should be able to access the cache_info properties of the compute function (docs) rather than keeping track of calls ourselves. up to you if you want to adopt that or keep this.

jeffseif · 2026-11-11T21:29:57Z

    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


I feel like this comment doesn't add much beyond the line below -- thoughts on dropping it?

Suggested change

# ZarrCacher's `complete` marker sits in the store root, which zarr warns about

jeffseif · 2026-11-11T21:35:34Z

+                # Another process may finish this entry first; its result is as good as ours
+                with contextlib.suppress(ZarrExistsAndCompleteException):
+                    cacher.serialize(func(*args, **kwargs))


It looks like we are immediately handling this scenario within the cacher. What do you think about starting with the simpler set of logic (= letting ZarrExistsAndCompleteException raise) initially and then adding in more logic if we find it is needed?

Suggested change

# Another process may finish this entry first; its result is as good as ours

with contextlib.suppress(ZarrExistsAndCompleteException):

cacher.serialize(func(*args, **kwargs))

cacher.serialize(func(*args, **kwargs))

jeffseif · 2026-11-11T21:38:05Z

This forces us to mark all existing zarrs as complete, even if they failed partway.
But since the result would match today's behavior (= reading those partial zarrs) this seems fine.

danohu marked this pull request as draft October 15, 2026 18:38

danohu marked this pull request as ready for review October 15, 2026 19:35

jeffseif requested changes Oct 22, 2026

View reviewed changes

Comment thread jdluc/storage.py Resolved

Comment thread jdluc/storage.py Resolved

Comment thread jdluc/__tests__/conftest.py Resolved

danohu added 3 commits November 11, 2026 15:56

danohu force-pushed the zarr-cache-complete-marker branch from d033289 to f9daffd Compare November 11, 2026 14:58

danohu requested a review from jeffseif November 11, 2026 17:27

jeffseif approved these changes Nov 11, 2026

View reviewed changes

Navigation Menu

Repository navigation

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

storage: a zarr cache entry counts only once its write has finished#16
danohu wants to merge 4 commits into
mainfrom
zarr-cache-complete-marker

danohu commented Oct 15, 2026 •

edited

Loading

Uh oh!

jeffseif left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

danohu commented Nov 11, 2026

Uh oh!

jeffseif left a comment

Uh oh!

jeffseif Nov 11, 2026

Uh oh!

jeffseif Nov 11, 2026

Uh oh!

jeffseif Nov 11, 2026

Uh oh!

jeffseif Nov 11, 2026

Uh oh!

jeffseif Nov 11, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def assert_reads_back(cacher: storage.ZarrCacher) -> None:
		xarray.testing.assert_equal(



		def test_cache_to_zarr_handles_each_state(scratch_root: pathlib.Path) -> None:
		calls = []

Conversation

danohu commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes

Push plan

Caveats

Uh oh!

jeffseif left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

danohu commented Nov 11, 2026

Uh oh!

jeffseif left a comment

Choose a reason for hiding this comment

Uh oh!

jeffseif Nov 11, 2026

Choose a reason for hiding this comment

Uh oh!

jeffseif Nov 11, 2026

Choose a reason for hiding this comment

Uh oh!

jeffseif Nov 11, 2026

Choose a reason for hiding this comment

Uh oh!

jeffseif Nov 11, 2026

Choose a reason for hiding this comment

Uh oh!

jeffseif Nov 11, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

danohu commented Oct 15, 2026 •

edited

Loading