How it works#

This page describes the architecture of aifs_modal, focusing on how AIFS forecasts are run within computational pipelines, highlighting what is executed locally versus on Modal and how data inputs and outputs are stored at each step.

Note that this set up is not intended for running an operational AIFS forecasting platform but rather as a playground to run AIFS forecasts for scientific experiments, where the pre-processing, forecasts and post-processing can be run interactively within the same local notebook, with only the forecasts running on a serverless GPU environment and with a task-based approach that prevents unnecessary duplication of work when re-running cells or iterating on the analysis (see the Re-running existing forecasts and reproducibility section below). Idempotency operates at two levels: run_forecast checks for existing outputs on the Modal side before ingesting ICs or starting inference; forecast_exists provides the same check locally, so the ephemeral app is never spun up when the forecast is already complete. See the example notebooks for practical example applications of this set up.

Key concepts#

AIFS (Artificial Intelligence Forecasting System): ECMWF’s machine-learning weather model. Two checkpoints are available: AIFS-Single (deterministic) and AIFS-ENS (ensemble, stochastic perturbations via random seeds).
Custom checkpoints: any anemoi-inference-compatible checkpoint can override the built-in defaults by passing a checkpoint dict to run_forecast or run_inference (e.g. {"huggingface": "org/fine-tuned-model"} or a path on a Modal Volume). aifs-modal reads variable_to_input_tensor_index directly from the checkpoint and silently drops any IC variables the model does not expect, so the IC source does not need to be trimmed manually. Requirements: (1) the checkpoint must be in anemoi-inference format (a .ckpt file produced by anemoi-training); (2) it must operate on the N320 reduced Gaussian grid — the output regridder is hardcoded to N320 → 0.25° lat/lon; (3) the IC source must cover at minimum the full set of variables the checkpoint requires.
Modal: serverless GPU platform. Forecast inference runs on Modal containers with GPU access. Ingestion and orchestration also run on Modal (CPU), with workers co-located with IC sources for in-region bandwidth.
Icechunk: versioned array storage engine backed by S3-compatible object storage. Used for forecast outputs, providing git-like branching and commit history.
Initial conditions: analysis fields ingested from an IC source (Brightband, ARCO-ERA5, or ECMWF open data), regridded to the N320 reduced Gaussian grid and written to a shared Modal IC Volume. AIFS requires two consecutive 6-hourly analyses (t−6h and t) as input. The Volume is a temporary per-run cache; ICs are deleted after successful inference unless keep_ics=True.
IC sources: four sources are supported — ifs-arraylake (default, Brightband ECMWF IFS on Earthmover ArrayLake, co-located in us-east), era5-arco (ARCO-ERA5 on GCS, co-located in us-central1), ifs-ekd (ECMWF open-data S3), and era5-cds (Copernicus CDS).

The diagrams below use a consistent four-zone color scheme:

Zone	Color	Description
Local	grey	Notebook or script on your machine
Modal CPU	blue	Orchestration and IC ingestion on Modal
Modal GPU	orange	AIFS inference on Modal
Object storage	green	IC source and Icechunk output repository on S3

Deterministic forecast (`run_forecast`)#

A single AIFS-Single run on one GPU. run_forecast is a Modal CPU function that acts as the orchestrator: it checks whether ICs are already on the IC Volume, dispatches ingestion to a co-located Modal CPU worker if they are missing, and then dispatches inference to a GPU container.

        flowchart LR
    subgraph local ["Local"]
        A["configure\nrun_forecast.remote()"]
        J[read outputs]
    end
    subgraph modal_cpu ["Modal — CPU"]
        B["orchestrate\n+ ingest ICs → Volume"]
    end
    subgraph modal_gpu ["Modal — GPU"]
        G["AIFS-Single\ninference"]
    end
    subgraph storage ["Object storage"]
        IC[("IC source")]
        OUT[("Icechunk outputs")]
    end
    A --> B
    IC --> B
    B --> G
    G --> OUT
    OUT --> J

    style local fill:#555,stroke:#333,color:#fff
    style A fill:#777,stroke:#555,color:#fff
    style J fill:#777,stroke:#555,color:#fff
    style modal_cpu fill:#2c6fad,stroke:#1a4a7a,color:#fff
    style B fill:#4a90d9,stroke:#2c6fad,color:#fff
    style modal_gpu fill:#c44d1a,stroke:#8c3210,color:#fff
    style G fill:#ff6b35,stroke:#c44d1a,color:#fff
    style storage fill:#1a5c3a,stroke:#0d3d26,color:#fff
    style IC fill:#2e7d52,stroke:#1a5c3a,color:#fff
    style OUT fill:#2e7d52,stroke:#1a5c3a,color:#fff

