Spaces:
Running on Zero
Running on Zero
| """Prepare Varda-single-1.0 initial conditions in a standalone process. | |
| Fetches the KENDA-CH1 (regional, 1 km) and ECMWF open-data IFS (global) analyses at `init` and | |
| `init - 6h`, preprocesses them exactly like `anemoi-inference run` does, and saves: | |
| * the model input tensor (multi_step_input, variables, values) -> <out> | |
| * the mesh (lat, lon, lam mask) -> <mesh> (if missing) | |
| Running this out-of-process keeps the native GRIB/regridding libraries (eccodes, MIR, atlas…) | |
| out of the ZeroGPU main process, and lets the web app stay responsive meanwhile. | |
| Usage: python prep.py <config.json> <init ISO> <out.npy> <mesh.npz> | |
| """ | |
| import json | |
| import logging | |
| import os | |
| import sys | |
| import threading | |
| import time | |
| from concurrent.futures import Future, ThreadPoolExecutor, wait | |
| from datetime import datetime, timedelta | |
| from eccodes_cosmo_resources import get_definitions_path | |
| os.environ["ECCODES_DEFINITION_PATH"] = str(get_definitions_path()) | |
| os.environ.setdefault("ECCODES_VERSION_CHECK_OFF", "1") | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" # CPU only | |
| import earthkit.data as ekd # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import requests # noqa: E402 | |
| from anemoi_plugins_meteoswiss.inference.inputs import oper_kenda_opendata as kenda # noqa: E402 | |
| from anemoi_plugins_meteoswiss.transform.filters.vertical_interpolation import ModelToPressureLevel # noqa: E402 | |
| from anemoi.inference.config.run import RunConfiguration # noqa: E402 | |
| from anemoi.inference.runners import create_runner # noqa: E402 | |
| logging.basicConfig(level=logging.INFO, format="[prep] %(levelname)s %(name)s: %(message)s") | |
| LOG = logging.getLogger("prep") | |
| # The KENDA-CH1 server is throttled per connection (~1-2 MB/s seen from the Space) but supports HTTP | |
| # Range requests, so every file is fetched as CHUNK-sized byte ranges over many connections at once. | |
| CHUNK = 16 * 1024 * 1024 | |
| CONNECTIONS = 32 | |
| DOWNLOAD_RETRIES = 4 | |
| # KENDA-CH1 fields the Varda input needs at each analysis time (logged by the plugin as | |
| # "fetching ... variable(s)"). Only used to start downloads early: anything missing from this | |
| # list is still fetched by the plugin itself, an extra entry just costs bandwidth. | |
| KENDA_VARS = ["P", "PMSL", "PS", "QV", "T", "TD_2M", "T_2M", "U", "U_10M", "V", "V_10M"] | |
| INPUT_TIMES_H = (-6, 0) # the forecaster needs the analyses at init - 6h and init | |
| _pool = ThreadPoolExecutor(64) # one light coordinating thread per file | |
| _chunks = ThreadPoolExecutor(CONNECTIONS) # the actual transfers; FIFO, so files finish in queue order | |
| _inflight: dict[str, Future] = {} # destination path -> download, shared by prefetch and plugin | |
| _inflight_lock = threading.Lock() | |
| _session = threading.local() | |
| def _http() -> requests.Session: | |
| if not hasattr(_session, "s"): # one keep-alive session per thread | |
| _session.s = requests.Session() | |
| return _session.s | |
| def _retrying(what: str, fn): | |
| for attempt in range(1, DOWNLOAD_RETRIES + 1): | |
| try: | |
| return fn() | |
| except Exception as e: # noqa: BLE001 | |
| LOG.warning(" %s: attempt %d/%d failed: %s", what, attempt, DOWNLOAD_RETRIES, e) | |
| if attempt == DOWNLOAD_RETRIES: | |
| raise | |
| time.sleep(3 * attempt) | |
| def _get_range(href: str, fd: int, start: int, end: int) -> None: | |
| """Fetch bytes [start, end] of `href` into `fd` at the same offset.""" | |
| # (connect, read) timeouts: a transfer that sends nothing for 60 s is retried | |
| with _http().get(href, headers={"Range": f"bytes={start}-{end}"}, stream=True, timeout=(30, 60)) as r: | |
| r.raise_for_status() | |
| if r.status_code != 206: | |
| raise IOError(f"server ignored the Range header (HTTP {r.status_code})") | |
| pos = start | |
| for block in r.iter_content(chunk_size=1024 * 1024): | |
| os.pwrite(fd, block, pos) | |
| pos += len(block) | |
| if pos != end + 1: | |
| raise IOError(f"short range {start}-{end}: got {pos - start} bytes") | |
| def _size(href: str) -> int: | |
| # a 1-byte GET rather than HEAD: presigned URLs are often signed for GET only | |
| with _http().get(href, headers={"Range": "bytes=0-0"}, stream=True, timeout=(30, 60)) as r: | |
| r.raise_for_status() | |
| if r.status_code == 206: | |
| return int(r.headers["Content-Range"].rsplit("/", 1)[1]) | |
| return -1 # no Range support | |
| def _download_stream(href: str, part: str) -> None: | |
| with _http().get(href, stream=True, timeout=(30, 60)) as r, open(part, "wb") as f: | |
| r.raise_for_status() | |
| for block in r.iter_content(chunk_size=4 * 1024 * 1024): | |
| f.write(block) | |
| def _download(href: str, path: str) -> None: | |
| """Download `href` to `path` (via a .part file) as parallel byte ranges, each retried on failure.""" | |
| name = os.path.relpath(path, os.path.dirname(os.path.dirname(path))) | |
| part = path + ".part" | |
| t0 = time.time() | |
| size = _retrying(name, lambda: _size(href)) | |
| if size < 0: | |
| _retrying(name, lambda: _download_stream(href, part)) | |
| else: | |
| fd = os.open(part, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644) | |
| try: | |
| os.ftruncate(fd, size) | |
| ranges = [(s, min(s + CHUNK, size) - 1) for s in range(0, size, CHUNK)] | |
| futs = [ | |
| _chunks.submit(_retrying, f"{name}@{s}", lambda s=s, e=e: _get_range(href, fd, s, e)) | |
| for s, e in ranges | |
| ] | |
| wait(futs) # let every range finish before closing fd, even if one failed for good | |
| for f in futs: | |
| f.result() | |
| finally: | |
| os.close(fd) | |
| os.replace(part, path) | |
| dt = time.time() - t0 | |
| mb = os.path.getsize(path) / 1e6 | |
| LOG.info(" %s: %.0f MB in %.0fs (%.1f MB/s)", name, mb, dt, mb / max(dt, 1e-3)) | |
| def _fetch(href: str, path: str) -> Future | None: | |
| """Start downloading `href` to `path` unless it is cached or already being downloaded.""" | |
| with _inflight_lock: | |
| if path in _inflight: | |
| return _inflight[path] | |
| if os.path.exists(path): | |
| return None | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| fut = _inflight[path] = _pool.submit(_download, href, path) | |
| return fut | |
| def _parallel_cached_fields(cache_dir, subdir, hrefs_by_filename): | |
| """Drop-in for the plugin's `_cached_fields`: same on-disk cache layout, but missing KENDA-CH1 | |
| assets (~180-275 MB each for the 3-D fields) download concurrently, reusing prefetches in flight.""" | |
| if not cache_dir: | |
| return _orig_cached_fields(cache_dir, subdir, hrefs_by_filename) | |
| paths = {fn: os.path.join(cache_dir, subdir, fn) for fn in hrefs_by_filename} | |
| t0 = time.time() | |
| futs = [f for f in (_fetch(hrefs_by_filename[fn], p) for fn, p in paths.items()) if f is not None] | |
| for f in futs: | |
| f.result() | |
| if futs: | |
| LOG.info("KENDA-CH1 %s: waited %.0fs for %d download(s)", subdir, time.time() - t0, len(futs)) | |
| return ekd.from_source("file", list(paths.values())) | |
| _orig_cached_fields = kenda._cached_fields | |
| kenda._cached_fields = _parallel_cached_fields | |
| def _m2p_forward_parallel(self, data): | |
| """ModelToPressureLevel.forward with the analysis times interpolated in parallel threads (the GRIB | |
| decoding and numpy work release the GIL): ~1.6x faster than the plugin's serial loop, same output.""" | |
| groups = [ekd.SimpleFieldList(list(g)) for g in data.group_by("valid_datetime")] | |
| if len(groups) < 2: | |
| return _orig_m2p_forward(self, data) | |
| with ThreadPoolExecutor(len(groups)) as pool: | |
| outs = list(pool.map(lambda g: _orig_m2p_forward(self, g), groups)) | |
| # same order as the serial filter: every interpolated field first, then the surface passthrough | |
| fields = [f for o in outs for f in o] | |
| interp = [f for f in fields if f.metadata("typeOfLevel") == "isobaricInhPa"] | |
| rest = [f for f in fields if f.metadata("typeOfLevel") != "isobaricInhPa"] | |
| return ekd.SimpleFieldList(interp + rest) | |
| _orig_m2p_forward = ModelToPressureLevel.forward | |
| ModelToPressureLevel.forward = _m2p_forward_parallel | |
| def _prefetch_kenda(cache_dir: str, init: datetime) -> None: | |
| """Queue every KENDA-CH1 download the run needs (grid constants + both analysis times) at once, | |
| so the t download overlaps decoding t-6h instead of waiting for it. Best effort.""" | |
| try: | |
| for name, href in kenda._collection_static_assets().items(): | |
| _fetch(href, os.path.join(cache_dir, "constants", name)) | |
| for dh in INPUT_TIMES_H: | |
| t = init + timedelta(hours=dh) | |
| hrefs = kenda._hrefs_by_variable(kenda._stac_items(t)) | |
| for v in KENDA_VARS: | |
| if v in hrefs: | |
| _fetch(hrefs[v], os.path.join(cache_dir, "raw", t.strftime("%Y%m%dT%H%M%S"), f"{v}.grib2")) | |
| LOG.info("prefetch: %d KENDA-CH1 download(s) queued", sum(not f.done() for f in _inflight.values())) | |
| except Exception as e: # noqa: BLE001 | |
| LOG.warning("prefetch failed (the plugin will fetch on demand): %s", e) | |
| def _kenda_cache_dir(cfg: dict) -> str | None: | |
| for src in cfg["input"]["cutout"]: | |
| for name, opts in src.get("lam_0", {}).items(): | |
| if name.replace("_", "-") == "oper-kenda-opendata": | |
| return opts.get("cache_dir") | |
| return None | |
| def main(cfg_path: str, init: str, out: str, mesh: str) -> None: | |
| with open(cfg_path) as f: | |
| cfg = json.load(f) | |
| cfg["device"] = "cpu" | |
| date = datetime.fromisoformat(init) | |
| t0 = time.time() | |
| cache_dir = _kenda_cache_dir(cfg) | |
| if cache_dir: | |
| threading.Thread(target=_prefetch_kenda, args=(cache_dir, date), daemon=True).start() | |
| runner = create_runner(RunConfiguration.load(cfg)) | |
| ds = runner.dataset_names[0] | |
| prog = runner.prognostics_inputs[ds].create_input_state(date=date) | |
| LOG.info("prognostics fetched in %.0fs", time.time() - t0) | |
| const = runner.constant_forcings_inputs[ds].create_input_state(date=date) | |
| forc = runner.dynamic_forcings_inputs[ds].create_input_state(date=date) | |
| LOG.info("forcings ready at %.0fs", time.time() - t0) | |
| state = runner._combine_states(prog, const, forc) | |
| tensor = runner.tensor_handlers[ds].prepare_input_tensor(state) | |
| LOG.info("input tensor %s ready in %.0fs", tensor.shape, time.time() - t0) | |
| if not os.path.exists(mesh): | |
| np.savez( | |
| mesh + ".part.npz", | |
| lat=np.asarray(state["latitudes"], dtype=np.float64), | |
| lon=np.asarray(state["longitudes"], dtype=np.float64), | |
| lam=np.asarray(state["_mask"]["lam_0"], dtype=bool), | |
| ) | |
| os.replace(mesh + ".part.npz", mesh) | |
| np.save(out + ".part.npy", tensor.astype(np.float32)) | |
| os.replace(out + ".part.npy", out) | |
| _pool.shutdown(wait=False, cancel_futures=True) | |
| _chunks.shutdown(wait=False, cancel_futures=True) | |
| if __name__ == "__main__": | |
| main(*sys.argv[1:5]) | |