diff --git a/tpu/queuedResource/createQueuedResource.js b/tpu/queuedResource/createQueuedResource.js index b2d3bc99803..bb5e991dc96 100644 --- a/tpu/queuedResource/createQueuedResource.js +++ b/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] diff --git a/tpu/test/queuedResource.test.js b/tpu/test/queuedResource.test.js index 1e242344227..0303acebcfa 100644 --- a/tpu/test/queuedResource.test.js +++ b/tpu/test/queuedResource.test.js @@ -20,19 +20,44 @@ 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; + } catch (err) { + console.log('TPU node not ready'); + } + } +} + +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-c'; 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}`, { @@ -56,7 +81,10 @@ describe('TPU queued resource', async () => { ); }); - 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}`, {