{
"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 initial conditions and forecast outputs arrays on a [Tigris storage backend](https://www.tigrisdata.com)."
]
},
{
"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 plot_utils\n",
"import seaborn as sns\n",
"import xarray as xr\n",
"import zarr\n",
"\n",
"from aifs_modal import app, ingest, run_forecast"
]
},
{
"cell_type": "markdown",
"id": "2",
"metadata": {},
"source": [
"The main parameters that we need to define are the start date for the forecast and the desired lead time:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3",
"metadata": {
"tags": [
"parameters"
]
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Start date: 2026-03-16 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 = dt.datetime(2026, 3, 16, 0, 0, tzinfo=dt.timezone.utc)\n",
"print(f\"Start date: {start_date}\")\n",
"lead_time = 96 # in hours\n",
"\n",
"# storage parameters\n",
"storage_bucket = \"aifs-modal\"\n",
"# initial conditions\n",
"initial_conditions_prefix = \"aifs-initial-conditions\"\n",
"initial_conditions_branch = \"main\""
]
},
{
"cell_type": "markdown",
"id": "4",
"metadata": {},
"source": [
"## Ingesting initial conditions into an icechunk repository\n",
"\n",
"In order to be able to run a forecast, we first need to ensure that the initial conditions data for the given dates is available in the initial conditions icechunk repository. Let us first set up the connection to the repository, which is hosted to a [S3-like bucket from Tigris](https://www.tigrisdata.com):"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" \u001b[2m2026-03-18T13:08:29.900307Z\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_1773674412996/_build_env/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/aws-runtime-1.5.18/src/env_config/normalize.rs:121\n",
"\n"
]
}
],
"source": [
"initial_conditions_storage = icechunk.tigris_storage(\n",
" bucket=storage_bucket,\n",
" prefix=initial_conditions_prefix,\n",
" region=os.getenv(\"AWS_REGION\", None),\n",
" access_key_id=os.environ[\"AWS_ACCESS_KEY_ID\"],\n",
" secret_access_key=os.environ[\"AWS_SECRET_ACCESS_KEY\"],\n",
")\n",
"initial_conditions_repo = icechunk.Repository.open_or_create(initial_conditions_storage)"
]
},
{
"cell_type": "markdown",
"id": "6",
"metadata": {},
"source": [
"Now, we will ensure that the initial conditions data for the target and previous (6h before) date:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Initial conditions already ingested for 2026-03-15T18:00:00+00:00 (group: 2026-03-15/18z); skipping\n",
"Initial conditions already ingested for 2026-03-16T00:00:00+00:00 (group: 2026-03-16/00z); skipping\n"
]
}
],
"source": [
"for date in [start_date - dt.timedelta(hours=6), start_date]:\n",
" ingest.ensure_date_ingested(\n",
" date, initial_conditions_repo, branch=initial_conditions_branch\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "8",
"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 that depending on your connectivity, ingestion can take a while. However, as noted [in the \"Running AI Weather Forecasts in the Cloud with ECMWF AIFS, Earthmover, and Coiled\" webinar](https://earthmover.io/blog/the-3-key-optimizations-that-cut-the-cost-of-ai-weather-forecasts-by-90), most of this time is spent downloading and parsing GRIB files, a process **which does not require GPU use** and thus is run locally in the present notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Ran 96-hour deterministic forecast in 148.7s\n"
]
}
],
"source": [
"# run a single (deterministic) forecast on modal\n",
"single_outputs_prefix = \"aifs-outputs-single\"\n",
"single_outputs_branch = \"main\"\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",
" lead_time=lead_time,\n",
" initial_conditions_prefix=initial_conditions_prefix,\n",
" initial_conditions_branch=initial_conditions_branch,\n",
" outputs_prefix=single_outputs_prefix,\n",
" outputs_branch=single_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": "10",
"metadata": {},
"source": [
"Once the forecast is done, we can read it from the outputs icechunk repository:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "11",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" \u001b[2m2026-03-18T13:11:10.195998Z\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_1773674412996/_build_env/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/aws-runtime-1.5.18/src/env_config/normalize.rs:121\n",
"\n"
]
}
],
"source": [
"outputs_storage = icechunk.tigris_storage(\n",
" bucket=storage_bucket,\n",
" prefix=single_outputs_prefix,\n",
" region=os.getenv(\"AWS_REGION\", None),\n",
" access_key_id=os.environ[\"AWS_ACCESS_KEY_ID\"],\n",
" secret_access_key=os.environ[\"AWS_SECRET_ACCESS_KEY\"],\n",
")\n",
"outputs_repo = icechunk.Repository.open(outputs_storage)\n",
"# here we just need to read data\n",
"outputs_session = outputs_repo.readonly_session(single_outputs_branch)\n",
"outputs_root = zarr.open_group(outputs_session.store, mode=\"r\", zarr_format=3)"
]
},
{
"cell_type": "markdown",
"id": "12",
"metadata": {},
"source": [
"We can first see the dates for which we have forecast outputs stored:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "13",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"['2026-03-16']\n"
]
}
],
"source": [
"# forecasts = [\"00z\", \"06z\", \"12z\", \"18z\"]\n",
"dates = list(outputs_root.group_keys())\n",
"dates.sort(reverse=True)\n",
"print(dates)"
]
},
{
"cell_type": "markdown",
"id": "14",
"metadata": {},
"source": [
"read the data into an xarray data set:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "15",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"