Skip to content

feat(tpu): Updating TPU samples to use TPUv5e #4028

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

Merged
glasnt merged 3 commits into from
Mar 14, 2025
Merged

feat(tpu): Updating TPU samples to use TPUv5e #4028

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
86ca966
feat(tpu): Updating TPU samples to use TPUv5e
m-strzelczyk Mar 13, 2025
0a8424d
Fixing test
m-strzelczyk Mar 13, 2025
37e5f34
Merge branch 'main' into strzelczyk/tpu-update
glasnt Mar 14, 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
4 changes: 2 additions & 2 deletions tpu/createStartupScriptVM.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,11 @@ async function main(tpuClient) {
// The accelerator type that specifies the version and size of the Cloud TPU you want to create.
// For more information about supported accelerator types for each TPU version,
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
const tpuType = 'v2-8';
const tpuType = 'v5litepod-4';
Comment thread
m-strzelczyk marked this conversation as resolved.

// Software version that specifies the version of the TPU runtime to install. For more information,
// see https://cloud.google.com/tpu/docs/runtimes
const tpuSoftwareVersion = 'tpu-vm-tf-2.17.0-pod-pjrt';
const tpuSoftwareVersion = 'v2-tpuv5-litepod';
Comment thread
m-strzelczyk marked this conversation as resolved.

async function callCreateTpuVMStartupScript() {
// Create a node
Expand Down
2 changes: 1 addition & 1 deletion tpu/createTopologyVM.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ async function main(tpuClient) {

// Software version that specifies the version of the TPU runtime to install. For more information,
// see https://cloud.google.com/tpu/docs/runtimes
const tpuSoftwareVersion = 'tpu-vm-tf-2.17.0-pod-pjrt';
const tpuSoftwareVersion = 'v2-tpuv5-litepod';

// The version of the Cloud TPU you want to create.
// Available options: TYPE_UNSPECIFIED = 0, V2 = 2, V3 = 4, V4 = 7
Expand Down
4 changes: 2 additions & 2 deletions tpu/createVM.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,11 +49,11 @@ async function main(tpuClient) {
// The accelerator type that specifies the version and size of the Cloud TPU you want to create.
// For more information about supported accelerator types for each TPU version,
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
const tpuType = 'v2-8';
const tpuType = 'v5litepod-4';
Comment thread
m-strzelczyk marked this conversation as resolved.

// Software version that specifies the version of the TPU runtime to install. For more information,
// see https://cloud.google.com/tpu/docs/runtimes
const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
const tpuSoftwareVersion = 'v2-tpuv5-litepod';
Comment thread
m-strzelczyk marked this conversation as resolved.

async function callCreateTpuVM() {
// Create a node
Expand Down
6 changes: 3 additions & 3 deletions tpu/queuedResources/createQueuedResource.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,16 +49,16 @@ async function main(tpuClient) {
// The zone in which to create the node.
// For more information about supported TPU types for specific zones,
// see https://cloud.google.com/tpu/docs/regions-zones
const zone = `${region}-f`;
const zone = `${region}-a`;

// The accelerator type that specifies the version and size of the node you want to create.
// For more information about supported accelerator types for each TPU version,
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
const tpuType = 'v2-8';
const tpuType = 'v5litepod-4';
Comment thread
m-strzelczyk marked this conversation as resolved.

// Software version that specifies the version of the node runtime to install. For more information,
// see https://cloud.google.com/tpu/docs/runtimes
const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
const tpuSoftwareVersion = 'v2-tpuv5-litepod';
Comment thread
m-strzelczyk marked this conversation as resolved.

async function callCreateQueuedResource() {
// Create a node
Expand Down
2 changes: 1 addition & 1 deletion tpu/queuedResources/deleteQueuedResource.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ async function main(tpuClient) {
const queuedResourceName = 'queued-resource-1';

// The zone of your queued resource.
const zone = 'us-central1-f';
const zone = 'us-central1-a';
Comment thread
m-strzelczyk marked this conversation as resolved.

async function callDeleteTpuVM(nodeName) {
const request = {
Expand Down
2 changes: 1 addition & 1 deletion tpu/queuedResources/forceDeleteQueuedResource.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ async function main(tpuClient) {
const queuedResourceName = 'queued-resource-1';

// The zone of your queued resource.
const zone = 'us-central1-f';
const zone = 'us-central1-a';
Comment thread
m-strzelczyk marked this conversation as resolved.

async function callForceDeleteQueuedResource() {
const request = {
Expand Down
2 changes: 1 addition & 1 deletion tpu/queuedResources/getQueuedResource.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ async function main(tpuClient) {
const queuedResourceName = 'queued-resource-1';

// The zone of your queued resource.
const zone = 'us-central1-f';
const zone = 'us-central1-a';
Comment thread
m-strzelczyk marked this conversation as resolved.

async function callGetQueuedResource() {
const request = {
Expand Down
2 changes: 1 addition & 1 deletion tpu/queuedResources/getQueuedResourcesList.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ async function main(tpuClient) {
const projectId = await tpuClient.getProjectId();

// The zone from which the Queued Resources are retrived.
const zone = 'us-central1-f';
const zone = 'us-central1-a';
Comment thread
m-strzelczyk marked this conversation as resolved.

async function callGetQueuedResourcesList() {
const request = {
Expand Down
2 changes: 1 addition & 1 deletion tpu/test/forceDeleteQueuedResource.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ const forceDeleteQueuedResource = require('../queuedResources/forceDeleteQueuedR

describe('TPU queued resource force deletion', async () => {
const queuedResourceName = 'queued-resource-1';
const zone = 'us-central1-f';
const zone = 'us-central1-a';
Comment thread
m-strzelczyk marked this conversation as resolved.
const projectId = 'project_id';
let tpuClientMock;

Expand Down
2 changes: 1 addition & 1 deletion tpu/test/queuedResource.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ const deleteQueuedResource = require('../queuedResources/deleteQueuedResource.js
describe('TPU queued resource', () => {
const queuedResourceName = 'queued-resource-1';
const nodeName = 'node-name-1';
const zone = 'us-central1-f';
const zone = 'us-central1-a';
Comment thread
m-strzelczyk marked this conversation as resolved.
const projectId = 'project_id';
let tpuClientMock;

Expand Down