Run AI NWP forecasts hassle-free, serverless in the cloud!
ai-models
For AllThis package boot-straps on top of the fantastic ai-models
library to build a serverless application to generate "pure AI
NWP" weather forecasts on Modal. Users can run their own
historical re-forecasts using either PanguWeather,
FourCastNet, or GraphCast,
and save the outputs to their own cloud storage provider for further use.
The initial release of this application is fully-featured, with some limitations:
This application relies on the fantastic ecmwf-labs/ai-models
package to automate a lot of the traditional MLOps that are necessary to run this type
of product in a semi-production context. ai-models
handles acquiring data to use as
inputs for model inference (read: generate a forecast from initial conditions) by
providing an as-needed interface with the Copernicus Data Store and MARS API, it
provides pre-trained model weights, it implements a simple ONNX-based interface
for performing model inference, and it outputs a well-formed (albeit in GRIB) output
file that can be fed into downstream workflows (e.g. model visualization). We don't
anticipate replacing this package, but we may contribute improvements and features
upstream (e.g. a high priority is writing a NetCDF output adapter that writes timesliced
files per model step, with metadata following CF conventions) as they mature here.
Your feedback to [email protected] or @danrothenberg would be greatly appreciated!
If you use this application, please give credit to Daniel Rothenberg ([email protected] or @danrothenberg), as well as the incredible team at ECMWF Lab and the publishers of any forecast model you use.
NOTE THAT EACH FORECAST MODEL PROVIDED BY AI-MODELS HAS ITS OWN LICENSE AND RESTRICTIONS.
This package may only be used in a manner compliant with the licenses and terms of all the libraries, model weights, and application platforms/services upon which it is built. The forecasts generated by the AI models and the software which power them are experimental in nature and may break or fail unexpectedly during normal use.
Set up accounts (if you don't already have them) for:
Complete the .env
file with your CDS API credentials, GCS service account keys, and
a bucket name where model outputs will be uploaded. You should create this bucket
before running the application!
From a terminal, login with the modal-client
Navigate to the repository on-disk and execute the command,
$ modal run ai-models-modal.main \
--model-name {panguweather,fourcastnetv2-small,graphcast} \
--model-init 2023-07-01T00:00:00 \
--lead-time 12 \
[--use-gfs]
The first time you run this, it will take a few minutes to build an image and set up assets on Modal. Then, the model will run remotely on Modal infrastructure, and you can monitor its progress via the logs streamed to your terminal. The bracketed CLI args are the defaults that will be used if you don't provide any.
Download the model output from Google Cloud Storage at gs://{GCS_BUCKET_NAME} as
provided via the .env
file.
Install required dependencies onto your machine using the requirements.txt file (pip install -r requirements.txt)
We've implemented the ability for users to fetch initial conditions from an
archived GFS forecast cycle. In the current implementation, we make some assumptions
about how to process and map the GFS data to the ERA-5 data that the ai-models
package typically tries to fetch:
We use the gfs.tHHz.pgrb2.0p25.f000
output files to pull the initial conditions. These
are available in near-real-time (unlike the final GDAS analyses, which are lagged by
about one model cycle). We may provide the option to use the .anl
analysis files, too
or hot-start initial conditions, based on feedback from the community/users. Converting
to these files simply requires building a new set of mappers from the corresponding
ERA-5 fields.
To tell ai-models-for-all
to use GFS initial conditions, simply pass the command line
flag "--use-gfs
" and initialize a model run as usual.
$ modal run ai-models-modal.main \
--model-name {panguweather,fourcastnetv2-small,graphcast} \
--model-init 2023-07-01T00:00:00 \
--lead-time 12 \
--use-gfs \
The package will automatically download and process the GFS data to use for you, as well as archive it for future reference. Please note that the model setup process (where assets are downloaded and cached for future runs) may take much longer than usual as we also take the liberty of generating independent copies of the ERA-5 template files used to process the GFS data. Given the current quota restrictions on the CDS-API, this may take a very long time (luckily, the stub functions which perform this process are super cheap to run and will cost pennies even if they get stuck for several hours).
For your convenience, we've saved pre-computed data templates for you to use; for a
typical .env
setup described below, you can locally run the following Google Cloud
SDK command to copy over the input templates so that ai-models-for-all
will
automatically discover them:
$ gcloud storage cp \
gs://ai-models-for-all-input-templates/era5-to-gfs-f000 \
gs://${GCS_BUCKET_NAME}
To use this demo, you'll need accounts set up on Google Cloud, Modal, and the Copernicus Data Store. Don't worry - even though you do need to supply them ith credit card information, this demo should cost virtually nothing to run; we'll use very limited storage on Google Cloud Storage for forecast model outputs that we generate (a few cents per month if you become a prolific user), and Modal offers new individual users a startup package which includes $30/month of free compute - so you could run this application for about 8 hours straight before you'd incur any fees (A100's cost about $3.73 on Modal at the time of writing).
If you're very new to cloud computing, the following sections will help walk you through
the handful of steps necessary to get started with this demo. Experienced users can
quickly skim through to see how they need to modify the provided .env
to set up the
necessary credentials for the application to work.
The current version of this library ships with a single storage handler - a tool to upload to Google Cloud Storage. Without an external storage mechanism, there isn't a simple way to access the forecasts you generate using this tool. Thankfully, a basic cloud storage setup is easy to create, extremely cheap, and will serve you many purposes!
There are two main steps that you need to take here. We assume you already have an account on Google Cloud Platform (it should be trivial to setup from http://console.cloud.google.com).
Navigate to your project's Cloud Storage control panel. From here, you should see a listing of buckets that you own.
Find the button labeled Create near the top of the page. On the form that
loads, you should only need to provide a name for your bucket. Your name needs
to be globally unique - buckets across different projects must have different
names. We recommend the simple naming scheme <username>-ai-models-for-all
.
Keep all the default settings after inputting your bucket name and submit the form with the blue Create button at the bottom of the page.
Finally, navigate to the .env
file in this repo; set the GCS_BUCKET_NAME
variable to the bucket name you chose previously. You do not need quotes around
the bucket name.
We need a mechanism so that your app running on Modal can authenticate with Google Cloud in order to use its resources and APIs. To do this, we're going to create a service account and set the limited permissions needed to run this application.
From your Google Cloud Console, navigate to the IAM & Admin panel and then select Service Accounts from the menu. On the resulting page, click the button near the top that says Create Service Account.
On the form that pops up, use the following information:
The Service account ID field should automatically fill out for you. Click Create and Continue and you should move to Step 2, "Grant this service account access to project (optional)". Here, we will give permissions for the service account to access Cloud Storage resources (so that it can be used to upload and list objects in the bucket you previously created).
From the drop-down labeled "Select a role", search for "Storage Object Admin" (you may want to use the filter). Add this role then click Continue. You shouldn't need to grant any specific user permissions, assuming you're the owner of the Google Cloud account in which you're setting this up. Click Done.
Finally, we need to access the credentials for this new account. Navigate back to the page IAM & Admin > Service Accounts, and click the name in the table with the "model-ai-models-for-all" service account you just created. At the top of the page, navigate to the Keys tab and click the ADD KEY button on the middle of hte page. This will generate and download a new private key that you'll use. Select the "JSON" option from the pop-up and download the file by clicking Create.
The credentials you created will be downloaded to disk automatically. Open that
JSON file in your favorite text editor; you'll see a mapping of many different
keys to values, including "private_key". We need to pass this entire JSON object
whenever we want to authenticate with Google Cloud Storage. To do this, we'll
save it in the same .env
file under the name GCS_SERVICE_ACCOUNT_INFO.
Unfortunately we can't just copy/paste - we need to "stringify" the data. You should probably do this in Python or your preferred programming language by reading in the JSON file you saved, serializing to a string, and outputting.
cdsapi
We need access to the Copernicus Data Store
to retrieve historical ERA-5 data to use when initializing our forecasts. The
easiest way to set this up would be to have the user retrieve their credentials
from here and save them to a
local file, ~/.cdsapirc
. But that's a tad inconvenient to build into our
application image. Instead, we can just set the environment variables
CDSAPI_URL and CDSAPI_KEY. Note that we still create a stub RC file during
image generation, but this is a shortcut so that users only need to modify a single
file with their credentials.
The following major projects are slated for Q1'24:
Operational AI Model Forecasts - We will begin running pseudo-operational forecasts for all included models in early Q1 using GFS initial conditions in near-real-time, and disseminating the outputs in a publicly available Google Cloud Storage bucket (per model licensing restrictions).
Post-processing / Visualization - We will implement some simple (optional) routines to post-process the default GRIB outputs into more standard ARCO formats, and generate a complete set of visualizations that users can review as a stand-alone gallery. Pending collaborations, we will try to make these available on popular model visualization websites (contact @darothen if you're interested in hosting).
Porting to earth2mip
- Although we've used ecmwf-labs/ai-models for the initial
development, this package's extremely tight coupling with ECMWF infrastructure and
the climetlab library pose considerable development challenges. Therefore, we aim
to re-write this library using the NVIDIA/earth2mip framework. This is a far more
comprehensive and extensible framework for preparing a variety of modeling and learning
tasks related to AI-NWP, and provides access to the very large library of AI models
NVIDIA is collecting for their model zoo. This will likely be built in a stand-alone
package, e.g. earth2mip-for-all
, but the goal is to provide the same accessibility
and ease-of-use for users who simply want to run these models to create their own
forecasts with limited engineering/infrastructure investment.