Quickstart Guide for NVIDIA Earth-2 FourCastNet NIM#

Use this documentation to get started with NVIDIA Earth-2 FourCastNet NIM.

Important

Before you can use this documentation, you must satisfy all prerequisites.

Launching the NIM#

  1. Pull the NIM container with the following command.

    Note

    The container is ~34GB (uncompressed) and the download time depends on internet connection speeds.

    docker pull nvcr.io/nim/nvidia/fourcastnet:1.0.0

  2. Run the NIM container with the following command. This command starts the NIM container and exposes port 8000 for the user to interact with the NIM. It pulls the model on the local filesystem.

    Note

    The model is ~6GB and the download time depends on internet connection speeds.

    export NGC_API_KEY=<NGC API Key>

docker run --rm --runtime=nvidia --gpus all --shm-size 4g \
    -p 8000:8000 \
    -e NGC_API_KEY \
    -t nvcr.io/nim/nvidia/fourcastnet:1.0.0

    After the NIM is running, verify that you see output similar to the following:

    I0829 05:18:42.152983 108 grpc_server.cc:2466] "Started GRPCInferenceService at 0.0.0.0:8001"
I0829 05:18:42.153277 108 http_server.cc:4638] "Started HTTPService at 0.0.0.0:8090"
I0829 05:18:42.222628 108 http_server.cc:320] "Started Metrics Service at 0.0.0.0:8002"

Checking NIM Health#

Bash#

  1. Open a new terminal, leaving the current terminal open with the launched service.

  2. In the new terminal, wait until the health check end point returns {"status":"ready"} before proceeding. This might take a couple of minutes. Use the following command to query the health check.

    curl -X 'GET' \
    'http://localhost:8000/v1/health/ready' \
    -H 'accept: application/json'

Python#

import requests

r = requests.get("http://localhost:8000/v1/health/ready")
if r.status_code == 200:
   print("NIM is healthy!")
else:
   print("NIM is not ready!")

Fetching Input Data#

The default profile used in the NIM is Spherical Harmonic Neural Operator, which uses 73 channels of ERA5 as the input. The specific input/output channels of this model can be found on it’s model card. In this guide, Earth-2 Studio is used to fetch the input data needed for the model. The following Python code fetches an initial state for the FourCastNet NIM from ARCO ERA5 and saves it into a NumPy file:

import numpy as np
from datetime import datetime
from earth2studio.data import ARCO
from earth2studio.models.px.sfno import VARIABLES

ds = ARCO()
da = ds(time=datetime(2023, 1, 1), variable=VARIABLES)
np.save("fcn_inputs.npy", da.to_numpy()[None].astype('float32'))

This will create a NumPy array with size (1,1,73,721,1440), which corresponds to the required input fields specified in the model card. More broadly, these dimensions map to (batch size, lead time, variables, latitudes, and longitudes) as discussed in the API documentation.

Tip

See the ARCO datasource API documentation for additional details on this particular data source.

Inference Requests#

Performing inference with the FourCastNet NIM can standard synchronous post request. For complete documentation on the API specification of the NIM, visit the API documentation page.

Bash#

curl -X POST \
        -F "input_array=@fcn_inputs.npy" \
        -F "input_time=2023-01-01T00:00:00Z" \
        -F "simulation_length=4" \
        -o output.tar \
        http://localhost:8000/v1/infer

Python#

import requests

url = "http://localhost:8000/v1/infer"
files = {
    "input_array": ("input_array", open("fcn_inputs.npy", "rb")),
}
data = {
    "input_time": "2023-01-01T00:00:00Z",
    "simulation_length": 4,
}
headers = {
    "accept": "application/x-tar",
}
print("Sending post request to NIM")

r = requests.post(url, headers=headers, data=data, files=files, timeout=180)

if r.status_code != 200:
    raise Exception(r.content)
else:
    # Dump response to file
    with open("output.tar", "wb") as tar:
        tar.write(r.content)

Results#

The results are inside a tar file that can be explored using:

tar -tvf output.tar

-rw-r--r-- 0/0       303166208 1970-01-01 00:00 000_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 006_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 012_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 018_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 024_000.npy

Tip

The tar archive is populated with NumPy arrays that have the naming convention {lead time}_{batch index}.npy.

Warning

For most inference calls, it is important to specify a longer timeout period for a request to this NIM.

Streaming Responses#

The FourCastNet NIM generates a time-series and streams each time-step back to the client as its generated. This can be extremely useful when creating a pipeline that runs following processes on each time-step. The follow snippet can be used to access each time-step, in memory, from the NIM as it is generated.

import io
import tarfile
from pathlib import Path

import numpy as np
import requests
import tqdm

url = "http://localhost:8000/v1/infer"
files = {
    "input_array": ("input_array", open("fcn_inputs.npy", "rb")),
}
data = {
    "input_time": "2023-01-01T00:00:00Z",
    "simulation_length": 4,
}
headers = {
    "accept": "application/x-tar",
}

pbar = tqdm.tqdm(range(data["simulation_length"] + 1), desc="Forecast steps")
with (
    requests.post(
        url,
        headers=headers,
        files=files,
        data=data,
        timeout=300,
        stream=True,
    )
) as resp:
    resp.raise_for_status()
    with tarfile.open(fileobj=resp.raw, mode="r|") as tar_stream:
        # Loop over file members from tar stream
        for member in tar_stream:
            arr_file = io.BytesIO()
            arr_file.write(tar_stream.extractfile(member).read())
            arr_file.seek(0)
            data = np.load(arr_file)
            # Array names are on the form <LEAD_TIME>_<BATCH_IDX>.npy
            arr_lead_time, arr_batch_idx = (
                int(x) for x in Path(member.name).stem.split("_", maxsplit=1)
            )

            pbar.write(f"Received data for lead {arr_lead_time} batch {arr_batch_idx}")
            pbar.write(f"Output numpy {member.name} with shape {data.shape}")
            pbar.update(1)
pbar.close()

Post Processing Results#

The last step is to visualize the results produced from the NIM. While there are a number of different tools you could use to create different visuals, this guide demonstrates using Matplotlib.

import tarfile
import io

import numpy as np
import matplotlib.pyplot as plt

fig, ax = plt.subplots(2, 2)
ax = ax.flatten()

with tarfile.open("output.tar") as tar:
    for i, member in enumerate(tar.getmembers()):
        arr_file = io.BytesIO()
        arr_file.write(tar.extractfile(member).read())
        arr_file.seek(0)
        data = np.load(arr_file)

        ax[i].set_title(f"Lead time: {6*i}hrs")
        ax[i].imshow(data[0,0,0], vmin=-10, vmax=10, cmap="coolwarm")
        if i == 3:
            break

plt.savefig("output.png")

The result is a plot of the predicted wind speed fields at 10 meters.

fcn output

On this page