From 31a2a86e2bfb20562252d50c707a6c3ac9f952c2 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Thu, 15 Oct 2026 19:38:57 +0200 Subject: [PATCH 1/2] 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 03a630e1b8f35d1a1eaa5d8c1c0fb21d7b16a97e Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Thu, 15 Oct 2026 20:50:29 +0200 Subject: [PATCH 2/2] 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: