Skip to content

Commit c2c14fb

Browse files
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>
1 parent 4a947e6 commit c2c14fb

17 files changed

+59
-59
lines changed

tpu/create_tpu.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ def create_cloud_tpu(
2020
project_id: str,
2121
zone: str,
2222
tpu_name: str,
23-
tpu_type: str = "v2-8",
24-
runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
23+
tpu_type: str = "v5litepod-4",
24+
runtime_version: str = "v2-tpuv5-litepod",
2525
) -> Node:
2626
"""Creates a Cloud TPU node.
2727
Args:
@@ -38,10 +38,10 @@ def create_cloud_tpu(
3838

3939
# TODO(developer): Update and un-comment below lines
4040
# project_id = "your-project-id"
41-
# zone = "us-central1-b"
41+
# zone = "us-central1-a"
4242
# tpu_name = "tpu-name"
43-
# tpu_type = "v2-8"
44-
# runtime_version = "tpu-vm-tf-2.17.0-pjrt"
43+
# tpu_type = "v5litepod-4"
44+
# runtime_version = "v2-tpuv5-litepod"
4545

4646
# Create a TPU node
4747
node = tpu_v2.Node()
@@ -65,7 +65,7 @@ def create_cloud_tpu(
6565
print(response)
6666
# Example response:
6767
# name: "projects/[project_id]/locations/[zone]/nodes/my-tpu"
68-
# accelerator_type: "v2-8"
68+
# accelerator_type: "v5litepod-4"
6969
# state: READY
7070
# ...
7171

@@ -75,5 +75,5 @@ def create_cloud_tpu(
7575

7676
if __name__ == "__main__":
7777
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
78-
ZONE = "us-central1-b"
78+
ZONE = "us-central1-a"
7979
create_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")

tpu/create_tpu_spot.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ def create_tpu_with_spot(
2020
project_id: str,
2121
zone: str,
2222
tpu_name: str,
23-
tpu_type: str = "v2-8",
24-
runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
23+
tpu_type: str = "v5litepod-4",
24+
runtime_version: str = "v2-tpuv5-litepod",
2525
) -> Node:
2626
"""Creates a Cloud TPU node.
2727
Args:
@@ -40,8 +40,8 @@ def create_tpu_with_spot(
4040
# project_id = "your-project-id"
4141
# zone = "us-central1-b"
4242
# tpu_name = "tpu-name"
43-
# tpu_type = "v2-8"
44-
# runtime_version = "tpu-vm-tf-2.17.0-pjrt"
43+
# tpu_type = "v5litepod-4"
44+
# runtime_version = "v2-tpuv5-litepod"
4545

4646
# Create a TPU node
4747
node = tpu_v2.Node()
@@ -75,5 +75,5 @@ def create_tpu_with_spot(
7575

7676
if __name__ == "__main__":
7777
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
78-
ZONE = "us-central1-b"
78+
ZONE = "us-central1-a"
7979
create_tpu_with_spot(PROJECT_ID, ZONE, "tpu-with-spot")

tpu/create_tpu_with_script.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,18 +20,18 @@ def create_cloud_tpu_with_script(
2020
project_id: str,
2121
zone: str,
2222
tpu_name: str,
23-
tpu_type: str = "v2-8",
24-
runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
23+
tpu_type: str = "v5litepod-4",
24+
runtime_version: str = "v2-tpuv5-litepod",
2525
) -> Node:
2626
# [START tpu_vm_create_startup_script]
2727
from google.cloud import tpu_v2
2828

2929
# TODO(developer): Update and un-comment below lines
3030
# project_id = "your-project-id"
31-
# zone = "us-central1-b"
31+
# zone = "us-central1-a"
3232
# tpu_name = "tpu-name"
33-
# tpu_type = "v2-8"
34-
# runtime_version = "tpu-vm-tf-2.17.0-pjrt"
33+
# tpu_type = "v5litepod-4"
34+
# runtime_version = "v2-tpuv5-litepod"
3535

3636
node = tpu_v2.Node()
3737
node.accelerator_type = tpu_type
@@ -72,5 +72,5 @@ def create_cloud_tpu_with_script(
7272

7373
if __name__ == "__main__":
7474
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
75-
ZONE = "us-central1-b"
75+
ZONE = "us-central1-a"
7676
create_cloud_tpu_with_script(PROJECT_ID, ZONE, "tpu-name")

tpu/delete_tpu.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,5 +44,5 @@ def delete_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") ->
4444

4545
if __name__ == "__main__":
4646
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
47-
ZONE = "us-central1-b"
47+
ZONE = "us-central1-a"
4848
delete_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")

tpu/get_tpu.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,5 +50,5 @@ def get_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> Nod
5050

5151
if __name__ == "__main__":
5252
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
53-
ZONE = "us-central1-b"
53+
ZONE = "us-central1-a"
5454
get_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")

tpu/list_tpu.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,5 +50,5 @@ def list_cloud_tpu(project_id: str, zone: str) -> ListNodesPager:
5050

5151
if __name__ == "__main__":
5252
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
53-
ZONE = "us-central1-b"
53+
ZONE = "us-central1-a"
5454
list_cloud_tpu(PROJECT_ID, ZONE)

tpu/queued_resources_create.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,19 @@ def create_queued_resource(
2020
project_id: str,
2121
zone: str,
2222
tpu_name: str,
23-
tpu_type: str = "v2-8",
24-
runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
23+
tpu_type: str = "v5litepod-4",
24+
runtime_version: str = "v2-tpuv5-litepod",
2525
queued_resource_name: str = "resource-name",
2626
) -> Node:
2727
# [START tpu_queued_resources_create]
2828
from google.cloud import tpu_v2alpha1
2929

3030
# TODO(developer): Update and un-comment below lines
3131
# project_id = "your-project-id"
32-
# zone = "us-central1-b"
32+
# zone = "us-central1-a"
3333
# tpu_name = "tpu-name"
34-
# tpu_type = "v2-8"
35-
# runtime_version = "tpu-vm-tf-2.17.0-pjrt"
34+
# tpu_type = "v5litepod-4"
35+
# runtime_version = "v2-tpuv5-litepod"
3636
# queued_resource_name = "resource-name"
3737

3838
node = tpu_v2alpha1.Node()
@@ -71,7 +71,7 @@ def create_queued_resource(
7171

7272
if __name__ == "__main__":
7373
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
74-
ZONE = "us-central1-b"
74+
ZONE = "us-central1-a"
7575
create_queued_resource(
7676
project_id=PROJECT_ID,
7777
zone=ZONE,

tpu/queued_resources_create_network.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ def create_queued_resource_network(
2020
project_id: str,
2121
zone: str,
2222
tpu_name: str,
23-
tpu_type: str = "v2-8",
24-
runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
23+
tpu_type: str = "v5litepod-4",
24+
runtime_version: str = "v2-tpuv5-litepod",
2525
queued_resource_name: str = "resource-name",
2626
network: str = "default",
2727
) -> Node:
@@ -30,10 +30,10 @@ def create_queued_resource_network(
3030

3131
# TODO(developer): Update and un-comment below lines
3232
# project_id = "your-project-id"
33-
# zone = "us-central1-b"
33+
# zone = "us-central1-a"
3434
# tpu_name = "tpu-name"
35-
# tpu_type = "v2-8"
36-
# runtime_version = "tpu-vm-tf-2.17.0-pjrt"
35+
# tpu_type = "v5litepod-4"
36+
# runtime_version = "v2-tpuv5-litepod"
3737
# queued_resource_name = "resource-name"
3838
# network = "default"
3939

@@ -81,7 +81,7 @@ def create_queued_resource_network(
8181

8282
if __name__ == "__main__":
8383
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
84-
ZONE = "us-central1-b"
84+
ZONE = "us-central1-a"
8585
create_queued_resource_network(
8686
project_id=PROJECT_ID,
8787
zone=ZONE,

tpu/queued_resources_create_spot.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,19 @@ def create_queued_resource_spot(
2020
project_id: str,
2121
zone: str,
2222
tpu_name: str,
23-
tpu_type: str = "v2-8",
24-
runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
23+
tpu_type: str = "v5litepod-4",
24+
runtime_version: str = "v2-tpuv5-litepod",
2525
queued_resource_name: str = "resource-name",
2626
) -> Node:
2727
# [START tpu_queued_resources_create_spot]
2828
from google.cloud import tpu_v2alpha1
2929

3030
# TODO(developer): Update and un-comment below lines
3131
# project_id = "your-project-id"
32-
# zone = "us-central1-b"
32+
# zone = "us-central1-a"
3333
# tpu_name = "tpu-name"
34-
# tpu_type = "v2-8"
35-
# runtime_version = "tpu-vm-tf-2.17.0-pjrt"
34+
# tpu_type = "v5litepod-4"
35+
# runtime_version = "v2-tpuv5-litepod"
3636
# queued_resource_name = "resource-name"
3737

3838
node = tpu_v2alpha1.Node()
@@ -73,7 +73,7 @@ def create_queued_resource_spot(
7373

7474
if __name__ == "__main__":
7575
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
76-
ZONE = "us-central1-b"
76+
ZONE = "us-central1-a"
7777
create_queued_resource_spot(
7878
project_id=PROJECT_ID,
7979
zone=ZONE,

tpu/queued_resources_create_startup_script.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,19 @@ def create_queued_resource_startup_script(
2020
project_id: str,
2121
zone: str,
2222
tpu_name: str,
23-
tpu_type: str = "v2-8",
24-
runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
23+
tpu_type: str = "v5litepod-4",
24+
runtime_version: str = "v2-tpuv5-litepod",
2525
queued_resource_name: str = "resource-name",
2626
) -> Node:
2727
# [START tpu_queued_resources_startup_script]
2828
from google.cloud import tpu_v2alpha1
2929

3030
# TODO(developer): Update and un-comment below lines
3131
# project_id = "your-project-id"
32-
# zone = "us-central1-b"
32+
# zone = "us-central1-a"
3333
# tpu_name = "tpu-name"
34-
# tpu_type = "v2-8"
35-
# runtime_version = "tpu-vm-tf-2.17.0-pjrt"
34+
# tpu_type = "v5litepod-4"
35+
# runtime_version = "v2-tpuv5-litepod"
3636
# queued_resource_name = "resource-name"
3737

3838
node = tpu_v2alpha1.Node()
@@ -84,7 +84,7 @@ def create_queued_resource_startup_script(
8484

8585
if __name__ == "__main__":
8686
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
87-
ZONE = "us-central1-b"
87+
ZONE = "us-central1-a"
8888
create_queued_resource_startup_script(
8989
project_id=PROJECT_ID,
9090
zone=ZONE,

0 commit comments

Comments
 (0)