feat: tpu_vm_delete

Author: Joanna Grycz

compute/test/tpu.test.js

@@ -18,26 +18,19 @@
 
 const path = require('path');
 const assert = require('node:assert/strict');
-const {describe, xit} = require('mocha');
+const {describe, it} = require('mocha');
 const cp = require('child_process');
-// const {TpuClient} = require('@google-cloud/tpu').v2;
 
 const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
 const cwd = path.join(__dirname, '..');
 
 describe('Compute tpu', async () => {
   const nodeName = `node-name-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`;
-  const zone = 'us-central1-a';
-  const tpuType = 'v2-32';
-  const tpuSoftwareVersion = 'tpu-vm-base';
-  //   const tpuClient = new TpuClient();
-  //   let projectId;
+  const zone = 'us-central1-b';
+  const tpuType = 'v2-8';
+  const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
 
-  //   before(async () => {
-  //     projectId = await tpuClient.getProjectId();
-  //   });
-
-  xit('should create a new tpu node', () => {
+  it('should create a new tpu node', () => {
     const response = execSync(
       `node ./tpu/vmCreate.js ${nodeName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
       {
@@ -48,19 +41,29 @@ describe('Compute tpu', async () => {
     assert(response.includes(`TPU VM: ${nodeName} created.`));
   });
 
-  xit('should return tpu node', () => {
+  it('should return tpu node', () => {
     const response = execSync(`node ./tpu/vmGet.js ${nodeName} ${zone}`, {
       cwd,
     });
 
     assert(response.includes(`Node: ${nodeName} retrived.`));
   });
 
-  xit('should return list of tpu nodes', () => {
-    const response = execSync(`node ./tpu/vmList.js ${zone}`, {
+  it('should return list of tpu nodes', () => {
+    const response = JSON.parse(
+      execSync(`node ./tpu/vmList.js ${zone}`, {
+        cwd,
+      })
+    );
+
+    assert(Array.isArray(response));
+  });
+
+  it('should delete tpu node', () => {
+    const response = execSync(`node ./tpu/vmDelete.js ${nodeName} ${zone}`, {
       cwd,
     });
 
-    assert(Array.isArray(response));
+    assert(response.includes(`Node: ${nodeName} deleted.`));
   });
 });

compute/tpu/vmCreate.js

@@ -37,16 +37,16 @@ async function main(nodeName, zone, tpuType, tpuSoftwareVersion) {
   // The zone in which to create the TPU.
   // For more information about supported TPU types for specific zones,
   // see https://cloud.google.com/tpu/docs/regions-zones
-  // zone = 'us-central1-a';
+  // zone = 'us-central1-b';
 
   // The accelerator type that specifies the version and size of the Cloud TPU 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-32';
+  // tpuType = 'v2-8';
 
   // Software version that specifies the version of the TPU runtime to install. For more information,
   // see https://cloud.google.com/tpu/docs/runtimes
-  // tpuSoftwareVersion = 'tpu-vm-base';
+  // tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
 
   async function callCreateTpuVM() {
     // Create a node
@@ -55,7 +55,7 @@ async function main(nodeName, zone, tpuType, tpuSoftwareVersion) {
       zone,
       apiVersion: tpu.Node.ApiVersion.V2,
       acceleratorType: tpuType,
-      // Ensure that the tpuSoftwareVersion you're using (e.g., 'tpu-vm-base') is a valid and supported TensorFlow version for the selected tpuType and zone.
+      // Ensure that the tpuSoftwareVersion you're using (e.g., 'tpu-vm-tf-2.14.1') is a valid and supported TensorFlow version for the selected tpuType and zone.
       // You can find a list of supported versions in the Cloud TPU documentation: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions
       tensorflowVersion: tpuSoftwareVersion,
       networkConfig: tpu.NetworkConfig({enableExternalIps: true}),
@@ -66,12 +66,14 @@ async function main(nodeName, zone, tpuType, tpuSoftwareVersion) {
     const parent = `projects/${projectId}/locations/${zone}`;
     const request = {parent, node, nodeId: nodeName};
 
-    const [response] = await tpuClient.createNode(request);
+    const [operation] = await tpuClient.createNode(request);
+
+    // Wait for the delete operation to complete.
+    const [response] = await operation.promise();
 
-    console.log(`TPU VM: ${nodeName} created.`);
     console.log(JSON.stringify(response));
+    console.log(`TPU VM: ${nodeName} created.`);
   }
-
   await callCreateTpuVM();
   // [END tpu_vm_create]
 }

compute/tpu/vmDelete.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(nodeName, zone) {
+  // [START tpu_vm_delete]
+  // Import the TPU library
+  const tpuLib = require('@google-cloud/tpu');
+
+  // Instantiate a tpuClient
+  const tpuClient = new tpuLib.TpuClient();
+
+  /**
+   * TODO(developer): Update/uncomment these variables before running the sample.
+   */
+  // Project ID or project number of the Google Cloud project you want to delete a node.
+  const projectId = await tpuClient.getProjectId();
+
+  // The name of TPU to delete.
+  // nodeName = 'node-name-1';
+
+  // The zone, where the TPU is created.
+  // zone = 'us-central1-a';
+
+  async function callDeleteTpuVM() {
+    const request = {
+      name: `projects/${projectId}/locations/${zone}/nodes/${nodeName}`,
+    };
+
+    const [operation] = await tpuClient.deleteNode(request);
+
+    // Wait for the delete operation to complete.
+    await operation.promise();
+
+    console.log(`Node: ${nodeName} deleted.`);
+  }
+
+  await callDeleteTpuVM();
+  // [END tpu_vm_delete]
+}
+
+main(...process.argv.slice(2)).catch(err => {
+  console.error(err);
+  process.exitCode = 1;
+});

compute/tpu/vmList.js

@@ -40,7 +40,6 @@ async function main(zone) {
 
     const [response] = await tpuClient.listNodes(request);
 
-    console.log('VM list retrived.');
     console.log(JSON.stringify(response));
   }