Working with the model on Gemini Enterprise Agent Platform

WeatherNext 2 is Google's medium-range probabilistic forecasting model. Instead of consuming pre-computed, static forecast datasets, WeatherNext 2 on Gemini Enterprise Agent Platform lets you build custom, on-demand applications tuned to your technical requirements.

Note: Before you begin, your Google Cloud project must be added to the allowlist and meet the prerequisites described on the access page.

Find the model in Model Garden

To locate the model card of WeatherNext 2, use Model Garden in Gemini Enterprise Agent Platform:

1. Go to Model Garden.
2. Search for WeatherNext 2 inside the catalog.
3. Click the WeatherNext 2 card to open its details panel.
4. Select Open Notebook. This initializes a managed Colab Enterprise environment pre-configured to handle and run the model's inference notebooks.

Quota increase request

WeatherNext 2 runs distributed, multi-host inference on high-performance accelerators. By default, the initial quota for these top-tier GPUs is 0 in your Google Cloud project. Ensure that you have sufficient GPU quota allocated before you provision your infrastructure.

Quota requirements depend on whether you use ephemeral Serverless (Flex Start, Spot or Standard) workflows or dedicated Persistent Resources:

Quota for ephemeral/serverless jobs (Flex Start, Spot, Standard)

If you are submitting standalone custom jobs to the shared pool, request quota for preemptible resources:

• Service: Agent Platform API
  • Resource Name: Custom model training preemptible Nvidia A100 80GB GPUs per region OR Custom model training preemptible Nvidia H100 GPUs per region

Quota for dedicated clusters (Persistent Resources)

If you plan to deploy a long-running, dedicated cluster to host your weather modeling suite, you must request standard (non-preemptible) training quota:

• Service: Agent Platform API
• Resource Names:
  • Custom model training Nvidia A100 80GB GPUs per region
  • OR Custom model training Nvidia H100 GPUs per region

Make sure you have sufficient GPU quota allocated for the inference configuration (i.e. num_samples) before running Gemini Enterprise Agent Platform Jobs or provisioning a Persistent Resource. Otherwise, some Gemini Enterprise Agent Platform jobs may run while others will fail which would produce incomplete results. By default, the quota for GPUs is 0. You can request a higher quota by following the instructions at Request a higher quota.

You will need to request quota for either NVIDIA H100 80GB GPUs or NVIDIA A100 80GB GPUs in your selected region. The total number of GPUs you request must be sufficient for your largest planned forecast (i.e., num_samples) or your provisioned Persistent Resource capacity.

Depending on whether you use a Persistent Resource or run standard custom jobs, you should request the following quota under Service Enterprise Gemini Agent Platform API:

• Option 1: Persistent Resource
  • Name: Persistent resource Nvidia A100 80GB GPUs per region OR Persistent resource Nvidia H100 GPUs per region
• Option 2: Standard Custom Model Training (Preemptible)
  • Name: Custom model training preemptible Nvidia A100 80GB GPUs per region OR Custom model training preemptible Nvidia H100 GPUs per region"

Scheduling strategies: Flex Start versus Persistent Resources

Before running WeatherNext 2 workloads, choose how to provision your high-performance compute infrastructure based on your operational SLAs and budget:

Serverless and dynamic scaling

• FLEX_START: Instructs the training service to schedule your multi-host distributed jobs dynamically as capacity becomes available in the shared pool. This optimizes cost and hardware allocation across massive parallel workflows.
• SPOT: Uses preemptible Google Cloud resources. It offers a heavily discounted rate but can be reclaimed if system resources are required elsewhere.
• STANDARD: Provisions standard dedicated resources immediately based on availability, ensuring predictable start times for operational critical workflows.

Dedicated infrastructure by using Persistent Resources

Instead of running standard ephemeral custom jobs, you can deploy your forecasting suite onto a Persistent Resource. A persistent resource is a long-running, dedicated cluster created specifically to execute your serverless training and modeling workflows.

We recommend using persistent resources for WeatherNext 2 in the following scenarios:

• Guaranteed capacity availability: Reserves the required A100 or H100 multi-node hardware continuously to avoid capacity shortages during critical weather modeling cycles or peak seasons.
• Faster startup times: Reduces initialization times to near zero between subsequent ensemble simulations. By comparison, standard serverless jobs spend several minutes spinning up and provisioning hardware.
• Advanced autoscaling: Dynamically scales between a minimum and maximum replica count (min_replica_count and max_replica_count) to absorb parallel ensemble spikes and to scale down when idle.

If you create your persistent resource using the Google Cloud Console, on the Compute settings tab, make sure that "Use default Agent Platform Custom Code Service Agent" is checked in the Advanced options.

