Spaces:
Running on Zero
Running on Zero
File size: 3,566 Bytes
80f643f 6a64990 80f643f 6a64990 80f643f 8ae4133 80f643f 6a64990 8ae4133 6a64990 80f643f 8ae4133 80f643f 6a64990 80f643f 8ae4133 6a64990 80f643f 6a64990 80f643f 6a64990 8ae4133 80f643f 6a64990 80f643f 6a64990 80f643f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | 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
|