jtlevine/gencast-benchmark

GenCast Benchmark API

Single-purpose Space for running Google DeepMind's GenCast 1.0° (full checkpoint, 12-member ensemble) on historical ERA5 initial conditions for the LastMileBench weather advisory benchmark.

Requires 4× A100 80GB hardware (a100-large ×4 on HF Spaces). Returns per-station probabilistic 0–72h cumulative rainfall (ensemble + quantiles + threshold-exceedance probabilities) rather than a single deterministic forecast.

API

GET /health — returns model + GPU status.

POST /forecast_batch

{
  "date": "2018-08-15",
  "points": [
    {"station_id": "KL_TVM", "lat": 8.4833, "lon": 76.95},
    {"station_id": "KL_COK", "lat": 9.95,   "lon": 76.27}
  ]
}

Response:

{
  "forecasts": {
    "KL_TVM": {
      "rain_72h_ensemble": [12 floats, mm],
      "rain_72h_p10": 4.2,
      "rain_72h_p50": 11.3,
      "rain_72h_p90": 28.1,
      "rain_72h_prob_5mm":  0.92,
      "rain_72h_prob_10mm": 0.75,
      "rain_72h_prob_15mm": 0.42,
      "rain_72h_prob_20mm": 0.25
    },
    "...": "..."
  },
  "model": "gencast_1p0_full",
  "checkpoint_name": "gencast/params/GenCast 1p0deg <2019.npz",
  "ensemble_size": 12,
  "init_time": "2018-08-15",
  "chosen_era5_date": "2018-08-15",
  "rollout_wall_s": 118.7,
  "inference_time_s": 118.7,
  "n_points": 2
}

Design notes

• —72h window (0–72h cumulative): steps 0–5 of the 12h rollout, summed. Different from the production-client default of 24–48h — the advisory benchmark scores 72h cumulative rainfall against per-station thresholds.
• —Ensemble: 12 members = 3 batches × 4 pmap devices. Batch seeds derived via jax.random.fold_in(PRNGKey(0), batch_idx) so runs are reproducible.
• —ERA5 source: ARCO ERA5 Zarr (public GCS bucket), coarsened 0.25° → 1.0° before inference to avoid the mesh2grid OOM that 0.25° causes on a 1.0° model.
• —ERA5T lag handling: if date is too recent for ARCO, walks back up to 12 days — chosen_era5_date reports what was actually used.
• —JIT cost: ~4–5 min on first inference call, ~2 min steady-state.

Based on weather-ai-pipeline/src/gencast_client.py (vendored verbatim to keep the Space self-contained).