{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Running your first serverless AIFS forecast on Modal\n", "\n", "This notebook shows a serverless approach to run [the Artificial Intelligence Forecasting System (AIFS) model](https://huggingface.co/ecmwf/aifs-single-1.1) on [Modal](https://modal.com), using [Icechunk](https://github.com/earth-mover/icechunk) as storage engine to manage forecast outputs on a [Tigris storage backend](https://www.tigrisdata.com).\n", "\n", "```{note}\n", "Before running this notebook, complete the [setup steps](index.md): Modal account and CLI, a Tigris bucket, an Earthmover account with Brightband subscription (for the default `ifs-arraylake` IC source), and the three Modal secrets (`aws-credentials`, `arraylake-api-token`, `huggingface-secret`). Then run `modal run -m aifs_modal._app` once to create the required Modal Volumes.\n", "```" ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "import datetime as dt\n", "import os\n", "import time\n", "\n", "import icechunk\n", "import matplotlib.pyplot as plt\n", "import seaborn as sns\n", "import xarray as xr\n", "import zarr\n", "\n", "import plot_utils\n", "from aifs_modal import app, run_forecast\n", "\n", "\n", "def add_valid_time(ds):\n", " valid_time = ds[\"init_time\"].isel(init_time=0) + ds[\"lead_time\"]\n", " return ds.assign_coords(valid_time=(\"lead_time\", valid_time.values))" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "The main parameters to define are the start date, lead time, and initial-conditions source. The default IC source is `ifs-arraylake` (ECMWF IFS data from the Brightband dataset on the Earthmover marketplace — fastest and most reliable). For alternatives including free no-credentials sources, see the [initial conditions notebook](initial-conditions.ipynb)." ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": { "tags": [ "parameters" ] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Start date: 2026-05-12 00:00:00+00:00\n" ] } ], "source": [ "# start_date = dt.datetime.fromisoformat(\"2026-02-16T06:00:00+00:00\")\n", "# start_date = OpendataClient().latest().replace(tzinfo=dt.UTC)\n", "# start date for IFS Brightband/Earthmover initial conditions: 1 day ago (safely within\n", "# the rolling window)\n", "start_date = dt.datetime.now(dt.UTC).replace(\n", " hour=0, minute=0, second=0, microsecond=0\n", ") - dt.timedelta(days=1)\n", "print(f\"Start date: {start_date}\")\n", "lead_time = 96 # in hours\n", "# n ensemble members\n", "n_members = 10\n", "# Brightband source repo (ECMWF IFS initial conditions on the Earthmover marketplace)\n", "ifs_source_repo = \"martibosch/ecmwf-ifs-hres-ics-open\"\n", "\n", "# storage parameters\n", "storage_bucket = \"aifs-modal-unibe\"\n", "# deterministic forecast outputs\n", "det_outputs_prefix = \"aifs-outputs-det\"\n", "det_outputs_branch = \"main\"\n", "# ensemble forecast outputs\n", "ens_outputs_prefix = \"aifs-outputs-ens\"\n", "ens_outputs_branch = \"main\"" ] }, { "cell_type": "markdown", "id": "4", "metadata": {}, "source": [ "## Deterministic forecast (AIFS-Single)\n", "\n", "The [AIFS-Single](https://huggingface.co/ecmwf/aifs-single-1.1) model produces a single deterministic forecast. We run it on Modal and store the outputs in the `aifs-outputs` prefix.\n", "\n", "```{note}\n", "`run_forecast` is idempotent: before ingesting ICs or starting inference it checks whether the forecast output already exists on the Modal side, and skips if so. To avoid spinning up the ephemeral app altogether, call `aifs_modal.forecast_exists(date, storage_bucket, ...)` locally first — it reads directly from the output store without touching Modal.\n", "```" ] }, { "cell_type": "code", "execution_count": null, "id": "5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Ran 96-hour deterministic forecast in 168.6s\n" ] } ], "source": [ "# run a single (deterministic) forecast on modal\n", "\n", "# with modal.enable_output():\n", "with app.run():\n", " _start = time.perf_counter()\n", " run_forecast.remote(\n", " start_date,\n", " storage_bucket,\n", " ic_source_repo=ifs_source_repo,\n", " lead_time=lead_time,\n", " outputs_prefix=det_outputs_prefix,\n", " outputs_branch=det_outputs_branch,\n", " )\n", " print(\n", " f\"Ran {lead_time}-hour deterministic forecast in \"\n", " f\"{time.perf_counter() - _start:.1f}s\"\n", " )" ] }, { "cell_type": "markdown", "id": "6", "metadata": {}, "source": [ "Once the forecast is done, we can read it from the outputs icechunk repository:" ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " \u001b[2m2026-05-16T08:13:09.720690Z\u001b[0m \u001b[33m WARN\u001b[0m \u001b[1;33maws_runtime::env_config::normalize\u001b[0m\u001b[33m: \u001b[33mprofile [plugins] ignored; sections in the AWS config file (other than [default]) must have a prefix i.e. [profile my-profile]\u001b[0m\n", " \u001b[2;3mat\u001b[0m /home/conda/feedstock_root/build_artifacts/icechunk_1776344638779/_build_env/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/aws-runtime-1.7.2/src/env_config/normalize.rs:121\n", "\n" ] }, { "data": { "text/html": [ "
<xarray.Dataset> Size: 2GB\n",
"Dimensions: (init_time: 1, lead_time: 16, lat: 721, lon: 1440, pressure: 13)\n",
"Coordinates:\n",
" * init_time (init_time) datetime64[ns] 8B 2026-05-12\n",
" * lead_time (lead_time) timedelta64[us] 128B 0 days 06:00:00 ... 4 days 0...\n",
" valid_time (lead_time) datetime64[ns] 128B 2026-05-12T06:00:00 ... 2026-...\n",
" * lat (lat) float64 6kB 90.0 89.75 89.5 89.25 ... -89.5 -89.75 -90.0\n",
" * lon (lon) float64 12kB 0.0 0.25 0.5 0.75 ... 359.0 359.2 359.5 359.8\n",
" * pressure (pressure) int64 104B 50 100 150 200 250 ... 700 850 925 1000\n",
"Data variables: (12/24)\n",
" 100u (init_time, lead_time, lat, lon) float32 66MB ...\n",
" 10u (init_time, lead_time, lat, lon) float32 66MB ...\n",
" 10v (init_time, lead_time, lat, lon) float32 66MB ...\n",
" 2t (init_time, lead_time, lat, lon) float32 66MB ...\n",
" 2d (init_time, lead_time, lat, lon) float32 66MB ...\n",
" cp (init_time, lead_time, lat, lon) float32 66MB ...\n",
" ... ...\n",
" swvl1 (init_time, lead_time, lat, lon) float32 66MB ...\n",
" tcc (init_time, lead_time, lat, lon) float32 66MB ...\n",
" tp (init_time, lead_time, lat, lon) float32 66MB ...\n",
" swvl2 (init_time, lead_time, lat, lon) float32 66MB ...\n",
" tcw (init_time, lead_time, lat, lon) float32 66MB ...\n",
" strd (init_time, lead_time, lat, lon) float32 66MB ...<xarray.Dataset> Size: 15GB\n",
"Dimensions: (ensemble_member: 10, init_time: 1, lead_time: 16,\n",
" lat: 721, lon: 1440, pressure: 13)\n",
"Coordinates:\n",
" * ensemble_member (ensemble_member) int64 80B 0 1 2 3 4 5 6 7 8 9\n",
" * init_time (init_time) datetime64[ns] 8B 2026-05-12\n",
" * lead_time (lead_time) timedelta64[us] 128B 0 days 06:00:00 ... 4 d...\n",
" valid_time (lead_time) datetime64[ns] 128B 2026-05-12T06:00:00 ... ...\n",
" * lat (lat) float64 6kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0\n",
" * lon (lon) float64 12kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8\n",
" * pressure (pressure) int64 104B 50 100 150 200 ... 700 850 925 1000\n",
"Data variables: (12/22)\n",
" 100u (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" 100v (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" 10v (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" 10u (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" cp (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" 2t (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" ... ...\n",
" sp (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" ssrd (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" tp (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" strd (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" tcc (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...\n",
" tcw (ensemble_member, init_time, lead_time, lat, lon) float32 664MB ...