aifs/era5_env.py

"""
aifs.era5_env
=============
Bootstraps and drives the isolated environment that talks to EarthMover's
ERA5 Icechunk store (see :mod:`aifs.era5_worker`).

Why isolated: ``icechunk`` requires ``zarr>=3``, but ``anemoi-datasets``
(already required for AIFS inference in this Space) pins ``zarr<=2.18``.
pip cannot satisfy both in one environment. Instead we ``pip install
--target=`` a private directory on first use, and run the actual ERA5
reads in a subprocess whose ``PYTHONPATH`` is prepended with that
directory — the subprocess resolves ``zarr`` to the isolated v3 install
regardless of what's importable from the main env's site-packages
(verified: this is plain CPython import-order semantics, not a hack
specific to zarr).
"""

from __future__ import annotations

import json
import os
import platform
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np

_HERE = Path(__file__).resolve().parent
WORKER_SCRIPT = _HERE / "era5_worker.py"
ISOLATED_DIR = _HERE.parent / "era5_env_isolated"
SENTINEL = ISOLATED_DIR / ".bootstrap_ok"

ISOLATED_PACKAGES = ["icechunk>=2.1", "zarr>=3", "pcodec"]


def _fingerprint() -> str:
    """
    Identifies the machine/interpreter the isolated install's compiled
    wheels (numpy, icechunk, pcodec) are built for. Compared against the
    sentinel on every call so a directory built on one machine (e.g. a
    dev sandbox) is never reused on another (e.g. the actual Space host)
    — compiled extensions are platform- and Python-version-specific, and
    a mismatch fails with a confusing "numpy C-extensions" ImportError
    rather than anything that points at the real cause.
    """
    return f"{platform.system()}-{platform.machine()}-py{sys.version_info.major}.{sys.version_info.minor}"


def ensure_bootstrapped(log=lambda msg: None) -> None:
    """Install the isolated zarr>=3 / icechunk stack if not already present for this machine."""
    fingerprint = _fingerprint()
    if SENTINEL.exists() and SENTINEL.read_text().strip() == fingerprint:
        return

    if ISOLATED_DIR.exists():
        log("🔄  Isolated ERA5 environment was built for a different machine — reinstalling…")
        shutil.rmtree(ISOLATED_DIR)

    ISOLATED_DIR.mkdir(parents=True, exist_ok=True)
    log(f"📦  Setting up isolated ERA5 environment (one-time, ~30s)…")
    result = subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q", "--target", str(ISOLATED_DIR), *ISOLATED_PACKAGES],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(
            f"Failed to bootstrap the isolated ERA5 environment:\n{result.stderr[-2000:]}"
        )
    SENTINEL.write_text(fingerprint)
    log("✅  Isolated ERA5 environment ready.")


def _subprocess_env() -> dict:
    env = dict(os.environ)
    existing = env.get("PYTHONPATH", "")
    env["PYTHONPATH"] = f"{ISOLATED_DIR}{os.pathsep}{existing}" if existing else str(ISOLATED_DIR)
    return env


def fetch_era5_fields(requests: list[dict], log=lambda msg: None, timeout: int = 600) -> tuple[dict, dict]:
    """
    Run a batch of ERA5 reads in the isolated subprocess.

    ``requests`` — list of ``{"group", "var", "level", "time_idx"}`` dicts.
    Returns ``(arrays, meta)`` where ``arrays`` maps request index (int) to
    a ``(721, 1440)`` float32 grid, and ``meta`` has ``"errors"`` (index ->
    message, for requests that failed) and ``"resolved_levels"``.
    """
    ensure_bootstrapped(log)

    with tempfile.TemporaryDirectory() as tmp:
        request_path = Path(tmp) / "request.json"
        response_prefix = Path(tmp) / "response"
        request_path.write_text(json.dumps(requests))

        proc = subprocess.Popen(
            [sys.executable, str(WORKER_SCRIPT), str(request_path), str(response_prefix)],
            env=_subprocess_env(),
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1,
        )
        stderr_tail = []
        try:
            for line in proc.stdout:
                line = line.rstrip()
                stderr_tail.append(line)
                if line.startswith("PROGRESS") or line.startswith("RETRY"):
                    log(f"📡  {line}")
            proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            proc.kill()
            raise RuntimeError("Timed out waiting for the ERA5 worker subprocess.")

        if proc.returncode != 0:
            raise RuntimeError(
                "ERA5 worker subprocess failed:\n" + "\n".join(stderr_tail[-30:])
            )

        npz_path = f"{response_prefix}.npz"
        meta_path = f"{response_prefix}.meta.json"
        with np.load(npz_path) as npz:
            arrays = {int(k): npz[k] for k in npz.files}
        meta = json.loads(Path(meta_path).read_text())
        meta["errors"] = {int(k): v for k, v in meta["errors"].items()}
        meta["resolved_levels"] = {int(k): v for k, v in meta["resolved_levels"].items()}
        return arrays, meta