aifs/era5_worker.py

"""
aifs.era5_worker
================
Standalone subprocess entry point — reads EarthMover's public ERA5
Icechunk/Zarr-v3 store on S3 (anonymous access, no account needed).

This script is deliberately NOT imported by the rest of the app. It is
invoked via ``subprocess`` with ``PYTHONPATH`` pointed at an isolated
``pip install --target`` directory (see :mod:`aifs.era5_env`), because
``icechunk`` requires ``zarr>=3`` while ``anemoi-datasets`` (already a
hard dependency of this Space, for AIFS inference) pins ``zarr<=2.18``.
The two cannot coexist in one interpreter's import path, so this script
only ever runs in a separate process with its own isolated zarr install.

Only stdlib + numpy + icechunk + zarr are imported here — keep it that
way so it never accidentally picks up the main env's (incompatible)
zarr via a transitive import.

Usage
-----
    python3 era5_worker.py <request.json> <response_prefix>

``request.json`` is a list of ``{"group", "var", "level", "time_idx"}``
dicts. Writes ``<response_prefix>.npz`` (arrays keyed by request index,
one per successfully-read request) and ``<response_prefix>.meta.json``
(``{"errors": {...}, "resolved_levels": {...}}``). Progress is reported
as ``PROGRESS <i>/<n> <group>/<var>`` lines on stdout.
"""

import json
import sys
import time

import numpy as np

BUCKET = "earthmover-icechunk-era5"
PREFIX = "icechunkV2"
REGION = "us-east-1"

MAX_RETRIES = 6
_RETRIABLE_KEYWORDS = (
    "429", "rate limit", "too many requests", "timeout", "connection reset",
    "503", "service unavailable", "throughput", "streaming error", "i/o error",
)


def _retriable(exc: Exception) -> bool:
    msg = str(exc).lower()
    return any(k in msg for k in _RETRIABLE_KEYWORDS)


def _with_retry(fn, label: str):
    for attempt in range(MAX_RETRIES):
        try:
            return fn()
        except Exception as exc:
            if attempt < MAX_RETRIES - 1 and _retriable(exc):
                wait = min(3 * (2 ** attempt), 30)
                print(f"RETRY {label}: {exc} (attempt {attempt + 2}/{MAX_RETRIES}, waiting {wait}s)", flush=True)
                time.sleep(wait)
            else:
                raise


def _open_store():
    import icechunk

    def _open():
        storage = icechunk.s3_storage(bucket=BUCKET, prefix=PREFIX, region=REGION, anonymous=True)
        repo = icechunk.Repository.open(storage)
        session = repo.readonly_session("main")
        return session.store

    return _with_retry(_open, "open icechunk repo")


def main():
    request_path, response_prefix = sys.argv[1], sys.argv[2]
    with open(request_path) as f:
        requests = json.load(f)

    import zarr

    store = _open_store()
    groups = {}
    level_coords = {}

    results = {}
    errors = {}
    resolved_levels = {}

    n = len(requests)
    for i, req in enumerate(requests):
        group, var, level, time_idx = req["group"], req["var"], req.get("level"), req["time_idx"]
        print(f"PROGRESS {i + 1}/{n} {group}/{var}" + (f"@{level}hPa" if level is not None else ""), flush=True)
        try:
            if group not in groups:
                groups[group] = _with_retry(
                    lambda g=group: zarr.open_group(store, mode="r", path=f"{g}/spatial"),
                    f"open group {group}",
                )
            g = groups[group]

            if var not in g.array_keys():
                raise KeyError(f"'{var}' not found in ERA5 group '{group}' (have: {sorted(g.array_keys())})")

            arr = g[var]

            def _read():
                if level is not None:
                    if group not in level_coords:
                        level_coords[group] = np.asarray(g["pressure_level"][:])
                    levels = level_coords[group]
                    pos = int(np.argmin(np.abs(levels - level)))
                    resolved_levels[str(i)] = float(levels[pos])
                    return arr[time_idx, pos]
                return arr[time_idx]

            data = _with_retry(_read, f"read {group}/{var}")
            results[str(i)] = np.asarray(data, dtype=np.float32)
        except Exception as exc:
            errors[str(i)] = str(exc)

    np.savez_compressed(f"{response_prefix}.npz", **results)
    with open(f"{response_prefix}.meta.json", "w") as f:
        json.dump({"errors": errors, "resolved_levels": resolved_levels}, f)


if __name__ == "__main__":
    main()