Fix samples and tests

Author: Joanna Grycz

tpu/queuedResource/createQueuedResource.js

@@ -96,7 +96,7 @@ async function main(
         ],
       },
       // TODO(developer): Uncomment next line if you want to specify reservation.
-      // reservationName: 'reservation-name'
+      // reservationName: 'reservation-name/ Before deleting the queued resource it is required to delete the TPU VM.'
     });
 
     const request = {
@@ -110,25 +110,10 @@ async function main(
     // Wait for the create operation to complete.
     await operation.promise();
 
+    // If you also want to wait for create operation of TPU Node,
+    // you can use `tpu_vm_get` sample to check current status of the node
+    // and wait until it is READY.
     console.log(`Queued resource ${queuedResourceName} created.`);
-
-    const getNodeRequest = {
-      name: `projects/${projectId}/locations/${zone}/nodes/${nodeName}`,
-    };
-
-    console.log(`Waiting for TPU node ${nodeName} to become ready...`);
-
-    // Poll for TPU node state every 30 seconds
-    const intervalId = setInterval(async () => {
-      const [node] = await tpuClient.getNode(getNodeRequest);
-
-      if (node.state === 'READY') {
-        clearInterval(intervalId);
-        console.log(`TPU node ${nodeName} is ready.`);
-      } else {
-        console.log(`TPU node ${nodeName} is in state: ${node.state}`);
-      }
-    }, 30000);
   }
   await callCreateQueuedResource();
   // [END tpu_queued_resources_create]

tpu/test/queuedResource.test.js

@@ -20,26 +20,54 @@ const path = require('path');
 const assert = require('node:assert/strict');
 const {describe, it} = require('mocha');
 const cp = require('child_process');
+const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
 
 const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
 const cwd = path.join(__dirname, '..');
 
-describe('TPU queued resource', async () => {
+async function waitForTPUCreation(nodeName, zone) {
+  const tpuClient = new TpuClient();
+  const projectId = await tpuClient.getProjectId();
+
+  // Give a time to start process of creating TPU Node
+  await new Promise(resolve => setTimeout(resolve, 60000));
+
+  const getNodeRequest = {
+    name: `projects/${projectId}/locations/${zone}/nodes/${nodeName}`,
+  };
+
+  console.log(`Waiting for TPU node ${nodeName} to become ready...`);
+
+  let state;
+
+  while (state !== 'READY') {
+    try {
+      state = (await tpuClient.getNode(getNodeRequest))[0].state;
+      // Wait another minute to try
+      await new Promise(resolve => setTimeout(resolve, 60000));
+    } catch (err) {
+      console.log('TPU node not ready');
+      // Wait another minute to try
+      await new Promise(resolve => setTimeout(resolve, 60000));
+    }
+  }
+}
+
+describe('TPU queued resource', () => {
   const queuedResourceName = `queued-resource-name-1a2sdf${Math.floor(Math.random() * 1000 + 1)}`;
   const nodePrefix = 'node-name-2a2b3c';
   const nodeName = `${nodePrefix}${Math.floor(Math.random() * 1000 + 1)}`;
-  const zone = 'europe-west4-a';
+  const zone = 'us-central1-f';
   const tpuType = 'v2-8';
   const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
 
-  it('should create a new queued resource', () => {
+  it('should create a new queued resource', async () => {
     const response = execSync(
       `node ./queuedResource/createQueuedResource.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
       {
         cwd,
       }
     );
-
     assert(response.includes(`Queued resource ${queuedResourceName} created.`));
   });
 
@@ -50,20 +78,20 @@ describe('TPU queued resource', async () => {
         cwd,
       }
     );
-
     assert(
       response.includes(`Queued resource ${queuedResourceName} retrived.`)
     );
   });
 
-  it('should delete queued resource', () => {
+  it('should delete queued resource', async () => {
+    // Wait until queued resource is ready to delete.
+    await waitForTPUCreation(nodeName, zone);
     const response = execSync(
       `node ./queuedResource/deleteQueuedResource.js ${queuedResourceName} ${zone}`,
       {
         cwd,
       }
     );
-
     assert(response.includes(`Queued resource ${queuedResourceName} deleted.`));
   });
 });