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

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))

jdluc/__tests__/storage_test.py

@@ -0,0 +1,44 @@
+import pathlib
+
+import numpy
+import pytest
+import xarray
+
+from jdluc import storage
+
+
+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)

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]):