Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Mar 11, 2025
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f0a3479
feat(tpu): Updating TPU samples to feature TPUv5e
m-strzelczyk Mar 6, 2025
272d891
Update tpu/create_tpu_with_script.py
m-strzelczyk Mar 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions tpu/create_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# 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()
Expand All @@ -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
# ...

Expand All @@ -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")
10 changes: 5 additions & 5 deletions tpu/create_tpu_spot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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()
Expand Down Expand Up @@ -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")
10 changes: 5 additions & 5 deletions tpu/create_tpu_with_script.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ 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
Expand All @@ -30,8 +30,8 @@ def create_cloud_tpu_with_script(
# 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"

node = tpu_v2.Node()
node.accelerator_type = tpu_type
Expand Down Expand Up @@ -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")
2 changes: 1 addition & 1 deletion tpu/delete_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
2 changes: 1 addition & 1 deletion tpu/get_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
2 changes: 1 addition & 1 deletion tpu/list_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
12 changes: 6 additions & 6 deletions tpu/queued_resources_create.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# 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()
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 6 additions & 6 deletions tpu/queued_resources_create_network.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# 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"

Expand Down Expand Up @@ -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,
Expand Down
12 changes: 6 additions & 6 deletions tpu/queued_resources_create_spot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# 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()
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 6 additions & 6 deletions tpu/queued_resources_create_startup_script.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# 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()
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 6 additions & 6 deletions tpu/queued_resources_create_time_bound.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# 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()
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion tpu/queued_resources_delete.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
2 changes: 1 addition & 1 deletion tpu/queued_resources_delete_force.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
2 changes: 1 addition & 1 deletion tpu/queued_resources_get.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
2 changes: 1 addition & 1 deletion tpu/queued_resources_list.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
4 changes: 2 additions & 2 deletions tpu/start_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# tpu_name = "tpu-name"

client = tpu_v2.TpuClient()
Expand All @@ -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")
4 changes: 2 additions & 2 deletions tpu/stop_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider updating the zone in the comment to match the default value.

# tpu_name = "tpu-name"

client = tpu_v2.TpuClient()
Expand All @@ -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")