utils: an optional SOURCE_CACHE_ROOT shares large source downloads

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

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"),
+        )

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

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
 

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
 

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}"
 

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