fix: Point to new reservation for GCE for jax_functional_tests DAG (#1143)

Since the reservation in project `tpu-prod-env-automated` does not exist for zone east5 anymore. It was need it to create a new subnetwork and point to a zone where is enough v5p chips reserved. Point to zone `europe-west4-b`

Author: camiloCienet

dags/common/vm_resource.py

@@ -23,6 +23,9 @@
 V5_NETWORKS = f"{V5_NETWORKS_PREFIX}/global/networks/mas-test"
 V5E_SUBNETWORKS = f"{V5_NETWORKS_PREFIX}/regions/us-east1/subnetworks/mas-test"
 V5P_SUBNETWORKS = f"{V5_NETWORKS_PREFIX}/regions/us-east5/subnetworks/mas-test"
+V5P_SUBNETWORKS_V2 = (
+    f"{V5_NETWORKS_PREFIX}/regions/europe-west4/subnetworks/mas-test-v2"
+)
 V6E_SUBNETWORKS = (
     f"{V5_NETWORKS_PREFIX}/regions/us-central2/subnetworks/mas-test"
 )

dags/multipod/jax_functional_tests.py

@@ -17,7 +17,7 @@
 import datetime
 from airflow import models
 from dags import composer_env
-from dags.common.vm_resource import DockerImage, TpuVersion, Zone, Project, V5_NETWORKS, V5P_SUBNETWORKS, RuntimeVersion, XpkClusters
+from dags.common.vm_resource import DockerImage, TpuVersion, Zone, Project, V5_NETWORKS, V5P_SUBNETWORKS_V2, RuntimeVersion, XpkClusters
 from dags.multipod.configs import jax_tests_gce_config, jax_tests_gke_config
 from dags.multipod.configs.common import SetupMode
 
@@ -34,7 +34,7 @@
   default_test_name = "jax-distributed-initialize"
   v5p_project_name = Project.TPU_PROD_ENV_AUTOMATED.value
   v5p_network = V5_NETWORKS
-  v5p_subnetwork = V5P_SUBNETWORKS
+  v5p_subnetwork = V5P_SUBNETWORKS_V2
   v5p_runtime_version = RuntimeVersion.V2_ALPHA_TPUV5.value
   test_modes_with_docker_images = [
       (SetupMode.STABLE, None),
@@ -50,7 +50,7 @@
       jax_gce_v4_8 = jax_tests_gce_config.get_jax_distributed_initialize_config(
           tpu_version=TpuVersion.V4,
           tpu_cores=8,
-          tpu_zone=Zone.US_CENTRAL2_B.value,
+          tpu_zone=Zone.EUROPE_WEST4_B.value,
           time_out_in_min=60,
           is_tpu_reserved=False,
           num_slices=num_slices,