feat(tpu): Updating TPU samples to feature TPUv5e (#13211)

* feat(tpu): Updating TPU samples to feature TPUv5e * Update tpu/create_tpu_with_script.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
GoogleCloudPlatform · Mar 11, 2025 · c2c14fb · c2c14fb
1 parent 4a947e6
commit c2c14fb
Show file tree

Hide file tree

Showing 17 changed files with 59 additions and 59 deletions.
diff --git a/tpu/create_tpu.py b/tpu/create_tpu.py
@@ -20,8 +20,8 @@ def create_cloud_tpu(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
 ) -> Node:
     """Creates a Cloud TPU node.
     Args:
@@ -38,10 +38,10 @@ def create_cloud_tpu(
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
 
     # Create a TPU node
     node = tpu_v2.Node()
@@ -65,7 +65,7 @@ def create_cloud_tpu(
     print(response)
     # Example response:
     # name: "projects/[project_id]/locations/[zone]/nodes/my-tpu"
-    # accelerator_type: "v2-8"
+    # accelerator_type: "v5litepod-4"
     # state: READY
     # ...
 
@@ -75,5 +75,5 @@ def create_cloud_tpu(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")
diff --git a/tpu/create_tpu_spot.py b/tpu/create_tpu_spot.py
@@ -20,8 +20,8 @@ def create_tpu_with_spot(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
 ) -> Node:
     """Creates a Cloud TPU node.
     Args:
@@ -40,8 +40,8 @@ def create_tpu_with_spot(
     # project_id = "your-project-id"
     # zone = "us-central1-b"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
 
     # Create a TPU node
     node = tpu_v2.Node()
@@ -75,5 +75,5 @@ def create_tpu_with_spot(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_tpu_with_spot(PROJECT_ID, ZONE, "tpu-with-spot")
diff --git a/tpu/create_tpu_with_script.py b/tpu/create_tpu_with_script.py
@@ -20,18 +20,18 @@ def create_cloud_tpu_with_script(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
 ) -> Node:
     # [START tpu_vm_create_startup_script]
     from google.cloud import tpu_v2
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
 
     node = tpu_v2.Node()
     node.accelerator_type = tpu_type
@@ -72,5 +72,5 @@ def create_cloud_tpu_with_script(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_cloud_tpu_with_script(PROJECT_ID, ZONE, "tpu-name")
diff --git a/tpu/delete_tpu.py b/tpu/delete_tpu.py
@@ -44,5 +44,5 @@ def delete_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") ->
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     delete_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")
diff --git a/tpu/get_tpu.py b/tpu/get_tpu.py
@@ -50,5 +50,5 @@ def get_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> Nod
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     get_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")
diff --git a/tpu/list_tpu.py b/tpu/list_tpu.py
@@ -50,5 +50,5 @@ def list_cloud_tpu(project_id: str, zone: str) -> ListNodesPager:
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     list_cloud_tpu(PROJECT_ID, ZONE)
diff --git a/tpu/queued_resources_create.py b/tpu/queued_resources_create.py
@@ -20,19 +20,19 @@ def create_queued_resource(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
     queued_resource_name: str = "resource-name",
 ) -> Node:
     # [START tpu_queued_resources_create]
     from google.cloud import tpu_v2alpha1
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
     # queued_resource_name = "resource-name"
 
     node = tpu_v2alpha1.Node()
@@ -71,7 +71,7 @@ def create_queued_resource(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_queued_resource(
         project_id=PROJECT_ID,
         zone=ZONE,

diff --git a/tpu/queued_resources_create_network.py b/tpu/queued_resources_create_network.py
@@ -20,8 +20,8 @@ def create_queued_resource_network(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
     queued_resource_name: str = "resource-name",
     network: str = "default",
 ) -> Node:
@@ -30,10 +30,10 @@ def create_queued_resource_network(
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
     # queued_resource_name = "resource-name"
     # network = "default"
 
@@ -81,7 +81,7 @@ def create_queued_resource_network(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_queued_resource_network(
         project_id=PROJECT_ID,
         zone=ZONE,

diff --git a/tpu/queued_resources_create_spot.py b/tpu/queued_resources_create_spot.py
@@ -20,19 +20,19 @@ def create_queued_resource_spot(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
     queued_resource_name: str = "resource-name",
 ) -> Node:
     # [START tpu_queued_resources_create_spot]
     from google.cloud import tpu_v2alpha1
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
     # queued_resource_name = "resource-name"
 
     node = tpu_v2alpha1.Node()
@@ -73,7 +73,7 @@ def create_queued_resource_spot(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_queued_resource_spot(
         project_id=PROJECT_ID,
         zone=ZONE,

diff --git a/tpu/queued_resources_create_startup_script.py b/tpu/queued_resources_create_startup_script.py
@@ -20,19 +20,19 @@ def create_queued_resource_startup_script(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
     queued_resource_name: str = "resource-name",
 ) -> Node:
     # [START tpu_queued_resources_startup_script]
     from google.cloud import tpu_v2alpha1
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
     # queued_resource_name = "resource-name"
 
     node = tpu_v2alpha1.Node()
@@ -84,7 +84,7 @@ def create_queued_resource_startup_script(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_queued_resource_startup_script(
         project_id=PROJECT_ID,
         zone=ZONE,

diff --git a/tpu/queued_resources_create_time_bound.py b/tpu/queued_resources_create_time_bound.py
@@ -20,19 +20,19 @@ def create_queued_resource_time_bound(
     project_id: str,
     zone: str,
     tpu_name: str,
-    tpu_type: str = "v2-8",
-    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+    tpu_type: str = "v5litepod-4",
+    runtime_version: str = "v2-tpuv5-litepod",
     queued_resource_name: str = "resource-name",
 ) -> Node:
     # [START tpu_queued_resources_time_bound]
     from google.cloud import tpu_v2alpha1
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
-    # tpu_type = "v2-8"
-    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+    # tpu_type = "v5litepod-4"
+    # runtime_version = "v2-tpuv5-litepod"
     # queued_resource_name = "resource-name"
 
     node = tpu_v2alpha1.Node()
@@ -81,7 +81,7 @@ def create_queued_resource_time_bound(
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     create_queued_resource_time_bound(
         project_id=PROJECT_ID,
         zone=ZONE,

diff --git a/tpu/queued_resources_delete.py b/tpu/queued_resources_delete.py
@@ -43,5 +43,5 @@ def delete_queued_resource(
 
 if __name__ == "__main__":
     PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     delete_queued_resource(PROJECT_ID, ZONE, "resource-name")
diff --git a/tpu/queued_resources_delete_force.py b/tpu/queued_resources_delete_force.py
@@ -44,5 +44,5 @@ def delete_force_queued_resource(
 
 if __name__ == "__main__":
     PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     delete_force_queued_resource(PROJECT_ID, ZONE, "resource-name")
diff --git a/tpu/queued_resources_get.py b/tpu/queued_resources_get.py
@@ -44,5 +44,5 @@ def get_queued_resource(
 
 if __name__ == "__main__":
     PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     get_queued_resource(PROJECT_ID, ZONE, "resource-name")
diff --git a/tpu/queued_resources_list.py b/tpu/queued_resources_list.py
@@ -40,5 +40,5 @@ def list_queued_resources(project_id: str, zone: str) -> ListQueuedResourcesPage
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     list_queued_resources(PROJECT_ID, ZONE)
diff --git a/tpu/start_tpu.py b/tpu/start_tpu.py
@@ -31,7 +31,7 @@ def start_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> N
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
 
     client = tpu_v2.TpuClient()
@@ -57,5 +57,5 @@ def start_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> N
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     start_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")
diff --git a/tpu/stop_tpu.py b/tpu/stop_tpu.py
@@ -32,7 +32,7 @@ def stop_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> No
 
     # TODO(developer): Update and un-comment below lines
     # project_id = "your-project-id"
-    # zone = "us-central1-b"
+    # zone = "us-central1-a"
     # tpu_name = "tpu-name"
 
     client = tpu_v2.TpuClient()
@@ -58,5 +58,5 @@ def stop_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> No
 
 if __name__ == "__main__":
     PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
-    ZONE = "us-central1-b"
+    ZONE = "us-central1-a"
     stop_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")