From 78c99126656caddd71006ab37a728b99743d1f72 Mon Sep 17 00:00:00 2001 From: Joanna Grycz Date: Wed, 23 Oct 2024 13:57:24 +0200 Subject: [PATCH 1/6] feat: tpu_queued_resources_startup_script --- CODEOWNERS | 1 + .../createQueuedResourceStartupScript.js | 130 ++++++++++++++++++ .../forceDeleteQueuedResource.js | 59 ++++++++ .../createQueuedResourceStartupScript.test.js | 58 ++++++++ 4 files changed, 248 insertions(+) create mode 100644 tpu/queuedResources/createQueuedResourceStartupScript.js create mode 100644 tpu/queuedResources/forceDeleteQueuedResource.js create mode 100644 tpu/test/createQueuedResourceStartupScript.test.js diff --git a/CODEOWNERS b/CODEOWNERS index d679277d5d..c7d66c17cf 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -20,6 +20,7 @@ compute @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/nodejs-samples-revie iam @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/nodejs-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers kms @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/nodejs-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers orgpolicy @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/nodejs-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers +tpu @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/nodejs-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers recaptcha_enterprise @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/nodejs-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers recaptcha_enterprise/demosite @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/recaptcha-customer-obsession-reviewers @GoogleCloudPlatform/nodejs-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers secret-manager @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/nodejs-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers @GoogleCloudPlatform/cloud-secrets-team diff --git a/tpu/queuedResources/createQueuedResourceStartupScript.js b/tpu/queuedResources/createQueuedResourceStartupScript.js new file mode 100644 index 0000000000..4ee2decb20 --- /dev/null +++ b/tpu/queuedResources/createQueuedResourceStartupScript.js @@ -0,0 +1,130 @@ +/* + * 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( + nodeName, + queuedResourceName, + zone, + tpuType, + tpuSoftwareVersion +) { + // [START tpu_queued_resources_startup_script] + // Import the TPU library + 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): Update/uncomment these variables before running the sample. + */ + // Project ID or project number of the Google Cloud project, where you want to create queued resource. + const projectId = await tpuClient.getProjectId(); + + // The name of the network you want the node to connect to. The network should be assigned to your project. + const networkName = 'compute-tpu-network'; + + // The region of the network, that you want the node to connect to. + const region = 'europe-west4'; + + // The name for your queued resource. + // queuedResourceName = 'queued-resource-1'; + + // The name for your node. + // 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'; + + // 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'; + + // 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'; + + async function callCreateQueuedResourceStartupScript() { + // Create a node + const node = new Node({ + name: nodeName, + zone, + acceleratorType: tpuType, + runtimeVersion: tpuSoftwareVersion, + // Define network + networkConfig: new NetworkConfig({ + enableExternalIps: true, + network: `projects/${projectId}/global/networks/${networkName}`, + subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`, + }), + queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`, + metadata: { + // The script updates numpy to the latest version and logs the output to a file. + 'startup-script': `#!/bin/bash + echo "Hello World" > /var/log/hello.log + sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1`, + }, + }); + + // Define parent for requests + const parent = `projects/${projectId}/locations/${zone}`; + + // Create queued resource + const queuedResource = new QueuedResource({ + name: queuedResourceName, + tpu: { + nodeSpec: [ + { + parent, + node, + nodeId: nodeName, + }, + ], + }, + }); + + const request = { + parent: `projects/${projectId}/locations/${zone}`, + queuedResource, + queuedResourceId: queuedResourceName, + }; + + const [operation] = await tpuClient.createQueuedResource(request); + + // Wait for the create operation to complete. + await operation.promise(); + + // You can wait until TPU Node is READY, + // and check its status using getTpuVm() from `tpu_vm_get` sample. + console.log( + `Queued resource ${queuedResourceName} with start-up script created.` + ); + } + await callCreateQueuedResourceStartupScript(); + // [END tpu_queued_resources_startup_script] +} + +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 new file mode 100644 index 0000000000..e842c39e26 --- /dev/null +++ b/tpu/queuedResources/forceDeleteQueuedResource.js @@ -0,0 +1,59 @@ +/* + * 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/createQueuedResourceStartupScript.test.js b/tpu/test/createQueuedResourceStartupScript.test.js new file mode 100644 index 0000000000..d82f770474 --- /dev/null +++ b/tpu/test/createQueuedResourceStartupScript.test.js @@ -0,0 +1,58 @@ +/* + * 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'; + +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, '..'); + +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, + } + ); + }); + + it('should create queued resource with start-up script', () => { + const response = execSync( + `node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, + { + cwd, + } + ); + + assert( + response.includes( + `Queued resource ${queuedResourceName} with start-up script created.` + ) + ); + }); +}); From 50cc2e8e1b2859dbadce2a9e79fec803a88b62d8 Mon Sep 17 00:00:00 2001 From: Joanna Grycz Date: Thu, 24 Oct 2024 13:44:32 +0200 Subject: [PATCH 2/6] feat: tpu_queued_resources_create_network --- .../createQueuedResourceNetwork.js | 127 ++++++++++++++++++ .../createQueuedResourceStartupScript.js | 7 +- tpu/test/createQueuedResourceNetwork.test.js | 72 ++++++++++ .../createQueuedResourceStartupScript.test.js | 6 + 4 files changed, 209 insertions(+), 3 deletions(-) create mode 100644 tpu/queuedResources/createQueuedResourceNetwork.js create mode 100644 tpu/test/createQueuedResourceNetwork.test.js diff --git a/tpu/queuedResources/createQueuedResourceNetwork.js b/tpu/queuedResources/createQueuedResourceNetwork.js new file mode 100644 index 0000000000..1a7a1f5a67 --- /dev/null +++ b/tpu/queuedResources/createQueuedResourceNetwork.js @@ -0,0 +1,127 @@ +/* + * 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( + nodeName, + queuedResourceName, + zone, + tpuType, + tpuSoftwareVersion +) { + // [START tpu_queued_resources_create_network] + // Import the TPU library + 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): Update/uncomment these variables before running the sample. + */ + // Project ID or project number of the Google Cloud project, where you want to create queued resource. + const projectId = await tpuClient.getProjectId(); + + // The name of the network you want the node to connect to. The network should be assigned to your project. + const networkName = 'compute-tpu-network'; + + // The region of the network, that you want the node to connect to. + const region = 'europe-west4'; + + // The name for your queued resource. + // queuedResourceName = 'queued-resource-1'; + + // The name for your node. + // 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'; + + // 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'; + + // 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'; + + async function callCreateQueuedResourceNetwork() { + // Specify the network and subnetwork that you want to connect your TPU to. + const networkConfig = new NetworkConfig({ + enableExternalIps: true, + network: `projects/${projectId}/global/networks/${networkName}`, + subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`, + }); + + // Create a node + const node = new Node({ + name: nodeName, + zone, + acceleratorType: tpuType, + runtimeVersion: tpuSoftwareVersion, + networkConfig, + queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`, + }); + + // Define parent for requests + const parent = `projects/${projectId}/locations/${zone}`; + + // Create queued resource + const queuedResource = new QueuedResource({ + name: queuedResourceName, + tpu: { + nodeSpec: [ + { + parent, + node, + nodeId: nodeName, + }, + ], + }, + }); + + const request = { + parent: `projects/${projectId}/locations/${zone}`, + queuedResource, + queuedResourceId: queuedResourceName, + }; + + const [operation] = await tpuClient.createQueuedResource(request); + + // Wait for the create operation to complete. + const [response] = await operation.promise(); + + // You can wait until TPU Node is READY, + // and check its status using getTpuVm() from `tpu_vm_get` sample. + console.log( + `Queued resource ${queuedResourceName} with specified network created.` + ); + console.log(JSON.stringify(response)); + } + await callCreateQueuedResourceNetwork(); + // [END tpu_queued_resources_create_network] +} + +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 4ee2decb20..b47fb8b506 100644 --- a/tpu/queuedResources/createQueuedResourceStartupScript.js +++ b/tpu/queuedResources/createQueuedResourceStartupScript.js @@ -81,8 +81,8 @@ async function main( metadata: { // The script updates numpy to the latest version and logs the output to a file. 'startup-script': `#!/bin/bash - echo "Hello World" > /var/log/hello.log - sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1`, + echo "Hello World" > /var/log/hello.log + sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1`, }, }); @@ -112,13 +112,14 @@ async function main( const [operation] = await tpuClient.createQueuedResource(request); // Wait for the create operation to complete. - await operation.promise(); + const [response] = await operation.promise(); // You can wait until TPU Node is READY, // and check its status using getTpuVm() from `tpu_vm_get` sample. console.log( `Queued resource ${queuedResourceName} with start-up script created.` ); + console.log(JSON.stringify(response)); } await callCreateQueuedResourceStartupScript(); // [END tpu_queued_resources_startup_script] diff --git a/tpu/test/createQueuedResourceNetwork.test.js b/tpu/test/createQueuedResourceNetwork.test.js new file mode 100644 index 0000000000..60045994c7 --- /dev/null +++ b/tpu/test/createQueuedResourceNetwork.test.js @@ -0,0 +1,72 @@ +/* + * 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'; + +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, '..'); + +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; + + before(async () => { + const tpuClient = new TpuClient(); + projectId = await tpuClient.getProjectId(); + }); + + after(() => { + // Delete queued resource + execSync( + `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`, + { + cwd, + } + ); + }); + + it('should create queued resource with specified network', () => { + const networkConfig = { + network: `projects/${projectId}/global/networks/compute-tpu-network`, + subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`, + enableExternalIps: true, + }; + + const response = execSync( + `node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, + { + cwd, + } + ); + + assert( + response.includes( + `Queued resource ${queuedResourceName} with specified network created.` + ) + ); + assert(response.includes(JSON.stringify(networkConfig))); + }); +}); diff --git a/tpu/test/createQueuedResourceStartupScript.test.js b/tpu/test/createQueuedResourceStartupScript.test.js index d82f770474..22235eb3c5 100644 --- a/tpu/test/createQueuedResourceStartupScript.test.js +++ b/tpu/test/createQueuedResourceStartupScript.test.js @@ -42,6 +42,11 @@ describe('TPU queued resource with start-up script', async () => { }); it('should create queued resource with start-up script', () => { + 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', + }; + const response = execSync( `node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, { @@ -54,5 +59,6 @@ describe('TPU queued resource with start-up script', async () => { `Queued resource ${queuedResourceName} with start-up script created.` ) ); + assert(response.includes(JSON.stringify(metadata))); }); }); From c9929032b7c0ba03dded18e1e4aa052f90446dc8 Mon Sep 17 00:00:00 2001 From: Joanna Grycz Date: Fri, 25 Oct 2024 10:24:11 +0200 Subject: [PATCH 3/6] feat: tpu_queued_resources_time_bound --- .../createQueuedResourceNetwork.js | 3 - .../createQueuedResourceStartupScript.js | 3 - .../createQueuedResourceTimeBound.js | 155 ++++++++++++++++++ tpu/test/createQueuedResourceNetwork.test.js | 22 +-- .../createQueuedResourceStartupScript.test.js | 21 +-- .../createQueuedResourceTimeBound.test.js | 59 +++++++ 6 files changed, 234 insertions(+), 29 deletions(-) create mode 100644 tpu/queuedResources/createQueuedResourceTimeBound.js create mode 100644 tpu/test/createQueuedResourceTimeBound.test.js diff --git a/tpu/queuedResources/createQueuedResourceNetwork.js b/tpu/queuedResources/createQueuedResourceNetwork.js index 1a7a1f5a67..1ec9a9b7ed 100644 --- a/tpu/queuedResources/createQueuedResourceNetwork.js +++ b/tpu/queuedResources/createQueuedResourceNetwork.js @@ -112,9 +112,6 @@ async function main( // You can wait until TPU Node is READY, // and check its status using getTpuVm() from `tpu_vm_get` sample. - console.log( - `Queued resource ${queuedResourceName} with specified network created.` - ); console.log(JSON.stringify(response)); } await callCreateQueuedResourceNetwork(); diff --git a/tpu/queuedResources/createQueuedResourceStartupScript.js b/tpu/queuedResources/createQueuedResourceStartupScript.js index b47fb8b506..ce0c2080ca 100644 --- a/tpu/queuedResources/createQueuedResourceStartupScript.js +++ b/tpu/queuedResources/createQueuedResourceStartupScript.js @@ -116,9 +116,6 @@ async function main( // You can wait until TPU Node is READY, // and check its status using getTpuVm() from `tpu_vm_get` sample. - console.log( - `Queued resource ${queuedResourceName} with start-up script created.` - ); console.log(JSON.stringify(response)); } await callCreateQueuedResourceStartupScript(); diff --git a/tpu/queuedResources/createQueuedResourceTimeBound.js b/tpu/queuedResources/createQueuedResourceTimeBound.js new file mode 100644 index 0000000000..92e1889c2d --- /dev/null +++ b/tpu/queuedResources/createQueuedResourceTimeBound.js @@ -0,0 +1,155 @@ +/* + * 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( + nodeName, + queuedResourceName, + zone, + tpuType, + tpuSoftwareVersion +) { + // [START tpu_queued_resources_time_bound] + // Import the TPU library + 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): Update/uncomment these variables before running the sample. + */ + // Project ID or project number of the Google Cloud project, where you want to create queued resource. + const projectId = await tpuClient.getProjectId(); + + // The name of the network you want the node to connect to. The network should be assigned to your project. + const networkName = 'compute-tpu-network'; + + // The region of the network, that you want the node to connect to. + const region = 'europe-west4'; + + // The name for your queued resource. + // queuedResourceName = 'queued-resource-1'; + + // The name for your node. + // 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'; + + // 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'; + + // 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'; + + async function callCreateQueuedResourceTimeBound() { + // Create a node + const node = new Node({ + name: nodeName, + zone, + acceleratorType: tpuType, + runtimeVersion: tpuSoftwareVersion, + // Define network + networkConfig: new NetworkConfig({ + enableExternalIps: true, + network: `projects/${projectId}/global/networks/${networkName}`, + subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`, + }), + queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`, + }); + + // Define parent for requests + const parent = `projects/${projectId}/locations/${zone}`; + + // Create queued resource + const queuedResource = new QueuedResource({ + name: queuedResourceName, + tpu: { + nodeSpec: [ + { + parent, + node, + nodeId: nodeName, + }, + ], + }, + queueingPolicy: new QueuedResource.QueueingPolicy({ + // You can specify a duration after which a resource should be allocated. + validAfterDuration: { + // format: hour * 3600s + seconds: 6 * 3600, + }, + // You can specify how long a queued resource request remains valid. + // validUntilDuration: { + // // format: hour * 3600s + // seconds: 6 * 3600, + // }, + // You can specify a time after which a resource should be allocated. + // validAfterTime: { + // // format: new Date('YOUR_TIMESTAMP').getTime() / 1000 + // seconds: new Date('2024-10-25T11:45:00Z').getTime() / 1000, + // }, + // You can specify a time before which the resource should be allocated. + // validUntilTime: { + // // format: new Date('YOUR_TIMESTAMP').getTime() / 1000 + // seconds: new Date('2024-10-25T11:45:00Z').getTime() / 1000, + // }, + // You can specify an allocation interval. `startTime` specifies the beginning of the allocation interval + // and `endTime` specifies the end of the allocation interval. + // validInterval: { + // // format: new Date('YOUR_TIMESTAMP').getTime() / 1000 + // startTime: { + // seconds: new Date('2024-10-25T15:45:00Z').getTime() / 1000, + // }, + // endTime: { + // seconds: new Date('2024-10-26T11:45:00Z').getTime() / 1000, + // }, + // }, + }), + }); + + const request = { + parent: `projects/${projectId}/locations/${zone}`, + queuedResource, + queuedResourceId: queuedResourceName, + }; + + const [operation] = await tpuClient.createQueuedResource(request); + + // Wait for the create operation to complete. + const [response] = await operation.promise(); + + // You can wait until TPU Node is READY, + // and check its status using getTpuVm() from `tpu_vm_get` sample. + console.log(JSON.stringify(response)); + } + await callCreateQueuedResourceTimeBound(); + // [END tpu_queued_resources_time_bound] +} + +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 60045994c7..898ffe9be1 100644 --- a/tpu/test/createQueuedResourceNetwork.test.js +++ b/tpu/test/createQueuedResourceNetwork.test.js @@ -48,25 +48,25 @@ describe('TPU queued resource with specified network', async () => { ); }); - it('should create queued resource with specified network', () => { + 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, }; - const response = execSync( - `node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, - { - cwd, - } + const response = JSON.parse( + execSync( + `node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, + { + cwd, + } + ) ); - assert( - response.includes( - `Queued resource ${queuedResourceName} with specified network created.` - ) + assert.deepEqual( + response.tpu.nodeSpec[0].node.networkConfig, + networkConfig ); - assert(response.includes(JSON.stringify(networkConfig))); }); }); diff --git a/tpu/test/createQueuedResourceStartupScript.test.js b/tpu/test/createQueuedResourceStartupScript.test.js index 22235eb3c5..abf368f467 100644 --- a/tpu/test/createQueuedResourceStartupScript.test.js +++ b/tpu/test/createQueuedResourceStartupScript.test.js @@ -41,24 +41,21 @@ describe('TPU queued resource with start-up script', async () => { ); }); - it('should create queued resource with start-up script', () => { + 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', }; - const response = execSync( - `node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, - { - cwd, - } - ); - - assert( - response.includes( - `Queued resource ${queuedResourceName} with start-up script created.` + const response = JSON.parse( + execSync( + `node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, + { + cwd, + } ) ); - assert(response.includes(JSON.stringify(metadata))); + + assert.deepEqual(response.tpu.nodeSpec[0].node.metadata, metadata); }); }); diff --git a/tpu/test/createQueuedResourceTimeBound.test.js b/tpu/test/createQueuedResourceTimeBound.test.js new file mode 100644 index 0000000000..2a70a2dcd3 --- /dev/null +++ b/tpu/test/createQueuedResourceTimeBound.test.js @@ -0,0 +1,59 @@ +/* + * 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'; + +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, '..'); + +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, + } + ); + }); + + it('should create queued resource', () => { + const response = JSON.parse( + execSync( + `node ./queuedResources/createQueuedResourceTimeBound.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`, + { + cwd, + } + ) + ); + + assert.ok(response.queueingPolicy); + assert.ok(response.queueingPolicy.validAfterTime); + assert(typeof response.queueingPolicy.validAfterTime.seconds, 'string'); + assert(typeof response.queueingPolicy.validAfterTime.nano, 'number'); + }); +}); From a9ee1f7ab5cd8f1417285dcff8a97ed615d01c47 Mon Sep 17 00:00:00 2001 From: Joanna Grycz Date: Tue, 12 Nov 2024 18:05:46 +0100 Subject: [PATCH 4/6] Use mocked TPUClient in tests --- .../createQueuedResourceNetwork.js | 47 +++++----- .../createQueuedResourceStartupScript.js | 43 +++++---- .../createQueuedResourceTimeBound.js | 43 +++++---- .../forceDeleteQueuedResource.js | 59 ------------ tpu/test/createQueuedResourceNetwork.test.js | 91 ++++++++++--------- .../createQueuedResourceStartupScript.test.js | 84 ++++++++++------- .../createQueuedResourceTimeBound.test.js | 76 +++++++++------- 7 files changed, 208 insertions(+), 235 deletions(-) delete mode 100644 tpu/queuedResources/forceDeleteQueuedResource.js diff --git a/tpu/queuedResources/createQueuedResourceNetwork.js b/tpu/queuedResources/createQueuedResourceNetwork.js index 1ec9a9b7ed..2b36291d3a 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 -) { - // [START tpu_queued_resources_create_network] - // Import the TPU library - const {TpuClient} = require('@google-cloud/tpu').v2alpha1; +async function main(tpuClient) { + // [START tpu_queued_resources_network] + // 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. @@ -42,27 +38,27 @@ async function main( const networkName = 'compute-tpu-network'; // The region of the network, that you want the node to connect to. - const region = 'europe-west4'; + const region = 'us-central1'; // 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 = `${zone}-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(); - // [END tpu_queued_resources_create_network] + return await callCreateQueuedResourceNetwork(); + // [END tpu_queued_resources_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 ce0c2080ca..7e5b499c1f 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. @@ -42,27 +38,27 @@ async function main( const networkName = 'compute-tpu-network'; // The region of the network, that you want the node to connect to. - const region = 'europe-west4'; + const region = 'us-central1'; // 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 = `${zone}-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 92e1889c2d..535d8d5803 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. @@ -42,27 +38,27 @@ async function main( const networkName = 'compute-tpu-network'; // The region of the network, that you want the node to connect to. - const region = 'europe-west4'; + const region = 'us-central1'; // 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 = `${zone}-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 e842c39e26..0000000000 --- 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 898ffe9be1..3c18558794 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 = 'us-central1-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 abf368f467..60c587e303 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 = 'us-central1-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 2a70a2dcd3..43b6dc8d2f 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 = 'us-central1-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)); }); }); From b8b5fa5b9932a114d9d6fc866a5cd4ed834f6913 Mon Sep 17 00:00:00 2001 From: BigBlackWolf Date: Mon, 9 Dec 2024 17:53:51 +0100 Subject: [PATCH 5/6] correct zone definition --- tpu/queuedResources/createQueuedResourceNetwork.js | 2 +- tpu/queuedResources/createQueuedResourceStartupScript.js | 2 +- tpu/queuedResources/createQueuedResourceTimeBound.js | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tpu/queuedResources/createQueuedResourceNetwork.js b/tpu/queuedResources/createQueuedResourceNetwork.js index 2b36291d3a..150392c743 100644 --- a/tpu/queuedResources/createQueuedResourceNetwork.js +++ b/tpu/queuedResources/createQueuedResourceNetwork.js @@ -49,7 +49,7 @@ 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 = `${zone}-a`; + 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, diff --git a/tpu/queuedResources/createQueuedResourceStartupScript.js b/tpu/queuedResources/createQueuedResourceStartupScript.js index 7e5b499c1f..51c55225fa 100644 --- a/tpu/queuedResources/createQueuedResourceStartupScript.js +++ b/tpu/queuedResources/createQueuedResourceStartupScript.js @@ -49,7 +49,7 @@ 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 = `${zone}-a`; + 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, diff --git a/tpu/queuedResources/createQueuedResourceTimeBound.js b/tpu/queuedResources/createQueuedResourceTimeBound.js index 535d8d5803..ce2c9a276b 100644 --- a/tpu/queuedResources/createQueuedResourceTimeBound.js +++ b/tpu/queuedResources/createQueuedResourceTimeBound.js @@ -49,7 +49,7 @@ 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 = `${zone}-a`; + 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, From 954e4ec75b1f7be6413148f43fdfab996f592aaa Mon Sep 17 00:00:00 2001 From: BigBlackWolf Date: Tue, 10 Dec 2024 16:33:22 +0100 Subject: [PATCH 6/6] correct the region network in test --- tpu/test/createQueuedResourceNetwork.test.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tpu/test/createQueuedResourceNetwork.test.js b/tpu/test/createQueuedResourceNetwork.test.js index 3c18558794..df460f13c4 100644 --- a/tpu/test/createQueuedResourceNetwork.test.js +++ b/tpu/test/createQueuedResourceNetwork.test.js @@ -64,7 +64,7 @@ describe('TPU queued resource with specified network', async () => { node: { networkConfig: { network: `projects/${projectId}/global/networks/compute-tpu-network`, - subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`, + subnetwork: `projects/${projectId}/regions/us-central1/subnetworks/compute-tpu-network`, enableExternalIps: true, }, },