From 77b2af09cec2dc14e98506e068a96560bf79bdf2 Mon Sep 17 00:00:00 2001 From: Damodar Panigrahi Date: Fri, 7 Aug 2026 10:27:35 -0700 Subject: [PATCH] feat: WN2 with GPU GA (#4563) --- notebooks/community/CODEOWNERS | 2 + notebooks/community/weathernext/README.md | 5 +- .../weathernext/weathernext_2_dws.ipynb | 577 +++++++++++ .../weathernext/weathernext_2_ic_pc.ipynb | 970 ++++++++++++++++++ 4 files changed, 1551 insertions(+), 3 deletions(-) create mode 100644 notebooks/community/weathernext/weathernext_2_dws.ipynb create mode 100644 notebooks/community/weathernext/weathernext_2_ic_pc.ipynb diff --git a/notebooks/community/CODEOWNERS b/notebooks/community/CODEOWNERS index 5846e51f9..eabdf619f 100644 --- a/notebooks/community/CODEOWNERS +++ b/notebooks/community/CODEOWNERS @@ -30,6 +30,8 @@ /notebooks/community/alphagenome/cloudai_alphagenome_finetune.ipynb @dpanigra /notebooks/community/weathernext/weathernext_2_early_access_program.ipynb @dpanigra /notebooks/community/weathernext/weathernext_2_ic_early_access_program.ipynb @dpanigra +/notebooks/community/weathernext/weathernext_2_dws.ipynb @dpanigra +/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb @dpanigra /notebooks/community/ml_ops/stage2/get_started_with_visionapi_and_automl.ipynb @mansari /notebooks/community/neo4j/graph_paysim.ipynb @benofben @laeg /notebooks/community/ml_ops/stage1/get_started_with_visionapi_and_vertex_datasets.ipynb @mansari diff --git a/notebooks/community/weathernext/README.md b/notebooks/community/weathernext/README.md index f18423670..e11d9b5dc 100644 --- a/notebooks/community/weathernext/README.md +++ b/notebooks/community/weathernext/README.md @@ -4,10 +4,9 @@ ## Overview **Disclaimer:** -*Experimental*\ -This product is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the [Service Specific Terms](https://cloud.google.com/terms/service-terms#1). Pre-GA products are available "as is" and might have limited support. For more information, see the [launch stage descriptions](https://cloud.google.com/products#product-launch-stages). +This product is subject to the General Service Terms section of the [Service Specific Terms](https://cloud.google.com/terms/service-terms#1). GA products follows standard support. For more information, see the [launch stage descriptions](https://cloud.google.com/products#product-launch-stages). -Access to the forecasting capabilities requires application and approval. Users must be added to an allowlist to generate forecasts using this service. Review pricing details at [Vertex AI Custom Training pricing,](https://cloud.google.com/vertex-ai/pricing?hl=en&e=48754805#custom-trained-models) [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) before running. +Please contact your Google Cloud sales person for using WeatherNext 2 on Agent Platform. In case you do not have a salesperson you work with, please submit the form using the Request access button in this model card and a Google Cloud representative will contact you. **Overview** diff --git a/notebooks/community/weathernext/weathernext_2_dws.ipynb b/notebooks/community/weathernext/weathernext_2_dws.ipynb new file mode 100644 index 000000000..0a16f4a83 --- /dev/null +++ b/notebooks/community/weathernext/weathernext_2_dws.ipynb @@ -0,0 +1,577 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "RirpY96_3zL0" + }, + "outputs": [], + "source": [ + "# Copyright 2026 Google LLC\n", + "\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "52d0938f" + }, + "source": [ + "# WeatherNext 2 (Using DWS)\n", + "\n", + " \n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Google
Open in Colab\n", + " \n", + " 		\n", + " \n", + " \"Google
Open in Colab Enterprise\n", + " \n", + " 		\n", + " \n", + " \"Vertex
Open in Vertex AI Workbench\n", + " \n", + " 		\n", + " \n", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "DGFzy9MTuYL_" + }, + "source": [ + "## Overview\n", + "\n", + "This notebook demonstrates running [WeatherNext 2 inference on Google Cloud Vertex AI](https://developers.google.com/weathernext/guides/access-vmg). WeatherNext 2 is Google's latest medium-range probabilistic forecasting model, principally an operational version the FGN model ([published June 2025](https://arxiv.org/abs/2506.10772)). More information is available in the [WeatherNext documentation](https://developers.google.com/weathernext).\n", + "\n", + "### Objective\n", + "\n", + "- Configure the model inputs for distributed, multi-host inference on H100 or A100 GPUs.\n", + "- Run WeatherNext 2 model forecasts in parallel.\n", + "- Visualize forecast results.\n", + "\n", + "### Costs\n", + "\n", + "This uses billable components of Google Cloud:\n", + "\n", + "* Vertex AI\n", + "* Cloud Storage\n", + "\n", + "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n", + "\n", + "\n", + "## Before you begin\n", + "\n", + "### Request For GPU Quota\n", + "\n", + "**WARNING:** Make sure you have sufficient GPU quota allocated for the inference configuration (i.e. `num_samples`) before running Vertex Jobs. Otherwise, some Vertex jobs may run while others will fail which would produce\n", + "incomplete results.\n", + "\n", + "\n", + "By default, the quota for GPUs is 0. You can request a higher quota by following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n", + "\n", + "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`).\n", + "\n", + "You should request for the following quota:\n", + "\n", + "- Service: `Vertex AI API`\n", + "- Name: `Custom model training preemptible Nvidia A100 80GB GPUs per region` OR `Custom model training preemptible Nvidia H100 GPUs per region`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "yt4GxkKcDD7Y" + }, + "outputs": [], + "source": [ + "# @title Install python packages\n", + "\n", + "# Note that you may need to restart the kernel after this step.\n", + "# If so, continue to the next cell after restarting.\n", + "\n", + "print(\"Installing python packages.\")\n", + "\n", + "! pip3 install \\\n", + " google-cloud-aiplatform==1.129.0 \\\n", + " xarray[complete]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dO2mF4CPfpHW1ZKbIsut8r69" + }, + "outputs": [], + "source": [ + "# @title Setup Google Cloud project\n", + "\n", + "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "\n", + "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched.\n", + "\n", + "\n", + "BUCKET_URI = \"gs://my-bucket\" # @param {type:\"string\"}\n", + "\n", + "# @markdown 3. Select a region that has the required GPUs available.\n", + "\n", + "REGION = \"us-central1\" # @param {type:\"string\"}\n", + "\n", + "# Import the necessary packages\n", + "import datetime\n", + "import os\n", + "import re\n", + "\n", + "from google.cloud import aiplatform\n", + "\n", + "# Get the default cloud project id.\n", + "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n", + "\n", + "# Enable the Vertex AI API and Compute Engine API, if not already.\n", + "print(\"Enabling Vertex AI API and Compute Engine API.\")\n", + "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n", + "\n", + "# Cloud Storage bucket for storing the experiment artifacts.\n", + "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n", + "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n", + "\n", + "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n", + " raise ValueError(\"GCS Bucket URI is invalid!\")\n", + "else:\n", + " assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n", + " shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n", + " bucket_region = shell_output[0].strip().lower()\n", + " if bucket_region != REGION:\n", + " raise ValueError(\n", + " f\"Bucket region {bucket_region} is different from notebook region {REGION}\"\n", + " )\n", + "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n", + "\n", + "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n", + "# Initialize Vertex AI API.\n", + "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n", + "\n", + "# Utility functions\n", + "\n", + "\n", + "def get_job_name_with_datetime(prefix: str) -> str:\n", + " return prefix + datetime.datetime.now().strftime(\"_%Y%m%d_%H%M%S\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "22BW43yjps8D" + }, + "outputs": [], + "source": [ + "# @title Configure Model Parameters\n", + "# @markdown Configure the hardware and input parameters for the WeatherNext 2 forecast.\n", + "\n", + "\n", + "# @markdown ### Hardware Configuration for Distributed Inference\n", + "# @markdown - **`machine_type`**: Select a valid machine type. `a3-highgpu` series use NVIDIA H100 80GB GPUs. `a2-ultragpu` series use NVIDIA A100 80GB GPUs.\n", + "# @markdown - **`num_samples`**: The total number of ensemble members to generate.\n", + "# @markdown The number of machine replicas will be calculated automatically (`num_samples` / GPUs per machine). **Therefore, `num_samples` must be a multiple of the number of GPUs in your selected `machine_type`.**\n", + "# @markdown - **`scheduling_strategy`**: The [strategy](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/CustomJobSpec#Strategy) used to acquire machines for the job. Defaults to [Dynamic Workload Scheduler](https://docs.cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws) (FLEX_START).\n", + "machine_type = \"a3-highgpu-1g\" # @param [\"a3-highgpu-1g\", \"a3-highgpu-2g\", \"a3-highgpu-4g\", \"a3-highgpu-8g\", \"a2-ultragpu-1g\", \"a2-ultragpu-2g\", \"a2-ultragpu-4g\", \"a2-ultragpu-8g\"]\n", + "num_samples = 8 # @param {type:\"integer\"}\n", + "scheduling_strategy = \"FLEX_START\" # @param [\"FLEX_START\", \"SPOT\", \"STANDARD\"]\n", + "\n", + "# @markdown ### Forecast Configuration\n", + "# @markdown - **`forecast_init_time`**: The starting time for the forecast in ISO 8601 format (e.g., `2025-09-21T00:00:00Z`). Models are available for dates from 2024 onwards.\n", + "# @markdown - **`horizon_hrs`**: The desired length of the forecast in hours (e.g., 240 for a 10-day forecast).\n", + "# @markdown - **`model_seed`**: Choose a specific model seed (1-4) or select \"all\" to run inference with all four seeds in parallel for improved accuracy.\n", + "# @markdown - **`enable_hourly_prediction`**: If checked, the model will generate 1-hour predictions.\n", + "forecast_init_time = \"2025-11-20T00:00:00Z\" # @param {type:\"string\"}\n", + "horizon_hrs = 72 # @param {type:\"integer\"}\n", + "model_seed = \"all\" # @param [\"1\", \"2\", \"3\", \"4\", \"all\"]\n", + "enable_hourly_prediction = True # @param {type:\"boolean\"}\n", + "\n", + "# --- Parameter Validation and Configuration ---\n", + "\n", + "# Derive accelerator type and count from the chosen machine type\n", + "if machine_type.startswith(\"a3-highgpu\"):\n", + " accelerator_type = \"NVIDIA_H100_80GB\"\n", + "elif machine_type.startswith(\"a2-ultragpu\"):\n", + " accelerator_type = \"NVIDIA_A100_80GB\"\n", + "else:\n", + " raise ValueError(f\"Invalid machine type selected: {machine_type}.\")\n", + "\n", + "try:\n", + " # Extract the number of GPUs from the machine type string, e.g., 'a3-highgpu-4g' -> 4\n", + " accelerators_per_machine = int(re.search(r\"-(\\d+)g$\", machine_type).group(1))\n", + "except (AttributeError, ValueError):\n", + " raise ValueError(\n", + " f\"Could not determine accelerator count from machine type: {machine_type}\"\n", + " )\n", + "\n", + "seeds_to_run = [1, 2, 3, 4] if model_seed == \"all\" else [int(model_seed)]\n", + "num_seeds_to_run = len(seeds_to_run)\n", + "\n", + "num_samples_per_seed = num_samples\n", + "if len(seeds_to_run) > 1:\n", + " if num_samples % num_seeds_to_run != 0:\n", + " raise ValueError(\n", + " f\"`num_samples` ({num_samples}) is not divisible by the number of seeds to run ({num_seeds_to_run}.\"\n", + " )\n", + " num_samples_per_seed = num_samples // num_seeds_to_run\n", + "\n", + "# Validate that num_samples is a multiple of accelerators_per_machine\n", + "if num_samples_per_seed % accelerators_per_machine != 0:\n", + " raise ValueError(\n", + " f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n", + " )\n", + "\n", + "# Calculate the number of replicas per seed\n", + "replica_count_per_seed = num_samples_per_seed // accelerators_per_machine\n", + "\n", + "# Ensure that there are enough samples\n", + "if num_samples_per_seed % accelerators_per_machine != 0:\n", + " raise ValueError(\n", + " f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n", + " )\n", + "\n", + "# Calculate total GPUs needed for all jobs\n", + "total_gpus_needed = num_samples * (4 if model_seed == \"all\" else 1)\n", + "\n", + "# Set Docker URI\n", + "WEATHERNEXT2_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/weather-next-2-inference.gpu.0-1:latest\"\n", + "\n", + "print(\"--- Job Configuration Summary ---\")\n", + "print(f\"Total Samples: {num_samples}\")\n", + "print(f\"Machine Type: {machine_type}\")\n", + "print(f\"Accelerator Type: {accelerator_type}\")\n", + "print(f\"GPUs per Machine: {accelerators_per_machine}\")\n", + "print(f\"Total number seeds to run: {num_seeds_to_run}\")\n", + "print(f\"Total number samples per seed: {num_samples_per_seed}\")\n", + "print(f\"Calculated Machine Replicas Per Seed: {replica_count_per_seed}\")\n", + "print(f\"Total GPUs per Job: {num_samples}\")\n", + "print(f\"Total GPUs across all Jobs (ensure sufficient quota): {total_gpus_needed}\")\n", + "print(f\"Docker Image: {WEATHERNEXT2_DOCKER_URI}\")\n", + "print(\"---------------------------------\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "VSuGx2jmpwP0" + }, + "outputs": [], + "source": [ + "# @title Run Forecasts\n", + "# @markdown This section creates and runs one or more Vertex AI Custom Training Jobs to generate the forecasts.\n", + "# @markdown **This operation is asynchronous.** The jobs will be submitted and this cell will complete quickly.\n", + "# @markdown You must monitor the job progress in the Google Cloud Console (https://console.cloud.google.com/vertex-ai/training/custom-jobs).\n", + "\n", + "from google.cloud.aiplatform.compat.types import \\\n", + " custom_job as gca_custom_job_compat\n", + "\n", + "print(f\"Submitting {len(seeds_to_run)} job(s) to run in parallel.\")\n", + "\n", + "launched_jobs = []\n", + "output_dirs = {}\n", + "\n", + "if scheduling_strategy == \"FLEX_START\":\n", + " SCHEDULLING_STRATEGY = gca_custom_job_compat.Scheduling.Strategy.FLEX_START\n", + "elif scheduling_strategy == \"SPOT\":\n", + " SCHEDULLING_STRATEGY = gca_custom_job_compat.Scheduling.Strategy.SPOT\n", + "else:\n", + " SCHEDULLING_STRATEGY = gca_custom_job_compat.Scheduling.Strategy.STANDARD\n", + "\n", + "for seed in seeds_to_run:\n", + " output_dir = os.path.join(BUCKET_URI, \"weathernext2_outputs\")\n", + " output_dirs[seed] = output_dir\n", + "\n", + " docker_args_list = [\n", + " f\"--pred_root_dir={output_dir}\",\n", + " f\"--num_samples={num_samples_per_seed}\",\n", + " f\"--horizon_hrs={horizon_hrs}\",\n", + " f\"--forecast_init_time={forecast_init_time}\",\n", + " f\"--model_seed={seed}\",\n", + " f\"--enable_hourly_prediction={enable_hourly_prediction}\",\n", + " ]\n", + "\n", + " JOB_NAME = get_job_name_with_datetime(\n", + " prefix=f\"wn2-forecast-s{seed}-n{num_samples_per_seed}\"\n", + " )\n", + " print(f\"\\n--- Submitting Job for Seed {seed} ---\")\n", + " print(f\"JOB_NAME: {JOB_NAME}\")\n", + "\n", + " job = aiplatform.CustomContainerTrainingJob(\n", + " display_name=JOB_NAME,\n", + " container_uri=WEATHERNEXT2_DOCKER_URI,\n", + " )\n", + "\n", + " job.run(\n", + " args=docker_args_list,\n", + " replica_count=replica_count_per_seed,\n", + " machine_type=machine_type,\n", + " accelerator_type=accelerator_type,\n", + " accelerator_count=accelerators_per_machine,\n", + " scheduling_strategy=SCHEDULLING_STRATEGY,\n", + " # Change this to True if you need to debug why the job hasn't started\n", + " sync=False,\n", + " )\n", + " launched_jobs.append(job)\n", + " print(\n", + " \"--> Job submitted successfully. Monitor it in the Google Cloud Console at https://console.cloud.google.com/vertex-ai/training/custom-jobs\"\n", + " )\n", + "\n", + "print(\"\\nAll forecast jobs have been submitted.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "x6ZVqWRopWcI" + }, + "outputs": [], + "source": [ + "# @title Visualize Forecasts (Unified)\n", + "# @markdown Select which forecast output you want to visualize. This single component\n", + "# @markdown can handle both the standard 6-hourly predictions and the datasets\n", + "# @markdown with 1-hour model (which have a 'subtime' dimension).\n", + "# @markdown If you run into `Error loading Zarr store: unrecognized engine 'zarr'...` try restarting the runtime session and reruning this cell.\n", + "\n", + "# @markdown ---\n", + "# @markdown ### Visualization Settings\n", + "# @markdown - **`model_seed_to_visualize`**: Choose a specific model seed (1-4) to visualize. This should be one of the model seeds selected in the **Forecast Configuration** above.\n", + "# @markdown - **`time_steps_to_visualize`**: Choose to visualize 1-hourly or 6-hourly forecasts. If 1-hourly is selected, ensure `enable_hourly_prediction` was selected in the **Forecast Configuration** above.\n", + "# @markdown - **`variable_to_visualize`**: Choose the weather variable to visualize. See the [WeatherNext documentation](https://developers.google.com/weathernext/guides/model-specs-vmg) for variable names and descriptions.\n", + "# @markdown - **`sample_to_visualize`**: Choose the sample (ensemble member) to visualize.\n", + "# @markdown - **`plot_size`**: Choose the size of the plot to generate.\n", + "model_seed_to_visualize = \"4\" # @param [\"1\", \"2\", \"3\", \"4\"]\n", + "time_steps_to_visualize = \"6-Hourly\" # @param [\"6-Hourly\", \"1-Hourly\"]\n", + "variable_to_visualize = \"2m_temperature\" # @param {type:\"string\"}\n", + "sample_to_visualize = 0 # @param {type:\"integer\"}\n", + "plot_size = 8 # @param {type:\"number\"}\n", + "level_to_visualize = None\n", + "# @markdown ---\n", + "\n", + "\n", + "# --- Helper Functions ---\n", + "\n", + "\n", + "def init_time_to_folder_path(init_time: str) -> str:\n", + " \"\"\"\n", + " Convert init time to expected GCS folder path.\n", + " \"\"\"\n", + " init_date, init_time = init_time.split(\"T\")\n", + " return f\"{init_date.replace('-', '')}_{init_time[0:2]}hr\"\n", + "\n", + "\n", + "# override these if you'd like to visualize a different set of forecasts\n", + "visualize_bucket = BUCKET_URI\n", + "visualize_init_date = forecast_init_time\n", + "\n", + "# set paths based on chosen model seed, bucket, and init date\n", + "path_to_6hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n", + "path_to_1hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}_hourly/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n", + "\n", + "\n", + "import datetime\n", + "from typing import Optional\n", + "\n", + "import matplotlib\n", + "import matplotlib.animation as animation\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import xarray\n", + "from IPython.display import HTML\n", + "\n", + "matplotlib.rcParams[\"animation.embed_limit\"] = 500\n", + "\n", + "\n", + "# --- Helper Functions ---\n", + "\n", + "\n", + "def select_data(\n", + " data: xarray.Dataset,\n", + " variable: str,\n", + " level: Optional[int] = None,\n", + ") -> xarray.Dataset:\n", + " \"\"\"Selects a variable from the dataset and optionally a level.\"\"\"\n", + " data = data[variable]\n", + " if \"batch\" in data.dims:\n", + " data = data.isel(batch=0)\n", + " if level is not None and \"level\" in data.coords:\n", + " data = data.sel(level=level)\n", + " return data\n", + "\n", + "\n", + "def scale_data(\n", + " data: xarray.Dataset,\n", + " center: Optional[float] = None,\n", + " robust: bool = False,\n", + ") -> tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n", + " \"\"\"Scales the data for visualization.\"\"\"\n", + " vmin = np.nanpercentile(data.values, (2 if robust else 0))\n", + " vmax = np.nanpercentile(data.values, (98 if robust else 100))\n", + " if center is not None:\n", + " diff = max(vmax - center, center - vmin)\n", + " vmin = center - diff\n", + " vmax = center + diff\n", + " return (\n", + " data,\n", + " matplotlib.colors.Normalize(vmin, vmax),\n", + " (\"RdBu_r\" if center is not None else \"viridis\"),\n", + " )\n", + "\n", + "\n", + "def create_forecast_animation(\n", + " dataset: xarray.Dataset,\n", + " fig_title: str,\n", + " plot_size: float = 5,\n", + " robust: bool = False,\n", + ") -> HTML:\n", + " \"\"\"\n", + " Creates a forecast animation from an xarray Dataset.\n", + " It intelligently handles datasets with or without a 'subtime' dimension.\n", + " \"\"\"\n", + " # --- Data Preparation ---\n", + " # Check if the data still has 'subtime'). If so, stack dimensions.\n", + " # Otherwise, just rename the 'time' dimension for consistency.\n", + " if \"subtime\" in dataset.dims:\n", + " print(\"Detected 'subtime' dimension. Stacking for hourly animation.\")\n", + " # Stack 'time' and 'subtime' into a single animation dimension\n", + " plot_data = dataset.stack(animation_step=(\"time\", \"subtime\")).transpose(\n", + " \"animation_step\", \"lat\", \"lon\"\n", + " )\n", + " else:\n", + " print(\"No 'subtime' dimension found. Using 'time' for 6-hourly animation.\")\n", + " # Use 'time' as the animation dimension\n", + " plot_data = dataset.rename({\"time\": \"animation_step\"})\n", + "\n", + " # Now, the animation dimension is always called 'animation_step'\n", + " max_steps = plot_data.sizes[\"animation_step\"]\n", + " init_time = plot_data.coords[\"init_time\"].values\n", + "\n", + " # Scale the data for color mapping\n", + " scaled_data, norm, cmap = scale_data(plot_data, robust=robust)\n", + "\n", + " # --- Plotting Setup ---\n", + " figure = plt.figure(figsize=(plot_size * 2, plot_size))\n", + " ax = figure.add_subplot(1, 1, 1)\n", + " ax.set_xticks([])\n", + " ax.set_yticks([])\n", + " figure.suptitle(fig_title, fontsize=16)\n", + " figure.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust for title\n", + "\n", + " im = ax.imshow(\n", + " scaled_data.isel(animation_step=0), norm=norm, origin=\"lower\", cmap=cmap\n", + " )\n", + "\n", + " plt.colorbar(\n", + " mappable=im,\n", + " ax=ax,\n", + " orientation=\"vertical\",\n", + " pad=0.02,\n", + " aspect=16,\n", + " shrink=0.75,\n", + " cmap=cmap,\n", + " extend=(\"both\" if robust else \"neither\"),\n", + " )\n", + "\n", + " # --- Animation Update Function ---\n", + " def update(frame):\n", + " # Get the coordinates for the current frame\n", + " step_coords = plot_data[\"animation_step\"][frame].coords\n", + "\n", + " # Calculate total offset and valid time based on available coordinates\n", + " if \"subtime\" in step_coords: # Hourly data\n", + " total_offset = step_coords[\"time\"].values + step_coords[\"subtime\"].values\n", + " else: # 6-hourly data\n", + " total_offset = step_coords[\"animation_step\"].values\n", + "\n", + " total_hours = total_offset / np.timedelta64(1, \"h\")\n", + " valid_time = init_time + total_offset\n", + " valid_time_str = np.datetime_as_string(valid_time, unit=\"m\").replace(\"T\", \" \")\n", + "\n", + " new_title = (\n", + " f\"{fig_title}\\n\"\n", + " f\"Valid: {valid_time_str} UTC (Forecast: +{total_hours:.1f}h)\"\n", + " )\n", + " figure.suptitle(new_title, fontsize=16)\n", + " im.set_data(scaled_data.isel(animation_step=frame))\n", + "\n", + " # --- Create and Display Animation ---\n", + " ani = animation.FuncAnimation(\n", + " fig=figure, func=update, frames=max_steps, interval=250\n", + " )\n", + " plt.close(figure.number)\n", + " return HTML(ani.to_html5_video())\n", + "\n", + "\n", + "# --- Main Visualization Logic ---\n", + "\n", + "# 1. Select the correct path based on the user's dropdown choice\n", + "if time_steps_to_visualize == \"6-Hourly\":\n", + " path_to_zarr = path_to_6hr_zarr\n", + "elif time_steps_to_visualize == \"1-Hourly\":\n", + " path_to_zarr = path_to_1hr_zarr\n", + "else:\n", + " raise ValueError(\"Invalid visualization target selected.\")\n", + "\n", + "print(f\"Loading data from: {path_to_zarr}\")\n", + "\n", + "# 2. Load the dataset\n", + "try:\n", + " full_dataset = xarray.open_zarr(path_to_zarr)\n", + "except Exception as e:\n", + " print(f\"Error loading Zarr store: {e}\")\n", + " # This is a common point of failure, so we exit gracefully.\n", + "else:\n", + " # 3. Select the specific data slice for visualization\n", + " data_for_vis = full_dataset.isel(sample=sample_to_visualize)\n", + " variable_data = select_data(data_for_vis, variable_to_visualize, level_to_visualize)\n", + "\n", + " # 4. Generate the title\n", + " title = f\"{variable_to_visualize} (Sample {sample_to_visualize})\"\n", + " if level_to_visualize:\n", + " title += f\" at {level_to_visualize} hPa\"\n", + "\n", + " # 5. Create and display the animation\n", + " display(create_forecast_animation(variable_data, title, plot_size, robust=True))" + ] + } + ], + "metadata": { + "colab": { + "name": "weathernext_2_dws.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb b/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb new file mode 100644 index 000000000..814e5b046 --- /dev/null +++ b/notebooks/community/weathernext/weathernext_2_ic_pc.ipynb @@ -0,0 +1,970 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "1a737002", + "metadata": { + "id": "RirpY96_3zL0" + }, + "outputs": [], + "source": [ + "# Copyright 2026 Google LLC\n", + "\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "id": "bf59250d", + "metadata": { + "id": "52d0938f" + }, + "source": [ + "# WeatherNext 2\n", + "\n", + " \n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Google
Open in Colab\n", + " \n", + " 		\n", + " \n", + " \"Google
Open in Colab Enterprise\n", + " \n", + " 		\n", + " \n", + " \"Enterprise
Open in Enterprise Gemini Agent Platform Workbench\n", + " \n", + " 		\n", + " \n", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "id": "3c096b65", + "metadata": { + "id": "DGFzy9MTuYL_" + }, + "source": [ + "## Overview\n", + "\n", + "This notebook demonstrates running [WeatherNext 2 inference on Google Cloud Enterprise Gemini Agent Platform](https://developers.google.com/weathernext/guides/access-vmg) using **customer-provided initial conditions** (custom inputs). WeatherNext 2 is Google's latest medium-range probabilistic forecasting model, principally an operational version the FGN model ([published June 2025](https://arxiv.org/abs/2506.10772)). More information is available in the [WeatherNext documentation](https://developers.google.com/weathernext).\n", + "\n", + "WeatherNext 2 supports customer-provided initial conditions for inference. Instead of using the default ECMWF HRES real-time data, customers can supply their own input Zarr files (e.g., from GFS or their own analysis systems) to generate forecasts with WeatherNext models.\n", + "\n", + "> **Important**: The model is **not fine-tuned** on custom input data. Forecast performance when using custom inputs is **not guaranteed** to match the quality achieved with the default ECMWF HRES inputs. Customers should perform their own evaluation of output quality.\n", + "\n", + "### Objective\n", + "\n", + "- Configure the model inputs for distributed, multi-host inference on H100 or A100 GPUs.\n", + "- Provide custom initial conditions (Zarr files) for model inference.\n", + "- Run WeatherNext 2 model forecasts in parallel.\n", + "- Visualize forecast results.\n", + "\n", + "### Costs\n", + "\n", + "This uses billable components of Google Cloud:\n", + "\n", + "* [Gemini Enterprise Agent Platform]( https://docs.cloud.google.com/gemini-enterprise-agent-platform)\n", + "* [Cloud Storage](https://cloud.google.com/storage/docs)\n", + "* [Gemini Enterprise Agent Platform Persistent Resource](https://docs.cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/persistent-resource-create)\n", + "\n", + "Learn about [Gemini Enterprise Agent Platform pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [Gemini Enterprise Platform Persistent Resource](https://cloud.google.com/products/gemini-enterprise-agent-platform/pricing#custom-trained-models) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n", + "\n", + "\n", + "## Before you begin\n", + "\n", + "### Request For GPU Quota\n", + "\n", + "**WARNING:** 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\n", + "incomplete results.\n", + "\n", + "\n", + "By default, the quota for GPUs is 0. You can request a higher quota by following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n", + "\n", + "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.\n", + "\n", + "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`:\n", + "\n", + "#### Option 1: Persistent Resource\n", + "- Name: `Persistent resource Nvidia A100 80GB GPUs per region` OR `Persistent resource Nvidia H100 GPUs per region`\n", + "\n", + "#### Option 2: Standard Custom Model Training (Preemptible)\n", + "- Name: `Custom model training preemptible Nvidia A100 80GB GPUs per region` OR `Custom model training preemptible Nvidia H100 GPUs per region`\n", + "\n", + "### Custom Inputs Guide\n", + "\n", + "```\n", + "The customer provides input Zarr files in a GCS bucket path via the\n", + "`--cns_data_dir` flag. \n", + "\n", + "#### How to Use\n", + "\n", + "##### Flag: `--input_data_gcs_dir`\n", + "\n", + "`--input_data_gcs_dir` flag for specifying custom input data:\n", + "\n", + "```\n", + "--input_data_gcs_dir=gs://customer-bucket/path-to-input-data\n", + "```\n", + "\n", + "##### GCS Bucket Setup\n", + "\n", + "We recommend using the **same GCS bucket** for both input data and output\n", + "predictions.\n", + "\n", + "The customer should place their input Zarr files under a path within their\n", + "existing output bucket, e.g.:\n", + "\n", + "```\n", + "gs://customer-bucket/custom-inputs/ ← input Zarr files go here\n", + "gs://customer-bucket/outputs/ ← model predictions are written here\n", + "```\n", + "\n", + "#### Input File Format\n", + "\n", + "##### Zarr V3 Format Requirement\n", + "\n", + "**All custom input Zarr files MUST be in Zarr V3 format.**\n", + "\n", + "When creating custom input files, ensure they are saved in Zarr V3 format:\n", + "\n", + "```python\n", + "import xarray as xa\n", + "\n", + "# When creating input files, ensure they are saved in Zarr V3 format\n", + "dataset.to_zarr(\"path/to/output.zarr\", zarr_format=3)\n", + "```\n", + "\n", + "If V2 format files are provided as custom input, the job raises\n", + "an error when attempting to read them.\n", + "\n", + "##### Example Dataset\n", + "\n", + "Opening an example input Zarr file with xarray (init time 2026-03-18 00:00 UTC):\n", + "\n", + "```python\n", + ">>> import xarray as xa\n", + ">>> ds = xa.open_zarr(\"2026_A1D03180000031800011.zarr\")\n", + ">>> ds\n", + " Size: ...\n", + "Dimensions: (isobaricInhPa: 13, latitude: 721, longitude: 1440)\n", + "Coordinates:\n", + " * isobaricInhPa (isobaricInhPa) float64 13 ...\n", + " * latitude (latitude) float64 721 ...\n", + " * longitude (longitude) float64 1440 ...\n", + " number int64 ...\n", + " step int64 ...\n", + " surface float64 ...\n", + " time int64 ...\n", + " valid_time int64 ...\n", + "Data variables: (13 total)\n", + " msl (latitude, longitude) float32 ...\n", + " q (isobaricInhPa, latitude, longitude) float32 ...\n", + " sst (latitude, longitude) float32 ...\n", + " t (isobaricInhPa, latitude, longitude) float32 ...\n", + " t2m (latitude, longitude) float32 ...\n", + " u (isobaricInhPa, latitude, longitude) float32 ...\n", + " u10 (latitude, longitude) float32 ...\n", + " u100 (latitude, longitude) float32 ...\n", + " v (isobaricInhPa, latitude, longitude) float32 ...\n", + " v10 (latitude, longitude) float32 ...\n", + " v100 (latitude, longitude) float32 ...\n", + " w (isobaricInhPa, latitude, longitude) float32 ...\n", + " z (isobaricInhPa, latitude, longitude) float32 ...\n", + "```\n", + "\n", + "##### Variables\n", + "\n", + "All data variables are `float32`.\n", + "\n", + "Only the variables and levels needed for inference are listed here. Your zarr\n", + "may contain additional variables and levels.\n", + "\n", + "###### Pressure Level Variables (3D: `isobaricInhPa` × `latitude` × `longitude`)\n", + "\n", + "Shape: `[13, 721, 1440]`\n", + "\n", + "These variables are used for input at the following pressure levels (hPa):\n", + "`50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000`\n", + "\n", + "| Short Name | Description |\n", + "| :--- | :--- |\n", + "| `q` | specific humidity |\n", + "| `t` | temperature |\n", + "| `u` | u component of wind |\n", + "| `v` | v component of wind |\n", + "| `w` | vertical velocity |\n", + "| `z` | geopotential |\n", + "\n", + "###### Surface Variables (2D: `latitude` × `longitude`)\n", + "\n", + "Shape: `[721, 1440]`\n", + "\n", + "| Short Name | Description |\n", + "| :--- | :--- |\n", + "| `msl` | mean sea level pressure |\n", + "| `sst` | sea surface temperature |\n", + "| `t2m` | 2m temperature |\n", + "| `u10` | 10m u component of wind |\n", + "| `u100` | 100m u component of wind |\n", + "| `v10` | 10m v component of wind |\n", + "| `v100` | 100m v component of wind |\n", + "\n", + "###### Scalar Coordinates\n", + "\n", + "| Name | dtype | Description |\n", + "|--------------|-----------|----------------------------------------------|\n", + "| `number` | `int64` | Ensemble member number |\n", + "| `step` | `int64` | Forecast step |\n", + "| `surface` | `float64` | Surface level indicator |\n", + "| `time` | `int64` | Timestamp (units: days since init time, calendar: proleptic_gregorian) |\n", + "| `valid_time` | `int64` | Validity time |\n", + "\n", + "##### Dimension Coordinates\n", + "\n", + "| Coordinate | dtype | Shape | Description |\n", + "|-----------------|-----------|----------|--------------------------------|\n", + "| `isobaricInhPa` | `float64` | `[13]` | 13 pressure levels in hPa |\n", + "| `latitude` | `float64` | `[721]` | 0.25° resolution, 721 points |\n", + "| `longitude` | `float64` | `[1440]` | 0.25° resolution, 1440 points |\n", + "\n", + "##### Spatial Resolution\n", + "\n", + "The data is at **0.25° resolution** globally:\n", + "\n", + "- Latitude: 721 points (90°N to 90°S)\n", + "- Longitude: 1440 points (0° to 359.75°E)\n", + "\n", + "#### File Naming and Structure\n", + "\n", + "##### File Naming Convention\n", + "\n", + "The inference binary expects input Zarr files to follow a specific file naming\n", + "convention. Each file corresponds to a specific forecast initialization time and\n", + "uses the following format:\n", + "\n", + "```\n", + "_.zarr\n", + "```\n", + "\n", + "Where:\n", + "\n", + "- ``: 4-digit year (e.g., `2025`)\n", + "- ``: Data config name (default: `A1`)\n", + "- ``: `D` for 00/12 UTC init times, `S` for 06/18 UTC\n", + "- First `MMDDHHMM`: Month, day, hour, minute of the forecast init time\n", + "- Second `MMDDHHMM`: Month, day, hour, minute of the validity time\n", + "- `EE`: Experiment version (default: `1`)\n", + "\n", + "The validity minute is hardcoded to `01`.\n", + "\n", + "###### Examples\n", + "\n", + "For a forecast initialized at **2026-03-18 00:00 UTC** using fc0:\n", + "```\n", + "2026_A1D03180000031800011.zarr\n", + "```\n", + "\n", + "For a forecast initialized at **2026-03-17 06:00 UTC** using fc0:\n", + "```\n", + "2026_A1S03170600031706011.zarr\n", + "```\n", + "\n", + "##### Success Sentinels\n", + "\n", + "Each Zarr file directory must contain a `success` sentinel file to signal that\n", + "the data is complete and ready for reading:\n", + "\n", + "```\n", + "gs:///custom-inputs/\n", + "├── 2026_A1D03180000031800011.zarr/\n", + "│ ├── .zmetadata\n", + "│ ├── \n", + "│ └── success ← required sentinel file\n", + "├── 2026_A1D03171200031712011.zarr/\n", + "│ ├── .zmetadata\n", + "│ ├── \n", + "│ └── success\n", + "```\n", + "\n", + "The sentinel is a zero-byte file named `success` placed inside each `.zarr`\n", + "directory. The job will not proceed with inference until all required input\n", + "sentinels exist.\n", + "\n", + "##### Number of Input Files\n", + "\n", + "The model typically requires **2 input timestamps** (the forecast init time and\n", + "6 hours prior). For example, for a forecast initialized at 2026-03-18 12:00 UTC,\n", + "the binary expects:\n", + "\n", + "1. `2026_A1D03181200031812011.zarr` (init time: 12:00 UTC)\n", + "2. `2026_A1D03180600031806011.zarr` (6 hours prior: 06:00 UTC)\n", + "\n", + "#### Caveats and Limitations\n", + "\n", + "1. **No fine-tuning guarantee**: The model is trained on ECMWF HRES\n", + " data. Using custom inputs from a different source may degrade\n", + " forecast quality, especially for features sensitive to the initial\n", + " condition source.\n", + "\n", + "2. **Variable completeness**: All variables listed above must be present in the\n", + " custom input files. Missing variables will cause the inference to fail.\n", + "\n", + "3. **Temporal alignment**: Custom input timestamps must align with valid HRES\n", + " forecast hours (00, 06, 12, or 18 UTC).\n", + "\n", + "4. **File format**: Only Zarr V3 format is accepted." + ] + }, + { + "cell_type": "markdown", + "id": "d656ed2b", + "metadata": { + "id": "9349a5bece68" + }, + "source": [ + "## Install packages - Restart the kernel after the installation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cbfab540", + "metadata": { + "id": "yt4GxkKcDD7Y" + }, + "outputs": [], + "source": [ + "# @title Install python packages\n", + "\n", + "# Note that you may need to restart the kernel after this step.\n", + "# If so, continue to the next cell after restarting.\n", + "\n", + "print(\"Installing python packages.\")\n", + "\n", + "! pip3 install \\\n", + " google-cloud-aiplatform==1.129.0 \\\n", + " xarray[complete]" + ] + }, + { + "cell_type": "markdown", + "id": "c123b19f", + "metadata": { + "id": "48a20b368e8c" + }, + "source": [ + "## Authenticate to Google Cloud Platform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3b7e5c45", + "metadata": { + "id": "pw5fTRvrDPQv" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + " from google.colab import auth\n", + "\n", + " auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "id": "7abff806", + "metadata": { + "id": "fd0ba250117b" + }, + "source": [ + "## Set Google Cloud Project Pertinent Variables" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b0805e58", + "metadata": { + "id": "iklx3Yb3Dmke" + }, + "outputs": [], + "source": [ + "# Replace my_gcp_project with your gcp project\n", + "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "# Alternatively collect the default cloud project id from the OS env variable\n", + "# PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n", + "\n", + "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "# @markdown 2. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched.\n", + "\n", + "# Replace my_wn_bucket with your bucket\n", + "BUCKET_URI = \"gs://\" # @param {type:\"string\"}\n", + "\n", + "# @markdown 3. Select a region that has the required GPUs available.\n", + "\n", + "# Select a ***US*** region region\n", + "REGION = \"us-central1\" # @param {type:\"string\"}\n", + "\n", + "\n", + "# @markdown 4. Provide the Persistence Resource ID\n", + "PERSISTENT_RESOURCE_ID = \"\" # @param {type:\"string\"}\n", + "\n", + "# Set Docker URI\n", + "# @markdown 5. Set the Image URI\n", + "WEATHERNEXT2_DOCKER_URI = \"us-central1-docker.pkg.dev/weathernext-1/wn25-private-preview-launch/weather-next-ic-2-inference.gpu.0-1:latest\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9af9574a", + "metadata": { + "id": "KHbl4756DjMU" + }, + "outputs": [], + "source": [ + "!gcloud config set project $PROJECT_ID" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d914b484", + "metadata": { + "id": "dO2mF4CPfpHW1ZKbIsut8r69" + }, + "outputs": [], + "source": [ + "# Import the necessary packages\n", + "import datetime\n", + "import os\n", + "import re\n", + "\n", + "from google.cloud import aiplatform\n", + "\n", + "# Enable the Enterprise Gemini Agent Platform API and Compute Engine API, if not already.\n", + "print(\"Enabling Enterprise Gemini Agent Platform API and Compute Engine API.\")\n", + "# ! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n", + "\n", + "# Cloud Storage bucket for storing the experiment artifacts.\n", + "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n", + "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n", + "\n", + "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n", + " raise ValueError(\"GCS Bucket URI is invalid!\")\n", + "else:\n", + " assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n", + " shell_output = !gcloud storage buckets describe {BUCKET_NAME} --format=\"value(location)\" | tr '[:upper:]' '[:lower:]'\n", + " bucket_region = shell_output[0].strip().lower()\n", + " if bucket_region != REGION:\n", + " raise ValueError(\n", + " f\"Bucket region {bucket_region} is different from notebook region {REGION}\"\n", + " )\n", + "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n", + "\n", + "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n", + "# Initialize Enterprise Gemini Agent Platform API.\n", + "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n", + "\n", + "# Utility functions\n", + "\n", + "\n", + "def get_job_name_with_datetime(prefix: str) -> str:\n", + " return prefix + datetime.datetime.now().strftime(\"_%Y%m%d_%H%M%S\")" + ] + }, + { + "cell_type": "markdown", + "id": "95c31932", + "metadata": { + "id": "cb667a3392f7" + }, + "source": [ + "## Configure Model Parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f4854c78", + "metadata": { + "id": "22BW43yjps8D" + }, + "outputs": [], + "source": [ + "# @markdown ### Hardware Configuration for Distributed Inference\n", + "# @markdown - **`machine_type`**: Select a valid machine type. `a3-highgpu` series use NVIDIA H100 80GB GPUs. `a2-ultragpu` series use NVIDIA A100 80GB GPUs.\n", + "# @markdown - **`num_samples`**: The total number of ensemble members to generate.\n", + "# @markdown The number of machine replicas will be calculated automatically (`num_samples` / GPUs per machine). **Therefore, `num_samples` must be a multiple of the number of GPUs in your selected `machine_type`.**\n", + "machine_type = \"a2-ultragpu-1g\" # @param [\"a3-highgpu-1g\", \"a3-highgpu-2g\", \"a3-highgpu-4g\", \"a3-highgpu-8g\", \"a2-ultragpu-1g\", \"a2-ultragpu-2g\", \"a2-ultragpu-4g\", \"a2-ultragpu-8g\"]\n", + "num_samples = 4 # @param {type:\"integer\"}\n", + "\n", + "# @markdown ### Forecast Configuration\n", + "# @markdown - **`forecast_init_time`**: The starting time for the forecast in ISO 8601 format (e.g., `2025-09-21T00:00:00Z`). Models are available for dates from 2024 onwards.\n", + "# @markdown - **`horizon_hrs`**: The desired length of the forecast in hours (e.g., 240 for a 10-day forecast).\n", + "# @markdown - **`model_seed`**: Choose a specific model seed (1-4) or select \"all\" to run inference with all four seeds in parallel for improved accuracy.\n", + "# @markdown - **`enable_hourly_prediction`**: If checked, the model will generate 1-hour predictions.\n", + "forecast_init_time = \"2026-05-16T12:00:00Z\"\n", + "\n", + "horizon_hrs = 72 # @param {type:\"integer\"}\n", + "model_seed = \"all\" # @param [\"1\", \"2\", \"3\", \"4\", \"all\"]\n", + "enable_hourly_prediction = True # @param {type:\"boolean\"}\n", + "\n", + "# --- Parameter Validation and Configuration ---\n", + "\n", + "# Derive accelerator type and count from the chosen machine type\n", + "if machine_type.startswith(\"a3-highgpu\"):\n", + " accelerator_type = \"NVIDIA_H100_80GB\"\n", + "elif machine_type.startswith(\"a2-ultragpu\"):\n", + " accelerator_type = \"NVIDIA_A100_80GB\"\n", + "else:\n", + " raise ValueError(f\"Invalid machine type selected: {machine_type}.\")\n", + "\n", + "try:\n", + " # Extract the number of GPUs from the machine type string, e.g., 'a3-highgpu-4g' -> 4\n", + " accelerators_per_machine = int(re.search(r\"-(\\d+)g$\", machine_type).group(1))\n", + "except (AttributeError, ValueError):\n", + " raise ValueError(\n", + " f\"Could not determine accelerator count from machine type: {machine_type}\"\n", + " )\n", + "\n", + "seeds_to_run = [1, 2, 3, 4] if model_seed == \"all\" else [int(model_seed)]\n", + "num_seeds_to_run = len(seeds_to_run)\n", + "\n", + "num_samples_per_seed = num_samples\n", + "if len(seeds_to_run) > 1:\n", + " if num_samples % num_seeds_to_run != 0:\n", + " raise ValueError(\n", + " f\"`num_samples` ({num_samples}) is not divisible by the number of seeds to run ({num_seeds_to_run}.\"\n", + " )\n", + " num_samples_per_seed = num_samples // num_seeds_to_run\n", + "\n", + "# Validate that num_samples is a multiple of accelerators_per_machine\n", + "if num_samples_per_seed % accelerators_per_machine != 0:\n", + " raise ValueError(\n", + " f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n", + " )\n", + "\n", + "# Calculate the number of replicas per seed\n", + "replica_count_per_seed = num_samples_per_seed // accelerators_per_machine\n", + "\n", + "# Ensure that there are enough samples\n", + "if num_samples_per_seed % accelerators_per_machine != 0:\n", + " raise ValueError(\n", + " f\"`num_samples_per_seed` ({num_samples_per_seed}) must be a multiple of the GPUs per machine ({accelerators_per_machine} for {machine_type}).\"\n", + " )\n", + "\n", + "# Calculate total GPUs needed for all jobs\n", + "total_gpus_needed = num_samples * (4 if model_seed == \"all\" else 1)\n", + "\n", + "print(\"--- Job Configuration Summary ---\")\n", + "print(f\"Total Samples: {num_samples}\")\n", + "print(f\"Machine Type: {machine_type}\")\n", + "print(f\"Accelerator Type: {accelerator_type}\")\n", + "print(f\"GPUs per Machine: {accelerators_per_machine}\")\n", + "print(f\"Total number seeds to run: {num_seeds_to_run}\")\n", + "print(f\"Total number samples per seed: {num_samples_per_seed}\")\n", + "print(f\"Calculated Machine Replicas Per Seed: {replica_count_per_seed}\")\n", + "print(f\"Total GPUs per Job: {num_samples}\")\n", + "print(f\"Total GPUs across all Jobs (ensure sufficient quota): {total_gpus_needed}\")\n", + "print(f\"Docker Image: {WEATHERNEXT2_DOCKER_URI}\")\n", + "print(\"---------------------------------\")" + ] + }, + { + "cell_type": "markdown", + "id": "0ba7d73c", + "metadata": { + "id": "fe8b458ce060" + }, + "source": [ + "## Run forecast jobs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "42329e44", + "metadata": { + "id": "jVTDLIM87QSr" + }, + "outputs": [], + "source": [ + "# @markdown This section creates and runs one or more Enterprise Gemini Agent Platform Custom Training Jobs to generate the forecasts.\n", + "# @markdown **This operation is asynchronous.** The jobs will be submitted and this cell will complete quickly.\n", + "# @markdown You must monitor the job progress in the Google Cloud Console (https://console.cloud.google.com/vertex-ai/training/custom-jobs).\n", + "\n", + "import time\n", + "\n", + "print(\n", + " f\"Submitting {len(seeds_to_run)} job(s) to target persistent resource: {PERSISTENT_RESOURCE_ID}\"\n", + ")\n", + "\n", + "launched_jobs = []\n", + "output_dirs = {}\n", + "\n", + "for seed in seeds_to_run:\n", + " output_dir = os.path.join(BUCKET_URI, \"weathernext2_outputs\")\n", + " output_dirs[seed] = output_dir\n", + "\n", + " docker_args_list = [\n", + " f\"--pred_root_dir={output_dir}\",\n", + " f\"--num_samples={num_samples_per_seed}\",\n", + " f\"--horizon_hrs={horizon_hrs}\",\n", + " f\"--forecast_init_time={forecast_init_time}\",\n", + " f\"--model_seed={seed}\",\n", + " f\"--enable_hourly_prediction={enable_hourly_prediction}\",\n", + " f\"--input_data_gcs_dir={BUCKET_URI}/custom_inputs/\",\n", + " ]\n", + "\n", + " # Define the worker pool spec for the custom job matching the warm pool specs\n", + " worker_pool_specs = [\n", + " {\n", + " \"machine_spec\": {\n", + " \"machine_type\": machine_type,\n", + " \"accelerator_type\": accelerator_type,\n", + " \"accelerator_count\": accelerators_per_machine,\n", + " },\n", + " \"disk_spec\": {\n", + " \"boot_disk_type\": \"pd-standard\",\n", + " \"boot_disk_size_gb\": 100,\n", + " },\n", + " \"replica_count\": 1,\n", + " \"container_spec\": {\n", + " \"image_uri\": WEATHERNEXT2_DOCKER_URI,\n", + " \"args\": docker_args_list,\n", + " },\n", + " }\n", + " ]\n", + "\n", + " JOB_NAME = get_job_name_with_datetime(prefix=f\"wn2-persistent-run-s{seed}\")\n", + " print(f\"\\n--- Submitting Job for Seed {seed} ---\")\n", + " print(f\"JOB_NAME: {JOB_NAME}\")\n", + "\n", + " custom_job = aiplatform.CustomJob(\n", + " display_name=JOB_NAME,\n", + " worker_pool_specs=worker_pool_specs,\n", + " project=PROJECT_ID,\n", + " location=REGION,\n", + " staging_bucket=f\"{BUCKET_URI}/custom_job_staging\",\n", + " )\n", + "\n", + " # Run the job asynchronously targeting the warm persistent resource pool!\n", + " # Since the 4-GPU persistent resource was created without explicit custom SA runtime permissions,\n", + " # we run using the default Enterprise Gemini Agent Platform Service Agent (which already has GCS and registry reader permissions).\n", + " custom_job.run(\n", + " persistent_resource_id=PERSISTENT_RESOURCE_ID,\n", + " disable_retries=True,\n", + " sync=False,\n", + " )\n", + "\n", + " # Wait for GCA resource generation details\n", + " print(\"Waiting for CustomJob details...\")\n", + " for _ in range(15):\n", + " try:\n", + " if custom_job.resource_name:\n", + " break\n", + " except Exception:\n", + " pass\n", + " time.sleep(1)\n", + "\n", + " job_id = custom_job.name.split(\"/\")[-1]\n", + " print(f\"--> Job submitted successfully. ID: {job_id}\")\n", + " print(\n", + " f\"Console Link: https://console.cloud.google.com/vertex-ai/locations/{REGION}/training/custom-jobs/{job_id}?project={PROJECT_ID}\"\n", + " )\n", + " launched_jobs.append(custom_job)\n", + "\n", + "print(\"\\nAll forecast jobs have been submitted. Starting real-time status monitor...\")\n", + "\n", + "# Monitoring loop\n", + "completed_jobs = set()\n", + "while len(completed_jobs) < len(launched_jobs):\n", + " for job in launched_jobs:\n", + " if job.name in completed_jobs:\n", + " continue\n", + "\n", + " job._sync_gca_resource()\n", + " state = job.state.value if hasattr(job.state, \"value\") else int(job.state)\n", + " timestamp = datetime.datetime.now().strftime(\"%H:%M:%S\")\n", + " print(\n", + " f\"[{timestamp}] Job '{job.display_name}': {job.state.name} (Code: {state})\"\n", + " )\n", + "\n", + " # Final states: 4 is SUCCEEDED, 5 is FAILED, 7 is CANCELLED\n", + " if state in [4, 5, 7]:\n", + " completed_jobs.add(job.name)\n", + " if state == 4:\n", + " print(f\"\\n🎉 SUCCESS: Job '{job.display_name}' completed successfully!\")\n", + " else:\n", + " print(\n", + " f\"\\n⚠️ FAILURE/CANCEL: Job '{job.display_name}' exited with code {state}. Error: {getattr(job, 'error', 'N/A')}\"\n", + " )\n", + "\n", + " if len(completed_jobs) < len(launched_jobs):\n", + " time.sleep(20)\n", + "\n", + "print(\"\\nAll forecast runs are complete!\")" + ] + }, + { + "cell_type": "markdown", + "id": "d11c21c4", + "metadata": { + "id": "112398663144" + }, + "source": [ + "## Visualize Forecasts (Unified) " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b04bd730", + "metadata": { + "id": "x6ZVqWRopWcI" + }, + "outputs": [], + "source": [ + "# @markdown Select which forecast output you want to visualize. This single component\n", + "# @markdown can handle both the standard 6-hourly predictions and the datasets\n", + "# @markdown with 1-hour model (which have a 'subtime' dimension).\n", + "# @markdown If you run into `Error loading Zarr store: unrecognized engine 'zarr'...` try restarting the runtime session and reruning this cell.\n", + "\n", + "# @markdown ---\n", + "# @markdown ### Visualization Settings\n", + "# @markdown - **`model_seed_to_visualize`**: Choose a specific model seed (1-4) to visualize. This should be one of the model seeds selected in the **Forecast Configuration** above.\n", + "# @markdown - **`time_steps_to_visualize`**: Choose to visualize 1-hourly or 6-hourly forecasts. If 1-hourly is selected, ensure `enable_hourly_prediction` was selected in the **Forecast Configuration** above.\n", + "# @markdown - **`variable_to_visualize`**: Choose the weather variable to visualize. See the [WeatherNext documentation](https://developers.google.com/weathernext/guides/model-specs-vmg) for variable names and descriptions.\n", + "# @markdown - **`sample_to_visualize`**: Choose the sample (ensemble member) to visualize.\n", + "# @markdown - **`plot_size`**: Choose the size of the plot to generate.\n", + "model_seed_to_visualize = \"4\" # @param [\"1\", \"2\", \"3\", \"4\"]\n", + "time_steps_to_visualize = \"6-Hourly\" # @param [\"6-Hourly\", \"1-Hourly\"]\n", + "variable_to_visualize = \"2m_temperature\" # @param {type:\"string\"}\n", + "sample_to_visualize = 0 # @param {type:\"integer\"}\n", + "plot_size = 8 # @param {type:\"number\"}\n", + "level_to_visualize = None\n", + "# @markdown ---\n", + "\n", + "\n", + "# --- Helper Functions ---\n", + "\n", + "\n", + "def init_time_to_folder_path(init_time: str) -> str:\n", + " \"\"\"\n", + " Convert init time to expected GCS folder path.\n", + " \"\"\"\n", + " init_date, init_time = init_time.split(\"T\")\n", + " return f\"{init_date.replace('-', '')}_{init_time[0:2]}hr\"\n", + "\n", + "\n", + "# override these if you'd like to visualize a different set of forecasts\n", + "visualize_bucket = BUCKET_URI\n", + "visualize_init_date = forecast_init_time\n", + "\n", + "# set paths based on chosen model seed, bucket, and init date\n", + "path_to_6hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n", + "path_to_1hr_zarr = f\"{BUCKET_URI}/weathernext2_outputs/weathernext_2_seed_{model_seed_to_visualize}_hourly/{init_time_to_folder_path(visualize_init_date)}_01_preds/predictions.zarr/\"\n", + "\n", + "\n", + "from typing import Optional\n", + "\n", + "import matplotlib\n", + "import matplotlib.animation as animation\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import xarray\n", + "from IPython.display import HTML\n", + "\n", + "matplotlib.rcParams[\"animation.embed_limit\"] = 500\n", + "\n", + "\n", + "# --- Helper Functions ---\n", + "\n", + "\n", + "def select_data(\n", + " data: xarray.Dataset,\n", + " variable: str,\n", + " level: Optional[int] = None,\n", + ") -> xarray.Dataset:\n", + " \"\"\"Selects a variable from the dataset and optionally a level.\"\"\"\n", + " data = data[variable]\n", + " if \"batch\" in data.dims:\n", + " data = data.isel(batch=0)\n", + " if level is not None and \"level\" in data.coords:\n", + " data = data.sel(level=level)\n", + " return data\n", + "\n", + "\n", + "def scale_data(\n", + " data: xarray.Dataset,\n", + " center: Optional[float] = None,\n", + " robust: bool = False,\n", + ") -> tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n", + " \"\"\"Scales the data for visualization.\"\"\"\n", + " vmin = np.nanpercentile(data.values, (2 if robust else 0))\n", + " vmax = np.nanpercentile(data.values, (98 if robust else 100))\n", + " if center is not None:\n", + " diff = max(vmax - center, center - vmin)\n", + " vmin = center - diff\n", + " vmax = center + diff\n", + " return (\n", + " data,\n", + " matplotlib.colors.Normalize(vmin, vmax),\n", + " (\"RdBu_r\" if center is not None else \"viridis\"),\n", + " )\n", + "\n", + "\n", + "def create_forecast_animation(\n", + " dataset: xarray.Dataset,\n", + " fig_title: str,\n", + " plot_size: float = 5,\n", + " robust: bool = False,\n", + ") -> HTML:\n", + " \"\"\"\n", + " Creates a forecast animation from an xarray Dataset.\n", + " It intelligently handles datasets with or without a 'subtime' dimension.\n", + " \"\"\"\n", + " # --- Data Preparation ---\n", + " # Check if the data still has 'subtime'). If so, stack dimensions.\n", + " # Otherwise, just rename the 'time' dimension for consistency.\n", + " if \"subtime\" in dataset.dims:\n", + " print(\"Detected 'subtime' dimension. Stacking for hourly animation.\")\n", + " # Stack 'time' and 'subtime' into a single animation dimension\n", + " plot_data = dataset.stack(animation_step=(\"time\", \"subtime\")).transpose(\n", + " \"animation_step\", \"lat\", \"lon\"\n", + " )\n", + " else:\n", + " print(\"No 'subtime' dimension found. Using 'time' for 6-hourly animation.\")\n", + " # Use 'time' as the animation dimension\n", + " plot_data = dataset.rename({\"time\": \"animation_step\"})\n", + "\n", + " # Now, the animation dimension is always called 'animation_step'\n", + " max_steps = plot_data.sizes[\"animation_step\"]\n", + " init_time = plot_data.coords[\"init_time\"].values\n", + "\n", + " # Scale the data for color mapping\n", + " scaled_data, norm, cmap = scale_data(plot_data, robust=robust)\n", + "\n", + " # --- Plotting Setup ---\n", + " figure = plt.figure(figsize=(plot_size * 2, plot_size))\n", + " ax = figure.add_subplot(1, 1, 1)\n", + " ax.set_xticks([])\n", + " ax.set_yticks([])\n", + " figure.suptitle(fig_title, fontsize=16)\n", + " figure.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust for title\n", + "\n", + " im = ax.imshow(\n", + " scaled_data.isel(animation_step=0), norm=norm, origin=\"lower\", cmap=cmap\n", + " )\n", + "\n", + " plt.colorbar(\n", + " mappable=im,\n", + " ax=ax,\n", + " orientation=\"vertical\",\n", + " pad=0.02,\n", + " aspect=16,\n", + " shrink=0.75,\n", + " cmap=cmap,\n", + " extend=(\"both\" if robust else \"neither\"),\n", + " )\n", + "\n", + " # --- Animation Update Function ---\n", + " def update(frame):\n", + " # Get the coordinates for the current frame\n", + " step_coords = plot_data[\"animation_step\"][frame].coords\n", + "\n", + " # Calculate total offset and valid time based on available coordinates\n", + " if \"subtime\" in step_coords: # Hourly data\n", + " total_offset = step_coords[\"time\"].values + step_coords[\"subtime\"].values\n", + " else: # 6-hourly data\n", + " total_offset = step_coords[\"animation_step\"].values\n", + "\n", + " total_hours = total_offset / np.timedelta64(1, \"h\")\n", + " valid_time = init_time + total_offset\n", + " valid_time_str = np.datetime_as_string(valid_time, unit=\"m\").replace(\"T\", \" \")\n", + "\n", + " new_title = (\n", + " f\"{fig_title}\\n\"\n", + " f\"Valid: {valid_time_str} UTC (Forecast: +{total_hours:.1f}h)\"\n", + " )\n", + " figure.suptitle(new_title, fontsize=16)\n", + " im.set_data(scaled_data.isel(animation_step=frame))\n", + "\n", + " # --- Create and Display Animation ---\n", + " ani = animation.FuncAnimation(\n", + " fig=figure, func=update, frames=max_steps, interval=250\n", + " )\n", + " plt.close(figure.number)\n", + " return HTML(ani.to_html5_video())\n", + "\n", + "\n", + "# --- Main Visualization Logic ---\n", + "\n", + "# 1. Select the correct path based on the user's dropdown choice\n", + "if time_steps_to_visualize == \"6-Hourly\":\n", + " path_to_zarr = path_to_6hr_zarr\n", + "elif time_steps_to_visualize == \"1-Hourly\":\n", + " path_to_zarr = path_to_1hr_zarr\n", + "else:\n", + " raise ValueError(\"Invalid visualization target selected.\")\n", + "\n", + "print(f\"Loading data from: {path_to_zarr}\")\n", + "\n", + "# 2. Load the dataset\n", + "try:\n", + " full_dataset = xarray.open_zarr(path_to_zarr)\n", + "except Exception as e:\n", + " print(f\"Error loading Zarr store: {e}\")\n", + " # This is a common point of failure, so we exit gracefully.\n", + "else:\n", + " # 3. Select the specific data slice for visualization\n", + " data_for_vis = full_dataset.isel(sample=sample_to_visualize)\n", + " variable_data = select_data(data_for_vis, variable_to_visualize, level_to_visualize)\n", + "\n", + " # 4. Generate the title\n", + " title = f\"{variable_to_visualize} (Sample {sample_to_visualize})\"\n", + " if level_to_visualize:\n", + " title += f\" at {level_to_visualize} hPa\"\n", + "\n", + " # 5. Create and display the animation\n", + " display(create_forecast_animation(variable_data, title, plot_size, robust=True))" + ] + } + ], + "metadata": { + "colab": { + "name": "weathernext_2_ic_pc.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +}