diff --git a/tpu/package.json b/tpu/package.json index 81edf2e8829..9006d2f0209 100644 --- a/tpu/package.json +++ b/tpu/package.json @@ -3,21 +3,22 @@ "license": "Apache-2.0", "author": "Google Inc.", "engines": { - "node": ">=16.0.0" + "node": ">=16.0.0" }, "repository": "googleapis/nodejs-tpu", "private": true, "files": [ - "*.js" + "*.js" ], "scripts": { - "test": "c8 mocha -p -j 2 test --timeout 1200000" + "test": "c8 mocha -p -j 2 test --timeout 1200000" }, "dependencies": { - "@google-cloud/tpu": "^3.5.0" + "@google-cloud/tpu": "^3.5.0", + "sinon": "^19.0.2" }, "devDependencies": { - "c8": "^10.0.0", - "mocha": "^10.0.0" + "c8": "^10.0.0", + "mocha": "^10.0.0" } - } \ No newline at end of file +} diff --git a/tpu/queuedResources/createQueuedResourceNetwork.js b/tpu/queuedResources/createQueuedResourceNetwork.js index 1ec9a9b7edf..6c9dbcbe320 100644 --- a/tpu/queuedResources/createQueuedResourceNetwork.js +++ b/tpu/queuedResources/createQueuedResourceNetwork.js @@ -16,21 +16,17 @@ 'use strict'; -async function main( - nodeName, - queuedResourceName, - zone, - tpuType, - tpuSoftwareVersion -) { +async function main(tpuClient) { // [START tpu_queued_resources_create_network] - // Import the TPU library - const {TpuClient} = require('@google-cloud/tpu').v2alpha1; + // Import the TPUClient + // TODO(developer): Uncomment below line before running the sample. + // const {TpuClient} = require('@google-cloud/tpu').v2alpha1; const {Node, NetworkConfig, QueuedResource} = require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1; // Instantiate a tpuClient - const tpuClient = new TpuClient(); + // TODO(developer): Uncomment below line before running the sample. + // tpuClient = new TpuClient(); /** * TODO(developer): Update/uncomment these variables before running the sample. @@ -45,24 +41,24 @@ async function main( const region = 'europe-west4'; // The name for your queued resource. - // queuedResourceName = 'queued-resource-1'; + const queuedResourceName = 'queued-resource-1'; // The name for your node. - // nodeName = 'node-name-1'; + const nodeName = 'node-name-1'; // 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 - // zone = 'europe-west4-a'; + const zone = 'europe-west4-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. - // tpuType = 'v2-8'; + const tpuType = 'v2-8'; // Software version that specifies the version of the node runtime to install. For more information, // see https://cloud.google.com/tpu/docs/runtimes - // tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; + const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; async function callCreateQueuedResourceNetwork() { // Specify the network and subnetwork that you want to connect your TPU to. @@ -112,13 +108,16 @@ async function main( // You can wait until TPU Node is READY, // and check its status using getTpuVm() from `tpu_vm_get` sample. - console.log(JSON.stringify(response)); + return response; } - await callCreateQueuedResourceNetwork(); + return await callCreateQueuedResourceNetwork(); // [END tpu_queued_resources_create_network] } -main(...process.argv.slice(2)).catch(err => { - console.error(err); - process.exitCode = 1; -}); +module.exports = main; + +// TODO(developer): Uncomment below lines before running the sample. +// main(...process.argv.slice(2)).catch(err => { +// console.error(err); +// process.exitCode = 1; +// }); diff --git a/tpu/queuedResources/createQueuedResourceStartupScript.js b/tpu/queuedResources/createQueuedResourceStartupScript.js index ce0c2080caf..f7509a5a41c 100644 --- a/tpu/queuedResources/createQueuedResourceStartupScript.js +++ b/tpu/queuedResources/createQueuedResourceStartupScript.js @@ -16,21 +16,17 @@ 'use strict'; -async function main( - nodeName, - queuedResourceName, - zone, - tpuType, - tpuSoftwareVersion -) { +async function main(tpuClient) { // [START tpu_queued_resources_startup_script] - // Import the TPU library - const {TpuClient} = require('@google-cloud/tpu').v2alpha1; + // Import the TPUClient + // TODO(developer): Uncomment below line before running the sample. + // const {TpuClient} = require('@google-cloud/tpu').v2alpha1; const {Node, NetworkConfig, QueuedResource} = require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1; // Instantiate a tpuClient - const tpuClient = new TpuClient(); + // TODO(developer): Uncomment below line before running the sample. + // tpuClient = new TpuClient(); /** * TODO(developer): Update/uncomment these variables before running the sample. @@ -45,24 +41,24 @@ async function main( const region = 'europe-west4'; // The name for your queued resource. - // queuedResourceName = 'queued-resource-1'; + const queuedResourceName = 'queued-resource-1'; // The name for your node. - // nodeName = 'node-name-1'; + const nodeName = 'node-name-1'; // 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 - // zone = 'europe-west4-a'; + const zone = 'europe-west4-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. - // tpuType = 'v2-8'; + const tpuType = 'v2-8'; // Software version that specifies the version of the node runtime to install. For more information, // see https://cloud.google.com/tpu/docs/runtimes - // tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; + const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; async function callCreateQueuedResourceStartupScript() { // Create a node @@ -116,13 +112,16 @@ async function main( // You can wait until TPU Node is READY, // and check its status using getTpuVm() from `tpu_vm_get` sample. - console.log(JSON.stringify(response)); + return response; } - await callCreateQueuedResourceStartupScript(); + return await callCreateQueuedResourceStartupScript(); // [END tpu_queued_resources_startup_script] } -main(...process.argv.slice(2)).catch(err => { - console.error(err); - process.exitCode = 1; -}); +module.exports = main; + +// TODO(developer): Uncomment below lines before running the sample. +// main(...process.argv.slice(2)).catch(err => { +// console.error(err); +// process.exitCode = 1; +// }); diff --git a/tpu/queuedResources/createQueuedResourceTimeBound.js b/tpu/queuedResources/createQueuedResourceTimeBound.js index 92e1889c2da..206568f1756 100644 --- a/tpu/queuedResources/createQueuedResourceTimeBound.js +++ b/tpu/queuedResources/createQueuedResourceTimeBound.js @@ -16,21 +16,17 @@ 'use strict'; -async function main( - nodeName, - queuedResourceName, - zone, - tpuType, - tpuSoftwareVersion -) { +async function main(tpuClient) { // [START tpu_queued_resources_time_bound] - // Import the TPU library - const {TpuClient} = require('@google-cloud/tpu').v2alpha1; + // Import the TPUClient + // TODO(developer): Uncomment below line before running the sample. + // const {TpuClient} = require('@google-cloud/tpu').v2alpha1; const {Node, NetworkConfig, QueuedResource} = require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1; // Instantiate a tpuClient - const tpuClient = new TpuClient(); + // TODO(developer): Uncomment below line before running the sample. + // tpuClient = new TpuClient(); /** * TODO(developer): Update/uncomment these variables before running the sample. @@ -45,24 +41,24 @@ async function main( const region = 'europe-west4'; // The name for your queued resource. - // queuedResourceName = 'queued-resource-1'; + const queuedResourceName = 'queued-resource-1'; // The name for your node. - // nodeName = 'node-name-1'; + const nodeName = 'node-name-1'; // 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 - // zone = 'europe-west4-a'; + const zone = 'europe-west4-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. - // tpuType = 'v2-8'; + const tpuType = 'v2-8'; // Software version that specifies the version of the node runtime to install. For more information, // see https://cloud.google.com/tpu/docs/runtimes - // tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; + const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; async function callCreateQueuedResourceTimeBound() { // Create a node @@ -143,13 +139,16 @@ async function main( // You can wait until TPU Node is READY, // and check its status using getTpuVm() from `tpu_vm_get` sample. - console.log(JSON.stringify(response)); + return response; } - await callCreateQueuedResourceTimeBound(); + return await callCreateQueuedResourceTimeBound(); // [END tpu_queued_resources_time_bound] } -main(...process.argv.slice(2)).catch(err => { - console.error(err); - process.exitCode = 1; -}); +module.exports = main; + +// TODO(developer): Uncomment below lines before running the sample. +// main(...process.argv.slice(2)).catch(err => { +// console.error(err); +// process.exitCode = 1; +// }); diff --git a/tpu/queuedResources/forceDeleteQueuedResource.js b/tpu/queuedResources/forceDeleteQueuedResource.js deleted file mode 100644 index e842c39e269..00000000000 --- a/tpu/queuedResources/forceDeleteQueuedResource.js +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright 2024 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -'use strict'; - -async function main(queuedResourceName, zone) { - // [START tpu_queued_resources_delete_force] - // Import the TPU library - const {TpuClient} = require('@google-cloud/tpu').v2alpha1; - - // Instantiate a tpuClient - const tpuClient = new TpuClient(); - - /** - * TODO(developer): Update/uncomment these variables before running the sample. - */ - // Project ID or project number of the Google Cloud project, where you want to delete node. - const projectId = await tpuClient.getProjectId(); - - // The name of queued resource. - // queuedResourceName = 'queued-resource-1'; - - // The zone of your queued resource. - // zone = 'europe-west4-a'; - - async function callForceDeleteQueuedResource() { - const request = { - name: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`, - force: true, - }; - - const [operation] = await tpuClient.deleteQueuedResource(request); - - // Wait for the delete operation to complete. - await operation.promise(); - - console.log(`Queued resource ${queuedResourceName} deletion forced.`); - } - await callForceDeleteQueuedResource(); - // [END tpu_queued_resources_delete_force] -} - -main(...process.argv.slice(2)).catch(err => { - console.error(err); - process.exitCode = 1; -}); diff --git a/tpu/test/createQueuedResourceNetwork.test.js b/tpu/test/createQueuedResourceNetwork.test.js index 898ffe9be15..840e1afa86e 100644 --- a/tpu/test/createQueuedResourceNetwork.test.js +++ b/tpu/test/createQueuedResourceNetwork.test.js @@ -16,57 +16,66 @@ 'use strict'; -const path = require('path'); const assert = require('node:assert/strict'); -const {after, before, 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, '..'); +const {beforeEach, afterEach, describe, it} = require('mocha'); +const sinon = require('sinon'); +const createQueuedResourceNetwork = require('../queuedResources/createQueuedResourceNetwork.js'); describe('TPU queued resource with specified network', async () => { - const queuedResourceName = `queued-resource-with-network-${Math.floor(Math.random() * 1000 + 1)}`; - const nodeName = `node-with-network-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`; - const zone = 'us-south1-a'; - const tpuType = 'v5litepod-1'; - const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; - let projectId; + const queuedResourceName = 'queued-resource-1'; + const nodeName = 'node-name-1'; + const zone = 'europe-west4-a'; + const projectId = 'project_id'; + let tpuClientMock; - before(async () => { - const tpuClient = new TpuClient(); - projectId = await tpuClient.getProjectId(); + beforeEach(() => { + tpuClientMock = { + getProjectId: sinon.stub().resolves(projectId), + }; }); - after(() => { - // Delete queued resource - execSync( - `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`, - { - cwd, - } - ); + afterEach(() => { + sinon.restore(); }); - it('should create queued resource', () => { - const networkConfig = { - network: `projects/${projectId}/global/networks/compute-tpu-network`, - subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`, - enableExternalIps: true, - }; + it('should create queued resource', async () => { + tpuClientMock.createQueuedResource = sinon.stub().resolves([ + { + promise: sinon.stub().resolves([ + { + name: queuedResourceName, + }, + ]), + }, + ]); - const response = JSON.parse( - execSync( - `node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, - { - cwd, - } - ) - ); + const response = await createQueuedResourceNetwork(tpuClientMock); - assert.deepEqual( - response.tpu.nodeSpec[0].node.networkConfig, - networkConfig + sinon.assert.calledWith( + tpuClientMock.createQueuedResource, + sinon.match({ + parent: `projects/${projectId}/locations/${zone}`, + queuedResource: { + name: queuedResourceName, + tpu: { + nodeSpec: [ + sinon.match({ + parent: `projects/${projectId}/locations/${zone}`, + node: { + networkConfig: { + network: `projects/${projectId}/global/networks/compute-tpu-network`, + subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`, + enableExternalIps: true, + }, + }, + nodeId: nodeName, + }), + ], + }, + }, + queuedResourceId: queuedResourceName, + }) ); + assert(response.name.includes(queuedResourceName)); }); }); diff --git a/tpu/test/createQueuedResourceStartupScript.test.js b/tpu/test/createQueuedResourceStartupScript.test.js index abf368f467b..b7132418081 100644 --- a/tpu/test/createQueuedResourceStartupScript.test.js +++ b/tpu/test/createQueuedResourceStartupScript.test.js @@ -16,46 +16,62 @@ 'use strict'; -const path = require('path'); const assert = require('node:assert/strict'); -const {after, describe, it} = require('mocha'); -const cp = require('child_process'); - -const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'}); -const cwd = path.join(__dirname, '..'); +const {beforeEach, afterEach, describe, it} = require('mocha'); +const sinon = require('sinon'); +const createQueuedResourceStartupScript = require('../queuedResources/createQueuedResourceStartupScript.js'); describe('TPU queued resource with start-up script', async () => { - const queuedResourceName = `queued-resource-startup-script-${Math.floor(Math.random() * 1000 + 1)}`; - const nodeName = `node-startup-script-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`; - const zone = 'us-east1-d'; - const tpuType = 'v3-32'; - const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; - - after(() => { - // Delete queued resource - execSync( - `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`, - { - cwd, - } - ); - }); + const queuedResourceName = 'queued-resource-1'; + const nodeName = 'node-name-1'; + const zone = 'europe-west4-a'; + const projectId = 'project_id'; + let tpuClientMock; - it('should create queued resource', () => { - const metadata = { - 'startup-script': - '#!/bin/bash\n echo "Hello World" > /var/log/hello.log\n sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1', + beforeEach(() => { + tpuClientMock = { + getProjectId: sinon.stub().resolves(projectId), }; + }); + + afterEach(() => { + sinon.restore(); + }); + + it('should create queued resource', async () => { + tpuClientMock.createQueuedResource = sinon.stub().resolves([ + { + promise: sinon.stub().resolves([ + { + name: queuedResourceName, + }, + ]), + }, + ]); - const response = JSON.parse( - execSync( - `node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, - { - cwd, - } - ) - ); + const response = await createQueuedResourceStartupScript(tpuClientMock); - assert.deepEqual(response.tpu.nodeSpec[0].node.metadata, metadata); + sinon.match({ + parent: `projects/${projectId}/locations/${zone}`, + queuedResource: { + name: queuedResourceName, + tpu: { + nodeSpec: [ + { + parent: `projects/${projectId}/locations/${zone}`, + node: { + metadata: { + 'startup-script': + '#!/bin/bash\n echo "Hello World" > /var/log/hello.log\n sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1', + }, + }, + nodeId: nodeName, + }, + ], + }, + }, + queuedResourceId: queuedResourceName, + }); + assert(response.name.includes(queuedResourceName)); }); }); diff --git a/tpu/test/createQueuedResourceTimeBound.test.js b/tpu/test/createQueuedResourceTimeBound.test.js index 2a70a2dcd3a..0e4d15fee05 100644 --- a/tpu/test/createQueuedResourceTimeBound.test.js +++ b/tpu/test/createQueuedResourceTimeBound.test.js @@ -16,44 +16,54 @@ 'use strict'; -const path = require('path'); const assert = require('node:assert/strict'); -const {after, describe, it} = require('mocha'); -const cp = require('child_process'); - -const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'}); -const cwd = path.join(__dirname, '..'); +const {beforeEach, afterEach, describe, it} = require('mocha'); +const sinon = require('sinon'); +const createQueuedResourceTimeBound = require('../queuedResources/createQueuedResourceTimeBound.js'); describe('TPU time bound queued resource', async () => { - const queuedResourceName = `queued-resource-time-bound-${Math.floor(Math.random() * 1000 + 1)}`; - const nodeName = `node-time-bound-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`; - const zone = 'us-west4-a'; - const tpuType = 'v5litepod-1'; - const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; - - after(() => { - // Delete queued resource - execSync( - `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`, - { - cwd, - } - ); + const queuedResourceName = 'queued-resource-1'; + const zone = 'europe-west4-a'; + const projectId = 'project_id'; + let tpuClientMock; + + beforeEach(() => { + tpuClientMock = { + getProjectId: sinon.stub().resolves(projectId), + }; }); - it('should create queued resource', () => { - const response = JSON.parse( - execSync( - `node ./queuedResources/createQueuedResourceTimeBound.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, - { - cwd, - } - ) - ); + afterEach(() => { + sinon.restore(); + }); + + it('should create queued resource', async () => { + tpuClientMock.createQueuedResource = sinon.stub().resolves([ + { + promise: sinon.stub().resolves([ + { + name: queuedResourceName, + }, + ]), + }, + ]); - assert.ok(response.queueingPolicy); - assert.ok(response.queueingPolicy.validAfterTime); - assert(typeof response.queueingPolicy.validAfterTime.seconds, 'string'); - assert(typeof response.queueingPolicy.validAfterTime.nano, 'number'); + const response = await createQueuedResourceTimeBound(tpuClientMock); + + sinon.assert.calledWith( + tpuClientMock.createQueuedResource, + sinon.match({ + parent: `projects/${projectId}/locations/${zone}`, + queuedResource: { + queueingPolicy: { + validAfterDuration: { + seconds: 6 * 3600, + }, + }, + }, + queuedResourceId: 'queued-resource-1', + }) + ); + assert(response.name.includes(queuedResourceName)); }); });