Flow:

Orchestrate (Modal CPU, run_forecast): check whether the forecast output already exists; if so, return immediately without ingesting ICs or starting inference. Otherwise, check the IC Volume for the target date. If ICs are missing, dispatch ingest_ifs_arraylake (Modal us-east) or ingest_era5_arco (Modal us-central1) and wait for completion; inline sources (ifs-ekd, era5-cds) run directly inside the orchestrator.
Ingest (Modal CPU, co-located): download and regrid the IC fields to N320; write them to the IC Volume.
Inference (Modal GPU, run_inference): load initial conditions from the IC Volume, load the checkpoint into SimpleRunner, drop any IC variables not required by the checkpoint, run the AIFS model, regrid output fields from N320 to 0.25° lat/lon, and stream each 6-hourly step to the outputs Icechunk repository.
Post-process (local, CPU): open the outputs repository with xarray and analyze them.

Ensemble forecast (`run_forecast` with `n_members`)#

Runs multiple AIFS-ENS members one after another on a single GPU. Each member uses a different random seed for stochastic perturbations. Ingestion and orchestration are identical to the deterministic case; the difference is that run_inference loops over members on the same GPU.

        flowchart LR
    subgraph local ["Local"]
        A["configure\nrun_forecast.remote(n_members=k)"]
        J[read outputs]
    end
    subgraph modal_cpu ["Modal — CPU"]
        B["orchestrate\n+ ingest ICs → Volume"]
    end
    subgraph modal_gpu ["Modal — GPU"]
        G["AIFS-ENS\nmember 0 → 1 → … → k−1"]
    end
    subgraph storage ["Object storage"]
        IC[("IC source")]
        OUT[("Icechunk outputs")]
    end
    A --> B
    IC --> B
    B --> G
    G --> OUT
    OUT --> J

    style local fill:#555,stroke:#333,color:#fff
    style A fill:#777,stroke:#555,color:#fff
    style J fill:#777,stroke:#555,color:#fff
    style modal_cpu fill:#2c6fad,stroke:#1a4a7a,color:#fff
    style B fill:#4a90d9,stroke:#2c6fad,color:#fff
    style modal_gpu fill:#c44d1a,stroke:#8c3210,color:#fff
    style G fill:#ff6b35,stroke:#c44d1a,color:#fff
    style storage fill:#1a5c3a,stroke:#0d3d26,color:#fff
    style IC fill:#2e7d52,stroke:#1a5c3a,color:#fff
    style OUT fill:#2e7d52,stroke:#1a5c3a,color:#fff

Each member’s output is appended along the ensemble_member dimension. This mode is simpler and cheaper (one GPU) but slower for large ensembles.

Re-running existing forecasts and reproducibility#

The set up of aifs-modal is designed to be used in computational pipelines. Accordingly a key feature is that every step checks whether its output already exists before doing any work, so interrupted or repeated runs pick up where they left off without duplicating effort or GPU cost.

Idempotency operates at two levels. The local check (forecast_exists) runs in the notebook before app.run() — if it returns True, the ephemeral app is never created. The Modal-side check (run_forecast) runs inside the orchestrator — if the output already exists, ingestion and inference are skipped without consuming a GPU.

Step	Where	Skip condition	Override
`forecast_exists`	local	Output group exists and has enough members	—
`run_forecast` output check	Modal	Same as above — prevents IC ingestion and inference	`overwrite=True`
IC ingestion in `run_forecast`	Modal	Both IC dates already present on the IC Volume	—

Ensemble reproducibility. Each ensemble member fixes the PyTorch random seed to its member index (torch.manual_seed(member_id)) before running AIFS-ENS. This means that member k always produces the same trajectory regardless of how many members are in the ensemble or how they are distributed across containers, making it straightforward to extend an existing ensemble run by increasing n_members (existing members are skipped, only the new ones are computed).

How it works#

Key concepts#

Deterministic forecast (run_forecast)#

Ensemble forecast (run_forecast with n_members)#

Re-running existing forecasts and reproducibility#

Deterministic forecast (`run_forecast`)#

Ensemble forecast (`run_forecast` with `n_members`)#