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 1 commit
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';

// 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';

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';

// 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';

async function callCreateTpuVM() {
// Create a node
Expand Down
4 changes: 2 additions & 2 deletions tpu/queuedResources/createQueuedResource.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,11 @@ async function main(tpuClient) {
// 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';

// 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';

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';

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';

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';

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';

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';
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';
const projectId = 'project_id';
let tpuClientMock;

Expand Down
Loading