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

* feat(tpu): Updating TPU samples to use TPUv5e

* Fixing test

---------

Co-authored-by: Katie McLaughlin <[email protected]>

Author: m-strzelczyk

tpu/createStartupScriptVM.js

@@ -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

tpu/createTopologyVM.js

@@ -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

tpu/createVM.js

@@ -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

tpu/queuedResources/createQueuedResource.js

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

tpu/queuedResources/deleteQueuedResource.js

@@ -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 = {

tpu/queuedResources/forceDeleteQueuedResource.js

@@ -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 = {

tpu/queuedResources/getQueuedResource.js

@@ -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 = {

tpu/queuedResources/getQueuedResourcesList.js

@@ -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 = {

tpu/test/forceDeleteQueuedResource.test.js

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

tpu/test/queuedResource.test.js

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