aifs/forecast.py

import os
import datetime
import queue
import threading
import time

# Apply flash-attn shim before any anemoi import
import aifs.compat  # noqa: F401
from aifs.device import get_device, device_label

DEFAULT_CHECKPOINT = "aifs-single-2.0"

CHECKPOINTS = {
    DEFAULT_CHECKPOINT: {"huggingface": f"ecmwf/{DEFAULT_CHECKPOINT}"},
}

HEARTBEAT_INTERVAL = 10  # seconds between "still computing" messages


def run_forecast(
    fields: dict,
    date: datetime.datetime,
    lead_time: int = 24,
    num_chunks: int = 16,
    checkpoint: str = DEFAULT_CHECKPOINT,
    verbose: bool = True,
):
    """
    Generator that runs the AIFS forecast and streams progress.

    Yields
    ------
    ("log", str)       -- accumulated progress log
    ("result", list)   -- final list of state dicts (last item yielded)
    """
    log_lines: list[str] = []

    def log(msg: str):
        log_lines.append(msg)
        return "log", "\n".join(log_lines)

    if lead_time % 6 != 0:
        raise ValueError(f"lead_time must be a multiple of 6, got {lead_time}")

    from anemoi.inference.runners.simple import SimpleRunner

    device = get_device()

    if device == "cuda":
        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    os.environ["ANEMOI_INFERENCE_NUM_CHUNKS"] = str(num_chunks)

    total_steps = lead_time // 6

    if verbose:
        yield log(f"🖥️   Device  : {device_label()}")
        yield log(f"📦  Checkpoint: {checkpoint}")
        yield log(f"⏱️   Lead time : {lead_time} h  ({total_steps} steps)")

    ckpt = CHECKPOINTS.get(checkpoint, checkpoint)

    if verbose:
        yield log("🤖  Loading model …")

    runner = SimpleRunner(ckpt)
    log_lines.clear()

    if verbose:
        yield log(f"🌍  Running inference — {total_steps} steps…")

    states: list[dict] = []
    input_state = {"fields": fields, "date": date}

    # Run runner.run() in a worker thread so we can send heartbeat messages
    # to the UI while each step is computing (steps can take tens of seconds).
    result_queue: queue.Queue = queue.Queue()
    last_step_time = [time.time()]

    def _worker():
        try:
            for state in runner.run(input_states=input_state, lead_time=lead_time):
                result_queue.put(("state", state))
                last_step_time[0] = time.time()
            result_queue.put(("done", None))
        except Exception as exc:
            result_queue.put(("error", exc))

    thread = threading.Thread(target=_worker, daemon=True)
    thread.start()

    while True:
        try:
            kind, payload = result_queue.get(timeout=HEARTBEAT_INTERVAL)
        except queue.Empty:
            elapsed = time.time() - last_step_time[0]
            step_num = len(states) + 1
            if verbose:
                yield log(f"⏳  Step {step_num}/{total_steps} still computing… ({elapsed:.0f}s)")
            continue

        if kind == "state":
            states.append({
                "date": payload["date"],
                "fields": {k: v.copy() for k, v in payload["fields"].items()},
                "latitudes": payload["latitudes"],
                "longitudes": payload["longitudes"],
            })
            if verbose:
                yield log(f"✓  Step {len(states)}/{total_steps}: {payload['date']}")
        elif kind == "done":
            break
        elif kind == "error":
            raise payload

    thread.join(timeout=5.0)

    if verbose:
        yield log(f"✅  Done — {len(states)} steps produced.")

    yield "result", states