From d642b6dc053680e0db2d85a2377d9d264f5f1f24 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Thu, 30 Jul 2026 12:11:04 +0100 Subject: [PATCH 01/13] switch to source.coop for downloads, fronted by Cloudflare --- CHANGES.md | 16 +++ README.md | 44 +++--- docs/architecture.rst | 23 +-- docs/index.rst | 44 +++--- geotessera/cli.py | 10 +- geotessera/registry.py | 271 +++++++++++++++++++----------------- geotessera/registry_cli.py | 4 +- geotessera/store.py | 4 +- geotessera/visualization.py | 4 +- geotessera/zarr.py | 12 +- pyproject.toml | 2 - tests/v11.t | 4 +- uv.lock | 61 +------- 13 files changed, 247 insertions(+), 252 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 76e34c9..5de3438 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,5 +1,21 @@ ## Unreleased +### Breaking Changes + +- **All downloads now come from Source Cooperative**, fronted by CloudFlare. + Embeddings, landmasks, manifests, and the zarr store are served from the public + `https://data.source.coop/tessera/tessera` repository over HTTPS, + replacing the retired `tessera-embeddings` AWS S3 bucket. The repository + is organised by media type (`npy/{version}/`, `landmasks/{version}/`, + `zarr/{version}/`) with per-version `manifest.parquet` and + `landmasks.parquet` files colocated with their data, and carries one + embedding tree per version. (@avsm, @mtelvers) +- **Dependencies removed**: `botocore` and `awscrt` are no longer required. + Downloads use the standard library with retry/backoff, `If-Modified-Since` + conditional-GET caching, and integrity verification against the response + `Content-Length` plus a streamed MD5 whenever the server's `ETag` is a + content MD5 (single-part uploads) (@avsm) + ### New Features - **`geotessera-registry zarr-consolidate`**: New subcommand that diff --git a/README.md b/README.md index 183b034..8efb27a 100644 --- a/README.md +++ b/README.md @@ -136,7 +136,7 @@ User Request (lat/lon bbox) ↓ Parquet Registry Lookup (find available tiles from manifest.parquet) ↓ -Anonymous S3 Downloads to Output Directory (CRC64NVMe verified) +HTTPS Downloads from Source Cooperative to Output Directory (integrity verified) ├── embedding.npy (quantized) → output dir └── embedding_scales.npy → output dir ↓ @@ -498,9 +498,11 @@ GeoTessera uses a Parquet-based registry system to efficiently manage and access - **Fast queries**: Uses pandas DataFrames for efficient spatial and temporal filtering - **Block-based organization**: Internal 5×5 degree geographic blocks for efficient queries - **Minimal storage**: Manifest files are ~few MB each and cached locally -- **Integrity checking**: End-to-end CRC64NVMe checksums verified against S3's - `x-amz-checksum-crc64nvme` response header during each download - - **Always enforced** for data integrity — a checksum mismatch (or a missing checksum header) rejects the download +- **Integrity checking**: Every download is verified against the response + `Content-Length`, and against an MD5 computed over the streamed body whenever + the server's `ETag` is a content MD5 (single-part uploads) + - A mismatch rejects the download and triggers a retry, so corrupt or + truncated files never reach the cache ### Dataset Versions and Variants @@ -523,7 +525,7 @@ The registry can be loaded from multiple sources (in priority order): 1. **Local file** (via `registry_path` parameter) 2. **Local directory** (via `--registry-dir` or `registry_dir` parameter, looks for `manifest.parquet`, falling back to the legacy `registry.parquet`) 3. **Remote URL** (via `registry_url` parameter) -4. **Default remote** (from `https://s3.us-west-2.amazonaws.com/tessera-embeddings/{version}/manifest.parquet`) +4. **Default remote** (from `https://data.source.coop/tessera/tessera/npy/{version}/manifest.parquet`) ```python # Use local manifest file @@ -556,7 +558,7 @@ print(manifest.head()) 2. **Request tiles for bbox** → Query DataFrame for tiles in region 3. **Filter by year and variant** → Select tiles matching the requested year/variant 4. **Find available tiles** → Return list of matching tiles -5. **Anonymous S3 download** → Fetch tiles on demand into the output directory, verified with CRC64NVMe +5. **HTTPS download** → Fetch tiles on demand from the Source Cooperative mirror into the output directory, with integrity checks 6. **Persist** → Downloaded tiles stay in the output directory and are skipped on rerun ## Data Organization @@ -564,19 +566,23 @@ print(manifest.head()) ### Tessera Data Structure ``` -Remote Server (https://s3.us-west-2.amazonaws.com/tessera-embeddings) -├── v1/ # Dataset version 1.0 -│ ├── manifest.parquet # Per-version tile manifest -│ ├── landmasks.parquet # Landmask manifest -│ ├── global_0.1_degree_representation/ # vultr variant (default) +Remote Server (https://data.source.coop/tessera/tessera) +├── npy/ # NPY embeddings + scales +│ ├── v1/ # Dataset version 1.0 +│ │ ├── manifest.parquet # Per-version tile manifest │ │ └── 2024/grid_0.15_52.05/grid_0.15_52.05{,_scales}.npy -│ └── global_0.1_degree_tiff_all/ -│ └── grid_0.15_52.05.tiff # Landmask with projection info -└── v1.1/ # Dataset version 1.1 - ├── manifest.parquet - ├── landmasks.parquet - └── global_0.1_degree_representation.cambridge/ - └── 2024/grid_0.15_52.05/grid_0.15_52.05{,_scales}.npy +│ └── v1.1/ # Dataset version 1.1 +│ ├── manifest.parquet +│ └── 2024/grid_0.15_52.05/grid_0.15_52.05{,_scales}.npy +├── landmasks/ # Landmask TIFFs +│ ├── v1/ +│ │ ├── landmasks.parquet # Landmask manifest +│ │ └── grid_0.15_52.05.tiff # Landmask with projection info +│ └── v1.1/ +│ ├── landmasks.parquet +│ └── grid_0.15_52.05.tiff +└── zarr/ # Cloud-native zarr store + └── v1/ ``` ### Local Cache Structure @@ -634,7 +640,7 @@ When `cache_dir` is not specified, the registry is cached in platform-appropriat ## Hash Verification -GeoTessera verifies end-to-end CRC64NVMe checksums for all downloaded files (embeddings, scales, and landmasks) against S3's `x-amz-checksum-crc64nvme` response header to ensure data integrity. This check is always enforced: a download whose checksum does not match — or whose S3 object is missing the checksum header — is rejected rather than used, so corrupt or truncated files never reach the cache. +GeoTessera verifies every downloaded file (embeddings, scales, and landmasks) against the response `Content-Length`, and additionally against an MD5 computed over the streamed body whenever the server's `ETag` is a content MD5 (i.e. a single-part upload — this covers landmask TIFFs and scales files; large multipart-uploaded embedding tiles carry a composite ETag that is not a content hash, so they are length-checked only). A mismatch rejects the download and triggers a retry with backoff, so corrupt or truncated files never reach the cache. ## Contributing diff --git a/docs/architecture.rst b/docs/architecture.rst index b59ede6..1fb955f 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -31,14 +31,14 @@ The library follows a layered architecture: └── Visualization (rendering and web maps) ↓ Data Access Layer - ├── Anonymous S3 downloads (botocore) + ├── HTTPS downloads (urllib, no cloud SDK) ├── Zarr v3 store (cloud-native streaming) ├── Rasterio (GeoTIFF I/O) └── GeoPandas (geospatial operations) ↓ Storage Layer - ├── Remote servers (https://s3.us-west-2.amazonaws.com/tessera-embeddings) - ├── Zarr store (https://s3.us-west-2.amazonaws.com/tessera-embeddings/v1/zarr) + ├── Source Cooperative repository (https://data.source.coop/tessera/tessera) + ├── Zarr store (https://data.source.coop/tessera/tessera/zarr/v1) └── Local cache (~/.cache/geotessera/{v1,v1.1}/manifest.parquet) Coordinate System and Grid @@ -201,7 +201,8 @@ The registry uses one **Parquet manifest per dataset version** for efficient data discovery and querying. The manifest is filtered by ``(version, variant)`` at load time, then queried by lat/lon/year: -**Manifest Structure** (``s3://tessera-embeddings/{v1,v1.1}/manifest.parquet``): +**Manifest Structure** +(``data.source.coop/tessera/tessera/npy/{v1,v1.1}/manifest.parquet``): .. code-block:: @@ -284,18 +285,20 @@ The manifest can be loaded from multiple sources: Data Access Layer ----------------- -S3 Downloads -~~~~~~~~~~~~ +HTTPS Downloads +~~~~~~~~~~~~~~~ -GeoTessera streams tiles directly from the public S3 bucket using anonymous -(unsigned) ``botocore`` requests: +GeoTessera streams everything — manifests, embedding tiles, and landmasks — +over plain HTTPS (stdlib ``urllib``, no cloud SDK) from the public Source +Cooperative repository: **Features**: - **Per-output-dir mirroring**: Tiles land in the user-supplied ``--output`` directory and persist there for re-use across runs -- **Integrity checking**: End-to-end CRC64NVMe verified against S3's - ``x-amz-checksum-crc64nvme`` response header during the body stream +- **Integrity checking**: Every download is verified against the response + ``Content-Length``, and against a streamed MD5 whenever the server's + ``ETag`` is a content MD5 (single-part uploads) - **Conditional caching**: Per-version manifests use ``If-None-Match`` / ``ETag`` sidecars so unchanged manifests yield a 304 with zero body - **Progress callbacks**: Real-time download feedback with speed and size info diff --git a/docs/index.rst b/docs/index.rst index db7348e..7d1e8d1 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -189,7 +189,7 @@ Data Flow ↓ Per-version Manifest Lookup (filter manifest.parquet by year/lon/lat/variant) ↓ - Anonymous S3 Downloads (with CRC64NVMe verification on the wire) + HTTPS Downloads from Source Cooperative (with integrity checks on the wire) ├── embedding.npy (int8 quantized) → output_dir └── embedding_scales.npy (float32 scale factors) → output_dir ↓ @@ -209,7 +209,8 @@ Manifest System GeoTessera uses a Parquet-based per-version manifest for efficient data access: -* **One manifest per dataset version**: ``s3://tessera-embeddings/{v1,v1.1}/manifest.parquet``. +* **One manifest per dataset version**: + ``data.source.coop/tessera/tessera/npy/{v1,v1.1}/manifest.parquet``. Each carries the file-scan inventory schema (``year, lon, lat, grid_size, scales_size, grid_path, ...``) plus explicit ``version`` and ``variant`` columns so a single file covers every variant in that version. @@ -218,8 +219,9 @@ GeoTessera uses a Parquet-based per-version manifest for efficient data access: * **Conditional fetches**: Per-version ETag sidecars enable ``If-None-Match`` conditional GETs — refetches only happen when the bucket's ETag actually changes; otherwise the server returns 304 with no body. -* **Integrity checking**: End-to-end CRC64NVMe verification using S3's - ``x-amz-checksum-crc64nvme`` response header on every download. +* **Integrity checking**: Every download is verified against the response + ``Content-Length``, and against a streamed MD5 whenever the server's + ``ETag`` is a content MD5 (single-part uploads). The manifest can be loaded from multiple sources: @@ -257,10 +259,10 @@ GeoTessera ships embeddings under two orthogonal axes: Vultr); ``cambridge`` is a test deployment by the Cambridge team for the 1.1 line. -Currently published combinations on ``s3://tessera-embeddings/``: +Currently published combinations on ``data.source.coop/tessera/tessera``: +-------------+------------------+--------------+----------------+----------------------------------------------------------------+ -| ``version`` | ``S3 path`` | ``variant`` | Years | Notes | +| ``version`` | ``path`` | ``variant`` | Years | Notes | +=============+==================+==============+================+================================================================+ | ``1.0`` | ``v1/`` | ``vultr`` | 2017–2025 | Legacy production line. Frozen — no new years will be added. | +-------------+------------------+--------------+----------------+----------------------------------------------------------------+ @@ -358,21 +360,25 @@ layer toggles. See the CLI reference for the full flag set. Data Organization ----------------- -**Remote Server Structure** (S3, ``us-west-2``):: +**Remote Server Structure** (Source Cooperative):: - https://s3.us-west-2.amazonaws.com/tessera-embeddings/ - ├── v1/ # Dataset version 1.0 - │ ├── manifest.parquet # Per-version tile manifest - │ ├── landmasks.parquet # Landmask manifest - │ ├── global_0.1_degree_representation/ # vultr variant (default) + https://data.source.coop/tessera/tessera/ + ├── npy/ # NPY embeddings + scales + │ ├── v1/ # Dataset version 1.0 + │ │ ├── manifest.parquet # Per-version tile manifest │ │ └── 2024/grid_0.15_52.05/grid_0.15_52.05{,_scales}.npy - │ └── global_0.1_degree_tiff_all/ - │ └── grid_0.15_52.05.tiff # Landmask TIFF - └── v1.1/ # Dataset version 1.1 - ├── manifest.parquet - ├── landmasks.parquet # Copy of v1's (same grid) - └── global_0.1_degree_representation.cambridge/ - └── 2024/grid_0.15_52.05/grid_0.15_52.05{,_scales}.npy + │ └── v1.1/ # Dataset version 1.1 + │ ├── manifest.parquet + │ └── 2024/grid_0.15_52.05/grid_0.15_52.05{,_scales}.npy + ├── landmasks/ # Landmask TIFFs + │ ├── v1/ + │ │ ├── landmasks.parquet # Landmask manifest + │ │ └── grid_0.15_52.05.tiff + │ └── v1.1/ + │ ├── landmasks.parquet # Per-version (superset of v1) + │ └── grid_0.15_52.05.tiff + └── zarr/ # Cloud-native zarr store + └── v1/ Each ``manifest.parquet`` is scoped to one version and lists every ``(year, lon, lat)`` tile available for that version's variants. The diff --git a/geotessera/cli.py b/geotessera/cli.py index f90876a..8e7f96c 100644 --- a/geotessera/cli.py +++ b/geotessera/cli.py @@ -833,15 +833,15 @@ def country_progress_callback(current: int, total: int, status: str = None): region_file_to_use = country_geojson_file if by_source: - # Multi-source render. With per-version manifests on S3 we - # download one manifest per requested dataset version and - # concat them in the renderer. + # Multi-source render. With per-version manifests we download + # one manifest per requested dataset version and concat them + # in the renderer. from geotessera.visualization import visualize_sources_coverage from geotessera.registry import ( _parse_dataset_version, download_file_to_temp, KNOWN_VERSIONS, - TESSERA_BASE_URL, + manifest_url, ) if version_spec.lower() == "all": @@ -860,7 +860,7 @@ def country_progress_callback(current: int, total: int, status: str = None): # Already downloaded during GeoTessera init. manifest_paths.append(gt.registry.manifest_path) continue - url = f"{TESSERA_BASE_URL}/{vp}/manifest.parquet" + url = manifest_url(vp) try: manifest_paths.append( Path(download_file_to_temp(url, cache_path=cache_path)) diff --git a/geotessera/registry.py b/geotessera/registry.py index 2826f06..1a037b3 100644 --- a/geotessera/registry.py +++ b/geotessera/registry.py @@ -16,12 +16,6 @@ import numpy as np import time from datetime import datetime, timezone -from urllib.parse import urlparse - -import botocore.session -from botocore import UNSIGNED -from botocore.config import Config -from botocore.exceptions import ClientError try: import pandas as pd @@ -81,7 +75,7 @@ def _version_path_from_norm(norm: str) -> str: return f"v{major}.{minor}" -# Well-known dataset versions on the public bucket. Used by the client when a +# Well-known published dataset versions. Used by the client when a # multi-version operation (e.g. ``coverage --by-source --dataset-version=all``) # needs to enumerate manifests without a separate listing call. Extend this as # new versions are published. @@ -105,7 +99,6 @@ def write_tessera_metadata( downstream tools can recover which dataset was downloaded. """ import json - from datetime import datetime, timezone version_path, version_norm = _parse_dataset_version(dataset_version) payload: Dict[str, object] = { @@ -114,9 +107,7 @@ def write_tessera_metadata( "dataset_variant": dataset_variant, "embeddings_subdir": EMBEDDINGS_DIR_NAME, "s3_embeddings_subdir": _variant_subdir(dataset_variant), - "source_url_prefix": ( - f"{TESSERA_BASE_URL}/{version_path}/{_variant_subdir(dataset_variant)}/" - ), + "source_url_prefix": f"{TESSERA_NPY_MIRROR_URL}/{version_path}/", "generated_at": datetime.now(timezone.utc).isoformat(), } if extra: @@ -354,66 +345,35 @@ def tile_to_bounds(lon: float, lat: float) -> Tuple[float, float, float, float]: return (lon - 0.05, lat - 0.05, lon + 0.05, lat + 0.05) -# Base URL for Tessera data downloads -TESSERA_BASE_URL = "https://s3.us-west-2.amazonaws.com/tessera-embeddings" - -# Directory structure constants (mirrors remote structure) +# All Tessera data is served from the public Source Cooperative repository, +# fetched with plain HTTPS. Layout: one tree per media type, one subdir per +# dataset version: +# npy/{version_path}/{year}/grid_.../grid_....npy embeddings + scales +# npy/{version_path}/manifest.parquet per-version manifest +# landmasks/{version_path}/grid_....tiff landmask TIFFs +# landmasks/{version_path}/landmasks.parquet landmask registry +# Unlike the retired S3 bucket there is no +# global_0.1_degree_representation[.] level: the repository carries +# one embedding tree per version. Versions appear incrementally as they are +# uploaded (npy has v1/v1.1/v2; landmasks has v1 with v1.1 to follow). +TESSERA_MIRROR_URL = "https://data.source.coop/tessera/tessera" +TESSERA_NPY_MIRROR_URL = f"{TESSERA_MIRROR_URL}/npy" +TESSERA_LANDMASKS_MIRROR_URL = f"{TESSERA_MIRROR_URL}/landmasks" + +# Directory structure constants (local embeddings_dir layout, kept from the +# original bucket structure) EMBEDDINGS_DIR_NAME = "global_0.1_degree_representation" # NPY embeddings and scales LANDMASKS_DIR_NAME = "global_0.1_degree_tiff_all" # Landmask TIFFs -# Note: Default manifest URLs are constructed with version in Registry.__init__ -# Format: {TESSERA_BASE_URL}/{version}/manifest.parquet - -# Anonymous (unsigned) S3 clients, one per region, built lazily and reused. -# The Tessera bucket is public-read, so requests are not signed. botocore's -# "standard" retry mode replaces the old hand-rolled exponential backoff. -_S3_CLIENTS: Dict[str, object] = {} +def manifest_url(version_path: str) -> str: + """Default URL of the embeddings manifest for *version_path* (e.g. ``"v1"``).""" + return f"{TESSERA_NPY_MIRROR_URL}/{version_path}/manifest.parquet" -def _s3_client(region: str): - """Return a cached anonymous S3 client for *region*.""" - client = _S3_CLIENTS.get(region) - if client is None: - client = botocore.session.get_session().create_client( - "s3", - region_name=region, - config=Config( - signature_version=UNSIGNED, - retries={"mode": "standard", "total_max_attempts": 5}, - ), - ) - _S3_CLIENTS[region] = client - return client - - -def _parse_s3_url(url: str) -> Tuple[str, str, str]: - """Parse an S3 HTTPS URL into ``(region, bucket, key)``. - - Handles both path-style (``s3..amazonaws.com//``) and - virtual-hosted (``.s3..amazonaws.com/``) URLs — all - geotessera ever builds from :data:`TESSERA_BASE_URL`. Raises ``ValueError`` - for anything that is not a regional S3 endpoint; the data path is S3-only, - so use ``registry_path`` for local files. - """ - parsed = urlparse(url) - host = parsed.netloc - path = parsed.path.lstrip("/") - m = re.match( - r"^(?:(?P[^.]+)\.)?s3[.-](?P[a-z0-9-]+)\.amazonaws\.com$", - host, - ) - if not m: - raise ValueError(f"Not a recognized regional S3 URL: {url!r}") - region = m.group("region") - bucket = m.group("bucket") - if bucket: - key = path # virtual-hosted: the whole path is the key - else: - bucket, _, key = path.partition("/") # path-style: first segment is bucket - if not bucket or not key: - raise ValueError(f"Could not extract bucket/key from S3 URL: {url!r}") - return region, bucket, key +def landmasks_parquet_url(version_path: str) -> str: + """Default URL of the landmasks registry for *version_path* (e.g. ``"v1"``).""" + return f"{TESSERA_LANDMASKS_MIRROR_URL}/{version_path}/landmasks.parquet" def format_bytes(num_bytes: float) -> str: @@ -430,17 +390,20 @@ def download_file_to_temp( progress_callback: Optional[Callable[[int, int, str], None]] = None, cache_path: Optional[Path] = None, ) -> str: - """Download an object from the public Tessera S3 bucket, with caching. - - Fetches via an anonymous botocore S3 client. Integrity is verified - end-to-end against the object's CRC64NVMe checksum: ``ChecksumMode=ENABLED`` - makes botocore validate it as the body is streamed and raise on a mismatch. - When *cache_path* already exists, an ``If-Modified-Since`` conditional GET - (keyed on the cached file's mtime, which was set from the previous - ``Last-Modified``) short-circuits to the cached copy on a 304. + """Download a file over HTTPS with caching, retries, and integrity checks. + + Integrity is verified against the response ``Content-Length`` (a short read + raises and triggers a retry) and, when the server's ``ETag`` is a plain MD5 + (single-part S3 upload — 32 hex digits with no multipart ``-`` + suffix), against an MD5 computed over the streamed body. When *cache_path* + already exists, an ``If-Modified-Since`` conditional GET (keyed on the + cached file's mtime, which was set from the previous ``Last-Modified``) + short-circuits to the cached copy on a 304. Transient failures (429/5xx, + connection errors, truncation, checksum mismatch) are retried with + exponential backoff. Args: - url: HTTPS S3 URL (path-style or virtual-hosted). + url: HTTPS URL. progress_callback: Optional callback(bytes_downloaded, total_bytes, status). cache_path: Optional destination. When given, the file is written here atomically and reused on later calls; when omitted it goes to a @@ -451,48 +414,119 @@ def download_file_to_temp( cache hit, or a temporary path when ``cache_path`` is None. Raises: - botocore.exceptions.ClientError: On non-304 S3 errors (after retries). - botocore.exceptions.FlexibleChecksumError: On CRC64NVMe mismatch. - ValueError: If the URL is not S3 or the object has no CRC64NVMe checksum. + urllib.error.HTTPError: On non-304 HTTP errors (after retries). + OSError: On a truncated, corrupted, or failed download (after retries). """ - import tempfile + import urllib.error - region, bucket, key = _parse_s3_url(url) - client = _s3_client(region) + attempts = 4 + for attempt in range(attempts): + try: + return _download_once(url, progress_callback, cache_path) + except urllib.error.HTTPError as e: + if e.code not in (429, 500, 502, 503, 504) or attempt == attempts - 1: + raise + except (urllib.error.URLError, OSError): + if attempt == attempts - 1: + raise + time.sleep(2**attempt) - get_kwargs = {"Bucket": bucket, "Key": key, "ChecksumMode": "ENABLED"} - # Conditional GET: the cached file's mtime was set from the object's prior - # Last-Modified, so If-Modified-Since lets S3 answer 304 when it is - # unchanged. (This replaces the old ETag ``.etag`` sidecar, which predated - # the move to S3 and is no longer needed.) +def _download_once( + url: str, + progress_callback: Optional[Callable[[int, int, str], None]], + cache_path: Optional[Path], +) -> str: + """Single download attempt (see :func:`download_file_to_temp`).""" + import hashlib + import urllib.error + import urllib.request + from email.utils import formatdate, parsedate_to_datetime + + from . import __version__ + + # Cloudflare (fronting the Source Cooperative mirror) rejects the default + # Python-urllib User-Agent with 403, so identify as geotessera. + request = urllib.request.Request( + url, headers={"User-Agent": f"geotessera/{__version__}"} + ) if cache_path and cache_path.exists(): - get_kwargs["IfModifiedSince"] = datetime.fromtimestamp( - cache_path.stat().st_mtime, tz=timezone.utc + request.add_header( + "If-Modified-Since", + formatdate(cache_path.stat().st_mtime, usegmt=True), ) try: - response = client.get_object(**get_kwargs) - except ClientError as e: - meta = e.response.get("ResponseMetadata", {}) - code = e.response.get("Error", {}).get("Code") - if code == "304" or meta.get("HTTPStatusCode") == 304: + response = urllib.request.urlopen(request, timeout=60) + except urllib.error.HTTPError as e: + if e.code == 304: # 304 Not Modified — the cached copy is current. if progress_callback: progress_callback(0, 0, "Cache is current") return str(cache_path) raise - # botocore silently returns an *unvalidated* body when the object carries no - # CRC64NVMe checksum header, so require it explicitly rather than trust an - # unverified download. - if not response.get("ChecksumCRC64NVME"): - raise ValueError( - f"S3 object {key!r} returned no CRC64NVMe checksum; " - "refusing to use an unverified download" + try: + total_size = int(response.headers.get("Content-Length") or 0) + + last_modified = None + lm_header = response.headers.get("Last-Modified") + if lm_header: + try: + last_modified = parsedate_to_datetime(lm_header) + except (TypeError, ValueError): + pass + + # A plain 32-hex-digit ETag is the object's content MD5 (single-part + # S3 upload). Multipart uploads get a "-" suffix and are not a + # content hash, so those downloads are only length-checked. + etag = (response.headers.get("ETag") or "").strip('"') + md5 = hashlib.md5() if re.fullmatch(r"[0-9a-f]{32}", etag) else None + + def chunks(): + downloaded = 0 + while True: + chunk = response.read(8192) + if not chunk: + break + downloaded += len(chunk) + if md5 is not None: + md5.update(chunk) + yield chunk + # A dropped connection can end the stream early without an + # exception; catch it so a truncated file is never cached. + if total_size and downloaded != total_size: + raise OSError( + f"Truncated download from {url}: " + f"got {downloaded} of {total_size} bytes" + ) + if md5 is not None and md5.hexdigest() != etag: + raise OSError( + f"Checksum mismatch for {url}: " + f"MD5 {md5.hexdigest()} != ETag {etag}" + ) + + return _stream_to_file( + chunks(), total_size, last_modified, cache_path, progress_callback ) + finally: + response.close() + - total_size = response.get("ContentLength", 0) +def _stream_to_file( + chunks, + total_size: int, + last_modified, + cache_path: Optional[Path], + progress_callback: Optional[Callable[[int, int, str], None]], +) -> str: + """Stream *chunks* to *cache_path* (or a temp file) atomically. + + Writes to a sibling temp file with progress reporting, stamps the mtime + from *last_modified* (a datetime) so later ``If-Modified-Since`` requests + work, then moves the file into place. Returns the final path as ``str``. + """ + import tempfile if cache_path: cache_path.parent.mkdir(parents=True, exist_ok=True) @@ -507,8 +541,6 @@ def download_file_to_temp( temp_file = tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=".npy") temp_path = Path(temp_file.name) - - body = response["Body"] success = False try: downloaded = 0 @@ -519,9 +551,7 @@ def download_file_to_temp( size_str = format_bytes(total_size) if total_size > 0 else "unknown size" progress_callback(0, total_size, f"Starting ({size_str})") - # Reading to end-of-stream triggers botocore's CRC64NVMe validation, - # which raises FlexibleChecksumError on a mismatch. - for chunk in body.iter_chunks(8192): + for chunk in chunks: temp_file.write(chunk) downloaded += len(chunk) @@ -546,8 +576,6 @@ def download_file_to_temp( temp_file.close() # Set file mtime from Last-Modified so the next If-Modified-Since works. - # boto hands LastModified back as a datetime, so there's no parsing. - last_modified = response.get("LastModified") if last_modified is not None: try: ts = last_modified.timestamp() @@ -576,7 +604,6 @@ def download_file_to_temp( return str(final_path) finally: - body.close() # Remove the partial temp file on any failure, including # KeyboardInterrupt/SystemExit (a BaseException, which a plain # `except Exception` misses) — this is what left stray @@ -690,7 +717,7 @@ def __init__( # matching its dataset_version and filters by variant on load. self._registry_gdf: Optional[gpd.GeoDataFrame] = None self._registry_url = ( - registry_url or f"{TESSERA_BASE_URL}/{self._version_path}/manifest.parquet" + registry_url or manifest_url(self._version_path) ) self._registry_path = Path(registry_path) if registry_path else None @@ -698,7 +725,7 @@ def __init__( self._landmasks_df: Optional[pd.DataFrame] = None self._landmasks_registry_url = ( landmasks_registry_url - or f"{TESSERA_BASE_URL}/{self._version_path}/landmasks.parquet" + or landmasks_parquet_url(self._version_path) ) self._landmasks_registry_path = ( Path(landmasks_registry_path) if landmasks_registry_path else None @@ -913,8 +940,8 @@ def _load_landmasks_registry(self): return # Validate landmasks registry structure. Hash columns are no longer - # required — integrity is verified via the S3 CRC64NVMe header at - # download time. + # required — integrity is verified against Content-Length and the + # MD5 ETag at download time. if self._landmasks_df is not None: required_columns = {"lat", "lon", "file_size"} if not required_columns.issubset(self._landmasks_df.columns): @@ -1148,14 +1175,11 @@ def fetch( # Use existing local file return str(local_path) - # Download to embeddings_dir. Integrity is verified end-to-end against - # the S3 x-amz-checksum-crc64nvme response header inside the downloader. + # Download to embeddings_dir from the Source Cooperative repository, + # which carries one embedding tree per version (no variant subdir). # Use as_posix() to ensure forward slashes in URL even on Windows path_str = path.as_posix() if isinstance(path, Path) else path - url = ( - f"{TESSERA_BASE_URL}/{self._version_path}/" - f"{self._embeddings_subdir}/{path_str}" - ) + url = f"{TESSERA_NPY_MIRROR_URL}/{self._version_path}/{path_str}" downloaded_path = download_file_to_temp( url, progress_callback=progress_callback, @@ -1200,18 +1224,15 @@ def fetch_landmask( # Use existing local file return str(local_path) - # Download to embeddings_dir. Integrity is verified end-to-end against - # the S3 x-amz-checksum-crc64nvme response header inside the downloader. - url = f"{TESSERA_BASE_URL}/{self._version_path}/{LANDMASKS_DIR_NAME}/{filename}" - downloaded_path = download_file_to_temp( + # Download to embeddings_dir from the Source Cooperative mirror (which + # flattens the LANDMASKS_DIR_NAME subdir away). + url = f"{TESSERA_LANDMASKS_MIRROR_URL}/{self._version_path}/{filename}" + return download_file_to_temp( url, progress_callback=progress_callback, cache_path=local_path, ) - # Return path to saved file - return downloaded_path - @property def available_embeddings(self) -> List[Tuple[int, float, float]]: """Get list of available embeddings.""" diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index a7ac0f1..126ab93 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -4431,7 +4431,7 @@ def main(): ) verify_parser.add_argument( "--store", - default="https://s3.us-west-2.amazonaws.com/tessera-embeddings/v1/zarr", + default="https://data.source.coop/tessera/tessera/zarr/v1", help="Zarr store URL", ) verify_parser.set_defaults(func=verify_tile_command) @@ -4446,7 +4446,7 @@ def main(): print_parser.add_argument("--year", type=int, required=True, help="Year") print_parser.add_argument( "--store", - default="https://s3.us-west-2.amazonaws.com/tessera-embeddings/v1/zarr", + default="https://data.source.coop/tessera/tessera/zarr/v1", help="Zarr store URL", ) print_parser.set_defaults(func=print_command) diff --git a/geotessera/store.py b/geotessera/store.py index 40e9afa..e8c1c7c 100644 --- a/geotessera/store.py +++ b/geotessera/store.py @@ -34,7 +34,7 @@ log = logging.getLogger(__name__) -DEFAULT_STORE = "https://s3.us-west-2.amazonaws.com/tessera-embeddings/v1/zarr" +DEFAULT_STORE = "https://data.source.coop/tessera/tessera/zarr/v1" # Shard-aligned chunk sizes so dask tasks match zarr shards SHARD_CHUNKS = {"time": 1, "band": 128, "y": 4096, "x": 4096} @@ -341,7 +341,7 @@ class GeoTesseraZarr: Args: store_url: Zarr store URL or local path. Defaults to the public - TESSERA store at ``s3.us-west-2.amazonaws.com/tessera-embeddings``. + TESSERA store at ``data.source.coop/tessera/tessera/zarr``. Example:: diff --git a/geotessera/visualization.py b/geotessera/visualization.py index d5ad919..9da1955 100644 --- a/geotessera/visualization.py +++ b/geotessera/visualization.py @@ -146,8 +146,8 @@ def visualize_sources_coverage( Args: manifest_path: Local path to a manifest parquet — or a list of paths which are concatenated before rendering. With per-version manifests - on S3 (``s3://tessera-embeddings/{v}/manifest.parquet``), pass a - list to compare versions on a single map. + (``data.source.coop/tessera/tessera/npy/{v}/manifest.parquet``), + pass a list to compare versions on a single map. output_path: Output PNG path. year: Optional year filter (applies to all sources). width_pixels: Output image width in pixels. diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 55eae51..17a7aa7 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -275,8 +275,10 @@ def gather_tile_infos( f"Landmask directory not found for {registry._version_path}: " f"expected {lm_s3_mirror}. Landmasks are per-version and " f"cannot be reused across versions. Fetch them with:\n" - f" aws s3 sync s3://tessera-embeddings/" - f"{registry._version_path}/{LANDMASKS_DIR_NAME}/ {lm_s3_mirror}/" + f" aws s3 sync --no-sign-request " + f"--endpoint-url https://data.source.coop " + f"s3://tessera/tessera/landmasks/" + f"{registry._version_path}/ {lm_s3_mirror}/" ) else: base_emb = str(registry._embeddings_dir / EMBEDDINGS_DIR_NAME) @@ -286,8 +288,10 @@ def gather_tile_infos( raise FileNotFoundError( f"Landmask directory not found: expected {lm_flat}. " f"Fetch them with:\n" - f" aws s3 sync s3://tessera-embeddings/" - f"{registry._version_path}/{LANDMASKS_DIR_NAME}/ {lm_flat}/" + f" aws s3 sync --no-sign-request " + f"--endpoint-url https://data.source.coop " + f"s3://tessera/tessera/landmasks/" + f"{registry._version_path}/ {lm_flat}/" ) zones_dict: Dict[int, List[TileInfo]] = {} transformer_cache: Dict[int, ProjTransformer] = {} diff --git a/pyproject.toml b/pyproject.toml index 4001b6c..f6d42ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,8 +41,6 @@ dependencies = [ "aiohttp", "geozarr-toolkit", "contextily", - "botocore>=1.43.14", - "awscrt>=0.33.0", ] [project.urls] diff --git a/tests/v11.t b/tests/v11.t index 6dcc29d..6f17b69 100644 --- a/tests/v11.t +++ b/tests/v11.t @@ -117,8 +117,8 @@ Confirm the consumer fetches from the per-version path on S3: > print('v1 URL :', r1._registry_url) > print('v1.1 :', r2._registry_url) > " 2>&1 | grep -E '^v1' - v1 URL : https://s3.us-west-2.amazonaws.com/tessera-embeddings/v1/manifest.parquet - v1.1 : https://s3.us-west-2.amazonaws.com/tessera-embeddings/v1.1/manifest.parquet + v1 URL : https://data.source.coop/tessera/tessera/npy/v1/manifest.parquet + v1.1 : https://data.source.coop/tessera/tessera/npy/v1.1/manifest.parquet Test: S3 Embeddings Subdir Reflects Variant ------------------------------------------- diff --git a/uv.lock b/uv.lock index c1fbcc6..2aa1c23 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -177,38 +177,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] -[[package]] -name = "awscrt" -version = "0.34.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ce/33/ed6d3c91d7b136a91eeea3bea1023e818f5f99d4fcbd8956c645a2dc6006/awscrt-0.34.1.tar.gz", hash = "sha256:a3ae8e35c3a3eefdb2a15859887a05b926d0456d21ccba1b49861cfe46bcc8c2", size = 36975094, upload-time = "2026-06-04T19:03:14.475Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/90/34/5737d2c0d1288caa76cb8ae867d4c09a86e31fc03e4f5a9e595308763415/awscrt-0.34.1-cp311-abi3-macosx_10_15_universal2.whl", hash = "sha256:c7c9c2ef4ec0db16d871fa099158b0617b1b8aa53335b73a88830a6630945c5a", size = 5089351, upload-time = "2026-06-04T19:02:11.912Z" }, - { url = "https://files.pythonhosted.org/packages/68/0b/79c55c606d83308cfa6d78528497ab72a170494a196261d23bd92ac733e8/awscrt-0.34.1-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cc7b55bcc719c2ecf6876b8c13a4808a8dda3b19b8936329c711072adc35053c", size = 3919082, upload-time = "2026-06-04T19:02:14.239Z" }, - { url = "https://files.pythonhosted.org/packages/4b/aa/494385f153a90e734ac04597e1c69ca6f66d8579bb0383f6fc8c5aa126d6/awscrt-0.34.1-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c696acaa421e74d7dd44962099099ae967c64c8491f797fe1ec7c32fea1c113", size = 4209970, upload-time = "2026-06-04T19:02:15.698Z" }, - { url = "https://files.pythonhosted.org/packages/5e/93/e63b09fbdef626e19edfe1bf28ac037153270bceee24327756c83c53c940/awscrt-0.34.1-cp311-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:4ab254a7761bb80594eba4ecdae65ba55608dad7ae08f5c0e6876cdf339e8ffe", size = 3827246, upload-time = "2026-06-04T19:02:17.148Z" }, - { url = "https://files.pythonhosted.org/packages/6a/43/edc9a8704f84aaffc3471b69d8da565fa834e8a7f565ced659a78373eaf4/awscrt-0.34.1-cp311-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:48b24b3b3de9c3392238181f4529b8d3511fee695f1852807849042681bb4922", size = 4066655, upload-time = "2026-06-04T19:02:18.411Z" }, - { url = "https://files.pythonhosted.org/packages/b9/e7/6c49160a8cc84c438cac59c2bd52f80eaf6db05da036495133349b68a351/awscrt-0.34.1-cp311-abi3-win32.whl", hash = "sha256:8356df183947a02a015b8b77ee69a1b24bddf2177db93fb063e6215265ecc33a", size = 4089039, upload-time = "2026-06-04T19:02:20.023Z" }, - { url = "https://files.pythonhosted.org/packages/74/49/8ae6a3bb2aa32e17c3b8e7ad14f9666edfa924052c22e0098ca525b110c9/awscrt-0.34.1-cp311-abi3-win_amd64.whl", hash = "sha256:40326b5fadafe679c0729f622c4a56cbcb5e91f1b0d2b84fc95f016427b11214", size = 4251132, upload-time = "2026-06-04T19:02:21.308Z" }, - { url = "https://files.pythonhosted.org/packages/80/d4/dd647d569189bc97f5110339f7ec8ee669d135408c3e2207ce74f951f284/awscrt-0.34.1-cp313-abi3-macosx_10_15_universal2.whl", hash = "sha256:3750622e9efa1af389a619573a93cf1f7e816454edbdd9c922ba78e6c292668c", size = 5088559, upload-time = "2026-06-04T19:02:22.937Z" }, - { url = "https://files.pythonhosted.org/packages/d1/c6/6ae127a3ca3ac9e7e4bd2c351be78856d0fe4b9a094cf9163e5f9ba30d76/awscrt-0.34.1-cp313-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6eccf12282f773f2d29ef0bb16c3b973976c578ea617d84a8c1d523509d2b5a4", size = 3909549, upload-time = "2026-06-04T19:02:24.62Z" }, - { url = "https://files.pythonhosted.org/packages/e4/8f/93df8d004dee2afb938111de203e4534101326c1affc1452eb8d6fb73da9/awscrt-0.34.1-cp313-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a39a0c2fed3b7354053b70a4a62428fa40834c9e9ac8570c8b39d3868645bc7", size = 4202624, upload-time = "2026-06-04T19:02:26.003Z" }, - { url = "https://files.pythonhosted.org/packages/70/f0/b0d96a536c582724aa5ebcfbe03a0c0283adb9379a31465076b8988416cd/awscrt-0.34.1-cp313-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:775405da040f8f5223870c389e6b37ad2909849d76a1f028fc713ad1c07889f8", size = 3818941, upload-time = "2026-06-04T19:02:27.398Z" }, - { url = "https://files.pythonhosted.org/packages/f1/16/92c10101b3f4c10e8d8d6b3ce8b6bd1254cff4c1925244c4e872026c17a4/awscrt-0.34.1-cp313-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:d88cc48538efabbf2b6643c13a8b63a758dcf40eeb0405f6ba4066ae4d361866", size = 4061435, upload-time = "2026-06-04T19:02:28.824Z" }, - { url = "https://files.pythonhosted.org/packages/57/cf/9f1610d5a4d8ed6c5062fea66d9c5bfda6fa354690b6c2121966b1be24b5/awscrt-0.34.1-cp313-abi3-win32.whl", hash = "sha256:534e0d9b415c50395b0602bb5bbc3b506fc36184859d4b6de18642384e46cb17", size = 4087600, upload-time = "2026-06-04T19:02:30.431Z" }, - { url = "https://files.pythonhosted.org/packages/4e/0f/5ac6d4b372cddc9dee02892e2a0ed676fe099d1740b2e5de253f6d9d4825/awscrt-0.34.1-cp313-abi3-win_amd64.whl", hash = "sha256:c680afe517ca09846d6b2abeeb44d37b935f575013afa34d8295fee95cce0b78", size = 4245466, upload-time = "2026-06-04T19:02:31.918Z" }, - { url = "https://files.pythonhosted.org/packages/de/1e/ce97fee2fe477e382f68f70f42846ca8d5b1d908f66e900b45334b9549b8/awscrt-0.34.1-cp313-cp313t-macosx_10_15_universal2.whl", hash = "sha256:b320bbbd70599d55c719f6a1201f0db51fbbdfc08d885ee17d2e553c06d2770d", size = 5097627, upload-time = "2026-06-04T19:02:33.41Z" }, - { url = "https://files.pythonhosted.org/packages/a1/c0/dc5bba8ae0c16965680ca1bf9fffb4914a43076ae8f9a5012db5b822610d/awscrt-0.34.1-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:3f58fb121531fd9d85ff3f1975cabe8b7de2f31d460d6de4f5a624177b4005d7", size = 3958448, upload-time = "2026-06-04T19:02:35.085Z" }, - { url = "https://files.pythonhosted.org/packages/6f/ac/71d009408421a48185f6672f830c49909f3dd076de1984f449dd3085a121/awscrt-0.34.1-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:d96bf355138073d5ecd8df68521c73181968fea82ee38e8aed523eb94ff73131", size = 4198705, upload-time = "2026-06-04T19:02:36.489Z" }, - { url = "https://files.pythonhosted.org/packages/2e/e0/e6cbbbcd71d9498a2cf3a9cce5df14320110d73e4aa6e3180c9029883eb5/awscrt-0.34.1-cp313-cp313t-win32.whl", hash = "sha256:9229d146f94c576778e5f4b41c0254800c8c6ad005f25b20fe07da9df976ce96", size = 4142609, upload-time = "2026-06-04T19:02:38.1Z" }, - { url = "https://files.pythonhosted.org/packages/8f/8a/c5b32222debb90cb5aa2858dd7766bb060fe3b6ce4356349b23f71570020/awscrt-0.34.1-cp313-cp313t-win_amd64.whl", hash = "sha256:f8b54977cf53539e68ff2c9ad2f51bb4bee7f2e72060037482262913e75499b0", size = 4300595, upload-time = "2026-06-04T19:02:39.821Z" }, - { url = "https://files.pythonhosted.org/packages/62/d7/40d6eba06ae5e4bb2e68125009a701a6f74159b96c921cc0998050c76cc4/awscrt-0.34.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:e519f4ba6622a9f5d745e2f49268690f3fac7bad7655911458c80e1cc30bd108", size = 5097638, upload-time = "2026-06-04T19:02:41.438Z" }, - { url = "https://files.pythonhosted.org/packages/22/14/3ddc4d76ca23c892ca82926d73575913377bf59a31a8b20d4c0c329a6a85/awscrt-0.34.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:185f3e899a24961a1bd6e203a0fdeab6a06a972ddf4a4e3a32b20e31705edd95", size = 4039424, upload-time = "2026-06-04T19:02:42.85Z" }, - { url = "https://files.pythonhosted.org/packages/e6/2f/954ac7be390d04160c8570810faeb9512447422752cac2bb277cbb62ab49/awscrt-0.34.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ce386fdaf86f07122305a14a84f5c592b7010d7032c6b5a7dcbf4a72b2704998", size = 4328101, upload-time = "2026-06-04T19:02:44.349Z" }, - { url = "https://files.pythonhosted.org/packages/f8/b0/bd30e6cca8a9b696866dffb587ac1d5bb09dcba2b62941063b02a7b747ae/awscrt-0.34.1-cp314-cp314t-win32.whl", hash = "sha256:dc68aefe66cd419c3d2a0fa0ecaeb9b72d882ad08f752767dc85740052bbf212", size = 4223490, upload-time = "2026-06-04T19:02:45.687Z" }, - { url = "https://files.pythonhosted.org/packages/3c/f5/4cce09065b37c775ab22dc9f179e03704f3bb9bb5a44cae6bc70b849596d/awscrt-0.34.1-cp314-cp314t-win_amd64.whl", hash = "sha256:1421ee18988c946bb058819b69707c85aca0de395520fcbf1ef9ace0b9e96c62", size = 4400512, upload-time = "2026-06-04T19:02:47.242Z" }, -] - [[package]] name = "babel" version = "2.18.0" @@ -218,20 +186,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/f5/21d2de20e8b8b0408f0681956ca2c69f1320a3848ac50e6e7f39c6159675/babel-2.18.0-py3-none-any.whl", hash = "sha256:e2b422b277c2b9a9630c1d7903c2a00d0830c409c59ac8cae9081c92f1aeba35", size = 10196845, upload-time = "2026-02-01T12:30:53.445Z" }, ] -[[package]] -name = "botocore" -version = "1.43.27" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "jmespath" }, - { name = "python-dateutil" }, - { name = "urllib3" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fa/4e/db50ef135f1d9ffc85e209a124004a5829d8f12f4a7a0afdf380cb19866d/botocore-1.43.27.tar.gz", hash = "sha256:2093c316c24214e50e18640b1869513b759bb8cc48b95b004a8306cb9f0d6703", size = 15504242, upload-time = "2026-06-10T19:38:25.389Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2d/46/05b227b34e434b54867c2c942b0bfbbe2fe41789c18bb15ef787d03e9a56/botocore-1.43.27-py3-none-any.whl", hash = "sha256:4976544e652d5a1d8eca135da019f8e1c2d749efa2f9a31a8fb8c76f1895a40b", size = 15190293, upload-time = "2026-06-10T19:38:22.298Z" }, -] - [[package]] name = "certifi" version = "2026.5.20" @@ -693,8 +647,6 @@ version = "0.9.1" source = { editable = "." } dependencies = [ { name = "aiohttp" }, - { name = "awscrt" }, - { name = "botocore" }, { name = "contextily" }, { name = "cram" }, { name = "dask" }, @@ -726,8 +678,6 @@ dev = [ [package.metadata] requires-dist = [ { name = "aiohttp" }, - { name = "awscrt", specifier = ">=0.33.0" }, - { name = "botocore", specifier = ">=1.43.14" }, { name = "contextily" }, { name = "cram", specifier = ">=0.7" }, { name = "dask" }, @@ -846,15 +796,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] -[[package]] -name = "jmespath" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, -] - [[package]] name = "joblib" version = "1.5.3" From 44b453ae34d0583b7ed20eb2681f165d41940c50 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Thu, 30 Jul 2026 17:56:36 +0100 Subject: [PATCH 02/13] cleanups --- CHANGES.md | 9 ++ README.md | 2 +- docs/architecture.rst | 7 +- geotessera/cli.py | 5 +- geotessera/core.py | 12 +- geotessera/registry.py | 238 +++++++++++++++++----------------- geotessera/registry_cli.py | 250 +++++++++++++++++++++++------------- geotessera/store.py | 4 +- geotessera/visualization.py | 6 +- geotessera/zarr.py | 21 +-- tests/v11.t | 2 +- 11 files changed, 317 insertions(+), 239 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 5de3438..393b1e0 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -21,6 +21,15 @@ - **`geotessera-registry zarr-consolidate`**: New subcommand that re-consolidates a store's root metadata after in-place changes. Mostly only for repairs and not regular use. +- **`geotessera-registry s3scan` scans Source Cooperative**: listings are + path-style against the S3-compatible endpoint given by `--endpoint-url` + (default `https://data.source.coop`); the AWS virtual-hosted addressing + for the retired bucket is gone, and `--region` with it. Discovery + understands the flat `{version}/{year}/` layout with the variant supplied + via `--variant`, and `--landmasks-uri` points the landmask scan at a + separate tree (e.g. `s3://tessera/tessera/landmasks/`), including + landmasks-only runs. This allows manifests and landmask registries to be + regenerated directly from the Source Cooperative repository. (@avsm) ## v0.9.0 (2026-06-09) diff --git a/README.md b/README.md index 8efb27a..bea778d 100644 --- a/README.md +++ b/README.md @@ -640,7 +640,7 @@ When `cache_dir` is not specified, the registry is cached in platform-appropriat ## Hash Verification -GeoTessera verifies every downloaded file (embeddings, scales, and landmasks) against the response `Content-Length`, and additionally against an MD5 computed over the streamed body whenever the server's `ETag` is a content MD5 (i.e. a single-part upload — this covers landmask TIFFs and scales files; large multipart-uploaded embedding tiles carry a composite ETag that is not a content hash, so they are length-checked only). A mismatch rejects the download and triggers a retry with backoff, so corrupt or truncated files never reach the cache. +GeoTessera verifies every downloaded file (embeddings, scales, and landmasks) against the response `Content-Length`, and additionally against an MD5 computed over the streamed body whenever the server's `ETag` is a content MD5 (a single-part upload; this covers landmask TIFFs and scales files). Large multipart-uploaded embedding tiles carry a composite ETag that is not a content hash, so they are length-checked only. A mismatch rejects the download and triggers a retry with backoff, so corrupt or truncated files never reach the cache. ## Contributing diff --git a/docs/architecture.rst b/docs/architecture.rst index 1fb955f..0ad97ee 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -288,7 +288,7 @@ Data Access Layer HTTPS Downloads ~~~~~~~~~~~~~~~ -GeoTessera streams everything — manifests, embedding tiles, and landmasks — +GeoTessera streams all data (manifests, embedding tiles, and landmasks) over plain HTTPS (stdlib ``urllib``, no cloud SDK) from the public Source Cooperative repository: @@ -299,8 +299,9 @@ Cooperative repository: - **Integrity checking**: Every download is verified against the response ``Content-Length``, and against a streamed MD5 whenever the server's ``ETag`` is a content MD5 (single-part uploads) -- **Conditional caching**: Per-version manifests use ``If-None-Match`` / - ``ETag`` sidecars so unchanged manifests yield a 304 with zero body +- **Conditional caching**: Per-version manifests are refetched with an + ``If-Modified-Since`` conditional GET keyed on the cached file's mtime, + so unchanged manifests yield a 304 with zero body - **Progress callbacks**: Real-time download feedback with speed and size info - **Resumable**: Existing files in the output dir are skipped on rerun diff --git a/geotessera/cli.py b/geotessera/cli.py index 8e7f96c..62bd3a4 100644 --- a/geotessera/cli.py +++ b/geotessera/cli.py @@ -833,9 +833,8 @@ def country_progress_callback(current: int, total: int, status: str = None): region_file_to_use = country_geojson_file if by_source: - # Multi-source render. With per-version manifests we download - # one manifest per requested dataset version and concat them - # in the renderer. + # Multi-source render. Download one manifest per requested + # dataset version and concatenate them in the renderer. from geotessera.visualization import visualize_sources_coverage from geotessera.registry import ( _parse_dataset_version, diff --git a/geotessera/core.py b/geotessera/core.py index 35cf6c1..abdeac4 100644 --- a/geotessera/core.py +++ b/geotessera/core.py @@ -102,12 +102,12 @@ def __init__( Args: dataset_version: Tessera dataset version. Accepts ``"v1"`` / - ``"1.0"`` or ``"v1.1"`` / ``"1.1"`` (the legacy S3 layout uses + ``"1.0"`` or ``"v1.1"`` / ``"1.1"`` (the repository uses ``v1/`` for the 1.0 series). dataset_variant: Variant of the embeddings to load (default - ``"vultr"``). Other published variants — e.g. ``"cambridge"`` - — are produced by different model runs and live in - ``global_0.1_degree_representation.`` dirs on S3. + ``"vultr"``). Other published variants (e.g. ``"cambridge"``) + are produced by different model runs and are selected by + filtering the manifest. cache_dir: Directory for caching registry files only (not embedding data) embeddings_dir: Directory containing pre-downloaded embedding tiles. Defaults to current working directory if not specified. @@ -164,11 +164,11 @@ def version(self) -> str: @property def embeddings_subdir(self) -> str: - """Variant-aware embeddings subdirectory name (mirrors S3 layout). + """Variant-aware embeddings subdirectory name. Equals ``global_0.1_degree_representation`` for the default ``vultr`` variant and ``global_0.1_degree_representation.`` otherwise. - Used to construct local mirror paths consistent with the bucket. + Recorded in the ``tessera_metadata.json`` sidecar. """ return self.registry.embeddings_subdir diff --git a/geotessera/registry.py b/geotessera/registry.py index 1a037b3..873dbfe 100644 --- a/geotessera/registry.py +++ b/geotessera/registry.py @@ -12,10 +12,14 @@ import os import math import re +import hashlib import logging import numpy as np import time +import urllib.error +import urllib.request from datetime import datetime, timezone +from email.utils import formatdate, parsedate_to_datetime try: import pandas as pd @@ -32,18 +36,19 @@ # Constants for block-based registry management BLOCK_SIZE = 5 # 5x5 degree blocks -# Default dataset variant. The bare ``global_0.1_degree_representation`` dir on -# S3 corresponds to this variant; named variants get a ``.`` suffix. +# Default dataset variant. The bare ``global_0.1_degree_representation`` +# subdirectory corresponds to this variant; named variants get a ``.`` +# suffix. DEFAULT_VARIANT = "vultr" def _parse_dataset_version(spec: str) -> Tuple[str, str]: """Parse a flexible dataset-version spec. - Returns ``(s3_path_component, normalized_version)``. Accepts inputs like + Returns ``(version_path, normalized_version)``. Accepts inputs like ``"v1"``, ``"1"``, ``"1.0"``, ``"v1.0"`` (all → ``("v1", "1.0")``) and - ``"v1.1"``, ``"1.1"`` (→ ``("v1.1", "1.1")``). The legacy S3 layout uses - ``v1/`` for the 1.0 series — `.0` minors are stripped from the path. + ``"v1.1"``, ``"1.1"`` (→ ``("v1.1", "1.1")``). The repository uses + ``v1/`` for the 1.0 series, so ``.0`` minors are stripped from the path. """ s = spec.strip() if s.startswith("v"): @@ -57,7 +62,7 @@ def _parse_dataset_version(spec: str) -> Tuple[str, str]: def _variant_subdir(variant: str) -> str: - """Map a variant name to its embeddings-dir name on S3.""" + """Map a variant name to its embeddings subdirectory name.""" if variant == DEFAULT_VARIANT: return EMBEDDINGS_DIR_NAME return f"{EMBEDDINGS_DIR_NAME}.{variant}" @@ -66,7 +71,7 @@ def _variant_subdir(variant: str) -> str: def _version_path_from_norm(norm: str) -> str: """Inverse of ``_parse_dataset_version`` for the path component. - ``"1.0"`` → ``"v1"`` (legacy S3 layout uses ``v1/`` for the 1.0 series); + ``"1.0"`` → ``"v1"`` (the repository uses ``v1/`` for the 1.0 series); ``"1.1"`` → ``"v1.1"``; ``"2.0"`` → ``"v2"``; etc. """ major, _, minor = norm.partition(".") @@ -75,10 +80,9 @@ def _version_path_from_norm(norm: str) -> str: return f"v{major}.{minor}" -# Well-known published dataset versions. Used by the client when a -# multi-version operation (e.g. ``coverage --by-source --dataset-version=all``) -# needs to enumerate manifests without a separate listing call. Extend this as -# new versions are published. +# Published dataset versions. Multi-version operations (e.g. ``coverage +# --by-source --dataset-version=all``) enumerate manifests from this list +# rather than issuing a listing call. Extend it as new versions are published. KNOWN_VERSIONS = ("v1", "v1.1") # Sidecar filename written into output directories alongside downloaded tiles @@ -106,6 +110,8 @@ def write_tessera_metadata( "dataset_version_path": version_path, "dataset_variant": dataset_variant, "embeddings_subdir": EMBEDDINGS_DIR_NAME, + # Variant-aware subdirectory name, kept for downstream readers. The + # repository layout itself has no variant level. "s3_embeddings_subdir": _variant_subdir(dataset_variant), "source_url_prefix": f"{TESSERA_NPY_MIRROR_URL}/{version_path}/", "generated_at": datetime.now(timezone.utc).isoformat(), @@ -345,23 +351,25 @@ def tile_to_bounds(lon: float, lat: float) -> Tuple[float, float, float, float]: return (lon - 0.05, lat - 0.05, lon + 0.05, lat + 0.05) -# All Tessera data is served from the public Source Cooperative repository, -# fetched with plain HTTPS. Layout: one tree per media type, one subdir per -# dataset version: -# npy/{version_path}/{year}/grid_.../grid_....npy embeddings + scales +# All Tessera data is served over plain HTTPS from the public Source +# Cooperative repository. The layout is one tree per media type with one +# subdirectory per dataset version: +# npy/{version_path}/{year}/grid_.../grid_....npy embeddings and scales # npy/{version_path}/manifest.parquet per-version manifest # landmasks/{version_path}/grid_....tiff landmask TIFFs # landmasks/{version_path}/landmasks.parquet landmask registry -# Unlike the retired S3 bucket there is no -# global_0.1_degree_representation[.] level: the repository carries -# one embedding tree per version. Versions appear incrementally as they are -# uploaded (npy has v1/v1.1/v2; landmasks has v1 with v1.1 to follow). -TESSERA_MIRROR_URL = "https://data.source.coop/tessera/tessera" +# zarr/{version_path}/tessera.zarr zarr store +# Each version holds a complete embedding tree with no variant subdirectory. +# New versions appear as they are uploaded. The repository is also reachable +# as an S3-compatible endpoint at TESSERA_MIRROR_ENDPOINT with bucket path +# TESSERA_MIRROR_REPO. +TESSERA_MIRROR_ENDPOINT = "https://data.source.coop" +TESSERA_MIRROR_REPO = "tessera/tessera" +TESSERA_MIRROR_URL = f"{TESSERA_MIRROR_ENDPOINT}/{TESSERA_MIRROR_REPO}" TESSERA_NPY_MIRROR_URL = f"{TESSERA_MIRROR_URL}/npy" TESSERA_LANDMASKS_MIRROR_URL = f"{TESSERA_MIRROR_URL}/landmasks" -# Directory structure constants (local embeddings_dir layout, kept from the -# original bucket structure) +# Subdirectory names used in the local embeddings_dir layout. EMBEDDINGS_DIR_NAME = "global_0.1_degree_representation" # NPY embeddings and scales LANDMASKS_DIR_NAME = "global_0.1_degree_tiff_all" # Landmask TIFFs @@ -376,6 +384,21 @@ def landmasks_parquet_url(version_path: str) -> str: return f"{TESSERA_LANDMASKS_MIRROR_URL}/{version_path}/landmasks.parquet" +def embedding_url(version_path: str, path: str) -> str: + """URL of the embedding file at *path* within *version_path*.""" + return f"{TESSERA_NPY_MIRROR_URL}/{version_path}/{path}" + + +def landmask_url(version_path: str, filename: str) -> str: + """URL of the landmask TIFF *filename* for *version_path*.""" + return f"{TESSERA_LANDMASKS_MIRROR_URL}/{version_path}/{filename}" + + +def zarr_store_url(version_path: str) -> str: + """Default URL of the zarr store for *version_path* (e.g. ``"v1"``).""" + return f"{TESSERA_MIRROR_URL}/zarr/{version_path}" + + def format_bytes(num_bytes: float) -> str: """Format a byte count as a human-readable string (e.g. ``"1.5 GB"``).""" for unit in ["B", "KB", "MB", "GB"]: @@ -392,15 +415,12 @@ def download_file_to_temp( ) -> str: """Download a file over HTTPS with caching, retries, and integrity checks. - Integrity is verified against the response ``Content-Length`` (a short read - raises and triggers a retry) and, when the server's ``ETag`` is a plain MD5 - (single-part S3 upload — 32 hex digits with no multipart ``-`` - suffix), against an MD5 computed over the streamed body. When *cache_path* - already exists, an ``If-Modified-Since`` conditional GET (keyed on the - cached file's mtime, which was set from the previous ``Last-Modified``) - short-circuits to the cached copy on a 304. Transient failures (429/5xx, - connection errors, truncation, checksum mismatch) are retried with - exponential backoff. + Integrity is verified against the response ``Content-Length``, and also + against an MD5 of the streamed body when the server's ``ETag`` is a plain + 32-hex-digit MD5. When *cache_path* already exists, an ``If-Modified-Since`` + conditional GET keyed on the cached file's mtime returns the cached copy on + a 304. Transient failures (429/5xx, connection errors, truncation, checksum + mismatch) are retried with exponential backoff. Args: url: HTTPS URL. @@ -417,8 +437,6 @@ def download_file_to_temp( urllib.error.HTTPError: On non-304 HTTP errors (after retries). OSError: On a truncated, corrupted, or failed download (after retries). """ - import urllib.error - attempts = 4 for attempt in range(attempts): try: @@ -438,15 +456,10 @@ def _download_once( cache_path: Optional[Path], ) -> str: """Single download attempt (see :func:`download_file_to_temp`).""" - import hashlib - import urllib.error - import urllib.request - from email.utils import formatdate, parsedate_to_datetime - from . import __version__ - # Cloudflare (fronting the Source Cooperative mirror) rejects the default - # Python-urllib User-Agent with 403, so identify as geotessera. + # The Source Cooperative CDN rejects the default Python urllib + # User-Agent with a 403, so identify as geotessera. request = urllib.request.Request( url, headers={"User-Agent": f"geotessera/{__version__}"} ) @@ -460,74 +473,50 @@ def _download_once( response = urllib.request.urlopen(request, timeout=60) except urllib.error.HTTPError as e: if e.code == 304: - # 304 Not Modified — the cached copy is current. + # A 304 response means the cached copy is current. if progress_callback: progress_callback(0, 0, "Cache is current") return str(cache_path) raise try: - total_size = int(response.headers.get("Content-Length") or 0) - - last_modified = None - lm_header = response.headers.get("Last-Modified") - if lm_header: - try: - last_modified = parsedate_to_datetime(lm_header) - except (TypeError, ValueError): - pass - - # A plain 32-hex-digit ETag is the object's content MD5 (single-part - # S3 upload). Multipart uploads get a "-" suffix and are not a - # content hash, so those downloads are only length-checked. - etag = (response.headers.get("ETag") or "").strip('"') - md5 = hashlib.md5() if re.fullmatch(r"[0-9a-f]{32}", etag) else None - - def chunks(): - downloaded = 0 - while True: - chunk = response.read(8192) - if not chunk: - break - downloaded += len(chunk) - if md5 is not None: - md5.update(chunk) - yield chunk - # A dropped connection can end the stream early without an - # exception; catch it so a truncated file is never cached. - if total_size and downloaded != total_size: - raise OSError( - f"Truncated download from {url}: " - f"got {downloaded} of {total_size} bytes" - ) - if md5 is not None and md5.hexdigest() != etag: - raise OSError( - f"Checksum mismatch for {url}: " - f"MD5 {md5.hexdigest()} != ETag {etag}" - ) - - return _stream_to_file( - chunks(), total_size, last_modified, cache_path, progress_callback - ) + return _stream_to_file(response, url, cache_path, progress_callback) finally: response.close() def _stream_to_file( - chunks, - total_size: int, - last_modified, + response, + url: str, cache_path: Optional[Path], progress_callback: Optional[Callable[[int, int, str], None]], ) -> str: - """Stream *chunks* to *cache_path* (or a temp file) atomically. + """Stream *response* to *cache_path* (or a temporary file) atomically. - Writes to a sibling temp file with progress reporting, stamps the mtime - from *last_modified* (a datetime) so later ``If-Modified-Since`` requests - work, then moves the file into place. Returns the final path as ``str``. + Writes to a sibling temporary file with progress reporting, verifies the + byte count against ``Content-Length`` and, when available, the MD5 ETag, + stamps the file mtime from ``Last-Modified`` so later + ``If-Modified-Since`` requests work, then moves the file into place. + Returns the final path as ``str``. """ import tempfile + total_size = int(response.headers.get("Content-Length") or 0) + + last_modified = None + lm_header = response.headers.get("Last-Modified") + if lm_header: + try: + last_modified = parsedate_to_datetime(lm_header) + except (TypeError, ValueError): + pass + + # A plain 32-hex-digit ETag is the object's content MD5. Multipart + # uploads carry a "-" suffix and are not a content hash, so + # those downloads are only length-checked. + etag = (response.headers.get("ETag") or "").strip('"') + md5 = hashlib.md5() if re.fullmatch(r"[0-9a-f]{32}", etag) else None + if cache_path: cache_path.parent.mkdir(parents=True, exist_ok=True) temp_file = tempfile.NamedTemporaryFile( @@ -551,7 +540,9 @@ def _stream_to_file( size_str = format_bytes(total_size) if total_size > 0 else "unknown size" progress_callback(0, total_size, f"Starting ({size_str})") - for chunk in chunks: + while chunk := response.read(256 * 1024): + if md5 is not None: + md5.update(chunk) temp_file.write(chunk) downloaded += len(chunk) @@ -573,6 +564,17 @@ def _stream_to_file( progress_callback(downloaded, total_size, status) last_update_time = current_time + # A dropped connection can end the stream early without raising, + # so check the length to keep a truncated file out of the cache. + if total_size and downloaded != total_size: + raise OSError( + f"Truncated download from {url}: got {downloaded} of {total_size} bytes" + ) + if md5 is not None and md5.hexdigest() != etag: + raise OSError( + f"Checksum mismatch for {url}: MD5 {md5.hexdigest()} != ETag {etag}" + ) + temp_file.close() # Set file mtime from Last-Modified so the next If-Modified-Since works. @@ -605,9 +607,9 @@ def _stream_to_file( finally: # Remove the partial temp file on any failure, including - # KeyboardInterrupt/SystemExit (a BaseException, which a plain - # `except Exception` misses) — this is what left stray - # ``._tmp_*`` files behind on interrupted downloads. + # KeyboardInterrupt and SystemExit, which a plain `except Exception` + # would miss. This keeps interrupted downloads from leaving stray + # ``._tmp_*`` files behind. if not success: temp_file.close() if temp_path.exists(): @@ -643,10 +645,8 @@ def __init__( Args: version: Dataset version. Accepts ``"v1"``/``"1.0"`` or ``"v1.1"``/``"1.1"``. - variant: Dataset variant to filter the manifest by. The default - ``"vultr"`` corresponds to the bare - ``global_0.1_degree_representation`` directory on S3; other - values (e.g. ``"cambridge"``) map to the ``.`` suffix. + variant: Dataset variant to filter the manifest by (default + ``"vultr"``). cache_dir: Optional directory for caching Parquet registries only (not data files) embeddings_dir: Directory for storing embedding tiles (defaults to current directory). Expected structure: global_0.1_degree_representation[.]/{year}/ @@ -659,14 +659,14 @@ def __init__( landmasks_registry_path: Local path to existing landmasks Parquet registry file logger: Optional logger instance. If not provided, creates a new one """ - # Resolve version into S3 path component and normalised numeric form. + # Resolve version into a path component and a normalised numeric form. self._version_path, self._version_norm = _parse_dataset_version(version) self._variant = variant self._embeddings_subdir = _variant_subdir(variant) # Preserve the original kwarg for callers that still read .version. self.version = self._version_path - # Public read-only view of the variant-aware subdir name used both in - # local mirrors and S3 URLs. + # Public read-only view of the variant-aware subdirectory name, + # recorded in the tessera_metadata.json sidecar. self.embeddings_subdir = self._embeddings_subdir self.variant = self._variant # Populated by _load_registry() with the local path to the manifest @@ -713,19 +713,16 @@ def __init__( # Embeddings manifest (GeoDataFrame with spatial index). The wire format # mirrors the file-scan inventory schema; geometry is derived from # lon/lat at load time if the parquet was written as a plain DataFrame. - # One manifest per version on S3; the consumer fetches the manifest + # Manifests are per version; the consumer fetches the manifest # matching its dataset_version and filters by variant on load. self._registry_gdf: Optional[gpd.GeoDataFrame] = None - self._registry_url = ( - registry_url or manifest_url(self._version_path) - ) + self._registry_url = registry_url or manifest_url(self._version_path) self._registry_path = Path(registry_path) if registry_path else None - # Landmasks Parquet registry (still per-version on S3). + # Landmasks Parquet registry (per version). self._landmasks_df: Optional[pd.DataFrame] = None - self._landmasks_registry_url = ( - landmasks_registry_url - or landmasks_parquet_url(self._version_path) + self._landmasks_registry_url = landmasks_registry_url or landmasks_parquet_url( + self._version_path ) self._landmasks_registry_path = ( Path(landmasks_registry_path) if landmasks_registry_path else None @@ -939,9 +936,9 @@ def _load_landmasks_registry(self): self._landmasks_df = None return - # Validate landmasks registry structure. Hash columns are no longer - # required — integrity is verified against Content-Length and the - # MD5 ETag at download time. + # Validate landmasks registry structure. Hash columns are not + # required because integrity is verified against Content-Length and + # the MD5 ETag at download time. if self._landmasks_df is not None: required_columns = {"lat", "lon", "file_size"} if not required_columns.issubset(self._landmasks_df.columns): @@ -1175,20 +1172,16 @@ def fetch( # Use existing local file return str(local_path) - # Download to embeddings_dir from the Source Cooperative repository, - # which carries one embedding tree per version (no variant subdir). - # Use as_posix() to ensure forward slashes in URL even on Windows + # Download to embeddings_dir. Use as_posix() so the URL uses forward + # slashes on Windows. path_str = path.as_posix() if isinstance(path, Path) else path - url = f"{TESSERA_NPY_MIRROR_URL}/{self._version_path}/{path_str}" - downloaded_path = download_file_to_temp( + url = embedding_url(self._version_path, path_str) + return download_file_to_temp( url, progress_callback=progress_callback, cache_path=local_path, ) - # Return path to saved file - return downloaded_path - def fetch_landmask( self, filename: Optional[str] = None, @@ -1224,9 +1217,8 @@ def fetch_landmask( # Use existing local file return str(local_path) - # Download to embeddings_dir from the Source Cooperative mirror (which - # flattens the LANDMASKS_DIR_NAME subdir away). - url = f"{TESSERA_LANDMASKS_MIRROR_URL}/{self._version_path}/{filename}" + # Download to embeddings_dir. + url = landmask_url(self._version_path, filename) return download_file_to_temp( url, progress_callback=progress_callback, diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 126ab93..52acf3b 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -38,10 +38,12 @@ ) from .registry import ( + TESSERA_MIRROR_ENDPOINT, block_from_world, block_to_embeddings_registry_filename, block_to_landmasks_registry_filename, parse_grid_name, + zarr_store_url, ) from ._terminal import console, emoji @@ -2446,21 +2448,24 @@ def _parse_s3_uri(uri: str) -> Tuple[str, str]: def _s3_list( bucket: str, prefix: str, - region: str, + endpoint_url: str, delimiter: Optional[str] = None, ) -> Iterator[Tuple[str, int, str]]: """Yield (key, size, last_modified) for every object under prefix. - Uses anonymous HTTPS calls to the S3 ListObjectsV2 endpoint. Follows - continuation tokens until the listing is exhausted. If delimiter is - set, yields CommonPrefixes as (prefix, 0, "") tuples instead of object - contents. + Issues anonymous path-style HTTPS calls to the S3-compatible + ListObjectsV2 endpoint at *endpoint_url* (e.g. + ``https://data.source.coop``). Follows continuation tokens until the + listing is exhausted. If delimiter is set, yields CommonPrefixes as + (prefix, 0, "") tuples instead of object contents. """ import xml.etree.ElementTree as ET from urllib.parse import urlencode - from urllib.request import urlopen + from urllib.request import Request, urlopen - base = f"https://{bucket}.s3.{region}.amazonaws.com/" + from . import __version__ + + base = f"{endpoint_url.rstrip('/')}/{bucket}/" token: Optional[str] = None while True: params = {"list-type": "2", "prefix": prefix} @@ -2468,8 +2473,11 @@ def _s3_list( params["delimiter"] = delimiter if token is not None: params["continuation-token"] = token - url = base + "?" + urlencode(params) - with urlopen(url, timeout=60) as resp: + request = Request( + base + "?" + urlencode(params), + headers={"User-Agent": f"geotessera/{__version__}"}, + ) + with urlopen(request, timeout=60) as resp: root = ET.parse(resp).getroot() if delimiter is not None: @@ -2514,13 +2522,23 @@ def _normalize_version(version_dir: str) -> str: def _discover_scan_units( - bucket: str, root_prefix: str, region: str, console: "Console" + bucket: str, + root_prefix: str, + endpoint_url: str, + console: "Console", + flat_variant: str = _DEFAULT_VARIANT, ) -> List[Tuple[str, str, int, str]]: """Walk an S3 prefix and return a list of (version, variant, year, year_prefix). Auto-detects the level of ``root_prefix``: bucket root, version dir, or variant dir. Version/variant components already present in the supplied prefix are inferred from the path and don't need to be discovered again. + + Supports two layouts: the variant-subdir layout + (``{version}/global_0.1_degree_representation[.]/{year}/``) and + the flat Source Cooperative layout (``{version}/{year}/``). The flat + layout carries no variant in the path, so its rows are attributed to + *flat_variant*. """ # Infer any version/variant already encoded in the path. pre_version = None @@ -2538,7 +2556,7 @@ def _discover_scan_units( def walk(prefix: str, version: Optional[str], variant: Optional[str], indent: int): pad = " " * indent - for sp, _, _ in _s3_list(bucket, prefix, region, delimiter="/"): + for sp, _, _ in _s3_list(bucket, prefix, endpoint_url, delimiter="/"): tail = sp[len(prefix) :].rstrip("/") if version is None: m = _VERSION_RE.match(tail) @@ -2555,8 +2573,9 @@ def walk(prefix: str, version: Optional[str], variant: Optional[str], indent: in walk(sp, version, var, indent + 1) continue m = _YEAR_RE.match(tail) - if m and version is not None and variant is not None: - units.append((version, variant, int(m.group(1)), sp)) + if m and version is not None: + # A year directly under a version is the flat layout. + units.append((version, variant or flat_variant, int(m.group(1)), sp)) walk(root_prefix, pre_version, pre_variant, 1) return units @@ -2567,14 +2586,18 @@ def walk(prefix: str, version: Optional[str], variant: Optional[str], indent: in def _discover_landmask_prefixes( - bucket: str, root_prefix: str, region: str, console: "Console" + bucket: str, + root_prefix: str, + endpoint_url: str, + console: "Console", ) -> List[Tuple[str, str, str]]: """Return ``(version_norm, version_path, landmask_prefix)`` per version. - Skips versions that don't have a ``global_0.1_degree_tiff_all/`` dir (e.g. - v1.1 currently ships only the cambridge variant, no landmasks). When - ``root_prefix`` already points at or inside a version dir, only that - version is probed. + Supports both layouts: TIFFs under ``{version}/global_0.1_degree_tiff_all/`` + and TIFFs directly under the version directory (the Source Cooperative + ``landmasks/{version}/`` tree). Versions with no TIFFs at either location + are skipped. When ``root_prefix`` already points at or inside a version + dir, only that version is probed. """ candidates: List[Tuple[str, str]] = [] # (version_norm, version_prefix) @@ -2596,7 +2619,7 @@ def _discover_landmask_prefixes( version_prefix = root_prefix candidates.append((pre_version, version_prefix)) else: - for sp, _, _ in _s3_list(bucket, root_prefix, region, delimiter="/"): + for sp, _, _ in _s3_list(bucket, root_prefix, endpoint_url, delimiter="/"): tail = sp[len(root_prefix) :].rstrip("/") m = _VERSION_RE.match(tail) if m: @@ -2604,18 +2627,19 @@ def _discover_landmask_prefixes( from .registry import _version_path_from_norm + from itertools import islice + out: List[Tuple[str, str, str]] = [] for version_norm, version_prefix in candidates: - lm_prefix = version_prefix + _LANDMASK_DIR + "/" - # Probe with a single key — landmask dirs are flat and large; a quick - # "is anything there?" listing is cheap. - has_any = any(True for _ in _s3_list(bucket, lm_prefix, region)) version_path = _version_path_from_norm(version_norm) - if has_any: - console.print( - f" Landmasks: [green]{version_path}/{_LANDMASK_DIR}/[/green]" - ) - out.append((version_norm, version_path, lm_prefix)) + # Probe a few keys at each candidate location. TIFF keys sort before + # any parquet sidecars, so the first page suffices. + for lm_prefix in (version_prefix + _LANDMASK_DIR + "/", version_prefix): + probe = islice(_s3_list(bucket, lm_prefix, endpoint_url), 20) + if any(key.endswith(".tiff") for key, _, _ in probe): + console.print(f" Landmasks: [green]{lm_prefix}[/green]") + out.append((version_norm, version_path, lm_prefix)) + break return out @@ -2647,7 +2671,10 @@ def s3scan_command(args): console.print(f"[red]Error: {e}[/red]") return 1 - region = args.region + endpoint = args.endpoint_url + if not endpoint: + console.print("[red]Error: --endpoint-url must not be empty[/red]") + return 1 # Output is a *directory* under which per-version manifests are written # as ``{output_dir}/{version_path}/manifest.parquet``. This mirrors the S3 @@ -2659,26 +2686,40 @@ def s3scan_command(args): f"[bold blue]🔍 Spidering S3 for Embedding Tiles[/bold blue]\n" f"📦 Bucket: {bucket}\n" f"🔑 Prefix: {prefix or '(root)'}\n" - f"🌐 Region: {region}\n" + f"🌐 Endpoint: {endpoint}\n" f"📂 Output dir: {output_dir}", style="blue", ) ) - console.print("\n[cyan]Discovering versions, variants, and years...[/cyan]") - try: - scan_units = _discover_scan_units(bucket, prefix, region, console) - except Exception as e: - console.print(f"[red]Error listing S3 prefix: {e}[/red]") - return 1 + # With an explicit --landmasks-uri the landmask scan can proceed even + # when the embeddings prefix yields nothing (a landmasks-only run). + landmasks_only_ok = not args.no_landmasks and args.landmasks_uri is not None - if not scan_units: - console.print( - "[yellow]No (version, variant, year) units found! " - "Expected layout: /v(.)/global_0.1_degree_representation" - "(.)//grid_*[/yellow]" - ) - return 1 + scan_units: List[Tuple[str, str, int, str]] = [] + if landmasks_only_ok and args.landmasks_uri == args.s3_uri: + # A pure landmasks run. Skip embedding discovery: walking a flat + # landmask tree pages through every TIFF key to find no year dirs. + console.print("\n[cyan]Landmasks-only run; skipping embedding scan.[/cyan]") + else: + console.print("\n[cyan]Discovering versions, variants, and years...[/cyan]") + try: + scan_units = _discover_scan_units( + bucket, prefix, endpoint, console, flat_variant=args.variant + ) + except Exception as e: + console.print(f"[red]Error listing S3 prefix: {e}[/red]") + return 1 + + if not scan_units: + console.print( + "[yellow]No (version, variant, year) units found! " + "Expected layout: /v(.)/" + "[global_0.1_degree_representation(.)/]/grid_*[/yellow]" + ) + if not landmasks_only_ok: + return 1 + console.print("[yellow]Continuing with the landmask scan only.[/yellow]") grid_re = re.compile(r"grid_(-?\d+\.\d+)_(-?\d+\.\d+)(_scales)?\.npy$") @@ -2706,7 +2747,7 @@ def _scan_shard( """List one (version, variant, year, lon-shard) and return tile records.""" tiles: Dict[Tuple[float, float], Dict[str, Tuple[str, int, str]]] = {} listed = 0 - for key, size, lm in _s3_list(bucket, year_prefix + shard, region): + for key, size, lm in _s3_list(bucket, year_prefix + shard, endpoint): listed += 1 # Batch UI updates: every 50 objects is plenty smooth and avoids # contention on the Progress lock during big listings. @@ -2828,50 +2869,60 @@ def _scan_shard( if not records: console.print("[yellow]No embedding tiles found![/yellow]") - return 1 + if not landmasks_only_ok: + return 1 from .registry import _version_path_from_norm - df = pd.DataFrame(records) - df = df.sort_values(["version", "variant", "year", "lon", "lat"]) - df["lon_i"] = (df["lon"] * 100).round().astype(np.int32) - df["lat_i"] = (df["lat"] * 100).round().astype(np.int32) - - # One parquet per dataset version. Layout mirrors S3 so the whole tree - # can be uploaded with ``aws s3 cp --recursive / s3:///``. + # One parquet per dataset version. Layout mirrors the remote tree so the + # whole output can be uploaded with + # ``aws s3 cp --recursive / s3:///``. written_files: List[Path] = [] try: - for version_norm, group_df in df.groupby("version", sort=True): - version_path = _version_path_from_norm(str(version_norm)) - # Defensive dedupe: bucket-side misfiling (e.g. a tile filed under - # the wrong grid directory) can surface the same key twice via - # different lon-shards. - before_dedupe = len(group_df) - group_df = group_df.drop_duplicates( - subset=["version", "variant", "year", "lon", "lat"], keep="first" - ) - if len(group_df) != before_dedupe: + if records: + df = pd.DataFrame(records) + df = df.sort_values(["version", "variant", "year", "lon", "lat"]) + df["lon_i"] = (df["lon"] * 100).round().astype(np.int32) + df["lat_i"] = (df["lat"] * 100).round().astype(np.int32) + for version_norm, group_df in df.groupby("version", sort=True): + version_path = _version_path_from_norm(str(version_norm)) + # Defensive dedupe: bucket-side misfiling (e.g. a tile filed + # under the wrong grid directory) can surface the same key + # twice via different lon-shards. + before_dedupe = len(group_df) + group_df = group_df.drop_duplicates( + subset=["version", "variant", "year", "lon", "lat"], + keep="first", + ) + if len(group_df) != before_dedupe: + console.print( + f"[yellow] Dropped {before_dedupe - len(group_df):,} " + f"duplicate (variant, year, lon, lat) rows for " + f"{version_path}[/yellow]" + ) + out_file = output_dir / version_path / "manifest.parquet" + out_file.parent.mkdir(parents=True, exist_ok=True) console.print( - f"[yellow] Dropped {before_dedupe - len(group_df):,} duplicate " - f"(variant, year, lon, lat) rows for {version_path}[/yellow]" + f"[cyan]Writing {len(group_df):,} tiles to {out_file}...[/cyan]" ) - out_file = output_dir / version_path / "manifest.parquet" - out_file.parent.mkdir(parents=True, exist_ok=True) - console.print( - f"[cyan]Writing {len(group_df):,} tiles to {out_file}...[/cyan]" - ) - _atomic_write_parquet(group_df, out_file) - written_files.append(out_file) + _atomic_write_parquet(group_df, out_file) + written_files.append(out_file) - # Landmask scan: one parquet per version that has a landmasks dir. - # Versions that share the 0.1° grid but don't ship their own landmasks - # dir (e.g. v1.1 → reuses v1's landmasks) get a copy of the first - # scanned landmasks parquet. + # Landmask scan: one parquet per version that has landmask TIFFs. landmask_files_by_version: Dict[str, Path] = {} if not args.no_landmasks: + # The landmask tree may live outside the embeddings prefix (on + # data.source.coop it is the sibling landmasks/ tree); + # --landmasks-uri points the scan there. + if args.landmasks_uri: + lm_bucket, lm_root = _parse_s3_uri(args.landmasks_uri) + else: + lm_bucket, lm_root = bucket, prefix console.print("\n[cyan]Discovering landmask directories...[/cyan]") try: - lm_units = _discover_landmask_prefixes(bucket, prefix, region, console) + lm_units = _discover_landmask_prefixes( + lm_bucket, lm_root, endpoint, console + ) except Exception as e: console.print(f"[yellow]Could not list landmasks: {e}[/yellow]") lm_units = [] @@ -2885,7 +2936,9 @@ def _scan_shard( def _scan_lm_shard(shard: str): out_local = [] - for key, size, lm in _s3_list(bucket, lm_prefix + shard, region): + for key, size, lm in _s3_list( + lm_bucket, lm_prefix + shard, endpoint + ): name = key.rsplit("/", 1)[-1] m = _LANDMASK_RE.match(name) if not m: @@ -2902,7 +2955,7 @@ def _scan_lm_shard(shard: str): "mtime": datetime.fromisoformat( lm.replace("Z", "+00:00") ), - "key": f"s3://{bucket}/{key}", + "key": f"s3://{lm_bucket}/{key}", } ) return out_local @@ -2970,8 +3023,8 @@ def _scan_lm_shard(shard: str): missing = sorted(embedding_version_paths - landmask_files_by_version.keys()) for vpath in missing: console.print( - f"[yellow]Warning: {vpath} has embeddings but no " - f"global_0.1_degree_tiff_all/ on S3. No landmasks.parquet " + f"[yellow]Warning: {vpath} has embeddings but no landmask " + f"TIFFs at the scanned location. No landmasks.parquet " f"will be written for {vpath}.[/yellow]" ) @@ -3930,8 +3983,8 @@ def main(): # Specify custom output path geotessera-registry file-scan /path/to/embeddings --output /path/to/inventory.parquet - # Spider an S3 bucket prefix to build an inventory parquet (no AWS creds required for public buckets) - geotessera-registry s3scan s3://tessera-embeddings/v1/global_0.1_degree_representation + # Spider the Source Cooperative repository to rebuild a per-version manifest + geotessera-registry s3scan s3://tessera/tessera/npy/v1.1/ --variant cambridge --landmasks-uri s3://tessera/tessera/landmasks/v1.1/ # This will: # - List year subprefixes (e.g. 2024/, 2023/) via anonymous ListObjectsV2 HTTPS calls @@ -4089,15 +4142,25 @@ def main(): ) s3scan_parser.add_argument( "s3_uri", - help="S3 URI at any level: bucket root (discovers all versions + " - "variants), version dir (e.g. s3://bucket/v1.1/), or variant dir " + help="S3 URI at any level: tree root (discovers all versions), " + "version dir (e.g. s3://tessera/tessera/npy/v1.1/), or variant dir " "(e.g. s3://bucket/v1.1/global_0.1_degree_representation.cambridge/)", ) s3scan_parser.add_argument( - "--region", + "--endpoint-url", type=str, - default="us-west-2", - help="AWS region of the bucket (default: us-west-2)", + default=TESSERA_MIRROR_ENDPOINT, + help="S3-compatible endpoint for path-style listing requests " + f"(default: {TESSERA_MIRROR_ENDPOINT})", + ) + s3scan_parser.add_argument( + "--variant", + type=str, + default="vultr", + help="Variant recorded in the manifest when the layout has no " + "variant directory, i.e. years sit directly under the version dir " + "as on data.source.coop (default: vultr). Layouts that encode the " + "variant in the path ignore this.", ) s3scan_parser.add_argument( "--workers", @@ -4119,6 +4182,15 @@ def main(): action="store_true", help="Skip scanning landmask TIFFs and writing landmasks.parquet", ) + s3scan_parser.add_argument( + "--landmasks-uri", + type=str, + default=None, + help="S3 URI of the landmask tree when it is not under the embeddings " + "prefix, e.g. s3://tessera/tessera/landmasks/ or a version dir within " + "it. With this set, the landmask scan runs even when the embeddings " + "prefix yields no tiles, so a landmasks-only regeneration is possible.", + ) s3scan_parser.set_defaults(func=s3scan_command) # File-check command @@ -4431,7 +4503,7 @@ def main(): ) verify_parser.add_argument( "--store", - default="https://data.source.coop/tessera/tessera/zarr/v1", + default=zarr_store_url("v1"), help="Zarr store URL", ) verify_parser.set_defaults(func=verify_tile_command) @@ -4446,7 +4518,7 @@ def main(): print_parser.add_argument("--year", type=int, required=True, help="Year") print_parser.add_argument( "--store", - default="https://data.source.coop/tessera/tessera/zarr/v1", + default=zarr_store_url("v1"), help="Zarr store URL", ) print_parser.set_defaults(func=print_command) diff --git a/geotessera/store.py b/geotessera/store.py index e8c1c7c..b772461 100644 --- a/geotessera/store.py +++ b/geotessera/store.py @@ -32,9 +32,11 @@ from pyproj import Transformer from rich.progress import track +from .registry import zarr_store_url + log = logging.getLogger(__name__) -DEFAULT_STORE = "https://data.source.coop/tessera/tessera/zarr/v1" +DEFAULT_STORE = zarr_store_url("v1") # Shard-aligned chunk sizes so dask tasks match zarr shards SHARD_CHUNKS = {"time": 1, "band": 128, "y": 4096, "x": 4096} diff --git a/geotessera/visualization.py b/geotessera/visualization.py index 9da1955..260e111 100644 --- a/geotessera/visualization.py +++ b/geotessera/visualization.py @@ -144,10 +144,10 @@ def visualize_sources_coverage( legend. Args: - manifest_path: Local path to a manifest parquet — or a list of paths - which are concatenated before rendering. With per-version manifests + manifest_path: Local path to a manifest parquet, or a list of paths + that are concatenated before rendering. Manifests are per version (``data.source.coop/tessera/tessera/npy/{v}/manifest.parquet``), - pass a list to compare versions on a single map. + so pass a list to compare versions on a single map. output_path: Output PNG path. year: Optional year filter (applies to all sources). width_pixels: Output image width in pixels. diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 17a7aa7..0fe4252 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -202,6 +202,8 @@ def gather_tile_infos( from .registry import ( EMBEDDINGS_DIR_NAME, LANDMASKS_DIR_NAME, + TESSERA_MIRROR_ENDPOINT, + TESSERA_MIRROR_REPO, tile_to_embedding_paths, tile_to_landmask_filename, ) @@ -255,6 +257,14 @@ def gather_tile_infos( # Landmasks are STRICTLY per-version — no cross-version fallback. Each # Tessera version has its own landmask grid and mixing them silently # would corrupt water masking. + def landmask_sync_hint(dest: Path) -> str: + return ( + f" aws s3 sync --no-sign-request " + f"--endpoint-url {TESSERA_MIRROR_ENDPOINT} " + f"s3://{TESSERA_MIRROR_REPO}/landmasks/" + f"{registry._version_path}/ {dest}/" + ) + emb_candidate = ( registry._embeddings_dir / registry._version_path / EMBEDDINGS_DIR_NAME ) @@ -275,10 +285,7 @@ def gather_tile_infos( f"Landmask directory not found for {registry._version_path}: " f"expected {lm_s3_mirror}. Landmasks are per-version and " f"cannot be reused across versions. Fetch them with:\n" - f" aws s3 sync --no-sign-request " - f"--endpoint-url https://data.source.coop " - f"s3://tessera/tessera/landmasks/" - f"{registry._version_path}/ {lm_s3_mirror}/" + + landmask_sync_hint(lm_s3_mirror) ) else: base_emb = str(registry._embeddings_dir / EMBEDDINGS_DIR_NAME) @@ -287,11 +294,7 @@ def gather_tile_infos( else: raise FileNotFoundError( f"Landmask directory not found: expected {lm_flat}. " - f"Fetch them with:\n" - f" aws s3 sync --no-sign-request " - f"--endpoint-url https://data.source.coop " - f"s3://tessera/tessera/landmasks/" - f"{registry._version_path}/ {lm_flat}/" + f"Fetch them with:\n" + landmask_sync_hint(lm_flat) ) zones_dict: Dict[int, List[TileInfo]] = {} transformer_cache: Dict[int, ProjTransformer] = {} diff --git a/tests/v11.t b/tests/v11.t index 6f17b69..bbc341d 100644 --- a/tests/v11.t +++ b/tests/v11.t @@ -102,7 +102,7 @@ manifest must fail loudly, not silently render an empty dataset: Test: Manifest URLs Are Per-Version ------------------------------------ -Confirm the consumer fetches from the per-version path on S3: +Confirm the consumer fetches from the per-version manifest path: $ uv run python -c " > from geotessera.registry import Registry From 08198ec03f6735e2541f1ce3d5d46e1f559b2695 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Thu, 30 Jul 2026 19:21:01 +0100 Subject: [PATCH 03/13] updates --- CHANGES.md | 5 +++++ README.md | 17 +++++++++-------- docs/cli_reference.rst | 9 +++++---- docs/index.rst | 12 ++++++------ geotessera/cli.py | 33 ++++++++++++++++++++++----------- geotessera/core.py | 14 ++++++++------ geotessera/registry.py | 20 +++++++++++++++----- geotessera/registry_cli.py | 30 ++++++++++++++++++++---------- tests/v11.t | 25 +++++++++++++++++++++++++ 9 files changed, 115 insertions(+), 50 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 393b1e0..7a28195 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -18,6 +18,11 @@ ### New Features +- **Per-version default variant**: omitting `dataset_variant` / + `--dataset-variant` now selects the version's published variant (`vultr` + for v1, `cambridge` for v1.1) instead of always `vultr`, so + `GeoTessera(dataset_version="v1.1")` works without an explicit variant. + (@avsm) - **`geotessera-registry zarr-consolidate`**: New subcommand that re-consolidates a store's root metadata after in-place changes. Mostly only for repairs and not regular use. diff --git a/README.md b/README.md index bea778d..5df9f5f 100644 --- a/README.md +++ b/README.md @@ -101,9 +101,9 @@ The Tessera embeddings use a **0.1-degree grid system**: ### File Structure and Downloads -When you request embeddings, GeoTessera downloads files from the public S3 -bucket (using anonymous, unsigned requests) into the output directory you -specify, where they persist for re-use: +When you request embeddings, GeoTessera downloads files over HTTPS from the +public Source Cooperative repository into the output directory you specify, +where they persist for re-use: #### Embedding Files (via `fetch_embedding`) 1. **Quantized embeddings** (`grid_X.XX_Y.YY.npy`): @@ -379,7 +379,7 @@ Options: -f, --format TEXT Output format: 'tiff' or 'npy' (default: tiff) --year INT Year of embeddings (default: 2024) --dataset-version TEXT Tessera dataset version (e.g. v1, v1.1) - --dataset-variant TEXT Tessera dataset variant (default: vultr) + --dataset-variant TEXT Tessera dataset variant (default: the version's published variant) --bands TEXT Comma-separated band indices (default: all 128) --compress TEXT Compression for TIFF format (default: lzw) --dry-run Calculate total download size without downloading @@ -449,7 +449,7 @@ Options: --tile TEXT Single tile by any point within it: 'lon,lat' --by-source Render each (version, variant) source in a distinct colour --dataset-version TEXT Tessera dataset version (e.g. v1, v1.1; or 'all' with --by-source) - --dataset-variant TEXT Tessera dataset variant (default: vultr; or 'all' with --by-source) + --dataset-variant TEXT Tessera dataset variant (default: the version's published variant; or 'all' with --by-source) --region-file PATH GeoJSON/Shapefile to focus on specific region --country TEXT Country name to focus on (e.g., 'United Kingdom') --tile-color TEXT Color for tiles (default: red) @@ -483,7 +483,7 @@ geotessera info [OPTIONS] Options: --tiles PATH Analyze tile files/directory (GeoTIFF or NPY format) --dataset-version TEXT Tessera dataset version (e.g. v1, v1.1) - --dataset-variant TEXT Tessera dataset variant (default: vultr) + --dataset-variant TEXT Tessera dataset variant (default: the version's published variant) -v, --verbose Verbose output ``` @@ -507,8 +507,9 @@ GeoTessera uses a Parquet-based registry system to efficiently manage and access ### Dataset Versions and Variants Tessera embeddings are published as dataset *versions* (e.g. `v1`, `v1.1`) and, -within a version, as *variants* produced by different model runs (e.g. the -default `vultr`, or `cambridge`). Select them on the CLI with `--dataset-version` +within a version, as *variants* produced by different model runs (`vultr` for +v1, `cambridge` for v1.1; each is the version's default). Select them on the +CLI with `--dataset-version` and `--dataset-variant`, or in Python: ```python diff --git a/docs/cli_reference.rst b/docs/cli_reference.rst index 8c35477..c9df57c 100644 --- a/docs/cli_reference.rst +++ b/docs/cli_reference.rst @@ -11,8 +11,9 @@ dataset-selection options:: --dataset-version TEXT Tessera dataset version (default: v1). Accepts v1, 1.0, v1.0, v1.1, 1.1 etc. - --dataset-variant TEXT Tessera dataset variant (default: vultr). - Known variants: vultr (1.0 default), cambridge (1.1). + --dataset-variant TEXT Tessera dataset variant (default: the + version's published variant: vultr for 1.0, + cambridge for 1.1). --verbose, -v Enable verbose output --help Show help message @@ -88,7 +89,7 @@ Download embeddings for a region in numpy or GeoTIFF format. * ``--year INT`` - Year of embeddings (default: 2024) * ``--dataset-version TEXT`` - Tessera dataset version (``v1`` / ``1.0`` / ``v1.1`` / ``1.1``; default ``v1``). Pick **once per project**. -* ``--dataset-variant TEXT`` - Tessera dataset variant (default: ``vultr``). +* ``--dataset-variant TEXT`` - Tessera dataset variant (default: the version's published variant). Pass ``cambridge`` for v1.1 test embeddings. **Other Options**: @@ -433,7 +434,7 @@ Display information about GeoTIFF files or the library. * ``--tiles PATH`` - Analyze tile files/directory (GeoTIFF or NPY format) * ``--geotiffs PATH`` - Alias for --tiles (deprecated) * ``--dataset-version TEXT`` - Tessera dataset version (default: ``v1``) -* ``--dataset-variant TEXT`` - Tessera dataset variant (default: ``vultr``) +* ``--dataset-variant TEXT`` - Tessera dataset variant (default: the version's published variant) * ``-v, --verbose`` - Verbose output **Examples**:: diff --git a/docs/index.rst b/docs/index.rst index 7d1e8d1..5ab8bd8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -255,9 +255,9 @@ GeoTessera ships embeddings under two orthogonal axes: Different versions have *different 128-channel feature spaces*: a feature vector from one version is **not comparable** to a vector from another. * **dataset variant** — for a given version, an independent model run / - release channel. The default is ``vultr`` (the production hosting on - Vultr); ``cambridge`` is a test deployment by the Cambridge team for the - 1.1 line. + release channel. ``vultr`` (the production hosting on Vultr) is the + variant for the 1.0 line; ``cambridge`` is the Cambridge team's deployment + for the 1.1 line. Each version's sole published variant is its default. Currently published combinations on ``data.source.coop/tessera/tessera``: @@ -314,9 +314,9 @@ accepts both flags:: internal normalised form (used in manifests and the metadata sidecar) is ``1.0`` / ``1.1``; the S3 path component is ``v1`` / ``v1.1``. -``--dataset-variant`` defaults to ``vultr`` so unflagged commands keep -working against the legacy line; pass ``cambridge`` (or any other -published variant) explicitly. +``--dataset-variant`` defaults to the version's published variant +(``vultr`` for 1.0, ``cambridge`` for 1.1), so unflagged commands work for +either version; pass a variant explicitly to override. **Python API**:: diff --git a/geotessera/cli.py b/geotessera/cli.py index 62bd3a4..de4abf1 100644 --- a/geotessera/cli.py +++ b/geotessera/cli.py @@ -270,11 +270,13 @@ def info( ), ] = "v1", dataset_variant: Annotated[ - str, + Optional[str], typer.Option( - "--dataset-variant", help="Tessera dataset variant (default: vultr)" + "--dataset-variant", + help="Tessera dataset variant (default: the version's published " + "variant, e.g. vultr for v1 and cambridge for v1.1)", ), - ] = "vultr", + ] = None, verbose: Annotated[ bool, typer.Option("--verbose", "-v", help="Verbose output") ] = False, @@ -557,9 +559,10 @@ def coverage( Optional[str], typer.Option( "--dataset-variant", - help="Tessera dataset variant. Defaults to vultr for the " - "single-source view; with --by-source, omitting this means " - "'all variants'. Pass 'all' explicitly to force the multi-variant view.", + help="Tessera dataset variant. Defaults to the version's " + "published variant for the single-source view; with --by-source, " + "omitting this means 'all variants'. Pass 'all' explicitly to " + "force the multi-variant view.", ), ] = None, cache_dir: Annotated[ @@ -755,8 +758,14 @@ def country_progress_callback(current: int, total: int, status: str = None): version_spec = dataset_version if dataset_version is not None else "all" variant_spec = dataset_variant if dataset_variant is not None else "all" else: + from geotessera.registry import _parse_dataset_version, default_variant + version_spec = dataset_version if dataset_version is not None else "v1" - variant_spec = dataset_variant if dataset_variant is not None else "vultr" + variant_spec = ( + dataset_variant + if dataset_variant is not None + else default_variant(_parse_dataset_version(version_spec)[1]) + ) # The placeholder GeoTessera below initialises one Registry to get its # cache paths; additional manifests are downloaded into the same cache @@ -1185,11 +1194,13 @@ def download( ), ] = "v1", dataset_variant: Annotated[ - str, + Optional[str], typer.Option( - "--dataset-variant", help="Tessera dataset variant (default: vultr)" + "--dataset-variant", + help="Tessera dataset variant (default: the version's published " + "variant, e.g. vultr for v1 and cambridge for v1.1)", ), - ] = "vultr", + ] = None, cache_dir: Annotated[ Optional[Path], typer.Option("--cache-dir", help="Cache directory") ] = None, @@ -1744,7 +1755,7 @@ def mark_file_complete(file_key): sidecar = write_tessera_metadata( output, dataset_version=dataset_version, - dataset_variant=dataset_variant, + dataset_variant=gt.dataset_variant, extra={ "format": format, "year": year, diff --git a/geotessera/core.py b/geotessera/core.py index abdeac4..23acdc6 100644 --- a/geotessera/core.py +++ b/geotessera/core.py @@ -91,7 +91,7 @@ class GeoTessera: def __init__( self, dataset_version: str = "v1", - dataset_variant: str = "vultr", + dataset_variant: Optional[str] = None, cache_dir: Optional[Union[str, Path]] = None, embeddings_dir: Optional[Union[str, Path]] = None, registry_url: Optional[str] = None, @@ -104,10 +104,10 @@ def __init__( dataset_version: Tessera dataset version. Accepts ``"v1"`` / ``"1.0"`` or ``"v1.1"`` / ``"1.1"`` (the repository uses ``v1/`` for the 1.0 series). - dataset_variant: Variant of the embeddings to load (default - ``"vultr"``). Other published variants (e.g. ``"cambridge"``) - are produced by different model runs and are selected by - filtering the manifest. + dataset_variant: Variant of the embeddings to load. Defaults to + the version's published variant (``"vultr"`` for v1, + ``"cambridge"`` for v1.1). Variants are produced by different + model runs and are selected by filtering the manifest. cache_dir: Directory for caching registry files only (not embedding data) embeddings_dir: Directory containing pre-downloaded embedding tiles. Defaults to current working directory if not specified. @@ -135,7 +135,6 @@ def __init__( registry_dir: Directory containing registry.parquet and landmasks.parquet files """ self.dataset_version = dataset_version - self.dataset_variant = dataset_variant # Initialize logger self.logger = logging.getLogger(__name__) @@ -156,6 +155,9 @@ def __init__( registry_dir=registry_dir, logger=self.logger, ) + # The resolved variant (per-version default applied when the caller + # passed None). + self.dataset_variant = self.registry.variant @property def version(self) -> str: diff --git a/geotessera/registry.py b/geotessera/registry.py index 873dbfe..f3e32e2 100644 --- a/geotessera/registry.py +++ b/geotessera/registry.py @@ -41,6 +41,15 @@ # suffix. DEFAULT_VARIANT = "vultr" +# The variant each published version ships. Versions currently publish a +# single variant each, so this doubles as the per-version default. +VERSION_DEFAULT_VARIANTS = {"1.0": "vultr", "1.1": "cambridge"} + + +def default_variant(version_norm: str) -> str: + """Default variant for *version_norm* (e.g. ``"1.1"`` → ``"cambridge"``).""" + return VERSION_DEFAULT_VARIANTS.get(version_norm, DEFAULT_VARIANT) + def _parse_dataset_version(spec: str) -> Tuple[str, str]: """Parse a flexible dataset-version spec. @@ -631,7 +640,7 @@ class Registry: def __init__( self, version: str, - variant: str = DEFAULT_VARIANT, + variant: Optional[str] = None, cache_dir: Optional[Union[str, Path]] = None, embeddings_dir: Optional[Union[str, Path]] = None, registry_url: Optional[str] = None, @@ -645,8 +654,9 @@ def __init__( Args: version: Dataset version. Accepts ``"v1"``/``"1.0"`` or ``"v1.1"``/``"1.1"``. - variant: Dataset variant to filter the manifest by (default - ``"vultr"``). + variant: Dataset variant to filter the manifest by. Defaults to + the version's published variant (``"vultr"`` for v1, + ``"cambridge"`` for v1.1). cache_dir: Optional directory for caching Parquet registries only (not data files) embeddings_dir: Directory for storing embedding tiles (defaults to current directory). Expected structure: global_0.1_degree_representation[.]/{year}/ @@ -661,8 +671,8 @@ def __init__( """ # Resolve version into a path component and a normalised numeric form. self._version_path, self._version_norm = _parse_dataset_version(version) - self._variant = variant - self._embeddings_subdir = _variant_subdir(variant) + self._variant = variant or default_variant(self._version_norm) + self._embeddings_subdir = _variant_subdir(self._variant) # Preserve the original kwarg for callers that still read .version. self.version = self._version_path # Public read-only view of the variant-aware subdirectory name, diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 52acf3b..345ff45 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -42,6 +42,7 @@ block_from_world, block_to_embeddings_registry_filename, block_to_landmasks_registry_filename, + default_variant, parse_grid_name, zarr_store_url, ) @@ -2526,7 +2527,7 @@ def _discover_scan_units( root_prefix: str, endpoint_url: str, console: "Console", - flat_variant: str = _DEFAULT_VARIANT, + flat_variant: Optional[str] = None, ) -> List[Tuple[str, str, int, str]]: """Walk an S3 prefix and return a list of (version, variant, year, year_prefix). @@ -2538,7 +2539,7 @@ def _discover_scan_units( (``{version}/global_0.1_degree_representation[.]/{year}/``) and the flat Source Cooperative layout (``{version}/{year}/``). The flat layout carries no variant in the path, so its rows are attributed to - *flat_variant*. + *flat_variant*, or to the version's published variant when that is None. """ # Infer any version/variant already encoded in the path. pre_version = None @@ -2575,7 +2576,8 @@ def walk(prefix: str, version: Optional[str], variant: Optional[str], indent: in m = _YEAR_RE.match(tail) if m and version is not None: # A year directly under a version is the flat layout. - units.append((version, variant or flat_variant, int(m.group(1)), sp)) + var = variant or flat_variant or default_variant(version) + units.append((version, var, int(m.group(1)), sp)) walk(root_prefix, pre_version, pre_variant, 1) return units @@ -3245,7 +3247,8 @@ def _detect_dataset_metadata( If the user passes the flags explicitly we honour them; otherwise we look for a ``tessera_metadata.json`` sidecar (written by the CLI download - flow) and use what it recorded. Falls back to v1/vultr. + flow) and use what it recorded. Falls back to v1 and the version's + published variant. """ if explicit_version and explicit_variant: return explicit_version, explicit_variant @@ -3267,7 +3270,11 @@ def _detect_dataset_metadata( except (OSError, ValueError): pass - return explicit_version or "v1", explicit_variant or "vultr" + from .registry import _parse_dataset_version + + version = explicit_version or "v1" + variant = explicit_variant or default_variant(_parse_dataset_version(version)[1]) + return version, variant def zarr_init_command(args): @@ -4156,11 +4163,12 @@ def main(): s3scan_parser.add_argument( "--variant", type=str, - default="vultr", + default=None, help="Variant recorded in the manifest when the layout has no " "variant directory, i.e. years sit directly under the version dir " - "as on data.source.coop (default: vultr). Layouts that encode the " - "variant in the path ignore this.", + "as on data.source.coop (default: the version's published variant, " + "e.g. vultr for v1 and cambridge for v1.1). Layouts that encode " + "the variant in the path ignore this.", ) s3scan_parser.add_argument( "--workers", @@ -4244,7 +4252,8 @@ def main(): type=str, default=None, help="Tessera dataset variant (e.g. vultr, cambridge). " - "Default: read from tessera_metadata.json in base_dir, else vultr.", + "Default: read from tessera_metadata.json in base_dir, else the " + "version's published variant.", ) zarr_init_parser.set_defaults(func=zarr_init_command) @@ -4298,7 +4307,8 @@ def main(): type=str, default=None, help="Tessera dataset variant (e.g. vultr, cambridge). " - "Default: read from tessera_metadata.json in base_dir, else vultr.", + "Default: read from tessera_metadata.json in base_dir, else the " + "version's published variant.", ) zarr_fill_parser.set_defaults(func=zarr_fill_command) diff --git a/tests/v11.t b/tests/v11.t index bbc341d..728b63d 100644 --- a/tests/v11.t +++ b/tests/v11.t @@ -80,6 +80,31 @@ The numeric form ``1.1`` is equivalent: > " [2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022, 2023, 2024, 2025] +Test: Variant Defaults Are Per Version +-------------------------------------- + +Each published version ships a single variant, which is its default: + + $ uv run python -c " + > from geotessera.registry import default_variant + > print(default_variant('1.0'), default_variant('1.1'), default_variant('9.9')) + > " + vultr cambridge vultr + +Omitting ``dataset_variant`` for v1.1 selects cambridge: + + $ uv run python -c " + > from geotessera import GeoTessera + > import tempfile, os + > with tempfile.TemporaryDirectory() as d: + > gt = GeoTessera(dataset_version='v1.1', + > cache_dir=os.path.join(d, 'c'), embeddings_dir=os.path.join(d, 'e')) + > print(gt.dataset_variant) + > print(min(gt.registry.get_available_years()), max(gt.registry.get_available_years())) + > " + cambridge + 2015 2025 + Test: Bad Variant Raises a Clear ValueError ------------------------------------------- From fb2932827f3e8870f94cd81019ea38f6c0a68055 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Fri, 31 Jul 2026 16:56:08 +0100 Subject: [PATCH 04/13] flesh out the zarr fill commands with s3 support and make s3fs optional --- CHANGES.md | 55 +- docs/architecture.rst | 137 ++++ geotessera/registry_cli.py | 646 ++++++++++++++--- geotessera/remote.py | 407 +++++++++++ geotessera/zarr.py | 1347 ++++++++++++++++++++++++++++-------- pyproject.toml | 6 + uv.lock | 135 ++++ 7 files changed, 2335 insertions(+), 398 deletions(-) create mode 100644 geotessera/remote.py diff --git a/CHANGES.md b/CHANGES.md index 7a28195..a14da08 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -23,9 +23,60 @@ for v1, `cambridge` for v1.1) instead of always `vultr`, so `GeoTessera(dataset_version="v1.1")` works without an explicit variant. (@avsm) +- **Remote zarr builds**: `zarr-init` and `zarr-fill` now take locations + rather than paths — both the tile source and the output store may be + fsspec URLs (`s3://bucket/prefix`), so a store on one S3 node can be + filled from tiles on another with no local mirror. Remote tiles are read + with byte-range GETs sized to the rows each shard needs (an `.npy` is a + header plus a flat C-ordered buffer), so no scratch disk is involved. + `--source-*` and `--store-*` flags configure the two endpoints + independently; credentials come from the environment, a named profile + (`--store-profile`), or an instance role rather than argv, and + `--store-acl` stamps a canned ACL such as `bucket-owner-full-control` on + every object written. `s3://` locations need the new optional + `s3` extra (`pip install 'geotessera[s3]'`), which is what pulls in + `s3fs` and `botocore`; the core install stays free of both and `https://` + sources work without it. (@avsm) +- **Parallel per-zone fills**: `zarr-fill --zones N` is now safe to run as + one process per UTM zone against a shared store. Ingestion tracking is one + object per zone/year, each zone/year takes an advisory lock + (`--force-lock` to take over a dead run's), and root-metadata + consolidation is skipped by default for a zone-restricted fill. See + the architecture guide for the sweep recipe. (@avsm) +- **The store now contains only Zarr**: build bookkeeping — the ingestion + registry, fill locks and global-preview resume markers — moved out of the + store into a sibling location, `.build` by default and relocatable + with `--state-url`. Previously these sat at the store root, where every + hierarchy listing and `consolidate_metadata` call warned about + unrecognised objects and readers saw non-Zarr entries. A `_registry.parquet` + left inside an older store is still read, so existing stores resume + correctly; nothing is written back into them. (@avsm) +- **`geotessera-registry zarr-extend`**: New subcommand that appends years + to an existing store's time axis, so a new year can be added without + rebuilding. Time is chunked one year per chunk, making this a + metadata-only edit — existing chunks are never rewritten — and the new + slice reads back with the same sentinels a freshly initialised year has. + Years may only be appended (inserting an earlier one would renumber every + chunk, so it is refused), and it will not run while a fill lock is held. + (@avsm) - **`geotessera-registry zarr-consolidate`**: New subcommand that - re-consolidates a store's root metadata after in-place changes. - Mostly only for repairs and not regular use. + re-consolidates a store's root metadata after in-place changes, and + merges the per-zone ingestion registries into `_registry.parquet`. This + is the single-writer step that finishes a parallel sweep; also useful for + repairs. Accepts a local path or a remote store URL. + +### Bug Fixes + +- **Incremental fills no longer erase neighbouring tiles**: a shard write + replaces the whole shard, so a fill that touched a shard already holding + data would zero out the tiles it did not re-read. Touched shards are now + rebuilt from every tile overlapping them. (@avsm) +- **Failed shards are no longer recorded as written**: tiles are only added + to the ingestion registry once every shard covering them succeeded, so + re-running a fill retries exactly the unfinished work. A fill with any + failed shard now reports an error. (@avsm) +- **`geotessera-registry` propagates exit status**: command return codes + were discarded, so failures reported success to the shell. (@avsm) - **`geotessera-registry s3scan` scans Source Cooperative**: listings are path-style against the S3-compatible endpoint given by `--endpoint-url` (default `https://data.source.coop`); the AWS virtual-hosted addressing diff --git a/docs/architecture.rst b/docs/architecture.rst index 0ad97ee..d285cde 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -516,6 +516,143 @@ group. The store automatically routes geographic queries to the correct zone:: Datasets are cached per zone for the lifetime of the ``GeoTesseraZarr`` instance. +Building a Zarr Store +~~~~~~~~~~~~~~~~~~~~~ + +Stores are built by the maintainer-facing ``geotessera-registry`` CLI in +three steps: ``zarr-init`` lays out the (metadata-only) hierarchy from the +landmask registry, ``zarr-fill`` writes tile data into it, and +``zarr-consolidate`` refreshes the root metadata that HTTP readers depend on. + +**Locations, not paths**: the tile source and the output store are each +either a local directory or an fsspec URL. A store on one S3 node can be +filled from tiles on another with no local mirror:: + + geotessera-registry zarr-fill \ + s3://source-bucket/tessera \ + s3://dest-bucket/tessera.zarr \ + --year 2024 --zones 30 \ + --source-endpoint-url https://data.source.coop --source-anon \ + --store-endpoint-url https://s3.example.org + +Credentials come from the environment (``AWS_ACCESS_KEY_ID``, +``AWS_SECRET_ACCESS_KEY``, ``AWS_PROFILE``, instance roles) rather than +flags, so they never appear in a process listing. ``--source-*`` and +``--store-*`` flags configure the two endpoints independently, so a named +AWS CLI profile is selected with ``--store-profile``. + +Where the destination bucket belongs to another account, ``--store-acl`` +stamps a canned ACL on every object written — the equivalent of the AWS +CLI's ``--acl``:: + + --store-profile sc-writer --store-acl bucket-owner-full-control + +It applies to the store's Zarr chunks and metadata as well as the sidecar +parquet and lock objects, and is filtered out of read requests. + +``s3://`` locations need the optional ``s3`` extra, which pulls in ``s3fs`` +and ``botocore``:: + + pip install 'geotessera[s3]' + +The core install stays free of both; ``https://`` sources (including the +public Source Cooperative front) work without the extra, since fsspec reads +those over the ``aiohttp`` already required. + +**Streaming reads**: a ``.npy`` tile is a short header followed by a flat +C-ordered buffer, so the rows a shard needs are one contiguous byte range. +Remote tiles are read with a single ranged GET per shard overlap rather +than downloaded whole, which means no scratch disk and no cache eviction +policy to tune. + +Adding a Year to an Existing Store +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The time axis is fixed at ``zarr-init``, but it can be grown afterwards. +Because time is chunked one year per chunk, appending is a **metadata-only** +edit — existing chunks keep their keys and are never rewritten, however large +the store:: + + # 1. grow every zone's time axis (no fills in flight) + geotessera-registry zarr-extend s3://dest-bucket/tessera.zarr --years 2026 + + # 2. fill the new year, one process per zone as usual + geotessera-registry zarr-fill s3://source-bucket/tessera \ + s3://dest-bucket/tessera.zarr --year 2026 --zones 30 + +The new slice reads back exactly like a freshly initialised year — +embeddings at 0, scales at ``+inf`` ("land, no data yet") — so ``zarr-fill`` +treats it no differently from the original ones, and the per-zone ingestion +registry keys on ``(zone, year)`` so earlier years are untouched. + +Two constraints: + +* **Append only.** Adding a year *earlier* than the current maximum would + renumber every existing chunk's time index, i.e. rewrite the store. It is + refused rather than done silently. +* **Single writer.** Unlike a fill, this rewrites array metadata for every + zone, so it refuses to run while any fill lock is held, and it *does* + re-consolidate afterwards (readers cannot see the new year until it has). + +.. _zarr-parallel-sweep: + +Parallel Per-Zone Sweeps +~~~~~~~~~~~~~~~~~~~~~~~~ + +A UTM zone's pixels live entirely within its own ``utm{zone}`` group, and +shards never straddle zones. That makes ``--zones N`` the natural unit of +parallelism: one process per zone, all writing to the same store. + +Everything a fill mutates is keyed by ``(zone, year)``, and none of it lives +inside the store — build bookkeeping goes to a sibling location so the +published hierarchy contains only Zarr: + +.. code-block:: text + + tessera.zarr/ + zarr.json # shared — consolidation only + utm30/, utm31/, ... # one zone per process + + tessera.zarr.build/ # --state-url to relocate + _registry/utm30_2024.parquet # per-zone ingestion tracking + _registry.parquet # merged view, written by consolidate + _locks/utm30_2024.json # advisory fill lock + +* **Ingestion tracking** is one object per zone/year, so no two jobs + read-modify-write the same file. It records which tiles have already been + written, which is what makes a fill resumable and lets a later run pick up + tiles the manifest has gained since. It is build state, not published + data — a reader of the store never needs it — so it lives in the state + sibling. Stores built before this split kept a ``_registry.parquet`` + inside the hierarchy; that is still read, so they resume correctly. +* **An advisory lock** is taken for the duration of a zone/year fill. It + catches the same zone being launched twice — the case that would silently + corrupt data, because a shard write replaces the whole shard. Object + stores offer no atomic create, so the lock is advisory; ``--force-lock`` + takes over one left behind by a dead run. +* **Consolidation is skipped** by default when ``--zones`` is given, since + the root ``zarr.json`` is the one object all jobs share. + +A sweep therefore looks like:: + + # fan out — one process per zone, in parallel + parallel -j8 geotessera-registry zarr-fill \ + s3://source-bucket/tessera s3://dest-bucket/tessera.zarr \ + --year 2024 --zones {} ::: $(seq 1 60) + + # single-writer finish: merge the per-zone registries, refresh the root + geotessera-registry zarr-consolidate s3://dest-bucket/tessera.zarr + +Each zone job exits non-zero if any of its shards failed, and tiles are only +recorded as written once every shard covering them succeeded — so re-running +the same command retries exactly the unfinished work. + +The one thing that is *not* safe is splitting a single zone across +processes: whole-shard writes mean two jobs with different tile subsets +would erase each other's pixels. Within one job this is handled by +rebuilding each touched shard from all of its tiles, including ones an +earlier run already wrote. + Future Extensions ~~~~~~~~~~~~~~~~~ diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 345ff45..01017a8 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -43,6 +43,7 @@ block_to_embeddings_registry_filename, block_to_landmasks_registry_filename, default_variant, + format_bytes, parse_grid_name, zarr_store_url, ) @@ -3277,32 +3278,362 @@ def _detect_dataset_metadata( return version, variant +def _add_source_args(parser) -> None: + """Register the dataset-selection flags shared by zarr-init and zarr-fill.""" + parser.add_argument( + "--registry-dir", + type=str, + default=None, + help="Local directory containing manifest.parquet / landmasks.parquet " + "(default: auto-detected from base_dir and parents, or pulled from a " + "URL base_dir)", + ) + parser.add_argument( + "--manifest-url", + type=str, + default=None, + help="Location of the embeddings manifest, for mirrors that do not " + "follow the npy//manifest.parquet layout", + ) + parser.add_argument( + "--landmasks-url", + type=str, + default=None, + help="Location of the landmasks registry, for mirrors that do not " + "follow the landmasks//landmasks.parquet layout", + ) + parser.add_argument( + "--dataset-version", + type=str, + default=None, + help="Tessera dataset version (e.g. v1, v1.1). " + "Default: read from tessera_metadata.json in base_dir, else v1.", + ) + parser.add_argument( + "--dataset-variant", + type=str, + default=None, + help="Tessera dataset variant (e.g. vultr, cambridge). " + "Default: read from tessera_metadata.json in base_dir, else the " + "version's published variant.", + ) + + +def _add_state_arg(parser) -> None: + """Register ``--state-url`` for the commands that keep build bookkeeping.""" + parser.add_argument( + "--state-url", + type=str, + default=None, + help="Where to keep build state (ingestion registry, fill locks). " + "Default: .build alongside the store. Kept outside the store " + "so the published Zarr hierarchy contains only Zarr.", + ) + + +def _add_storage_args(parser, prefix: str, label: str, writable: bool = False) -> None: + """Register ``--{prefix}-*`` object-store flags on *parser*. + + Access keys are deliberately absent: they come from the environment + (``AWS_ACCESS_KEY_ID``/``AWS_SECRET_ACCESS_KEY``), a named profile, or an + instance role, so they never appear in a process listing. + + Args: + writable: Also offer ``--{prefix}-acl``, which only affects writes. + """ + group = parser.add_argument_group(f"{label} object-store options") + group.add_argument( + f"--{prefix}-endpoint-url", + type=str, + default=None, + help=f"S3-compatible endpoint for the {label.lower()} " + f"(default: $AWS_ENDPOINT_URL, else the AWS default)", + ) + group.add_argument( + f"--{prefix}-anon", + action="store_true", + help=f"Access the {label.lower()} anonymously (public buckets)", + ) + group.add_argument( + f"--{prefix}-region", + type=str, + default=None, + help=f"Region for the {label.lower()} (default: $AWS_DEFAULT_REGION)", + ) + group.add_argument( + f"--{prefix}-profile", + type=str, + default=None, + help=f"Shared-config profile for the {label.lower()} " + f"(default: $AWS_PROFILE)", + ) + group.add_argument( + f"--{prefix}-requester-pays", + action="store_true", + help=f"Send requester-pays headers to the {label.lower()}", + ) + if writable: + from .remote import OBJECT_ACLS + + group.add_argument( + f"--{prefix}-acl", + type=str, + default=None, + choices=OBJECT_ACLS, + metavar="ACL", + help=f"Canned ACL applied to every object written to the " + f"{label.lower()} (e.g. bucket-owner-full-control when the " + f"bucket belongs to another account)", + ) + + +def _object_store_errors() -> Tuple[type, ...]: + """Exception types worth reporting as a message rather than a traceback. + + Empty when botocore is absent (i.e. the s3 extra is not installed), in + which case no such error can be raised anyway. + """ + try: + from botocore.exceptions import BotoCoreError, ClientError + + return (BotoCoreError, ClientError) + except ImportError: + return () + + +def _report_store_error(e: Exception, console: "Console") -> int: + """Print an object-store or missing-backend failure and return exit 1. + + Error text is escaped: it routinely contains bracketed fragments (and + our own ``geotessera[s3]`` hint) that Rich would otherwise eat as markup. + """ + from rich.markup import escape + + console.print(f"[red]{emoji('❌ ')}{escape(str(e))}[/red]") + if isinstance(e, _object_store_errors() or ()): + console.print( + "Check the endpoint and credentials: --source-endpoint-url / " + "--store-endpoint-url, --*-profile, --*-anon, or the AWS_* " + "environment variables." + ) + return 1 + + +def _storage_options_for( + args, prefix: str, location: str +) -> Optional[Dict[str, Any]]: + """Build fsspec storage options from ``--{prefix}-*`` flags. + + Returns None for local paths so plain filesystem access never picks up + stray ``AWS_*`` environment variables. + """ + from .remote import build_storage_options, is_url + + if not is_url(location): + return None + + def opt(name): + return getattr(args, f"{prefix}_{name}", None) + + return build_storage_options( + endpoint_url=opt("endpoint_url"), + anon=bool(opt("anon")), + region=opt("region"), + profile=opt("profile"), + requester_pays=bool(opt("requester_pays")), + acl=opt("acl"), + ) + + +def _remote_registry_cache_dir(root: str, version_path: str) -> Path: + """Local cache directory for registries pulled from a remote mirror. + + Keyed by a digest of the source root so mirrors with different contents + never share a cached manifest. + """ + if os.name == "nt": + base = Path(os.environ.get("LOCALAPPDATA", "~")).expanduser() + else: + base = Path(os.environ.get("XDG_CACHE_HOME", "~/.cache")).expanduser() + digest = hashlib.sha256(root.encode()).hexdigest()[:12] + return base / "geotessera" / "mirrors" / digest / version_path + + +def _sync_remote_file( + loc: str, + dest: Path, + storage_options: Optional[Dict[str, Any]], + console: "Console", + optional: bool = False, +) -> Optional[Path]: + """Download *loc* to *dest* unless the cached copy already matches. + + Freshness is judged on the object's size and last-modified stamp, which + every S3-compatible endpoint reports on a HEAD — far cheaper than + re-pulling a multi-hundred-megabyte manifest each run. + + A local path is used where it lies; only URLs are ever copied. + """ + import json + + from . import remote + + if not remote.is_url(loc): + local = Path(loc) + if local.exists(): + return local + if optional: + return None + raise FileNotFoundError(loc) + + fs = remote.get_fs(loc, storage_options) + try: + info = fs.info(str(loc)) + except FileNotFoundError: + if optional: + return None + raise + except Exception as e: + if optional: + console.print(f"[yellow]Could not stat {loc}: {e}[/yellow]") + return None + raise + + stamp = { + "size": info.get("size"), + "mtime": str(info.get("LastModified") or info.get("mtime") or ""), + } + meta = dest.with_name(dest.name + ".meta.json") + + if dest.exists() and meta.exists(): + try: + if json.loads(meta.read_text()) == stamp: + console.print(f"[dim]Using cached {dest.name} ({loc})[/dim]") + return dest + except (OSError, ValueError): + pass + + size_str = format_bytes(stamp["size"]) if stamp["size"] else "unknown size" + console.print(f"Downloading {loc} ({size_str})...") + remote.download(loc, dest, storage_options) + meta.write_text(json.dumps(stamp, default=str)) + return dest + + +def _resolve_source(args, console: "Console"): + """Resolve the tile source and its registry for a zarr-init/fill run. + + ``args.base_dir`` is either a local mirror directory (the historical + behaviour) or a URL such as ``s3://bucket/prefix`` pointing at a + repository in the published layout. In the remote case the manifest and + landmask registries are pulled from that same root — or from + ``--manifest-url``/``--landmasks-url`` when a mirror deviates — and + cached locally, so no local tile mirror is needed at all. + + Returns ``(registry, source, dataset_version, dataset_variant)`` where + ``source`` is a :class:`~geotessera.zarr.TileSource`, or None to mean + "use the registry's local mirror". + """ + from .registry import Registry, _parse_dataset_version + from .remote import is_url, join + from .zarr import TileSource + + base_dir = args.base_dir + + def override(loc, name, optional=False): + """Resolve an explicit --manifest-url / --landmasks-url, if given.""" + if not loc: + return None + return _sync_remote_file( + loc, + _remote_registry_cache_dir(str(base_dir), "override") / name, + _storage_options_for(args, "source", loc), + console, + optional=optional, + ) + + if not is_url(base_dir): + base_dir, registry_dir = _find_registry(base_dir, args.registry_dir) + dataset_version, dataset_variant = _detect_dataset_metadata( + base_dir, args.dataset_version, args.dataset_variant + ) + console.print( + f"[cyan]Using dataset version={dataset_version}, " + f"variant={dataset_variant}[/cyan]" + ) + registry = Registry( + version=dataset_version, + variant=dataset_variant, + embeddings_dir=base_dir, + registry_dir=registry_dir, + registry_path=override(args.manifest_url, "manifest.parquet"), + landmasks_registry_path=override( + args.landmasks_url, "landmasks.parquet", optional=True + ), + ) + return registry, None, dataset_version, dataset_variant + + # Remote mirror. + dataset_version = args.dataset_version or "v1" + version_path, version_norm = _parse_dataset_version(dataset_version) + dataset_variant = args.dataset_variant or default_variant(version_norm) + storage_options = _storage_options_for(args, "source", base_dir) + + console.print( + f"[cyan]Streaming tiles from {base_dir} " + f"(version={version_path}, variant={dataset_variant})[/cyan]" + ) + + if args.registry_dir: + registry = Registry( + version=dataset_version, + variant=dataset_variant, + registry_dir=args.registry_dir, + ) + else: + cache_dir = _remote_registry_cache_dir(base_dir, version_path) + manifest_loc = args.manifest_url or join( + base_dir, "npy", version_path, "manifest.parquet" + ) + landmasks_loc = args.landmasks_url or join( + base_dir, "landmasks", version_path, "landmasks.parquet" + ) + manifest_path = _sync_remote_file( + manifest_loc, + cache_dir / "manifest.parquet", + _storage_options_for(args, "source", manifest_loc), + console, + ) + landmasks_path = _sync_remote_file( + landmasks_loc, + cache_dir / "landmasks.parquet", + _storage_options_for(args, "source", landmasks_loc), + console, + optional=True, + ) + registry = Registry( + version=dataset_version, + variant=dataset_variant, + registry_path=manifest_path, + landmasks_registry_path=landmasks_path, + ) + + source = TileSource.for_url(base_dir, version_path, storage_options) + return registry, source, dataset_version, dataset_variant + + def zarr_init_command(args): """Create an empty tessera store with time dimension.""" from rich.console import Console - from .registry import Registry from .zarr import init_store console = Console() - base_dir = args.base_dir - base_dir, registry_dir = _find_registry(base_dir, args.registry_dir) - dataset_version, dataset_variant = _detect_dataset_metadata( - base_dir, args.dataset_version, args.dataset_variant - ) - console.print( - f"[cyan]Using dataset version={dataset_version}, variant={dataset_variant}[/cyan]" - ) - - registry = Registry( - version=dataset_version, - variant=dataset_variant, - embeddings_dir=base_dir, - registry_dir=registry_dir, - ) + registry, _source, _version, _variant = _resolve_source(args, console) years = _parse_int_range(args.years) - output = Path(args.output) + output = args.output + store_options = _storage_options_for(args, "store", output) try: import importlib.metadata @@ -3327,10 +3658,14 @@ def zarr_init_command(args): geotessera_version=version, model_version=model_version, console=console, + storage_options=store_options, + state_url=args.state_url, ) except FileExistsError as e: console.print(f"[red]Error:[/red] {e}") return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) return 0 @@ -3339,46 +3674,95 @@ def zarr_fill_command(args): """Incrementally fill a tessera store with tile data.""" import warnings from rich.console import Console - from .registry import Registry from .zarr import fill_store warnings.filterwarnings("ignore", message="Object at .* is not recognized") console = Console() - base_dir = args.base_dir - base_dir, registry_dir = _find_registry(base_dir, args.registry_dir) - dataset_version, dataset_variant = _detect_dataset_metadata( - base_dir, args.dataset_version, args.dataset_variant - ) - console.print( - f"[cyan]Using dataset version={dataset_version}, variant={dataset_variant}[/cyan]" - ) + registry, source, _version, _variant = _resolve_source(args, console) - registry = Registry( - version=dataset_version, - variant=dataset_variant, - embeddings_dir=base_dir, - registry_dir=registry_dir, - ) - - store_path = Path(args.store_path) + store_path = args.store_path + store_options = _storage_options_for(args, "store", store_path) year = args.year zones = _parse_int_range(args.zones) if args.zones else None - n = fill_store( - registry, - store_path, - year=year, - zones=zones, - console=console, - workers=args.workers, - ) + consolidate = None + if args.consolidate: + consolidate = True + elif args.no_consolidate: + consolidate = False + + try: + n = fill_store( + registry, + store_path, + year=year, + zones=zones, + console=console, + workers=args.workers, + storage_options=store_options, + source=source, + consolidate=consolidate, + force_lock=args.force_lock, + state_url=args.state_url, + ) + except RuntimeError as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) console.print(f"\n{emoji('✅ ')}{n} shards written") return 0 +def zarr_extend_command(args): + """Append new years to an existing store's time axis.""" + import warnings + from rich.console import Console + + from .zarr import extend_store + + warnings.filterwarnings("ignore", message="Object at .* is not recognized") + + console = Console() + store_options = _storage_options_for(args, "store", args.store_path) + years = _parse_int_range(args.years) + zones = _parse_int_range(args.zones) if args.zones else None + + try: + n = extend_store( + args.store_path, + years, + console=console, + storage_options=store_options, + zones=zones, + consolidate=not args.no_consolidate, + force=args.force, + state_url=args.state_url, + ) + except (ValueError, RuntimeError) as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) + + if n == 0: + console.print( + f"{emoji('✅ ')}Nothing to do — every zone already has " + f"{', '.join(str(y) for y in years)}" + ) + return 0 + + console.print( + f"{emoji('✅ ')}{n} zone(s) extended. " + f"Fill them with: geotessera-registry zarr-fill " + f"{args.store_path} --year {years[0]} --zones " + ) + return 0 + + def zarr_consolidate_command(args): """Re-consolidate store metadata after in-place changes.""" from rich.console import Console @@ -3386,22 +3770,44 @@ def zarr_consolidate_command(args): from .zarr import consolidate_store console = Console() + store_options = _storage_options_for(args, "store", args.store_path) try: - n = consolidate_store(args.store_path, console=console) + n = consolidate_store( + args.store_path, + console=console, + storage_options=store_options, + merge_registry=not args.no_merge_registry, + state_url=args.state_url, + ) except FileNotFoundError as e: console.print(f"[red]{emoji('❌ ')}Error: {e}[/red]") - raise SystemExit(1) - except ImportError as e: - console.print( - f"[red]{emoji('❌ ')}Error: {e} " - f"(remote store URLs need the matching fsspec backend, " - f"e.g. s3fs for s3://)[/red]" - ) - raise SystemExit(1) + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) console.print(f"{emoji('✅ ')}Consolidated metadata for {n} nodes") return 0 +def _require_local_store(store_path: str, command: str, console: "Console") -> bool: + """Reject a URL store for the commands that still need a local path. + + Without this the URL is silently mangled by ``Path()`` into something + like ``s3:/bucket/store`` and fails much later with a confusing error. + """ + from .remote import is_url + + if not is_url(store_path): + return True + console.print( + f"[red]{emoji('❌ ')}{command} needs a local store path; " + f"got {store_path}.[/red]\n" + f"Copy the store locally (or mount it) and run it there — only " + f"zarr-init, zarr-fill and zarr-consolidate work against remote " + f"locations." + ) + return False + + def zarr_global_preview_command(args): """Build global EPSG:4326 RGB pyramid from zone-level embeddings.""" import warnings @@ -3411,6 +3817,8 @@ def zarr_global_preview_command(args): warnings.filterwarnings("ignore", message="Object at .* is not recognized") console = Console() + if not _require_local_store(args.store_path, "zarr-global-preview", console): + return 1 store_path = Path(args.store_path) zones = _parse_int_range(args.zones) if args.zones else None @@ -3438,6 +3846,8 @@ def zarr_stretch_command(args): warnings.filterwarnings("ignore", message="Object at .* is not recognized") console = Console() + if not _require_local_store(args.store_path, "zarr-stretch", console): + return 1 store_path = Path(args.store_path) zones = _parse_int_range(args.zones) if args.zones else None @@ -4220,7 +4630,8 @@ def main(): ) zarr_init_parser.add_argument( "base_dir", - help="Base directory containing downloaded tile data", + help="Base directory containing downloaded tile data, or a URL of a " + "repository in the published layout (e.g. s3://bucket/tessera)", ) zarr_init_parser.add_argument( "--years", @@ -4231,45 +4642,34 @@ def main(): "--output", required=True, type=str, - help="Output store path (e.g. tessera.zarr)", - ) - zarr_init_parser.add_argument( - "--registry-dir", - type=str, - default=None, - help="Directory containing manifest.parquet / landmasks.parquet " - "(default: auto-detected from base_dir and parents)", - ) - zarr_init_parser.add_argument( - "--dataset-version", - type=str, - default=None, - help="Tessera dataset version (e.g. v1, v1.1). " - "Default: read from tessera_metadata.json in base_dir, else v1.", - ) - zarr_init_parser.add_argument( - "--dataset-variant", - type=str, - default=None, - help="Tessera dataset variant (e.g. vultr, cambridge). " - "Default: read from tessera_metadata.json in base_dir, else the " - "version's published variant.", + help="Output store path or URL (e.g. tessera.zarr, " + "s3://bucket/tessera.zarr)", ) + _add_source_args(zarr_init_parser) + _add_storage_args(zarr_init_parser, "source", "Tile source") + _add_state_arg(zarr_init_parser) + _add_storage_args(zarr_init_parser, "store", "Output store", writable=True) zarr_init_parser.set_defaults(func=zarr_init_command) # Zarr-fill command zarr_fill_parser = subparsers.add_parser( "zarr-fill", help="Incrementally fill a tessera store with tile data", + description="Fill a tessera store, optionally streaming tiles from a " + "remote bucket into a store on another. Restricting a run to one UTM " + "zone with --zones makes it safe to run many fills concurrently " + "against the same store: each zone owns its own group, tracking file " + "and lock. Run zarr-consolidate once after the sweep.", ) zarr_fill_parser.add_argument( "base_dir", - help="Base directory containing downloaded tile data", + help="Base directory containing downloaded tile data, or a URL of a " + "repository in the published layout (e.g. s3://bucket/tessera)", ) zarr_fill_parser.add_argument( "store_path", type=str, - help="Path to existing tessera store", + help="Path or URL of an existing tessera store", ) zarr_fill_parser.add_argument( "--year", @@ -4280,7 +4680,8 @@ def main(): zarr_fill_parser.add_argument( "--zones", default=None, - help="Zone numbers to fill (e.g. 29-34). Default: all initialised zones", + help="Zone numbers to fill (e.g. 30 or 29-34). " + "Default: all initialised zones", ) zarr_fill_parser.add_argument( "--workers", @@ -4289,33 +4690,79 @@ def main(): help="Number of parallel workers (default: 4)", ) zarr_fill_parser.add_argument( - "--registry-dir", - type=str, - default=None, - help="Directory containing manifest.parquet / landmasks.parquet " - "(default: auto-detected from base_dir and parents)", + "--consolidate", + action="store_true", + help="Rewrite the root consolidated metadata when done. Unsafe while " + "sibling zone fills are running. Default: on for a whole-store fill, " + "off when --zones is given.", ) zarr_fill_parser.add_argument( - "--dataset-version", - type=str, - default=None, - help="Tessera dataset version (e.g. v1, v1.1). " - "Default: read from tessera_metadata.json in base_dir, else v1.", + "--no-consolidate", + action="store_true", + help="Never rewrite the root consolidated metadata.", ) zarr_fill_parser.add_argument( - "--dataset-variant", + "--force-lock", + action="store_true", + help="Take over a zone/year lock left behind by a dead run. Only use " + "this when no other fill is touching the same zone.", + ) + _add_source_args(zarr_fill_parser) + _add_storage_args(zarr_fill_parser, "source", "Tile source") + _add_state_arg(zarr_fill_parser) + _add_storage_args(zarr_fill_parser, "store", "Output store", writable=True) + zarr_fill_parser.set_defaults(func=zarr_fill_command) + + # Zarr-extend command + zarr_extend_parser = subparsers.add_parser( + "zarr-extend", + help="Append new years to an existing store's time axis", + description="Grow a store's time dimension so a new year can be " + "filled. The time axis is chunked one year per chunk, so this is a " + "metadata-only edit: existing data is never rewritten. Years may " + "only be appended to the end. Run it with no fills in flight, then " + "zarr-fill the new year.", + ) + zarr_extend_parser.add_argument( + "store_path", type=str, + help="Path or URL of an existing tessera store", + ) + zarr_extend_parser.add_argument( + "--years", + required=True, + help="Year(s) to append (e.g. 2026 or 2026-2027)", + ) + zarr_extend_parser.add_argument( + "--zones", default=None, - help="Tessera dataset variant (e.g. vultr, cambridge). " - "Default: read from tessera_metadata.json in base_dir, else the " - "version's published variant.", + help="Restrict to these zones (e.g. 29-34). Default: every zone. " + "Leaving zones behind makes the store's time axes disagree, so only " + "use this to finish an interrupted run.", ) - zarr_fill_parser.set_defaults(func=zarr_fill_command) + zarr_extend_parser.add_argument( + "--no-consolidate", + action="store_true", + help="Skip re-consolidating the root metadata. Array metadata has " + "changed, so readers need a consolidate before they see the new year.", + ) + zarr_extend_parser.add_argument( + "--force", + action="store_true", + help="Proceed even if fill locks are present (only when they are stale)", + ) + _add_state_arg(zarr_extend_parser) + _add_storage_args(zarr_extend_parser, "store", "Store", writable=True) + zarr_extend_parser.set_defaults(func=zarr_extend_command) # Zarr-consolidate command zarr_consolidate_parser = subparsers.add_parser( "zarr-consolidate", help="Re-consolidate store metadata after in-place changes", + description="Rewrite a store's root consolidated metadata and merge " + "the per-zone ingestion registries. This is the single-writer step " + "that finishes a parallel zone sweep — run it once, with no fills in " + "flight.", ) zarr_consolidate_parser.add_argument( "store_path", @@ -4323,6 +4770,13 @@ def main(): help="Path or URL of an existing tessera store " "(e.g. tessera.zarr or s3://bucket/store.zarr)", ) + zarr_consolidate_parser.add_argument( + "--no-merge-registry", + action="store_true", + help="Skip merging _registry/ per-zone files into _registry.parquet", + ) + _add_state_arg(zarr_consolidate_parser) + _add_storage_args(zarr_consolidate_parser, "store", "Store", writable=True) zarr_consolidate_parser.set_defaults(func=zarr_consolidate_command) # Zarr-global-preview command @@ -4539,8 +4993,10 @@ def main(): parser.print_help() return - # Execute the command - args.func(args) + # Execute the command. Commands return a shell exit status; propagate it + # so a failed zone in a parallel sweep is visible to the orchestrator + # rather than silently reported as success. + raise SystemExit(args.func(args) or 0) if __name__ == "__main__": diff --git a/geotessera/remote.py b/geotessera/remote.py new file mode 100644 index 0000000..0c05189 --- /dev/null +++ b/geotessera/remote.py @@ -0,0 +1,407 @@ +"""Location-transparent I/O for local paths and remote object stores. + +Every path in the Zarr build pipeline — the tile inputs (NPY embeddings, +scales, landmask GeoTIFFs), the output store, and the store's own tracking +parquet — can be either a local filesystem path or an fsspec URL such as +``s3://bucket/prefix``. This module is the single place that knows the +difference. + +Two design points matter for the remote case: + +* **Byte-range reads, no scratch disk.** A ``.npy`` file is a short header + followed by a flat C-ordered buffer, so the rows a shard needs are one + contiguous range. :func:`read_npy_window` parses the header (one small + ranged GET, memoised per process) and then fetches exactly those rows. + A 158 MB tile costs only the bytes actually written into the shard. + +* **Credentials never go through argv by default.** :func:`build_storage_options` + assembles an fsspec option dict from explicit arguments, falling back to + the standard ``AWS_*`` environment variables, so a sweep can run with keys + supplied by the environment or an instance profile. +""" + +from __future__ import annotations + +import io +import json +import logging +import os +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +import numpy as np + +logger = logging.getLogger(__name__) + +# Protocols we can read/write through fsspec. Anything else without a +# scheme is treated as a local filesystem path. +_URL_MARKER = "://" + + +def is_url(loc: str | Path) -> bool: + """True if *loc* is an fsspec URL rather than a local filesystem path.""" + return _URL_MARKER in str(loc) + + +def protocol_of(loc: str | Path) -> str: + """Return the fsspec protocol for *loc* (``"file"`` for local paths).""" + s = str(loc) + if _URL_MARKER not in s: + return "file" + return s.split(_URL_MARKER, 1)[0] + + +def join(base: str | Path, *parts: str) -> str: + """Join path components, using ``/`` for URLs and os.sep for local paths. + + Returns a ``str`` in both cases so callers can pass the result straight + to zarr, fsspec, or ``open()``. + """ + if is_url(base): + out = str(base).rstrip("/") + for p in parts: + p = str(p).strip("/") + if p: + out = f"{out}/{p}" + return out + return str(Path(base).joinpath(*[str(p) for p in parts])) + + +# Canned ACLs S3 accepts on an object (PutObject's x-amz-acl). Validated up +# front so a typo fails before a long fill rather than on the first PUT. +OBJECT_ACLS = ( + "private", + "public-read", + "public-read-write", + "authenticated-read", + "aws-exec-read", + "bucket-owner-read", + "bucket-owner-full-control", +) + + +def build_storage_options( + endpoint_url: Optional[str] = None, + anon: bool = False, + region: Optional[str] = None, + profile: Optional[str] = None, + requester_pays: bool = False, + acl: Optional[str] = None, + extra: Optional[Dict[str, Any]] = None, +) -> Optional[Dict[str, Any]]: + """Assemble an fsspec storage-options dict for an S3-compatible endpoint. + + Unset arguments fall back to the conventional environment variables + (``AWS_ENDPOINT_URL``, ``AWS_DEFAULT_REGION``/``AWS_REGION``, + ``AWS_PROFILE``), which is how credentials should normally be supplied + so keys never appear in a process listing. Access keys themselves are + left entirely to botocore's own resolution chain (environment, shared + config, instance metadata). + + Args: + acl: Canned ACL to stamp on every object written, e.g. + ``bucket-owner-full-control`` for a bucket owned by another + account. s3fs filters it per operation, so reads are unaffected. + + Returns ``None`` when nothing needs configuring, so callers can pass the + result through to zarr unchanged. + """ + options: Dict[str, Any] = {} + + endpoint_url = endpoint_url or os.environ.get("AWS_ENDPOINT_URL") or None + if endpoint_url: + options["endpoint_url"] = endpoint_url + + if anon: + options["anon"] = True + + region = ( + region + or os.environ.get("AWS_DEFAULT_REGION") + or os.environ.get("AWS_REGION") + or None + ) + if region: + options["client_kwargs"] = {"region_name": region} + + profile = profile or os.environ.get("AWS_PROFILE") or None + if profile and not anon: + options["profile"] = profile + + if requester_pays: + options["requester_pays"] = True + + if acl: + if acl not in OBJECT_ACLS: + raise ValueError( + f"Unknown canned ACL {acl!r}. Expected one of: " + + ", ".join(OBJECT_ACLS) + ) + options["s3_additional_kwargs"] = {"ACL": acl} + + if extra: + options.update(extra) + + return options or None + + +def _options_key(storage_options: Optional[Dict[str, Any]]) -> str: + """Stable hashable key for a storage-options dict.""" + if not storage_options: + return "" + return json.dumps(storage_options, sort_keys=True, default=str) + + +@lru_cache(maxsize=8) +def _filesystem_cached(protocol: str, options_key: str): + """Build (and memoise) one filesystem per protocol/options pair. + + Worker processes call this on every tile read, so the cache keeps a + single connection pool per process rather than one per file. + """ + import fsspec + + options = json.loads(options_key) if options_key else {} + if protocol in ("file", "local"): + # Object stores have no directories to create; a file:// store still + # needs them, and fsspec will not make them unless asked. + options.setdefault("auto_mkdir", True) + if protocol in ("http", "https"): + # The Source Cooperative CDN rejects some default user agents. + from . import __version__ + + options.setdefault( + "client_kwargs", {"headers": {"User-Agent": f"geotessera/{__version__}"}} + ) + try: + return fsspec.filesystem(protocol, **options) + except ImportError as e: + hint = ( + "install the s3 extra: pip install 'geotessera[s3]'" + if protocol == "s3" + else f"install the fsspec backend for {protocol}://" + ) + raise ImportError( + f"{protocol}:// locations need an fsspec backend — {hint} ({e})" + ) from e + + +def get_fs(loc: str | Path, storage_options: Optional[Dict[str, Any]] = None): + """Return an fsspec filesystem for *loc*, or ``None`` for local paths.""" + if not is_url(loc): + return None + return _filesystem_cached(protocol_of(loc), _options_key(storage_options)) + + +def exists(loc: str | Path, storage_options: Optional[Dict[str, Any]] = None) -> bool: + """True if *loc* exists, locally or remotely.""" + fs = get_fs(loc, storage_options) + if fs is None: + return Path(loc).exists() + return bool(fs.exists(str(loc))) + + +def read_bytes( + loc: str | Path, storage_options: Optional[Dict[str, Any]] = None +) -> bytes: + """Read a whole object/file into memory.""" + fs = get_fs(loc, storage_options) + if fs is None: + return Path(loc).read_bytes() + return fs.cat_file(str(loc)) + + +def write_bytes( + loc: str | Path, + data: bytes, + storage_options: Optional[Dict[str, Any]] = None, +) -> None: + """Write *data* to *loc*, creating parent directories for local paths.""" + fs = get_fs(loc, storage_options) + if fs is None: + path = Path(loc) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return + with fs.open(str(loc), "wb") as f: + f.write(data) + + +def remove(loc: str | Path, storage_options: Optional[Dict[str, Any]] = None) -> None: + """Delete *loc*, ignoring a missing target.""" + fs = get_fs(loc, storage_options) + try: + if fs is None: + Path(loc).unlink(missing_ok=True) + else: + fs.rm_file(str(loc)) + except FileNotFoundError: + pass + except Exception as e: # pragma: no cover - best-effort cleanup + logger.warning(f"Could not remove {loc}: {e}") + + +def listdir( + loc: str | Path, storage_options: Optional[Dict[str, Any]] = None +) -> list[str]: + """List the immediate children of a directory/prefix, as full locations. + + Returns an empty list when the directory does not exist. + """ + fs = get_fs(loc, storage_options) + if fs is None: + path = Path(loc) + if not path.is_dir(): + return [] + return [str(p) for p in sorted(path.iterdir())] + if not fs.exists(str(loc)): + return [] + protocol = protocol_of(loc) + out = [] + for entry in sorted(fs.ls(str(loc), detail=False)): + # fsspec strips the protocol from listing results; restore it so the + # entries are usable as standalone locations. Both shapes round-trip: + # "bucket/key" -> "s3://bucket/key", "/abs/path" -> "file:///abs/path". + out.append(entry if is_url(entry) else f"{protocol}://{entry}") + return out + + +def download( + loc: str | Path, + dest: Path, + storage_options: Optional[Dict[str, Any]] = None, +) -> Path: + """Copy a remote object to a local file, writing atomically.""" + dest = Path(dest) + dest.parent.mkdir(parents=True, exist_ok=True) + tmp = dest.with_name(f".{dest.name}.tmp") + tmp.write_bytes(read_bytes(loc, storage_options)) + tmp.replace(dest) + return dest + + +# --------------------------------------------------------------------------- +# NPY windowed reads +# --------------------------------------------------------------------------- + +# Per-process memo of parsed .npy headers: loc -> (shape, dtype, data_offset). +# Each tile is read once per shard it overlaps, so this saves a small ranged +# GET on the second and subsequent reads. +_npy_header_cache: Dict[str, Tuple[Tuple[int, ...], np.dtype, int]] = {} + +# Enough for any 1.0/2.0 header; the spec pads them to a 64-byte boundary and +# real Tessera tiles use 128 bytes. +_HEADER_PROBE_BYTES = 4096 + + +def _parse_npy_header(head: bytes) -> Tuple[Tuple[int, ...], np.dtype, int]: + """Parse a .npy header prefix into (shape, dtype, data_offset). + + Raises ValueError for Fortran-ordered arrays, whose rows are not + contiguous and so cannot be range-read a row band at a time. + """ + f = io.BytesIO(head) + version = np.lib.format.read_magic(f) + if version == (1, 0): + shape, fortran_order, dtype = np.lib.format.read_array_header_1_0(f) + elif version == (2, 0): + shape, fortran_order, dtype = np.lib.format.read_array_header_2_0(f) + else: + raise ValueError(f"Unsupported .npy format version {version}") + if fortran_order: + raise ValueError("Fortran-ordered .npy files are not supported") + return tuple(shape), dtype, f.tell() + + +def _npy_header(loc: str, storage_options: Optional[Dict[str, Any]]): + """Fetch and memoise the header of a remote .npy file.""" + cached = _npy_header_cache.get(loc) + if cached is not None: + return cached + + fs = get_fs(loc, storage_options) + # Keyword arguments are mandatory: S3FileSystem._cat_file takes version_id + # as its second positional parameter, so a positional (start, end) would + # silently read from the wrong offset. + head = fs.cat_file(loc, start=0, end=_HEADER_PROBE_BYTES) + parsed = _parse_npy_header(head) + _npy_header_cache[loc] = parsed + return parsed + + +def read_npy_window( + loc: str | Path, + row_start: int, + row_end: int, + col_start: int, + col_end: int, + storage_options: Optional[Dict[str, Any]] = None, +) -> np.ndarray: + """Read ``arr[row_start:row_end, col_start:col_end, ...]`` from a .npy file. + + Local files are memory-mapped exactly as before. Remote files are served + by a single ranged GET covering the requested rows: rows are contiguous in + a C-ordered array, so the read is one range even though the column window + is strided. Full-width windows therefore transfer no wasted bytes, and a + narrow column window costs at most the row band it sits in. + + The returned array may be a read-only view of the fetched buffer; callers + that mutate it must copy first. + """ + fs = get_fs(loc, storage_options) + if fs is None: + arr = np.load(str(loc), mmap_mode="r") + return arr[row_start:row_end, col_start:col_end, ...] + + loc = str(loc) + shape, dtype, offset = _npy_header(loc, storage_options) + row_end = min(row_end, shape[0]) + if row_end <= row_start: + return np.empty((0, max(col_end - col_start, 0)) + shape[2:], dtype=dtype) + + row_bytes = int(np.prod(shape[1:])) * dtype.itemsize + start = offset + row_start * row_bytes + end = offset + row_end * row_bytes + buf = fs.cat_file(loc, start=start, end=end) + expected = (row_end - row_start) * row_bytes + if len(buf) != expected: + raise OSError( + f"Short read from {loc}: got {len(buf)} of {expected} bytes " + f"for rows {row_start}:{row_end}" + ) + + band = np.frombuffer(buf, dtype=dtype).reshape((row_end - row_start,) + shape[1:]) + return band[:, col_start:col_end, ...] + + +def read_tiff_window( + loc: str | Path, + row_start: int, + row_end: int, + col_start: int, + col_end: int, + band: int = 1, + storage_options: Optional[Dict[str, Any]] = None, +) -> np.ndarray: + """Read a window from a single-band GeoTIFF, locally or remotely. + + Landmask TIFFs are small (tens of KB compressed), so the remote path + fetches the whole object once and opens it from memory rather than + paying GDAL's multi-request /vsicurl dance. + """ + import rasterio + from rasterio.windows import Window + + window = Window.from_slices((row_start, row_end), (col_start, col_end)) + + fs = get_fs(loc, storage_options) + if fs is None: + with rasterio.open(str(loc)) as src: + return src.read(band, window=window) + + from rasterio.io import MemoryFile + + data = fs.cat_file(str(loc)) + with MemoryFile(data) as memfile, memfile.open() as src: + return src.read(band, window=window) diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 0fe4252..199c8fe 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -10,7 +10,15 @@ x # float64 (W,) y # float64 (H,) band # int32 (B,) - _registry.parquet # tile ingestion tracking + +The store contains nothing but Zarr. Build-time bookkeeping lives in a +sibling location, ``.build`` by default: + + tessera.zarr.build/ + _registry/utm{zone:02d}_{year}.parquet # per-zone ingestion tracking + _registry.parquet # merged tracking (written by consolidate) + _locks/utm{zone:02d}_{year}.json # advisory fill locks + _preview/zone_{zone}_done # global-preview resume markers Dimension order: (time, band, y, x) — ML-standard NCHW. Inner chunks: (1, 128, 32, 32), Shards: (1, 128, 4096, 4096). @@ -19,10 +27,28 @@ NaN = water (permanent, from landmask) +inf = land, no data yet (set at init, replaced by real scale on fill) finite = valid data + +Locations +--------- +The store and the tile inputs are addressed as *locations*: either local +filesystem paths or fsspec URLs (``s3://bucket/prefix``). :mod:`geotessera.remote` +resolves the difference, so a store on one S3 node can be filled from tiles on +another without any local mirror. + +Parallel fills +-------------- +A UTM zone's pixels live entirely within its own ``utm{zone}`` group, and +shards never straddle zones, so one process per zone can fill the same store +concurrently. All per-fill state is keyed by (zone, year) — the ingestion +registry and the advisory lock — so no two zone jobs touch the same object. +The one shared object, the root ``zarr.json``, is only rewritten by +consolidation, which a zone fill skips by default when ``zones`` is set; +run ``zarr-consolidate`` once after the sweep instead. """ from __future__ import annotations +import io import logging import math import os @@ -119,6 +145,272 @@ class ShardSpec: time_index: int = 0 +# --------------------------------------------------------------------------- +# Locations: the store and the tile inputs +# --------------------------------------------------------------------------- + + +@dataclass +class StoreLocation: + """A Zarr store addressed by local path or fsspec URL. + + Wraps the handful of operations the build pipeline needs beyond the Zarr + API itself — existence checks and reading/writing build-state objects — + so callers never branch on local-vs-remote. + """ + + url: str + storage_options: Optional[Dict[str, Any]] = None + state_url: Optional[str] = None + state_storage_options: Optional[Dict[str, Any]] = None + + @classmethod + def resolve( + cls, + store: "str | Path | StoreLocation", + storage_options: Optional[Dict[str, Any]] = None, + state_url: Optional[str] = None, + state_storage_options: Optional[Dict[str, Any]] = None, + ) -> "StoreLocation": + """Coerce a path, URL, or existing location into a StoreLocation.""" + if isinstance(store, StoreLocation): + return store + return cls(str(store), storage_options, state_url, state_storage_options) + + @property + def state(self) -> "StoreLocation": + """Where build-time state lives — a sibling of the store, not inside it. + + The ingestion registry and fill locks are the builder's bookkeeping, + not published data. Keeping them out of the store leaves a clean Zarr + hierarchy: nothing for readers to trip over and nothing for + ``consolidate_metadata`` to warn about. Defaults to ``.build`` + alongside the store, so a sweep on another host still finds it. + """ + return StoreLocation( + self.state_url or f"{self.url.rstrip('/')}.build", + self.state_storage_options + if self.state_storage_options is not None + else self.storage_options, + ) + + @property + def is_remote(self) -> bool: + from . import remote + + return remote.is_url(self.url) + + def join(self, *parts: str) -> str: + from . import remote + + return remote.join(self.url, *parts) + + def exists(self, *parts: str) -> bool: + from . import remote + + return remote.exists(self.join(*parts), self.storage_options) + + def read_bytes(self, *parts: str) -> bytes: + from . import remote + + return remote.read_bytes(self.join(*parts), self.storage_options) + + def write_bytes(self, data: bytes, *parts: str) -> None: + from . import remote + + remote.write_bytes(self.join(*parts), data, self.storage_options) + + def remove(self, *parts: str) -> None: + from . import remote + + remote.remove(self.join(*parts), self.storage_options) + + def listdir(self, *parts: str) -> List[str]: + from . import remote + + return remote.listdir(self.join(*parts), self.storage_options) + + def _ensure_backend(self) -> None: + """Fail early, with an actionable message, if the backend is missing. + + zarr raises its own import error deep inside store construction; this + surfaces ours (which names the ``geotessera[s3]`` extra) first. + """ + if self.is_remote: + from . import remote + + remote.get_fs(self.url, self.storage_options) + + def as_zarr_store(self, read_only: bool = False): + """Return something ``zarr`` accepts as a store. + + Local paths pass through as strings; remote URLs become an + ``FsspecStore`` so credentials/endpoint reach APIs like + ``consolidate_metadata`` that take no ``storage_options``. + """ + if not self.is_remote: + return self.url + self._ensure_backend() + from zarr.storage import FsspecStore + + return FsspecStore.from_url( + self.url, storage_options=self.storage_options, read_only=read_only + ) + + def open_group( + self, + mode: str = "r", + path: Optional[str] = None, + zarr_format: Optional[int] = None, + use_consolidated: Optional[bool] = False, + ) -> "zarr.Group": + """Open the store (or a group within it) with the right backend.""" + import zarr + + self._ensure_backend() + kwargs: Dict[str, Any] = { + "mode": mode, + "use_consolidated": use_consolidated, + } + if path is not None: + kwargs["path"] = path + if zarr_format is not None: + kwargs["zarr_format"] = zarr_format + if self.storage_options: + kwargs["storage_options"] = self.storage_options + return zarr.open_group(self.url, **kwargs) + + def __str__(self) -> str: + return self.url + + +@dataclass +class TileSource: + """Where the NPY tiles and landmask GeoTIFFs for a fill live. + + ``embeddings_root`` contains ``{year}/grid_{lon}_{lat}/grid_{lon}_{lat}.npy`` + and its ``_scales.npy`` sibling; ``landmasks_root`` contains + ``grid_{lon}_{lat}.tiff``. Both are locations, so a fill can stream from a + remote bucket with no local mirror. + """ + + embeddings_root: str + landmasks_root: str + storage_options: Optional[Dict[str, Any]] = None + + @property + def is_remote(self) -> bool: + from . import remote + + return remote.is_url(self.embeddings_root) + + def embedding_locations(self, lon: float, lat: float, year: int) -> Tuple[str, str]: + """Return (embedding, scales) locations for a tile.""" + from . import remote + from .registry import tile_to_embedding_paths + + emb_rel, scales_rel = tile_to_embedding_paths(lon, lat, year) + return ( + remote.join(self.embeddings_root, emb_rel.as_posix()), + remote.join(self.embeddings_root, scales_rel.as_posix()), + ) + + def landmask_location(self, lon: float, lat: float) -> str: + """Return the landmask location for a tile.""" + from . import remote + from .registry import tile_to_landmask_filename + + return remote.join(self.landmasks_root, tile_to_landmask_filename(lon, lat)) + + @classmethod + def for_url( + cls, + root: str, + version_path: str, + storage_options: Optional[Dict[str, Any]] = None, + ) -> "TileSource": + """Build a source from a repository root in the published layout. + + The Source Cooperative repository — and any mirror of it — lays tiles + out as ``{root}/npy/{version}/{year}/grid_.../`` with landmasks under + ``{root}/landmasks/{version}/``. ``root`` may be an ``s3://`` URL for + a credentialed mirror or an ``https://`` URL for the public front. + """ + from . import remote + + return cls( + embeddings_root=remote.join(root, "npy", version_path), + landmasks_root=remote.join(root, "landmasks", version_path), + storage_options=storage_options, + ) + + @classmethod + def for_local_mirror(cls, registry: "Registry") -> "TileSource": + """Build a source from a registry's local ``embeddings_dir``. + + The local mirror can be in two shapes: + + * flat: ``/global_0.1_degree_representation//...`` (what + the geotessera-download CLI writes — variant info in a sidecar) + * S3-mirror: ``//global_0.1_degree_representation/ + /...`` (what ``aws s3 cp --recursive`` produces) + + The S3-mirror layout is preferred when present so users keeping a + multi-version mirror under one root can point zarr-fill at the top. + Landmasks are STRICTLY per-version — no cross-version fallback. Each + Tessera version has its own landmask grid and mixing them silently + would corrupt water masking. + """ + from .registry import ( + EMBEDDINGS_DIR_NAME, + LANDMASKS_DIR_NAME, + TESSERA_MIRROR_ENDPOINT, + TESSERA_MIRROR_REPO, + ) + + def landmask_sync_hint(dest: Path) -> str: + return ( + f" aws s3 sync --no-sign-request " + f"--endpoint-url {TESSERA_MIRROR_ENDPOINT} " + f"s3://{TESSERA_MIRROR_REPO}/landmasks/" + f"{registry._version_path}/ {dest}/" + ) + + emb_candidate = ( + registry._embeddings_dir / registry._version_path / EMBEDDINGS_DIR_NAME + ) + lm_s3_mirror = ( + registry._embeddings_dir / registry._version_path / LANDMASKS_DIR_NAME + ) + lm_flat = registry._embeddings_dir / LANDMASKS_DIR_NAME + + if emb_candidate.exists(): + base_emb = str(emb_candidate) + # When embeddings are in S3-mirror layout, landmasks must match. + # Don't fall back to the flat layout — that would silently pick up + # the wrong version's landmasks. + if lm_s3_mirror.exists(): + base_lm = str(lm_s3_mirror) + else: + raise FileNotFoundError( + f"Landmask directory not found for {registry._version_path}: " + f"expected {lm_s3_mirror}. Landmasks are per-version and " + f"cannot be reused across versions. Fetch them with:\n" + + landmask_sync_hint(lm_s3_mirror) + ) + else: + base_emb = str(registry._embeddings_dir / EMBEDDINGS_DIR_NAME) + if lm_flat.exists(): + base_lm = str(lm_flat) + else: + raise FileNotFoundError( + f"Landmask directory not found: expected {lm_flat}. " + f"Fetch them with:\n" + landmask_sync_hint(lm_flat) + ) + + return cls(embeddings_root=base_emb, landmasks_root=base_lm) + + # --------------------------------------------------------------------------- # UTM helpers # --------------------------------------------------------------------------- @@ -161,23 +453,25 @@ def _load_landmask_slice( row_end: int, col_start: int, col_end: int, + storage_options: Optional[Dict[str, Any]] = None, ) -> np.ndarray: - """Load a sub-region of a landmask GeoTIFF using a rasterio window. + """Load a sub-region of a landmask GeoTIFF, locally or from a remote store. Returns a 2D uint8 array where 0 = water. If the landmask cannot be read (missing file, shape mismatch, etc.) returns all-ones (all land) so that no pixels are masked. """ - import rasterio - from rasterio.windows import Window + from . import remote try: - with rasterio.open(landmask_path) as src: - window = Window.from_slices( - (row_start, row_end), - (col_start, col_end), - ) - return src.read(1, window=window) + return remote.read_tiff_window( + landmask_path, + row_start, + row_end, + col_start, + col_end, + storage_options=storage_options, + ) except Exception as e: logger.warning(f"Failed to read landmask slice from {landmask_path}: {e}") return np.ones((row_end - row_start, col_end - col_start), dtype=np.uint8) @@ -193,20 +487,18 @@ def gather_tile_infos( year: int, zones: Optional[List[int]] = None, console: Optional["rich.console.Console"] = None, + source: Optional[TileSource] = None, ) -> Dict[int, List[TileInfo]]: """Gather tile metadata and group by UTM zone. Computes grid info deterministically from coordinates (no file I/O). + + Args: + source: Where the tile inputs live. Defaults to the registry's local + mirror; pass a :class:`TileSource` built with + :meth:`TileSource.for_url` to stream from a remote bucket. """ from rasterio.transform import Affine - from .registry import ( - EMBEDDINGS_DIR_NAME, - LANDMASKS_DIR_NAME, - TESSERA_MIRROR_ENDPOINT, - TESSERA_MIRROR_REPO, - tile_to_embedding_paths, - tile_to_landmask_filename, - ) # Get tiles for this year from MultiIndex, filtering to those with data gdf = registry._registry_gdf @@ -247,66 +539,18 @@ def gather_tile_infos( # Build TileInfos using computed grid (no file I/O) from pyproj import Transformer as ProjTransformer - # The local mirror can be in two shapes: - # * flat: /global_0.1_degree_representation//... (what - # the geotessera-download CLI writes — variant info in sidecar) - # * S3-mirror: //global_0.1_degree_representation/ - # /... (what `aws s3 cp --recursive` produces) - # Prefer the S3-mirror layout when it exists so users who keep a - # multi-version mirror under one root can point zarr-fill at the top. - # Landmasks are STRICTLY per-version — no cross-version fallback. Each - # Tessera version has its own landmask grid and mixing them silently - # would corrupt water masking. - def landmask_sync_hint(dest: Path) -> str: - return ( - f" aws s3 sync --no-sign-request " - f"--endpoint-url {TESSERA_MIRROR_ENDPOINT} " - f"s3://{TESSERA_MIRROR_REPO}/landmasks/" - f"{registry._version_path}/ {dest}/" - ) + if source is None: + source = TileSource.for_local_mirror(registry) - emb_candidate = ( - registry._embeddings_dir / registry._version_path / EMBEDDINGS_DIR_NAME - ) - lm_s3_mirror = ( - registry._embeddings_dir / registry._version_path / LANDMASKS_DIR_NAME - ) - lm_flat = registry._embeddings_dir / LANDMASKS_DIR_NAME - - if emb_candidate.exists(): - base_emb = str(emb_candidate) - # When embeddings are in S3-mirror layout, landmasks must match. - # Don't fall back to the flat layout — that would silently pick up - # the wrong version's landmasks. - if lm_s3_mirror.exists(): - base_lm = str(lm_s3_mirror) - else: - raise FileNotFoundError( - f"Landmask directory not found for {registry._version_path}: " - f"expected {lm_s3_mirror}. Landmasks are per-version and " - f"cannot be reused across versions. Fetch them with:\n" - + landmask_sync_hint(lm_s3_mirror) - ) - else: - base_emb = str(registry._embeddings_dir / EMBEDDINGS_DIR_NAME) - if lm_flat.exists(): - base_lm = str(lm_flat) - else: - raise FileNotFoundError( - f"Landmask directory not found: expected {lm_flat}. " - f"Fetch them with:\n" + landmask_sync_hint(lm_flat) - ) zones_dict: Dict[int, List[TileInfo]] = {} transformer_cache: Dict[int, ProjTransformer] = {} pixel_size = 10.0 for tile_year, tile_lon, tile_lat in tiles: - emb_rel, scales_rel = tile_to_embedding_paths(tile_lon, tile_lat, tile_year) - emb_path = os.path.join(base_emb, emb_rel) - scales_path = os.path.join(base_emb, scales_rel) - landmask_path = os.path.join( - base_lm, tile_to_landmask_filename(tile_lon, tile_lat) + emb_path, scales_path = source.embedding_locations( + tile_lon, tile_lat, tile_year ) + landmask_path = source.landmask_location(tile_lon, tile_lat) # Compute EPSG and zone from coordinates zone_num = int(math.floor((tile_lon + 180) / 6)) + 1 @@ -443,6 +687,18 @@ def _execute(pool): # --------------------------------------------------------------------------- +def _preview_marker_path(store_path: Path, zone_num: int) -> Path: + """Resume marker for a zone's global-preview reprojection. + + Kept in the state sibling (``.build/_preview/``) rather than the + store, for the same reason as the ingestion registry: the published Zarr + hierarchy should contain only Zarr. + """ + return Path(f"{str(store_path).rstrip('/')}.build") / "_preview" / ( + f"zone_{zone_num}_done" + ) + + def _zone_output_bounds( zone_epsg: int, zone_transform: list, @@ -695,12 +951,36 @@ def _tile_pixel_offset( # --------------------------------------------------------------------------- +def shard_coords_for_tiles( + tile_infos: List[TileInfo], + grid: UnifiedZoneGrid, +) -> set: + """Return the set of (shard_row, shard_col) covered by these tiles.""" + coords = set() + for ti in tile_infos: + row, col = _tile_pixel_offset(ti, grid) + for sr in range(row // SHARD_SIZE, (row + ti.height - 1) // SHARD_SIZE + 1): + for sc in range(col // SHARD_SIZE, (col + ti.width - 1) // SHARD_SIZE + 1): + coords.add((sr, sc)) + return coords + + def build_shard_index( tile_infos: List[TileInfo], grid: UnifiedZoneGrid, time_index: int, + restrict_to: Optional[set] = None, ) -> List[ShardSpec]: - """Build shard index for one year's tiles against a unified zone grid.""" + """Build shard index for one year's tiles against a unified zone grid. + + Args: + restrict_to: Optional set of (shard_row, shard_col) to emit specs for. + A shard write replaces the whole shard, so an incremental fill has + to pass *every* tile overlapping the shards it rewrites — not just + the new ones — or previously written neighbours are zeroed out. + Callers get that by passing all of the zone's tiles here along with + the shard coordinates the new tiles touch. + """ shard_map: Dict[Tuple[int, int], List[ShardTileOverlap]] = {} for ti in tile_infos: @@ -714,6 +994,8 @@ def build_shard_index( for sr in range(sr_start, sr_end + 1): for sc in range(sc_start, sc_end + 1): + if restrict_to is not None and (sr, sc) not in restrict_to: + continue shard_top = sr * SHARD_SIZE shard_left = sc * SHARD_SIZE @@ -838,12 +1120,14 @@ def _compute_zone_grid_from_landmask( def init_store( registry: "Registry", - output_path: Path, + output_path: "str | Path | StoreLocation", years: List[int], geotessera_version: str = "unknown", model_version: str = "1.0", console: Optional["rich.console.Console"] = None, -) -> Path: + storage_options: Optional[Dict[str, Any]] = None, + state_url: Optional[str] = None, +) -> str: """Create a tessera store with time dimension from the landmask registry. Creates all UTM zones that have landmask coverage. For each zone, the @@ -855,18 +1139,22 @@ def init_store( - +inf = land, no data yet (replaced by real scale values during fill) No embedding data is written. The embeddings array stays at fill_value (0). + + ``output_path`` may be a local path or an fsspec URL such as + ``s3://bucket/tessera.zarr``; the whole store is metadata-only at this + point, so initialising directly on the target object store is cheap. """ import zarr - output_path = Path(output_path) - if output_path.exists(): - raise FileExistsError(f"Store already exists: {output_path}") + store = StoreLocation.resolve(output_path, storage_options, state_url) + if store.exists("zarr.json") or (not store.is_remote and Path(store.url).exists()): + raise FileExistsError(f"Store already exists: {store}") years = sorted(years) T = len(years) if console: - console.print(f"Initialising store at [bold]{output_path}[/bold]") + console.print(f"Initialising store at [bold]{store}[/bold]") console.print(f" Years: {years[0]}-{years[-1]} ({T} time steps)") # Get landmask coverage grouped by UTM zone @@ -880,7 +1168,7 @@ def init_store( # Create root group via zarr API (not manual JSON) so consolidation # preserves attributes correctly. - root = zarr.open_group(str(output_path), mode="w", zarr_format=3) + root = store.open_group(mode="w", zarr_format=3, use_consolidated=None) root.attrs.update( { "zarr_conventions": [GEOEMB_CONVENTION], @@ -926,39 +1214,37 @@ def init_store( f"{n_shards_x}x{n_shards_y} shards[/dim]" ) - _create_zone_group(grid, output_path) + _create_zone_group(grid, store) - # Create empty tile registry - _init_tile_registry(output_path) + # Nothing else is written into the store: ingestion tracking and locks + # are build state and live in the state sibling, created on first fill. # Consolidate metadata so HTTP readers can discover the hierarchy import warnings with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Consolidated metadata") - zarr.consolidate_metadata(str(output_path)) + warnings.filterwarnings("ignore", message="Object at .* is not recognized") + zarr.consolidate_metadata(store.as_zarr_store()) if console: console.print( " [green]Store initialised (metadata only, no data written)[/green]" ) - return output_path + return store.url def _create_zone_group( grid: UnifiedZoneGrid, - store_path: Path, + store_location: StoreLocation, ) -> "zarr.Group": """Create a zone group with empty (T, B, H, W) arrays.""" - import zarr from zarr.codecs import BloscCodec zone_group = _zone_group_name(grid.zone) - root_reopen = zarr.open_group( - str(store_path), mode="r+", zarr_format=3, use_consolidated=False - ) + root_reopen = store_location.open_group(mode="r+", zarr_format=3) store = root_reopen.create_group(zone_group) T = len(grid.years) @@ -1054,18 +1340,32 @@ def _create_zone_group( # --------------------------------------------------------------------------- # Tile registry (GeoParquet tracking which tiles are written) # --------------------------------------------------------------------------- +# This is build bookkeeping, so it lives in the state sibling rather than the +# store: an incremental fill needs to know which tiles it already wrote, but a +# reader of the published store never does. Tracking is sharded by (zone, +# year) under ``_registry/`` so that concurrent per-zone fills never +# read-modify-write the same object, and ``consolidate_store`` merges the +# parts into one ``_registry.parquet``. +# +# Stores built before the split keep a single ``_registry.parquet`` at the +# store root; it is still read so those stores resume correctly, but nothing +# is written back into the store any more. + +REGISTRY_DIR_NAME = "_registry" +MERGED_REGISTRY_NAME = "_registry.parquet" +LEGACY_REGISTRY_NAME = MERGED_REGISTRY_NAME -def _registry_path(store_path: Path) -> Path: - return store_path / "_registry.parquet" +def _zone_registry_name(zone: int, year: int) -> str: + return f"{_zone_group_name(zone)}_{year}.parquet" -def _init_tile_registry(store_path: Path) -> None: - """Create an empty tile registry parquet file.""" +def _empty_tile_registry() -> "geopandas.GeoDataFrame": + """An empty registry frame with the canonical schema.""" import geopandas as gpd import pandas as pd - schema = gpd.GeoDataFrame( + return gpd.GeoDataFrame( { "year": pd.array([], dtype="int32"), "zone": pd.array([], dtype="int32"), @@ -1075,35 +1375,229 @@ def _init_tile_registry(store_path: Path) -> None: "geometry": gpd.array.GeometryArray( gpd.points_from_xy([], []), ), - } + }, + crs="EPSG:4326", ) - schema.to_parquet(str(_registry_path(store_path))) -def _load_tile_registry(store_path: Path) -> "geopandas.GeoDataFrame": - """Load the tile registry, or create it if missing.""" +def _read_parquet_at(store: StoreLocation, *parts: str): + """Read a GeoParquet object from the store, or None if absent.""" import geopandas as gpd - path = _registry_path(store_path) - if path.exists(): - return gpd.read_parquet(str(path)) - _init_tile_registry(store_path) - return gpd.read_parquet(str(path)) + if not store.exists(*parts): + return None + try: + return gpd.read_parquet(io.BytesIO(store.read_bytes(*parts))) + except Exception as e: + logger.warning(f"Could not read {store.join(*parts)}: {e}") + return None + + +def _write_parquet_at(store: StoreLocation, gdf, *parts: str) -> None: + """Write a GeoParquet object into the store (local path or remote URL).""" + buf = io.BytesIO() + gdf.to_parquet(buf) + store.write_bytes(buf.getvalue(), *parts) + + +_MERGED_UNSET = object() + + +def load_merged_registry(store: StoreLocation): + """Read the merged ingestion registry, from the state dir or an old store. + + Returns None when neither exists (a store nobody has filled yet). + """ + merged = _read_parquet_at(store.state, MERGED_REGISTRY_NAME) + if merged is not None: + return merged + # Stores built before the split kept it inside the Zarr hierarchy. + return _read_parquet_at(store, LEGACY_REGISTRY_NAME) -def _save_tile_registry(store_path: Path, gdf: "geopandas.GeoDataFrame") -> None: - """Save the tile registry.""" - gdf.to_parquet(str(_registry_path(store_path))) +def _get_written_tiles( + store: StoreLocation, + year: int, + zone: int, + merged=_MERGED_UNSET, +) -> set: + """Return set of (tile_lon, tile_lat) already written for a year/zone. + Reads this zone/year's own tracking file, then unions in any rows the + merged registry holds for the same zone/year — including one left inside + an older store, so those resume correctly too. -def _get_written_tiles(store_path: Path, year: int, zone: int) -> set: - """Return set of (tile_lon, tile_lat) already written for a year/zone.""" - gdf = _load_tile_registry(store_path) - if gdf.empty: - return set() - mask = (gdf["year"] == year) & (gdf["zone"] == zone) - subset = gdf[mask] - return set(zip(subset["tile_lon"], subset["tile_lat"])) + Args: + merged: A pre-loaded merged registry frame (or None if there isn't + one). It covers the whole store, so a multi-zone fill should read + it once and pass it in rather than re-fetching it per zone/year. + """ + written: set = set() + + zone_gdf = _read_parquet_at( + store.state, REGISTRY_DIR_NAME, _zone_registry_name(zone, year) + ) + if zone_gdf is not None and not zone_gdf.empty: + written |= set(zip(zone_gdf["tile_lon"], zone_gdf["tile_lat"])) + + if merged is _MERGED_UNSET: + merged = load_merged_registry(store) + if merged is not None and not merged.empty: + mask = (merged["year"] == year) & (merged["zone"] == zone) + subset = merged[mask] + written |= set(zip(subset["tile_lon"], subset["tile_lat"])) + + return written + + +def _record_written_tiles( + store: StoreLocation, + tile_infos: List[TileInfo], + year: int, + zone: int, +) -> None: + """Append newly written tiles to this zone/year's tracking file. + + Single-writer by construction: only the process filling (zone, year) + touches this object, so parallel zone sweeps need no locking here. + """ + import geopandas as gpd + import pandas as pd + from shapely.geometry import Point + + now = pd.Timestamp.now(tz="UTC") + rows = [ + { + "year": np.int32(year), + "zone": np.int32(zone), + "tile_lon": ti.lon, + "tile_lat": ti.lat, + "written_at": now, + "geometry": Point(ti.lon, ti.lat), + } + for ti in tile_infos + ] + if not rows: + return + + new_gdf = gpd.GeoDataFrame(rows, crs="EPSG:4326") + name = _zone_registry_name(zone, year) + existing = _read_parquet_at(store.state, REGISTRY_DIR_NAME, name) + + if existing is not None and not existing.empty: + combined = gpd.GeoDataFrame( + pd.concat([existing, new_gdf], ignore_index=True), crs="EPSG:4326" + ).drop_duplicates(subset=["year", "zone", "tile_lon", "tile_lat"], keep="last") + else: + combined = new_gdf + + _write_parquet_at(store.state, combined, REGISTRY_DIR_NAME, name) + + +def merge_tile_registry( + store: StoreLocation, + console: Optional["rich.console.Console"] = None, +) -> int: + """Merge every per-zone tracking file into one ``_registry.parquet``. + + Both the parts and the result live in the state sibling, never inside the + Zarr hierarchy. Run this once after a parallel sweep, when no fill is in + flight — it is the only step that rewrites a store-wide object. Returns + the total row count in the merged registry. + """ + import geopandas as gpd + import pandas as pd + from . import remote + + state = store.state + frames = [] + previous = load_merged_registry(store) + if previous is not None and not previous.empty: + frames.append(previous) + + n_parts = 0 + for entry in state.listdir(REGISTRY_DIR_NAME): + if not entry.endswith(".parquet"): + continue + try: + data = remote.read_bytes(entry, state.storage_options) + part = gpd.read_parquet(io.BytesIO(data)) + except Exception as e: + logger.warning(f"Skipping unreadable registry part {entry}: {e}") + continue + n_parts += 1 + if not part.empty: + frames.append(part) + + if not frames: + merged = _empty_tile_registry() + else: + merged = gpd.GeoDataFrame( + pd.concat(frames, ignore_index=True), crs="EPSG:4326" + ).drop_duplicates(subset=["year", "zone", "tile_lon", "tile_lat"], keep="last") + + _write_parquet_at(state, merged, MERGED_REGISTRY_NAME) + + if console: + console.print( + f" Merged {n_parts} per-zone registry file(s) into " + f"{state.join(MERGED_REGISTRY_NAME)}: {len(merged):,} tiles" + ) + return len(merged) + + +# --------------------------------------------------------------------------- +# Advisory zone locks +# --------------------------------------------------------------------------- +# Two processes filling the same (zone, year) would each rewrite whole shards +# from their own tile subset and silently erase each other's pixels. Object +# stores give us no atomic create, so this is advisory only — it catches the +# common accident (the same zone launched twice) rather than enforcing +# mutual exclusion. + +LOCK_DIR_NAME = "_locks" + + +def _lock_name(zone: int, year: int) -> str: + return f"{_zone_group_name(zone)}_{year}.json" + + +def _acquire_zone_lock( + store: StoreLocation, zone: int, year: int, force: bool = False +) -> None: + """Claim (zone, year) for this process, or raise if someone else holds it.""" + import json + import socket + import pandas as pd + + state = store.state + name = _lock_name(zone, year) + if not force and state.exists(LOCK_DIR_NAME, name): + try: + held = json.loads(state.read_bytes(LOCK_DIR_NAME, name)) + except Exception: + held = {} + raise RuntimeError( + f"Zone {zone} year {year} is locked by " + f"{held.get('host', '?')}:{held.get('pid', '?')} " + f"since {held.get('acquired_at', 'unknown time')}. " + f"Another fill is in progress, or a previous one died. " + f"Re-run with --force-lock to take it over." + ) + + payload = { + "zone": zone, + "year": year, + "host": socket.gethostname(), + "pid": os.getpid(), + "acquired_at": pd.Timestamp.now(tz="UTC").isoformat(), + } + state.write_bytes(json.dumps(payload).encode(), LOCK_DIR_NAME, name) + + +def _release_zone_lock(store: StoreLocation, zone: int, year: int) -> None: + """Drop this process's claim on (zone, year).""" + store.state.remove(LOCK_DIR_NAME, _lock_name(zone, year)) # --------------------------------------------------------------------------- @@ -1111,24 +1605,41 @@ def _get_written_tiles(store_path: Path, year: int, zone: int) -> set: # --------------------------------------------------------------------------- _worker_store = None +_worker_source_options: Optional[Dict[str, Any]] = None -def _init_shard_worker(store_path: str, zone_group: str) -> None: - """Process pool initializer: open the zone group once per worker.""" - global _worker_store - import zarr +def _init_shard_worker( + store_url: str, + zone_group: str, + store_options: Optional[Dict[str, Any]] = None, + source_options: Optional[Dict[str, Any]] = None, +) -> None: + """Process pool initializer: open the zone group once per worker. - _worker_store = zarr.open_group( - store_path, - mode="r+", - path=zone_group, - zarr_format=3, - use_consolidated=False, + Both option dicts are plain picklable mappings, so a worker rebuilds its + own filesystem connections rather than inheriting an unforkable client. + """ + global _worker_store, _worker_source_options + + _worker_store = StoreLocation(store_url, store_options).open_group( + mode="r+", path=zone_group, zarr_format=3 ) + _worker_source_options = source_options + + +def _write_one_shard( + spec: ShardSpec, + store: "zarr.Group", + source_options: Optional[Dict[str, Any]] = None, +) -> bool: + """Write one shard in NCHW layout: (T, B, H, W). + Tile reads go through :mod:`geotessera.remote`, so ``spec`` may reference + either local paths or remote URLs — a remote tile costs one ranged GET for + the rows this shard needs, not the whole 150 MB object. + """ + from . import remote -def _write_one_shard(spec: ShardSpec, store: "zarr.Group") -> bool: - """Write one shard in NCHW layout: (T, B, H, W).""" t = spec.time_index S = SHARD_SIZE @@ -1141,12 +1652,14 @@ def _write_one_shard(spec: ShardSpec, store: "zarr.Group") -> bool: has_data = False for ov in spec.tiles: # Read HWB tile, transpose to BHW - emb = np.load(ov.embedding_path, mmap_mode="r") - tile_slice = emb[ - ov.t_row_start : ov.t_row_end, - ov.t_col_start : ov.t_col_end, - :, - ] + tile_slice = remote.read_npy_window( + ov.embedding_path, + ov.t_row_start, + ov.t_row_end, + ov.t_col_start, + ov.t_col_end, + storage_options=source_options, + ) emb_buf[ :, ov.s_row_start : ov.s_row_end, @@ -1154,11 +1667,17 @@ def _write_one_shard(spec: ShardSpec, store: "zarr.Group") -> bool: ] = tile_slice.transpose(2, 0, 1) # Scales - scales_mmap = np.load(ov.scales_path, mmap_mode="r") - s = scales_mmap[ - ov.t_row_start : ov.t_row_end, - ov.t_col_start : ov.t_col_end, - ].copy() + s = np.array( + remote.read_npy_window( + ov.scales_path, + ov.t_row_start, + ov.t_row_end, + ov.t_col_start, + ov.t_col_end, + storage_options=source_options, + ), + dtype=np.float32, + ) # Landmask lm = _load_landmask_slice( @@ -1167,6 +1686,7 @@ def _write_one_shard(spec: ShardSpec, store: "zarr.Group") -> bool: ov.t_row_end, ov.t_col_start, ov.t_col_end, + storage_options=source_options, ) s[lm == 0] = np.float32("nan") s[~np.isfinite(s)] = np.float32("nan") @@ -1185,7 +1705,7 @@ def _write_one_shard(spec: ShardSpec, store: "zarr.Group") -> bool: def _write_one_shard_worker(spec: ShardSpec) -> bool: """Picklable wrapper for process pool.""" - return _write_one_shard(spec, _worker_store) + return _write_one_shard(spec, _worker_store, _worker_source_options) # --------------------------------------------------------------------------- @@ -1193,49 +1713,224 @@ def _write_one_shard_worker(spec: ShardSpec) -> bool: # --------------------------------------------------------------------------- +def _store_years(store: StoreLocation, zones: Optional[List[int]] = None) -> List[int]: + """Read the store's year axis from a zone's ``time`` coordinate array. + + Tries the requested zones first so a single-zone fill against a remote + store needs no listing of the whole hierarchy. + """ + root = store.open_group(mode="r") + + def years_from(name: str) -> Optional[List[int]]: + try: + return [int(v) for v in root[name]["time"][:]] + except Exception: + return None + + tried = set() + for zone in zones or []: + name = _zone_group_name(zone) + tried.add(name) + years = years_from(name) + if years: + return years + + # Fall back to a listing only if the requested zones told us nothing — + # enumerating the hierarchy is a round trip we can usually skip. + for name in _member_names(root): + if not name.startswith("utm") or name in tried: + continue + years = years_from(name) + if years: + return years + return [] + + +def _member_names(group: "zarr.Group") -> List[str]: + """Sorted member names of a group, without the sidecar warnings. + + A Tessera store deliberately keeps non-Zarr objects at its root (the + ingestion registry, the lock directory), and zarr warns about each one + every time the hierarchy is enumerated. They are expected here. + """ + import warnings + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Object at .* is not recognized") + return sorted(group.keys()) + + +def _zone_group_names( + store: StoreLocation, zones: Optional[List[int]] = None +) -> List[str]: + """Names of the store's UTM zone groups, optionally filtered.""" + import re + + root = store.open_group(mode="r") + pattern = re.compile(r"^utm(\d{2})$") + names = [] + for name in _member_names(root): + m = pattern.match(name) + if m and (zones is None or int(m.group(1)) in zones): + names.append(name) + return names + + +def extend_store( + store_path: "str | Path | StoreLocation", + years: List[int], + console: Optional["rich.console.Console"] = None, + storage_options: Optional[Dict[str, Any]] = None, + zones: Optional[List[int]] = None, + consolidate: bool = True, + force: bool = False, + state_url: Optional[str] = None, +) -> int: + """Append new years to an existing store's time axis. + + The time dimension is chunked one year per chunk, so growing it is a + metadata-only edit: existing chunks keep their keys and are never + rewritten, and the new slice reads back with the same sentinels a + freshly initialised year has (embeddings at 0, scales at +inf). Fill it + afterwards with ``zarr-fill --year ``. + + Years must extend the axis at the end. Inserting an earlier year would + shift every existing chunk's time index — a full rewrite of the store — + so it is refused rather than done silently. + + This is a single-writer operation: it rewrites array metadata for every + zone, so no fill may be in flight. Returns the number of zone groups + extended. + """ + store = StoreLocation.resolve(store_path, storage_options, state_url) + years = sorted(set(int(y) for y in years)) + if not years: + raise ValueError("No years given to add") + + held = [Path(p).name for p in store.state.listdir(LOCK_DIR_NAME)] + if held and not force: + raise RuntimeError( + f"{len(held)} fill lock(s) present ({', '.join(sorted(held)[:4])}" + f"{'...' if len(held) > 4 else ''}). Extending rewrites array " + f"metadata for every zone, so wait for the sweep to finish. " + f"Use --force if these are stale." + ) + + zone_names = _zone_group_names(store, zones) + if not zone_names: + raise ValueError(f"No UTM zone groups found in {store}") + + if console: + console.print(f"Extending [bold]{store}[/bold] with years {years}") + console.print(f" {len(zone_names)} zone group(s)") + + extended = 0 + skipped = 0 + for name in zone_names: + group = store.open_group(mode="r+", path=name, zarr_format=3) + existing = [int(v) for v in group["time"][:]] + + missing = [y for y in years if y not in existing] + if not missing: + skipped += 1 + continue + + earliest_new = min(missing) + if existing and earliest_new <= max(existing): + raise ValueError( + f"{name}: cannot add {earliest_new} to a time axis ending at " + f"{max(existing)}. Years may only be appended — inserting an " + f"earlier one would renumber every existing chunk." + ) + + old_t = len(existing) + new_t = old_t + len(missing) + + # Order matters only for crash-safety: grow the data arrays before + # advertising the year on the time axis, so a run interrupted midway + # never leaves a year readers can select but not read. + emb = group["embeddings"] + scales = group["scales"] + emb.resize((new_t,) + tuple(emb.shape[1:])) + scales.resize((new_t,) + tuple(scales.shape[1:])) + + time_arr = group["time"] + time_arr.resize((new_t,)) + time_arr[old_t:new_t] = np.array(missing, dtype=time_arr.dtype) + + extended += 1 + if console: + console.print( + f" {name}: {old_t} -> {new_t} time steps " + f"[dim](added {', '.join(str(y) for y in missing)})[/dim]" + ) + + if console and skipped: + console.print(f" {skipped} zone(s) already had every year") + + # Array metadata changed, so the consolidated root is now stale — unlike + # a fill, this step genuinely requires re-consolidation. + if consolidate and extended: + consolidate_store(store, console=console) + + return extended + + def fill_store( registry: "Registry", - store_path: Path, + store_path: "str | Path | StoreLocation", year: Optional[int] = None, zones: Optional[List[int]] = None, console: Optional["rich.console.Console"] = None, workers: Optional[int] = None, + storage_options: Optional[Dict[str, Any]] = None, + source: Optional[TileSource] = None, + consolidate: Optional[bool] = None, + force_lock: bool = False, + state_url: Optional[str] = None, ) -> int: """Incrementally fill a store with tile data. Reads the tile registry to skip already-written tiles. Returns the number of shards written. - """ - import warnings - import zarr - from concurrent.futures import ProcessPoolExecutor, as_completed - store_path = Path(store_path) + Args: + store_path: Local path or fsspec URL of an initialised store. + zones: Restrict the fill to these UTM zones. One process per zone + can run concurrently against the same store. + storage_options: fsspec options for the store (endpoint, credentials). + source: Where the tile inputs live; defaults to the registry's local + mirror. + consolidate: Rewrite the root consolidated metadata when done. + Defaults to True for a whole-store fill and False when ``zones`` + is set, because the root object is the one thing parallel zone + jobs share — run ``zarr-consolidate`` once after the sweep. + force_lock: Take over a (zone, year) lock held by another process. + """ + store = StoreLocation.resolve(store_path, storage_options, state_url) if workers is None: workers = DEFAULT_WORKERS + if consolidate is None: + consolidate = zones is None - root = zarr.open_group(str(store_path), mode="r", use_consolidated=False) - # Derive years from the first zone's time coordinate array - all_years: list[int] = [] - for member_name in sorted(root.keys()): - if member_name.startswith("utm"): - try: - time_arr = root[member_name]["time"][:] - all_years = [int(v) for v in time_arr] - break - except Exception: - continue - + all_years = _store_years(store, zones) if not all_years: - raise ValueError("Store has no years (checked root attrs and zone time coords)") + raise ValueError("Store has no years (checked zone time coords)") fill_years = [year] if year is not None else all_years if console: - console.print(f"Filling store at [bold]{store_path}[/bold]") + console.print(f"Filling store at [bold]{store}[/bold]") console.print(f" Years to fill: {fill_years}") + if source is not None and source.is_remote: + console.print(f" Streaming tiles from [bold]{source.embeddings_root}[/bold]") total_shards_written = 0 + total_shards_failed = 0 + + # The merged registry spans the whole store, so fetch it once rather + # than per zone and year. + merged_registry = load_merged_registry(store) for fill_year in fill_years: if fill_year not in all_years: @@ -1245,21 +1940,19 @@ def fill_store( ) continue - time_index = all_years.index(fill_year) - # Gather tiles for this year year_tiles = gather_tile_infos( registry, fill_year, zones=zones, console=console, + source=source, ) for zone_num, tile_infos in sorted(year_tiles.items()): zone_group = _zone_group_name(zone_num) - zone_path = store_path / zone_group - if not zone_path.exists(): + if not store.exists(zone_group): if console: console.print( f" [yellow]Zone {zone_num} not initialised, skipping[/yellow]" @@ -1267,7 +1960,9 @@ def fill_store( continue # Check which tiles are already written - written = _get_written_tiles(store_path, fill_year, zone_num) + written = _get_written_tiles( + store, fill_year, zone_num, merged=merged_registry + ) remaining = [ti for ti in tile_infos if (ti.lon, ti.lat) not in written] if not remaining: @@ -1285,12 +1980,22 @@ def fill_store( ) # Read the zone grid from store metadata - zone_store = zarr.open_group( - str(store_path), - mode="r", - path=zone_group, - use_consolidated=False, - ) + zone_store = store.open_group(mode="r", path=zone_group) + + # Resolve the time index against *this* zone's own axis. An + # interrupted zarr-extend can leave zones with different lengths, + # and a store-wide index would then address the wrong year. + zone_years = [int(v) for v in zone_store["time"][:]] + if fill_year not in zone_years: + if console: + console.print( + f" [yellow]Zone {zone_num} has no {fill_year} on its " + f"time axis ({zone_years}); run zarr-extend first. " + f"Skipping.[/yellow]" + ) + continue + time_index = zone_years.index(fill_year) + zone_attrs = dict(zone_store.attrs) transform = zone_attrs["spatial:transform"] shape = zone_attrs["spatial:shape"] @@ -1305,175 +2010,213 @@ def fill_store( height_px=shape[0], ) - # Build shard index - shard_specs = build_shard_index(remaining, grid, time_index) + # A shard write replaces the whole shard, so every shard we touch + # must be rebuilt from all of its tiles — including ones an + # earlier run already wrote, which would otherwise be zeroed. + touched = shard_coords_for_tiles(remaining, grid) + shard_specs = build_shard_index( + tile_infos, grid, time_index, restrict_to=touched + ) if console: + n_rewritten = sum(len(s.tiles) for s in shard_specs) - len(remaining) + extra = ( + f", {n_rewritten} previously-written tile(s) rewritten" + if n_rewritten > 0 + else "" + ) console.print( - f" {len(shard_specs)} shards to write ({workers} workers)" + f" {len(shard_specs)} shards to write " + f"({workers} workers{extra})" ) - # Write shards via process pool - zone_store_path = str(store_path) - written_count = 0 - n_shards = len(shard_specs) - - if console: - from rich.progress import ( - Progress, - BarColumn, - TextColumn, - MofNCompleteColumn, - TimeElapsedColumn, - TimeRemainingColumn, - SpinnerColumn, + _acquire_zone_lock(store, zone_num, fill_year, force=force_lock) + try: + written_count, failed = _write_shards( + store=store, + zone_group=zone_group, + shard_specs=shard_specs, + workers=workers, + source_options=source.storage_options if source else None, + label=f" Zone {zone_num} y{fill_year}", + console=console, ) - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - MofNCompleteColumn(), - TimeElapsedColumn(), - TimeRemainingColumn(), - console=console, - ) as progress: - task = progress.add_task( - f" Zone {zone_num} y{fill_year}", - total=n_shards, + total_shards_written += written_count + total_shards_failed += len(failed) + + if console: + console.print( + f" [green]{written_count}/{len(shard_specs)} " + f"shards written[/green]" ) - with ProcessPoolExecutor( - max_workers=workers, - initializer=_init_shard_worker, - initargs=(zone_store_path, zone_group), - ) as pool: - futures = { - pool.submit(_write_one_shard_worker, spec): spec - for spec in shard_specs - } - for future in as_completed(futures): - try: - if future.result(): - written_count += 1 - except Exception as e: - spec = futures[future] - logger.warning( - f"Shard ({spec.sr},{spec.sc}) failed: {e}" - ) - progress.advance(task) - else: - with ProcessPoolExecutor( - max_workers=workers, - initializer=_init_shard_worker, - initargs=(zone_store_path, zone_group), - ) as pool: - futures = { - pool.submit(_write_one_shard_worker, spec): spec - for spec in shard_specs - } - for future in as_completed(futures): - try: - if future.result(): - written_count += 1 - except Exception as e: - spec = futures[future] - logger.warning(f"Shard ({spec.sr},{spec.sc}) failed: {e}") + if failed: + console.print( + f" [red]{len(failed)} shard(s) failed[/red]" + ) + + # Only record tiles whose shards all landed, so a retry picks + # up exactly the work a transient failure left behind. + if failed: + recorded = [ + ti + for ti in remaining + if not (shard_coords_for_tiles([ti], grid) & failed) + ] + else: + recorded = remaining + _record_written_tiles(store, recorded, fill_year, zone_num) + finally: + _release_zone_lock(store, zone_num, fill_year) + + # A failed shard leaves its tiles unrecorded, so re-running finishes the + # job. Surface it as an error rather than a quiet partial success — a + # sweep orchestrator has no other way to tell the zone needs a retry. + if total_shards_failed: + raise RuntimeError( + f"{total_shards_failed} shard(s) failed " + f"({total_shards_written} written). Re-run the same command to " + f"retry only the unfinished tiles." + ) - total_shards_written += written_count + # Re-consolidate metadata after filling. Skipped for a zone-restricted + # fill: the root object is shared with any sibling zone jobs. + if total_shards_written > 0: + if consolidate: + consolidate_store(store, console=console) + elif console: + console.print( + " [dim]Skipped consolidation (zone-restricted fill). " + "Run `geotessera-registry zarr-consolidate` once the sweep " + "finishes.[/dim]" + ) - if console: - console.print( - f" [green]{written_count}/{n_shards} shards written[/green]" - ) + return total_shards_written - # Update tile registry - _record_written_tiles(store_path, remaining, fill_year, zone_num) - # Re-consolidate metadata after filling - if total_shards_written > 0: - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", message="Consolidated metadata") - zarr.consolidate_metadata(str(store_path)) +def _write_shards( + store: StoreLocation, + zone_group: str, + shard_specs: List[ShardSpec], + workers: int, + source_options: Optional[Dict[str, Any]], + label: str, + console: Optional["rich.console.Console"], +) -> Tuple[int, set]: + """Run the shard writes through a process pool. - return total_shards_written + Returns (shards written, set of (sr, sc) that failed). + """ + from concurrent.futures import ProcessPoolExecutor, as_completed + + written_count = 0 + failed: set = set() + initargs = (store.url, zone_group, store.storage_options, source_options) + + def _drain(pool, advance=None): + nonlocal written_count + futures = { + pool.submit(_write_one_shard_worker, spec): spec for spec in shard_specs + } + for future in as_completed(futures): + spec = futures[future] + try: + if future.result(): + written_count += 1 + except Exception as e: + logger.warning(f"Shard ({spec.sr},{spec.sc}) failed: {e}") + failed.add((spec.sr, spec.sc)) + if advance is not None: + advance() + + if console: + from rich.progress import ( + Progress, + BarColumn, + TextColumn, + MofNCompleteColumn, + TimeElapsedColumn, + TimeRemainingColumn, + SpinnerColumn, + ) + + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TimeElapsedColumn(), + TimeRemainingColumn(), + console=console, + ) as progress: + task = progress.add_task(label, total=len(shard_specs)) + with ProcessPoolExecutor( + max_workers=workers, + initializer=_init_shard_worker, + initargs=initargs, + ) as pool: + _drain(pool, advance=lambda: progress.advance(task)) + else: + with ProcessPoolExecutor( + max_workers=workers, + initializer=_init_shard_worker, + initargs=initargs, + ) as pool: + _drain(pool) + + return written_count, failed def consolidate_store( - store_path: str | Path, + store_path: "str | Path | StoreLocation", console: Optional["rich.console.Console"] = None, + storage_options: Optional[Dict[str, Any]] = None, + merge_registry: bool = True, + state_url: Optional[str] = None, ) -> int: """Re-consolidate a store's root metadata after in-place changes. - ``fill_store`` only re-consolidates when it writes at least one shard, - so a metadata-only change to an existing store (e.g. rewriting an - array with a different compressor) leaves the consolidated metadata - in the root ``zarr.json`` stale. HTTP readers cannot list a store and - trust consolidated metadata exclusively, so a stale root breaks them. + ``fill_store`` skips consolidation for zone-restricted fills so parallel + zone jobs never contend for the root ``zarr.json``, and a metadata-only + change to an existing store (e.g. rewriting an array with a different + compressor) leaves the consolidated metadata stale too. HTTP readers + cannot list a store and trust consolidated metadata exclusively, so a + stale root breaks them. This is the single-writer step that fixes both. Accepts a local store path or a remote fsspec URL such as ``s3://bucket/store.zarr``. Remote URLs need the matching fsspec backend installed (``s3fs`` for S3) and write credentials for the final root ``zarr.json`` upload. + Args: + merge_registry: Also fold the per-zone ingestion files under + ``_registry/`` into the root ``_registry.parquet``. + Returns the number of consolidated nodes. """ import warnings import zarr - store = str(store_path) - if "://" not in store and not Path(store).exists(): - raise FileNotFoundError(f"store not found: {store}") + store = StoreLocation.resolve(store_path, storage_options, state_url) + if not store.is_remote and not Path(store.url).exists(): + raise FileNotFoundError(f"store not found: {store.url}") if console: console.print(f"Consolidating metadata at [bold]{store}[/bold]") + if merge_registry: + merge_tile_registry(store, console=console) + with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Consolidated metadata") # Tessera stores carry non-zarr marker objects (tile registry, # zone-completion flags) that the consolidation walk would warn about. - warnings.filterwarnings( - "ignore", message="Object at .* is not recognized" - ) - group = zarr.consolidate_metadata(store) + warnings.filterwarnings("ignore", message="Object at .* is not recognized") + group = zarr.consolidate_metadata(store.as_zarr_store()) return len(group.metadata.consolidated_metadata.flattened_metadata) -def _record_written_tiles( - store_path: Path, - tile_infos: List[TileInfo], - year: int, - zone: int, -) -> None: - """Append newly written tiles to the registry.""" - import geopandas as gpd - import pandas as pd - from shapely.geometry import Point - - now = pd.Timestamp.now(tz="UTC") - rows = [] - for ti in tile_infos: - rows.append( - { - "year": np.int32(year), - "zone": np.int32(zone), - "tile_lon": ti.lon, - "tile_lat": ti.lat, - "written_at": now, - "geometry": Point(ti.lon, ti.lat), - } - ) - - new_gdf = gpd.GeoDataFrame(rows, crs="EPSG:4326") - existing = _load_tile_registry(store_path) - - combined = gpd.GeoDataFrame( - pd.concat([existing, new_gdf], ignore_index=True), - crs="EPSG:4326", - ) - _save_tile_registry(store_path, combined) - - # --------------------------------------------------------------------------- # RGB preview generation (NCHW layout) # --------------------------------------------------------------------------- @@ -1822,7 +2565,7 @@ def compute_global_stretch( # Find time_index for the requested year via the first zone's time coord. time_index = None - for member_name in sorted(root.keys()): + for member_name in _member_names(root): if member_name.startswith("utm"): try: time_arr = root[member_name]["time"][:] @@ -1839,7 +2582,7 @@ def compute_global_stretch( zone_pattern = re.compile(r"^utm(\d{2})$") all_shards: List[Tuple[str, int, int]] = [] zones_visited = set() - for name in sorted(root.keys()): + for name in _member_names(root): m = zone_pattern.match(name) if not m: continue @@ -2481,8 +3224,9 @@ def _reproject_zone( chunk_row_start = row_start // GLOBAL_CHUNK chunk_col_start = col_start // GLOBAL_CHUNK - # Resume check - marker = store_path / f".zone_{zone_num}_done" + # Resume check. The marker lives in the state sibling, not the store, so + # the published hierarchy stays free of non-Zarr objects. + marker = _preview_marker_path(store_path, zone_num) if marker.exists(): if force: marker.unlink() @@ -2572,6 +3316,7 @@ def _reproject_zone( except Exception as e: logger.warning(f"Reproject chunk failed: {e}") + marker.parent.mkdir(parents=True, exist_ok=True) marker.write_text( f"zone={zone_num} chunks={chunks_total} written={chunks_written}\n" ) @@ -2617,7 +3362,7 @@ def build_global_preview( # Derive years from first zone's time coordinate all_years: list[int] = [] - for member_name in sorted(root.keys()): + for member_name in _member_names(root): if member_name.startswith("utm"): try: time_arr = root[member_name]["time"][:] @@ -2647,7 +3392,7 @@ def build_global_preview( zone_pattern = re.compile(r"^utm(\d{2})$") zone_infos: Dict[int, dict] = {} - for name in sorted(root.keys()): + for name in _member_names(root): m = zone_pattern.match(name) if not m: continue diff --git a/pyproject.toml b/pyproject.toml index f6d42ec..df62e7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,12 @@ dependencies = [ "contextily", ] +[project.optional-dependencies] +# s3:// locations for zarr-init/zarr-fill/zarr-consolidate. Kept optional +# because s3fs pulls in botocore (~26 MB), which the core install +# deliberately does not depend on — https:// reads work without it. +s3 = ["s3fs"] + [project.urls] Homepage = "https://github.com/ucam-eo/geotessera" Documentation = "https://geotessera.readthedocs.io" diff --git a/uv.lock b/uv.lock index 2aa1c23..a2c7947 100644 --- a/uv.lock +++ b/uv.lock @@ -19,6 +19,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0b/f7/85273299ab57117850cc0a936c64151171fac4da49bc6fba0dad984a7c5f/affine-2.4.0-py3-none-any.whl", hash = "sha256:8a3df80e2b2378aef598a83c1392efd47967afec4242021a0b06b4c7cbc61a92", size = 15662, upload-time = "2023-01-19T23:44:28.833Z" }, ] +[[package]] +name = "aiobotocore" +version = "3.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "aioitertools" }, + { name = "botocore" }, + { name = "jmespath" }, + { name = "multidict" }, + { name = "python-dateutil" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d8/a7/bc31b7046c610471f0630819ca5d2a57ac4efa8d47135cb53e43f2785390/aiobotocore-3.8.0.tar.gz", hash = "sha256:80a1eb64ea915f3af3c1518669975bae74a17b2f37c14eb0fa2f83b915974670", size = 131368, upload-time = "2026-07-17T03:10:30.258Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/f4/5a7d76dc844d3ff8ed1f1a043158aa393794aebb787d3e2f8c0fe87f674f/aiobotocore-3.8.0-py3-none-any.whl", hash = "sha256:8bc605132cadfe844a3f334635a0a64fa5e360a4a206e915d99d53db5b6deeba", size = 91169, upload-time = "2026-07-17T03:10:28.771Z" }, +] + [[package]] name = "aiohappyeyeballs" version = "2.6.2" @@ -128,6 +146,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/71/6e22be134a4061ada85a92951b842f2657f17d926b727f3f94c56ae963d6/aiohttp-3.14.1-cp314-cp314t-win_arm64.whl", hash = "sha256:90d53f1609c29ccc2193945ef732428382a28f78d0456ae4d3daf0d48b74f0f6", size = 469640, upload-time = "2026-06-07T21:09:33.028Z" }, ] +[[package]] +name = "aioitertools" +version = "0.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/3c/53c4a17a05fb9ea2313ee1777ff53f5e001aefd5cc85aa2f4c2d982e1e38/aioitertools-0.13.0.tar.gz", hash = "sha256:620bd241acc0bbb9ec819f1ab215866871b4bbd1f73836a55f799200ee86950c", size = 19322, upload-time = "2025-11-06T22:17:07.609Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl", hash = "sha256:0be0292b856f08dfac90e31f4739432f4cb6d7520ab9eb73e143f4f2fa5259be", size = 24182, upload-time = "2025-11-06T22:17:06.502Z" }, +] + [[package]] name = "aiosignal" version = "1.4.0" @@ -186,6 +213,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/f5/21d2de20e8b8b0408f0681956ca2c69f1320a3848ac50e6e7f39c6159675/babel-2.18.0-py3-none-any.whl", hash = "sha256:e2b422b277c2b9a9630c1d7903c2a00d0830c409c59ac8cae9081c92f1aeba35", size = 10196845, upload-time = "2026-02-01T12:30:53.445Z" }, ] +[[package]] +name = "botocore" +version = "1.43.46" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7d/f1/1917891851ac5ac09bb9f4862b8fc9252a009d7c24e8688bb67e4383d9e7/botocore-1.43.46.tar.gz", hash = "sha256:59f2e1ac3cdc66d191cae91c0804bc41847ce817dc8147cf43eaada8f76a5533", size = 15694635, upload-time = "2026-07-10T19:32:00.437Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/f2/4bd8f2f419088feb3ce55f0ca91040ff902f402edfd197450b20a2e1d533/botocore-1.43.46-py3-none-any.whl", hash = "sha256:cb673891e623ae6e6a1bf24d94ef169504f3eb02584adb5d5bee2f6aae819b60", size = 15380350, upload-time = "2026-07-10T19:31:57.616Z" }, +] + [[package]] name = "certifi" version = "2026.5.20" @@ -669,6 +710,11 @@ dependencies = [ { name = "zarr" }, ] +[package.optional-dependencies] +s3 = [ + { name = "s3fs" }, +] + [package.dev-dependencies] dev = [ { name = "pytest" }, @@ -692,6 +738,7 @@ requires-dist = [ { name = "rasterio" }, { name = "rich" }, { name = "rioxarray" }, + { name = "s3fs", marker = "extra == 's3'" }, { name = "scikit-image", specifier = ">=0.25.2" }, { name = "scikit-learn", specifier = ">=1.7.1" }, { name = "sphinx", specifier = ">=8.2.3" }, @@ -699,6 +746,7 @@ requires-dist = [ { name = "xarray" }, { name = "zarr" }, ] +provides-extras = ["s3"] [package.metadata.requires-dev] dev = [ @@ -796,6 +844,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + [[package]] name = "joblib" version = "1.5.3" @@ -1985,6 +2042,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/19/016553f86f207450aebebc2b2b5088d086b901cc8186c02ac4284db3bd88/ruff-0.15.16-py3-none-win_arm64.whl", hash = "sha256:8cd61783afb39638a7133ef0d2dfb1e91277593962f81b5a8423eb0b888a6121", size = 11134555, upload-time = "2026-06-04T16:33:00.136Z" }, ] +[[package]] +name = "s3fs" +version = "2026.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiobotocore" }, + { name = "aiohttp" }, + { name = "fsspec" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cb/d8/76f3dc1558bdf4494b117a9f7a9cc0a5d9d34edadc9e5d7ceabc5a6a7c37/s3fs-2026.4.0.tar.gz", hash = "sha256:5bdce0abb00b0435ee150807a45fea727451dbc22de4cbc116464f8504ab9d37", size = 85986, upload-time = "2026-04-29T20:52:51.748Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/a4/9d1ea10ebc9e028a289a72fec84da170689549a8102c8aacfcad26bc5035/s3fs-2026.4.0-py3-none-any.whl", hash = "sha256:de0d2a1f33cdf03831fd2382d278c6e4e31fe57c3bf2f703c61f8aec6b703e2a", size = 32392, upload-time = "2026-04-29T20:52:50.295Z" }, +] + [[package]] name = "scikit-image" version = "0.26.0" @@ -2396,6 +2467,70 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] +[[package]] +name = "wrapt" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/b0/c1f5a970721f06b85c0cd5142e0ff8fe067708abd779b0c4f4be7d61d09f/wrapt-2.3.0.tar.gz", hash = "sha256:681a2d0eefd721998f90642762b8e75c2159ec531b20ad5e437245ea7b06a107", size = 131509, upload-time = "2026-07-28T06:06:14.895Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/4a/d17a0fad1bf1c5f2c887ff71fef75654141b0880bff71d157d955b5bec3a/wrapt-2.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:0a45ffae742ce91a16e11cb6c7cd71e7f9994f3cbd283b962ab093f5c6dcf525", size = 82139, upload-time = "2026-07-28T06:04:35.082Z" }, + { url = "https://files.pythonhosted.org/packages/6e/55/51b92daaf6defb57f4dc56bdcce985400f75c6984a03ca5e78ccac717028/wrapt-2.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:69e477046f2237ef0bc6547544ee73008dc764ca26eff44f09e976d221b34d5d", size = 82723, upload-time = "2026-07-28T06:04:36.502Z" }, + { url = "https://files.pythonhosted.org/packages/28/7f/cfd9bc4b1f5e424eeea83d0493e43f3b1b02707ce8e50c47945873982bd5/wrapt-2.3.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:5d221a6e6ddd302b8397433184e96b59f259f50024b854db1c411a881586b6b8", size = 172381, upload-time = "2026-07-28T06:04:37.674Z" }, + { url = "https://files.pythonhosted.org/packages/cb/89/ff7814f6eb6856b479946117d1138a2fbb46cdb6b1f379db359056c69743/wrapt-2.3.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:392158c9a7f2ab1b8699418bfc0fe6f83548788c418b27d7bf2019ad3405cebb", size = 174120, upload-time = "2026-07-28T06:04:38.987Z" }, + { url = "https://files.pythonhosted.org/packages/12/1e/8eded8615d39e3ce81f626937a3a87b280a2a86239a2bf14a4b4bb345034/wrapt-2.3.0-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e5301c35cf75655eb33498f2bd6ae8703ca19940e3167dc9cdf740c712a39c60", size = 163035, upload-time = "2026-07-28T06:04:40.361Z" }, + { url = "https://files.pythonhosted.org/packages/35/ea/a0af2d9da62897af2a055484920de05dade30d2ba2c0d65cbdea875d3d8b/wrapt-2.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:418f54bb09d1762db02c7009b4051149893af3153a87f92d70356703c11eea02", size = 171887, upload-time = "2026-07-28T06:04:41.614Z" }, + { url = "https://files.pythonhosted.org/packages/7e/dd/63cd4c864c65ef4906df64bd2d378f4a62b54f28063f282dfb3bf93caead/wrapt-2.3.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:1598becd30f8f2777d18564064eb4f4dbe1ab0e05a8f09786d0ef505ac782bf3", size = 161113, upload-time = "2026-07-28T06:04:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ee/82f1fc9e431b5c2c5a6d201aa865dbeae3984c311c6d11a185f0c8367cf6/wrapt-2.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3da470536bf9645143323dd41b32db55c6f4304ad382094c1a1da8a92061e10d", size = 170530, upload-time = "2026-07-28T06:04:44.212Z" }, + { url = "https://files.pythonhosted.org/packages/37/a5/5dc590e863a419930d988f8b7ca3e75a6befcfb10b6003b3a152f3d5f732/wrapt-2.3.0-cp312-cp312-win32.whl", hash = "sha256:fb8e2e6704a1e0b1b989546c69e2688371ef4a07fa5f61bde3eb6211186f5ac1", size = 78323, upload-time = "2026-07-28T06:04:45.484Z" }, + { url = "https://files.pythonhosted.org/packages/51/f9/4a6925a07951df56394f7e6ebe14f69f1c5ef9d87aa63e0839acf15aa63a/wrapt-2.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:cdc021cb0b62471d6aac7f2bd92f3b4658073775f9ee7fcd325c511129e7bcc8", size = 81180, upload-time = "2026-07-28T06:04:47.021Z" }, + { url = "https://files.pythonhosted.org/packages/a8/4f/8b5de0395b2a72216751d41c9861df6facaeb611b619d8810ed2b3b23eb2/wrapt-2.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:67bfe2485f50368c3fcd2275fc1fd100e350d601e0058921a7c82678a465aeab", size = 80155, upload-time = "2026-07-28T06:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/8e/6e/0f88a072483e76b881e3fdcd6b6ffb4a5791002514fe541e72b1b73c859a/wrapt-2.3.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0d3fb71e65b001adfc42684522eeccd9c21d8ba679945abc993439567b66e59f", size = 81960, upload-time = "2026-07-28T06:04:49.622Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ff/b7e2776e7c294075eb712cc9ef573d1b818f393006d09787262b8fc871c4/wrapt-2.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:51a7a4181c1295774812271fbcd7c909df372bc25579d4ed9eb875caaf0ae86f", size = 82435, upload-time = "2026-07-28T06:04:50.9Z" }, + { url = "https://files.pythonhosted.org/packages/d8/90/343bb5d0f1f9669bc252a6073f085b4abf862511bd5c9c9eaec754341f1d/wrapt-2.3.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9045917809c63fdf7abe3a2ceaed3d670b8ee4500ddd9291192d30aeb34467c5", size = 170350, upload-time = "2026-07-28T06:04:52.187Z" }, + { url = "https://files.pythonhosted.org/packages/59/f8/13b79a392930bd0dd6b86cbfbfe1c40944110456e1dc6d809e5c46ece904/wrapt-2.3.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:54ca1d5573f69b5fe1d74f1f65799c68015e82f685efec9fd8cfa40a094c44d0", size = 170022, upload-time = "2026-07-28T06:04:53.599Z" }, + { url = "https://files.pythonhosted.org/packages/b2/fc/4f1b6918f5290db959d6e0c07f77385d87cede29c39c9cf8f145e9c82954/wrapt-2.3.0-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:242b60c21e30866e6a2fa606c612b47c553fa60c0eaeeeb7797fb842ac0ce609", size = 161043, upload-time = "2026-07-28T06:04:54.936Z" }, + { url = "https://files.pythonhosted.org/packages/01/e1/45d3cf74414780bdff6d0380467e003f6eb0f028b6c9403db868dbc7209c/wrapt-2.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3f3d7ec0a51fbfe00d3aef047641ff2c58b25565b4717fc1f90e050be01cba8", size = 168576, upload-time = "2026-07-28T06:04:56.261Z" }, + { url = "https://files.pythonhosted.org/packages/f3/73/2fa58dd97f191c997755e2c6d569a68f0c433db4e4b36099bdd7227b6cac/wrapt-2.3.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:261f53870cd4fb2bf38f9f972c56c728fd224cb7c65721307de59d9e7e6741ae", size = 159140, upload-time = "2026-07-28T06:04:57.754Z" }, + { url = "https://files.pythonhosted.org/packages/29/a8/08a56e2000a8816d449dcbad8c8b081697acbbd490821ceca0f9d8e8d20c/wrapt-2.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8159ec0b0cb7608175eb150de94c19e34f4d47ac655f5ca9baf45df6b688ffd3", size = 169263, upload-time = "2026-07-28T06:04:59.161Z" }, + { url = "https://files.pythonhosted.org/packages/9e/d4/354e1725e35a73b2af4fa70a3e024c7a5d1bf1802dfb862dcb668aae0253/wrapt-2.3.0-cp313-cp313-win32.whl", hash = "sha256:10461884b3014fbfc8eb7d09a93c5f246363e6711d9d881f95eb8c27fdef049f", size = 78241, upload-time = "2026-07-28T06:05:00.507Z" }, + { url = "https://files.pythonhosted.org/packages/6c/7e/34c87fa2174848dfee820322aaa318bab08913998ccecc8d2f57b4ad4639/wrapt-2.3.0-cp313-cp313-win_amd64.whl", hash = "sha256:ac870cc97b73bb00ac353329e9559a4bebc47c4c86792ed9b23b58c15b6ad838", size = 81113, upload-time = "2026-07-28T06:05:01.839Z" }, + { url = "https://files.pythonhosted.org/packages/11/86/fcc9a530579e008c9478bb565a6cdfbfd33536660f069c8b91a6607c5050/wrapt-2.3.0-cp313-cp313-win_arm64.whl", hash = "sha256:a65e8db2b4e90c2e7ade931086351c98ef420bf7a94ee08c95ac8a3cbbc43579", size = 80182, upload-time = "2026-07-28T06:05:03.152Z" }, + { url = "https://files.pythonhosted.org/packages/96/50/3864848b95b28ef73e17551fc8dccbff2628a834f52cf26a57f9c419fb83/wrapt-2.3.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:fd1f2f557dd3491fe75905e578f4db967393d40d1a8f468edc4d40ac7f2d5944", size = 83921, upload-time = "2026-07-28T06:05:04.476Z" }, + { url = "https://files.pythonhosted.org/packages/3b/4c/3d1921a60c3e8c71c540ff136e6a47a1fbccf7f671e818394889f7871d9c/wrapt-2.3.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:9f5d2aec29dfc76c37e23897dee92766a3fd4f3bff3ae7fc9c6b4bf37d8c1360", size = 84412, upload-time = "2026-07-28T06:05:05.921Z" }, + { url = "https://files.pythonhosted.org/packages/fa/1a/4a796ff7adb26ada6d4b758c94d47a38320b085e7099afc088efbbcdb006/wrapt-2.3.0-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:646d20d413ffcd1b0a2f700076e2d0252d872dcb7754860a73e45a59ea883614", size = 207168, upload-time = "2026-07-28T06:05:07.256Z" }, + { url = "https://files.pythonhosted.org/packages/1d/3e/d7777776806c579b761bac2f91721dda9f04c7a1b380213c5935cc750ae6/wrapt-2.3.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:379f670f45b7bb8993edd9f6fc36c6cc65edb81cffa0b504be34acb0303fff0a", size = 214351, upload-time = "2026-07-28T06:05:08.945Z" }, + { url = "https://files.pythonhosted.org/packages/63/27/2d64d394df7bf181955b3bb562bf33c4492fb4be113f53071106d43ad8b5/wrapt-2.3.0-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:6208f302f110295d64b22a7ac96500c791bf492dce4366e622e4912b077c9687", size = 199020, upload-time = "2026-07-28T06:05:10.418Z" }, + { url = "https://files.pythonhosted.org/packages/3e/3d/fb31d3db7d9834d265fb1a27a2adf0ddf51557c67458c97b22439ad6ae3d/wrapt-2.3.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:ed635a9ca4f3a5a2b900c10c69e823373bc00ebc114b459383596d3487da3570", size = 209969, upload-time = "2026-07-28T06:05:11.983Z" }, + { url = "https://files.pythonhosted.org/packages/1f/d1/8724b5da582e62070dc9bf4d8bf1972f317297eefd7ba1f2b5c6393ccf6c/wrapt-2.3.0-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:e3b9eaa742ae7a0aaaaad4ca4b69469d757af2d6e6663ef1dadc47adec0aeb41", size = 196324, upload-time = "2026-07-28T06:05:13.557Z" }, + { url = "https://files.pythonhosted.org/packages/0d/5c/3d9ef411149543016ee6bcf3af707f787cebd946527452b94bf122e9b7b4/wrapt-2.3.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:d0f7284f88f4833705132d06d3b425a43095c2cbd07c58166aac3ab646ba12a4", size = 202610, upload-time = "2026-07-28T06:05:15.048Z" }, + { url = "https://files.pythonhosted.org/packages/13/9b/4fc042ceb757866dd4a5fc057b3b736f2b360d3703ce9f830d83dc9226e0/wrapt-2.3.0-cp313-cp313t-win32.whl", hash = "sha256:7ebb274aba688b043429eb1500ff8a76ce0cb8ac0812ca3e301f06247b8722b3", size = 79178, upload-time = "2026-07-28T06:05:16.469Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ff/b94878f8eed809ca042685276bcea9f24e8c2ca7c9653bb80bbb920a68a5/wrapt-2.3.0-cp313-cp313t-win_amd64.whl", hash = "sha256:c4bded758ad6f03b965830944a2f0bc5b2eb3767fe5a7310134315d1a6610e98", size = 82634, upload-time = "2026-07-28T06:05:18.026Z" }, + { url = "https://files.pythonhosted.org/packages/80/fb/663e1de5332a71685a729754312d327d4cada767c36e1c5a2db4c8de49e6/wrapt-2.3.0-cp313-cp313t-win_arm64.whl", hash = "sha256:d2cc64539da63e39ffb9c7ede849b6e8ddaaf7b3876b5cfb04efd85a5f3f4eb6", size = 81387, upload-time = "2026-07-28T06:05:19.417Z" }, + { url = "https://files.pythonhosted.org/packages/58/10/b073beaea89bc0d3670a75ff51139430a54b6af7ba7796507730634536dd/wrapt-2.3.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ea52a0d0f08c584943d5764be0e84efa912c8da23c23e1e285ff2f5641c18fcc", size = 81978, upload-time = "2026-07-28T06:05:21.133Z" }, + { url = "https://files.pythonhosted.org/packages/b3/31/0916d9cebf848ed3f1a0c1888faee421747df77331e4db2bc527a9a85988/wrapt-2.3.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:fd85b0aa88efdb189d6ae2f35f4526943a8f091c38599c9c31478241c819e6a1", size = 82518, upload-time = "2026-07-28T06:05:22.562Z" }, + { url = "https://files.pythonhosted.org/packages/f5/73/31c1bf0f3384062751c2094dadb314916d70aa9b6bfd26d994b4a7b393fa/wrapt-2.3.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:141ed6211286a9660d8d6702de598b43f0934b4f0eda16393f100a80f501d945", size = 170187, upload-time = "2026-07-28T06:05:23.904Z" }, + { url = "https://files.pythonhosted.org/packages/ed/25/fce087d54b79b8905f3c3c9dd5f454bbd8d8acb80b960c4a6aee5b4659b3/wrapt-2.3.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e49885a62ec4ee854d1b9e6371fda6afd219917225752abf729a3f36d4df9a5", size = 169288, upload-time = "2026-07-28T06:05:25.378Z" }, + { url = "https://files.pythonhosted.org/packages/c7/30/0d09e6dddc6b7a7230ac77f50254b5980ab4fcd22976f72f8cc8a0404458/wrapt-2.3.0-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:1d6159c9b2fefec02314e1332dbbbfaf960e369dfd26bcf7f8b258b5732065b3", size = 160932, upload-time = "2026-07-28T06:05:27.022Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ca/0913af0d2ec0c43865d32d615f518fea66c13c5c930e489e9b0de248e9a8/wrapt-2.3.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:24da48596326ef8e448cfa837b454f638713d3531262375f00e5a9681682fc07", size = 169017, upload-time = "2026-07-28T06:05:28.501Z" }, + { url = "https://files.pythonhosted.org/packages/c3/f2/3d1e47ea81b822210f5df1bf942fd90780a75c055243d569b664529dea88/wrapt-2.3.0-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:cd3a2edf0427013736b8127955cec62608c56e53ea47e82812ea32059cda407f", size = 159065, upload-time = "2026-07-28T06:05:30.01Z" }, + { url = "https://files.pythonhosted.org/packages/43/a5/ef2066ced8e5fca204e2b361e9708e36555b40949c583d997ea3b590817d/wrapt-2.3.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:4fa0df3bff4e7ce45759f33fd39335fe2f60477bb9ecf7b8aa41e7d07ee36a23", size = 168821, upload-time = "2026-07-28T06:05:31.649Z" }, + { url = "https://files.pythonhosted.org/packages/d5/e1/016104650d4e572fa91506eb396b3dd8efbccc9284fdc1c9479c3d21db28/wrapt-2.3.0-cp314-cp314-win32.whl", hash = "sha256:2935d5454b3f179a29b12cf390ee47246740ba2c3a7545b1b46ba31a5f2a4a0b", size = 78700, upload-time = "2026-07-28T06:05:33.391Z" }, + { url = "https://files.pythonhosted.org/packages/3d/97/6fdc20a9f2ca304748b3f0819cbf377d55260562777bf0b615431bc3c181/wrapt-2.3.0-cp314-cp314-win_amd64.whl", hash = "sha256:cc2cea812e5cb179a796b766747e7d3b21088760d8deb95676d482b8c8e6fa7d", size = 81422, upload-time = "2026-07-28T06:05:34.774Z" }, + { url = "https://files.pythonhosted.org/packages/5e/a4/9cbd53bf05746bea2c392af39cb052427a8ec95cbd494d930733d8f44681/wrapt-2.3.0-cp314-cp314-win_arm64.whl", hash = "sha256:22cc5c0a717bd4da87018ae0bffd4c19c6fb679d3ff357216ba566ab26c76cab", size = 80639, upload-time = "2026-07-28T06:05:36.228Z" }, + { url = "https://files.pythonhosted.org/packages/43/bb/6c5e4a0f66ea0d2b2dd267e8dd05a0014eea56840b3c8595d40b0a5d1f91/wrapt-2.3.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:a6b5984cd65dd639546f0eb4b8eacf1c31cb2fe9fb5c27bffe240987cdb2cf84", size = 84030, upload-time = "2026-07-28T06:05:37.714Z" }, + { url = "https://files.pythonhosted.org/packages/6a/eb/a1aedf03283bc9cbf8a1783995ddc54e3c5a86878f19002d2c428494f4c5/wrapt-2.3.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c88abcf53daef80e01a75c7530e727fa6e2c1888fe83e3dcdba4c96216a1f5c7", size = 84419, upload-time = "2026-07-28T06:05:39.131Z" }, + { url = "https://files.pythonhosted.org/packages/63/61/50d511c0dc5105563849e86daa3e16ac7feef699f79fb05af45ea70107d5/wrapt-2.3.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:85de890ff968196e92dd1ae73a9fb8970495e7650a457b1c9ef0ac3dd550bce2", size = 207171, upload-time = "2026-07-28T06:05:40.69Z" }, + { url = "https://files.pythonhosted.org/packages/3f/59/9b538cf7795217e810699d16bc88b96a830d9b5c403eb2ec2db6b5f2ae81/wrapt-2.3.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:50f416b74d092bb9f41b424e90dd457f365f7ba4b11de62a23679769a21bd85c", size = 214329, upload-time = "2026-07-28T06:05:42.287Z" }, + { url = "https://files.pythonhosted.org/packages/b3/28/9935d62b1499e5c8b3d191e99ba4eb31ca237a0b699142011a837e9dc7ea/wrapt-2.3.0-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:39febbee6d77301d31da6996b152ce52452da7c7ef72aba10c2fa976dff9c295", size = 199079, upload-time = "2026-07-28T06:05:43.958Z" }, + { url = "https://files.pythonhosted.org/packages/2b/01/4446b80fa2ffa47a3449b250d004ba1c1937f07f64a179608fec735df866/wrapt-2.3.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:93513bec052c6cd987f9f580c3df068c8bc4ebae6543736be3ca7ec5959cafcd", size = 209992, upload-time = "2026-07-28T06:05:45.677Z" }, + { url = "https://files.pythonhosted.org/packages/d4/07/56f26c9f9979586a021e8148747004aba4498f49458c90b0502969b904e1/wrapt-2.3.0-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:729126e667da34d251b8ebf8a45ef0c5ddadc21542b3d6e1abf4259ece6508df", size = 196334, upload-time = "2026-07-28T06:05:47.608Z" }, + { url = "https://files.pythonhosted.org/packages/8b/41/6d7bcc895b0f28b2250e10908f060687b9165429dcd7f22ddb3d4c031b74/wrapt-2.3.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:626b69db2021aa01671ec7bbc9740e558522bd44c18cf2ce69bf3d666a014109", size = 202644, upload-time = "2026-07-28T06:05:49.183Z" }, + { url = "https://files.pythonhosted.org/packages/cd/25/7860927edba06b758b8852a6f02e832be715563c67a6795d94350bc81099/wrapt-2.3.0-cp314-cp314t-win32.whl", hash = "sha256:629d73378082c00a8173031f9fb30a3ac6abbc894a5bfdfae71fabc60642d501", size = 79685, upload-time = "2026-07-28T06:05:50.976Z" }, + { url = "https://files.pythonhosted.org/packages/c4/0f/270bafe92fde3b069a39bc01e39ee79340895b335640df861d43d2a51885/wrapt-2.3.0-cp314-cp314t-win_amd64.whl", hash = "sha256:42869085687f0aefd57c0f636c3f9354f8ffb321a8ba9cb52d19beb796e561c5", size = 83104, upload-time = "2026-07-28T06:05:52.405Z" }, + { url = "https://files.pythonhosted.org/packages/55/b3/af176d79a8515a8a720eccdad9a96f6e31a30abf2865430c8c42adf2fd13/wrapt-2.3.0-cp314-cp314t-win_arm64.whl", hash = "sha256:b1e5aa486e269b00ed35e64771c7d0ab8096cfd2643405ca8cd60ebedc099a51", size = 81774, upload-time = "2026-07-28T06:05:53.902Z" }, + { url = "https://files.pythonhosted.org/packages/00/39/3daf9f47be208606586de4568ba6713db53ebc8fd7a575aea1fe57983b69/wrapt-2.3.0-py3-none-any.whl", hash = "sha256:d8c7ed08477429752b8c44991f40ad7838b18332a160698740a6bfbc10d998a2", size = 61866, upload-time = "2026-07-28T06:06:12.9Z" }, +] + [[package]] name = "xarray" version = "2026.4.0" From a3259f64b762eed0ae38f68d890beb126b632c90 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Fri, 31 Jul 2026 19:05:24 +0100 Subject: [PATCH 05/13] more reliability fixes for zarr-init zarr-init used mode="w", which deletes the destination prefix before creating the group -- pointed at an existing store it would have wiped it. Use mode="w-" instead: create, never clobber, and no list or delete permission needed. Object stores answer 403, not 404, for a missing key when the caller cannot list the prefix, so existence probes cannot distinguish "absent" from "not allowed". remote.exists()/listdir() take an on_denied argument; init warns and proceeds, registry and lock reads treat it as absent, and the zone check in zarr-fill opens the group (GetObject) rather than probing the prefix (ListBucket). Reading a parquet no longer probes first, halving its round trips. s3fs surfaces a 403 as a builtin PermissionError rather than a botocore exception, so the CLI handler missed it and printed a traceback; catch the OSError family too and explain the 403-instead-of-404 behaviour. Document that Source Cooperative serves reads from data.source.coop but takes writes only on the backing bucket with no --endpoint-url, which is the likeliest cause of an AccessDenied here. Also adds the zarr test suite omitted from the previous commit. --- docs/architecture.rst | 17 ++ geotessera/registry_cli.py | 33 ++- geotessera/remote.py | 54 +++- geotessera/zarr.py | 80 ++++-- tests/zarr.t | 79 ++++++ tests/zarr_remote_check.py | 507 +++++++++++++++++++++++++++++++++++++ 6 files changed, 736 insertions(+), 34 deletions(-) create mode 100644 tests/zarr.t create mode 100644 tests/zarr_remote_check.py diff --git a/docs/architecture.rst b/docs/architecture.rst index d285cde..be5b383 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -550,6 +550,23 @@ CLI's ``--acl``:: It applies to the store's Zarr chunks and metadata as well as the sidecar parquet and lock objects, and is filtered out of read requests. +.. note:: + + **Source Cooperative reads and writes use different endpoints.** + ``https://data.source.coop`` is the read-only gateway: anonymous reads + work, but any write (and even a listing with write credentials) returns + ``AccessDenied``. Writes go directly to the backing AWS bucket with *no* + ``--endpoint-url`` at all:: + + # read from the gateway, write to the backing bucket + --source-endpoint-url https://data.source.coop --source-anon + --output s3://us-west-2.opendata.source.coop///zarr/v1 + --store-profile --store-region us-west-2 + --store-acl bucket-owner-full-control + + Passing ``--store-endpoint-url https://data.source.coop`` is the common + mistake and produces a 403 that looks like a credentials problem. + ``s3://`` locations need the optional ``s3`` extra, which pulls in ``s3fs`` and ``botocore``:: diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 01017a8..2f94717 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3390,15 +3390,20 @@ def _add_storage_args(parser, prefix: str, label: str, writable: bool = False) - def _object_store_errors() -> Tuple[type, ...]: """Exception types worth reporting as a message rather than a traceback. - Empty when botocore is absent (i.e. the s3 extra is not installed), in - which case no such error can be raised anyway. + fsspec translates most S3 failures into builtin OSErrors — a 403 arrives + as PermissionError, a 404 as FileNotFoundError — so catching only + botocore's own types misses the common cases. botocore's are included for + the errors raised before fsspec gets a chance to translate them (bad + endpoint, unresolvable credentials); the tuple is short when the s3 + extra is not installed. """ + types: Tuple[type, ...] = (PermissionError, OSError) try: from botocore.exceptions import BotoCoreError, ClientError - return (BotoCoreError, ClientError) + return types + (BotoCoreError, ClientError) except ImportError: - return () + return types def _report_store_error(e: Exception, console: "Console") -> int: @@ -3409,12 +3414,22 @@ def _report_store_error(e: Exception, console: "Console") -> int: """ from rich.markup import escape - console.print(f"[red]{emoji('❌ ')}{escape(str(e))}[/red]") - if isinstance(e, _object_store_errors() or ()): + detail = str(e) or type(e).__name__ + console.print(f"[red]{emoji('❌ ')}{escape(detail)}[/red]") + + if isinstance(e, ImportError): + return 1 + + console.print( + "Check the endpoint and credentials: --source-endpoint-url / " + "--store-endpoint-url, --*-profile, --*-anon, or the AWS_* " + "environment variables." + ) + if isinstance(e, PermissionError): console.print( - "Check the endpoint and credentials: --source-endpoint-url / " - "--store-endpoint-url, --*-profile, --*-anon, or the AWS_* " - "environment variables." + "A 403 on a key that does not exist yet usually means the " + "credentials lack s3:ListBucket on the prefix — S3 then hides " + "whether the object is there rather than answering 404." ) return 1 diff --git a/geotessera/remote.py b/geotessera/remote.py index 0c05189..f1b85e6 100644 --- a/geotessera/remote.py +++ b/geotessera/remote.py @@ -194,12 +194,36 @@ def get_fs(loc: str | Path, storage_options: Optional[Dict[str, Any]] = None): return _filesystem_cached(protocol_of(loc), _options_key(storage_options)) -def exists(loc: str | Path, storage_options: Optional[Dict[str, Any]] = None) -> bool: - """True if *loc* exists, locally or remotely.""" +def exists( + loc: str | Path, + storage_options: Optional[Dict[str, Any]] = None, + on_denied: Optional[bool] = None, +) -> bool: + """True if *loc* exists, locally or remotely. + + Args: + on_denied: What to answer when the store refuses to say. S3 returns + 403, not 404, for a key that does not exist when the caller lacks + ``s3:ListBucket`` on the prefix — it will not reveal whether the + object is there. Write-scoped credentials hit this constantly, so + probes that can treat "don't know" as "absent" pass ``False`` + here. ``None`` (the default) re-raises the PermissionError. + """ fs = get_fs(loc, storage_options) if fs is None: return Path(loc).exists() - return bool(fs.exists(str(loc))) + try: + return bool(fs.exists(str(loc))) + except PermissionError: + if on_denied is None: + raise + logger.warning( + f"Permission denied probing {loc}; assuming " + f"{'present' if on_denied else 'absent'}. This usually means the " + f"credentials lack s3:ListBucket on the prefix, which makes a " + f"missing key answer 403 instead of 404." + ) + return on_denied def read_bytes( @@ -243,11 +267,17 @@ def remove(loc: str | Path, storage_options: Optional[Dict[str, Any]] = None) -> def listdir( - loc: str | Path, storage_options: Optional[Dict[str, Any]] = None + loc: str | Path, + storage_options: Optional[Dict[str, Any]] = None, + on_denied: Optional[list] = None, ) -> list[str]: """List the immediate children of a directory/prefix, as full locations. Returns an empty list when the directory does not exist. + + Args: + on_denied: Value to return when listing is refused (no + ``s3:ListBucket``). ``None`` re-raises. """ fs = get_fs(loc, storage_options) if fs is None: @@ -255,11 +285,23 @@ def listdir( if not path.is_dir(): return [] return [str(p) for p in sorted(path.iterdir())] - if not fs.exists(str(loc)): + + try: + entries = sorted(fs.ls(str(loc), detail=False)) + except FileNotFoundError: return [] + except PermissionError: + if on_denied is None: + raise + logger.warning( + f"Permission denied listing {loc}; treating it as empty. The " + f"credentials need s3:ListBucket on this prefix." + ) + return on_denied + protocol = protocol_of(loc) out = [] - for entry in sorted(fs.ls(str(loc), detail=False)): + for entry in entries: # fsspec strips the protocol from listing results; restore it so the # entries are usable as standalone locations. Both shapes round-trip: # "bucket/key" -> "s3://bucket/key", "/abs/path" -> "file:///abs/path". diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 199c8fe..7864134 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -205,10 +205,12 @@ def join(self, *parts: str) -> str: return remote.join(self.url, *parts) - def exists(self, *parts: str) -> bool: + def exists(self, *parts: str, on_denied: Optional[bool] = None) -> bool: from . import remote - return remote.exists(self.join(*parts), self.storage_options) + return remote.exists( + self.join(*parts), self.storage_options, on_denied=on_denied + ) def read_bytes(self, *parts: str) -> bytes: from . import remote @@ -225,10 +227,12 @@ def remove(self, *parts: str) -> None: remote.remove(self.join(*parts), self.storage_options) - def listdir(self, *parts: str) -> List[str]: + def listdir(self, *parts: str, on_denied: Optional[List[str]] = None) -> List[str]: from . import remote - return remote.listdir(self.join(*parts), self.storage_options) + return remote.listdir( + self.join(*parts), self.storage_options, on_denied=on_denied + ) def _ensure_backend(self) -> None: """Fail early, with an actionable message, if the backend is missing. @@ -1147,8 +1151,25 @@ def init_store( import zarr store = StoreLocation.resolve(output_path, storage_options, state_url) - if store.exists("zarr.json") or (not store.is_remote and Path(store.url).exists()): - raise FileExistsError(f"Store already exists: {store}") + if not store.is_remote: + if Path(store.url).exists(): + raise FileExistsError(f"Store already exists: {store}") + else: + # Write-scoped credentials commonly cannot list the prefix, and S3 + # then answers 403 rather than 404 for the key we are probing. That + # is not a reason to refuse to create a store — say so and carry on. + try: + if store.exists("zarr.json"): + raise FileExistsError(f"Store already exists: {store}") + except PermissionError: + if console: + console.print( + " [yellow]Could not check whether a store already exists " + "here (permission denied on HEAD — credentials without " + "s3:ListBucket get 403 instead of 404 for a missing key). " + "Creating it; an existing store's root would be " + "overwritten.[/yellow]" + ) years = sorted(years) T = len(years) @@ -1167,8 +1188,16 @@ def init_store( console.print(f" {len(landmask_by_zone)} zone(s) with land coverage") # Create root group via zarr API (not manual JSON) so consolidation - # preserves attributes correctly. - root = store.open_group(mode="w", zarr_format=3, use_consolidated=None) + # preserves attributes correctly. Mode "w-" creates but never clobbers: + # "w" would delete the destination prefix first, which needs list and + # delete permissions and would destroy an existing store if our + # pre-check above could not see it. + from zarr.errors import ContainsArrayError, ContainsGroupError + + try: + root = store.open_group(mode="w-", zarr_format=3, use_consolidated=None) + except (ContainsGroupError, ContainsArrayError) as e: + raise FileExistsError(f"Store already exists: {store}") from e root.attrs.update( { "zarr_conventions": [GEOEMB_CONVENTION], @@ -1381,17 +1410,28 @@ def _empty_tile_registry() -> "geopandas.GeoDataFrame": def _read_parquet_at(store: StoreLocation, *parts: str): - """Read a GeoParquet object from the store, or None if absent.""" + """Read a GeoParquet object from the store, or None if absent. + + Reads straight through rather than probing first: it halves the round + trips, and an existence probe is unreliable anyway against credentials + without list permission, where a missing key answers 403 not 404. + """ import geopandas as gpd - if not store.exists(*parts): - return None try: - return gpd.read_parquet(io.BytesIO(store.read_bytes(*parts))) + data = store.read_bytes(*parts) + except (FileNotFoundError, PermissionError): + return None except Exception as e: logger.warning(f"Could not read {store.join(*parts)}: {e}") return None + try: + return gpd.read_parquet(io.BytesIO(data)) + except Exception as e: + logger.warning(f"Could not parse {store.join(*parts)}: {e}") + return None + def _write_parquet_at(store: StoreLocation, gdf, *parts: str) -> None: """Write a GeoParquet object into the store (local path or remote URL).""" @@ -1516,7 +1556,7 @@ def merge_tile_registry( frames.append(previous) n_parts = 0 - for entry in state.listdir(REGISTRY_DIR_NAME): + for entry in state.listdir(REGISTRY_DIR_NAME, on_denied=[]): if not entry.endswith(".parquet"): continue try: @@ -1572,7 +1612,7 @@ def _acquire_zone_lock( state = store.state name = _lock_name(zone, year) - if not force and state.exists(LOCK_DIR_NAME, name): + if not force and state.exists(LOCK_DIR_NAME, name, on_denied=False): try: held = json.loads(state.read_bytes(LOCK_DIR_NAME, name)) except Exception: @@ -1807,7 +1847,7 @@ def extend_store( if not years: raise ValueError("No years given to add") - held = [Path(p).name for p in store.state.listdir(LOCK_DIR_NAME)] + held = [Path(p).name for p in store.state.listdir(LOCK_DIR_NAME, on_denied=[])] if held and not force: raise RuntimeError( f"{len(held)} fill lock(s) present ({', '.join(sorted(held)[:4])}" @@ -1952,7 +1992,12 @@ def fill_store( for zone_num, tile_infos in sorted(year_tiles.items()): zone_group = _zone_group_name(zone_num) - if not store.exists(zone_group): + # Open the group rather than probing for the prefix: reading + # utm{n}/zarr.json needs only GetObject, whereas an existence + # check on a prefix needs list permission the writer may lack. + try: + zone_store = store.open_group(mode="r", path=zone_group) + except Exception: if console: console.print( f" [yellow]Zone {zone_num} not initialised, skipping[/yellow]" @@ -1979,9 +2024,6 @@ def fill_store( f"{len(remaining)}/{len(tile_infos)} tiles to write" ) - # Read the zone grid from store metadata - zone_store = store.open_group(mode="r", path=zone_group) - # Resolve the time index against *this* zone's own axis. An # interrupted zarr-extend can leave zones with different lengths, # and a store-wide index would then address the wrong year. diff --git a/tests/zarr.t b/tests/zarr.t new file mode 100644 index 0000000..8b8aa5a --- /dev/null +++ b/tests/zarr.t @@ -0,0 +1,79 @@ +Zarr Remote Store Tests +======================= + +These cover the pieces that let `geotessera-registry zarr-fill` write to a +remote object store while streaming tiles from another one, and the state +that makes per-zone fills safe to run in parallel. + +Setup +----- + + $ export TERM=dumb + +Test: Location-transparent I/O and parallel-fill state +------------------------------------------------------ + +Everything here runs offline against a temporary directory, using `file://` +URLs to exercise the same fsspec path an `s3://` store takes: + + $ python "$TESTDIR/zarr_remote_check.py" | tail -1 + all checks passed + +Test: zarr-fill accepts remote locations +----------------------------------------- + +Both the tile source and the store may be URLs: + + $ geotessera-registry zarr-fill --help | grep -c 'Path or URL of an existing tessera store' + 1 + +Per-zone sweeps are documented on the command itself: + + $ geotessera-registry zarr-fill --help | grep -c 'concurrently' + 1 + +Object-store credentials come from flags or the environment, never from +positional arguments: + + $ geotessera-registry zarr-fill --help | grep -oE '\-\-(source|store)-endpoint-url' | sort -u + --source-endpoint-url + --store-endpoint-url + +Consolidation is opt-out for a zone-restricted fill: + + $ geotessera-registry zarr-fill --help | grep -oE '\-\-(no-)?consolidate' | sort -u + --consolidate + --no-consolidate + +A dead sibling job's lock can be taken over explicitly: + + $ geotessera-registry zarr-fill --help | grep -o '\-\-force-lock' | sort -u + --force-lock + +Test: zarr-extend grows the time axis +-------------------------------------- + +Adding a year is a metadata-only edit, and only ever appends: + + $ geotessera-registry zarr-extend --help | grep -c 'metadata-only edit' + 1 + + $ geotessera-registry zarr-extend --help | grep -o '\-\-years YEARS' | sort -u + --years YEARS + +Test: zarr-consolidate finishes a sweep +---------------------------------------- + + $ geotessera-registry zarr-consolidate --help | grep -c 'single-writer step' + 1 + + $ geotessera-registry zarr-consolidate --help | grep -o '\-\-no-merge-registry' | sort -u + --no-merge-registry + +Test: exit status propagates +----------------------------- + +A failing command must report failure so a sweep orchestrator notices: + + $ geotessera-registry zarr-consolidate /nonexistent/store.zarr > /dev/null 2>&1 + [1] diff --git a/tests/zarr_remote_check.py b/tests/zarr_remote_check.py new file mode 100644 index 0000000..3c7b29c --- /dev/null +++ b/tests/zarr_remote_check.py @@ -0,0 +1,507 @@ +"""Checks for location-transparent I/O and parallel-safe zarr fill state. + +Run by tests/zarr.t. Every check works on small in-memory arrays and a +temporary directory, using ``file://`` URLs to drive the same fsspec code +path a remote ``s3://`` store takes — so the suite stays offline and fast +while still covering the byte-range reader, the per-zone tracking files and +the advisory locks. + +Prints one ``ok - `` line per check and exits non-zero on the first +failure. +""" + +import logging +import os +import sys +import tempfile +from pathlib import Path + +import numpy as np + +from geotessera import remote +from geotessera.zarr import ( + LOCK_DIR_NAME, + REGISTRY_DIR_NAME, + StoreLocation, + TileInfo, + TileSource, + UnifiedZoneGrid, + _acquire_zone_lock, + _get_written_tiles, + _record_written_tiles, + _release_zone_lock, + build_shard_index, + merge_tile_registry, + shard_coords_for_tiles, +) + +TMP = Path(tempfile.mkdtemp(prefix="gt-zarr-check-")) +FAILED = [] + + +def check(name, condition): + if condition: + print(f"ok - {name}") + else: + print(f"FAIL - {name}") + FAILED.append(name) + + +def url(path): + return f"file://{Path(path).resolve()}" + + +# --------------------------------------------------------------------------- +# Byte-range .npy reads +# --------------------------------------------------------------------------- + +emb = np.random.default_rng(0).integers(-128, 127, size=(40, 30, 8), dtype=np.int8) +scales = np.random.default_rng(1).random((40, 30)).astype(np.float32) +np.save(TMP / "emb.npy", emb) +np.save(TMP / "scales.npy", scales) + +windows = [(0, 40, 0, 30), (5, 12, 0, 30), (7, 33, 4, 19), (38, 40, 29, 30)] +for r0, r1, c0, c1 in windows: + local = remote.read_npy_window(TMP / "emb.npy", r0, r1, c0, c1) + over_url = remote.read_npy_window(url(TMP / "emb.npy"), r0, r1, c0, c1) + check( + f"npy window {r0}:{r1},{c0}:{c1} matches source", + np.array_equal(local, emb[r0:r1, c0:c1]), + ) + check( + f"npy window {r0}:{r1},{c0}:{c1} identical over url", + np.array_equal(local, over_url), + ) + +check( + "2-D npy window (scales) reads over url", + np.array_equal( + remote.read_npy_window(url(TMP / "scales.npy"), 3, 9, 2, 11), + scales[3:9, 2:11], + ), +) + +# A window past the end of the array clamps rather than over-reading. +clamped = remote.read_npy_window(url(TMP / "emb.npy"), 35, 60, 0, 30) +check("npy window past EOF clamps to array height", clamped.shape[0] == 5) + +# --------------------------------------------------------------------------- +# GeoTIFF window reads +# --------------------------------------------------------------------------- + +import rasterio # noqa: E402 +from rasterio.transform import from_origin # noqa: E402 + +mask = np.ones((40, 30), dtype=np.uint8) +mask[:6, :6] = 0 +with rasterio.open( + TMP / "lm.tiff", + "w", + driver="GTiff", + height=40, + width=30, + count=1, + dtype="uint8", + crs="EPSG:4326", + transform=from_origin(0, 52, 0.01, 0.01), +) as dst: + dst.write(mask, 1) + +check( + "tiff window identical local and over url", + np.array_equal( + remote.read_tiff_window(TMP / "lm.tiff", 2, 10, 1, 9), + remote.read_tiff_window(url(TMP / "lm.tiff"), 2, 10, 1, 9), + ), +) +check( + "tiff window matches source", + np.array_equal( + remote.read_tiff_window(url(TMP / "lm.tiff"), 0, 6, 0, 6), + np.zeros((6, 6), dtype=np.uint8), + ), +) + +# --------------------------------------------------------------------------- +# StoreLocation basics +# --------------------------------------------------------------------------- + +for label, loc in [ + ("local", StoreLocation(str(TMP / "store_local"))), + ("url", StoreLocation(url(TMP / "store_url"))), +]: + check(f"{label} store reports remote correctly", loc.is_remote == (label == "url")) + check(f"{label} store missing object absent", not loc.exists("nope.bin")) + loc.write_bytes(b"hello", "sub", "a.bin") + check(f"{label} store write creates parents", loc.exists("sub", "a.bin")) + check(f"{label} store read round-trips", loc.read_bytes("sub", "a.bin") == b"hello") + listed = loc.listdir("sub") + check(f"{label} store lists children", len(listed) == 1) + # Listed entries must be usable as standalone locations — merge_tile_registry + # reads the per-zone files straight from a listing. + check( + f"{label} store listing entries are readable locations", + remote.read_bytes(listed[0], loc.storage_options) == b"hello", + ) + loc.remove("sub", "a.bin") + check(f"{label} store remove deletes", not loc.exists("sub", "a.bin")) + loc.remove("sub", "a.bin") # removing twice must not raise + check(f"{label} store remove is idempotent", True) + +# --------------------------------------------------------------------------- +# Permission-denied probes +# --------------------------------------------------------------------------- +# S3 answers 403, not 404, for a key that does not exist when the caller +# lacks s3:ListBucket. Write-scoped credentials hit this on every probe, so +# callers that can treat "don't know" as "absent" must be able to say so. + + +class _DenyingFS: + def exists(self, path): + raise PermissionError("Forbidden") + + def ls(self, path, detail=False): + raise PermissionError("Forbidden") + + def cat_file(self, path, start=None, end=None): + raise PermissionError("Forbidden") + + +# The tolerant paths log a warning by design; silence it for the run. +logging.getLogger("geotessera.remote").setLevel(logging.ERROR) + +_real_get_fs = remote.get_fs +remote.get_fs = lambda loc, so=None: ( + _DenyingFS() if remote.is_url(loc) else _real_get_fs(loc, so) +) +try: + denied = "s3://bucket/store.zarr/zarr.json" + try: + remote.exists(denied) + check("denied probe raises by default", False) + except PermissionError: + check("denied probe raises by default", True) + + check( + "denied probe can assume absent", + remote.exists(denied, on_denied=False) is False, + ) + check( + "denied probe can assume present", remote.exists(denied, on_denied=True) is True + ) + + try: + remote.listdir("s3://bucket/store.zarr.build/_registry") + check("denied listing raises by default", False) + except PermissionError: + check("denied listing raises by default", True) + + check( + "denied listing can fall back to empty", + remote.listdir("s3://bucket/x", on_denied=[]) == [], + ) + # A registry read that is refused must look like "no registry yet", not + # crash a fill that is otherwise fine. + from geotessera.zarr import _read_parquet_at # noqa: E402 + + check( + "denied registry read reads as absent", + _read_parquet_at(StoreLocation("s3://bucket/store.zarr.build"), "x.parquet") + is None, + ) +finally: + remote.get_fs = _real_get_fs + + +# --------------------------------------------------------------------------- +# Per-zone ingestion registry +# --------------------------------------------------------------------------- + + +def tile(lon, lat, zone=31): + return TileInfo( + lon=lon, + lat=lat, + year=2024, + epsg=32600 + zone, + transform=None, + height=10, + width=10, + landmask_path="", + embedding_path="", + scales_path="", + ) + + +store = StoreLocation(str(TMP / "reg_store")) +check("registry empty before any fill", _get_written_tiles(store, 2024, 31) == set()) + +_record_written_tiles(store, [tile(0.05, 52.05), tile(0.15, 52.05)], 2024, 31) +check( + "registry records this zone/year", + _get_written_tiles(store, 2024, 31) == {(0.05, 52.05), (0.15, 52.05)}, +) +check("registry isolates other zones", _get_written_tiles(store, 2024, 30) == set()) +check("registry isolates other years", _get_written_tiles(store, 2023, 31) == set()) + +# A second zone writes its own object — the file a sibling job owns is +# untouched, which is what makes concurrent zone fills safe. +_record_written_tiles(store, [tile(-0.05, 52.05, zone=30)], 2024, 30) +parts = sorted(Path(p).name for p in store.state.listdir(REGISTRY_DIR_NAME)) +check( + "one registry object per zone/year", + parts == ["utm30_2024.parquet", "utm31_2024.parquet"], +) + +# Appending to a zone keeps earlier rows and dedupes repeats. +_record_written_tiles(store, [tile(0.15, 52.05), tile(0.25, 52.05)], 2024, 31) +check( + "registry append keeps and dedupes", + _get_written_tiles(store, 2024, 31) + == {(0.05, 52.05), (0.15, 52.05), (0.25, 52.05)}, +) + +n = merge_tile_registry(store) +check("merge folds every zone into the root registry", n == 4) +check( + "merged registry lands in the state sibling", + store.state.exists("_registry.parquet"), +) +check("nothing written into the store itself", not store.exists("_registry.parquet")) +check("state sibling sits next to the store", store.state.url == store.url + ".build") + +# The same merge must work through a URL location, since that is how a +# remote store is finished after a sweep. +url_store = StoreLocation(url(TMP / "reg_store_url")) +for zone, tiles in [(31, [tile(0.05, 52.05)]), (30, [tile(-0.05, 52.05, zone=30)])]: + _record_written_tiles(url_store, tiles, 2024, zone) +check("merge over a url location", merge_tile_registry(url_store) == 2) +check( + "url store resumes from its own registry", + _get_written_tiles(url_store, 2024, 31) == {(0.05, 52.05)}, +) + +# Stores built before the split kept the registry inside the hierarchy; +# it must still be read so those resume correctly. +legacy = StoreLocation(str(TMP / "legacy_store")) +legacy.write_bytes(store.state.read_bytes("_registry.parquet"), "_registry.parquet") +check( + "legacy root registry still resumes", + _get_written_tiles(legacy, 2024, 31) + == {(0.05, 52.05), (0.15, 52.05), (0.25, 52.05)}, +) + +# --------------------------------------------------------------------------- +# Advisory zone locks +# --------------------------------------------------------------------------- + +lock_store = StoreLocation(str(TMP / "lock_store")) +_acquire_zone_lock(lock_store, 31, 2024) +check( + "lock object created in the state sibling", + lock_store.state.exists(LOCK_DIR_NAME, "utm31_2024.json"), +) +check("no lock object inside the store", not lock_store.exists(LOCK_DIR_NAME)) + +try: + _acquire_zone_lock(lock_store, 31, 2024) + check("second acquire refused", False) +except RuntimeError as e: + check("second acquire refused", "locked by" in str(e)) + +# A different zone or year is a different lock, so sibling jobs proceed. +_acquire_zone_lock(lock_store, 30, 2024) +_acquire_zone_lock(lock_store, 31, 2023) +check("sibling zone/year locks independent", True) + +_acquire_zone_lock(lock_store, 31, 2024, force=True) +check("force takes over a stale lock", True) + +_release_zone_lock(lock_store, 31, 2024) +check( + "release removes the lock", not lock_store.exists(LOCK_DIR_NAME, "utm31_2024.json") +) + +# --------------------------------------------------------------------------- +# Shard index: rewriting a shard must carry its already-written neighbours +# --------------------------------------------------------------------------- + +from rasterio.transform import Affine # noqa: E402 + +grid = UnifiedZoneGrid( + zone=31, + years=[2024], + canonical_epsg=32631, + origin_x=0.0, + origin_y=100000.0, + width_px=8192, + height_px=8192, +) + + +def placed(lon, x_off): + t = tile(lon, 52.05) + t.transform = Affine(10.0, 0.0, x_off, 0.0, -10.0, 100000.0) + t.height, t.width = 1000, 1000 + return t + + +old = placed(0.05, 0.0) # shard (0, 0) +new = placed(0.15, 10000.0) # shard (0, 0), adjacent +far = placed(0.95, 60000.0) # shard (0, 1) + +check("shard coords derived per tile", shard_coords_for_tiles([old], grid) == {(0, 0)}) +check( + "distant tile lands in another shard", + shard_coords_for_tiles([far], grid) == {(0, 1)}, +) + +touched = shard_coords_for_tiles([new], grid) +specs = build_shard_index([old, new, far], grid, 0, restrict_to=touched) +check("only touched shards are rebuilt", [(s.sr, s.sc) for s in specs] == [(0, 0)]) +check( + "rebuilt shard carries the already-written neighbour", + len(specs[0].tiles) == 2, +) + +unrestricted = build_shard_index([old, new, far], grid, 0) +check("unrestricted index covers every shard", len(unrestricted) == 2) + +# --------------------------------------------------------------------------- +# Extending the time axis +# --------------------------------------------------------------------------- + +import zarr # noqa: E402 +from zarr.codecs import BloscCodec # noqa: E402 + +from geotessera.zarr import extend_store # noqa: E402 + +ext = StoreLocation(str(TMP / "ext_store")) +root = zarr.open_group(ext.url, mode="w", zarr_format=3) +for zname in ("utm30", "utm31"): + grp = root.create_group(zname) + grp.create_array( + "embeddings", + shape=(1, 4, 32, 32), + chunks=(1, 4, 16, 16), + shards=(1, 4, 32, 32), + dtype=np.int8, + fill_value=np.int8(0), + compressors=BloscCodec(cname="zstd"), + dimension_names=["time", "band", "y", "x"], + ) + grp.create_array( + "scales", + shape=(1, 32, 32), + chunks=(1, 16, 16), + shards=(1, 32, 32), + dtype=np.float32, + fill_value=np.float32("inf"), + dimension_names=["time", "y", "x"], + ) + t = grp.create_array("time", shape=(1,), dtype=np.int32, dimension_names=["time"]) + t[:] = [2024] +# Existing data that must survive the resize untouched. +root["utm31"]["embeddings"][0] = 7 +root["utm31"]["scales"][0] = 0.5 + +check("extend adds the year to every zone", extend_store(ext, [2026]) == 2) + +g31 = zarr.open_group(ext.url, mode="r", use_consolidated=False)["utm31"] +check("time axis grew", [int(v) for v in g31["time"][:]] == [2024, 2026]) +check("arrays grew along time", g31["embeddings"].shape[0] == 2) +check( + "existing year untouched by the resize", + bool((np.asarray(g31["embeddings"][0]) == 7).all()) + and bool((np.asarray(g31["scales"][0]) == 0.5).all()), +) +check( + "new year reads as freshly initialised", + bool((np.asarray(g31["embeddings"][1]) == 0).all()) + and bool(np.isinf(np.asarray(g31["scales"][1])).all()), +) + +check("extend is idempotent", extend_store(ext, [2026]) == 0) + +try: + extend_store(ext, [2020]) + check("inserting an earlier year refused", False) +except ValueError as e: + check("inserting an earlier year refused", "only be appended" in str(e)) + +ext.state.write_bytes(b"{}", LOCK_DIR_NAME, "utm30_2026.json") +try: + extend_store(ext, [2027]) + check("extend refuses while a fill lock is held", False) +except RuntimeError as e: + check("extend refuses while a fill lock is held", "fill lock" in str(e)) +check( + "extend --force overrides a stale lock", extend_store(ext, [2027], force=True) == 2 +) + +# --------------------------------------------------------------------------- +# Storage options and source layout +# --------------------------------------------------------------------------- + +for var in ("AWS_ENDPOINT_URL", "AWS_DEFAULT_REGION", "AWS_REGION", "AWS_PROFILE"): + os.environ.pop(var, None) + +check("no options when nothing configured", remote.build_storage_options() is None) +check( + "explicit endpoint and anon are passed through", + remote.build_storage_options(endpoint_url="https://s3.example", anon=True) + == {"endpoint_url": "https://s3.example", "anon": True}, +) + +os.environ["AWS_ENDPOINT_URL"] = "https://from-env.example" +check( + "endpoint falls back to the environment", + remote.build_storage_options()["endpoint_url"] == "https://from-env.example", +) +check( + "explicit endpoint beats the environment", + remote.build_storage_options(endpoint_url="https://explicit.example")[ + "endpoint_url" + ] + == "https://explicit.example", +) +os.environ.pop("AWS_ENDPOINT_URL") + +check( + "canned acl becomes an s3fs write kwarg", + remote.build_storage_options(acl="bucket-owner-full-control")[ + "s3_additional_kwargs" + ] + == {"ACL": "bucket-owner-full-control"}, +) +try: + remote.build_storage_options(acl="not-an-acl") + check("bad acl rejected up front", False) +except ValueError: + check("bad acl rejected up front", True) + +src = TileSource.for_url("s3://bucket/tessera", "v1", {"anon": True}) +e, s = src.embedding_locations(0.05, 52.05, 2024) +check( + "embedding location follows the published layout", + e == "s3://bucket/tessera/npy/v1/2024/grid_0.05_52.05/grid_0.05_52.05.npy", +) +check( + "scales location follows the published layout", + s == "s3://bucket/tessera/npy/v1/2024/grid_0.05_52.05/grid_0.05_52.05_scales.npy", +) +check( + "landmask location follows the published layout", + src.landmask_location(0.05, 52.05) + == "s3://bucket/tessera/landmasks/v1/grid_0.05_52.05.tiff", +) +check("url source reports remote", src.is_remote) + +import shutil # noqa: E402 + +shutil.rmtree(TMP, ignore_errors=True) + +if FAILED: + print(f"\n{len(FAILED)} check(s) failed: {', '.join(FAILED)}") + sys.exit(1) +print("\nall checks passed") From 96d6dbeff34d799e843e0960edf63c08de3f6ce7 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Sat, 1 Aug 2026 07:02:17 +0100 Subject: [PATCH 06/13] resume a crashed zarr-fill from the store itself, and add zarr-scan The ingestion registry is only written when a (zone, year) finishes, so a run killed partway -- an OOM kill leaves no traceback -- loses that year's bookkeeping even though the shards it wrote are safely in the store. The shard objects survive anything, so use them as the resume signal: zarr-fill --skip-existing-shards lists the shards already present for a (zone, year), skips them, and records their tiles so later runs need no flag. It assumes the tile inventory has not grown since those shards were written. Where the credentials cannot list the store it probes only the shards the run would touch, and an empty listing is distinguished from a denied one so a fresh zone costs no probes. zarr-scan reuses that scan to inventory a store against the manifest without writing anything. Shards are classified written, missing, or empty -- no manifest tiles fall in an empty one, so it is ocean or outside coverage and will never be filled. Keeping those separate makes the percentages meaningful: they are over land, not over each zone's bounding box. Prints per-zone/year and per-year summaries and can write the full per-shard index as parquet. Also warn when the worker count cannot fit in RAM. Each worker holds a full (128, 4096, 4096) int8 shard buffer plus scales, about 2.1 GiB, so --workers 16 needs 33 GiB resident and is silently OOM-killed on a smaller box. --- CHANGES.md | 17 ++ docs/architecture.rst | 34 ++++ geotessera/registry_cli.py | 96 ++++++++++ geotessera/zarr.py | 353 +++++++++++++++++++++++++++++++++++-- tests/zarr.t | 24 +++ 5 files changed, 511 insertions(+), 13 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index a14da08..11bf37d 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -51,6 +51,23 @@ unrecognised objects and readers saw non-Zarr entries. A `_registry.parquet` left inside an older store is still read, so existing stores resume correctly; nothing is written back into them. (@avsm) +- **`geotessera-registry zarr-scan`**: New subcommand that inventories a + store's shards against the manifest without writing anything, classifying + each as `written`, `missing`, or `empty` (no manifest tiles fall in it, so + it is ocean or outside coverage and will never be filled). Prints + per-zone/year and per-year summaries of how much is left to fill — + percentages are over land, not over each zone's bounding box — and + optionally writes the per-shard index as parquet. (@avsm) +- **`zarr-fill --skip-existing-shards`**: Resume from the store itself. The + ingestion registry is only written when a (zone, year) finishes, so a run + killed partway loses that year's bookkeeping even though its shards are + safely written; the shard objects survive anything. This scans for them, + skips them and records their tiles. Assumes the tile inventory has not + grown since. Falls back to probing just the shards in hand where the + credentials cannot list the store. (@avsm) +- **`zarr-fill` warns when the worker count will not fit in RAM**: each + worker holds a ~2.1 GiB shard buffer, so `--workers 16` needs 33 GiB and + an OOM kill leaves no traceback to diagnose. (@avsm) - **`geotessera-registry zarr-extend`**: New subcommand that appends years to an existing store's time axis, so a new year can be added without rebuilding. Time is chunked one year per chunk, making this a diff --git a/docs/architecture.rst b/docs/architecture.rst index be5b383..5f28075 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -650,6 +650,40 @@ published hierarchy contains only Zarr: * **Consolidation is skipped** by default when ``--zones`` is given, since the root ``zarr.json`` is the one object all jobs share. +Resuming After a Crash +~~~~~~~~~~~~~~~~~~~~~~ + +The ingestion registry is written when a (zone, year) finishes, so a run +that dies partway — an OOM kill leaves no traceback — loses that year's +bookkeeping even though the shards it wrote are safely in the store. + +The shard objects are the ground truth, and they survive anything:: + + geotessera-registry zarr-fill --zones 30 \ + --skip-existing-shards + +This lists the shard objects already present for each (zone, year), skips +them, and records their tiles so subsequent runs need no flag. It assumes +the tile inventory has not grown since those shards were written: a tile +added to the manifest afterwards falls inside an existing shard and would +be skipped rather than merged in. Where the credentials cannot list the +store, it falls back to probing only the shards the run would touch. + +To see what is outstanding before committing to a sweep:: + + geotessera-registry zarr-scan --output index.parquet + +Every shard is classified ``written``, ``missing``, or ``empty`` — the last +meaning no manifest tiles fall in it, so it is ocean or outside coverage and +will never be filled. Keeping those separate means the percentages are over +land, not over each zone's bounding box. The command prints per-zone/year +and per-year summaries and writes the full per-shard index as parquet. + +Note that worker memory, not cores, bounds the fill: each holds a full +``(128, 4096, 4096)`` int8 shard buffer plus its scales, about 2.1 GiB, so +``--workers 16`` needs 33 GiB resident. ``zarr-fill`` warns when the +requested count will not fit. + A sweep therefore looks like:: # fan out — one process per zone, in parallel diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 2f94717..49c95d7 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3721,6 +3721,7 @@ def zarr_fill_command(args): consolidate=consolidate, force_lock=args.force_lock, state_url=args.state_url, + skip_existing_shards=args.skip_existing_shards, ) except RuntimeError as e: console.print(f"[red]{emoji('❌ ')}{e}[/red]") @@ -3732,6 +3733,44 @@ def zarr_fill_command(args): return 0 +def zarr_scan_command(args): + """Report how much of a store still needs filling.""" + import warnings + from rich.console import Console + + from .zarr import scan_store, summarise_scan + + warnings.filterwarnings("ignore", message="Object at .* is not recognized") + + console = Console() + registry, source, _version, _variant = _resolve_source(args, console) + + store_options = _storage_options_for(args, "store", args.store_path) + years = _parse_int_range(args.years) if args.years else None + zones = _parse_int_range(args.zones) if args.zones else None + + try: + df = scan_store( + registry, + args.store_path, + years=years, + zones=zones, + console=console if args.verbose else None, + storage_options=store_options, + source=source, + state_url=args.state_url, + output=args.output, + ) + except ValueError as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) + + summarise_scan(df, console) + return 0 + + def zarr_extend_command(args): """Append new years to an existing store's time axis.""" import warnings @@ -4716,6 +4755,14 @@ def main(): action="store_true", help="Never rewrite the root consolidated metadata.", ) + zarr_fill_parser.add_argument( + "--skip-existing-shards", + action="store_true", + help="Treat a shard that is already in the store as done. Recovers a " + "run killed before it recorded progress, since the shard objects " + "outlive the bookkeeping. Assumes the tile inventory has not grown " + "since those shards were written.", + ) zarr_fill_parser.add_argument( "--force-lock", action="store_true", @@ -4728,6 +4775,55 @@ def main(): _add_storage_args(zarr_fill_parser, "store", "Output store", writable=True) zarr_fill_parser.set_defaults(func=zarr_fill_command) + # Zarr-scan command + zarr_scan_parser = subparsers.add_parser( + "zarr-scan", + help="Report how much of a store still needs filling", + description="Inventory a store's shards against the manifest without " + "writing anything. Each shard is classified written, missing, or " + "empty (no manifest tiles fall in it — ocean or outside coverage), " + "so the percentages are over land rather than over the zone's " + "bounding box. Prints per-zone/year and per-year summaries and can " + "write the full index as parquet.", + ) + zarr_scan_parser.add_argument( + "base_dir", + help="Base directory containing downloaded tile data, or a URL of a " + "repository in the published layout (used for the manifest)", + ) + zarr_scan_parser.add_argument( + "store_path", + type=str, + help="Path or URL of an existing tessera store", + ) + zarr_scan_parser.add_argument( + "--years", + default=None, + help="Years to scan (e.g. 2024 or 2017-2025). Default: every year " + "on the store's time axis", + ) + zarr_scan_parser.add_argument( + "--zones", + default=None, + help="Zones to scan (e.g. 30 or 29-34). Default: every zone", + ) + zarr_scan_parser.add_argument( + "--output", + type=str, + default=None, + help="Write the per-shard index here as parquet (path or URL)", + ) + zarr_scan_parser.add_argument( + "--verbose", + action="store_true", + help="Print progress for each zone/year as it is scanned", + ) + _add_source_args(zarr_scan_parser) + _add_state_arg(zarr_scan_parser) + _add_storage_args(zarr_scan_parser, "source", "Tile source") + _add_storage_args(zarr_scan_parser, "store", "Store") + zarr_scan_parser.set_defaults(func=zarr_scan_command) + # Zarr-extend command zarr_extend_parser = subparsers.add_parser( "zarr-extend", diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 7864134..ea589ad 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -60,6 +60,7 @@ if TYPE_CHECKING: import geopandas + import pandas import rich.console import zarr from rasterio.transform import Affine @@ -918,6 +919,47 @@ def _coarsen_tile( INNER_CHUNK = 32 # spatial pixels per inner chunk side DEFAULT_WORKERS = 4 # fewer workers due to larger shard buffers (~2GB each) +# Each shard worker holds a full (N_BANDS, SHARD_SIZE, SHARD_SIZE) int8 +# buffer plus its float32 scales — the dominant cost of a fill, and the +# reason the worker count is bounded by RAM rather than by cores. +WORKER_BUFFER_BYTES = N_BANDS * SHARD_SIZE * SHARD_SIZE + 4 * SHARD_SIZE * SHARD_SIZE + + +def _total_memory_bytes() -> Optional[int]: + """Physical RAM, or None where it cannot be determined.""" + try: + return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") + except (ValueError, OSError, AttributeError): + return None + + +def _warn_worker_memory(workers: int, console=None) -> None: + """Warn when the requested worker count cannot fit in RAM. + + A fill that is OOM-killed leaves no traceback, so the cause is easy to + miss; say it up front instead. + """ + needed = workers * WORKER_BUFFER_BYTES + total = _total_memory_bytes() + gib = 2**30 + message = ( + f"{workers} workers need at least {needed / gib:.0f} GiB " + f"({WORKER_BUFFER_BYTES / gib:.1f} GiB of shard buffer each)" + ) + if total is None: + logger.info(message) + return + if needed > 0.8 * total: + safe = max(1, int(0.8 * total // WORKER_BUFFER_BYTES)) + text = ( + f"{message}, but this machine has {total / gib:.0f} GiB. " + f"The fill will likely be OOM-killed — consider --workers {safe}." + ) + if console: + console.print(f" [yellow]{text}[/yellow]") + else: + logger.warning(text) + # --------------------------------------------------------------------------- # Data model @@ -1753,6 +1795,59 @@ def _write_one_shard_worker(spec: ShardSpec) -> bool: # --------------------------------------------------------------------------- +def _existing_shards( + store: StoreLocation, + zone_group: str, + time_index: int, + wanted: set, + console: Optional["rich.console.Console"] = None, +) -> set: + """Which of *wanted* shard coordinates already exist in the store. + + A written shard is a single object under the array's chunk prefix, so its + presence is proof the shard landed — bookkeeping that survives a ``kill + -9`` and needs no state file. Zarr v3's default chunk key encoding puts + the (time, band, row, col) chunk grid indices in the key, and the band + dimension is one chunk wide, so the shard for (sr, sc) at time t is + ``embeddings/c/{t}/0/{sr}/{sc}``. + + Prefers one listing of the time slice's prefix; falls back to probing + each wanted shard when listing is refused, which credentials without + ``s3:ListBucket`` will be. + """ + prefix = f"{zone_group}/embeddings/c/{time_index}/0" + + try: + rows = store.listdir(prefix) + except PermissionError: + # Credentials without s3:ListBucket. Probe each shard this fill would + # touch instead — bounded by the work in hand, not the whole zone. + if console: + console.print( + f" [dim]Cannot list the store; probing {len(wanted)} " + f"shard(s) individually[/dim]" + ) + return { + (sr, sc) + for sr, sc in sorted(wanted) + if store.exists(prefix, str(sr), str(sc), on_denied=False) + } + + # An empty listing means nothing has been written for this time slice, + # which is the common case on a fresh zone — no probing needed. + present = set() + for row_entry in rows: + sr_name = row_entry.rstrip("/").rsplit("/", 1)[-1] + if not sr_name.isdigit(): + continue + sr = int(sr_name) + for col_entry in store.listdir(prefix, sr_name, on_denied=[]): + sc_name = col_entry.rstrip("/").rsplit("/", 1)[-1] + if sc_name.isdigit() and (sr, int(sc_name)) in wanted: + present.add((sr, int(sc_name))) + return present + + def _store_years(store: StoreLocation, zones: Optional[List[int]] = None) -> List[int]: """Read the store's year axis from a zone's ``time`` coordinate array. @@ -1916,6 +2011,199 @@ def extend_store( return extended +def scan_store( + registry: "Registry", + store_path: "str | Path | StoreLocation", + years: Optional[List[int]] = None, + zones: Optional[List[int]] = None, + console: Optional["rich.console.Console"] = None, + storage_options: Optional[Dict[str, Any]] = None, + source: Optional[TileSource] = None, + state_url: Optional[str] = None, + output: Optional[str] = None, +) -> "pandas.DataFrame": + """Inventory a store's shards against the manifest, without writing data. + + Answers "how much is left to fill" from the store itself rather than from + bookkeeping, by listing the shard objects that exist and comparing them + with the shards the manifest says should exist. Each shard is classified: + + ``written`` + The shard object is in the store. + ``missing`` + The manifest has tiles here but no shard object exists — the work + still to do. + ``empty`` + No manifest tiles fall in this shard, so it is ocean or outside the + data's coverage and will never be filled. Reported separately so the + percentages are over land, not over the zone's bounding box. + + Returns a DataFrame with one row per (zone, year, shard), also written to + *output* as parquet when given. + """ + import pandas as pd + + store = StoreLocation.resolve(store_path, storage_options, state_url) + + all_years = _store_years(store, zones) + if not all_years: + raise ValueError("Store has no years (checked zone time coords)") + scan_years = [y for y in (years or all_years) if y in all_years] + + if console: + console.print(f"Scanning [bold]{store}[/bold]") + console.print(f" Years: {scan_years}") + + rows: List[Dict[str, Any]] = [] + + for scan_year in scan_years: + year_tiles = gather_tile_infos( + registry, scan_year, zones=zones, console=None, source=source + ) + + for zone_num, tile_infos in sorted(year_tiles.items()): + zone_group = _zone_group_name(zone_num) + try: + zone_store = store.open_group(mode="r", path=zone_group) + zone_years = [int(v) for v in zone_store["time"][:]] + except Exception: + continue + if scan_year not in zone_years: + continue + time_index = zone_years.index(scan_year) + + attrs = dict(zone_store.attrs) + transform = attrs["spatial:transform"] + shape = attrs["spatial:shape"] + grid = UnifiedZoneGrid( + zone=zone_num, + years=all_years, + canonical_epsg=int(attrs["proj:code"].split(":")[1]), + origin_x=transform[2], + origin_y=transform[5], + width_px=shape[1], + height_px=shape[0], + ) + + # Tiles per shard, so the index records how much work each holds. + per_shard: Dict[Tuple[int, int], int] = {} + for ti in tile_infos: + for coord in shard_coords_for_tiles([ti], grid): + per_shard[coord] = per_shard.get(coord, 0) + 1 + + expected = set(per_shard) + present = _existing_shards( + store, zone_group, time_index, expected, console=console + ) + + n_rows_grid = math.ceil(grid.height_px / SHARD_SIZE) + n_cols_grid = math.ceil(grid.width_px / SHARD_SIZE) + for sr in range(n_rows_grid): + for sc in range(n_cols_grid): + coord = (sr, sc) + n_tiles = per_shard.get(coord, 0) + if n_tiles == 0: + status = "empty" + elif coord in present: + status = "written" + else: + status = "missing" + rows.append( + { + "zone": zone_num, + "year": scan_year, + "shard_row": sr, + "shard_col": sc, + "n_tiles": n_tiles, + "status": status, + } + ) + + if console: + n_exp = len(expected) + n_have = len(present & expected) + pct = 100.0 * (n_exp - n_have) / n_exp if n_exp else 0.0 + console.print( + f" utm{zone_num:02d} {scan_year}: " + f"{n_have}/{n_exp} shards written, {pct:.1f}% to fill" + ) + + df = pd.DataFrame( + rows, + columns=["zone", "year", "shard_row", "shard_col", "n_tiles", "status"], + ) + + if output: + from . import remote + + buf = io.BytesIO() + df.to_parquet(buf, index=False) + remote.write_bytes(output, buf.getvalue(), storage_options) + if console: + console.print(f" Wrote shard index to [bold]{output}[/bold]") + + return df + + +def summarise_scan(df: "pandas.DataFrame", console: "rich.console.Console") -> None: + """Print per-zone/year and per-year fill summaries from a scan.""" + from rich.table import Table + + if df.empty: + console.print("[yellow]Nothing scanned.[/yellow]") + return + + def _pct(sub) -> float: + expected = int((sub["status"] != "empty").sum()) + missing = int((sub["status"] == "missing").sum()) + return 100.0 * missing / expected if expected else 0.0 + + detail = Table(title="Fill needed by zone and year") + for col in ("Zone", "Year", "Land shards", "Written", "Missing", "% to fill"): + detail.add_column(col, justify="right" if col != "Zone" else "left") + + for (zone, year), sub in df.groupby(["zone", "year"], sort=True): + expected = int((sub["status"] != "empty").sum()) + written = int((sub["status"] == "written").sum()) + missing = int((sub["status"] == "missing").sum()) + if expected == 0: + continue + detail.add_row( + f"utm{int(zone):02d}", + str(int(year)), + f"{expected:,}", + f"{written:,}", + f"{missing:,}", + f"{_pct(sub):.1f}%", + ) + console.print(detail) + + by_year = Table(title="Fill needed by year (all scanned zones)") + for col in ("Year", "Land shards", "Written", "Missing", "% to fill"): + by_year.add_column(col, justify="right" if col != "Year" else "left") + for year, sub in df.groupby("year", sort=True): + expected = int((sub["status"] != "empty").sum()) + if expected == 0: + continue + by_year.add_row( + str(int(year)), + f"{expected:,}", + f"{int((sub['status'] == 'written').sum()):,}", + f"{int((sub['status'] == 'missing').sum()):,}", + f"{_pct(sub):.1f}%", + ) + console.print(by_year) + + total_expected = int((df["status"] != "empty").sum()) + total_missing = int((df["status"] == "missing").sum()) + total_empty = int((df["status"] == "empty").sum()) + overall = 100.0 * total_missing / total_expected if total_expected else 0.0 + console.print( + f"Overall: {total_missing:,}/{total_expected:,} land shards to fill " + f"({overall:.1f}%); {total_empty:,} shard(s) are water/no-coverage." + ) + + def fill_store( registry: "Registry", store_path: "str | Path | StoreLocation", @@ -1928,6 +2216,7 @@ def fill_store( consolidate: Optional[bool] = None, force_lock: bool = False, state_url: Optional[str] = None, + skip_existing_shards: bool = False, ) -> int: """Incrementally fill a store with tile data. @@ -1946,6 +2235,12 @@ def fill_store( is set, because the root object is the one thing parallel zone jobs share — run ``zarr-consolidate`` once after the sweep. force_lock: Take over a (zone, year) lock held by another process. + skip_existing_shards: Treat a shard object that already exists as + done. Recovers a run killed before it could record progress, + since the objects outlive the bookkeeping. Assumes the tile + inventory has not grown since those shards were written — a tile + added to the manifest afterwards falls inside an existing shard + and would be skipped rather than merged in. """ store = StoreLocation.resolve(store_path, storage_options, state_url) if workers is None: @@ -1959,6 +2254,8 @@ def fill_store( fill_years = [year] if year is not None else all_years + _warn_worker_memory(workers, console) + if console: console.print(f"Filling store at [bold]{store}[/bold]") console.print(f" Years to fill: {fill_years}") @@ -2060,11 +2357,40 @@ def fill_store( tile_infos, grid, time_index, restrict_to=touched ) + # The shard objects in the store are the ground truth for what + # landed — unlike the ingestion registry they survive a kill -9, + # so a crashed run can be resumed by scanning for them. + skipped_specs: List[ShardSpec] = [] + if skip_existing_shards: + present = _existing_shards( + store, + zone_group, + time_index, + {(s.sr, s.sc) for s in shard_specs}, + console=console, + ) + if present: + skipped_specs = [ + s for s in shard_specs if (s.sr, s.sc) in present + ] + shard_specs = [ + s for s in shard_specs if (s.sr, s.sc) not in present + ] + if console: + console.print( + f" [cyan]{len(skipped_specs)} shard(s) already " + f"written, skipping[/cyan]" + ) + if console: - n_rewritten = sum(len(s.tiles) for s in shard_specs) - len(remaining) + pending = {(ti.lon, ti.lat) for ti in remaining} + rebuilt = { + (ov.embedding_path) for s in shard_specs for ov in s.tiles + } + n_rebuilt = max(0, len(rebuilt) - len(pending)) extra = ( - f", {n_rewritten} previously-written tile(s) rewritten" - if n_rewritten > 0 + f", {n_rebuilt} already-written tile(s) rebuilt" + if n_rebuilt > 0 else "" ) console.print( @@ -2097,16 +2423,17 @@ def fill_store( f" [red]{len(failed)} shard(s) failed[/red]" ) - # Only record tiles whose shards all landed, so a retry picks - # up exactly the work a transient failure left behind. - if failed: - recorded = [ - ti - for ti in remaining - if not (shard_coords_for_tiles([ti], grid) & failed) - ] - else: - recorded = remaining + # Record tiles whose shards all landed — counting the ones we + # skipped as landed, since they are already in the store — so + # a retry picks up exactly the work still outstanding. + done = {(s.sr, s.sc) for s in skipped_specs} | ( + {(s.sr, s.sc) for s in shard_specs} - failed + ) + recorded = [ + ti + for ti in remaining + if shard_coords_for_tiles([ti], grid) <= done + ] _record_written_tiles(store, recorded, fill_year, zone_num) finally: _release_zone_lock(store, zone_num, fill_year) diff --git a/tests/zarr.t b/tests/zarr.t index 8b8aa5a..f82fc6b 100644 --- a/tests/zarr.t +++ b/tests/zarr.t @@ -50,6 +50,30 @@ A dead sibling job's lock can be taken over explicitly: $ geotessera-registry zarr-fill --help | grep -o '\-\-force-lock' | sort -u --force-lock +Test: resume from the store itself +----------------------------------- + +Shard objects outlive the bookkeeping, so a killed run can be recovered by +scanning for them: + + $ geotessera-registry zarr-fill --help | grep -o '\-\-skip-existing-shards' | sort -u + --skip-existing-shards + +Test: zarr-scan reports outstanding work +----------------------------------------- + + $ geotessera-registry zarr-scan --help | grep -c 'written, missing, or' + 1 + +Shards with no manifest tiles are counted separately, so percentages are +over land rather than over the zone's bounding box: + + $ geotessera-registry zarr-scan --help | grep -c 'ocean or outside coverage' + 1 + + $ geotessera-registry zarr-scan --help | grep -o '\-\-output OUTPUT' | sort -u + --output OUTPUT + Test: zarr-extend grows the time axis -------------------------------------- From 01cf155a991a9b4f2bfa6112edc992f6f5f3c690 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Sat, 1 Aug 2026 07:08:45 +0100 Subject: [PATCH 07/13] update zarr-scan for optional source dir Scan a store on its own to report how much is left to fill. The store argument comes first and the tile mirror is now optional: what decides whether a shard can ever hold data is land coverage, and that comes from the landmask registry (~19 MB, fetched and cached) rather than the ~200 MB manifest. Supplying a mirror switches the denominator to each year's actual embedding coverage, which only matters where a year covers less than the full land area. Shards are classified written, missing, or empty -- no land falls in an empty one, so it is ocean or outside coverage and will never be filled. Keeping those separate makes the percentages meaningful: they are over land, not over each zone's bounding box, which for a coastal zone is mostly sea. Factors the tile-footprint projection out of gather_tile_infos so it can be applied to landmask tiles too, which have no embeddings to look up. --- CHANGES.md | 15 ++- docs/architecture.rst | 21 +++- geotessera/registry_cli.py | 25 +++- geotessera/zarr.py | 230 ++++++++++++++++++++++++++----------- tests/zarr.t | 9 ++ 5 files changed, 217 insertions(+), 83 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 11bf37d..304f703 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -52,12 +52,15 @@ left inside an older store is still read, so existing stores resume correctly; nothing is written back into them. (@avsm) - **`geotessera-registry zarr-scan`**: New subcommand that inventories a - store's shards against the manifest without writing anything, classifying - each as `written`, `missing`, or `empty` (no manifest tiles fall in it, so - it is ocean or outside coverage and will never be filled). Prints - per-zone/year and per-year summaries of how much is left to fill — - percentages are over land, not over each zone's bounding box — and - optionally writes the per-shard index as parquet. (@avsm) + store's shards without writing anything, classifying each as `written`, + `missing`, or `empty` (no land falls in it, so it is ocean or outside + coverage and will never be filled). Prints per-zone/year and per-year + summaries of how much is left to fill — percentages are over land, not + over each zone's bounding box — and optionally writes the per-shard index + as parquet. Takes the store alone: the land denominator comes from the + landmask registry (~19 MB, cached), so no tile mirror or manifest is + needed. An optional tile mirror switches the denominator to each year's + actual embedding coverage from the manifest. (@avsm) - **`zarr-fill --skip-existing-shards`**: Resume from the store itself. The ingestion registry is only written when a (zone, year) finishes, so a run killed partway loses that year's bookkeeping even though its shards are diff --git a/docs/architecture.rst b/docs/architecture.rst index 5f28075..bbb26de 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -669,15 +669,24 @@ added to the manifest afterwards falls inside an existing shard and would be skipped rather than merged in. Where the credentials cannot list the store, it falls back to probing only the shards the run would touch. -To see what is outstanding before committing to a sweep:: +To see what is outstanding before committing to a sweep, point ``zarr-scan`` +at the store on its own:: - geotessera-registry zarr-scan --output index.parquet + geotessera-registry zarr-scan s3://bucket/tessera.zarr --output index.parquet Every shard is classified ``written``, ``missing``, or ``empty`` — the last -meaning no manifest tiles fall in it, so it is ocean or outside coverage and -will never be filled. Keeping those separate means the percentages are over -land, not over each zone's bounding box. The command prints per-zone/year -and per-year summaries and writes the full per-shard index as parquet. +meaning no land falls in it, so it is ocean or outside coverage and will +never be filled. Keeping those separate means the percentages are over land, +not over each zone's bounding box, which for a coastal zone is mostly sea. +The command prints per-zone/year and per-year summaries and writes the full +per-shard index as parquet. + +No tile mirror or manifest is needed: the land denominator comes from the +landmask registry (~19 MB, fetched and cached), which is all that decides +whether a shard can ever hold data. Supplying a tile mirror as an optional +second argument switches the denominator to each year's actual embedding +coverage from the manifest (~200 MB) — worth it only where a year covers +less than the full land area. Note that worker memory, not cores, bounds the fill: each holds a full ``(128, 4096, 4096)`` int8 shard buffer plus its scales, about 2.1 GiB, so diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 49c95d7..cf91518 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3743,7 +3743,14 @@ def zarr_scan_command(args): warnings.filterwarnings("ignore", message="Object at .* is not recognized") console = Console() - registry, source, _version, _variant = _resolve_source(args, console) + + # The manifest is optional: land coverage alone says which shards can + # ever hold data, and that comes from the much smaller landmask registry. + if args.base_dir: + registry, source, dataset_version, _variant = _resolve_source(args, console) + else: + registry, source = None, None + dataset_version = args.dataset_version or "v1" store_options = _storage_options_for(args, "store", args.store_path) years = _parse_int_range(args.years) if args.years else None @@ -3760,6 +3767,8 @@ def zarr_scan_command(args): source=source, state_url=args.state_url, output=args.output, + dataset_version=dataset_version, + landmasks_path=args.landmasks_url, ) except ValueError as e: console.print(f"[red]{emoji('❌ ')}{e}[/red]") @@ -4786,16 +4795,20 @@ def main(): "bounding box. Prints per-zone/year and per-year summaries and can " "write the full index as parquet.", ) - zarr_scan_parser.add_argument( - "base_dir", - help="Base directory containing downloaded tile data, or a URL of a " - "repository in the published layout (used for the manifest)", - ) zarr_scan_parser.add_argument( "store_path", type=str, help="Path or URL of an existing tessera store", ) + zarr_scan_parser.add_argument( + "base_dir", + nargs="?", + default=None, + help="Optional tile mirror or repository URL. Without it, land " + "coverage comes from the landmask registry, which is all that is " + "needed and far smaller than the manifest. Give it to measure " + "against each year's actual embedding coverage instead.", + ) zarr_scan_parser.add_argument( "--years", default=None, diff --git a/geotessera/zarr.py b/geotessera/zarr.py index ea589ad..86fe5a1 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -447,6 +447,65 @@ def _zone_group_name(zone: int) -> str: return f"utm{zone:02d}" +def tile_zone(lon: float) -> int: + """UTM zone number (1-60) containing a longitude.""" + return max(1, min(60, int(math.floor((lon + 180) / 6)) + 1)) + + +def project_tile( + lon: float, + lat: float, + year: int = 0, + transformer_cache: Optional[Dict[int, Any]] = None, + landmask_path: str = "", + embedding_path: str = "", + scales_path: str = "", + pixel_size: float = 10.0, +) -> TileInfo: + """Compute a tile's UTM footprint from its centre coordinates alone. + + Deterministic — no file is opened — so it works for a landmask tile whose + embeddings do not exist as readily as for one that has data. + """ + from pyproj import Transformer as ProjTransformer + from rasterio.transform import Affine + + if transformer_cache is None: + transformer_cache = {} + + zone_num = tile_zone(lon) + epsg = (32700 if lat < 0 else 32600) + zone_num + + if epsg not in transformer_cache: + transformer_cache[epsg] = ProjTransformer.from_crs( + "EPSG:4326", f"EPSG:{epsg}", always_xy=True + ) + proj = transformer_cache[epsg] + + west, east = lon - 0.05, lon + 0.05 + south, north = lat - 0.05, lat + 0.05 + ul_e, ul_n = proj.transform(west, north) + ur_e, ur_n = proj.transform(east, north) + ll_e, ll_n = proj.transform(west, south) + lr_e, lr_n = proj.transform(east, south) + + origin_e = min(ul_e, ll_e) + origin_n = max(ul_n, ur_n) + + return TileInfo( + lon=lon, + lat=lat, + year=year, + epsg=epsg, + transform=Affine(pixel_size, 0.0, origin_e, 0.0, -pixel_size, origin_n), + height=round((origin_n - min(ll_n, lr_n)) / pixel_size), + width=round((max(ur_e, lr_e) - origin_e) / pixel_size), + landmask_path=landmask_path, + embedding_path=embedding_path, + scales_path=scales_path, + ) + + # --------------------------------------------------------------------------- # Landmask handling # --------------------------------------------------------------------------- @@ -503,8 +562,6 @@ def gather_tile_infos( mirror; pass a :class:`TileSource` built with :meth:`TileSource.for_url` to stream from a remote bucket. """ - from rasterio.transform import Affine - # Get tiles for this year from MultiIndex, filtering to those with data gdf = registry._registry_gdf try: @@ -542,62 +599,26 @@ def gather_tile_infos( ) # Build TileInfos using computed grid (no file I/O) - from pyproj import Transformer as ProjTransformer - if source is None: source = TileSource.for_local_mirror(registry) zones_dict: Dict[int, List[TileInfo]] = {} - transformer_cache: Dict[int, ProjTransformer] = {} - pixel_size = 10.0 + transformer_cache: Dict[int, Any] = {} for tile_year, tile_lon, tile_lat in tiles: - emb_path, scales_path = source.embedding_locations( - tile_lon, tile_lat, tile_year - ) - landmask_path = source.landmask_location(tile_lon, tile_lat) - - # Compute EPSG and zone from coordinates - zone_num = int(math.floor((tile_lon + 180) / 6)) + 1 - zone_num = max(1, min(60, zone_num)) + zone_num = tile_zone(tile_lon) if zone_set is not None and zone_num not in zone_set: continue - is_south = tile_lat < 0 - epsg = 32700 + zone_num if is_south else 32600 + zone_num - # Reuse cached transformer for this EPSG - if epsg not in transformer_cache: - transformer_cache[epsg] = ProjTransformer.from_crs( - "EPSG:4326", f"EPSG:{epsg}", always_xy=True - ) - proj = transformer_cache[epsg] - - # Project tile corners to UTM - west, east = tile_lon - 0.05, tile_lon + 0.05 - south, north = tile_lat - 0.05, tile_lat + 0.05 - ul_e, ul_n = proj.transform(west, north) - ur_e, ur_n = proj.transform(east, north) - ll_e, ll_n = proj.transform(west, south) - lr_e, lr_n = proj.transform(east, south) - - origin_e = min(ul_e, ll_e) - origin_n = max(ul_n, ur_n) - max_e = max(ur_e, lr_e) - min_n = min(ll_n, lr_n) - - width = round((max_e - origin_e) / pixel_size) - height = round((origin_n - min_n) / pixel_size) - tf_tuple = (pixel_size, 0.0, origin_e, 0.0, -pixel_size, origin_n) - - ti = TileInfo( - lon=tile_lon, - lat=tile_lat, - year=tile_year, - epsg=epsg, - transform=Affine(*tf_tuple), - height=height, - width=width, - landmask_path=landmask_path, + emb_path, scales_path = source.embedding_locations( + tile_lon, tile_lat, tile_year + ) + ti = project_tile( + tile_lon, + tile_lat, + tile_year, + transformer_cache, + landmask_path=source.landmask_location(tile_lon, tile_lat), embedding_path=emb_path, scales_path=scales_path, ) @@ -2011,8 +2032,62 @@ def extend_store( return extended +def load_landmask_tiles( + dataset_version: str = "v1", + landmasks_path: Optional[str] = None, + cache_dir: Optional[Path] = None, +) -> List[Tuple[float, float]]: + """Land tile centres for a dataset version, without touching the manifest. + + Land coverage is all that is needed to say which shards can ever hold + data, and the landmask registry is ~19 MB against the manifest's ~200 MB. + Fetched from the public mirror and cached unless *landmasks_path* points + at a local copy. + """ + import pandas as pd + + from .registry import _parse_dataset_version, download_file_to_temp + from .registry import landmasks_parquet_url + + version_path, _ = _parse_dataset_version(dataset_version) + + if landmasks_path and Path(landmasks_path).exists(): + path = Path(landmasks_path) + else: + if cache_dir is None: + if os.name == "nt": + base = Path(os.environ.get("LOCALAPPDATA", "~")).expanduser() + else: + base = Path(os.environ.get("XDG_CACHE_HOME", "~/.cache")).expanduser() + cache_dir = base / "geotessera" / version_path + cache_dir.mkdir(parents=True, exist_ok=True) + path = Path( + download_file_to_temp( + landmasks_parquet_url(version_path), + cache_path=cache_dir / "landmasks.parquet", + ) + ) + + df = pd.read_parquet(path, columns=["lon", "lat"]) + return list(zip(df["lon"].astype(float), df["lat"].astype(float))) + + +def _shards_from_tiles( + tile_coords: List[Tuple[float, float]], + grid: UnifiedZoneGrid, + transformer_cache: Dict[int, Any], +) -> Dict[Tuple[int, int], int]: + """Map tile centres to the shards they touch, with a tile count each.""" + per_shard: Dict[Tuple[int, int], int] = {} + for lon, lat in tile_coords: + ti = project_tile(lon, lat, transformer_cache=transformer_cache) + for coord in shard_coords_for_tiles([ti], grid): + per_shard[coord] = per_shard.get(coord, 0) + 1 + return per_shard + + def scan_store( - registry: "Registry", + registry: Optional["Registry"], store_path: "str | Path | StoreLocation", years: Optional[List[int]] = None, zones: Optional[List[int]] = None, @@ -2021,22 +2096,30 @@ def scan_store( source: Optional[TileSource] = None, state_url: Optional[str] = None, output: Optional[str] = None, + dataset_version: str = "v1", + landmasks_path: Optional[str] = None, ) -> "pandas.DataFrame": - """Inventory a store's shards against the manifest, without writing data. + """Inventory a store's shards, without writing data. Answers "how much is left to fill" from the store itself rather than from bookkeeping, by listing the shard objects that exist and comparing them - with the shards the manifest says should exist. Each shard is classified: + with the shards that could hold data. Each shard is classified: ``written`` The shard object is in the store. ``missing`` - The manifest has tiles here but no shard object exists — the work - still to do. + Land falls here but no shard object exists — the work still to do. ``empty`` - No manifest tiles fall in this shard, so it is ocean or outside the - data's coverage and will never be filled. Reported separately so the - percentages are over land, not over the zone's bounding box. + No land falls in this shard, so it is ocean or outside coverage and + will never be filled. Reported separately so the percentages are + over land, not over the zone's bounding box. + + The land denominator comes from the landmask registry (~19 MB, fetched + and cached automatically), so no tile mirror or manifest is needed — + scanning a remote store on its own is enough. Passing *registry* instead + narrows it to the tiles that version's manifest lists **for each year**, + which is exact where a year's embedding coverage is smaller than the land + area, at the cost of loading the much larger manifest. Returns a DataFrame with one row per (zone, year, shard), also written to *output* as parquet when given. @@ -2054,14 +2137,28 @@ def scan_store( console.print(f"Scanning [bold]{store}[/bold]") console.print(f" Years: {scan_years}") + transformer_cache: Dict[int, Any] = {} + land_by_zone: Dict[int, List[Tuple[float, float]]] = {} + if registry is None: + if console: + console.print(" Land coverage from the landmask registry") + for lon, lat in load_landmask_tiles(dataset_version, landmasks_path): + zone_num = tile_zone(lon) + if zones is None or zone_num in zones: + land_by_zone.setdefault(zone_num, []).append((lon, lat)) + rows: List[Dict[str, Any]] = [] for scan_year in scan_years: - year_tiles = gather_tile_infos( - registry, scan_year, zones=zones, console=None, source=source - ) + if registry is not None: + year_tiles = gather_tile_infos( + registry, scan_year, zones=zones, console=None, source=source + ) + zone_items: List[Tuple[int, Any]] = sorted(year_tiles.items()) + else: + zone_items = sorted(land_by_zone.items()) - for zone_num, tile_infos in sorted(year_tiles.items()): + for zone_num, coverage in zone_items: zone_group = _zone_group_name(zone_num) try: zone_store = store.open_group(mode="r", path=zone_group) @@ -2086,10 +2183,13 @@ def scan_store( ) # Tiles per shard, so the index records how much work each holds. - per_shard: Dict[Tuple[int, int], int] = {} - for ti in tile_infos: - for coord in shard_coords_for_tiles([ti], grid): - per_shard[coord] = per_shard.get(coord, 0) + 1 + if registry is not None: + per_shard = {} + for ti in coverage: + for coord in shard_coords_for_tiles([ti], grid): + per_shard[coord] = per_shard.get(coord, 0) + 1 + else: + per_shard = _shards_from_tiles(coverage, grid, transformer_cache) expected = set(per_shard) present = _existing_shards( diff --git a/tests/zarr.t b/tests/zarr.t index f82fc6b..c29b6de 100644 --- a/tests/zarr.t +++ b/tests/zarr.t @@ -74,6 +74,15 @@ over land rather than over the zone's bounding box: $ geotessera-registry zarr-scan --help | grep -o '\-\-output OUTPUT' | sort -u --output OUTPUT +The tile mirror is optional -- scanning a remote store needs only the +landmask registry for its land denominator: + + $ geotessera-registry zarr-scan --help | grep -c 'Optional tile mirror' + 1 + + $ geotessera-registry zarr-scan --help | grep -oE 'store_path \[base_dir\]' + store_path [base_dir] + Test: zarr-extend grows the time axis -------------------------------------- From 567f195878332a4bc4926291bf8f12c3dc2e8ff1 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Sat, 1 Aug 2026 08:01:22 +0100 Subject: [PATCH 08/13] let scan work with no local mirror Making the tile mirror optional also flipped zarr-scan's positionals to , which no other command takes. A script passing the usual therefore read the store URL as the tile source, went looking for a manifest inside the store, and signed the request with source credentials the caller had no reason to set -- surfacing as an unrelated NoCredentialsError. Restore with the source optional. argparse binds a lone positional to store_path, so both `zarr-scan ` and `zarr-scan ` do the obvious thing. A source that is really a store now says so, rather than failing later inside botocore. And _resolve_source sat outside the error handling in zarr-init, zarr-fill and zarr-scan, so any credential failure while resolving the manifest escaped as a traceback instead of the clean message. --- docs/architecture.rst | 2 +- geotessera/registry_cli.py | 60 +++++++++++++++++++++++++++++++++----- tests/zarr.t | 4 +-- 3 files changed, 55 insertions(+), 11 deletions(-) diff --git a/docs/architecture.rst b/docs/architecture.rst index bbb26de..2af505c 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -684,7 +684,7 @@ per-shard index as parquet. No tile mirror or manifest is needed: the land denominator comes from the landmask registry (~19 MB, fetched and cached), which is all that decides whether a shard can ever hold data. Supplying a tile mirror as an optional -second argument switches the denominator to each year's actual embedding +first argument switches the denominator to each year's actual embedding coverage from the manifest (~200 MB) — worth it only where a year covers less than the full land area. diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index cf91518..022925e 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3555,6 +3555,26 @@ def _resolve_source(args, console: "Console"): base_dir = args.base_dir + # A Zarr store passed where the tile source belongs otherwise fails much + # later, looking for a manifest inside the store and often with the wrong + # credentials. Say what happened instead. + from .remote import exists as _loc_exists + + try: + looks_like_store = _loc_exists( + join(base_dir, "zarr.json"), + _storage_options_for(args, "source", base_dir), + on_denied=False, + ) + except Exception: + looks_like_store = False + if looks_like_store: + raise ValueError( + f"{base_dir} looks like a Zarr store, not a tile source — it has " + f"a zarr.json. The tile source comes first and the store second: " + f" ." + ) + def override(loc, name, optional=False): """Resolve an explicit --manifest-url / --landmasks-url, if given.""" if not loc: @@ -3644,7 +3664,13 @@ def zarr_init_command(args): console = Console() - registry, _source, _version, _variant = _resolve_source(args, console) + try: + registry, _source, _version, _variant = _resolve_source(args, console) + except ValueError as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) years = _parse_int_range(args.years) output = args.output @@ -3695,7 +3721,13 @@ def zarr_fill_command(args): console = Console() - registry, source, _version, _variant = _resolve_source(args, console) + try: + registry, source, _version, _variant = _resolve_source(args, console) + except ValueError as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) store_path = args.store_path store_options = _storage_options_for(args, "store", store_path) @@ -3747,7 +3779,15 @@ def zarr_scan_command(args): # The manifest is optional: land coverage alone says which shards can # ever hold data, and that comes from the much smaller landmask registry. if args.base_dir: - registry, source, dataset_version, _variant = _resolve_source(args, console) + try: + registry, source, dataset_version, _variant = _resolve_source( + args, console + ) + except ValueError as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) else: registry, source = None, None dataset_version = args.dataset_version or "v1" @@ -4795,11 +4835,10 @@ def main(): "bounding box. Prints per-zone/year and per-year summaries and can " "write the full index as parquet.", ) - zarr_scan_parser.add_argument( - "store_path", - type=str, - help="Path or URL of an existing tessera store", - ) + # Same positional order as zarr-init/zarr-fill — source first, store + # second — with the source optional. argparse binds a lone argument to + # store_path, so `zarr-scan ` and `zarr-scan ` + # both do the obvious thing. zarr_scan_parser.add_argument( "base_dir", nargs="?", @@ -4809,6 +4848,11 @@ def main(): "needed and far smaller than the manifest. Give it to measure " "against each year's actual embedding coverage instead.", ) + zarr_scan_parser.add_argument( + "store_path", + type=str, + help="Path or URL of an existing tessera store", + ) zarr_scan_parser.add_argument( "--years", default=None, diff --git a/tests/zarr.t b/tests/zarr.t index c29b6de..101f73f 100644 --- a/tests/zarr.t +++ b/tests/zarr.t @@ -80,8 +80,8 @@ landmask registry for its land denominator: $ geotessera-registry zarr-scan --help | grep -c 'Optional tile mirror' 1 - $ geotessera-registry zarr-scan --help | grep -oE 'store_path \[base_dir\]' - store_path [base_dir] + $ geotessera-registry zarr-scan --help | grep -oE '\[base_dir\] store_path' + [base_dir] store_path Test: zarr-extend grows the time axis -------------------------------------- From 063e7f7b19e8ddb66be2b636335c801dc1da382a Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Sat, 1 Aug 2026 08:23:31 +0100 Subject: [PATCH 09/13] scan before fill The ingestion registry is written only when a (zone, year) finishes, so a run killed partway loses that year's bookkeeping and re-uploads everything. For a zone that is 97% done that means rebuilding 1,398 shards to add 48. The shard objects survive anything, and a shard is always written from every tile covering it, so their presence is proof of completeness. Scan for them before doing any work and skip what is already there. Where the credentials cannot list the store, fall back to probing just the shards the run would touch. --rewrite-existing-shards forces a rebuild, which is needed only when the tile inventory has grown, since a newly-added tile falls inside an existing shard and would otherwise be skipped rather than merged in. --skip-existing-shards stays accepted as a no-op so scripts keep working. Also exit cleanly on Ctrl-C. The pool was shut down with wait=True, so an interrupt blocked until every in-flight shard finished; with workers holding 2 GiB each that looks like a hang and invites a second Ctrl-C and a second traceback. Cancel what has not started, release the zone lock, and exit 130. --- CHANGES.md | 24 +++++++---- docs/architecture.rst | 21 ++++++---- geotessera/registry_cli.py | 23 ++++++++--- geotessera/zarr.py | 84 +++++++++++++++++++++----------------- tests/zarr.t | 13 ++++-- 5 files changed, 101 insertions(+), 64 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 304f703..13a1116 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -61,13 +61,23 @@ landmask registry (~19 MB, cached), so no tile mirror or manifest is needed. An optional tile mirror switches the denominator to each year's actual embedding coverage from the manifest. (@avsm) -- **`zarr-fill --skip-existing-shards`**: Resume from the store itself. The - ingestion registry is only written when a (zone, year) finishes, so a run - killed partway loses that year's bookkeeping even though its shards are - safely written; the shard objects survive anything. This scans for them, - skips them and records their tiles. Assumes the tile inventory has not - grown since. Falls back to probing just the shards in hand where the - credentials cannot list the store. (@avsm) +- **`zarr-fill` scans before writing and uploads only what is missing.** + The ingestion registry is written only when a (zone, year) finishes, so a + run killed partway loses that year's bookkeeping and would re-upload + everything — for a zone that is 97% done, rebuilding 1,398 shards to add + 48. The shard objects survive anything and a shard is always written from + every tile covering it, so their presence is proof of completeness. Fills + now scan for them by default and skip what is there; falls back to + probing just the shards in hand where the credentials cannot list the + store. `--rewrite-existing-shards` forces a rebuild, needed only when the + tile inventory has grown, since a newly-added tile falls inside an + existing shard. `--skip-existing-shards` is still accepted as a no-op. + (@avsm) +- **Ctrl-C during a fill exits cleanly** with status 130 instead of two + tracebacks. The process pool was shut down with `wait=True`, so an + interrupt blocked until every in-flight shard finished — with + multi-gigabyte workers that looks like a hang and invites a second + Ctrl-C. (@avsm) - **`zarr-fill` warns when the worker count will not fit in RAM**: each worker holds a ~2.1 GiB shard buffer, so `--workers 16` needs 33 GiB and an OOM kill leaves no traceback to diagnose. (@avsm) diff --git a/docs/architecture.rst b/docs/architecture.rst index 2af505c..a52d4dc 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -657,17 +657,20 @@ The ingestion registry is written when a (zone, year) finishes, so a run that dies partway — an OOM kill leaves no traceback — loses that year's bookkeeping even though the shards it wrote are safely in the store. -The shard objects are the ground truth, and they survive anything:: +The shard objects are the ground truth and they survive anything, so a fill +scans for them before doing any work and skips what is already there. That +is the default: re-running an interrupted fill uploads only what is missing. + +A shard is always written from every tile covering it, so its presence means +it is complete. The exception is a tile inventory that has grown since — +a newly-added tile falls inside an existing shard, which would then be +skipped rather than merged in. Force those shards to be rebuilt with:: geotessera-registry zarr-fill --zones 30 \ - --skip-existing-shards - -This lists the shard objects already present for each (zone, year), skips -them, and records their tiles so subsequent runs need no flag. It assumes -the tile inventory has not grown since those shards were written: a tile -added to the manifest afterwards falls inside an existing shard and would -be skipped rather than merged in. Where the credentials cannot list the -store, it falls back to probing only the shards the run would touch. + --rewrite-existing-shards + +Where the credentials cannot list the store, the scan falls back to probing +only the shards the run would touch. To see what is outstanding before committing to a sweep, point ``zarr-scan`` at the store on its own:: diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 022925e..7f62457 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3753,7 +3753,7 @@ def zarr_fill_command(args): consolidate=consolidate, force_lock=args.force_lock, state_url=args.state_url, - skip_existing_shards=args.skip_existing_shards, + skip_existing_shards=not args.rewrite_existing_shards, ) except RuntimeError as e: console.print(f"[red]{emoji('❌ ')}{e}[/red]") @@ -4807,10 +4807,15 @@ def main(): zarr_fill_parser.add_argument( "--skip-existing-shards", action="store_true", - help="Treat a shard that is already in the store as done. Recovers a " - "run killed before it recorded progress, since the shard objects " - "outlive the bookkeeping. Assumes the tile inventory has not grown " - "since those shards were written.", + help="Deprecated and now the default; accepted so existing scripts " + "keep working.", + ) + zarr_fill_parser.add_argument( + "--rewrite-existing-shards", + action="store_true", + help="Rebuild shards that are already in the store instead of " + "skipping them. Needed only when the tile inventory has grown, since " + "a newly-added tile falls inside an existing shard.", ) zarr_fill_parser.add_argument( "--force-lock", @@ -5164,7 +5169,13 @@ def main(): # Execute the command. Commands return a shell exit status; propagate it # so a failed zone in a parallel sweep is visible to the orchestrator # rather than silently reported as success. - raise SystemExit(args.func(args) or 0) + try: + raise SystemExit(args.func(args) or 0) + except KeyboardInterrupt: + # A long fill is routinely interrupted; report it as the shell + # convention rather than as two pages of traceback. + console.print(f"\n[yellow]{emoji('⚠️ ')}Interrupted.[/yellow]") + raise SystemExit(130) if __name__ == "__main__": diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 86fe5a1..91df039 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -2316,7 +2316,7 @@ def fill_store( consolidate: Optional[bool] = None, force_lock: bool = False, state_url: Optional[str] = None, - skip_existing_shards: bool = False, + skip_existing_shards: bool = True, ) -> int: """Incrementally fill a store with tile data. @@ -2335,12 +2335,15 @@ def fill_store( is set, because the root object is the one thing parallel zone jobs share — run ``zarr-consolidate`` once after the sweep. force_lock: Take over a (zone, year) lock held by another process. - skip_existing_shards: Treat a shard object that already exists as - done. Recovers a run killed before it could record progress, - since the objects outlive the bookkeeping. Assumes the tile - inventory has not grown since those shards were written — a tile - added to the manifest afterwards falls inside an existing shard - and would be skipped rather than merged in. + skip_existing_shards: Scan for shards already in the store and skip + them (the default). A shard is always written from every tile + covering it, so its presence means it is complete, and the + objects outlive the ingestion registry — which makes this both + the cheapest resume and the only one that survives a kill -9. + Set False to rebuild them, which is needed only when the tile + inventory has grown: a tile added to the manifest afterwards + falls inside an existing shard and would otherwise be skipped + rather than merged in. """ store = StoreLocation.resolve(store_path, storage_options, state_url) if workers is None: @@ -2598,40 +2601,45 @@ def _drain(pool, advance=None): if advance is not None: advance() - if console: - from rich.progress import ( - Progress, - BarColumn, - TextColumn, - MofNCompleteColumn, - TimeElapsedColumn, - TimeRemainingColumn, - SpinnerColumn, - ) + # Not a `with` block: on Ctrl-C the context manager's shutdown(wait=True) + # blocks until every in-flight shard finishes, which with a pool of + # multi-gigabyte workers looks like a hang and invites a second Ctrl-C + # (and a second traceback). Cancel what has not started and leave. + pool = ProcessPoolExecutor( + max_workers=workers, + initializer=_init_shard_worker, + initargs=initargs, + ) + try: + if console: + from rich.progress import ( + Progress, + BarColumn, + TextColumn, + MofNCompleteColumn, + TimeElapsedColumn, + TimeRemainingColumn, + SpinnerColumn, + ) - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - MofNCompleteColumn(), - TimeElapsedColumn(), - TimeRemainingColumn(), - console=console, - ) as progress: - task = progress.add_task(label, total=len(shard_specs)) - with ProcessPoolExecutor( - max_workers=workers, - initializer=_init_shard_worker, - initargs=initargs, - ) as pool: + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TimeElapsedColumn(), + TimeRemainingColumn(), + console=console, + ) as progress: + task = progress.add_task(label, total=len(shard_specs)) _drain(pool, advance=lambda: progress.advance(task)) - else: - with ProcessPoolExecutor( - max_workers=workers, - initializer=_init_shard_worker, - initargs=initargs, - ) as pool: + else: _drain(pool) + except KeyboardInterrupt: + pool.shutdown(wait=False, cancel_futures=True) + raise + else: + pool.shutdown(wait=True) return written_count, failed diff --git a/tests/zarr.t b/tests/zarr.t index 101f73f..b0cc263 100644 --- a/tests/zarr.t +++ b/tests/zarr.t @@ -53,11 +53,16 @@ A dead sibling job's lock can be taken over explicitly: Test: resume from the store itself ----------------------------------- -Shard objects outlive the bookkeeping, so a killed run can be recovered by -scanning for them: +Shard objects outlive the bookkeeping, so a fill scans for them and skips +what is already there. Rebuilding is the opt-in: - $ geotessera-registry zarr-fill --help | grep -o '\-\-skip-existing-shards' | sort -u - --skip-existing-shards + $ geotessera-registry zarr-fill --help | grep -o '\-\-rewrite-existing-shards' | sort -u + --rewrite-existing-shards + +The old flag stays accepted so existing scripts keep working: + + $ geotessera-registry zarr-fill --help | grep -c 'now the default' + 1 Test: zarr-scan reports outstanding work ----------------------------------------- From 701c3c5b5cefe1d369e03dd97d754d070642c7c1 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Sat, 1 Aug 2026 10:08:18 +0100 Subject: [PATCH 10/13] fix deadlocks --- CHANGES.md | 34 +++++++++++++++-- geotessera/registry_cli.py | 6 +++ geotessera/remote.py | 77 ++++++++++++++++++++++++++++++++++++++ geotessera/zarr.py | 74 ++++++++++++++++++++++++------------ 4 files changed, 165 insertions(+), 26 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 13a1116..ebad80a 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -73,14 +73,42 @@ tile inventory has grown, since a newly-added tile falls inside an existing shard. `--skip-existing-shards` is still accepted as a no-op. (@avsm) +- **`zarr-fill` reports the shard arithmetic**: `Shards: 1,373 land, 48 + recorded done, 43 found in store, 1,282 to write`. The previous line + showed only the last figure, which could not be reconciled with + `zarr-scan` — that counts every land shard, whereas a fill considers only + those covering tiles the registry has not already recorded. (@avsm) +- **Object-store libraries no longer log over the progress bar**: botocore + logs "Found credentials in shared credentials file" at INFO every time a + client is built, once per worker process, and the workers share a + terminal with the progress bar. botocore, boto3, aiobotocore, s3fs, + urllib3 and aiohttp are now capped at WARNING in both the parent and the + workers. (@avsm) +- **Fills no longer deadlock against an object store.** `s3fs` runs its + client on a background event-loop thread and `fsspec` caches both the loop + and its filesystem instances globally; a forked worker inherits those + objects but not the thread running the loop, so its first call to the + store waits forever (main thread in `futex_do_wait`, loop thread idle in + `ep_poll`). Workers are now started with "spawn" rather than the Linux + default of "fork", and reset any inherited fsspec state on startup. + (@avsm) +- **Workers die with their parent.** A fill killed outright left its + workers running, each holding gigabytes and accumulating across runs — + 19 orphans holding 8 GB were observed on one host. Workers now set + `PR_SET_PDEATHSIG` on Linux. (@avsm) - **Ctrl-C during a fill exits cleanly** with status 130 instead of two tracebacks. The process pool was shut down with `wait=True`, so an interrupt blocked until every in-flight shard finished — with multi-gigabyte workers that looks like a hang and invites a second Ctrl-C. (@avsm) -- **`zarr-fill` warns when the worker count will not fit in RAM**: each - worker holds a ~2.1 GiB shard buffer, so `--workers 16` needs 33 GiB and - an OOM kill leaves no traceback to diagnose. (@avsm) +- **`zarr-fill` warns when the worker count will not fit in memory**, using + `MemAvailable` rather than physical RAM since these hosts are usually + shared. The estimate covers the real peak — a worker holds a 2.1 GiB + shard buffer, then zarr's sharding codec compresses every inner chunk and + assembles the shard while s3fs holds the upload body, measured at 4.3 GiB + and observed being OOM-killed at 12 GiB — so it budgets ~6.2 GiB each + rather than the raw buffer. An OOM kill leaves no traceback, so the + warning is the only diagnosis. (@avsm) - **`geotessera-registry zarr-extend`**: New subcommand that appends years to an existing store's time axis, so a new year can be added without rebuilding. Time is chunked one year per chunk, making this a diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 7f62457..3a01811 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -4425,6 +4425,12 @@ def main(): else [logging.StreamHandler()], ) + # The object-store libraries log at INFO on every client build; left on, + # they interleave with the progress bars during a fill. + from .remote import quieten_dependency_logging + + quieten_dependency_logging() + parser = argparse.ArgumentParser( description="GeoTessera Registry Management Tool - Generate and maintain Pooch registry files", formatter_class=argparse.RawDescriptionHelpFormatter, diff --git a/geotessera/remote.py b/geotessera/remote.py index f1b85e6..a930c3c 100644 --- a/geotessera/remote.py +++ b/geotessera/remote.py @@ -26,6 +26,7 @@ import json import logging import os +import sys from functools import lru_cache from pathlib import Path from typing import Any, Dict, Optional, Tuple @@ -39,6 +40,82 @@ _URL_MARKER = "://" +# Chatty at INFO and of no interest to a user watching a fill. botocore in +# particular logs "Found credentials in shared credentials file" every time +# a client is built — once per worker process, which shreds a progress bar +# the workers share a terminal with. +_NOISY_LOGGERS = ( + "botocore", + "boto3", + "aiobotocore", + "s3fs", + "urllib3", + "aiohttp", +) + + +def quieten_dependency_logging(level: int = logging.WARNING) -> None: + """Raise the log level of the object-store libraries. + + Call in the parent before a progress bar starts, and again in each worker + process, since a worker that did not inherit the parent's configuration + would otherwise start logging into the shared terminal. + """ + for name in _NOISY_LOGGERS: + logging.getLogger(name).setLevel(level) + + +def reset_after_fork() -> None: + """Drop inherited fsspec/asyncio state in a freshly started worker. + + ``s3fs`` drives its client from a background event-loop thread and + ``fsspec`` caches both the loop and its filesystem instances globally. A + forked child inherits those objects but not the thread running the loop, + so the next call deadlocks: the worker's main thread waits on a future + the loop will never service (main in ``futex_do_wait``, loop thread idle + in ``ep_poll``). Workers are started with "spawn" to avoid this outright; + this is the belt-and-braces reset for anything that still leaks in. + """ + try: + import fsspec.asyn + + fsspec.asyn.reset_lock() + fsspec.asyn.iothread[0] = None + fsspec.asyn.loop[0] = None + except Exception as e: # pragma: no cover - best effort + logger.debug(f"Could not reset fsspec event loop: {e}") + + try: + import fsspec + + fsspec.AbstractFileSystem.clear_instance_cache() + except Exception as e: # pragma: no cover - best effort + logger.debug(f"Could not clear fsspec instance cache: {e}") + + _filesystem_cached.cache_clear() + + +def die_with_parent() -> None: + """Ask the kernel to kill this process when its parent goes away. + + Without it, a fill that is killed outright leaves its workers running: + each holds gigabytes and keeps writing, and they accumulate across runs + until the machine is out of memory. Linux only; a no-op elsewhere. + """ + if sys.platform != "linux": + return + try: + import ctypes + import signal + + PR_SET_PDEATHSIG = 1 + ctypes.CDLL("libc.so.6", use_errno=True).prctl( + PR_SET_PDEATHSIG, signal.SIGKILL + ) + except Exception as e: # pragma: no cover - best effort + logger.debug(f"Could not set parent-death signal: {e}") + + def is_url(loc: str | Path) -> bool: """True if *loc* is an fsspec URL rather than a local filesystem path.""" return _URL_MARKER in str(loc) diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 91df039..90a937c 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -945,9 +945,28 @@ def _coarsen_tile( # reason the worker count is bounded by RAM rather than by cores. WORKER_BUFFER_BYTES = N_BANDS * SHARD_SIZE * SHARD_SIZE + 4 * SHARD_SIZE * SHARD_SIZE +# Peak is several times that buffer, and planning against the raw figure +# gets fills OOM-killed. On top of the buffer, zarr's sharding codec +# compresses every inner chunk and assembles the shard, and s3fs holds the +# upload body. Measured 4.3 GiB per worker on a sparse three-tile shard; a +# dense one on a loaded host was killed holding 12 GiB. +WORKER_PEAK_BYTES = 3 * WORKER_BUFFER_BYTES + def _total_memory_bytes() -> Optional[int]: - """Physical RAM, or None where it cannot be determined.""" + """Memory a fill can realistically use, or None if undeterminable. + + Prefers MemAvailable over physical RAM: these hosts are often shared, + and what is free right now is what decides whether the kernel starts + killing workers. + """ + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) * 1024 + except (OSError, ValueError, IndexError): + pass try: return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") except (ValueError, OSError, AttributeError): @@ -960,20 +979,22 @@ def _warn_worker_memory(workers: int, console=None) -> None: A fill that is OOM-killed leaves no traceback, so the cause is easy to miss; say it up front instead. """ - needed = workers * WORKER_BUFFER_BYTES + needed = workers * WORKER_PEAK_BYTES total = _total_memory_bytes() gib = 2**30 message = ( - f"{workers} workers need at least {needed / gib:.0f} GiB " - f"({WORKER_BUFFER_BYTES / gib:.1f} GiB of shard buffer each)" + f"{workers} workers can peak around {needed / gib:.0f} GiB " + f"(~{WORKER_PEAK_BYTES / gib:.1f} GiB each: a " + f"{WORKER_BUFFER_BYTES / gib:.1f} GiB shard buffer plus compression " + f"and the upload body)" ) if total is None: logger.info(message) return if needed > 0.8 * total: - safe = max(1, int(0.8 * total // WORKER_BUFFER_BYTES)) + safe = max(1, int(0.8 * total // WORKER_PEAK_BYTES)) text = ( - f"{message}, but this machine has {total / gib:.0f} GiB. " + f"{message}, but only {total / gib:.0f} GiB is available here. " f"The fill will likely be OOM-killed — consider --workers {safe}." ) if console: @@ -1724,6 +1745,12 @@ def _init_shard_worker( """ global _worker_store, _worker_source_options + from . import remote + + remote.quieten_dependency_logging() + remote.reset_after_fork() + remote.die_with_parent() + _worker_store = StoreLocation(store_url, store_options).open_group( mode="r+", path=zone_group, zarr_format=3 ) @@ -2479,26 +2506,20 @@ def fill_store( shard_specs = [ s for s in shard_specs if (s.sr, s.sc) not in present ] - if console: - console.print( - f" [cyan]{len(skipped_specs)} shard(s) already " - f"written, skipping[/cyan]" - ) if console: - pending = {(ti.lon, ti.lat) for ti in remaining} - rebuilt = { - (ov.embedding_path) for s in shard_specs for ov in s.tiles - } - n_rebuilt = max(0, len(rebuilt) - len(pending)) - extra = ( - f", {n_rebuilt} already-written tile(s) rebuilt" - if n_rebuilt > 0 - else "" - ) + # Spell the arithmetic out. The count of shards to write is + # otherwise hard to reconcile with zarr-scan, which counts + # every land shard, whereas a fill only considers those + # covering tiles the registry has not already recorded. + n_land = len(shard_coords_for_tiles(tile_infos, grid)) + n_recorded = n_land - len(touched) console.print( - f" {len(shard_specs)} shards to write " - f"({workers} workers{extra})" + f" Shards: {n_land:,} land, " + f"{n_recorded:,} recorded done, " + f"{len(skipped_specs):,} found in store, " + f"[bold]{len(shard_specs):,} to write[/bold] " + f"({workers} workers)" ) _acquire_zone_lock(store, zone_num, fill_year, force=force_lock) @@ -2579,12 +2600,18 @@ def _write_shards( Returns (shards written, set of (sr, sc) that failed). """ + import multiprocessing from concurrent.futures import ProcessPoolExecutor, as_completed written_count = 0 failed: set = set() initargs = (store.url, zone_group, store.storage_options, source_options) + # "spawn", not the Linux default of "fork": a forked worker inherits the + # parent's fsspec event-loop object without the thread that runs it, and + # deadlocks the first time it talks to the store. + mp_context = multiprocessing.get_context("spawn") + def _drain(pool, advance=None): nonlocal written_count futures = { @@ -2609,6 +2636,7 @@ def _drain(pool, advance=None): max_workers=workers, initializer=_init_shard_worker, initargs=initargs, + mp_context=mp_context, ) try: if console: From 5e5027226fc626013a69616147814f1db37fc195 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Sun, 2 Aug 2026 12:27:25 +0100 Subject: [PATCH 11/13] zarr stats gathering during fill Computing the global RGB stretch by re-reading shards costs terabytes; every pixel it needs was already in a fill worker's buffer exactly once. Collect per-(zone, year) statistics at fill time instead, as six ordinary arrays in each zone group (spec: docs/specs/zarr-stretch-stats.md): stretch_stats_{count,sum,prod} exact N / Sum(x) / Sum(x x^T), additive stretch_sample{,_scales,_count} weighted raw-pixel sample, int8+scale The sums are the sufficient statistics for the covariance, so the global PCA derived from them is exact -- verified |cos| = 1.0 per component against a full-population fit -- and additive across zones. Quantiles are not additive and PC-space projections need the global axes, hence the raw sample (~2.6 MiB per zone-year at K=20000), drawn per shard proportional to valid-pixel count and merged by weighted reservoir. Zone groups gain no attributes: the geoemb: convention keeps those on the root, so the filled slot count is itself an array. zarr-stretch defaults to aggregating these (a few MiB of reads, works on remote stores, seconds); --from-shards keeps the legacy sampling path. A drift check compares the stats covariance against one refitted from the stored sample -- relative Frobenius distance with a sqrt(d/n) noise-aware limit, after the specified per-component |cos| check false-alarmed on near-degenerate eigenvalues -- and flags double-counted rewrites. zarr-fill --backfill-stretch-stats rebuilds a zone's arrays from its shards (the repair for pre-feature stores and interrupted fills), and zarr-extend refuses to desynchronise a group that lacks them. Also included from the same working set: --spill-dir to memory-map shard buffers (RssAnon 1.73 -> 0.35 GiB per worker, for memory-tight VMs); --source/--store-path-style for S3-compatible endpoints without wildcard DNS; and --source-npy-root/--source-landmask-root for mirrors in the legacy flat layout. --- CHANGES.md | 15 + docs/specs/zarr-stretch-stats.md | 235 +++++++++ geotessera/registry_cli.py | 204 +++++++- geotessera/remote.py | 9 + geotessera/zarr.py | 794 +++++++++++++++++++++++++++++-- tests/zarr.t | 24 + tests/zarr_remote_check.py | 113 +++++ 7 files changed, 1329 insertions(+), 65 deletions(-) create mode 100644 docs/specs/zarr-stretch-stats.md diff --git a/CHANGES.md b/CHANGES.md index ebad80a..53da45f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -73,6 +73,21 @@ tile inventory has grown, since a newly-added tile falls inside an existing shard. `--skip-existing-shards` is still accepted as a no-op. (@avsm) +- **Per-zone stretch statistics, collected at fill time** (see + `docs/specs/zarr-stretch-stats.md`): each zone group gains six arrays — + exact mean/covariance sufficient statistics per (zone, year), additive + across zones, plus a weighted 20k-pixel sample for quantiles — folded in + by `zarr-fill` from the shard buffers it already holds, at no extra I/O. + `zarr-stretch` now aggregates these by default: a few MiB of reads and an + *exact* global PCA (verified |cos| = 1.0 against a full-population fit) + instead of terabytes of shard re-reads, and it works against remote + stores. A drift check compares the stats-derived covariance against one + refitted from the stored sample and warns when rewritten shards have + double-counted. `--from-shards` keeps the legacy path; + `zarr-fill --backfill-stretch-stats` rebuilds statistics for stores + filled before this existed (and is required before `zarr-extend` will + touch such stores); `zarr-init --stretch-sample-size` tunes the sample. + (@avsm) - **`zarr-fill` reports the shard arithmetic**: `Shards: 1,373 land, 48 recorded done, 43 found in store, 1,282 to write`. The previous line showed only the last figure, which could not be reconciled with diff --git a/docs/specs/zarr-stretch-stats.md b/docs/specs/zarr-stretch-stats.md new file mode 100644 index 0000000..b204189 --- /dev/null +++ b/docs/specs/zarr-stretch-stats.md @@ -0,0 +1,235 @@ +# Per-zone stretch statistics and the global preview pipeline + +**Status:** Implemented (except `zarr-global-preview` remote support and the +two-phase mode, which remain proposed). Deviations from the draft are marked +**[as built]**. +**Scope:** `geotessera-registry` subcommands `zarr-init`, `zarr-fill`, `zarr-extend`, `zarr-stretch`, `zarr-global-preview`, `zarr-consolidate` +**Store convention:** additions to the per-zone group layout of the Tessera Zarr v3 store + +## Motivation + +`zarr-stretch` today computes the global RGB stretch (percentiles, optional histogram-equalisation CDF, optional 3-component PCA over all 128 bands) by reading randomly chosen shards from every zone until roughly 2 M valid pixels have been sampled. Each shard is a `(1, 128, 4096, 4096)` int8 object — 2 GiB — and typically a few hundred shards must be fetched to reach the sample target. That is on the order of a terabyte of embedding reads to derive a stretch object that serialises to a few kilobytes of root attributes. Against a remote store the cost is prohibitive and the command does not currently work at all. + +The waste is structural: every one of those pixels was already resident in a fill worker's shard buffer, exactly once, at `zarr-fill` time. This specification moves stretch-statistic collection into the fill, so that the global stretch is derived from a few MiB of per-zone summaries and never re-reads embeddings. + +Secondary goals folded into the same change: + +- `zarr-stretch` and `zarr-global-preview` must work against remote (S3/fsspec) stores, as `zarr-init`/`zarr-fill` already do. +- `zarr-global-preview` resume markers currently live at `_preview/zone_{N}_done` in the build-state sibling; the migration below moves any legacy in-store markers there too, preserving the rule that the store contains only Zarr. +- The preview's parallelism constraints are documented precisely, with a safe default and an explicit opt-in two-phase mode. + +## Design + +### New arrays + +Five arrays are added to every zone group `utm{zz}/`. They are ordinary Zarr v3 arrays with dimension names, consolidated with everything else — not sidecar files. `T` is the length of the zone's `time` axis; `K` is the per-zone-year sample capacity (default 20 000); `B` = `N_BANDS` = 128. + +| Array | Shape | Dtype | Dimensions | Semantics | +|---|---|---|---|---| +| `stretch_stats_count` | `(T,)` | int64 | `time` | N: valid pixels contributed to the sums | +| `stretch_stats_sum` | `(T, 128)` | float64 | `time, band` | S: per-band sum of dequantised values | +| `stretch_stats_prod` | `(T, 128, 128)` | float64 | `time, band, band2` | M: sum of outer products xxᵀ | +| `stretch_sample` | `(T, K, 128)` | int8 | `time, sample, band` | raw sampled embedding vectors | +| `stretch_sample_scales` | `(T, K)` | float32 | `time, sample` | per-sample dequant scale | + +**[as built]** The number of filled slots is a sixth array, `stretch_sample_count` `(T,)` int64, not a zone-group attribute: the `geoemb:` convention keeps its attributes on the root group only, and zone groups deliberately carry nothing beyond `proj:`/`spatial:`. Unfilled slots are zero-valued with scale `+inf`, matching the existing "not yet filled" sentinel, so a reader that ignores the count still cannot mistake padding for data. + +Samples are stored in **source representation** — the int8 embedding vector plus its float32 scale — which is lossless with respect to the store itself and lets any downstream statistic be recomputed exactly as if the pixels had been read from `embeddings`/`scales`. + +### Semantics of the additive statistics + +For each (zone, year), over every valid pixel x = int8 vector × scale (float64 accumulation): + +- `count` N — number of valid pixels (scale finite, i.e. not NaN water, not +inf unfilled), +- `sum` S = Σ x (length 128), +- `prod` M = Σ x xᵀ (128 × 128). + +These are the sufficient statistics for mean and covariance and are additive across zones and across shards within a zone. Global aggregation is exact: + +``` +μ = ΣS / ΣN +Σcov = ΣM / ΣN − μ μᵀ +``` + +### Size budget + +Per (zone, year): 8 + 1 024 + 131 072 bytes of statistics (≈ 0.13 MiB) plus 20 000 × 132 bytes of sample (≈ 2.5 MiB) ≈ **2.6 MiB**. For 60 zones × 9 years the whole-world total is ≈ **1.4 GiB**, negligible against the multi-terabyte store. For one year across all 60 zones, the additive statistics alone are ≈ 8 MiB; statistics plus samples ≈ 160 MiB — a few seconds from S3 either way. + +Chunking: one chunk per year on the `time` axis for all five arrays (matching the existing one-year-per-chunk convention), so a (zone, year) update touches exactly one chunk per array and `zarr-extend` can append years as a metadata-only edit. + +### Creation: at init, not lazily + +The arrays are created by `zarr-init`, alongside `embeddings` and `scales`. Lazy creation by `zarr-fill` was rejected: fills run in parallel across zones but the *decision* to create arrays would race with `zarr-consolidate` snapshots and would leave stores in mixed states that every reader must handle. With init-time creation, a store either predates the feature entirely (handled by backfill, below) or has the full set. `zarr-extend` grows all five arrays' `time` axis when appending years, exactly as it does for `embeddings`. + +## Statistical soundness + +The design rests on three empirically verified facts (all measured on real Tessera tile data): + +1. **Sufficient statistics give exact PCA.** Eigenvectors computed from summed (N, S, M) match a full-population PCA at |cos| = 1.000000 per component — exact to float64 — whereas a 20 k-pixel sampled PCA achieves |cos| = 0.9977. Summation beats sampling because no information is discarded: the covariance *is* the sufficient statistic. Every valid pixel in the store contributes. + +2. **Quantiles are not additive**, and worse, the quantities we need quantiles of — the principal-component projections — do not exist until the global PCA axes are known, which happens only after cross-zone aggregation. Hence the raw 128-dimensional samples: at stretch time they are projected onto the freshly derived PCs and percentiles/CDF are computed in PC space from the pooled sample. + +3. **Sampled quantiles are accurate enough.** Measured worst-case error of sampled p2/p98 against the full population of a real tile, expressed as a fraction of the stretch span over 20 trials: 5 000 samples → 3.89 %, 20 000 → 1.90 %, 50 000 → 2.28 % (no further improvement past ~20 k; the residual is population tail noise). Default **K = 20 000** per (zone, year), giving a pooled global sample of up to 1.2 M pixels — comparable to today's 2 M target, but drawn once and stored. + +**Weighting.** Samples must be drawn proportionally to each shard's valid-pixel count. Uniform per-shard draws would over-represent sparse coastal shards (a shard with 500 valid pixels would carry the same weight as a full 16.7 M-pixel interior shard). The implementation uses a weighted reservoir over the zone's shards: each worker draws a per-shard sub-sample proportional to its valid count, and the parent merges reservoirs with weights, so the final K samples approximate a uniform draw over all valid pixels of the (zone, year). + +## Fill-time collection + +Collection is free in I/O terms: workers already hold every decoded pixel in the 2 GiB shard buffer. + +**Worker side.** After assembling a shard and before writing it, the worker computes the shard's (n, s, m) triple over its valid pixels and draws its weighted pixel sub-sample (int8 vectors + scales). This adds one 128×128 GEMM-style pass per shard — negligible against tile decode — and a bounded few MiB of memory, preserving the ~6.2 GiB (or ~4.1 GiB with `--spill-dir`) per-worker budget. Partial results are returned to the parent over the existing result queue. + +**Parent side.** The parent process owns the (zone, year) advisory lock in the state dir for the duration of the fill. It sums the (n, s, m) triples, merges the reservoirs, and — once the (zone, year) sweep completes — performs a **single write per (zone, year)** into the five arrays plus the `geoemb:stretch_sample_count` attribute. Each write touches one chunk per array (one-year chunking), so the update is as atomic as Zarr offers; the lock guarantees no concurrent writer for that (zone, year). + +**Cross-zone safety.** Parallel `zarr-fill --zones N` processes touch disjoint zone groups, disjoint locks, and never the root metadata, exactly as today. The stats arrays live inside the zone group, so the existing parallelism contract is unchanged. + +**Resumed fills.** A fill that scans the store and skips existing shards contributes stats only for the shards it actually writes, and *adds* them to the stored arrays (read-modify-write under the zone lock). Shards written by a previous crashed run therefore may or may not be represented; see Caveats. + +## Aggregation: the `zarr-stretch` fast path + +When the stats arrays are present (detected from consolidated metadata), `zarr-stretch` defaults to: + +1. Read `stretch_stats_{count,sum,prod}` for the requested year from every selected zone (≈ 8 MiB for the world; seconds against S3). +2. Sum: Nᵍ = ΣN, Sᵍ = ΣS, Mᵍ = ΣM. Compute μ = Sᵍ/Nᵍ and Σcov = Mᵍ/Nᵍ − μμᵀ in float64. +3. Eigendecompose Σcov; take the top `--pca-components` eigenvectors (for `--mode pca`; `--mode bands` skips PCA and uses bands 0–2, with percentiles still computed from the pooled sample). +4. Read `stretch_sample`/`stretch_sample_scales` for the year (≈ 150 MiB for the world), dequantise, pool across zones (each zone's sample already approximates a uniform draw of that zone; pooling weights zones by their filled sample counts), project into PC space. +5. Compute `--p-low`/`--p-high` percentiles and, unless `--no-equalise`, the per-channel CDF with `--breakpoints` breakpoints. +6. Run the **drift check** (see Caveats) and write the result to the root attribute `geoemb:stretch.{year}`. + +Total runtime: seconds to low minutes, dominated by S3 GETs of small chunks. This is a **single-writer** step (it rewrites root attributes): run it after all zone fills have finished and before any preview, never concurrently with `zarr-consolidate`. + +`--from-shards` forces the legacy shard-sampling mode for stores without stats (and remains useful as an independent cross-check). If stats arrays are absent the command falls back to legacy mode with a warning. + +## Caveats and failure modes + +**Double-counting under shard rewrites.** `--rewrite-existing-shards` (or a fill re-run after the tile manifest gained tiles) rebuilds shards whose pixels already contributed to (N, S, M), adding them again. Per-shard keying of contributions was considered and **rejected**: it multiplies storage and bookkeeping by the shard count for a second-order benefit. Instead the design **accepts the drift**: a duplicate shard perturbs a 128 × 128 covariance built from ~10⁹ pixels negligibly, and the stored sample provides an independent estimator. `zarr-stretch` therefore always runs a **drift check**. **[as built]** the metric is the relative Frobenius distance between the stats-derived and sample-derived covariances, not per-component |cos|: eigenvector comparison false-alarms whenever eigenvalues are close (the vectors are then arbitrary rotations), which testing surfaced immediately. The alarm limit is `max(--drift-threshold, 3·√(128/n_eff))` — the sample covariance's own noise floor scales as √(d/n), so a small sample must not read as drift. Above the limit it warns and recommends `--from-shards` or a stats backfill (`zarr-fill --backfill-stretch-stats`, below) to rebuild the arrays from the store's actual contents. The sample arrays themselves are *overwritten*, not summed, on rewrite — the reservoir merge replaces slots — so they do not accumulate duplicates the same way, which is what makes them a valid cross-check. + +**Stores initialised before this feature.** Their zone groups lack the arrays. `zarr-fill` gains `--backfill-stretch-stats`: for the selected zone(s) it creates the five arrays (under the zone lock; a zone-group edit, not a root edit, so it composes with parallel fills of *other* zones) and populates them by scanning that zone's existing shards once — this is the one path that does re-read embeddings, but it is per-zone, opt-in, and runs at full shard-streaming bandwidth. A subsequent `zarr-consolidate` publishes the new arrays. + +**Interrupted fills.** If the parent dies before the end-of-year write, the shards are in the store but their stats are not. A resumed fill skips those shards and never accounts for them. The drift check catches gross discrepancies; `--backfill-stretch-stats` repairs them exactly. + +**`zarr-extend`.** Must grow the `time` axis of all five stats arrays in the same metadata-only edit as `embeddings`/`scales`/`time`. An extend that predates awareness of these arrays would desynchronise time axes within a group; the implementation must fail loudly if any of the five arrays is missing or of unexpected length, directing the operator to `--backfill-stretch-stats`. + +**Preview marker migration.** Older `zarr-global-preview` runs wrote `.zone_{N}_done` markers inside the store. Markers now live at `/preview/{year}/zone_{N}_done`. On startup the command migrates any legacy in-store markers into the state dir and deletes them from the store, restoring the "store contains only Zarr" invariant. + +## CLI surface + +The running example is the source.coop deployment. Tile source: the public mirror `s3://tessera/tessera` via the gateway (`--source-endpoint-url https://data.source.coop --source-anon`), or in-region anonymous reads of the backing bucket `s3://us-west-2.opendata.source.coop/tessera/tessera` (`--source-anon --source-region us-west-2` — preferred for large fills, no gateway hop). Store: `s3://us-west-2.opendata.source.coop/tessera/tessera/zarr/v1`, written with credentials (`--store-profile sc-writer --store-region us-west-2 --store-acl bucket-owner-full-control`). Note there is **no** `--store-endpoint-url`: `data.source.coop` is a read-only gateway and writes go straight to the backing bucket. + +```sh +STORE=s3://us-west-2.opendata.source.coop/tessera/tessera/zarr/v1 +SRC=s3://us-west-2.opendata.source.coop/tessera/tessera +SRCFLAGS="--source-anon --source-region us-west-2" +STOREFLAGS="--store-profile sc-writer --store-region us-west-2 --store-acl bucket-owner-full-control" +STATE="--state-url s3://sc-build-state/tessera-v1.build" +``` + +### `zarr-init` — changed + +Unchanged interface; now also creates the five stats arrays (year-chunked, `K` from `--stretch-sample-size`, default 20 000, recorded in zone attrs) in every zone group. + +```sh +geotessera-registry zarr-init "$SRC" --years 2017-2025 --output "$STORE" \ + $SRCFLAGS $STOREFLAGS $STATE +``` + +### `zarr-fill` — changed + +Existing behaviour and flags (`--year`, `--zones`, `--workers`, `--spill-dir`, `--rewrite-existing-shards`, `--force-lock`, `--state-url`, source/store storage flags) unchanged. New behaviour: accumulates (N, S, M) and the weighted sample in workers, aggregates in the parent, writes once per (zone, year) at sweep end. New flags: + +| Flag | Meaning | +|---|---| +| `--no-stretch-stats` | Skip collection (escape hatch; leaves arrays untouched) | +| `--backfill-stretch-stats` | Create/rebuild the stats arrays for the selected zones by scanning existing shards; implies no tile ingestion | + +```sh +# one zone per VM / process; safe in parallel across zones +geotessera-registry zarr-fill "$SRC" "$STORE" --zones 30 --workers 4 \ + --spill-dir /scratch/spill $SRCFLAGS $STOREFLAGS $STATE +``` + +### `zarr-stretch` — changed + +New default: stats fast path when the arrays are present (detected per zone; zones lacking them are reported and skipped, or the run aborts under `--strict`). Existing flags (`--year`, `--p-low`, `--p-high`, `--no-equalise`, `--breakpoints`, `--mode`, `--pca-components`, `--pca-total-bands`, `--pca-rgb-order`, `--zones`) keep their meaning. `--target-samples`/`--max-shards`/`--workers` apply only to legacy mode. Gains store storage flags so it works remotely. New flags: + +| Flag | Meaning | +|---|---| +| `--from-shards` | Force the legacy shard-sampling path | +| `--drift-threshold` | Maximum relative Frobenius distance between stats- and sample-derived covariances before warning (default 0.25; the limit is never below 3× the sample noise floor) | + +```sh +geotessera-registry zarr-stretch "$STORE" --year 2024 --mode pca $STOREFLAGS +``` + +(`zarr-stretch` keeps no build state, so it takes no `--state-url`.) + +### `zarr-global-preview` — changed + +Gains store storage flags and `--state-url`; resume markers move to the state dir (with in-store legacy migration). Zone-level parallelism over the shared pyramid is **unsafe**: adjacent zones share level-0 edge chunks (read-modify-write composites) and coarse levels overlap almost totally. Two modes: + +- **Default (safe): sequential zones**, exactly today's semantics, now remote-capable. +- **Two-phase (opt-in):** phase 1, `--reproject-only --zones N`, writes a zone's *interior* level-0 chunks and stages its edge-column contributions in the state dir; runnable in parallel across zones. Phase 2, `--coarsen-only`, run once after **all** zone level-0 writes are complete: composites the staged edge columns into the shared level-0 edge chunks, then builds levels 1–9. Running `--coarsen-only` before every zone has finished produces a silently incomplete pyramid; the command refuses unless every selected zone's phase-1 marker exists (`--force` overrides). + +```sh +# safe default +geotessera-registry zarr-global-preview "$STORE" --year 2024 \ + --gamma 0.7 --saturation 1.8 $STOREFLAGS $STATE + +# opt-in two-phase +for z in $(seq 1 60); do + geotessera-registry zarr-global-preview "$STORE" --year 2024 \ + --reproject-only --zones $z $STOREFLAGS $STATE & # parallel, one per slot +done; wait +geotessera-registry zarr-global-preview "$STORE" --year 2024 \ + --coarsen-only $STOREFLAGS $STATE # single-writer +``` + +### `zarr-consolidate` — unchanged + +Single-writer; merges per-zone registries and rewrites root consolidated metadata (which now includes the stats arrays). Run when no fill or preview is in flight. + +## End-to-end runbook + +Fresh store, 2017–2025, source.coop deployment. Parallel-safety and cost annotated per step. + +1. **Init** — single-writer, minutes. + ```sh + geotessera-registry zarr-init "$SRC" --years 2017-2025 --output "$STORE" \ + $SRCFLAGS $STOREFLAGS $STATE + ``` +2. **Fill, one zone per slot** — parallel-safe across zones; hours per zone, memory-bound (~6.2 GiB/worker, ~4.1 GiB with `--spill-dir`). On one large box: + ```sh + parallel -j 8 geotessera-registry zarr-fill "$SRC" "$STORE" \ + --zones {} --workers 4 --spill-dir /scratch/spill \ + $SRCFLAGS $STOREFLAGS $STATE ::: $(seq 1 60) + ``` + Across VMs, give each VM a zone range (`--zones 1-8`, `--zones 9-16`, …). Fills are resumable: re-running skips existing shards. +3. **Consolidate** — single-writer, minutes; publishes filled arrays and stats. + ```sh + geotessera-registry zarr-consolidate "$STORE" $STOREFLAGS $STATE + ``` +4. **Stretch (fast path)** — single-writer (root attrs), seconds–minutes per year. + ```sh + for y in $(seq 2017 2025); do + geotessera-registry zarr-stretch "$STORE" --year $y --mode pca $STOREFLAGS + done + ``` +5. **Global preview** — sequential default (hours; the remaining sequential bottleneck) or the two-phase opt-in above (phase 1 parallel, phase 2 single-writer). + ```sh + geotessera-registry zarr-global-preview "$STORE" --year 2024 \ + --gamma 0.7 --saturation 1.8 $STOREFLAGS $STATE + ``` +6. **Final consolidate** — single-writer, minutes; publishes the pyramid. + ```sh + geotessera-registry zarr-consolidate "$STORE" $STOREFLAGS $STATE + ``` + +Appending a year later: `zarr-extend --years 2026` (single-writer, no fills in flight; grows all arrays including stats) → step 2 for the new year → steps 3–6. + +## Open questions + +1. **K tuning.** 20 k/zone-year is justified by the p2/p98 error measurements, but equalisation CDFs with 257 breakpoints may benefit from larger pooled samples. Should `--stretch-sample-size` be raised for stores intended primarily for equalised previews? +2. **Sample refresh policy on partial re-fills.** The reservoir-replacement rule keeps the sample duplicate-free but means a small re-fill can churn slots contributed by unrelated shards. Is slot churn acceptable, or should re-fills merge weighted against the stored `geoemb:stretch_sample_count`? +3. **Water/land class statistics.** N/S/M cover valid land pixels only. Per-class counts (water, unfilled) per (zone, year) would be nearly free to collect and useful for coverage dashboards — worth adding now while the array set is being defined? +4. **Preview edge-staging format.** ~~Phase 1 stages edge-column composites in the state dir; the serialisation is unspecified.~~ **[likely resolved]**: staging may be unnecessary. The contended region is only the chunk column containing each zone boundary meridian (~1–2 columns; a zone's data pixels stay within its own 6° even though its iterated rectangle bulges far wider), so only *immediate* neighbours conflict and the conflict graph is 2-colourable. Scheduling zones in two waves (odd zones in parallel, then even) gives 30-way parallelism with no concurrent writer per chunk and no staged state — only the single-writer `--coarsen-only` barrier survives. To verify at implementation time: whether resampling smear at extreme latitudes can push a zone's data more than half a chunk past its meridian, which would require a third wave. +5. **Drift-check threshold.** ~~0.99 |cos|~~ **[resolved as built]**: the metric changed to covariance Frobenius distance with a noise-floor-aware limit, after eigenvector comparison false-alarmed on near-isotropic data. The 0.25 default should still be revisited once real rewrite workloads exist. diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index 3a01811..af7eb81 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3288,6 +3288,21 @@ def _add_source_args(parser) -> None: "(default: auto-detected from base_dir and parents, or pulled from a " "URL base_dir)", ) + parser.add_argument( + "--source-npy-root", + type=str, + default=None, + help="Location holding {year}/grid__/ tile directories, for " + "a mirror that does not use the published /npy/ " + "layout", + ) + parser.add_argument( + "--source-landmask-root", + type=str, + default=None, + help="Location holding grid__.tiff landmasks, for a mirror " + "that does not use the published /landmasks/ layout", + ) parser.add_argument( "--manifest-url", type=str, @@ -3372,6 +3387,13 @@ def _add_storage_args(parser, prefix: str, label: str, writable: bool = False) - action="store_true", help=f"Send requester-pays headers to the {label.lower()}", ) + group.add_argument( + f"--{prefix}-path-style", + action="store_true", + help=f"Address the {label.lower()} as endpoint/bucket/key rather " + f"than bucket.endpoint/key. Needed by most S3-compatible servers " + f"(MinIO, Ceph, and any endpoint without wildcard DNS).", + ) if writable: from .remote import OBJECT_ACLS @@ -3457,6 +3479,7 @@ def opt(name): profile=opt("profile"), requester_pays=bool(opt("requester_pays")), acl=opt("acl"), + path_style=bool(opt("path_style")), ) @@ -3653,14 +3676,24 @@ def override(loc, name, optional=False): landmasks_registry_path=landmasks_path, ) - source = TileSource.for_url(base_dir, version_path, storage_options) + if args.source_npy_root or args.source_landmask_root: + # A mirror that lays tiles out its own way: take the roots verbatim + # rather than appending the published npy/ convention. + default = TileSource.for_url(base_dir, version_path, storage_options) + source = TileSource( + embeddings_root=args.source_npy_root or default.embeddings_root, + landmasks_root=args.source_landmask_root or default.landmasks_root, + storage_options=storage_options, + ) + else: + source = TileSource.for_url(base_dir, version_path, storage_options) return registry, source, dataset_version, dataset_variant def zarr_init_command(args): """Create an empty tessera store with time dimension.""" from rich.console import Console - from .zarr import init_store + from .zarr import STRETCH_SAMPLE_K, init_store console = Console() @@ -3701,6 +3734,7 @@ def zarr_init_command(args): console=console, storage_options=store_options, state_url=args.state_url, + stretch_sample_size=args.stretch_sample_size or STRETCH_SAMPLE_K, ) except FileExistsError as e: console.print(f"[red]Error:[/red] {e}") @@ -3721,6 +3755,47 @@ def zarr_fill_command(args): console = Console() + # Backfill mode scans the store itself — no tile source, no manifest. + # Accept the store as the only positional (argparse binds it to base_dir + # when store_path is absent). + if args.backfill_stretch_stats: + from .zarr import backfill_stretch_stats + + store_path = args.store_path or args.base_dir + if store_path is None: + console.print(f"[red]{emoji('❌ ')}No store given.[/red]") + return 1 + store_options = _storage_options_for(args, "store", store_path) + zones = _parse_int_range(args.zones) if args.zones else None + years = [args.year] if args.year else None + try: + n = backfill_stretch_stats( + store_path, + zones=zones, + years=years, + console=console, + storage_options=store_options, + state_url=args.state_url, + force_lock=args.force_lock, + ) + except (ValueError, RuntimeError) as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) + console.print( + f"{emoji('✅ ')}{n} (zone, year) statistics rebuilt. " + f"Run zarr-consolidate to publish them." + ) + return 0 + + if args.store_path is None: + console.print( + f"[red]{emoji('❌ ')}Usage: zarr-fill " + f"(the store-only form is for --backfill-stretch-stats).[/red]" + ) + return 1 + try: registry, source, _version, _variant = _resolve_source(args, console) except ValueError as e: @@ -3754,6 +3829,8 @@ def zarr_fill_command(args): force_lock=args.force_lock, state_url=args.state_url, skip_existing_shards=not args.rewrite_existing_shards, + spill_dir=args.spill_dir, + collect_stretch_stats=not args.no_stretch_stats, ) except RuntimeError as e: console.print(f"[red]{emoji('❌ ')}{e}[/red]") @@ -3944,33 +4021,63 @@ def zarr_stretch_command(args): """Compute a global cross-zone RGB stretch and persist it to the store.""" import warnings from rich.console import Console - from .zarr import compute_global_stretch warnings.filterwarnings("ignore", message="Object at .* is not recognized") console = Console() - if not _require_local_store(args.store_path, "zarr-stretch", console): - return 1 - store_path = Path(args.store_path) zones = _parse_int_range(args.zones) if args.zones else None - compute_global_stretch( - store_path=store_path, - year=args.year, - target_samples=args.target_samples, - max_shards=args.max_shards, - p_low=args.p_low, - p_high=args.p_high, - workers=args.workers, - zones=zones, - equalise=not args.no_equalise, - equalise_breakpoints=args.breakpoints, - mode=args.mode, - pca_components=args.pca_components, - pca_total_bands=args.pca_total_bands, - pca_rgb_order=args.pca_rgb_order, - console=console, - ) + if args.from_shards: + # Legacy shard-sampling path: re-reads embeddings, local stores only. + from .zarr import compute_global_stretch + + if not _require_local_store(args.store_path, "zarr-stretch", console): + return 1 + compute_global_stretch( + store_path=Path(args.store_path), + year=args.year, + target_samples=args.target_samples, + max_shards=args.max_shards, + p_low=args.p_low, + p_high=args.p_high, + workers=args.workers, + zones=zones, + equalise=not args.no_equalise, + equalise_breakpoints=args.breakpoints, + mode=args.mode, + pca_components=args.pca_components, + pca_total_bands=args.pca_total_bands, + pca_rgb_order=args.pca_rgb_order, + console=console, + ) + return 0 + + # Fast path: aggregate the per-zone statistics collected at fill time. + # A few MiB of reads; works against local and remote stores alike. + from .zarr import compute_stretch_from_stats + + store_options = _storage_options_for(args, "store", args.store_path) + try: + compute_stretch_from_stats( + args.store_path, + year=args.year, + zones=zones, + p_low=args.p_low, + p_high=args.p_high, + equalise=not args.no_equalise, + equalise_breakpoints=args.breakpoints, + mode=args.mode, + pca_components=args.pca_components, + pca_rgb_order=args.pca_rgb_order, + drift_threshold=args.drift_threshold, + console=console, + storage_options=store_options, + ) + except (ValueError, RuntimeError) as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) return 0 @@ -4754,6 +4861,13 @@ def main(): help="Output store path or URL (e.g. tessera.zarr, " "s3://bucket/tessera.zarr)", ) + zarr_init_parser.add_argument( + "--stretch-sample-size", + type=int, + default=None, + help="Per-(zone, year) capacity of the raw pixel sample kept for the " + "global stretch quantiles (default: 20000)", + ) _add_source_args(zarr_init_parser) _add_storage_args(zarr_init_parser, "source", "Tile source") _add_state_arg(zarr_init_parser) @@ -4778,6 +4892,8 @@ def main(): zarr_fill_parser.add_argument( "store_path", type=str, + nargs="?", + default=None, help="Path or URL of an existing tessera store", ) zarr_fill_parser.add_argument( @@ -4823,6 +4939,32 @@ def main(): "skipping them. Needed only when the tile inventory has grown, since " "a newly-added tile falls inside an existing shard.", ) + zarr_fill_parser.add_argument( + "--no-stretch-stats", + action="store_true", + help="Skip folding stretch statistics into the zone group as shards " + "are written (escape hatch; the global stretch then needs a backfill " + "or the legacy shard-sampling path)", + ) + zarr_fill_parser.add_argument( + "--backfill-stretch-stats", + action="store_true", + help="Rebuild the selected zones' stretch statistics by scanning " + "their existing shards, instead of ingesting tiles. The repair path " + "for stores filled before fill-time collection existed, interrupted " + "fills, and suspected double-counting. With this flag the tile " + "source is not read; `zarr-fill --backfill-stretch-stats` " + "works with the store as the only positional.", + ) + zarr_fill_parser.add_argument( + "--spill-dir", + type=str, + default=None, + help="Memory-map each worker's shard buffers under this directory " + "instead of holding them in RAM. Trades disk for roughly 2 GiB of " + "anonymous memory per worker, which is what the OOM killer targets — " + "worth it on a memory-tight box with spare disk.", + ) zarr_fill_parser.add_argument( "--force-lock", action="store_true", @@ -5120,6 +5262,22 @@ def main(): "(PC1->R, PC2->G, PC3->B). Use '213' to swap R and G (PC2->R, " "PC1->G, PC3->B), '321' to fully reverse, etc.", ) + zarr_stretch_parser.add_argument( + "--from-shards", + action="store_true", + help="Force the legacy shard-sampling path (re-reads embeddings; " + "local stores only). Default: aggregate the per-zone statistics " + "collected at fill time — a few MiB of reads, remote-capable.", + ) + zarr_stretch_parser.add_argument( + "--drift-threshold", + type=float, + default=0.25, + help="Maximum relative Frobenius distance between the stats-derived " + "and sample-derived covariances before warning of stale statistics " + "(default: 0.25)", + ) + _add_storage_args(zarr_stretch_parser, "store", "Store", writable=True) zarr_stretch_parser.set_defaults(func=zarr_stretch_command) # Verify-tile command diff --git a/geotessera/remote.py b/geotessera/remote.py index a930c3c..66e7e9d 100644 --- a/geotessera/remote.py +++ b/geotessera/remote.py @@ -165,6 +165,7 @@ def build_storage_options( profile: Optional[str] = None, requester_pays: bool = False, acl: Optional[str] = None, + path_style: bool = False, extra: Optional[Dict[str, Any]] = None, ) -> Optional[Dict[str, Any]]: """Assemble an fsspec storage-options dict for an S3-compatible endpoint. @@ -180,6 +181,11 @@ def build_storage_options( acl: Canned ACL to stamp on every object written, e.g. ``bucket-owner-full-control`` for a bucket owned by another account. s3fs filters it per operation, so reads are unaffected. + path_style: Address buckets as ``endpoint/bucket/key`` rather than + ``bucket.endpoint/key``. Most S3-compatible servers (MinIO, + Ceph, and anything on a hostname without wildcard DNS) need + this; botocore's default would try to resolve + ``.``. Returns ``None`` when nothing needs configuring, so callers can pass the result through to zarr unchanged. @@ -209,6 +215,9 @@ def build_storage_options( if requester_pays: options["requester_pays"] = True + if path_style: + options["config_kwargs"] = {"s3": {"addressing_style": "path"}} + if acl: if acl not in OBJECT_ACLS: raise ValueError( diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 90a937c..06295f2 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -52,6 +52,7 @@ import logging import math import os +import shutil from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple @@ -938,6 +939,10 @@ def _coarsen_tile( SHARD_SIZE = 4096 # spatial pixels per shard side INNER_CHUNK = 32 # spatial pixels per inner chunk side + +# Default per-(zone, year) capacity of the raw pixel sample kept for the +# global stretch quantiles (docs/specs/zarr-stretch-stats.md). +STRETCH_SAMPLE_K = 20_000 DEFAULT_WORKERS = 4 # fewer workers due to larger shard buffers (~2GB each) # Each shard worker holds a full (N_BANDS, SHARD_SIZE, SHARD_SIZE) int8 @@ -952,6 +957,11 @@ def _coarsen_tile( # dense one on a loaded host was killed holding 12 GiB. WORKER_PEAK_BYTES = 3 * WORKER_BUFFER_BYTES +# With --spill-dir the shard buffers are memory-mapped, so their pages are +# reclaimable page cache rather than anonymous memory. Measured on Linux, +# that takes the buffer's contribution to RssAnon from 1.73 GiB to 0.35 GiB. +WORKER_PEAK_BYTES_SPILLED = WORKER_PEAK_BYTES - WORKER_BUFFER_BYTES + def _total_memory_bytes() -> Optional[int]: """Memory a fill can realistically use, or None if undeterminable. @@ -973,29 +983,36 @@ def _total_memory_bytes() -> Optional[int]: return None -def _warn_worker_memory(workers: int, console=None) -> None: +def _warn_worker_memory(workers: int, console=None, spilled: bool = False) -> None: """Warn when the requested worker count cannot fit in RAM. A fill that is OOM-killed leaves no traceback, so the cause is easy to miss; say it up front instead. """ - needed = workers * WORKER_PEAK_BYTES + per_worker = WORKER_PEAK_BYTES_SPILLED if spilled else WORKER_PEAK_BYTES + needed = workers * per_worker total = _total_memory_bytes() gib = 2**30 message = ( f"{workers} workers can peak around {needed / gib:.0f} GiB " - f"(~{WORKER_PEAK_BYTES / gib:.1f} GiB each: a " - f"{WORKER_BUFFER_BYTES / gib:.1f} GiB shard buffer plus compression " - f"and the upload body)" + f"(~{per_worker / gib:.1f} GiB each" + + ( + ", shard buffers spilled to disk)" + if spilled + else f": a {WORKER_BUFFER_BYTES / gib:.1f} GiB shard buffer plus " + f"compression and the upload body)" + ) ) if total is None: logger.info(message) return if needed > 0.8 * total: - safe = max(1, int(0.8 * total // WORKER_PEAK_BYTES)) + safe = max(1, int(0.8 * total // per_worker)) + hint = "" if spilled else " (or --spill-dir to cut ~2 GiB per worker)" text = ( f"{message}, but only {total / gib:.0f} GiB is available here. " - f"The fill will likely be OOM-killed — consider --workers {safe}." + f"The fill will likely be OOM-killed — consider " + f"--workers {safe}{hint}." ) if console: console.print(f" [yellow]{text}[/yellow]") @@ -1215,6 +1232,7 @@ def init_store( console: Optional["rich.console.Console"] = None, storage_options: Optional[Dict[str, Any]] = None, state_url: Optional[str] = None, + stretch_sample_size: int = STRETCH_SAMPLE_K, ) -> str: """Create a tessera store with time dimension from the landmask registry. @@ -1327,7 +1345,7 @@ def init_store( f"{n_shards_x}x{n_shards_y} shards[/dim]" ) - _create_zone_group(grid, store) + _create_zone_group(grid, store, stretch_sample_size) # Nothing else is written into the store: ingestion tracking and locks # are build state and live in the state sibling, created on first fill. @@ -1351,6 +1369,7 @@ def init_store( def _create_zone_group( grid: UnifiedZoneGrid, store_location: StoreLocation, + stretch_sample_size: int = STRETCH_SAMPLE_K, ) -> "zarr.Group": """Create a zone group with empty (T, B, H, W) arrays.""" from zarr.codecs import BloscCodec @@ -1412,6 +1431,11 @@ def _create_zone_group( ) store[name][:] = data + # Per-zone stretch statistics, populated by zarr-fill (see + # docs/specs/zarr-stretch-stats.md). Plain arrays: zone groups carry no + # geoemb: attributes. + create_stretch_arrays(store, T, stretch_sample_size) + # Use geozarr-toolkit for proj: and spatial: convention metadata from geozarr_toolkit import create_geozarr_attrs @@ -1450,6 +1474,221 @@ def _create_zone_group( return store +# --------------------------------------------------------------------------- +# Stretch statistics (per-zone, collected at fill time) +# --------------------------------------------------------------------------- +# The global RGB stretch needs a mean/covariance (for PCA) and quantiles (for +# the percentile stretch and equalisation CDF) over every valid pixel in the +# store. Recomputing those by re-reading shards costs terabytes; instead each +# fill records, per (zone, year), the exact sufficient statistics for the +# covariance — which are additive across zones — plus a weighted raw-pixel +# sample for the quantiles, which are not. +# +# These live as ordinary arrays inside each zone group. Zone groups carry no +# geoemb: attributes (the convention keeps those on the root), so the filled +# sample count is itself an array rather than an attr. + +STRETCH_ARRAY_NAMES = ( + "stretch_stats_count", + "stretch_stats_sum", + "stretch_stats_prod", + "stretch_sample", + "stretch_sample_scales", + "stretch_sample_count", +) + + +def create_stretch_arrays(group: "zarr.Group", n_years: int, k: int) -> None: + """Create the per-zone stretch-statistics arrays in *group*. + + One chunk per year on the time axis, so a (zone, year) update touches + exactly one chunk per array and ``zarr-extend`` grows them the same way + it grows ``embeddings``. + """ + from zarr.codecs import BloscCodec + + comp = BloscCodec(cname="zstd", clevel=3) + T = n_years + specs = [ + ("stretch_stats_count", (T,), (1,), np.int64, 0, ["time"]), + ("stretch_stats_sum", (T, N_BANDS), (1, N_BANDS), np.float64, 0.0, + ["time", "band"]), + ("stretch_stats_prod", (T, N_BANDS, N_BANDS), (1, N_BANDS, N_BANDS), + np.float64, 0.0, ["time", "band", "band2"]), + ("stretch_sample", (T, k, N_BANDS), (1, k, N_BANDS), np.int8, + np.int8(0), ["time", "sample", "band"]), + # +inf matches the "land, no data" sentinel, so padding slots can + # never be mistaken for real pixels even by a reader that ignores + # stretch_sample_count. + ("stretch_sample_scales", (T, k), (1, k), np.float32, + np.float32("inf"), ["time", "sample"]), + ("stretch_sample_count", (T,), (1,), np.int64, 0, ["time"]), + ] + for name, shape, chunks, dtype, fill, dims in specs: + group.create_array( + name, + shape=shape, + chunks=chunks, + dtype=dtype, + fill_value=fill, + compressors=comp, + dimension_names=dims, + ) + + +def _shard_sample_cap(k_slots: int, n_shards: int) -> int: + """Per-shard sample size: a few times K spread over the shards. + + Oversampling by 4x gives the weighted merge enough candidates to + approximate a uniform draw without ballooning the result queue. + """ + return min(k_slots, max(64, -(-4 * k_slots // max(1, n_shards)))) + + +def shard_stretch_stats( + emb_buf: np.ndarray, + scales_buf: np.ndarray, + sample_cap: int, + seed: Optional[int] = None, + block: int = 262_144, +) -> Optional[Dict[str, Any]]: + """Exact (n, S, M) sufficient statistics plus a pixel sample for one shard. + + Works on the shard buffers the fill already holds: ``emb_buf`` is + ``(B, S, S)`` int8, ``scales_buf`` ``(S, S)`` float32. Valid pixels are + those with finite scale. The sum-of-products matrix is accumulated in + float64 from float32 block GEMMs — each block sums ~2.6e5 terms of O(1) + magnitude, so the block partials carry ~7 significant digits and the + float64 accumulation loses nothing that a covariance of 1e9 pixels could + show. + + Returns None when the shard has no valid pixels. + """ + valid = np.isfinite(scales_buf) + flat = np.flatnonzero(valid.ravel()) + n = int(flat.size) + if n == 0: + return None + + emb_flat = emb_buf.reshape(emb_buf.shape[0], -1) + scales_flat = scales_buf.ravel() + + s = np.zeros(N_BANDS, dtype=np.float64) + m = np.zeros((N_BANDS, N_BANDS), dtype=np.float64) + for i in range(0, n, block): + idx = flat[i : i + block] + xb = emb_flat[:, idx].astype(np.float32) * scales_flat[idx].astype(np.float32) + s += xb.sum(axis=1, dtype=np.float64) + m += (xb @ xb.T).astype(np.float64) + + rng = np.random.default_rng(seed) + k = min(sample_cap, n) + pick = flat[rng.choice(n, size=k, replace=False)] + return { + "n": n, + "sum": s, + "prod": m, + "sample_emb": np.ascontiguousarray(emb_flat[:, pick].T), # (k, B) int8 + "sample_scales": scales_flat[pick].astype(np.float32), + # Each returned row stands for n/k pixels of the shard's population. + "sample_weight": n / k, + } + + +def merge_stretch_samples( + candidates: List[Tuple[np.ndarray, np.ndarray, float]], + k: int, + seed: Optional[int] = None, +) -> Tuple[np.ndarray, np.ndarray]: + """Draw K rows from weighted candidate pools (Efraimidis–Spirakis). + + ``candidates`` is a list of ``(emb (n, B) int8, scales (n,) f32, weight + per row)``. Rows are selected with probability proportional to their + weight, without replacement, so pooling per-shard samples of different + coverage reproduces a uniform draw over the union population. + """ + embs = [c[0] for c in candidates if len(c[0])] + if not embs: + return ( + np.zeros((0, N_BANDS), dtype=np.int8), + np.zeros((0,), dtype=np.float32), + ) + emb = np.concatenate(embs, axis=0) + scales = np.concatenate([c[1] for c in candidates if len(c[0])], axis=0) + weights = np.concatenate( + [np.full(len(c[0]), max(c[2], 1e-12)) for c in candidates if len(c[0])] + ) + + if len(emb) <= k: + return emb, scales + + rng = np.random.default_rng(seed) + keys = rng.random(len(emb)) ** (1.0 / weights) + top = np.argpartition(keys, -k)[-k:] + return emb[top], scales[top] + + +def weighted_percentile( + values: np.ndarray, weights: np.ndarray, qs: np.ndarray +) -> np.ndarray: + """Percentiles of a weighted sample (qs in 0..100).""" + order = np.argsort(values, kind="stable") + v = values[order] + w = weights[order].astype(np.float64) + cdf = np.cumsum(w) - 0.5 * w + cdf /= w.sum() + return np.interp(np.asarray(qs, dtype=np.float64) / 100.0, cdf, v) + + +def update_zone_stretch_stats( + zone_group: "zarr.Group", + time_index: int, + n: int, + s: np.ndarray, + m: np.ndarray, + sample_candidates: List[Tuple[np.ndarray, np.ndarray, float]], + seed: Optional[int] = None, +) -> None: + """Fold one fill run's statistics into a zone's arrays (read-modify-write). + + The additive triple is summed onto what is stored; the sample is re-drawn + from the stored sample and the new candidates together, weighted so the + result still approximates a uniform draw over all pixels either has seen. + Caller must hold the (zone, year) fill lock — this is the same + single-writer context the shard writes ran under. + """ + t = time_index + count_arr = zone_group["stretch_stats_count"] + prev_n = int(count_arr[t]) + + count_arr[t] = prev_n + n + zone_group["stretch_stats_sum"][t] = ( + np.asarray(zone_group["stretch_stats_sum"][t]) + s + ) + zone_group["stretch_stats_prod"][t] = ( + np.asarray(zone_group["stretch_stats_prod"][t]) + m + ) + + k = zone_group["stretch_sample"].shape[1] + stored_k = int(zone_group["stretch_sample_count"][t]) + pool = list(sample_candidates) + if stored_k > 0: + pool.append( + ( + np.asarray(zone_group["stretch_sample"][t, :stored_k]), + np.asarray(zone_group["stretch_sample_scales"][t, :stored_k]), + max(prev_n, 1) / stored_k, + ) + ) + emb, scales = merge_stretch_samples(pool, k, seed=seed) + + filled = len(emb) + if filled: + zone_group["stretch_sample"][t, :filled] = emb + zone_group["stretch_sample_scales"][t, :filled] = scales + zone_group["stretch_sample_count"][t] = filled + + # --------------------------------------------------------------------------- # Tile registry (GeoParquet tracking which tiles are written) # --------------------------------------------------------------------------- @@ -1730,6 +1969,8 @@ def _release_zone_lock(store: StoreLocation, zone: int, year: int) -> None: _worker_store = None _worker_source_options: Optional[Dict[str, Any]] = None +_worker_spill_dir: Optional[str] = None +_worker_sample_cap: int = 0 # 0 = stats collection off def _init_shard_worker( @@ -1737,13 +1978,16 @@ def _init_shard_worker( zone_group: str, store_options: Optional[Dict[str, Any]] = None, source_options: Optional[Dict[str, Any]] = None, + spill_dir: Optional[str] = None, + sample_cap: int = 0, ) -> None: """Process pool initializer: open the zone group once per worker. Both option dicts are plain picklable mappings, so a worker rebuilds its own filesystem connections rather than inheriting an unforkable client. """ - global _worker_store, _worker_source_options + global _worker_store, _worker_source_options, _worker_spill_dir + global _worker_sample_cap from . import remote @@ -1755,30 +1999,85 @@ def _init_shard_worker( mode="r+", path=zone_group, zarr_format=3 ) _worker_source_options = source_options + _worker_spill_dir = spill_dir + _worker_sample_cap = sample_cap def _write_one_shard( spec: ShardSpec, store: "zarr.Group", source_options: Optional[Dict[str, Any]] = None, -) -> bool: + spill_dir: Optional[str] = None, + sample_cap: int = 0, +) -> "bool | Dict[str, Any]": """Write one shard in NCHW layout: (T, B, H, W). Tile reads go through :mod:`geotessera.remote`, so ``spec`` may reference either local paths or remote URLs — a remote tile costs one ranged GET for the rows this shard needs, not the whole 150 MB object. """ + S = SHARD_SIZE + + # Allocate BHW buffer (bands-first for NCHW write). With a spill + # directory the two buffers are memory-mapped instead of anonymous: the + # pages become reclaimable page cache, so the kernel evicts them under + # pressure rather than the OOM killer taking the whole worker. Measured + # on Linux, the buffer's contribution to RssAnon drops from 1.73 GiB to + # 0.35 GiB. + spill = _open_spill(spill_dir) + if spill is None: + emb_buf = np.zeros((N_BANDS, S, S), dtype=np.int8) + # Start with +inf (land/nodata); landmask sets water to NaN, + # valid tiles overwrite with finite scales. + scales_buf = np.full((S, S), np.float32("inf")) + else: + emb_buf = np.memmap( + spill / "emb.buf", dtype=np.int8, mode="w+", shape=(N_BANDS, S, S) + ) + scales_buf = np.memmap( + spill / "scales.buf", dtype=np.float32, mode="w+", shape=(S, S) + ) + scales_buf[:] = np.float32("inf") + + try: + return _fill_and_write_shard( + spec, store, emb_buf, scales_buf, source_options, sample_cap + ) + finally: + del emb_buf, scales_buf + if spill is not None: + shutil.rmtree(spill, ignore_errors=True) + + +def _open_spill(spill_dir: Optional[str]): + """Create a per-shard scratch directory, or None to stay in memory.""" + if not spill_dir: + return None + import tempfile + + Path(spill_dir).mkdir(parents=True, exist_ok=True) + return Path(tempfile.mkdtemp(prefix="shard-", dir=spill_dir)) + + +def _fill_and_write_shard( + spec: ShardSpec, + store: "zarr.Group", + emb_buf: np.ndarray, + scales_buf: np.ndarray, + source_options: Optional[Dict[str, Any]] = None, + sample_cap: int = 0, +) -> "bool | Dict[str, Any]": + """Populate the shard buffers from their tiles and write them out. + + With ``sample_cap > 0`` the return value is the shard's stretch + statistics (see :func:`shard_stretch_stats`) — collected here because + this is the one moment every decoded pixel of the shard is in memory. + """ from . import remote t = spec.time_index S = SHARD_SIZE - # Allocate BHW buffer (bands-first for NCHW write) - emb_buf = np.zeros((N_BANDS, S, S), dtype=np.int8) - # Start with +inf (land/nodata); landmask sets water to NaN, - # valid tiles overwrite with finite scales. - scales_buf = np.full((S, S), np.float32("inf")) - has_data = False for ov in spec.tiles: # Read HWB tile, transpose to BHW @@ -1830,12 +2129,23 @@ def _write_one_shard( r, c = spec.row_px, spec.col_px store["embeddings"][t, :, r : r + S, c : c + S] = emb_buf store["scales"][t, r : r + S, c : c + S] = scales_buf + + if sample_cap > 0: + stats = shard_stretch_stats(emb_buf, scales_buf, sample_cap) + if stats is not None: + return stats return True -def _write_one_shard_worker(spec: ShardSpec) -> bool: +def _write_one_shard_worker(spec: ShardSpec) -> "bool | Dict[str, Any]": """Picklable wrapper for process pool.""" - return _write_one_shard(spec, _worker_store, _worker_source_options) + return _write_one_shard( + spec, + _worker_store, + _worker_source_options, + _worker_spill_dir, + _worker_sample_cap, + ) # --------------------------------------------------------------------------- @@ -2029,13 +2339,25 @@ def extend_store( old_t = len(existing) new_t = old_t + len(missing) + # Every time-indexed array must grow together, or the group's axes + # desynchronise. A zone from a store that predates the stretch + # statistics lacks those arrays; extending it would leave them + # permanently short, so refuse and point at the repair path. + absent = [a for a in STRETCH_ARRAY_NAMES if a not in group] + if absent: + raise ValueError( + f"{name}: missing stretch-statistics array(s) " + f"{', '.join(absent)} — this store predates fill-time stretch " + f"statistics. Run `zarr-fill --backfill-stretch-stats " + f"--zones {name[3:]}` first, then re-run zarr-extend." + ) + # Order matters only for crash-safety: grow the data arrays before # advertising the year on the time axis, so a run interrupted midway # never leaves a year readers can select but not read. - emb = group["embeddings"] - scales = group["scales"] - emb.resize((new_t,) + tuple(emb.shape[1:])) - scales.resize((new_t,) + tuple(scales.shape[1:])) + for arr_name in ("embeddings", "scales", *STRETCH_ARRAY_NAMES): + arr = group[arr_name] + arr.resize((new_t,) + tuple(arr.shape[1:])) time_arr = group["time"] time_arr.resize((new_t,)) @@ -2344,6 +2666,8 @@ def fill_store( force_lock: bool = False, state_url: Optional[str] = None, skip_existing_shards: bool = True, + spill_dir: Optional[str] = None, + collect_stretch_stats: bool = True, ) -> int: """Incrementally fill a store with tile data. @@ -2384,7 +2708,7 @@ def fill_store( fill_years = [year] if year is not None else all_years - _warn_worker_memory(workers, console) + _warn_worker_memory(workers, console, spilled=bool(spill_dir)) if console: console.print(f"Filling store at [bold]{store}[/bold]") @@ -2522,9 +2846,25 @@ def fill_store( f"({workers} workers)" ) + # Stretch statistics: collect only when the zone has the arrays + # (stores initialised before the feature lack them; repair with + # --backfill-stretch-stats). Per-shard cap sized so the expected + # candidate pool is a few times K without ballooning the result + # queue. + sample_cap = 0 + if collect_stretch_stats and "stretch_sample" in zone_store: + k_slots = zone_store["stretch_sample"].shape[1] + sample_cap = _shard_sample_cap(k_slots, len(shard_specs)) + elif collect_stretch_stats and console: + console.print( + f" [yellow]Zone {zone_num} has no stretch-statistics " + f"arrays (store predates them); skipping collection. " + f"Backfill later with --backfill-stretch-stats.[/yellow]" + ) + _acquire_zone_lock(store, zone_num, fill_year, force=force_lock) try: - written_count, failed = _write_shards( + written_count, failed, shard_stats = _write_shards( store=store, zone_group=zone_group, shard_specs=shard_specs, @@ -2532,11 +2872,33 @@ def fill_store( source_options=source.storage_options if source else None, label=f" Zone {zone_num} y{fill_year}", console=console, + spill_dir=spill_dir, + sample_cap=sample_cap, ) total_shards_written += written_count total_shards_failed += len(failed) + if shard_stats: + zone_rw = store.open_group(mode="r+", path=zone_group) + update_zone_stretch_stats( + zone_rw, + time_index, + n=sum(st["n"] for st in shard_stats), + s=sum(st["sum"] for st in shard_stats), + m=sum(st["prod"] for st in shard_stats), + sample_candidates=[ + (st["sample_emb"], st["sample_scales"], st["sample_weight"]) + for st in shard_stats + ], + ) + if console: + console.print( + f" [dim]Stretch stats: " + f"{sum(st['n'] for st in shard_stats):,} pixels " + f"folded in[/dim]" + ) + if console: console.print( f" [green]{written_count}/{len(shard_specs)} " @@ -2595,17 +2957,28 @@ def _write_shards( source_options: Optional[Dict[str, Any]], label: str, console: Optional["rich.console.Console"], -) -> Tuple[int, set]: + spill_dir: Optional[str] = None, + sample_cap: int = 0, +) -> Tuple[int, set, List[Dict[str, Any]]]: """Run the shard writes through a process pool. - Returns (shards written, set of (sr, sc) that failed). + Returns (shards written, set of (sr, sc) that failed, per-shard stretch + statistics — empty when collection is off or no shard had valid pixels). """ import multiprocessing from concurrent.futures import ProcessPoolExecutor, as_completed written_count = 0 failed: set = set() - initargs = (store.url, zone_group, store.storage_options, source_options) + stats_results: List[Dict[str, Any]] = [] + initargs = ( + store.url, + zone_group, + store.storage_options, + source_options, + spill_dir, + sample_cap, + ) # "spawn", not the Linux default of "fork": a forked worker inherits the # parent's fsspec event-loop object without the thread that runs it, and @@ -2620,8 +2993,11 @@ def _drain(pool, advance=None): for future in as_completed(futures): spec = futures[future] try: - if future.result(): + result = future.result() + if result: written_count += 1 + if isinstance(result, dict): + stats_results.append(result) except Exception as e: logger.warning(f"Shard ({spec.sr},{spec.sc}) failed: {e}") failed.add((spec.sr, spec.sc)) @@ -2669,7 +3045,7 @@ def _drain(pool, advance=None): else: pool.shutdown(wait=True) - return written_count, failed + return written_count, failed, stats_results def consolidate_store( @@ -2953,6 +3329,349 @@ def compute_stretch( _GLOBAL_STRETCH_ATTR = "geoemb:stretch" +def _parse_pca_perm(pca_rgb_order: str, pca_components: int) -> List[int]: + """Validate a PC→RGB permutation like '213' and return 0-based indices.""" + if len(pca_rgb_order) != pca_components or set(pca_rgb_order) != { + str(i + 1) for i in range(pca_components) + }: + raise ValueError( + f"pca_rgb_order must be a permutation of the digits " + f"1..{pca_components} (e.g. '123' or '213'), got {pca_rgb_order!r}" + ) + return [int(c) - 1 for c in pca_rgb_order] + + +def compute_stretch_from_stats( + store_path: "str | Path | StoreLocation", + year: int, + zones: Optional[List[int]] = None, + p_low: float = 2.0, + p_high: float = 98.0, + equalise: bool = True, + equalise_breakpoints: int = 257, + mode: str = "pca", + pca_components: int = 3, + pca_rgb_order: str = "123", + drift_threshold: float = 0.25, + console: Optional["rich.console.Console"] = None, + storage_options: Optional[Dict[str, Any]] = None, +) -> dict: + """Derive the global stretch from the per-zone ``stretch_*`` arrays. + + The fast path of ``zarr-stretch`` (docs/specs/zarr-stretch-stats.md): + reads a few MiB of per-zone summaries instead of terabytes of shards. + The PCA comes from the summed sufficient statistics and is exact — every + valid pixel in the store contributes. Quantiles come from the pooled + weighted samples, projected into PC space. + + Writes the result to the root ``geoemb:stretch.{year}`` attribute with + the same keys the legacy shard-sampling path produces, so + ``build_global_preview`` and other readers are unaffected. Works against + local and remote stores alike. + """ + if mode not in ("bands", "pca"): + raise ValueError(f"mode must be 'bands' or 'pca', got {mode!r}") + pca_perm = _parse_pca_perm(pca_rgb_order, pca_components) if mode == "pca" else None + + store = StoreLocation.resolve(store_path, storage_options) + zone_names = _zone_group_names(store, zones) + if not zone_names: + raise ValueError(f"No UTM zone groups found in {store}") + + n_total = 0 + s_total = np.zeros(N_BANDS, dtype=np.float64) + m_total = np.zeros((N_BANDS, N_BANDS), dtype=np.float64) + sample_parts: List[Tuple[np.ndarray, np.ndarray, float]] = [] + zones_used: List[str] = [] + zones_missing: List[str] = [] + + for name in zone_names: + group = store.open_group(mode="r", path=name) + if "stretch_stats_count" not in group: + zones_missing.append(name) + continue + try: + zone_years = [int(v) for v in group["time"][:]] + t = zone_years.index(year) + except (ValueError, KeyError): + continue + + n_z = int(group["stretch_stats_count"][t]) + if n_z == 0: + continue + n_total += n_z + s_total += np.asarray(group["stretch_stats_sum"][t]) + m_total += np.asarray(group["stretch_stats_prod"][t]) + + k_z = int(group["stretch_sample_count"][t]) + if k_z > 0: + sample_parts.append( + ( + np.asarray(group["stretch_sample"][t, :k_z]), + np.asarray(group["stretch_sample_scales"][t, :k_z]), + n_z / k_z, + ) + ) + zones_used.append(name) + + if zones_missing: + raise ValueError( + f"{len(zones_missing)} zone(s) have no stretch-statistics arrays " + f"({', '.join(zones_missing[:5])}{'...' if len(zones_missing) > 5 else ''}). " + f"Run `zarr-fill --backfill-stretch-stats` for them, or use " + f"--from-shards for the legacy path." + ) + if n_total == 0 or not sample_parts: + raise RuntimeError( + f"No stretch statistics recorded for year {year} — have the " + f"zone fills for this year run with stats collection on?" + ) + + if console: + console.print( + f"Stretch from stored statistics: {len(zones_used)} zone(s), " + f"{n_total:,} pixels in the exact covariance, " + f"{sum(len(p[0]) for p in sample_parts):,} sampled pixels for " + f"quantiles" + ) + + # Exact global mean/covariance from the summed sufficient statistics. + mu = s_total / n_total + cov = m_total / n_total - np.outer(mu, mu) + + # Pooled weighted sample, dequantised. + emb = np.concatenate([p[0] for p in sample_parts], axis=0) + scales = np.concatenate([p[1] for p in sample_parts], axis=0) + weights = np.concatenate( + [np.full(len(p[0]), p[2], dtype=np.float64) for p in sample_parts] + ) + x = emb.astype(np.float32) * scales[:, None] # (n, 128) + + pca_proj_components = None + pca_proj_mean = None + pca_evr = None + if mode == "pca": + eigvals, eigvecs = np.linalg.eigh(cov) + order = np.argsort(eigvals)[::-1][:pca_components] + components = eigvecs[:, order].T # (k, 128), eigenvalue-descending + evr = eigvals[order] / max(eigvals.sum(), 1e-30) + + # Drift check: re-estimate the covariance from the (independent) + # stored sample and compare. Rewritten shards double-count into the + # sums but replace sample slots, so divergence here flags stale + # statistics. The metric is relative Frobenius distance between the + # two covariances — comparing eigenvectors instead would false-alarm + # whenever eigenvalues are close, where the vectors are arbitrary. + cov_s = np.cov(x.T, aweights=weights) + drift = float( + np.linalg.norm(cov - cov_s) / max(np.linalg.norm(cov), 1e-30) + ) + # The sample covariance itself carries ~sqrt(d/n_eff) relative error, + # so the alarm floor scales with the effective sample size — a small + # sample must not read as drift. + n_eff = float(weights.sum() ** 2 / (weights**2).sum()) + noise = math.sqrt(N_BANDS / max(n_eff, 1.0)) + limit = max(drift_threshold, 3.0 * noise) + if console: + colour = "green" if drift <= limit else "red" + console.print( + f" Drift check: |cov_stats − cov_sample|/|cov_stats| = " + f"[{colour}]{drift:.4f}[/{colour}] " + f"(limit {limit:.2f}; sample noise floor {noise:.2f})" + ) + if drift > limit: + logger.warning( + f"Stretch statistics drift: {drift:.4f} > {limit:.2f}. " + f"The additive sums likely double-counted rewritten shards; " + f"rebuild with `zarr-fill --backfill-stretch-stats`, or " + f"cross-check with `zarr-stretch --from-shards`." + ) + + # Bake the PC→RGB permutation into the stored matrix, as the legacy + # path does, so the render path needs no extra swapping. + components = components[pca_perm] + evr = evr[pca_perm] + + channels = (x - mu.astype(np.float32)) @ components.T.astype(np.float32) + pca_proj_components = [[float(v) for v in row] for row in components] + pca_proj_mean = [float(v) for v in mu] + pca_evr = [float(v) for v in evr] + band_indices: Tuple[int, ...] = tuple(range(N_BANDS)) + else: + band_indices = RGB_PREVIEW_BANDS + channels = x[:, list(band_indices)] + + n_ch = channels.shape[1] + stretch_min = [ + float(weighted_percentile(channels[:, i], weights, np.array([p_low]))[0]) + for i in range(n_ch) + ] + stretch_max = [ + float(weighted_percentile(channels[:, i], weights, np.array([p_high]))[0]) + for i in range(n_ch) + ] + for i in range(n_ch): + if stretch_max[i] <= stretch_min[i]: + stretch_max[i] = stretch_min[i] + 1.0 + + cdf_breaks = None + if equalise: + n_break = max(64, int(equalise_breakpoints)) + qs = np.linspace(0.0, 100.0, n_break) + cdf_breaks = [] + for i in range(n_ch): + bks = weighted_percentile(channels[:, i], weights, qs) + for j in range(1, len(bks)): + if bks[j] <= bks[j - 1]: + bks[j] = bks[j - 1] + 1e-9 + cdf_breaks.append([float(v) for v in bks]) + + # Persist with the same key set as the legacy path so readers + # (_load_global_stretch, build_global_preview) are unaffected. + root_rw = store.open_group(mode="r+") + stretch_map = dict(root_rw.attrs.get(_GLOBAL_STRETCH_ATTR, {})) + method_prefix = "zone_stats_pca" if mode == "pca" else "zone_stats_percentile" + entry: Dict[str, Any] = { + "min": stretch_min, + "max": stretch_max, + "p_low": p_low, + "p_high": p_high, + "samples": int(channels.shape[0]), + "stats_pixels": int(n_total), + "zones_used": len(zones_used), + "bands": list(band_indices), + "method": f"{method_prefix}{'_equalised' if equalise else ''}", + "mode": mode, + } + if cdf_breaks is not None: + entry["cdf"] = cdf_breaks + if pca_proj_components is not None: + entry["pca_components"] = pca_proj_components + entry["pca_mean"] = pca_proj_mean + entry["pca_explained_variance_ratio"] = pca_evr + stretch_map[str(year)] = entry + root_rw.attrs[_GLOBAL_STRETCH_ATTR] = stretch_map + + if console: + console.print( + f"[green]Saved to {_GLOBAL_STRETCH_ATTR}.{year} on the store " + f"root.[/green] Run zarr-consolidate so consolidated-metadata " + f"readers see it." + ) + + return {"min": stretch_min, "max": stretch_max, "samples": int(channels.shape[0])} + + +def backfill_stretch_stats( + store_path: "str | Path | StoreLocation", + zones: Optional[List[int]] = None, + years: Optional[List[int]] = None, + sample_k: int = STRETCH_SAMPLE_K, + console: Optional["rich.console.Console"] = None, + storage_options: Optional[Dict[str, Any]] = None, + state_url: Optional[str] = None, + force_lock: bool = False, +) -> int: + """Rebuild a zone's stretch statistics by scanning its existing shards. + + The repair path for stores filled before fill-time collection existed, + for interrupted fills, and for suspected double-counting: it re-reads + the zone's shards once (the only stats path that touches embeddings) and + *sets* the arrays from what is actually in the store. Creates the arrays + if the zone predates them. Per-zone, so it composes with fills of other + zones; takes the same (zone, year) lock a fill would. + + Returns the number of (zone, year) slots rebuilt. + """ + store = StoreLocation.resolve(store_path, storage_options, state_url) + zone_names = _zone_group_names(store, zones) + if not zone_names: + raise ValueError(f"No UTM zone groups found in {store}") + + rebuilt = 0 + for name in zone_names: + group = store.open_group(mode="r+", path=name, zarr_format=3) + zone_years = [int(v) for v in group["time"][:]] + T = len(zone_years) + + if "stretch_sample" not in group: + create_stretch_arrays(group, T, sample_k) + if console: + console.print(f" {name}: created stretch-statistics arrays") + k_slots = group["stretch_sample"].shape[1] + + emb_arr = group["embeddings"] + scales_arr = group["scales"] + H, W = emb_arr.shape[2], emb_arr.shape[3] + all_coords = { + (sr, sc) + for sr in range(math.ceil(H / SHARD_SIZE)) + for sc in range(math.ceil(W / SHARD_SIZE)) + } + + for fill_year in years or zone_years: + if fill_year not in zone_years: + continue + t = zone_years.index(fill_year) + present = _existing_shards(store, name, t, all_coords, console=None) + if not present: + continue + + cap = _shard_sample_cap(k_slots, len(present)) + zone_num = int(name[3:]) + _acquire_zone_lock(store, zone_num, fill_year, force=force_lock) + try: + n_total, s_total = 0, np.zeros(N_BANDS, dtype=np.float64) + m_total = np.zeros((N_BANDS, N_BANDS), dtype=np.float64) + candidates: List[Tuple[np.ndarray, np.ndarray, float]] = [] + for i, (sr, sc) in enumerate(sorted(present)): + r0, c0 = sr * SHARD_SIZE, sc * SHARD_SIZE + r1, c1 = min(r0 + SHARD_SIZE, H), min(c0 + SHARD_SIZE, W) + st = shard_stretch_stats( + np.asarray(emb_arr[t, :, r0:r1, c0:c1]), + np.asarray(scales_arr[t, r0:r1, c0:c1]), + cap, + ) + if st is None: + continue + n_total += st["n"] + s_total += st["sum"] + m_total += st["prod"] + candidates.append( + (st["sample_emb"], st["sample_scales"], st["sample_weight"]) + ) + if console: + console.print( + f" {name} {fill_year}: shard {i + 1}/{len(present)} " + f"({st['n']:,} px)", + end="\r", + ) + + emb_s, scales_s = merge_stretch_samples(candidates, k_slots) + # Backfill SETS from actual contents (it is the repair for + # double-counting), unlike the fill's additive fold. + group["stretch_stats_count"][t] = n_total + group["stretch_stats_sum"][t] = s_total + group["stretch_stats_prod"][t] = m_total + full_emb = np.zeros((k_slots, N_BANDS), dtype=np.int8) + full_sc = np.full(k_slots, np.float32("inf"), dtype=np.float32) + full_emb[: len(emb_s)] = emb_s + full_sc[: len(emb_s)] = scales_s + group["stretch_sample"][t] = full_emb + group["stretch_sample_scales"][t] = full_sc + group["stretch_sample_count"][t] = len(emb_s) + rebuilt += 1 + if console: + console.print( + f" {name} {fill_year}: rebuilt from " + f"{len(present)} shard(s), {n_total:,} pixels " + ) + finally: + _release_zone_lock(store, zone_num, fill_year) + + return rebuilt + + def _sample_shard_task( store_path_str: str, zone_group: str, @@ -3054,19 +3773,10 @@ def compute_global_stretch( band_indices = tuple(range(pca_total_bands)) # Parse the pca_rgb_order permutation now so we fail fast on bad input. - pca_perm: Optional[List[int]] = None - if mode == "pca": - if len(pca_rgb_order) != pca_components or set(pca_rgb_order) != { - str(i + 1) for i in range(pca_components) - }: - raise ValueError( - f"pca_rgb_order must be a permutation of the digits " - f"1..{pca_components} (e.g. '123' or '213'), got " - f"{pca_rgb_order!r}" - ) - # 0-indexed permutation: pca_perm[k] = which PC ends up in output channel k. - # "123" -> [0, 1, 2] = identity; "213" -> [1, 0, 2] = swap R/G. - pca_perm = [int(c) - 1 for c in pca_rgb_order] + # pca_perm[k] = which PC ends up in output channel k ("213" swaps R/G). + pca_perm: Optional[List[int]] = ( + _parse_pca_perm(pca_rgb_order, pca_components) if mode == "pca" else None + ) # Find time_index for the requested year via the first zone's time coord. time_index = None diff --git a/tests/zarr.t b/tests/zarr.t index b0cc263..ccb69c2 100644 --- a/tests/zarr.t +++ b/tests/zarr.t @@ -88,6 +88,30 @@ landmask registry for its land denominator: $ geotessera-registry zarr-scan --help | grep -oE '\[base_dir\] store_path' [base_dir] store_path +Test: stretch statistics are collected at fill time +---------------------------------------------------- + +Fills fold each shard's stretch statistics into the zone group, so the +global stretch never re-reads embeddings: + + $ geotessera-registry zarr-fill --help | grep -o '\-\-no-stretch-stats' | sort -u + --no-stretch-stats + + $ geotessera-registry zarr-fill --help | grep -o '\-\-backfill-stretch-stats' | sort -u + --backfill-stretch-stats + + $ geotessera-registry zarr-init --help | grep -o '\-\-stretch-sample-size' | sort -u + --stretch-sample-size + +zarr-stretch aggregates them by default, remote-capable; the legacy +shard-sampling path is the opt-in: + + $ geotessera-registry zarr-stretch --help | grep -o '\-\-from-shards' | sort -u + --from-shards + + $ geotessera-registry zarr-stretch --help | grep -o '\-\-drift-threshold' | sort -u + --drift-threshold + Test: zarr-extend grows the time axis -------------------------------------- diff --git a/tests/zarr_remote_check.py b/tests/zarr_remote_check.py index 3c7b29c..efeae3c 100644 --- a/tests/zarr_remote_check.py +++ b/tests/zarr_remote_check.py @@ -405,7 +405,34 @@ def placed(lon, x_off): root["utm31"]["embeddings"][0] = 7 root["utm31"]["scales"][0] = 0.5 +from geotessera.zarr import ( # noqa: E402 + STRETCH_ARRAY_NAMES, + create_stretch_arrays, +) + +# A pre-stats store must be refused (extending it would leave the stretch +# arrays permanently short) and pointed at the backfill. +try: + extend_store(ext, [2026]) + check("extend refuses a store without stretch arrays", False) +except ValueError as e: + check( + "extend refuses a store without stretch arrays", + "backfill-stretch-stats" in str(e), + ) + +for zname in ("utm30", "utm31"): + create_stretch_arrays(root[zname], n_years=1, k=50) + check("extend adds the year to every zone", extend_store(ext, [2026]) == 2) +check( + "extend grows the stretch arrays too", + all( + root[z][a].shape[0] == 2 + for z in ("utm30", "utm31") + for a in STRETCH_ARRAY_NAMES + ), +) g31 = zarr.open_group(ext.url, mode="r", use_consolidated=False)["utm31"] check("time axis grew", [int(v) for v in g31["time"][:]] == [2024, 2026]) @@ -439,6 +466,92 @@ def placed(lon, x_off): "extend --force overrides a stale lock", extend_store(ext, [2027], force=True) == 2 ) +# --------------------------------------------------------------------------- +# Stretch statistics +# --------------------------------------------------------------------------- + +from geotessera.zarr import ( # noqa: E402 + merge_stretch_samples, + shard_stretch_stats, + update_zone_stretch_stats, + weighted_percentile, +) + +nprng = np.random.default_rng(11) +B = 128 +semb = nprng.integers(-128, 127, (B, 96, 96), dtype=np.int8) +ssc = (nprng.random((96, 96)).astype(np.float32) * 0.01 + 0.001) +ssc[:20, :20] = np.nan +ssc[80:, 80:] = np.inf + +sst = shard_stretch_stats(semb, ssc, sample_cap=200, seed=5) +svalid = np.isfinite(ssc) +sx = semb.reshape(B, -1)[:, svalid.ravel()].astype(np.float64) * ssc.ravel()[ + svalid.ravel() +] +check("stats count exact", sst["n"] == int(svalid.sum())) +check( + "stats sum matches population", + bool(np.allclose(sst["sum"], sx.sum(1), rtol=1e-5)), +) +_truth = sx @ sx.T +check( + "stats product matches population", + float(np.abs(sst["prod"] - _truth).max()) < 1e-5 * float(np.abs(_truth).max()), +) +check( + "sampled pixels are valid pixels", + bool(np.isfinite(sst["sample_scales"]).all()), +) + +ha = shard_stretch_stats(semb[:, :48, :], ssc[:48, :], 50, seed=1) +hb = shard_stretch_stats(semb[:, 48:, :], ssc[48:, :], 50, seed=2) +check("stats additive across shards", ha["n"] + hb["n"] == sst["n"]) +check( + "sums additive across shards", + bool(np.allclose(ha["sum"] + hb["sum"], sst["sum"], rtol=1e-5)), +) + +heavy = (np.ones((500, B), np.int8), np.ones(500, np.float32), 50.0) +light = (np.zeros((500, B), np.int8), np.zeros(500, np.float32), 1.0) +me, ms = merge_stretch_samples([heavy, light], 300, seed=4) +check("merge respects capacity", len(me) == 300) +check( + "merge favours high-weight rows", + float((me[:, 0] == 1).mean()) > 0.8, +) + +vv = nprng.normal(size=4000) +check( + "weighted percentile matches numpy under uniform weights", + float( + np.abs( + weighted_percentile(vv, np.ones(4000), np.array([2.0, 50.0, 98.0])) + - np.percentile(vv, [2, 50, 98]) + ).max() + ) + < 0.02, +) + +# Zone-array round trip: create, fold twice, contents accumulate. +import zarr # noqa: E402 + +zs = zarr.open_group(str(TMP / "stats.zarr"), mode="w", zarr_format=3) +create_stretch_arrays(zs, n_years=2, k=100) +check( + "stretch arrays created", + all(n in zs for n in STRETCH_ARRAY_NAMES), +) +cand = [(sst["sample_emb"], sst["sample_scales"], sst["sample_weight"])] +update_zone_stretch_stats(zs, 0, sst["n"], sst["sum"], sst["prod"], cand, seed=1) +update_zone_stretch_stats(zs, 0, sst["n"], sst["sum"], sst["prod"], cand, seed=2) +check("stats fold additively", int(zs["stretch_stats_count"][0]) == 2 * sst["n"]) +check( + "sample capacity bounded", + int(zs["stretch_sample_count"][0]) <= 100, +) +check("other year untouched", int(zs["stretch_stats_count"][1]) == 0) + # --------------------------------------------------------------------------- # Storage options and source layout # --------------------------------------------------------------------------- From f276bcfdd8c32245dabbaddd072ec45f7a704daf Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Sun, 2 Aug 2026 14:14:16 +0100 Subject: [PATCH 12/13] simplify for serverless deployment zarr-fill is now stateless: no build directory, no ingestion registry, no advisory locks. The store's shard objects are the only record of progress -- a shard is always written from every tile covering it, so its presence is proof of completion -- and the start-of-run scan is the resume. A spot instance killed at any moment leaves nothing to clean up or take over; relaunching the same command continues where the objects stop. Stretch statistics collect themselves by the same mechanism. A seventh per-zone array, stretch_stats_shards (T, shard_rows, shard_cols) uint8, records which shards are folded into the sums; the fill diffs it against the scan and reads back any shard whose statistics are missing as catch-up tasks in the same worker pool. That covers a crash between shard write and stats fold (sums are written before the mask, so the worst-timed kill re-folds one shard rather than silently dropping it), shards written by builds that predate collection (arrays are created on the zone's next fill), and sums whose provenance the missing mask makes unknowable (reset and recomputed). --backfill-stretch-stats survives only as the explicit repair for suspected double-counting. One fill per (zone, year) at a time remains the operating contract but is no longer enforced: identical concurrent fills write identical shards, so the failure mode is wasted work or a drift-detectable stats double-fold, not data corruption. --state-url and --force-lock are accepted as no-ops so existing scripts keep working; zarr-consolidate still reads --state-url to merge ingestion registries written by older builds, and skips the merge entirely rather than conjure a state directory for an empty result. --- CHANGES.md | 13 + docs/architecture.rst | 77 ++-- docs/specs/zarr-stretch-stats.md | 4 +- geotessera/registry_cli.py | 23 +- geotessera/zarr.py | 621 +++++++++++++++++-------------- tests/zarr_remote_check.py | 70 +--- 6 files changed, 429 insertions(+), 379 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 53da45f..0dcf970 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -73,6 +73,19 @@ tile inventory has grown, since a newly-added tile falls inside an existing shard. `--skip-existing-shards` is still accepted as a no-op. (@avsm) +- **Fills are stateless and stretch statistics collect themselves.** The + ingestion registry and advisory locks are gone from `zarr-fill`: the + store's shard objects are the only record of progress, so a preemptible + (spot) instance that dies mid-run leaves nothing to clean up or take + over — relaunching the same command scans and continues. A per-zone + coverage mask (`stretch_stats_shards`) records which shards are folded + into the stretch sums; the fill diffs it against the same scan and reads + back any shard whose statistics are missing, so interrupted runs and + stores from older builds converge automatically — no separate backfill + step, which now exists only as the explicit repair for suspected + double-counting. `--state-url` and `--force-lock` are accepted as no-ops + for existing scripts; `zarr-consolidate` still reads `--state-url` to + merge registries written by older builds. (@avsm) - **Per-zone stretch statistics, collected at fill time** (see `docs/specs/zarr-stretch-stats.md`): each zone group gains six arrays — exact mean/covariance sufficient statistics per (zone, year), additive diff --git a/docs/architecture.rst b/docs/architecture.rst index a52d4dc..1ac8061 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -547,8 +547,8 @@ CLI's ``--acl``:: --store-profile sc-writer --store-acl bucket-owner-full-control -It applies to the store's Zarr chunks and metadata as well as the sidecar -parquet and lock objects, and is filtered out of read requests. +It applies to every object written to the store and is filtered out of +read requests. .. note:: @@ -608,7 +608,7 @@ Two constraints: renumber every existing chunk's time index, i.e. rewrite the store. It is refused rather than done silently. * **Single writer.** Unlike a fill, this rewrites array metadata for every - zone, so it refuses to run while any fill lock is held, and it *does* + zone — do not run it while fills are in flight — and it *does* re-consolidate afterwards (readers cannot see the new year until it has). .. _zarr-parallel-sweep: @@ -620,51 +620,54 @@ A UTM zone's pixels live entirely within its own ``utm{zone}`` group, and shards never straddle zones. That makes ``--zones N`` the natural unit of parallelism: one process per zone, all writing to the same store. -Everything a fill mutates is keyed by ``(zone, year)``, and none of it lives -inside the store — build bookkeeping goes to a sibling location so the -published hierarchy contains only Zarr: +Fills are **stateless**: the store's own shard objects are the only record +of progress, so there is no registry to update, no lock to hold, and no +build directory at all. A zone job on a preemptible (spot) instance that +dies mid-run leaves nothing to clean up — relaunching the same command +scans the store and continues from wherever the objects stop: .. code-block:: text tessera.zarr/ zarr.json # shared — consolidation only - utm30/, utm31/, ... # one zone per process - - tessera.zarr.build/ # --state-url to relocate - _registry/utm30_2024.parquet # per-zone ingestion tracking - _registry.parquet # merged view, written by consolidate - _locks/utm30_2024.json # advisory fill lock - -* **Ingestion tracking** is one object per zone/year, so no two jobs - read-modify-write the same file. It records which tiles have already been - written, which is what makes a fill resumable and lets a later run pick up - tiles the manifest has gained since. It is build state, not published - data — a reader of the store never needs it — so it lives in the state - sibling. Stores built before this split kept a ``_registry.parquet`` - inside the hierarchy; that is still read, so they resume correctly. -* **An advisory lock** is taken for the duration of a zone/year fill. It - catches the same zone being launched twice — the case that would silently - corrupt data, because a shard write replaces the whole shard. Object - stores offer no atomic create, so the lock is advisory; ``--force-lock`` - takes over one left behind by a dead run. + utm30/, utm31/, ... # one zone per process; each + # holds its own stretch_* stats + +* **Resume is the scan**: a shard is always written from every tile + covering it, so its presence in the store is proof of completion; the + fill lists what exists and writes only what is missing. +* **Stretch statistics are self-catching-up**: each zone's coverage mask + records which shards are folded into its statistics, so the same scan + also finds shards whose stats are missing (a crash between write and + fold, or shards written by older builds) and reads them back. Statistics + converge without any separate backfill step. +* **One fill per (zone, year) at a time remains the operating contract**, + but it is no longer enforced by locks: two identical concurrent fills + write identical shards (wasted work, not corruption), and a stats + double-fold is caught by the drift check and repaired by + ``--backfill-stretch-stats``. ``--force-lock`` is accepted as a no-op + for older scripts. * **Consolidation is skipped** by default when ``--zones`` is given, since the root ``zarr.json`` is the one object all jobs share. +* ``--state-url`` is legacy: only ``zarr-consolidate`` still reads it, to + merge ingestion registries written by pre-stateless builds. Resuming After a Crash ~~~~~~~~~~~~~~~~~~~~~~ -The ingestion registry is written when a (zone, year) finishes, so a run -that dies partway — an OOM kill leaves no traceback — loses that year's -bookkeeping even though the shards it wrote are safely in the store. - -The shard objects are the ground truth and they survive anything, so a fill -scans for them before doing any work and skips what is already there. That -is the default: re-running an interrupted fill uploads only what is missing. - -A shard is always written from every tile covering it, so its presence means -it is complete. The exception is a tile inventory that has grown since — -a newly-added tile falls inside an existing shard, which would then be -skipped rather than merged in. Force those shards to be rebuilt with:: +There is nothing to resume *from* except the store itself, which is the +point: a fill keeps no state of its own, so a run killed at any moment — +including a spot-instance preemption — is continued by running the same +command again. The scan finds the shard objects that landed and writes the +rest; the stretch statistics' coverage mask finds any shard whose pixels +were written but not yet folded in and reads it back. Worst case for a +crash between a shard write and its stats fold is one re-read of that +shard. + +The one assumption resume makes is that the tile inventory has not grown +under an existing shard: a shard is complete with respect to the manifest +it was written from, so a tile added later falls inside an object the scan +skips. Force those shards to be rebuilt with:: geotessera-registry zarr-fill --zones 30 \ --rewrite-existing-shards diff --git a/docs/specs/zarr-stretch-stats.md b/docs/specs/zarr-stretch-stats.md index b204189..1fbf333 100644 --- a/docs/specs/zarr-stretch-stats.md +++ b/docs/specs/zarr-stretch-stats.md @@ -32,7 +32,9 @@ Five arrays are added to every zone group `utm{zz}/`. They are ordinary Zarr v3 | `stretch_sample` | `(T, K, 128)` | int8 | `time, sample, band` | raw sampled embedding vectors | | `stretch_sample_scales` | `(T, K)` | float32 | `time, sample` | per-sample dequant scale | -**[as built]** The number of filled slots is a sixth array, `stretch_sample_count` `(T,)` int64, not a zone-group attribute: the `geoemb:` convention keeps its attributes on the root group only, and zone groups deliberately carry nothing beyond `proj:`/`spatial:`. Unfilled slots are zero-valued with scale `+inf`, matching the existing "not yet filled" sentinel, so a reader that ignores the count still cannot mistake padding for data. +**[as built]** Two structural deviations. First, the number of filled slots is a sixth array, `stretch_sample_count` `(T,)` int64, not a zone-group attribute: the `geoemb:` convention keeps its attributes on the root group only, and zone groups deliberately carry nothing beyond `proj:`/`spatial:`. Unfilled slots are zero-valued with scale `+inf`, matching the existing "not yet filled" sentinel, so a reader that ignores the count still cannot mistake padding for data. + +**[as built]** Second, a seventh array `stretch_stats_shards` `(T, n_shard_rows, n_shard_cols)` uint8 — the **coverage mask**, 1 where a shard's pixels are folded into the sums. It is what makes collection *automatic*: the fill's normal store scan diffed against the mask yields the shards whose statistics are missing (crash between write and fold, or shards written by older builds), and they are read back as catch-up tasks in the same worker pool. This replaced both the fill-time lock/registry state and the manual backfill workflow: fills are now fully stateless (no `--state-url`, no locks — the store is the only record), `--backfill-stretch-stats` survives only as the explicit repair for suspected double-counting, and `ensure_stretch_arrays` heals pre-feature stores on their next fill (resetting sums whose provenance the missing mask makes unknowable). Samples are stored in **source representation** — the int8 embedding vector plus its float32 scale — which is lossless with respect to the store itself and lets any downstream statistic be recomputed exactly as if the pixels had been read from `embeddings`/`scales`. diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index af7eb81..b215476 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3335,14 +3335,15 @@ def _add_source_args(parser) -> None: def _add_state_arg(parser) -> None: - """Register ``--state-url`` for the commands that keep build bookkeeping.""" + """Register ``--state-url`` (legacy; fills are now stateless).""" parser.add_argument( "--state-url", type=str, default=None, - help="Where to keep build state (ingestion registry, fill locks). " - "Default: .build alongside the store. Kept outside the store " - "so the published Zarr hierarchy contains only Zarr.", + help="Legacy build-state location from before fills became " + "stateless. Only zarr-consolidate still reads it (to merge old " + "ingestion registries); accepted elsewhere for script " + "compatibility and ignored.", ) @@ -3775,8 +3776,6 @@ def zarr_fill_command(args): years=years, console=console, storage_options=store_options, - state_url=args.state_url, - force_lock=args.force_lock, ) except (ValueError, RuntimeError) as e: console.print(f"[red]{emoji('❌ ')}{e}[/red]") @@ -3826,8 +3825,6 @@ def zarr_fill_command(args): storage_options=store_options, source=source, consolidate=consolidate, - force_lock=args.force_lock, - state_url=args.state_url, skip_existing_shards=not args.rewrite_existing_shards, spill_dir=args.spill_dir, collect_stretch_stats=not args.no_stretch_stats, @@ -3919,8 +3916,6 @@ def zarr_extend_command(args): storage_options=store_options, zones=zones, consolidate=not args.no_consolidate, - force=args.force, - state_url=args.state_url, ) except (ValueError, RuntimeError) as e: console.print(f"[red]{emoji('❌ ')}{e}[/red]") @@ -4968,8 +4963,9 @@ def main(): zarr_fill_parser.add_argument( "--force-lock", action="store_true", - help="Take over a zone/year lock left behind by a dead run. Only use " - "this when no other fill is touching the same zone.", + help="Deprecated no-op: fills no longer take locks. The store's own " + "shard objects are the only state, so a dead run leaves nothing to " + "take over.", ) _add_source_args(zarr_fill_parser) _add_storage_args(zarr_fill_parser, "source", "Tile source") @@ -5070,7 +5066,8 @@ def main(): zarr_extend_parser.add_argument( "--force", action="store_true", - help="Proceed even if fill locks are present (only when they are stale)", + help="Deprecated no-op: fills no longer take locks. Do not extend " + "while a fill is in flight.", ) _add_state_arg(zarr_extend_parser) _add_storage_args(zarr_extend_parser, "store", "Store", writable=True) diff --git a/geotessera/zarr.py b/geotessera/zarr.py index 06295f2..fdf6654 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -1434,7 +1434,13 @@ def _create_zone_group( # Per-zone stretch statistics, populated by zarr-fill (see # docs/specs/zarr-stretch-stats.md). Plain arrays: zone groups carry no # geoemb: attributes. - create_stretch_arrays(store, T, stretch_sample_size) + create_stretch_arrays( + store, + T, + stretch_sample_size, + math.ceil(H / SHARD_SIZE), + math.ceil(W / SHARD_SIZE), + ) # Use geozarr-toolkit for proj: and spatial: convention metadata from geozarr_toolkit import create_geozarr_attrs @@ -1495,15 +1501,27 @@ def _create_zone_group( "stretch_sample", "stretch_sample_scales", "stretch_sample_count", + "stretch_stats_shards", ) -def create_stretch_arrays(group: "zarr.Group", n_years: int, k: int) -> None: +def create_stretch_arrays( + group: "zarr.Group", + n_years: int, + k: int, + n_shard_rows: int, + n_shard_cols: int, +) -> None: """Create the per-zone stretch-statistics arrays in *group*. One chunk per year on the time axis, so a (zone, year) update touches exactly one chunk per array and ``zarr-extend`` grows them the same way it grows ``embeddings``. + + ``stretch_stats_shards`` is the coverage mask: 1 where a shard's pixels + are folded into the sums. It is what lets a fill know, from its normal + store scan, which existing shards the statistics have not yet seen — so + catch-up is automatic and no separate backfill pass is needed. """ from zarr.codecs import BloscCodec @@ -1523,6 +1541,9 @@ def create_stretch_arrays(group: "zarr.Group", n_years: int, k: int) -> None: ("stretch_sample_scales", (T, k), (1, k), np.float32, np.float32("inf"), ["time", "sample"]), ("stretch_sample_count", (T,), (1,), np.int64, 0, ["time"]), + ("stretch_stats_shards", (T, n_shard_rows, n_shard_cols), + (1, n_shard_rows, n_shard_cols), np.uint8, np.uint8(0), + ["time", "shard_row", "shard_col"]), ] for name, shape, chunks, dtype, fill, dims in specs: group.create_array( @@ -1536,6 +1557,81 @@ def create_stretch_arrays(group: "zarr.Group", n_years: int, k: int) -> None: ) +def ensure_stretch_arrays( + group: "zarr.Group", + console: Optional["rich.console.Console"] = None, + sample_k: int = STRETCH_SAMPLE_K, +) -> None: + """Create any missing stretch arrays on an existing zone group. + + Called by the fill, so stores initialised before the feature (or before + the coverage mask) heal themselves on their next fill. If the coverage + mask is missing but the sums are non-zero, the sums' provenance is + unknowable — they were collected without shard tracking — so they are + reset and the fill's automatic catch-up recomputes them from the store. + """ + absent = [a for a in STRETCH_ARRAY_NAMES if a not in group] + if not absent: + return + + T = group["time"].shape[0] + H, W = group["embeddings"].shape[2], group["embeddings"].shape[3] + n_sr, n_sc = math.ceil(H / SHARD_SIZE), math.ceil(W / SHARD_SIZE) + k = ( + group["stretch_sample"].shape[1] + if "stretch_sample" in group + else sample_k + ) + + had_untracked_sums = ( + "stretch_stats_shards" in absent + and "stretch_stats_count" in group + and int(np.asarray(group["stretch_stats_count"][:]).sum()) > 0 + ) + + from zarr.codecs import BloscCodec + + comp = BloscCodec(cname="zstd", clevel=3) + all_specs = { + "stretch_stats_count": ((T,), (1,), np.int64, 0, ["time"]), + "stretch_stats_sum": ((T, N_BANDS), (1, N_BANDS), np.float64, 0.0, + ["time", "band"]), + "stretch_stats_prod": ((T, N_BANDS, N_BANDS), (1, N_BANDS, N_BANDS), + np.float64, 0.0, ["time", "band", "band2"]), + "stretch_sample": ((T, k, N_BANDS), (1, k, N_BANDS), np.int8, + np.int8(0), ["time", "sample", "band"]), + "stretch_sample_scales": ((T, k), (1, k), np.float32, + np.float32("inf"), ["time", "sample"]), + "stretch_sample_count": ((T,), (1,), np.int64, 0, ["time"]), + "stretch_stats_shards": ((T, n_sr, n_sc), (1, n_sr, n_sc), np.uint8, + np.uint8(0), ["time", "shard_row", + "shard_col"]), + } + for name in absent: + shape, chunks, dtype, fill, dims = all_specs[name] + group.create_array( + name, shape=shape, chunks=chunks, dtype=dtype, fill_value=fill, + compressors=comp, dimension_names=dims, + ) + + if had_untracked_sums: + for t in range(T): + group["stretch_stats_count"][t] = 0 + group["stretch_stats_sum"][t] = np.zeros(N_BANDS) + group["stretch_stats_prod"][t] = np.zeros((N_BANDS, N_BANDS)) + group["stretch_sample_count"][t] = 0 + if console: + console.print( + " [yellow]Existing stretch sums predate shard tracking; " + "reset — the fill recomputes them from the store as it " + "goes.[/yellow]" + ) + if console: + console.print( + f" [dim]Created stretch array(s): {', '.join(absent)}[/dim]" + ) + + def _shard_sample_cap(k_slots: int, n_shards: int) -> int: """Per-shard sample size: a few times K spread over the shards. @@ -1647,15 +1743,17 @@ def update_zone_stretch_stats( s: np.ndarray, m: np.ndarray, sample_candidates: List[Tuple[np.ndarray, np.ndarray, float]], + seen_coords: Optional[List[Tuple[int, int]]] = None, seed: Optional[int] = None, ) -> None: - """Fold one fill run's statistics into a zone's arrays (read-modify-write). + """Fold one run's statistics into a zone's arrays (read-modify-write). The additive triple is summed onto what is stored; the sample is re-drawn from the stored sample and the new candidates together, weighted so the result still approximates a uniform draw over all pixels either has seen. - Caller must hold the (zone, year) fill lock — this is the same - single-writer context the shard writes ran under. + ``seen_coords`` marks those shards in the coverage mask — written after + the sums, so a crash in between re-folds rather than silently drops. + One fill per (zone, year) at a time remains the operating contract. """ t = time_index count_arr = zone_group["stretch_stats_count"] @@ -1686,7 +1784,13 @@ def update_zone_stretch_stats( if filled: zone_group["stretch_sample"][t, :filled] = emb zone_group["stretch_sample_scales"][t, :filled] = scales - zone_group["stretch_sample_count"][t] = filled + zone_group["stretch_sample_count"][t] = filled + + if seen_coords: + mask = np.asarray(zone_group["stretch_stats_shards"][t]) + for sr, sc in seen_coords: + mask[sr, sc] = 1 + zone_group["stretch_stats_shards"][t] = mask # --------------------------------------------------------------------------- @@ -1893,11 +1997,16 @@ def merge_tile_registry( frames.append(part) if not frames: - merged = _empty_tile_registry() - else: - merged = gpd.GeoDataFrame( - pd.concat(frames, ignore_index=True), crs="EPSG:4326" - ).drop_duplicates(subset=["year", "zone", "tile_lon", "tile_lat"], keep="last") + # Lock-free fills keep no ingestion registry, so a store built + # entirely by them has nothing here — do not conjure a state dir + # just to hold an empty parquet. + if console: + console.print(" No legacy ingestion registry to merge") + return 0 + + merged = gpd.GeoDataFrame( + pd.concat(frames, ignore_index=True), crs="EPSG:4326" + ).drop_duplicates(subset=["year", "zone", "tile_lon", "tile_lat"], keep="last") _write_parquet_at(state, merged, MERGED_REGISTRY_NAME) @@ -1909,60 +2018,6 @@ def merge_tile_registry( return len(merged) -# --------------------------------------------------------------------------- -# Advisory zone locks -# --------------------------------------------------------------------------- -# Two processes filling the same (zone, year) would each rewrite whole shards -# from their own tile subset and silently erase each other's pixels. Object -# stores give us no atomic create, so this is advisory only — it catches the -# common accident (the same zone launched twice) rather than enforcing -# mutual exclusion. - -LOCK_DIR_NAME = "_locks" - - -def _lock_name(zone: int, year: int) -> str: - return f"{_zone_group_name(zone)}_{year}.json" - - -def _acquire_zone_lock( - store: StoreLocation, zone: int, year: int, force: bool = False -) -> None: - """Claim (zone, year) for this process, or raise if someone else holds it.""" - import json - import socket - import pandas as pd - - state = store.state - name = _lock_name(zone, year) - if not force and state.exists(LOCK_DIR_NAME, name, on_denied=False): - try: - held = json.loads(state.read_bytes(LOCK_DIR_NAME, name)) - except Exception: - held = {} - raise RuntimeError( - f"Zone {zone} year {year} is locked by " - f"{held.get('host', '?')}:{held.get('pid', '?')} " - f"since {held.get('acquired_at', 'unknown time')}. " - f"Another fill is in progress, or a previous one died. " - f"Re-run with --force-lock to take it over." - ) - - payload = { - "zone": zone, - "year": year, - "host": socket.gethostname(), - "pid": os.getpid(), - "acquired_at": pd.Timestamp.now(tz="UTC").isoformat(), - } - state.write_bytes(json.dumps(payload).encode(), LOCK_DIR_NAME, name) - - -def _release_zone_lock(store: StoreLocation, zone: int, year: int) -> None: - """Drop this process's claim on (zone, year).""" - store.state.remove(LOCK_DIR_NAME, _lock_name(zone, year)) - - # --------------------------------------------------------------------------- # Shard writing (NCHW layout) # --------------------------------------------------------------------------- @@ -1971,6 +2026,7 @@ def _release_zone_lock(store: StoreLocation, zone: int, year: int) -> None: _worker_source_options: Optional[Dict[str, Any]] = None _worker_spill_dir: Optional[str] = None _worker_sample_cap: int = 0 # 0 = stats collection off +_worker_time_index: int = 0 def _init_shard_worker( @@ -1980,6 +2036,7 @@ def _init_shard_worker( source_options: Optional[Dict[str, Any]] = None, spill_dir: Optional[str] = None, sample_cap: int = 0, + time_index: int = 0, ) -> None: """Process pool initializer: open the zone group once per worker. @@ -1987,7 +2044,7 @@ def _init_shard_worker( own filesystem connections rather than inheriting an unforkable client. """ global _worker_store, _worker_source_options, _worker_spill_dir - global _worker_sample_cap + global _worker_sample_cap, _worker_time_index from . import remote @@ -2001,6 +2058,7 @@ def _init_shard_worker( _worker_source_options = source_options _worker_spill_dir = spill_dir _worker_sample_cap = sample_cap + _worker_time_index = time_index def _write_one_shard( @@ -2132,8 +2190,10 @@ def _fill_and_write_shard( if sample_cap > 0: stats = shard_stretch_stats(emb_buf, scales_buf, sample_cap) - if stats is not None: - return stats + if stats is None: + stats = _empty_shard_stats() + stats["coord"] = (spec.sr, spec.sc) + return stats return True @@ -2148,6 +2208,47 @@ def _write_one_shard_worker(spec: ShardSpec) -> "bool | Dict[str, Any]": ) +def _empty_shard_stats() -> Dict[str, Any]: + """Zero-contribution statistics for a shard with no valid pixels. + + Folding zeros is harmless, and returning them (rather than nothing) + lets the parent mark the shard seen so it is never re-read. + """ + return { + "n": 0, + "sum": np.zeros(N_BANDS, dtype=np.float64), + "prod": np.zeros((N_BANDS, N_BANDS), dtype=np.float64), + "sample_emb": np.zeros((0, N_BANDS), dtype=np.int8), + "sample_scales": np.zeros(0, dtype=np.float32), + "sample_weight": 1.0, + } + + +def _stats_catchup_worker(coord: Tuple[int, int]) -> Dict[str, Any]: + """Compute stretch statistics for a shard already in the store. + + The catch-up half of automatic collection: reads the shard back from the + zone arrays (one full-shard read — the price of a shard written before + its statistics were) and returns the same result shape as a write task. + """ + sr, sc = coord + emb_arr = _worker_store["embeddings"] + H, W = emb_arr.shape[2], emb_arr.shape[3] + r0, c0 = sr * SHARD_SIZE, sc * SHARD_SIZE + r1, c1 = min(r0 + SHARD_SIZE, H), min(c0 + SHARD_SIZE, W) + t = _worker_time_index + + stats = shard_stretch_stats( + np.asarray(emb_arr[t, :, r0:r1, c0:c1]), + np.asarray(_worker_store["scales"][t, r0:r1, c0:c1]), + max(_worker_sample_cap, 1), + ) + if stats is None: + stats = _empty_shard_stats() + stats["coord"] = coord + return stats + + # --------------------------------------------------------------------------- # Fill orchestration (zarr-fill) # --------------------------------------------------------------------------- @@ -2276,8 +2377,6 @@ def extend_store( storage_options: Optional[Dict[str, Any]] = None, zones: Optional[List[int]] = None, consolidate: bool = True, - force: bool = False, - state_url: Optional[str] = None, ) -> int: """Append new years to an existing store's time axis. @@ -2295,20 +2394,11 @@ def extend_store( zone, so no fill may be in flight. Returns the number of zone groups extended. """ - store = StoreLocation.resolve(store_path, storage_options, state_url) + store = StoreLocation.resolve(store_path, storage_options) years = sorted(set(int(y) for y in years)) if not years: raise ValueError("No years given to add") - held = [Path(p).name for p in store.state.listdir(LOCK_DIR_NAME, on_denied=[])] - if held and not force: - raise RuntimeError( - f"{len(held)} fill lock(s) present ({', '.join(sorted(held)[:4])}" - f"{'...' if len(held) > 4 else ''}). Extending rewrites array " - f"metadata for every zone, so wait for the sweep to finish. " - f"Use --force if these are stale." - ) - zone_names = _zone_group_names(store, zones) if not zone_names: raise ValueError(f"No UTM zone groups found in {store}") @@ -2663,8 +2753,6 @@ def fill_store( storage_options: Optional[Dict[str, Any]] = None, source: Optional[TileSource] = None, consolidate: Optional[bool] = None, - force_lock: bool = False, - state_url: Optional[str] = None, skip_existing_shards: bool = True, spill_dir: Optional[str] = None, collect_stretch_stats: bool = True, @@ -2685,7 +2773,6 @@ def fill_store( Defaults to True for a whole-store fill and False when ``zones`` is set, because the root object is the one thing parallel zone jobs share — run ``zarr-consolidate`` once after the sweep. - force_lock: Take over a (zone, year) lock held by another process. skip_existing_shards: Scan for shards already in the store and skip them (the default). A shard is always written from every tile covering it, so its presence means it is complete, and the @@ -2696,7 +2783,7 @@ def fill_store( falls inside an existing shard and would otherwise be skipped rather than merged in. """ - store = StoreLocation.resolve(store_path, storage_options, state_url) + store = StoreLocation.resolve(store_path, storage_options) if workers is None: workers = DEFAULT_WORKERS if consolidate is None: @@ -2719,10 +2806,6 @@ def fill_store( total_shards_written = 0 total_shards_failed = 0 - # The merged registry spans the whole store, so fetch it once rather - # than per zone and year. - merged_registry = load_merged_registry(store) - for fill_year in fill_years: if fill_year not in all_years: if console: @@ -2755,26 +2838,6 @@ def fill_store( ) continue - # Check which tiles are already written - written = _get_written_tiles( - store, fill_year, zone_num, merged=merged_registry - ) - remaining = [ti for ti in tile_infos if (ti.lon, ti.lat) not in written] - - if not remaining: - if console: - console.print( - f" Zone {zone_num} year {fill_year}: " - f"all {len(tile_infos)} tiles already written" - ) - continue - - if console: - console.print( - f" Zone {zone_num} year {fill_year}: " - f"{len(remaining)}/{len(tile_infos)} tiles to write" - ) - # Resolve the time index against *this* zone's own axis. An # interrupted zarr-extend can leave zones with different lengths, # and a store-wide index would then address the wrong year. @@ -2803,18 +2866,17 @@ def fill_store( height_px=shape[0], ) - # A shard write replaces the whole shard, so every shard we touch - # must be rebuilt from all of its tiles — including ones an - # earlier run already wrote, which would otherwise be zeroed. - touched = shard_coords_for_tiles(remaining, grid) - shard_specs = build_shard_index( - tile_infos, grid, time_index, restrict_to=touched - ) - - # The shard objects in the store are the ground truth for what - # landed — unlike the ingestion registry they survive a kill -9, - # so a crashed run can be resumed by scanning for them. - skipped_specs: List[ShardSpec] = [] + # Every shard the manifest implies, rebuilt from all of its + # tiles. A shard write replaces the whole shard, so a spec always + # carries every overlapping tile — never a delta. + shard_specs = build_shard_index(tile_infos, grid, time_index) + n_land = len(shard_specs) + + # The store is the only state: shard objects that exist are done. + # No registry, no locks — a spot instance that dies mid-run left + # nothing that needs cleaning up, and the next run's scan resumes + # exactly where the objects stop. + present: set = set() if skip_existing_shards: present = _existing_shards( store, @@ -2823,106 +2885,97 @@ def fill_store( {(s.sr, s.sc) for s in shard_specs}, console=console, ) - if present: - skipped_specs = [ - s for s in shard_specs if (s.sr, s.sc) in present - ] - shard_specs = [ - s for s in shard_specs if (s.sr, s.sc) not in present - ] + shard_specs = [ + s for s in shard_specs if (s.sr, s.sc) not in present + ] - if console: - # Spell the arithmetic out. The count of shards to write is - # otherwise hard to reconcile with zarr-scan, which counts - # every land shard, whereas a fill only considers those - # covering tiles the registry has not already recorded. - n_land = len(shard_coords_for_tiles(tile_infos, grid)) - n_recorded = n_land - len(touched) - console.print( - f" Shards: {n_land:,} land, " - f"{n_recorded:,} recorded done, " - f"{len(skipped_specs):,} found in store, " - f"[bold]{len(shard_specs):,} to write[/bold] " - f"({workers} workers)" + # Stretch statistics. The coverage mask says which store shards + # are already folded into the sums; anything present but unseen + # is a catch-up read. That makes stats collection idempotent and + # crash-safe by the same scan that drives the fill itself: + # a crash between shard write and stats fold just leaves the + # shard present-but-unseen, and the next run reads it back. + sample_cap = 0 + catch_up: List[Tuple[int, int]] = [] + if collect_stretch_stats: + zone_rw = store.open_group(mode="r+", path=zone_group) + ensure_stretch_arrays(zone_rw, console=console) + seen = { + (int(r), int(c)) + for r, c in zip( + *np.nonzero( + np.asarray(zone_rw["stretch_stats_shards"][time_index]) + ) + ) + } + writing = {(s.sr, s.sc) for s in shard_specs} + catch_up = sorted(present - seen - writing) + k_slots = zone_rw["stretch_sample"].shape[1] + sample_cap = _shard_sample_cap( + k_slots, len(shard_specs) + len(catch_up) ) - # Stretch statistics: collect only when the zone has the arrays - # (stores initialised before the feature lack them; repair with - # --backfill-stretch-stats). Per-shard cap sized so the expected - # candidate pool is a few times K without ballooning the result - # queue. - sample_cap = 0 - if collect_stretch_stats and "stretch_sample" in zone_store: - k_slots = zone_store["stretch_sample"].shape[1] - sample_cap = _shard_sample_cap(k_slots, len(shard_specs)) - elif collect_stretch_stats and console: + if console: console.print( - f" [yellow]Zone {zone_num} has no stretch-statistics " - f"arrays (store predates them); skipping collection. " - f"Backfill later with --backfill-stretch-stats.[/yellow]" + f" Shards: {n_land:,} land, " + f"{len(present):,} in store, " + f"[bold]{len(shard_specs):,} to write[/bold]" + + ( + f", {len(catch_up):,} stats catch-up read(s)" + if catch_up + else "" + ) + + f" ({workers} workers)" ) - _acquire_zone_lock(store, zone_num, fill_year, force=force_lock) - try: - written_count, failed, shard_stats = _write_shards( - store=store, - zone_group=zone_group, - shard_specs=shard_specs, - workers=workers, - source_options=source.storage_options if source else None, - label=f" Zone {zone_num} y{fill_year}", - console=console, - spill_dir=spill_dir, - sample_cap=sample_cap, - ) + written_count, failed, shard_stats = _write_shards( + store=store, + zone_group=zone_group, + shard_specs=shard_specs, + workers=workers, + source_options=source.storage_options if source else None, + label=f" Zone {zone_num} y{fill_year}", + console=console, + spill_dir=spill_dir, + sample_cap=sample_cap, + stats_coords=catch_up, + time_index=time_index, + ) - total_shards_written += written_count - total_shards_failed += len(failed) - - if shard_stats: - zone_rw = store.open_group(mode="r+", path=zone_group) - update_zone_stretch_stats( - zone_rw, - time_index, - n=sum(st["n"] for st in shard_stats), - s=sum(st["sum"] for st in shard_stats), - m=sum(st["prod"] for st in shard_stats), - sample_candidates=[ - (st["sample_emb"], st["sample_scales"], st["sample_weight"]) - for st in shard_stats - ], - ) - if console: - console.print( - f" [dim]Stretch stats: " - f"{sum(st['n'] for st in shard_stats):,} pixels " - f"folded in[/dim]" - ) + total_shards_written += written_count + total_shards_failed += len(failed) + if collect_stretch_stats and shard_stats: + # Sums before mask: a crash in between re-folds those shards + # next run (double count, drift-detectable and repairable) + # rather than silently dropping them. + update_zone_stretch_stats( + zone_rw, + time_index, + n=sum(st["n"] for st in shard_stats), + s=sum(st["sum"] for st in shard_stats), + m=sum(st["prod"] for st in shard_stats), + sample_candidates=[ + (st["sample_emb"], st["sample_scales"], st["sample_weight"]) + for st in shard_stats + if st["n"] > 0 + ], + seen_coords=[st["coord"] for st in shard_stats], + ) if console: console.print( - f" [green]{written_count}/{len(shard_specs)} " - f"shards written[/green]" + f" [dim]Stretch stats: " + f"{sum(st['n'] for st in shard_stats):,} pixels " + f"folded in[/dim]" ) - if failed: - console.print( - f" [red]{len(failed)} shard(s) failed[/red]" - ) - # Record tiles whose shards all landed — counting the ones we - # skipped as landed, since they are already in the store — so - # a retry picks up exactly the work still outstanding. - done = {(s.sr, s.sc) for s in skipped_specs} | ( - {(s.sr, s.sc) for s in shard_specs} - failed + if console: + console.print( + f" [green]{written_count}/{len(shard_specs)} " + f"shards written[/green]" ) - recorded = [ - ti - for ti in remaining - if shard_coords_for_tiles([ti], grid) <= done - ] - _record_written_tiles(store, recorded, fill_year, zone_num) - finally: - _release_zone_lock(store, zone_num, fill_year) + if failed: + console.print(f" [red]{len(failed)} shard(s) failed[/red]") # A failed shard leaves its tiles unrecorded, so re-running finishes the # job. Surface it as an error rather than a quiet partial success — a @@ -2959,11 +3012,18 @@ def _write_shards( console: Optional["rich.console.Console"], spill_dir: Optional[str] = None, sample_cap: int = 0, + stats_coords: Optional[List[Tuple[int, int]]] = None, + time_index: int = 0, ) -> Tuple[int, set, List[Dict[str, Any]]]: - """Run the shard writes through a process pool. + """Run shard writes — and stats catch-up reads — through a process pool. + + ``stats_coords`` are shards already in the store whose statistics the + coverage mask has not seen; they are read back and folded alongside the + writes. A failed catch-up is only a warning (the mask stays unset, so + the next run retries it); a failed write is a hard error as before. - Returns (shards written, set of (sr, sc) that failed, per-shard stretch - statistics — empty when collection is off or no shard had valid pixels). + Returns (shards written, set of (sr, sc) writes that failed, per-shard + stretch statistics from both task kinds). """ import multiprocessing from concurrent.futures import ProcessPoolExecutor, as_completed @@ -2971,6 +3031,7 @@ def _write_shards( written_count = 0 failed: set = set() stats_results: List[Dict[str, Any]] = [] + stats_coords = stats_coords or [] initargs = ( store.url, zone_group, @@ -2978,6 +3039,7 @@ def _write_shards( source_options, spill_dir, sample_cap, + time_index, ) # "spawn", not the Linux default of "fork": a forked worker inherits the @@ -2988,19 +3050,29 @@ def _write_shards( def _drain(pool, advance=None): nonlocal written_count futures = { - pool.submit(_write_one_shard_worker, spec): spec for spec in shard_specs + pool.submit(_write_one_shard_worker, spec): ("write", spec) + for spec in shard_specs } + futures.update( + { + pool.submit(_stats_catchup_worker, coord): ("stats", coord) + for coord in stats_coords + } + ) for future in as_completed(futures): - spec = futures[future] + kind, item = futures[future] try: result = future.result() - if result: + if kind == "write" and result: written_count += 1 if isinstance(result, dict): stats_results.append(result) except Exception as e: - logger.warning(f"Shard ({spec.sr},{spec.sc}) failed: {e}") - failed.add((spec.sr, spec.sc)) + if kind == "write": + logger.warning(f"Shard ({item.sr},{item.sc}) failed: {e}") + failed.add((item.sr, item.sc)) + else: + logger.warning(f"Stats catch-up for shard {item} failed: {e}") if advance is not None: advance() @@ -3035,7 +3107,9 @@ def _drain(pool, advance=None): TimeRemainingColumn(), console=console, ) as progress: - task = progress.add_task(label, total=len(shard_specs)) + task = progress.add_task( + label, total=len(shard_specs) + len(stats_coords) + ) _drain(pool, advance=lambda: progress.advance(task)) else: _drain(pool) @@ -3569,21 +3643,19 @@ def backfill_stretch_stats( sample_k: int = STRETCH_SAMPLE_K, console: Optional["rich.console.Console"] = None, storage_options: Optional[Dict[str, Any]] = None, - state_url: Optional[str] = None, - force_lock: bool = False, ) -> int: """Rebuild a zone's stretch statistics by scanning its existing shards. - The repair path for stores filled before fill-time collection existed, - for interrupted fills, and for suspected double-counting: it re-reads - the zone's shards once (the only stats path that touches embeddings) and - *sets* the arrays from what is actually in the store. Creates the arrays - if the zone predates them. Per-zone, so it composes with fills of other - zones; takes the same (zone, year) lock a fill would. + Normally unnecessary — fills create the arrays and catch up on unseen + shards automatically via the coverage mask. This is the explicit repair + for suspected double-counting (the drift check pointing here): it + re-reads the zone's shards once and *sets* the arrays and mask from what + is actually in the store, discarding the running sums. Do not run it + while a fill is writing the same zone. Returns the number of (zone, year) slots rebuilt. """ - store = StoreLocation.resolve(store_path, storage_options, state_url) + store = StoreLocation.resolve(store_path, storage_options) zone_names = _zone_group_names(store, zones) if not zone_names: raise ValueError(f"No UTM zone groups found in {store}") @@ -3592,12 +3664,8 @@ def backfill_stretch_stats( for name in zone_names: group = store.open_group(mode="r+", path=name, zarr_format=3) zone_years = [int(v) for v in group["time"][:]] - T = len(zone_years) - if "stretch_sample" not in group: - create_stretch_arrays(group, T, sample_k) - if console: - console.print(f" {name}: created stretch-statistics arrays") + ensure_stretch_arrays(group, console=console, sample_k=sample_k) k_slots = group["stretch_sample"].shape[1] emb_arr = group["embeddings"] @@ -3618,56 +3686,55 @@ def backfill_stretch_stats( continue cap = _shard_sample_cap(k_slots, len(present)) - zone_num = int(name[3:]) - _acquire_zone_lock(store, zone_num, fill_year, force=force_lock) - try: - n_total, s_total = 0, np.zeros(N_BANDS, dtype=np.float64) - m_total = np.zeros((N_BANDS, N_BANDS), dtype=np.float64) - candidates: List[Tuple[np.ndarray, np.ndarray, float]] = [] - for i, (sr, sc) in enumerate(sorted(present)): - r0, c0 = sr * SHARD_SIZE, sc * SHARD_SIZE - r1, c1 = min(r0 + SHARD_SIZE, H), min(c0 + SHARD_SIZE, W) - st = shard_stretch_stats( - np.asarray(emb_arr[t, :, r0:r1, c0:c1]), - np.asarray(scales_arr[t, r0:r1, c0:c1]), - cap, - ) - if st is None: - continue - n_total += st["n"] - s_total += st["sum"] - m_total += st["prod"] - candidates.append( - (st["sample_emb"], st["sample_scales"], st["sample_weight"]) - ) - if console: - console.print( - f" {name} {fill_year}: shard {i + 1}/{len(present)} " - f"({st['n']:,} px)", - end="\r", - ) - - emb_s, scales_s = merge_stretch_samples(candidates, k_slots) - # Backfill SETS from actual contents (it is the repair for - # double-counting), unlike the fill's additive fold. - group["stretch_stats_count"][t] = n_total - group["stretch_stats_sum"][t] = s_total - group["stretch_stats_prod"][t] = m_total - full_emb = np.zeros((k_slots, N_BANDS), dtype=np.int8) - full_sc = np.full(k_slots, np.float32("inf"), dtype=np.float32) - full_emb[: len(emb_s)] = emb_s - full_sc[: len(emb_s)] = scales_s - group["stretch_sample"][t] = full_emb - group["stretch_sample_scales"][t] = full_sc - group["stretch_sample_count"][t] = len(emb_s) - rebuilt += 1 + n_total, s_total = 0, np.zeros(N_BANDS, dtype=np.float64) + m_total = np.zeros((N_BANDS, N_BANDS), dtype=np.float64) + candidates: List[Tuple[np.ndarray, np.ndarray, float]] = [] + for i, (sr, sc) in enumerate(sorted(present)): + r0, c0 = sr * SHARD_SIZE, sc * SHARD_SIZE + r1, c1 = min(r0 + SHARD_SIZE, H), min(c0 + SHARD_SIZE, W) + st = shard_stretch_stats( + np.asarray(emb_arr[t, :, r0:r1, c0:c1]), + np.asarray(scales_arr[t, r0:r1, c0:c1]), + cap, + ) + if st is None: + continue + n_total += st["n"] + s_total += st["sum"] + m_total += st["prod"] + candidates.append( + (st["sample_emb"], st["sample_scales"], st["sample_weight"]) + ) if console: console.print( - f" {name} {fill_year}: rebuilt from " - f"{len(present)} shard(s), {n_total:,} pixels " + f" {name} {fill_year}: shard {i + 1}/{len(present)} " + f"({st['n']:,} px)", + end="\r", ) - finally: - _release_zone_lock(store, zone_num, fill_year) + + emb_s, scales_s = merge_stretch_samples(candidates, k_slots) + # Backfill SETS from actual contents (it is the repair for + # double-counting), unlike the fill's additive fold. + group["stretch_stats_count"][t] = n_total + group["stretch_stats_sum"][t] = s_total + group["stretch_stats_prod"][t] = m_total + full_emb = np.zeros((k_slots, N_BANDS), dtype=np.int8) + full_sc = np.full(k_slots, np.float32("inf"), dtype=np.float32) + full_emb[: len(emb_s)] = emb_s + full_sc[: len(emb_s)] = scales_s + group["stretch_sample"][t] = full_emb + group["stretch_sample_scales"][t] = full_sc + group["stretch_sample_count"][t] = len(emb_s) + mask = np.zeros(group["stretch_stats_shards"].shape[1:], np.uint8) + for sr, sc in present: + mask[sr, sc] = 1 + group["stretch_stats_shards"][t] = mask + rebuilt += 1 + if console: + console.print( + f" {name} {fill_year}: rebuilt from " + f"{len(present)} shard(s), {n_total:,} pixels " + ) return rebuilt diff --git a/tests/zarr_remote_check.py b/tests/zarr_remote_check.py index efeae3c..3237d4e 100644 --- a/tests/zarr_remote_check.py +++ b/tests/zarr_remote_check.py @@ -20,16 +20,13 @@ from geotessera import remote from geotessera.zarr import ( - LOCK_DIR_NAME, REGISTRY_DIR_NAME, StoreLocation, TileInfo, TileSource, UnifiedZoneGrid, - _acquire_zone_lock, _get_written_tiles, _record_written_tiles, - _release_zone_lock, build_shard_index, merge_tile_registry, shard_coords_for_tiles, @@ -291,37 +288,6 @@ def tile(lon, lat, zone=31): == {(0.05, 52.05), (0.15, 52.05), (0.25, 52.05)}, ) -# --------------------------------------------------------------------------- -# Advisory zone locks -# --------------------------------------------------------------------------- - -lock_store = StoreLocation(str(TMP / "lock_store")) -_acquire_zone_lock(lock_store, 31, 2024) -check( - "lock object created in the state sibling", - lock_store.state.exists(LOCK_DIR_NAME, "utm31_2024.json"), -) -check("no lock object inside the store", not lock_store.exists(LOCK_DIR_NAME)) - -try: - _acquire_zone_lock(lock_store, 31, 2024) - check("second acquire refused", False) -except RuntimeError as e: - check("second acquire refused", "locked by" in str(e)) - -# A different zone or year is a different lock, so sibling jobs proceed. -_acquire_zone_lock(lock_store, 30, 2024) -_acquire_zone_lock(lock_store, 31, 2023) -check("sibling zone/year locks independent", True) - -_acquire_zone_lock(lock_store, 31, 2024, force=True) -check("force takes over a stale lock", True) - -_release_zone_lock(lock_store, 31, 2024) -check( - "release removes the lock", not lock_store.exists(LOCK_DIR_NAME, "utm31_2024.json") -) - # --------------------------------------------------------------------------- # Shard index: rewriting a shard must carry its already-written neighbours # --------------------------------------------------------------------------- @@ -422,7 +388,7 @@ def placed(lon, x_off): ) for zname in ("utm30", "utm31"): - create_stretch_arrays(root[zname], n_years=1, k=50) + create_stretch_arrays(root[zname], n_years=1, k=50, n_shard_rows=1, n_shard_cols=1) check("extend adds the year to every zone", extend_store(ext, [2026]) == 2) check( @@ -456,15 +422,7 @@ def placed(lon, x_off): except ValueError as e: check("inserting an earlier year refused", "only be appended" in str(e)) -ext.state.write_bytes(b"{}", LOCK_DIR_NAME, "utm30_2026.json") -try: - extend_store(ext, [2027]) - check("extend refuses while a fill lock is held", False) -except RuntimeError as e: - check("extend refuses while a fill lock is held", "fill lock" in str(e)) -check( - "extend --force overrides a stale lock", extend_store(ext, [2027], force=True) == 2 -) +check("extend appends a further year", extend_store(ext, [2027]) == 2) # --------------------------------------------------------------------------- # Stretch statistics @@ -480,15 +438,16 @@ def placed(lon, x_off): nprng = np.random.default_rng(11) B = 128 semb = nprng.integers(-128, 127, (B, 96, 96), dtype=np.int8) -ssc = (nprng.random((96, 96)).astype(np.float32) * 0.01 + 0.001) +ssc = nprng.random((96, 96)).astype(np.float32) * 0.01 + 0.001 ssc[:20, :20] = np.nan ssc[80:, 80:] = np.inf sst = shard_stretch_stats(semb, ssc, sample_cap=200, seed=5) svalid = np.isfinite(ssc) -sx = semb.reshape(B, -1)[:, svalid.ravel()].astype(np.float64) * ssc.ravel()[ - svalid.ravel() -] +sx = ( + semb.reshape(B, -1)[:, svalid.ravel()].astype(np.float64) + * ssc.ravel()[svalid.ravel()] +) check("stats count exact", sst["n"] == int(svalid.sum())) check( "stats sum matches population", @@ -537,15 +496,24 @@ def placed(lon, x_off): import zarr # noqa: E402 zs = zarr.open_group(str(TMP / "stats.zarr"), mode="w", zarr_format=3) -create_stretch_arrays(zs, n_years=2, k=100) +create_stretch_arrays(zs, n_years=2, k=100, n_shard_rows=2, n_shard_cols=3) check( "stretch arrays created", all(n in zs for n in STRETCH_ARRAY_NAMES), ) cand = [(sst["sample_emb"], sst["sample_scales"], sst["sample_weight"])] -update_zone_stretch_stats(zs, 0, sst["n"], sst["sum"], sst["prod"], cand, seed=1) -update_zone_stretch_stats(zs, 0, sst["n"], sst["sum"], sst["prod"], cand, seed=2) +update_zone_stretch_stats( + zs, 0, sst["n"], sst["sum"], sst["prod"], cand, seen_coords=[(0, 1)], seed=1 +) +update_zone_stretch_stats( + zs, 0, sst["n"], sst["sum"], sst["prod"], cand, seen_coords=[(1, 2)], seed=2 +) check("stats fold additively", int(zs["stretch_stats_count"][0]) == 2 * sst["n"]) +mask0 = np.asarray(zs["stretch_stats_shards"][0]) +check( + "coverage mask accumulates seen shards", + mask0[0, 1] == 1 and mask0[1, 2] == 1 and int(mask0.sum()) == 2, +) check( "sample capacity bounded", int(zs["stretch_sample_count"][0]) <= 100, From 0427e3d8b9b31f9e30f2622e0584c4f44448cf53 Mon Sep 17 00:00:00 2001 From: Anil Madhavapeddy Date: Fri, 7 Aug 2026 16:57:51 +0100 Subject: [PATCH 13/13] update for remote stretch --- geotessera/registry_cli.py | 89 +++- geotessera/remote.py | 11 +- geotessera/zarr.py | 830 ++++++++++++++++++++++++++----------- pyproject.toml | 2 +- tests/zarr_remote_check.py | 255 ++++++++++++ uv.lock | 151 +------ 6 files changed, 936 insertions(+), 402 deletions(-) diff --git a/geotessera/registry_cli.py b/geotessera/registry_cli.py index b215476..100a335 100644 --- a/geotessera/registry_cli.py +++ b/geotessera/registry_cli.py @@ -3976,9 +3976,9 @@ def _require_local_store(store_path: str, command: str, console: "Console") -> b console.print( f"[red]{emoji('❌ ')}{command} needs a local store path; " f"got {store_path}.[/red]\n" - f"Copy the store locally (or mount it) and run it there — only " - f"zarr-init, zarr-fill and zarr-consolidate work against remote " - f"locations." + f"Only the legacy shard-sampling path is local-only — the default " + f"stats-based stretch and zarr-global-preview --output both work " + f"against remote stores." ) return False @@ -3992,23 +3992,42 @@ def zarr_global_preview_command(args): warnings.filterwarnings("ignore", message="Object at .* is not recognized") console = Console() - if not _require_local_store(args.store_path, "zarr-global-preview", console): - return 1 - store_path = Path(args.store_path) - zones = _parse_int_range(args.zones) if args.zones else None - build_global_preview( - store_path=store_path, - year=args.year, - zones=zones, - num_levels=args.levels, - workers=args.workers, - gamma=args.gamma, - saturation=args.saturation, - console=console, - force=args.force, + zones = _parse_int_range(args.zones) if args.zones else None + store_options = _storage_options_for(args, "store", args.store_path) + output_options = ( + _storage_options_for(args, "output", args.output) if args.output else None ) + try: + build_global_preview( + store_path=args.store_path, + year=args.year, + zones=zones, + num_levels=args.levels, + workers=args.workers, + gamma=args.gamma, + saturation=args.saturation, + console=console, + force=args.force, + storage_options=store_options, + output_path=args.output, + output_storage_options=output_options, + state_url=args.state_url, + state_storage_options=( + _storage_options_for(args, "state", args.state_url) + if args.state_url + else None + ), + reproject_only=args.reproject_only, + coarsen_only=args.coarsen_only, + ) + except (ValueError, RuntimeError) as e: + console.print(f"[red]{emoji('❌ ')}{e}[/red]") + return 1 + except (ImportError, *_object_store_errors()) as e: + return _report_store_error(e, console) + return 0 @@ -5152,6 +5171,42 @@ def main(): "Try 1.5–2.5 if colours look washed out. Beyond ~3 most pixels " "start clipping at the colour-cube edges.", ) + zarr_gp_parser.add_argument( + "--output", + type=str, + default=None, + help="Path or URL of the store to hold the pyramid, with its own " + "--output-* credentials. Embeddings stream from the source store " + "read-only (sub-shard byte ranges, no copy), so this can write to a " + "bucket while reading from a read-only mirror. Default: write the " + "pyramid into the source store itself.", + ) + zarr_gp_parser.add_argument( + "--reproject-only", + action="store_true", + help="Render this zone's level 0 and coarsen only as far as the " + "depth where zones stay disjoint. Safe to run concurrently across " + "zones that share no level-0 chunks — an odd-numbered round then an " + "even-numbered one covers all 60. Finish with --coarsen-only.", + ) + zarr_gp_parser.add_argument( + "--coarsen-only", + action="store_true", + help="Build the remaining coarse levels in one global pass, without " + "reprojecting. Single-writer: run once after every zone's " + "--reproject-only has finished.", + ) + zarr_gp_parser.add_argument( + "--state-url", + type=str, + default=None, + help="Where per-zone resume markers live (default: .build, " + "alongside the pyramid). Point it at a local directory to keep " + "build bookkeeping out of a published bucket.", + ) + _add_storage_args(zarr_gp_parser, "store", "Store") + _add_storage_args(zarr_gp_parser, "output", "Pyramid output", writable=True) + _add_storage_args(zarr_gp_parser, "state", "Build state", writable=True) zarr_gp_parser.set_defaults(func=zarr_global_preview_command) # Zarr-stretch command diff --git a/geotessera/remote.py b/geotessera/remote.py index 66e7e9d..b51af14 100644 --- a/geotessera/remote.py +++ b/geotessera/remote.py @@ -50,9 +50,16 @@ "aiobotocore", "s3fs", "urllib3", - "aiohttp", ) +# These two emit "Unclosed client session"/"Unclosed connector" at ERROR +# from aiohttp object destructors, via asyncio's exception handler, whenever +# an fsspec event loop is torn down without a graceful session close — which +# fsspec does not do. Harmless garbage-collection noise, but at ERROR it +# survives a WARNING cap and lands all over the progress bars, so these two +# are silenced outright. Real S3 failures surface as exceptions, not logs. +_DESTRUCTOR_NOISE_LOGGERS = ("aiohttp", "asyncio") + def quieten_dependency_logging(level: int = logging.WARNING) -> None: """Raise the log level of the object-store libraries. @@ -63,6 +70,8 @@ def quieten_dependency_logging(level: int = logging.WARNING) -> None: """ for name in _NOISY_LOGGERS: logging.getLogger(name).setLevel(level) + for name in _DESTRUCTOR_NOISE_LOGGERS: + logging.getLogger(name).setLevel(logging.CRITICAL) def reset_after_fork() -> None: diff --git a/geotessera/zarr.py b/geotessera/zarr.py index fdf6654..e2d4586 100644 --- a/geotessera/zarr.py +++ b/geotessera/zarr.py @@ -85,6 +85,15 @@ GLOBAL_NUM_BANDS = 4 GLOBAL_DEFAULT_LEVELS = 10 +# Deepest pyramid level a per-zone pass may coarsen while other zones run +# concurrently. Each level halves the region, so zone strips that start far +# apart converge until they share a 512px chunk. Measured against the real +# written-shard footprints, an odd/even zone split has no overlapping pairs +# through level 6 and starts colliding at level 7 (9 pairs, then 21 at level +# 8 and 36 at level 9). Levels above this must be built by a single global +# pass — see ``--coarsen-only``. +COARSEN_PARALLEL_SAFE_LEVEL = 6 + # GeoZarr convention registration entries GEOEMB_CONVENTION = { "uuid": "61c12cc5-0e28-4056-999a-480cf3fb7e4c", @@ -94,6 +103,49 @@ "schema_url": "https://raw.githubusercontent.com/geo-embeddings/embeddings-zarr-convention/refs/tags/v1/schema.json", } +# Revisions of the shared conventions to stamp into ``zarr_conventions``. +# zarr-cm pins each revision to the spec commit that defined it, so these +# resolve where a tag-based URL does not. ``multiscales`` has no r3. +SPATIAL_REVISION = "r3" +PROJ_REVISION = "r3" +MULTISCALES_REVISION = "r2" + + +def _geo_convention_attrs( + dimensions: List[str], + crs: str, + bbox: List[float], + transform: Optional[List[float]] = None, + shape: Optional[List[int]] = None, + registration: str = "pixel", +) -> Dict[str, Any]: + """Build ``spatial:`` and ``proj:`` attrs plus their registrations. + + Arguments left as None are omitted rather than written as nulls, so a + group that has no single affine transform (the multiscale pyramid root, + whose geometry is per level) carries only the attributes it can state. + """ + from zarr_cm import geo_proj, spatial + + attrs: Dict[str, Any] = {} + attrs = spatial.insert( + attrs, + spatial.create( + revision=SPATIAL_REVISION, + dimensions=dimensions, + bbox=bbox, + transform_type="affine", + transform=transform, + shape=shape, + registration=registration, + ), + ) + attrs = geo_proj.insert( + attrs, + geo_proj.create(revision=PROJ_REVISION, code=crs), + ) + return attrs + # --------------------------------------------------------------------------- # Data types (shared) @@ -714,100 +766,153 @@ def _execute(pool): # --------------------------------------------------------------------------- -def _preview_marker_path(store_path: Path, zone_num: int) -> Path: - """Resume marker for a zone's global-preview reprojection. +def _preview_marker_parts(zone_num: int) -> Tuple[str, str]: + """Key of a zone's global-preview resume marker within the state area. - Kept in the state sibling (``.build/_preview/``) rather than the - store, for the same reason as the ingestion registry: the published Zarr - hierarchy should contain only Zarr. + Markers live in the state sibling (``.build/_preview/``) rather + than the store, for the same reason as the ingestion registry: the + published Zarr hierarchy should contain only Zarr. Returned relative to + :attr:`StoreLocation.state` so the same marker works on a local pyramid + and one on object storage. """ - return Path(f"{str(store_path).rstrip('/')}.build") / "_preview" / ( - f"zone_{zone_num}_done" - ) + return ("_preview", f"zone_{zone_num}_done") -def _zone_output_bounds( +def _chunks_for_shards( + present: set, zone_epsg: int, zone_transform: list, zone_shape: tuple, -) -> Tuple[int, int, int, int]: - """Compute the chunk-aligned output bounds for a zone in global grid pixels. - - Returns (row_start, row_end, col_start, col_end). +) -> Tuple[set, List[Tuple[int, int, int, int]]]: + """Output chunks that can actually receive data, from present shards. + + A zone's *bounding rectangle* back-projected to lon/lat is hopeless as a + work list: a high-latitude zone spans most longitudes at its top edge, + which for utm02 meant ~5.6 million candidate chunks to render 28 shards. + Projecting each present shard's footprint instead yields only the chunks + its data can touch — the same objects-are-the-truth move the fill makes. + + Returns (chunk (row, col) set, list of (row_start, row_end, col_start, + col_end) rectangles in level-0 pixels, chunk-aligned). The rectangles + cover the same chunks and drive the pyramid coarsening; see + :func:`_regions_for_chunks` for why there can be more than one. """ from pyproj import Transformer + if not present: + return set(), (0, 0, 0, 0) + src_pixel = zone_transform[0] - src_origin_e = zone_transform[2] - src_origin_n = zone_transform[5] - src_h, src_w = zone_shape[:2] + origin_e = zone_transform[2] + origin_n = zone_transform[5] + H, W = zone_shape[:2] + west, _s, _e, north = GLOBAL_BOUNDS + + to_4326 = Transformer.from_crs(f"EPSG:{zone_epsg}", "EPSG:4326", always_xy=True) + + chunks: set = set() + for sr, sc in present: + r0, c0 = sr * SHARD_SIZE, sc * SHARD_SIZE + r1, c1 = min(r0 + SHARD_SIZE, H), min(c0 + SHARD_SIZE, W) + es = [origin_e + c0 * src_pixel, origin_e + c1 * src_pixel] + ns = [origin_n - r0 * src_pixel, origin_n - r1 * src_pixel] + # Corners plus edge midpoints: enough to bound the curved footprint. + pts_e = [es[0], es[1], es[0], es[1], (es[0] + es[1]) / 2, + (es[0] + es[1]) / 2, es[0], es[1]] + pts_n = [ns[0], ns[0], ns[1], ns[1], ns[0], ns[1], + (ns[0] + ns[1]) / 2, (ns[0] + ns[1]) / 2] + lons, lats = to_4326.transform(pts_e, pts_n) + finite = [(lo, la) for lo, la in zip(lons, lats) + if math.isfinite(lo) and math.isfinite(la)] + if not finite: + continue + la_min = min(p[1] for p in finite) + la_max = max(p[1] for p in finite) + + # A shard straddling the antimeridian samples corners near -180 and + # +180, whose naive min/max spans the globe and would enqueue every + # chunk column at that latitude — for utm01/utm60 that was ~2.3M + # bogus chunks, a third of a year's work list. Re-measuring the span + # with longitudes shifted to [0, 360) detects the wrap, because that + # frame's discontinuity is at 0 rather than at the antimeridian; the + # footprint is then contiguous and splits into two column ranges. + # Near a pole a shard genuinely does span most longitudes and both + # frames stay wide, so the full range is kept. + lons_f = [p[0] for p in finite] + lo_min, lo_max = min(lons_f), max(lons_f) + shifted = [lo % 360.0 for lo in lons_f] + sh_min, sh_max = min(shifted), max(shifted) + if (lo_max - lo_min) > 180.0 and (sh_max - sh_min) < 180.0: + lon_ranges = [(sh_min, 180.0), (-180.0, sh_max - 360.0)] + else: + lon_ranges = [(lo_min, lo_max)] - west, _south, _east, north = GLOBAL_BOUNDS + cr0 = max(0, int((north - la_max) / GLOBAL_BASE_RES) // GLOBAL_CHUNK - 1) + cr1 = min( + GLOBAL_LEVEL0_H // GLOBAL_CHUNK, + int((north - la_min) / GLOBAL_BASE_RES) // GLOBAL_CHUNK + 2, + ) + # One chunk of margin absorbs footprint curvature between samples. + for lo_a, lo_b in lon_ranges: + cc0 = max(0, int((lo_a - west) / GLOBAL_BASE_RES) // GLOBAL_CHUNK - 1) + cc1 = min( + GLOBAL_LEVEL0_W // GLOBAL_CHUNK, + int((lo_b - west) / GLOBAL_BASE_RES) // GLOBAL_CHUNK + 2, + ) + for cr in range(cr0, cr1): + for cc in range(cc0, cc1): + chunks.add((cr, cc)) - to_4326 = Transformer.from_crs( - f"EPSG:{zone_epsg}", - "EPSG:4326", - always_xy=True, - ) - corners_utm = [ - (src_origin_e, src_origin_n), - (src_origin_e + src_w * src_pixel, src_origin_n), - (src_origin_e, src_origin_n - src_h * src_pixel), - (src_origin_e + src_w * src_pixel, src_origin_n - src_h * src_pixel), - ] - mid_e = src_origin_e + src_w * src_pixel / 2 - mid_n = src_origin_n - src_h * src_pixel / 2 - corners_utm += [ - (mid_e, src_origin_n), - (mid_e, src_origin_n - src_h * src_pixel), - (src_origin_e, mid_n), - (src_origin_e + src_w * src_pixel, mid_n), - ] - corners_4326 = [to_4326.transform(e, n) for e, n in corners_utm] - lons = [c[0] for c in corners_4326] - lats = [c[1] for c in corners_4326] + return chunks, _regions_for_chunks(chunks) - zlon_min, zlon_max = min(lons), max(lons) - zlat_min, zlat_max = min(lats), max(lats) - col_start = max( - 0, - ( - int(math.floor((zlon_min - west) / GLOBAL_BASE_RES)) - // GLOBAL_CHUNK - * GLOBAL_CHUNK - ), - ) - col_end = min( - GLOBAL_LEVEL0_W, - ( - (int(math.ceil((zlon_max - west) / GLOBAL_BASE_RES)) + GLOBAL_CHUNK - 1) - // GLOBAL_CHUNK - * GLOBAL_CHUNK - ), - ) - row_start = max( - 0, - ( - int(math.floor((north - zlat_max) / GLOBAL_BASE_RES)) - // GLOBAL_CHUNK - * GLOBAL_CHUNK - ), - ) - row_end = min( - GLOBAL_LEVEL0_H, - ( - (int(math.ceil((north - zlat_min) / GLOBAL_BASE_RES)) + GLOBAL_CHUNK - 1) - // GLOBAL_CHUNK - * GLOBAL_CHUNK - ), - ) +def _regions_for_chunks(chunks: set) -> List[Tuple[int, int, int, int]]: + """Chunk-aligned pixel rectangles covering *chunks*, for the coarsening. - return (row_start, row_end, col_start, col_end) + Normally one rectangle. A zone straddling the antimeridian holds chunks + at both edges of the grid and none between, and a single enclosing + rectangle then spans every column: 16.7M chunk slots for utm60's 38k real + chunks, each of which the coarsening pass reads and rewrites. Splitting on + a column gap wider than half the grid keeps that case to two tight + rectangles and leaves every other zone with exactly one. + """ + if not chunks: + return [] + + cols = sorted({c for _r, c in chunks}) + split_at = None + if len(cols) > 1: + gap, idx = max((cols[i + 1] - cols[i], i) for i in range(len(cols) - 1)) + if gap > (GLOBAL_LEVEL0_W // GLOBAL_CHUNK) // 2: + split_at = cols[idx + 1] + + if split_at is None: + groups = [chunks] + else: + groups = [ + {(r, c) for r, c in chunks if c < split_at}, + {(r, c) for r, c in chunks if c >= split_at}, + ] + + regions = [] + for group in groups: + if not group: + continue + rows = [r for r, _c in group] + gcols = [c for _r, c in group] + regions.append( + ( + min(rows) * GLOBAL_CHUNK, + (max(rows) + 1) * GLOBAL_CHUNK, + min(gcols) * GLOBAL_CHUNK, + (max(gcols) + 1) * GLOBAL_CHUNK, + ) + ) + return regions def _coarsen_zone_pyramid( - store_path: Path, + dest: "StoreLocation", row_start: int, row_end: int, col_start: int, @@ -815,18 +920,22 @@ def _coarsen_zone_pyramid( num_levels: int, workers: int, console: Optional["rich.console.Console"] = None, + start_level: int = 1, ) -> None: - """Update pyramid levels 1 through num_levels-1 for the affected region. + """Update pyramid levels ``start_level``..``num_levels``-1 for a region. Reads from the previous level and writes coarsened data to the current level, processing in 2D tiles parallelised with a thread pool. + + Levels below *start_level* are still walked, because each level's region + is derived by halving the one above it, but no data is touched there. + That is what lets a parallel zone sweep stop at + :data:`COARSEN_PARALLEL_SAFE_LEVEL` and a single global pass pick up the + rest. """ from concurrent.futures import ThreadPoolExecutor, as_completed - import zarr - root = zarr.open_group( - str(store_path), mode="r+", zarr_format=3, use_consolidated=False - ) + root = dest.open_group(mode="r+", zarr_format=3) prev_row_start, prev_row_end = row_start, row_end prev_col_start, prev_col_end = col_start, col_end @@ -856,6 +965,12 @@ def _coarsen_zone_pyramid( if lr_end <= lr_start or lc_end <= lc_start: break + if lvl < start_level: + # Walk the region down to the starting level without touching it. + prev_row_start, prev_row_end = lr_start, lr_end + prev_col_start, prev_col_end = lc_start, lc_end + continue + if console is not None: console.print( f" Level {lvl}: rows {lr_start}-{lr_end}, cols {lc_start}-{lc_end}" @@ -943,6 +1058,20 @@ def _coarsen_tile( # Default per-(zone, year) capacity of the raw pixel sample kept for the # global stretch quantiles (docs/specs/zarr-stretch-stats.md). STRETCH_SAMPLE_K = 20_000 + +# Scales above this are treated as nodata. Some published scales files carry +# a huge-finite sentinel (~FLT_MAX) that passes isfinite(): real per-pixel +# quantisation scales here are O(0.1), and a sentinel pixel squared +# overflows float32 — utm04's stored sum reached 1e37 and its product +# matrix went inf from a handful of such pixels. Six orders of magnitude of +# headroom above anything plausible, thirty-two below the sentinel. +MAX_VALID_SCALE = 1.0e6 + + +def valid_scale_mask(scales: np.ndarray) -> np.ndarray: + """True where a scale denotes real data: finite, positive, plausible.""" + with np.errstate(invalid="ignore"): + return np.isfinite(scales) & (scales > 0) & (scales < MAX_VALID_SCALE) DEFAULT_WORKERS = 4 # fewer workers due to larger shard buffers (~2GB each) # Each shard worker holds a full (N_BANDS, SHARD_SIZE, SHARD_SIZE) int8 @@ -1442,40 +1571,28 @@ def _create_zone_group( math.ceil(W / SHARD_SIZE), ) - # Use geozarr-toolkit for proj: and spatial: convention metadata - from geozarr_toolkit import create_geozarr_attrs - x_min = grid.origin_x x_max = grid.origin_x + W * grid.pixel_size y_max = grid.origin_y y_min = grid.origin_y - H * grid.pixel_size - geozarr_attrs = create_geozarr_attrs( - dimensions=["y", "x"], - crs=f"EPSG:{grid.canonical_epsg}", - transform=[ - grid.pixel_size, - 0.0, - grid.origin_x, - 0.0, - -grid.pixel_size, - grid.origin_y, - ], - bbox=[x_min, y_min, x_max, y_max], - shape=[H, W], - registration="pixel", - ) - - # Fix convention descriptions to match upstream schemas exactly - # (geozarr-toolkit has a bug: "Spatial coordinate and transformation - # information" instead of "Spatial coordinate information") - for conv in geozarr_attrs.get("zarr_conventions", []): - if conv.get("uuid") == "689b58e2-cf7b-45e0-9fff-9cfc0883d6b4": - conv["description"] = "Spatial coordinate information" - # Zone groups only carry proj: and spatial: conventions (geoemb: is on root) - - store.attrs.update(geozarr_attrs) + store.attrs.update( + _geo_convention_attrs( + dimensions=["y", "x"], + crs=f"EPSG:{grid.canonical_epsg}", + transform=[ + grid.pixel_size, + 0.0, + grid.origin_x, + 0.0, + -grid.pixel_size, + grid.origin_y, + ], + bbox=[x_min, y_min, x_max, y_max], + shape=[H, W], + ) + ) return store @@ -1660,7 +1777,7 @@ def shard_stretch_stats( Returns None when the shard has no valid pixels. """ - valid = np.isfinite(scales_buf) + valid = valid_scale_mask(scales_buf) flat = np.flatnonzero(valid.ravel()) n = int(flat.size) if n == 0: @@ -2176,7 +2293,9 @@ def _fill_and_write_shard( storage_options=source_options, ) s[lm == 0] = np.float32("nan") - s[~np.isfinite(s)] = np.float32("nan") + # Non-finite, non-positive and sentinel-huge scales are all nodata; + # storing them as NaN keeps every reader's isfinite() test honest. + s[~valid_scale_mask(s)] = np.float32("nan") scales_buf[ov.s_row_start : ov.s_row_end, ov.s_col_start : ov.s_col_end] = s has_data = True @@ -3221,7 +3340,7 @@ def _compute_rgb_chunk( """ h, w = scales_hw.shape rgba = np.zeros((4, h, w), dtype=np.uint8) - valid = np.isfinite(scales_hw) + valid = valid_scale_mask(scales_hw) scales_safe = np.where(valid, scales_hw, 0.0) # Build the per-channel float arrays. Two paths: @@ -3308,7 +3427,7 @@ def _sample_chunk_stats( c0, c1 = cj * shard_size, min(cj * shard_size + shard_size, W) scales_chunk = np.asarray(scales_arr[time_index, r0:r1, c0:c1]) - valid = np.isfinite(scales_chunk) + valid = valid_scale_mask(scales_chunk) if not np.any(valid): return None @@ -3429,9 +3548,14 @@ def compute_stretch_from_stats( drift_threshold: float = 0.25, console: Optional["rich.console.Console"] = None, storage_options: Optional[Dict[str, Any]] = None, + persist: bool = True, ) -> dict: """Derive the global stretch from the per-zone ``stretch_*`` arrays. + With ``persist=False`` the stretch entry is returned without touching + the store — how a read-only consumer (a preview against a store it + cannot write to) gets a stretch from the live statistics. + The fast path of ``zarr-stretch`` (docs/specs/zarr-stretch-stats.md): reads a few MiB of per-zone summaries instead of terabytes of shards. The PCA comes from the summed sufficient statistics and is exact — every @@ -3458,6 +3582,7 @@ def compute_stretch_from_stats( sample_parts: List[Tuple[np.ndarray, np.ndarray, float]] = [] zones_used: List[str] = [] zones_missing: List[str] = [] + zones_poisoned: List[str] = [] for name in zone_names: group = store.open_group(mode="r", path=name) @@ -3473,9 +3598,14 @@ def compute_stretch_from_stats( n_z = int(group["stretch_stats_count"][t]) if n_z == 0: continue + s_z = np.asarray(group["stretch_stats_sum"][t]) + m_z = np.asarray(group["stretch_stats_prod"][t]) + if not (np.isfinite(s_z).all() and np.isfinite(m_z).all()): + zones_poisoned.append(name) + continue n_total += n_z - s_total += np.asarray(group["stretch_stats_sum"][t]) - m_total += np.asarray(group["stretch_stats_prod"][t]) + s_total += s_z + m_total += m_z k_z = int(group["stretch_sample_count"][t]) if k_z > 0: @@ -3488,6 +3618,14 @@ def compute_stretch_from_stats( ) zones_used.append(name) + if zones_poisoned: + raise ValueError( + f"Stretch statistics for {', '.join(zones_poisoned)} contain " + f"non-finite sums — collected before sentinel-scale filtering " + f"existed, so nodata pixels with huge finite scales overflowed " + f"them. Rebuild with `zarr-fill --backfill-stretch-stats " + f"--zones {','.join(z[3:].lstrip('0') for z in zones_poisoned)}`." + ) if zones_missing: raise ValueError( f"{len(zones_missing)} zone(s) have no stretch-statistics arrays " @@ -3537,9 +3675,16 @@ def compute_stretch_from_stats( # two covariances — comparing eigenvectors instead would false-alarm # whenever eigenvalues are close, where the vectors are arbitrary. cov_s = np.cov(x.T, aweights=weights) - drift = float( - np.linalg.norm(cov - cov_s) / max(np.linalg.norm(cov), 1e-30) - ) + with np.errstate(invalid="ignore", over="ignore"): + drift = float( + np.linalg.norm(cov - cov_s) / max(np.linalg.norm(cov), 1e-30) + ) + if not math.isfinite(drift): + raise RuntimeError( + "Drift check is non-finite — the aggregated statistics are " + "corrupt. Rebuild them with `zarr-fill " + "--backfill-stretch-stats`." + ) # The sample covariance itself carries ~sqrt(d/n_eff) relative error, # so the alarm floor scales with the effective sample size — a small # sample must not read as drift. @@ -3600,10 +3745,8 @@ def compute_stretch_from_stats( bks[j] = bks[j - 1] + 1e-9 cdf_breaks.append([float(v) for v in bks]) - # Persist with the same key set as the legacy path so readers + # The same key set as the legacy path so readers # (_load_global_stretch, build_global_preview) are unaffected. - root_rw = store.open_group(mode="r+") - stretch_map = dict(root_rw.attrs.get(_GLOBAL_STRETCH_ATTR, {})) method_prefix = "zone_stats_pca" if mode == "pca" else "zone_stats_percentile" entry: Dict[str, Any] = { "min": stretch_min, @@ -3623,17 +3766,19 @@ def compute_stretch_from_stats( entry["pca_components"] = pca_proj_components entry["pca_mean"] = pca_proj_mean entry["pca_explained_variance_ratio"] = pca_evr - stretch_map[str(year)] = entry - root_rw.attrs[_GLOBAL_STRETCH_ATTR] = stretch_map - - if console: - console.print( - f"[green]Saved to {_GLOBAL_STRETCH_ATTR}.{year} on the store " - f"root.[/green] Run zarr-consolidate so consolidated-metadata " - f"readers see it." - ) + if persist: + root_rw = store.open_group(mode="r+") + stretch_map = dict(root_rw.attrs.get(_GLOBAL_STRETCH_ATTR, {})) + stretch_map[str(year)] = entry + root_rw.attrs[_GLOBAL_STRETCH_ATTR] = stretch_map + if console: + console.print( + f"[green]Saved to {_GLOBAL_STRETCH_ATTR}.{year} on the store " + f"root.[/green] Run zarr-consolidate so consolidated-metadata " + f"readers see it." + ) - return {"min": stretch_min, "max": stretch_max, "samples": int(channels.shape[0])} + return entry def backfill_stretch_stats( @@ -4082,7 +4227,11 @@ def compute_global_stretch( } -def _load_global_stretch(store_path: Path, year: int) -> Optional[dict]: +def _load_global_stretch( + store_path: "str | Path | StoreLocation", + year: int, + storage_options: Optional[Dict[str, Any]] = None, +) -> Optional[dict]: """Look up a previously-computed global stretch for ``year``. Returns ``{"min": [..], "max": [..], "cdf": [[..], ..], @@ -4090,9 +4239,7 @@ def _load_global_stretch(store_path: Path, year: int) -> Optional[dict]: only populated when the stretch was computed in ``mode='pca'``. Returns ``None`` if no stretch is stored for the year. """ - import zarr - - root = zarr.open_group(str(store_path), mode="r", use_consolidated=False) + root = StoreLocation.resolve(store_path, storage_options).open_group(mode="r") stretch_map = root.attrs.get(_GLOBAL_STRETCH_ATTR, {}) if not isinstance(stretch_map, dict): return None @@ -4118,26 +4265,30 @@ def _load_global_stretch(store_path: Path, year: int) -> Optional[dict]: # coarsening. -def _ensure_global_store(store_path: Path, num_levels: int) -> None: +def _ensure_global_store(dest: "StoreLocation", num_levels: int) -> None: """Create the global_rgb/ pyramid group within the store.""" - import zarr from zarr.codecs import BloscCodec - root = zarr.open_group( - str(store_path), mode="r+", zarr_format=3, use_consolidated=False - ) + root = dest.open_group(mode="r+", zarr_format=3) # Check if already exists with correct shape if "global_rgb/0/rgb" in root: shape = root["global_rgb/0/rgb"].shape if shape == (GLOBAL_LEVEL0_H, GLOBAL_LEVEL0_W, GLOBAL_NUM_BANDS): return + if dest.is_remote: + # Dropping the prefix could be millions of objects; deleting that + # implicitly is not something a build step should decide. + raise ValueError( + f"{dest} already holds a global_rgb pyramid of shape {shape}, " + f"which does not match the expected " + f"{(GLOBAL_LEVEL0_H, GLOBAL_LEVEL0_W, GLOBAL_NUM_BANDS)}. " + f"Delete the global_rgb/ prefix yourself and re-run." + ) import shutil - shutil.rmtree(str(store_path / "global_rgb")) - root = zarr.open_group( - str(store_path), mode="r+", zarr_format=3, use_consolidated=False - ) + shutil.rmtree(str(Path(dest.url) / "global_rgb")) + root = dest.open_group(mode="r+", zarr_format=3) # Create pyramid levels via zarr API global_grp = root.create_group("global_rgb") @@ -4167,18 +4318,12 @@ def _ensure_global_store(store_path: Path, num_levels: int) -> None: w //= 2 # Re-open the global_rgb group to ensure attrs write to the correct handle - root = zarr.open_group( - str(store_path), mode="r+", zarr_format=3, use_consolidated=False - ) + root = dest.open_group(mode="r+", zarr_format=3) global_grp = root["global_rgb"] # Build multiscale + spatial + proj metadata directly # (avoids depending on unstable topozarr API) - from geozarr_toolkit import ( - create_geozarr_attrs, - create_multiscales_layout, - ) - from geozarr_toolkit.conventions.multiscales import MultiscalesConventionMetadata + from zarr_cm import multiscales west, south, east, north_ = GLOBAL_BOUNDS actual_levels = len([k for k in global_grp.keys() if k.isdigit()]) @@ -4202,27 +4347,23 @@ def _ensure_global_store(store_path: Path, num_levels: int) -> None: w_lvl //= 2 res *= 2.0 - ms_layout = create_multiscales_layout(levels, resampling_method="mean") - - # Geospatial attrs (proj + spatial) - geozarr_attrs = create_geozarr_attrs( + # The pyramid root states no single transform or shape — each level + # carries its own in the layout entries above. Insert multiscales last so + # it joins the same ``zarr_conventions`` list rather than replacing it. + attrs = _geo_convention_attrs( dimensions=["lat", "lon"], crs="EPSG:4326", bbox=[west, south, east, north_], ) - - # Fix spatial description bug in geozarr-toolkit - for conv in geozarr_attrs.get("zarr_conventions", []): - if conv.get("uuid") == "689b58e2-cf7b-45e0-9fff-9cfc0883d6b4": - conv["description"] = "Spatial coordinate information" - - # Add multiscales convention registration - ms_conv = MultiscalesConventionMetadata() - geozarr_attrs["zarr_conventions"].insert(0, ms_conv.model_dump(exclude_none=True)) - - # Merge all attrs - geozarr_attrs.update(ms_layout) - global_grp.attrs.update(geozarr_attrs) + attrs = multiscales.insert( + attrs, + multiscales.create( + revision=MULTISCALES_REVISION, + layout=levels, + resampling_method="mean", + ), + ) + global_grp.attrs.update(attrs) # Per-worker state for reprojection @@ -4235,21 +4376,40 @@ def _ensure_global_store(store_path: Path, num_levels: int) -> None: def _init_reproj_worker( - store_path: str, + source_url: str, + source_options: Optional[Dict[str, Any]], + dest_url: str, + dest_options: Optional[Dict[str, Any]], zone_group: str, zone_epsg: int, time_index: int, stretch: dict, ) -> None: - """Process pool initializer: open stores and create transformer.""" + """Process pool initializer: open stores and create transformer. + + The zone embeddings are read from ``source_url`` and the pyramid written + to ``dest_url``; either may be a remote store, and they may live on + different endpoints with different credentials. Workers are spawned, so + remote filesystem state is built fresh here rather than inherited from a + fork. + """ global _reproj_global_arr, _reproj_emb_arr, _reproj_scales_arr global _reproj_to_utm, _reproj_time_index, _reproj_stretch - import zarr from pyproj import Transformer - root = zarr.open_group(store_path, mode="r+", zarr_format=3, use_consolidated=False) - _reproj_global_arr = root["global_rgb/0/rgb"] - zone = root[zone_group] + from . import remote + + remote.quieten_dependency_logging() + remote.reset_after_fork() + remote.die_with_parent() + + dest = StoreLocation(dest_url, dest_options).open_group( + mode="r+", zarr_format=3 + ) + _reproj_global_arr = dest["global_rgb/0/rgb"] + zone = StoreLocation(source_url, source_options).open_group( + mode="r", path=zone_group + ) _reproj_emb_arr = zone["embeddings"] _reproj_scales_arr = zone["scales"] _reproj_to_utm = Transformer.from_crs( @@ -4373,7 +4533,7 @@ def _reproject_chunk( # Compute RGB on the fly from embeddings + scales (no stored rgb array needed) scales_chunk = np.asarray(scales_arr[time_index, r_min:r_max, c_min:c_max]) - valid = np.isfinite(scales_chunk) + valid = valid_scale_mask(scales_chunk) if not np.any(valid): return False @@ -4470,7 +4630,8 @@ def _reproject_chunk( def _reproject_zone( - store_path: Path, + source: StoreLocation, + dest: "StoreLocation", zone_num: int, zone_group: str, zone_epsg: int, @@ -4479,64 +4640,64 @@ def _reproject_zone( time_index: int, stretch: dict, workers: int, + present: set, console: Optional["rich.console.Console"] = None, force: bool = False, -) -> Tuple[int, int, int, int, bool]: - """Reproject one zone's embeddings into global level 0 (computing RGB on the fly).""" +) -> Tuple[List[Tuple[int, int, int, int]], bool]: + """Reproject one zone's embeddings into global level 0. + + Embeddings are read from *source* and the pyramid's level 0 written into + *dest*; either may be local or remote. ``present`` is the zone's set of + existing shard coordinates for this year — the work list is derived from + their footprints, never from the zone's bounding box, which at high + latitude covers most longitudes and would enqueue millions of empty + chunks. + + Returns (rectangles the zone touched, whether work was done); the + rectangles feed the coarsening. + """ + import multiprocessing from concurrent.futures import ProcessPoolExecutor, as_completed + # Spawn, not fork: the source may be a remote store, and a forked + # worker inheriting fsspec's event loop without its thread deadlocks. + mp_context = multiprocessing.get_context("spawn") + src_pixel = zone_transform[0] src_origin_e = zone_transform[2] src_origin_n = zone_transform[5] src_h, src_w = zone_shape[:2] - row_start, row_end, col_start, col_end = _zone_output_bounds( - zone_epsg=zone_epsg, - zone_transform=zone_transform, - zone_shape=(src_h, src_w), + chunk_set, regions = _chunks_for_shards( + present, zone_epsg, zone_transform, (src_h, src_w) ) - - if col_end <= col_start or row_end <= row_start: + if not chunk_set: if console: - console.print(f" [yellow]Zone {zone_num}: no output region[/yellow]") - return (0, 0, 0, 0, False) - - n_chunk_rows = (row_end - row_start) // GLOBAL_CHUNK - n_chunk_cols = (col_end - col_start) // GLOBAL_CHUNK - chunk_row_start = row_start // GLOBAL_CHUNK - chunk_col_start = col_start // GLOBAL_CHUNK + console.print(f" [yellow]Zone {zone_num}: no data to render[/yellow]") + return ([], False) # Resume check. The marker lives in the state sibling, not the store, so # the published hierarchy stays free of non-Zarr objects. - marker = _preview_marker_path(store_path, zone_num) - if marker.exists(): + state = dest.state + marker = _preview_marker_parts(zone_num) + if state.exists(*marker, on_denied=False): if force: - marker.unlink() + state.remove(*marker) else: if console: console.print(f" Zone {zone_num:02d}: already complete, skipping") - return (row_start, row_end, col_start, col_end, False) + return (regions, False) - chunks_total = n_chunk_rows * n_chunk_cols + chunks_total = len(chunk_set) if console: console.print( - f" Zone {zone_num:02d}: {n_chunk_rows}x{n_chunk_cols} " - f"= {chunks_total} chunks" + f" Zone {zone_num:02d}: {chunks_total:,} candidate chunk(s) " + f"from {len(present)} shard footprint(s)" ) work_items = [ - ( - chunk_row_start + cr, - chunk_col_start + cc, - zone_epsg, - src_pixel, - src_origin_e, - src_origin_n, - src_h, - src_w, - ) - for cr in range(n_chunk_rows) - for cc in range(n_chunk_cols) + (cr, cc, zone_epsg, src_pixel, src_origin_e, src_origin_n, src_h, src_w) + for cr, cc in sorted(chunk_set) ] chunks_written = 0 @@ -4567,8 +4728,18 @@ def _reproject_zone( ) with ProcessPoolExecutor( max_workers=workers, + mp_context=mp_context, initializer=_init_reproj_worker, - initargs=(str(store_path), zone_group, zone_epsg, time_index, stretch), + initargs=( + source.url, + source.storage_options, + dest.url, + dest.storage_options, + zone_group, + zone_epsg, + time_index, + stretch, + ), ) as pool: futures = { pool.submit(_reproject_chunk_worker, item): item @@ -4585,8 +4756,18 @@ def _reproject_zone( else: with ProcessPoolExecutor( max_workers=workers, + mp_context=mp_context, initializer=_init_reproj_worker, - initargs=(str(store_path), zone_group, zone_epsg, time_index, stretch), + initargs=( + source.url, + source.storage_options, + dest.url, + dest.storage_options, + zone_group, + zone_epsg, + time_index, + stretch, + ), ) as pool: futures = { pool.submit(_reproject_chunk_worker, item): item for item in work_items @@ -4598,16 +4779,16 @@ def _reproject_zone( except Exception as e: logger.warning(f"Reproject chunk failed: {e}") - marker.parent.mkdir(parents=True, exist_ok=True) - marker.write_text( - f"zone={zone_num} chunks={chunks_total} written={chunks_written}\n" + state.write_bytes( + f"zone={zone_num} chunks={chunks_total} written={chunks_written}\n".encode(), + *marker, ) - return (row_start, row_end, col_start, col_end, True) + return (regions, True) def build_global_preview( - store_path: Path, + store_path: "str | Path | StoreLocation", year: int = 2024, zones: Optional[List[int]] = None, num_levels: int = GLOBAL_DEFAULT_LEVELS, @@ -4616,12 +4797,40 @@ def build_global_preview( saturation: float = 1.0, console: Optional["rich.console.Console"] = None, force: bool = False, + storage_options: Optional[Dict[str, Any]] = None, + output_path: "Optional[str | Path | StoreLocation]" = None, + output_storage_options: Optional[Dict[str, Any]] = None, + state_url: Optional[str] = None, + state_storage_options: Optional[Dict[str, Any]] = None, + reproject_only: bool = False, + coarsen_only: bool = False, ) -> None: """Build the global EPSG:4326 RGB pyramid from zone-level embeddings. - Computes RGB from embeddings+scales (bands 0-2) for the specified year, - reprojects from UTM to geographic coordinates and composites into the - pyramid. No pre-computed rgb array needed. + Computes RGB from embeddings+scales for the specified year, reprojects + from UTM to geographic coordinates and composites into the pyramid. + + Source and destination are independent locations, each local or remote + with its own credentials, so a pyramid can be written to a bucket while + embeddings stream anonymously from a read-only mirror — reads are + sub-shard byte ranges, so no copy of the source is ever made. Without + ``output_path`` the pyramid is written into the source store itself. If + the store carries no persisted ``geoemb:stretch`` for the year, one is + computed on the fly from the per-zone stretch statistics — a few MiB of + reads — so previewing a read-only store needs no prior ``zarr-stretch``. + + Zones are composited into shared level-0 chunks with a read-modify-write, + so by default separate invocations must not run **concurrently** against + the same destination. Running them one at a time (``--zones N`` per + invocation) is safe and resumable: each zone records a marker in the + state area, which ``state_url`` can place on local disk even when the + pyramid itself is remote. + + For a parallel sweep, *reproject_only* stops each zone's coarsening at + :data:`COARSEN_PARALLEL_SAFE_LEVEL`, below which zones that do not share + level-0 chunks stay disjoint at every level. Two rounds — odd zones, then + even — cover all 60 without any pair colliding. *coarsen_only* then + builds the remaining levels in one global pass. Args: gamma: Per-channel gamma applied after normalisation. ``< 1.0`` @@ -4639,8 +4848,38 @@ def build_global_preview( warnings.filterwarnings("ignore", message="Object at .* is not recognized") - store_path = Path(store_path) - root = zarr.open_group(str(store_path), mode="r", use_consolidated=False) + if reproject_only and coarsen_only: + raise ValueError("--reproject-only and --coarsen-only are mutually exclusive") + + source = StoreLocation.resolve(store_path, storage_options) + if output_path is not None: + dest = StoreLocation.resolve( + output_path, + output_storage_options + if output_storage_options is not None + else storage_options, + state_url, + state_storage_options, + ) + # The pyramid gets its own store; create the root on first use. + # Opened through the location rather than as_zarr_store() so a local + # destination (plain path or file:// URL) gets its directory made. + if not dest.exists("zarr.json", on_denied=False): + dest.open_group(mode="a", zarr_format=3) + else: + dest = StoreLocation.resolve( + source.url, source.storage_options, state_url, state_storage_options + ) + + # Each zone's coarsening stops short of the levels where zones converge + # when a parallel sweep is in play; ``--coarsen-only`` finishes them. + zone_coarsen_levels = ( + min(num_levels, COARSEN_PARALLEL_SAFE_LEVEL + 1) + if reproject_only + else num_levels + ) + + root = source.open_group(mode="r") # Derive years from first zone's time coordinate all_years: list[int] = [] @@ -4707,11 +4946,98 @@ def build_global_preview( console.print(f" {len(zone_infos)} zone(s) with data") # Ensure global pyramid structure exists - _ensure_global_store(store_path, num_levels) + _ensure_global_store(dest, num_levels) + + if coarsen_only: + # Single-writer finish for the levels where zones share chunks. The + # region is the whole grid: by this depth every zone's contribution + # overlaps, so there is nothing to restrict it to. + start = COARSEN_PARALLEL_SAFE_LEVEL + 1 + if start >= num_levels: + if console: + console.print( + f"[yellow]Nothing to do: levels beyond " + f"{COARSEN_PARALLEL_SAFE_LEVEL} are outside a " + f"{num_levels}-level pyramid.[/yellow]" + ) + return + if console: + console.print( + f"Coarsening levels {start}-{num_levels - 1} over the full grid" + ) + _coarsen_zone_pyramid( + dest=dest, + row_start=0, + row_end=GLOBAL_LEVEL0_H, + col_start=0, + col_end=GLOBAL_LEVEL0_W, + num_levels=num_levels, + workers=workers, + console=console, + start_level=start, + ) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Consolidated metadata") + zarr.consolidate_metadata(dest.as_zarr_store()) + if console: + console.print("\n [green]Coarsening complete[/green]") + return - # Prefer a pre-computed cross-zone stretch (written by `zarr-stretch`). + # Prefer a pre-computed cross-zone stretch (written by `zarr-stretch`); + # fall back to deriving one from the per-zone statistics right now. # Using one shared stretch eliminates inter-zone colour discontinuities. - global_stretch = _load_global_stretch(store_path, year) + # + # The destination is consulted first when it differs from the source. A + # read-only mirror is typically a cache in front of the very bucket the + # destination addresses, so a stretch written minutes ago may not have + # propagated to it yet; going direct avoids silently falling back to a + # derived stretch because the cached copy still lacks the attribute. + global_stretch = None + if dest.url != source.url: + try: + global_stretch = _load_global_stretch(dest, year) + except Exception: + global_stretch = None # A pyramid-only destination has no stretch. + if global_stretch is None: + global_stretch = _load_global_stretch(source, year) + if global_stretch is None: + try: + entry = compute_stretch_from_stats( + source, + year=year, + zones=zones, + console=console, + persist=False, + ) + global_stretch = { + "min": list(entry["min"]), + "max": list(entry["max"]), + "mode": entry.get("mode", "bands"), + } + if entry.get("cdf") is not None: + global_stretch["cdf"] = [list(c) for c in entry["cdf"]] + if entry.get("pca_components") is not None: + global_stretch["pca_components"] = [ + list(r) for r in entry["pca_components"] + ] + global_stretch["pca_mean"] = list(entry["pca_mean"]) + if console: + console.print( + "[cyan]Derived stretch from the store's per-zone " + "statistics (not persisted).[/cyan]" + ) + except (ValueError, RuntimeError) as e: + if source.is_remote: + raise RuntimeError( + f"No persisted stretch for {year} and none derivable " + f"from statistics ({e}). A remote source cannot fall " + f"back to shard sampling." + ) from e + if console: + console.print( + f"[yellow]No stretch statistics available ({e}); " + f"falling back to per-zone sampling.[/yellow]" + ) if global_stretch is not None: global_stretch["gamma"] = gamma global_stretch["saturation"] = saturation @@ -4739,13 +5065,7 @@ def build_global_preview( stretch = global_stretch else: # Fallback: per-zone stretch (produces seams at zone boundaries). - zone_store = zarr.open_group( - str(store_path), - mode="r", - path=info["zone_group"], - zarr_format=3, - use_consolidated=False, - ) + zone_store = source.open_group(mode="r", path=info["zone_group"]) if console: console.print(" Sampling stretch...") stretch = compute_stretch( @@ -4763,8 +5083,24 @@ def build_global_preview( f"gamma={gamma}, saturation={saturation}" ) - row_start, row_end, col_start, col_end, did_work = _reproject_zone( - store_path=store_path, + zone_h, zone_w = info["shape"] + all_coords = { + (sr, sc) + for sr in range(math.ceil(zone_h / SHARD_SIZE)) + for sc in range(math.ceil(zone_w / SHARD_SIZE)) + } + present = _existing_shards( + source, info["zone_group"], time_index, all_coords, console=None + ) + if not present: + if console: + console.print(f" no shards for {year}, skipping") + gc.collect() + continue + + regions, did_work = _reproject_zone( + source=source, + dest=dest, zone_num=zone_num, zone_group=info["zone_group"], zone_epsg=info["epsg"], @@ -4773,6 +5109,7 @@ def build_global_preview( time_index=time_index, stretch=stretch, workers=workers, + present=present, console=console, force=force, ) @@ -4780,23 +5117,24 @@ def build_global_preview( if did_work: if console: console.print(" Building pyramid...") - _coarsen_zone_pyramid( - store_path=store_path, - row_start=row_start, - row_end=row_end, - col_start=col_start, - col_end=col_end, - num_levels=num_levels, - workers=workers, - console=console, - ) + for row_start, row_end, col_start, col_end in regions: + _coarsen_zone_pyramid( + dest=dest, + row_start=row_start, + row_end=row_end, + col_start=col_start, + col_end=col_end, + num_levels=zone_coarsen_levels, + workers=workers, + console=console, + ) gc.collect() - # Consolidate + # Consolidate the pyramid store (always local). with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Consolidated metadata") - zarr.consolidate_metadata(str(store_path)) + zarr.consolidate_metadata(dest.as_zarr_store()) if console: console.print("\n [green]Global preview complete[/green]") diff --git a/pyproject.toml b/pyproject.toml index df62e7e..48a2250 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "dask", "fsspec", "aiohttp", - "geozarr-toolkit", + "zarr-cm>=0.4.1", "contextily", ] diff --git a/tests/zarr_remote_check.py b/tests/zarr_remote_check.py index 3237d4e..1781fb1 100644 --- a/tests/zarr_remote_check.py +++ b/tests/zarr_remote_check.py @@ -14,6 +14,7 @@ import os import sys import tempfile +import warnings from pathlib import Path import numpy as np @@ -492,6 +493,23 @@ def placed(lon, x_off): < 0.02, ) +# Sentinel scales: huge-finite nodata values must not count as data. Some +# published scales files carry ~FLT_MAX sentinels that pass isfinite() — +# they inflated N by 100x and overflowed the product sums to inf. +jemb = nprng.integers(-128, 127, (B, 64, 64), dtype=np.int8) +jsc = nprng.random((64, 64)).astype(np.float32) * 0.01 + 0.001 +jsc[:8, :8] = np.float32(3.4e38) # FLT_MAX-style sentinel +jsc[8, 8] = np.float32(0.0) # degenerate +jsc[9, 9] = np.float32(-1.0) # negative +jst = shard_stretch_stats(jemb, jsc, sample_cap=100, seed=3) +check("sentinel scales excluded from N", jst["n"] == 64 * 64 - 64 - 2) +check("junk-free sums stay finite", bool(np.isfinite(jst["sum"]).all())) +check("junk-free products stay finite", bool(np.isfinite(jst["prod"]).all())) +check( + "sampled scales all plausible", + float(jst["sample_scales"].max()) < 1.0, +) + # Zone-array round trip: create, fold twice, contents accumulate. import zarr # noqa: E402 @@ -520,6 +538,40 @@ def placed(lon, x_off): ) check("other year untouched", int(zs["stretch_stats_count"][1]) == 0) +# --------------------------------------------------------------------------- +# Preview work list from shard footprints +# --------------------------------------------------------------------------- + +from geotessera.zarr import GLOBAL_CHUNK, _chunks_for_shards # noqa: E402 + +# One shard at UTM zone 31's origin near (0.0E, ~0.9N): its footprint is +# ~41 km, so the candidate set must be a handful of chunks, not a +# bounding-box sweep. +transform31 = [10.0, 0.0, 166021.44, 0.0, -10.0, 100000.0] +chunks, regions = _chunks_for_shards({(0, 0)}, 32631, transform31, (8192, 8192)) +check("shard footprint yields a small chunk set", 0 < len(chunks) < 200) +check("an ordinary footprint needs one region", len(regions) == 1) +r0, r1, c0, c1 = regions[0] +check( + "footprint bounds are chunk-aligned and ordered", + r0 < r1 and c0 < c1 and r0 % GLOBAL_CHUNK == 0 and c1 % GLOBAL_CHUNK == 0, +) +check( + "empty shard set yields no work", + _chunks_for_shards(set(), 32631, transform31, (8192, 8192))[0] == set(), +) + +# Two far-apart shards must not fill the space between them: the union is +# exactly the two footprints, despite a row span of >1000 chunks. +near, _ = _chunks_for_shards({(0, 0)}, 32631, transform31, (1011712, 8192)) +far, _ = _chunks_for_shards({(200, 0)}, 32631, transform31, (1011712, 8192)) +both, _ = _chunks_for_shards({(0, 0), (200, 0)}, 32631, transform31, (1011712, 8192)) +rows = {c[0] for c in both} +check( + "sparse shards keep a sparse work list", + both == near | far and max(rows) - min(rows) > 1000, +) + # --------------------------------------------------------------------------- # Storage options and source layout # --------------------------------------------------------------------------- @@ -578,6 +630,209 @@ def placed(lon, x_off): ) check("url source reports remote", src.is_remote) + +# --------------------------------------------------------------------------- +# Preview work list: antimeridian footprints must not span the globe +# --------------------------------------------------------------------------- +# A shard straddling 180 samples corners near -180 and +180. Taking the naive +# min/max of those makes it claim every chunk column at its latitude, which +# for utm01/utm60 enqueued ~2.3M bogus chunks. + +from geotessera.zarr import ( # noqa: E402 + GLOBAL_CHUNK, + GLOBAL_LEVEL0_W, + _chunks_for_shards, +) + +N_COLS = GLOBAL_LEVEL0_W // GLOBAL_CHUNK + + +def _shard_cols(epsg, origin_x, origin_y): + """Columns and coarsening regions of one 4096px shard at a zone origin.""" + chunks, regions = _chunks_for_shards( + {(0, 0)}, + epsg, + [10.0, 0.0, origin_x, 0.0, -10.0, origin_y], + (4096, 4096), + ) + cols = {c for _r, c in chunks} + return cols, (max(cols) - min(cols)) if cols else 0, regions, chunks + + +# EPSG:32660 has central meridian 177E, so easting ~834000 sits on the +# antimeridian at the equator; this shard straddles it. +wrap_cols, wrap_span, wrap_regions, wrap_chunks = _shard_cols( + 32660, 810_000.0, 500_000.0 +) +check( + "antimeridian shard does not claim the whole grid width", + len(wrap_cols) < N_COLS // 10, +) +check( + "antimeridian shard reaches both grid edges", + min(wrap_cols) == 0 and max(wrap_cols) == N_COLS - 1, +) +# One enclosing rectangle would span every column and make the coarsening +# read and rewrite the entire grid width; two tight ones must not. +check("antimeridian footprint splits into two regions", len(wrap_regions) == 2) +wrap_slots = sum( + ((b - a) // GLOBAL_CHUNK) * ((d - c) // GLOBAL_CHUNK) for a, b, c, d in wrap_regions +) +wrap_rows = {r for r, _c in wrap_chunks} +enclosing = (max(wrap_rows) - min(wrap_rows) + 1) * ( + max(wrap_cols) - min(wrap_cols) + 1 +) +check( + "split regions stay near the real chunk count", + wrap_slots <= 2 * len(wrap_chunks), +) +check( + "split beats a single enclosing rectangle by orders of magnitude", + wrap_slots * 100 < enclosing, +) + +# A shard well inside the same zone must be unaffected by the wrap handling. +mid_cols, mid_span, mid_regions, _mid_chunks = _shard_cols(32660, 500_000.0, 500_000.0) +check("mid-zone shard stays contiguous", mid_span == len(mid_cols) - 1) +check("mid-zone shard spans few columns", len(mid_cols) < 20) +check("mid-zone shard needs one region", len(mid_regions) == 1) + + +# --------------------------------------------------------------------------- +# Global preview pyramid: destination may be local or remote +# --------------------------------------------------------------------------- +# The pyramid is written through a StoreLocation, so a file:// destination +# drives the same fsspec path an s3:// one takes. + +import zarr # noqa: E402 + +from geotessera.zarr import ( # noqa: E402 + GLOBAL_CHUNK, + GLOBAL_LEVEL0_H, + GLOBAL_LEVEL0_W, + GLOBAL_NUM_BANDS, + _coarsen_zone_pyramid, + _ensure_global_store, + _preview_marker_parts, +) + +for label, dest_loc in ( + ("local path", str(TMP / "pyr_local.zarr")), + ("file:// url", url(TMP / "pyr_url.zarr")), +): + dest = StoreLocation.resolve(dest_loc) + dest.open_group(mode="a", zarr_format=3) + _ensure_global_store(dest, 4) + + root = dest.open_group(mode="r+", zarr_format=3) + check( + f"pyramid level 0 has the global shape over {label}", + root["global_rgb/0/rgb"].shape + == (GLOBAL_LEVEL0_H, GLOBAL_LEVEL0_W, GLOBAL_NUM_BANDS), + ) + check(f"pyramid levels created over {label}", "global_rgb/3/rgb" in root) + check( + f"pyramid registers its conventions over {label}", + {c["name"] for c in root["global_rgb"].attrs["zarr_conventions"]} + == {"spatial:", "proj:", "multiscales"}, + ) + check( + f"pyramid keeps per-level geometry over {label}", + "spatial:shape" in root["global_rgb"].attrs["multiscales"]["layout"][1], + ) + + # A second ensure on a matching pyramid must not wipe what is there. + r0, c0 = 4 * GLOBAL_CHUNK, 6 * GLOBAL_CHUNK + root["global_rgb/0/rgb"][r0 : r0 + GLOBAL_CHUNK, c0 : c0 + GLOBAL_CHUNK, :] = ( + np.full((GLOBAL_CHUNK, GLOBAL_CHUNK, GLOBAL_NUM_BANDS), 200, dtype=np.uint8) + ) + _ensure_global_store(dest, 4) + check( + f"re-ensure keeps existing pyramid data over {label}", + int(dest.open_group(mode="r", zarr_format=3)["global_rgb/0/rgb"][r0, c0, 0]) + == 200, + ) + + _coarsen_zone_pyramid( + dest=dest, + row_start=r0, + row_end=r0 + GLOBAL_CHUNK, + col_start=c0, + col_end=c0 + GLOBAL_CHUNK, + num_levels=4, + workers=2, + ) + check( + f"coarsening writes level 1 over {label}", + int(dest.open_group(mode="r", zarr_format=3)["global_rgb/1/rgb"][r0 // 2, c0 // 2, 0]) + == 200, + ) + + # start_level walks the region down without touching the levels below it, + # which is what lets a parallel sweep stop early and a global pass finish. + g = dest.open_group(mode="r+", zarr_format=3) + g["global_rgb/1/rgb"][r0 // 2, c0 // 2, :] = 0 + g["global_rgb/2/rgb"][r0 // 4, c0 // 4, :] = 0 + _coarsen_zone_pyramid( + dest=dest, + row_start=r0, + row_end=r0 + GLOBAL_CHUNK, + col_start=c0, + col_end=c0 + GLOBAL_CHUNK, + num_levels=4, + workers=2, + start_level=2, + ) + g = dest.open_group(mode="r", zarr_format=3) + check( + f"start_level leaves shallower levels alone over {label}", + int(g["global_rgb/1/rgb"][r0 // 2, c0 // 2, 0]) == 0, + ) + check( + f"start_level still coarsens deeper levels over {label}", + int(g["global_rgb/2/rgb"][r0 // 4, c0 // 4, 0]) > 0, + ) + + # Markers follow an explicit state_url rather than the .build + # sibling, so a remote pyramid can keep its bookkeeping on local disk. + elsewhere = StoreLocation.resolve( + dest_loc, None, url(TMP / f"state_{label.split()[0]}") + ) + parts = _preview_marker_parts(7) + elsewhere.state.write_bytes(b"zone=7\n", *parts) + check( + f"state_url redirects markers off the store for {label}", + elsewhere.state.exists(*parts, on_denied=False) + and not dest.state.exists(*parts, on_denied=False), + ) + + # Resume markers belong to the state sibling, never the published store. + state, parts = dest.state, _preview_marker_parts(30) + state.write_bytes(b"zone=30\n", *parts) + check( + f"preview marker lands in the state sibling over {label}", + state.exists(*parts, on_denied=False) + and not dest.exists(*parts, on_denied=False), + ) + state.remove(*parts) + check( + f"preview marker is removable over {label}", + not state.exists(*parts, on_denied=False), + ) + +# Consolidation goes through the location too, so it works on either. +with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Consolidated metadata") + zarr.consolidate_metadata( + StoreLocation.resolve(url(TMP / "pyr_url.zarr")).as_zarr_store() + ) +check( + "pyramid consolidates over a url", + StoreLocation.resolve(url(TMP / "pyr_url.zarr")).exists( + "zarr.json", on_denied=False + ), +) + import shutil # noqa: E402 shutil.rmtree(TMP, ignore_errors=True) diff --git a/uv.lock b/uv.lock index a2c7947..86db73d 100644 --- a/uv.lock +++ b/uv.lock @@ -186,15 +186,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, ] -[[package]] -name = "annotated-types" -version = "0.7.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, -] - [[package]] name = "attrs" version = "26.1.0" @@ -694,7 +685,6 @@ dependencies = [ { name = "fsspec" }, { name = "geodatasets" }, { name = "geopandas" }, - { name = "geozarr-toolkit" }, { name = "matplotlib" }, { name = "numpy" }, { name = "pandas" }, @@ -708,6 +698,7 @@ dependencies = [ { name = "typer" }, { name = "xarray" }, { name = "zarr" }, + { name = "zarr-cm" }, ] [package.optional-dependencies] @@ -730,7 +721,6 @@ requires-dist = [ { name = "fsspec" }, { name = "geodatasets", specifier = ">=2024.8.0" }, { name = "geopandas" }, - { name = "geozarr-toolkit" }, { name = "matplotlib" }, { name = "numpy", specifier = ">=1.24.0" }, { name = "pandas" }, @@ -745,6 +735,7 @@ requires-dist = [ { name = "typer" }, { name = "xarray" }, { name = "zarr" }, + { name = "zarr-cm", specifier = ">=0.4.1" }, ] provides-extras = ["s3"] @@ -754,21 +745,6 @@ dev = [ { name = "ruff", specifier = ">=0.12.8" }, ] -[[package]] -name = "geozarr-toolkit" -version = "0.1.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pydantic" }, - { name = "pyproj" }, - { name = "structlog" }, - { name = "zarr" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/96/00/eb0acfacf70c23c26ec84730234657279645be24de289de1a865fcad8a7e/geozarr_toolkit-0.1.2.tar.gz", hash = "sha256:04b4196245210bfc854051729d0c3a92865bc47cf258f70227ae39a5ff40576b", size = 26324, upload-time = "2026-04-01T01:45:03.528Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/2b/fb54fdbad8041ca7ffa8332b29afd3697580c8884656540b7ab6a8c56096/geozarr_toolkit-0.1.2-py3-none-any.whl", hash = "sha256:6d4317b791da7fcb19fce9df345f8f9ee3a1a822eb3bb3aad448ec41b73b526e", size = 19719, upload-time = "2026-04-01T01:45:02.146Z" }, -] - [[package]] name = "google-crc32c" version = "1.8.0" @@ -1636,96 +1612,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/be/6f79d55816d5c22557cf27533543d5d70dfe692adfbee4b99f2760674f38/pyarrow-24.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:c91d00057f23b8d353039520dc3a6c09d8608164c692e9f59a175a42b2ae0c19", size = 28131282, upload-time = "2026-04-21T10:51:16.815Z" }, ] -[[package]] -name = "pydantic" -version = "2.13.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "annotated-types" }, - { name = "pydantic-core" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/18/a5/b60d21ac674192f8ab0ba4e9fd860690f9b4a6e51ca5df118733b487d8d6/pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6", size = 844775, upload-time = "2026-05-06T13:43:05.343Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/7b/122376b1fd3c62c1ed9dc80c931ace4844b3c55407b6fb2d199377c9736f/pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba", size = 472262, upload-time = "2026-05-06T13:43:02.641Z" }, -] - -[[package]] -name = "pydantic-core" -version = "2.46.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9d/56/921726b776ace8d8f5db44c4ef961006580d91dc52b803c489fafd1aa249/pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1", size = 471464, upload-time = "2026-05-06T13:37:06.98Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ce/8c/af022f0af448d7747c5154288d46b5f2bc5f17366eaa0e23e9aa04d59f3b/pydantic_core-2.46.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2", size = 2106158, upload-time = "2026-05-06T13:38:57.215Z" }, - { url = "https://files.pythonhosted.org/packages/19/95/6195171e385007300f0f5574592e467c568becce2d937a0b6804f218bc49/pydantic_core-2.46.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f", size = 1951724, upload-time = "2026-05-06T13:37:02.697Z" }, - { url = "https://files.pythonhosted.org/packages/8e/bc/f47d1ff9cbb1620e1b5b697eef06010035735f07820180e74178226b27b3/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7", size = 1975742, upload-time = "2026-05-06T13:37:09.448Z" }, - { url = "https://files.pythonhosted.org/packages/5b/11/9b9a5b0306345664a2da6410877af6e8082481b5884b3ddd78d47c6013ce/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7", size = 2052418, upload-time = "2026-05-06T13:37:38.234Z" }, - { url = "https://files.pythonhosted.org/packages/f1/b7/a65fec226f5d78fc39f4a13c4cc0c768c22b113438f60c14adc9d2865038/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712", size = 2232274, upload-time = "2026-05-06T13:38:27.753Z" }, - { url = "https://files.pythonhosted.org/packages/68/f0/92039db98b907ef49269a8271f67db9cb78ae2fc68062ef7e4e77adb5f61/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4", size = 2309940, upload-time = "2026-05-06T13:38:05.353Z" }, - { url = "https://files.pythonhosted.org/packages/5f/97/2aab507d3d00ca626e8e57c1eac6a79e4e5fbcc63eb99733ff55d1717f65/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce", size = 2094516, upload-time = "2026-05-06T13:39:10.577Z" }, - { url = "https://files.pythonhosted.org/packages/22/37/a8aca44d40d737dde2bc05b3c6c07dff0de07ce6f82e9f3167aeaf4d5dea/pydantic_core-2.46.4-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987", size = 2136854, upload-time = "2026-05-06T13:40:22.59Z" }, - { url = "https://files.pythonhosted.org/packages/24/99/fcef1b79238c06a8cbec70819ac722ba76e02bc8ada9b0fd66eba40da01b/pydantic_core-2.46.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b", size = 2180306, upload-time = "2026-05-06T13:40:10.666Z" }, - { url = "https://files.pythonhosted.org/packages/ae/6c/fc44000918855b42779d007ae63b0532794739027b2f417321cddbc44f6a/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458", size = 2190044, upload-time = "2026-05-06T13:40:43.231Z" }, - { url = "https://files.pythonhosted.org/packages/6b/65/d9cadc9f1920d7a127ad2edba16c1db7916e59719285cd6c94600b0080ba/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b", size = 2329133, upload-time = "2026-05-06T13:39:57.365Z" }, - { url = "https://files.pythonhosted.org/packages/d0/cf/c873d91679f3a30bcf5e7ac280ce5573483e72295307685120d0d5ad3416/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c", size = 2374464, upload-time = "2026-05-06T13:38:06.976Z" }, - { url = "https://files.pythonhosted.org/packages/47/bd/6f2fc8188f31bf10590f1e98e7b306336161fac930a8c514cd7bd828c7dc/pydantic_core-2.46.4-cp312-cp312-win32.whl", hash = "sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894", size = 1974823, upload-time = "2026-05-06T13:40:47.985Z" }, - { url = "https://files.pythonhosted.org/packages/40/8c/985c1d41ea1107c2534abd9870e4ed5c8e7669b5c308297835c001e7a1c4/pydantic_core-2.46.4-cp312-cp312-win_amd64.whl", hash = "sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89", size = 2072919, upload-time = "2026-05-06T13:39:21.153Z" }, - { url = "https://files.pythonhosted.org/packages/c4/ba/f463d006e0c47373ca7ec5e1a261c59dc01ef4d62b2657af925fb0deee3a/pydantic_core-2.46.4-cp312-cp312-win_arm64.whl", hash = "sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a", size = 2027604, upload-time = "2026-05-06T13:39:03.753Z" }, - { url = "https://files.pythonhosted.org/packages/51/a2/5d30b469c5267a17b39dec53208222f76a8d351dfac4af661888c5aee77d/pydantic_core-2.46.4-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008", size = 2106306, upload-time = "2026-05-06T13:37:48.029Z" }, - { url = "https://files.pythonhosted.org/packages/c1/81/4fa520eaffa8bd7d1525e644cd6d39e7d60b1592bc5b516693c7340b50f1/pydantic_core-2.46.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4", size = 1951906, upload-time = "2026-05-06T13:37:17.012Z" }, - { url = "https://files.pythonhosted.org/packages/03/d5/fd02da45b659668b05923b17ba3a0100a0a3d5541e3bd8fcc4ecb711309e/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76", size = 1976802, upload-time = "2026-05-06T13:37:35.113Z" }, - { url = "https://files.pythonhosted.org/packages/21/f2/95727e1368be3d3ed485eaab7adbd7dda408f33f7a36e8b48e0144002b91/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3", size = 2052446, upload-time = "2026-05-06T13:37:12.313Z" }, - { url = "https://files.pythonhosted.org/packages/9c/86/5d99feea3f77c7234b8718075b23db11532773c1a0dbd9b9490215dc2eeb/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76", size = 2232757, upload-time = "2026-05-06T13:39:01.149Z" }, - { url = "https://files.pythonhosted.org/packages/d2/3a/508ac615935ef7588cf6d9e9b91309fdc2da751af865e02a9098de88258c/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4", size = 2309275, upload-time = "2026-05-06T13:37:41.406Z" }, - { url = "https://files.pythonhosted.org/packages/07/f8/41db9de19d7987d6b04715a02b3b40aea467000275d9d758ffaa31af7d50/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a", size = 2094467, upload-time = "2026-05-06T13:39:18.847Z" }, - { url = "https://files.pythonhosted.org/packages/2c/e2/f35033184cb11d0052daf4416e8e10a502ea2ac006fc4f459aee872727d1/pydantic_core-2.46.4-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262", size = 2134417, upload-time = "2026-05-06T13:40:17.944Z" }, - { url = "https://files.pythonhosted.org/packages/7e/7b/6ceeb1cc90e193862f444ebe373d8fdf613f0a82572dde03fb10734c6c71/pydantic_core-2.46.4-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e", size = 2179782, upload-time = "2026-05-06T13:40:32.618Z" }, - { url = "https://files.pythonhosted.org/packages/5a/f2/c8d7773ede6af08036423a00ae0ceffce266c3c52a096c435d68c896083f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd", size = 2188782, upload-time = "2026-05-06T13:36:51.018Z" }, - { url = "https://files.pythonhosted.org/packages/59/31/0c864784e31f09f05cdd87606f08923b9c9e7f6e51dd27f20f62f975ce9f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be", size = 2328334, upload-time = "2026-05-06T13:40:37.764Z" }, - { url = "https://files.pythonhosted.org/packages/c2/eb/4f6c8a41efa30baa755590f4141abf3a8c370fab610915733e74134a7270/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d", size = 2372986, upload-time = "2026-05-06T13:39:34.152Z" }, - { url = "https://files.pythonhosted.org/packages/5b/24/b375a480d53113860c299764bfe9f349a3dc9108b3adc0d7f0d786492ebf/pydantic_core-2.46.4-cp313-cp313-win32.whl", hash = "sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb", size = 1973693, upload-time = "2026-05-06T13:37:55.072Z" }, - { url = "https://files.pythonhosted.org/packages/7e/e8/cff247591966f2d22ec8c003cd7587e27b7ba7b81ab2fb888e3ab75dc285/pydantic_core-2.46.4-cp313-cp313-win_amd64.whl", hash = "sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292", size = 2071819, upload-time = "2026-05-06T13:38:49.139Z" }, - { url = "https://files.pythonhosted.org/packages/c6/1a/f4aee670d5670e9e148e0c82c7db98d780be566c6e6a97ee8035528ca0b3/pydantic_core-2.46.4-cp313-cp313-win_arm64.whl", hash = "sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d", size = 2027411, upload-time = "2026-05-06T13:40:45.796Z" }, - { url = "https://files.pythonhosted.org/packages/8d/74/228a26ddad29c6672b805d9fd78e8d251cd04004fa7eed0e622096cd0250/pydantic_core-2.46.4-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb", size = 2102079, upload-time = "2026-05-06T13:38:41.019Z" }, - { url = "https://files.pythonhosted.org/packages/ad/1f/8970b150a4b4365623ae00fc88603491f763c627311ae8031e3111356d6e/pydantic_core-2.46.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462", size = 1952179, upload-time = "2026-05-06T13:36:59.812Z" }, - { url = "https://files.pythonhosted.org/packages/95/30/5211a831ae054928054b2f79731661087a2bc5c01e825c672b3a4a8f1b3e/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9", size = 1978926, upload-time = "2026-05-06T13:37:39.933Z" }, - { url = "https://files.pythonhosted.org/packages/57/e9/689668733b1eb67adeef047db3c2e8788fcf65a7fd9c9e2b46b7744fe245/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4", size = 2046785, upload-time = "2026-05-06T13:38:01.995Z" }, - { url = "https://files.pythonhosted.org/packages/60/d9/6715260422ff50a2109878fd24d948a6c3446bb2664f34ee78cd972b3acd/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914", size = 2228733, upload-time = "2026-05-06T13:40:50.371Z" }, - { url = "https://files.pythonhosted.org/packages/18/ae/fdb2f64316afca925640f8e70bb1a564b0ec2721c1389e25b8eb4bf9a299/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28", size = 2307534, upload-time = "2026-05-06T13:37:21.531Z" }, - { url = "https://files.pythonhosted.org/packages/89/1d/8eff589b45bb8190a9d12c49cfad0f176a5cbd1534908a6b5125e2886239/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b", size = 2099732, upload-time = "2026-05-06T13:39:31.942Z" }, - { url = "https://files.pythonhosted.org/packages/06/d5/ee5a3366637fee41dee51a1fc91562dcf12ddbc68fda34e6b253da2324bb/pydantic_core-2.46.4-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c", size = 2129627, upload-time = "2026-05-06T13:37:25.033Z" }, - { url = "https://files.pythonhosted.org/packages/94/33/2414be571d2c6a6c4d08be21f9292b6d3fdb08949a97b6dfe985017821db/pydantic_core-2.46.4-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb", size = 2179141, upload-time = "2026-05-06T13:37:14.046Z" }, - { url = "https://files.pythonhosted.org/packages/7b/79/7daa95be995be0eecc4cf75064cb33f9bbbfe3fe0158caf2f0d4a996a5c7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898", size = 2184325, upload-time = "2026-05-06T13:36:53.615Z" }, - { url = "https://files.pythonhosted.org/packages/9f/cb/d0a382f5c0de8a222dc61c65348e0ce831b1f68e0a018450d31c2cace3a5/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e", size = 2323990, upload-time = "2026-05-06T13:40:29.971Z" }, - { url = "https://files.pythonhosted.org/packages/05/db/d9ba624cc4a5aced1598e88c04fdbd8310c8a69b9d38b9a3d39ce3a61ed7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519", size = 2369978, upload-time = "2026-05-06T13:37:23.027Z" }, - { url = "https://files.pythonhosted.org/packages/f2/20/d15df15ba918c423461905802bfd2981c3af0bfa0e40d05e13edbfa48bc3/pydantic_core-2.46.4-cp314-cp314-win32.whl", hash = "sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4", size = 1966354, upload-time = "2026-05-06T13:38:03.499Z" }, - { url = "https://files.pythonhosted.org/packages/fc/b6/6b8de4c0a7d7ab3004c439c80c5c1e0a3e8d78bbae19379b01960383d9e5/pydantic_core-2.46.4-cp314-cp314-win_amd64.whl", hash = "sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac", size = 2072238, upload-time = "2026-05-06T13:39:40.807Z" }, - { url = "https://files.pythonhosted.org/packages/32/36/51eb763beec1f4cf59b1db243a7dcc39cbb41230f050a09b9d69faaf0a48/pydantic_core-2.46.4-cp314-cp314-win_arm64.whl", hash = "sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a", size = 2018251, upload-time = "2026-05-06T13:37:26.72Z" }, - { url = "https://files.pythonhosted.org/packages/e8/91/855af51d625b23aa987116a19e231d2aaef9c4a415273ddc189b79a45fee/pydantic_core-2.46.4-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0", size = 2099593, upload-time = "2026-05-06T13:39:47.682Z" }, - { url = "https://files.pythonhosted.org/packages/fb/1b/8784a54c65edb5f49f0a14d6977cf1b209bba85a4c77445b255c2de58ab3/pydantic_core-2.46.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d", size = 1935226, upload-time = "2026-05-06T13:40:40.428Z" }, - { url = "https://files.pythonhosted.org/packages/e8/e7/1955d28d1afc56dd4b3ad7cc0cf39df1b9852964cf16e5d13912756d6d6b/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b", size = 1974605, upload-time = "2026-05-06T13:37:32.029Z" }, - { url = "https://files.pythonhosted.org/packages/93/e2/3fedbf0ba7a22850e6e9fd78117f1c0f10f950182344d8a6c535d468fdd8/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000", size = 2030777, upload-time = "2026-05-06T13:38:55.239Z" }, - { url = "https://files.pythonhosted.org/packages/f8/61/46be275fcaaba0b4f5b9669dd852267ce1ff616592dccf7a7845588df091/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e", size = 2236641, upload-time = "2026-05-06T13:37:08.096Z" }, - { url = "https://files.pythonhosted.org/packages/60/db/12e93e46a8bac9988be3c016860f83293daea8c716c029c9ace279036f2f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd", size = 2286404, upload-time = "2026-05-06T13:40:20.221Z" }, - { url = "https://files.pythonhosted.org/packages/e2/4a/4d8b19008f38d31c53b8219cfedc2e3d5de5fe99d90076b7e767de29274f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3", size = 2109219, upload-time = "2026-05-06T13:38:12.153Z" }, - { url = "https://files.pythonhosted.org/packages/88/70/3cbc40978fefb7bb09c6708d40d4ad1a5d70fd7213c3d17f971de868ec1f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7", size = 2110594, upload-time = "2026-05-06T13:40:02.971Z" }, - { url = "https://files.pythonhosted.org/packages/9d/20/b8d36736216e29491125531685b2f9e61aa5b4b2599893f8268551da3338/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff", size = 2159542, upload-time = "2026-05-06T13:39:27.506Z" }, - { url = "https://files.pythonhosted.org/packages/1d/a2/367df868eb584dacf6bf82a389272406d7178e301c4ac82545ab98bc2dd9/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424", size = 2168146, upload-time = "2026-05-06T13:38:31.93Z" }, - { url = "https://files.pythonhosted.org/packages/c1/b8/4460f77f7e201893f649a29ab355dddd3beee8a97bcb1a320db414f9a06e/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6", size = 2306309, upload-time = "2026-05-06T13:37:44.717Z" }, - { url = "https://files.pythonhosted.org/packages/64/c4/be2639293acd87dc8ddbcec41a73cee9b2ebf996fe6d892a1a74e88ad3f7/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565", size = 2369736, upload-time = "2026-05-06T13:37:05.645Z" }, - { url = "https://files.pythonhosted.org/packages/30/a6/9f9f380dbb301f67023bf8f707aaa75daadf84f7152d95c410fd7e81d994/pydantic_core-2.46.4-cp314-cp314t-win32.whl", hash = "sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02", size = 1955575, upload-time = "2026-05-06T13:38:51.116Z" }, - { url = "https://files.pythonhosted.org/packages/40/1f/f1eb9eb350e795d1af8586289746f5c5677d16043040d63710e22abc43c9/pydantic_core-2.46.4-cp314-cp314t-win_amd64.whl", hash = "sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5", size = 2051624, upload-time = "2026-05-06T13:38:21.672Z" }, - { url = "https://files.pythonhosted.org/packages/f6/d2/42dd53d0a85c27606f316d3aa5d2869c4e8470a5ed6dec30e4a1abe19192/pydantic_core-2.46.4-cp314-cp314t-win_arm64.whl", hash = "sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596", size = 2017325, upload-time = "2026-05-06T13:40:52.723Z" }, - { url = "https://files.pythonhosted.org/packages/9d/1d/8987ad40f65ae1432753072f214fb5c74fe47ffbd0698bb9cbbb585664f8/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7", size = 2095527, upload-time = "2026-05-06T13:39:52.283Z" }, - { url = "https://files.pythonhosted.org/packages/64/d3/84c282a7eee1d3ac4c0377546ef5a1ea436ce26840d9ac3b7ed54a377507/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df", size = 1936024, upload-time = "2026-05-06T13:40:15.671Z" }, - { url = "https://files.pythonhosted.org/packages/d7/ca/eac61596cdeb4d7e174d3dc0bd8a6238f14f75f97a24e7b7db4c7e7340a0/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526", size = 1990696, upload-time = "2026-05-06T13:38:34.717Z" }, - { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, -] - [[package]] name = "pygments" version = "2.20.0" @@ -2374,15 +2260,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/a7/d2782e4e3f77c8450f727ba74a8f12756d5ba823d81b941f1b04da9d033a/sphinxcontrib_serializinghtml-2.0.0-py3-none-any.whl", hash = "sha256:6e2cb0eef194e10c27ec0023bfeb25badbbb5868244cf5bc5bdc04e4464bf331", size = 92072, upload-time = "2024-07-29T01:10:08.203Z" }, ] -[[package]] -name = "structlog" -version = "26.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5e/89/b4a0bcfdf4f71a3dea31379f095929613d7e4528a0996bca6aa964cd0dca/structlog-26.1.0.tar.gz", hash = "sha256:f63a716cbd1b1291cf7661de7794b455acfa4c43c5bcf1630e6ad5ddc1adb3b7", size = 1459881, upload-time = "2026-06-06T07:33:39.348Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a9/18/489c97b834dfff9cf2fc2507cede4bcd4b11e67f84bc462acd1992496f86/structlog-26.1.0-py3-none-any.whl", hash = "sha256:e081a26d6c373e6d201eca24eede26d8ffab07f88f477822e679183428d3d91e", size = 73764, upload-time = "2026-06-06T07:33:38.046Z" }, -] - [[package]] name = "threadpoolctl" version = "3.6.0" @@ -2437,18 +2314,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] -[[package]] -name = "typing-inspection" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, -] - [[package]] name = "tzdata" version = "2026.2" @@ -2652,3 +2517,15 @@ sdist = { url = "https://files.pythonhosted.org/packages/93/8d/aeb164004f87543b0 wheels = [ { url = "https://files.pythonhosted.org/packages/88/0a/469e2bd01be1490336e6c8707386845655d59261543315778a3ccc7e8019/zarr-3.2.1-py3-none-any.whl", hash = "sha256:f78cdd3d9687ad0e9f9cba2c5683b64f0c52589c19f685eeabe872e93cc0d2c7", size = 319617, upload-time = "2026-05-05T12:37:20.66Z" }, ] + +[[package]] +name = "zarr-cm" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f9/c6/51d38cafa07bdf725c77a67719647eea43255f925fda062b40799bbb361d/zarr_cm-0.4.1.tar.gz", hash = "sha256:693d24ca2b8e3a7230e1ed448c2c57f98833b81e5899607b6cb4a7ac31f94002", size = 50839, upload-time = "2026-06-21T19:20:36.277Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/14/3de8976647909b2f6764b8d0c9add0a961667ef4ab46f3f99e12d7a522fa/zarr_cm-0.4.1-py3-none-any.whl", hash = "sha256:2c7f36383af2e6f75eb274a797a13154fd5c5827390de7c79294dfeb9f0543ee", size = 33005, upload-time = "2026-06-21T19:20:35.05Z" }, +]