refactor for create and delete

Author: Joanna Grycz

tpu/queuedResource/createQueuedResource.js

@@ -108,10 +108,27 @@ async function main(
     const [operation] = await tpuClient.createQueuedResource(request);
 
     // Wait for the create operation to complete.
-    const [response] = await operation.promise();
+    await operation.promise();
 
-    console.log(JSON.stringify(response));
     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]

tpu/queuedResource/deleteQueuedResource.js

@@ -36,17 +36,36 @@ async function main(queuedResourceName, zone) {
   // The zone of your queued resource.
   // zone = 'europe-west4-a';
 
+  async function callDeleteTpuVM(nodeName) {
+    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.`);
+  }
+
   async function callDeleteQueuedResource() {
     const request = {
       name: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
     };
 
+    // Retrive node name
+    const [response] = await tpuClient.getQueuedResource(request);
+    const nodeName = response.tpu.nodeSpec[0].nodeId;
+
+    // Before deleting the queued resource it is required to delete the TPU VM.
+    await callDeleteTpuVM(nodeName);
+
     const [operation] = await tpuClient.deleteQueuedResource(request);
 
     // Wait for the delete operation to complete.
-    const [response] = await operation.promise();
+    await operation.promise();
 
-    console.log(JSON.stringify(response));
     console.log(`Queued resource ${queuedResourceName} deleted.`);
   }
   await callDeleteQueuedResource();

tpu/queuedResource/forceDeleteQueuedResource.js

@@ -45,9 +45,8 @@ async function main(queuedResourceName, zone) {
     const [operation] = await tpuClient.deleteQueuedResource(request);
 
     // Wait for the delete operation to complete.
-    const [response] = await operation.promise();
+    await operation.promise();
 
-    console.log(JSON.stringify(response));
     console.log(`Queued resource ${queuedResourceName} deletion forced.`);
   }
   await callForceDeleteQueuedResource();

tpu/test/forceDeleteQueuedResource.test.js

@@ -0,0 +1,57 @@
+/*
+ * 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 {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 force deletion', async () => {
+  const queuedResourceName = `queued-resource-force-delete-${Math.floor(Math.random() * 1000 + 1)}`;
+  const nodePrefix = 'node-force-delete-2a2b3c';
+  const nodeName = `${nodePrefix}${Math.floor(Math.random() * 1000 + 1)}`;
+  const zone = 'us-central1-c';
+  const tpuType = 'v2-8';
+  const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
+
+  it('should force queued resource deletion', () => {
+    // Create queued resource
+    execSync(
+      `node ./queuedResource/createQueuedResource.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
+      {
+        cwd,
+      }
+    );
+
+    const response = execSync(
+      `node ./queuedResource/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
+      {
+        cwd,
+      }
+    );
+
+    assert(
+      response.includes(
+        `Queued resource ${queuedResourceName} deletion forced.`
+      )
+    );
+  });
+});

tpu/test/queuedResource.test.js

@@ -56,18 +56,14 @@ describe('TPU queued resource', async () => {
     );
   });
 
-  it('should force queued resource deletion', () => {
+  it('should delete queued resource', () => {
     const response = execSync(
-      `node ./queuedResource/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
+      `node ./queuedResource/deleteQueuedResource.js ${queuedResourceName} ${zone}`,
       {
         cwd,
       }
     );
 
-    assert(
-      response.includes(
-        `Queued resource ${queuedResourceName} deletion forced.`
-      )
-    );
+    assert(response.includes(`Queued resource ${queuedResourceName} deleted.`));
   });
 });