Checking the "Use default Agent Platform Custom Code Service Agent" checkbox in the Console UI is programmatically identical on Python SDK (google.cloud.aiplatform): Setting enable_custom_service_account=False.

Using a custom service account is not supported when running WeatherNext 2, so if you don't check it, the job may fail.

For more details on creating and managing persistent resources, see the access page.

Run the provided Colab notebooks

First, open the appropriate "Getting Started" notebook. It contains all the code needed to authenticate, configure, and run your first forecast without any local setup.

Select the notebook based on how you are configuring capacity:

• Open the Getting Started Notebook for Persistent Resources
  • Open the Getting Started Notebook for DWS (Flex Start)

Follow these steps inside your Colab Enterprise environment to submit a forecasting payload:

1. Set up project and storage: Specify your Google Cloud project ID, target execution region, and a unique Cloud Storage bucket path where the output forecast structures are written.
2. Authenticate: Run the authentication cells to authorize your runtime session or service account with the following roles:
   • Agent Platform Viewer or Agent Platform Administrator
   • Storage Object Admin.
3. Configure Parameters: Modify the interactive form elements, such as hardware configurations, ensemble dimensions, and weather target parameters.
4. Execute Code Cells: Execute the setup cells sequentially to validate inputs, calculate machine replicas (num_samples divided by the GPUs available per machine), and trigger the training container pipeline.
5. Monitor and Visualize: Monitor live streaming logs directly inside the notebook or in the Custom jobs dashboard. Once complete, run the visualization block to render generated prediction maps.

General inference parameters

To perform general inference, adjust the primary configuration settings inside the notebook form:

Hardware configuration for distributed inference

• machine_type: Select a valid machine type depending on your quota.
  • a3-highgpu series utilize NVIDIA H100 80GB GPUs.
  • a2-ultragpu series utilize NVIDIA A100 80GB GPUs.
• num_samples: The total number of ensemble members to generate. The number of machine replicas is calculated automatically (num_samples / GPUs per machine).
• scheduling_strategy: Choose between FLEX_START, SPOT, or STANDARD scheduling. Alternatively, you can target your configured persistent resource ID cluster.

Forecast configurations

• horizon_hrs: The chosen length of the forecast in hours, such as 240 for a 10-day forecast.
• model_seed: Select a specific model seed (1 to 4) or select all to run inference across all four seeds in parallel for optimal prediction ensemble accuracy.
• enable_hourly_prediction: A checkbox parameter that determines whether to generate granular 1-hour interval prediction points rather than the standard.

Inference with initial conditions (IC)

For specialized simulations where global atmospheric contexts must be modified, use the WeatherNext 2 IC notebook. Inference with Initial Conditions lets you replace the default input with your custom input or initial condition data.

Additional IC parameters

• input_data_gcs_dir: The source directory or Cloud Storage bucket URI containing the custom weather state file.

Note: WeatherNext 2 is principally trained on ECMWF data and is not fine-tuned on custom inputs. Using custom starting data from different weather models (e.g., GFS) means forecast quality is not guaranteed to match the default benchmark.

Retrieve and use forecast outputs

After your job successfully finishes processing, WeatherNext 2 automatically writes the output datasets back into your target Cloud Storage bucket.

Output file format: Zarr arrays

The model encodes predictions in the Zarr format (.zarr), an open-source data format optimized for cloud-native workflows that stores multi-dimensional, typed numerical arrays in compressed, chunked file segments.

Depending on your notebook runtime settings, you receive up to two distinct Zarr output payloads:

• Standard Output: Contains the standard 6-hourly temporal resolution interval forecast data.
• Granular Output: Contains 1-hourly interval steps. This file is generated only if you enable the enable_hourly_prediction flag prior to job submission.

Key output dimensions and spatial resolution

The output files structure data over global geographic, temporal, and probabilistic matrixes:

• Spatial Grid: 0.25-degree global resolution (approximately a 28 km by 28 km structural cell grid matching equatorial terrain mapping).
• Ensemble Dimension (num_samples): Organized by coordinate blocks indexing each distinct generated ensemble path (such as slices indexed from 0 to 63 if configuring a 64-member simulation) to capture structural forecast spread and system uncertainty.
• Temporal Horizon (horizon_hrs): Time-series matrixes mapping predictions across your configured forecast window.

Key output variables

The WeatherNext 2.0 (FGN) model provides a wide range of forecast variables. See the model specifications and data schema page for more information.

Troubleshooting

Execution logs are generated for every inference job. All job executions are available at the Custom Jobs page in the Google Cloud Console.

Clicking on a job name provides details on the parameters used and access to the execution logs. If an execution fails, provide your Google account team with the information from the job screen and the error logs found in Cloud Monitoring.