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

Jump to bottom

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

merged 2 commits into from
Mar 11, 2025

Conversation

m-strzelczyk
Copy link
Contributor

Description

Fixes b/400926465

Updating the samples to use TPU v5e, as requested by the TPU team.

@product-auto-label product-auto-label bot added api: tpu Issues related to the Cloud TPU API. samples Issues that are directly related to samples. labels Mar 6, 2025
gemini-code-assist[bot]
gemini-code-assist bot reviewed Mar 6, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Hello @m-strzelczyk, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

This pull request updates the TPU samples to feature TPU v5e, as requested by the TPU team. The changes involve updating the tpu_type and runtime_version parameters in several sample files to reflect the new TPU version. Specifically, the tpu_type is changed from v2-8 to v5litepod-4, and the runtime_version is changed from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod. Additionally, the default zone is changed from us-central1-b to us-central1-a in all the samples.

Highlights

  • TPU Version Update: The samples are updated to use TPU v5e instead of v2-8.
  • Runtime Version Update: The runtime version is updated to v2-tpuv5-litepod.
  • Zone Update: The default zone is updated to us-central1-a.

Changelog

Click here to see the changelog
  • tpu/create_tpu.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 41 and 78).
  • tpu/create_tpu_spot.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 41 and 78).
  • tpu/create_tpu_with_script.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 31 and 75).
  • tpu/delete_tpu.py
    • Updated default zone from us-central1-b to us-central1-a (line 47).
  • tpu/get_tpu.py
    • Updated default zone from us-central1-b to us-central1-a (line 53).
  • tpu/list_tpu.py
    • Updated default zone from us-central1-b to us-central1-a (line 53).
  • tpu/queued_resources_create.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 32 and 74).
  • tpu/queued_resources_create_network.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 33 and 84).
  • tpu/queued_resources_create_spot.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 32 and 76).
  • tpu/queued_resources_create_startup_script.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 32 and 87).
  • tpu/queued_resources_create_time_bound.py
    • Updated tpu_type from v2-8 to v5litepod-4 (line 23).
    • Updated runtime_version from tpu-vm-tf-2.17.0-pjrt to v2-tpuv5-litepod (line 24).
    • Updated default zone from us-central1-b to us-central1-a (line 32 and 84).
  • tpu/queued_resources_delete.py
    • Updated default zone from us-central1-b to us-central1-a (line 46).
  • tpu/queued_resources_delete_force.py
    • Updated default zone from us-central1-b to us-central1-a (line 47).
  • tpu/queued_resources_get.py
    • Updated default zone from us-central1-b to us-central1-a (line 47).
  • tpu/queued_resources_list.py
    • Updated default zone from us-central1-b to us-central1-a (line 43).
  • tpu/start_tpu.py
    • Updated default zone from us-central1-b to us-central1-a (line 34 and 60).
  • tpu/stop_tpu.py
    • Updated default zone from us-central1-b to us-central1-a (line 35 and 61).
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in issue comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Did you know?

The first TPU was internally deployed in Google's data centers in 2015.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist[bot]
gemini-code-assist bot reviewed Mar 6, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

The pull request updates the TPU samples to feature TPUv5e. The changes involve updating the default TPU type and runtime version in the sample code. The changes look good overall, but I have a few suggestions.

Summary of Findings

  • Zone consistency: The zone is updated to us-central1-a in most files, but the create_tpu_with_script.py file still contains the old zone in the comments. This inconsistency should be addressed.
  • TODO comment in create_tpu_spot.py: The create_tpu_spot.py file contains a TODO comment regarding the update of the library to change preemptible to spot=True. It would be good to address this if possible, or at least provide more context on why it's still a TODO.

Merge Readiness

The pull request is almost ready for merging. Please address the zone inconsistency in create_tpu_with_script.py and consider providing more context for the TODO comment in create_tpu_spot.py. I am unable to directly approve this pull request, and recommend that others review and approve this code before merging. Given the medium severity comment, I recommend that the pull request not be merged until it is addressed.

tpu/create_tpu.py
@@ -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/queued_resources_create.py
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/queued_resources_create_network.py
@@ -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/queued_resources_create_spot.py
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/queued_resources_create_startup_script.py
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/queued_resources_create_time_bound.py
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/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"
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/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"
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.

@m-strzelczyk @gemini-code-assist
Update tpu/create_tpu_with_script.py
272d891 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@m-strzelczyk m-strzelczyk marked this pull request as ready for review March 6, 2025 15:50
@m-strzelczyk m-strzelczyk requested review from a team as code owners March 6, 2025 15:50
@m-strzelczyk m-strzelczyk merged commit c2c14fb into main Mar 11, 2025
13 of 14 checks passed
@m-strzelczyk m-strzelczyk deleted the strzelczyk/tpu-v5-update branch March 11, 2025 12:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

@kweinmeister kweinmeister kweinmeister approved these changes

Assignees

@Sita04 Sita04

Labels
api: tpu Issues related to the Cloud TPU API. samples Issues that are directly related to samples.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@m-strzelczyk @kweinmeister @Sita04