Skip to content

Commit

Permalink
Fix samples and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
gryczj committed Oct 22, 2024
1 parent bf5bab3 commit 39771e6
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 26 deletions.
23 changes: 4 additions & 19 deletions tpu/queuedResource/createQueuedResource.js
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand All @@ -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]
Expand Down
42 changes: 35 additions & 7 deletions tpu/test/queuedResource.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -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.`));
});

Expand All @@ -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.`));
});
});

0 comments on commit 39771e6

Please sign in to comment.