feat: WeatherNext2 Initial Updates

notebooks/community/CODEOWNERS

@@ -27,6 +27,7 @@
 /vertex_endpoints/nvidia-triton/nvidia-triton-custom-container-prediction.ipynb @RajeshThallam
 /vertex_endpoints/optimized_tensorflow_runtime @vlasenkoalexey
 /notebooks/community/alphagenome/cloudai_alphagenome_vai_quickstart.ipynb @dpanigra
+/notebooks/community/weathernext/weathernext_2_early_access_program.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

notebooks/community/weathernext/README.md

@@ -0,0 +1,78 @@
+# WeatherNext 2
+[**Overview**](#overview) | [**Use Cases**](#use-cases) | [**Documentation**](#documentation) | [**Pricing**](#pricing) | [**Quick start**](#quick-start)
+
+## 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). <!-- disableFinding(LINE_OVER_80) -->
+
+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. <!-- disableFinding(LINE_OVER_80) -->
+
+**Overview**
+
+WeatherNext is a family of state-of-the-art AI weather forecasting models from Google DeepMind and Google Research. WeatherNext models are faster and more efficient than traditional physics-based weather models and yield superior forecast reliability. <!-- disableFinding(LINE_OVER_80) -->
+
+WeatherNext 2
+Generates ensemble forecasts at a spatial resolution of 0.25 degrees. Forecast init times have 6 hour resolution (00z, 06z, 12z, 18z). Forecast lead times have 1 hour resolution and a lead time of 15 days. Additional information on the model is described in "[Skillful joint probabilistic weather forecasting from marginals](https://arxiv.org/abs/2506.10772)”. <!-- disableFinding(LINE_OVER_80) -->
+
+## Use Cases
+**Use Cases**
+
+*What use cases are best for WeatherNext 2?*\
+WeatherNext 2 offers a fast, high-resolution global medium-range forecasting
+model, serving as a powerful alternative to traditional systems like ECMWF ENS
+or NOAA GEFS. Key applications include medium-range forecasts (15 days),
+predicting severe events like tropical cyclones, feeding data into downstream
+systems such as flood models, and generating large ensembles for comprehensive
+uncertainty quantification
+
+## Documentation
+**Training data**
+
+ECMWF [ERA5](https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5) (1979-2018)
+
+ECMWF [HRES](https://www.ecmwf.int/en/forecasts/datasets/set-i)
+
+NOAA [IBTrACS](https://www.ncei.noaa.gov/products/international-best-track-archive)
+
+**Model Evaluation**
+
+Model training and evaluation methodology, and a complete set of model
+scorecards evaluating each across variables and lead times, are reported in
+the research paper associated with the WeatherNext 2 model:
+
+* [Skillful joint probabilistic weather forecasting from marginals](https://arxiv.org/abs/2506.10772)
+
+**Documentation**
+
+See [the WeatherNext developer’s guide](http://developers.google.com/weathernext) for documentation regarding the models.
+
+**Acknowledgements**
+
+The models communicate with and reference data and products of the European
+Centre for Medium-range Weather Forecasts (ECMWF), as modified by Google.
+
+**Contact**
+
+If you have any questions on using these models please contact: [weathernext-cloud@google.com](mailto:weathernext-cloud@google.com). Any information collected via email will be used in accordance with [Google's privacy policy](https://policies.google.com/privacy).
+
+**Links**
+
+* [WeatherNext developer’s guide](http://developers.google.com/weathernext)
+* [Research model paper](https://arxiv.org/abs/2506.10772)
+
+## Pricing
+Access to forecasting with the WeatherNext models on Vertex AI is currently
+restricted.
+To utilize these models for generating real-time forecasts via this service:
+
+* You must **apply for access**.
+* Your application will be reviewed, and if approved, you will be **added to an allowlist**.
+* Only allow listed users can access the forecasting service.
+* **Pricing information** for the forecasting service will be shared directly with users upon approval and placement on the allowlist.
+
+Running these models *will incur costs* for the GPUs and other Google Cloud resources used. Learn about [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/) to generate an estimate.
+
+## Quick start
+The quickest way to get started with the WeatherNext in Google Cloud Platform is to run [our example notebook](weathernext_2_early_access_program.ipynb) in [Google Colab](https://colab.research.google.com/).

notebooks/community/weathernext/weathernext_2_early_access_program.ipynb

@@ -0,0 +1,558 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "RirpY96_3zL0",
+      "metadata": {},
+      "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": "52d0938f",
+      "metadata": {},
+      "source": [
+        "# WeatherNext 2 Early Access Program\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_early_access_program.ipynb\">\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%weathernext%2Fweathernext_2_early_access_program.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_early_access_program.ipynb\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/weathernext/weathernext_2_early_access_program.ipynb\">\n",
+        "      <img width=\"32px\"src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "DGFzy9MTuYL_",
+      "metadata": {},
+      "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,
+      "id": "yt4GxkKcDD7Y",
+      "metadata": {},
+      "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,
+      "id": "dO2mF4CPfpHW1ZKbIsut8r69",
+      "metadata": {
+        "tags": []
+      },
+      "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 importlib\n",
+        "import os\n",
+        "import uuid\n",
+        "import glob\n",
+        "from google.cloud import aiplatform, storage\n",
+        "\n",
+        "import json\n",
+        "import math\n",
+        "import re\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(f\"Bucket region {bucket_region} is different from notebook region {REGION}\")\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",
+        "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,
+      "id": "22BW43yjps8D",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# @title Configure Model Parameters\n",
+        "# @markdown Configure the hardware and input parameters for the WeatherNext 2 forecast.\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(f\"Could not determine accelerator count from machine type: {machine_type}\")\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(f\"`num_samples` ({num_samples}) is not divisible by the number of seeds to run ({num_seeds_to_run}.\")\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(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",
+        "# 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(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",
+        "# 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,
+      "id": "VSuGx2jmpwP0",
+      "metadata": {},
+      "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(prefix=f\"wn2-forecast-s{seed}-n{num_samples_per_seed}\")\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(f\"--> Job submitted successfully. Monitor it in the Google Cloud Console at https://console.cloud.google.com/vertex-ai/training/custom-jobs\")\n",
+        "\n",
+        "print(\"\\nAll forecast jobs have been submitted.\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "x6ZVqWRopWcI",
+      "metadata": {
+        "cellView": "form"
+      },
+      "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",
+        "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 matplotlib\n",
+        "import xarray\n",
+        "from typing import Optional, Tuple\n",
+        "import matplotlib.pyplot as plt\n",
+        "import matplotlib.animation as animation\n",
+        "import math\n",
+        "from IPython.display import HTML\n",
+        "import numpy as np\n",
+        "import datetime\n",
+        "\n",
+        "matplotlib.rcParams['animation.embed_limit'] = 500\n",
+        "\n",
+        "\n",
+        "# --- Helper Functions ---\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",
+        "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 (data, matplotlib.colors.Normalize(vmin, vmax),\n",
+        "          (\"RdBu_r\" if center is not None else \"viridis\"))\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(\n",
+        "        animation_step=(\"time\", \"subtime\")\n",
+        "    ).transpose(\"animation_step\", \"lat\", \"lon\")\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",
+        "  plt.colorbar(\n",
+        "      mappable=im, ax=ax, orientation=\"vertical\", pad=0.02,\n",
+        "      aspect=16, shrink=0.75, cmap=cmap,\n",
+        "      extend=(\"both\" if robust else \"neither\"))\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",
+        "  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": "seedfix_scheduling_weathernext_2_early_access_program.ipynb",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.10.10"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}