From adb7502a1c7e5c003da557a8c36f31ada516b012 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 21:38:44 +0200 Subject: [PATCH] utils: an optional SOURCE_CACHE_ROOT shares large source downloads A download given a cache_key, with SOURCE_CACHE_ROOT set, is copied from {SOURCE_CACHE_ROOT}/{cache_key} when there, and left there once fetched, so later processes (other tiles, other pods) skip the publisher. Only immutable URLs get a key: the Zenodo records behind gpw_livestock and liao_gaced30. SOURCE_CACHE_ROOT is optional; a Config field with a default may now be left out of .env. --- .env.example | 3 + jdluc/__tests__/source_cache_test.py | 149 +++++++++++++++++++++++++++ jdluc/config.py | 3 + jdluc/datasets/gpw_livestock.py | 4 +- jdluc/datasets/liao_gaced30.py | 9 +- jdluc/storage.py | 8 ++ jdluc/utils.py | 20 ++++ 7 files changed, 190 insertions(+), 6 deletions(-) create mode 100644 jdluc/__tests__/source_cache_test.py diff --git a/.env.example b/.env.example index dff0bf2..498a3f8 100644 --- a/.env.example +++ b/.env.example @@ -7,3 +7,6 @@ NUMBER_OF_DASK_WORKERS = 12 # SCRATCH_ROOT = '/absolute/path/to/scratch' SCRATCH_ROOT = 'gs://cornerstone-scratch-us-central1' USDA_NASS_API_KEY = 'Replace me with value from https://quickstats.nass.usda.gov/api' +# SOURCE_CACHE_ROOT is optional: when set, large immutable source downloads are kept there and +# reused by later processes instead of being fetched again. Any fsspec URI, like the roots above. +# SOURCE_CACHE_ROOT = '/absolute/path/to/source-cache' diff --git a/jdluc/__tests__/source_cache_test.py b/jdluc/__tests__/source_cache_test.py new file mode 100644 index 0000000..89174ed --- /dev/null +++ b/jdluc/__tests__/source_cache_test.py @@ -0,0 +1,149 @@ +import contextlib +import dataclasses +import pathlib +import types +import typing + +import dotenv +import pytest + +from jdluc import config, utils + +URL = "https://publisher.example/records/1/files/source.zip/content" + +# The real .env parsing, past its functools.cache, captured before any test stubs it +FROM_DOT_ENV = config.Config.__dict__["from_dot_env"].__func__.__wrapped__ + + +def get_required_lines(leave_out: str = "") -> str: + """A .env line for every field `Config` requires, but `leave_out`.""" + types_ = typing.get_type_hints(config.Config) + return "".join( + f"{field.name.upper():s} = '{'2' if types_[field.name] is int else 'x'}'\n" + for field in dataclasses.fields(config.Config) + if field.default is dataclasses.MISSING and field.name != leave_out + ) + + +def read_dot_env( + monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path, text: str +) -> config.Config: + path_to_env = tmp_path / ".env" + path_to_env.write_text(text) + monkeypatch.setattr(dotenv, "find_dotenv", lambda **kwargs: str(path_to_env)) + return FROM_DOT_ENV(config.Config) + + +@pytest.fixture +def requested_urls(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Stub the publisher: every GET serves b"source"; return the URLs it was asked for.""" + urls: list[str] = [] + + @contextlib.contextmanager + def request(url: str, **kwargs: object): # type: ignore[no-untyped-def] + urls.append(url) + yield types.SimpleNamespace( + iter_content=lambda chunk_size: iter([b"source"]), + raise_for_status=lambda: None, + ) + + monkeypatch.setattr( + utils, "get_requests_session", lambda: types.SimpleNamespace(request=request) + ) + return urls + + +def set_source_cache_root( + monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path, root: str +) -> None: + configured = dataclasses.replace( + read_dot_env( + monkeypatch=monkeypatch, tmp_path=tmp_path, text=get_required_lines() + ), + source_cache_root=root, + ) + monkeypatch.setattr( + config.Config, "from_dot_env", classmethod(lambda cls: configured) + ) + + +def test_a_miss_downloads_and_stores_and_a_hit_skips_the_publisher( + monkeypatch: pytest.MonkeyPatch, + requested_urls: list[str], + tmp_path: pathlib.Path, +) -> None: + set_source_cache_root( + monkeypatch=monkeypatch, root=str(tmp_path / "cache"), tmp_path=tmp_path + ) + for name in ("first", "second"): + utils.save_remote_url_to_local_path( + cache_key="zenodo/1/source.zip", + local_path=str(tmp_path / name), + params={}, + remote_url=URL, + ) + assert (tmp_path / name).read_bytes() == b"source" + assert requested_urls == [URL] + assert (tmp_path / "cache/zenodo/1/source.zip").read_bytes() == b"source" + + +@pytest.mark.parametrize( + ("root", "cache_key"), + [ + pytest.param("", "zenodo/1/source.zip", id="root-unset"), + pytest.param("cache", None, id="no-key"), + ], +) +def test_without_a_root_or_a_key_every_call_downloads( + monkeypatch: pytest.MonkeyPatch, + requested_urls: list[str], + tmp_path: pathlib.Path, + root: str, + cache_key: str | None, +) -> None: + set_source_cache_root( + monkeypatch=monkeypatch, + root=str(tmp_path / root) if root else "", + tmp_path=tmp_path, + ) + for name in ("first", "second"): + utils.save_remote_url_to_local_path( + cache_key=cache_key, + local_path=str(tmp_path / name), + params={}, + remote_url=URL, + ) + assert requested_urls == [URL, URL] + assert not (tmp_path / "cache").exists() + + +@pytest.mark.parametrize( + ("extra_line", "expected"), + [ + pytest.param("", "", id="absent"), + pytest.param("SOURCE_CACHE_ROOT = '/sources'\n", "/sources", id="set"), + ], +) +def test_source_cache_root_may_be_left_out_of_dot_env( + monkeypatch: pytest.MonkeyPatch, + tmp_path: pathlib.Path, + extra_line: str, + expected: str, +) -> None: + configured = read_dot_env( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + text=get_required_lines() + extra_line, + ) + assert configured.source_cache_root == expected + + +def test_a_required_key_may_not_be_left_out_of_dot_env( + monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path +) -> None: + with pytest.raises(KeyError, match="SCRATCH_ROOT"): + read_dot_env( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + text=get_required_lines(leave_out="scratch_root"), + ) diff --git a/jdluc/config.py b/jdluc/config.py index 5eab326..48c1d02 100644 --- a/jdluc/config.py +++ b/jdluc/config.py @@ -16,6 +16,8 @@ class Config: number_of_dask_workers: int scratch_root: str usda_nass_api_key: str + # Optional: a field with a default may be left out of .env + source_cache_root: str = "" @classmethod @functools.cache @@ -29,6 +31,7 @@ def from_dot_env(cls) -> typing.Self: **{ field.name: types[field.name](configs[field.name.upper()]) for field in dataclasses.fields(cls) + if field.default is dataclasses.MISSING or field.name.upper() in configs } ) diff --git a/jdluc/datasets/gpw_livestock.py b/jdluc/datasets/gpw_livestock.py index 7b3e8f5..7e740bb 100644 --- a/jdluc/datasets/gpw_livestock.py +++ b/jdluc/datasets/gpw_livestock.py @@ -73,10 +73,12 @@ def get_path_to_source(species: Species, year: int) -> str: path_to_source = os.path.join( tempfile.mkdtemp(), f"gpw-livestock-{species:s}-{year:d}" ) + record = SPECIES_TO_ZENODO_RECORD[species] utils.save_remote_url_to_local_path( + cache_key=f"zenodo/{record:d}/{name:s}", local_path=path_to_source, params={}, - remote_url=f"https://zenodo.org/api/records/{SPECIES_TO_ZENODO_RECORD[species]:d}/files/{name:s}/content", + remote_url=f"https://zenodo.org/api/records/{record:d}/files/{name:s}/content", ) return path_to_source diff --git a/jdluc/datasets/liao_gaced30.py b/jdluc/datasets/liao_gaced30.py index 7ee63ec..f0fa682 100644 --- a/jdluc/datasets/liao_gaced30.py +++ b/jdluc/datasets/liao_gaced30.py @@ -72,15 +72,14 @@ def get_transform(max_lat: float, min_lon: float) -> rasterio.Affine: @utils.threadsafe_cache def get_path_to_archive_for_low_longitude(low_longitude: int) -> str: # The published names sign both bounds, as in lon_-006_to_+000.zip - name = urllib.parse.quote( - f"lon_{low_longitude:+04d}_to_{low_longitude + ARCHIVE_DEGREES:+04d}.zip" - ) + name = f"lon_{low_longitude:+04d}_to_{low_longitude + ARCHIVE_DEGREES:+04d}.zip" # NB: deliberately not a TemporaryDirectory, so that the archive outlives this call - path_to_archive = os.path.join(tempfile.mkdtemp(), name) + path_to_archive = os.path.join(tempfile.mkdtemp(), urllib.parse.quote(name)) utils.save_remote_url_to_local_path( + cache_key=f"zenodo/18199675/{name:s}", local_path=path_to_archive, params={}, - remote_url=f"https://zenodo.org/api/records/18199675/files/{name:s}/content", + remote_url=f"https://zenodo.org/api/records/18199675/files/{urllib.parse.quote(name):s}/content", ) return path_to_archive diff --git a/jdluc/storage.py b/jdluc/storage.py index 1159d1e..b992556 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -38,6 +38,14 @@ def to_gdal_path(uri: str) -> str: return uri +def get_file(uri: str, local_path: str) -> None: + import fsspec # type: ignore[import-untyped] + + logger.info(f"Reading file {uri=:s} to {local_path=:s}") + fs, path = fsspec.url_to_fs(url=uri) + fs.get_file(path, local_path) + + def join_uri(root: str, prefix: str) -> str: return f"{root.rstrip('/'):s}/{prefix:s}" diff --git a/jdluc/utils.py b/jdluc/utils.py index 61e7ca0..99186f1 100644 --- a/jdluc/utils.py +++ b/jdluc/utils.py @@ -12,6 +12,8 @@ import requests.adapters import xarray +from jdluc import config, storage + logger = logging.getLogger(__name__) @@ -42,7 +44,23 @@ def save_remote_url_to_local_path( # is the failure mode, and without this the read blocks forever read_timeout_seconds: int = 120, body_retries: int = 3, + cache_key: str | None = None, ) -> None: + """Stream `remote_url` to `local_path`. + + With a `cache_key`, and SOURCE_CACHE_ROOT set, the download is shared through + `{SOURCE_CACHE_ROOT}/{cache_key}`: a process that finds it there copies it rather than + asking the publisher, and one that does not downloads as usual and leaves a copy for the + next. Pass a key only for a URL whose content never changes, as nothing refreshes the copy. + """ + cache_uri: str | None = None + if cache_key is not None: + source_cache_root = config.Config.from_dot_env().source_cache_root + if source_cache_root: + cache_uri = storage.join_uri(prefix=cache_key, root=source_cache_root) + if storage.path_exists(uri=cache_uri): + storage.get_file(local_path=local_path, uri=cache_uri) + return for attempt in range(body_retries): logger.info(f"GET'ing from {remote_url=:s} with {params=:}") try: @@ -69,6 +87,8 @@ def save_remote_url_to_local_path( raise time.sleep(2**attempt) else: + if cache_uri is not None: + storage.put_file(local_path=local_path, uri=cache_uri) return