feat: tpu_vm_create_startup_script

Author: Joanna Grycz

tpu/createStartupScriptVM.js

@@ -0,0 +1,96 @@
+/*
+ * 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, tpuType, tpuSoftwareVersion) {
+  // [START tpu_vm_create_startup_script]
+  // Import the TPU library
+  const {TpuClient} = require('@google-cloud/tpu').v2;
+  const {Node, NetworkConfig} =
+    require('@google-cloud/tpu').protos.google.cloud.tpu.v2;
+
+  // 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 you want to create a node.
+  const projectId = await tpuClient.getProjectId();
+
+  // The name of the network you want the TPU 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 TPU node to connect to.
+  const region = 'europe-west4';
+
+  // The name for your TPU.
+  // nodeName = 'node-name-1';
+
+  // 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 = 'europe-west4-a';
+
+  // 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-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-tf-2.17.0-pod-pjrt';
+
+  async function callCreateTpuVMStartupScript() {
+    // 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}`,
+      }),
+      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`,
+      },
+    });
+
+    const parent = `projects/${projectId}/locations/${zone}`;
+    const request = {parent, node, nodeId: nodeName};
+
+    const [operation] = await tpuClient.createNode(request);
+
+    // Wait for the create operation to complete.
+    const [response] = await operation.promise();
+
+    console.log(JSON.stringify(response));
+  }
+  await callCreateTpuVMStartupScript();
+  // [END tpu_vm_create_startup_script]
+}
+
+main(...process.argv.slice(2)).catch(err => {
+  console.error(err);
+  process.exitCode = 1;
+});

tpu/vmCreateTopology.js renamed to tpu/createTopologyVM.js

@@ -52,7 +52,7 @@ async function main(nodeName, zone, tpuSoftwareVersion) {
 
   // The version of the Cloud TPU you want to create.
   // Available options: TYPE_UNSPECIFIED = 0, V2 = 2, V3 = 4, V4 = 7
-  const tpuVersion = AcceleratorConfig.Type.V2;
+  const tpuVersion = AcceleratorConfig.Type.V3;
 
   // The physical topology of your TPU slice.
   // For more information about topology for each TPU version,
@@ -87,7 +87,6 @@ async function main(nodeName, zone, tpuSoftwareVersion) {
     const [response] = await operation.promise();
 
     console.log(JSON.stringify(response));
-    console.log(`TPU VM: ${nodeName} created.`);
   }
   await callCreateTpuVMTopology();
   // [END tpu_vm_create_topology]

tpu/test/createStartupScriptVM.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 {getStaleNodes, deleteNode} = require('./util');
+
+const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
+const cwd = path.join(__dirname, '..');
+
+describe('Compute tpu', async () => {
+  const nodePrefix = 'node-name-startup-script-2a2b3c';
+  const nodeName = `${nodePrefix}${Math.floor(Math.random() * 1000 + 1)}`;
+  const zone = 'us-central1-b';
+  const tpuType = 'v3-8';
+  const tpuSoftwareVersion = 'tpu-vm-tf-2.17.0-pod-pjrt';
+
+  after(async () => {
+    // Clean-up resources
+    const nodes = await getStaleNodes(nodePrefix);
+    await Promise.all(nodes.map(node => deleteNode(node.zone, node.nodeName)));
+  });
+
+  it('should create a new tpu with startup 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 = JSON.parse(
+      execSync(
+        `node ./createStartupScriptVM.js ${nodeName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
+        {
+          cwd,
+        }
+      )
+    );
+
+    assert.deepEqual(response.metadata, metadata);
+  });
+});

tpu/test/vmTopology.test.js renamed to tpu/test/createTopologyVM.test.js

@@ -18,7 +18,7 @@
 
 const path = require('path');
 const assert = require('node:assert/strict');
-const {before, after, describe, it} = require('mocha');
+const {after, describe, it} = require('mocha');
 const cp = require('child_process');
 const {getStaleNodes, deleteNode} = require('./util');
 
@@ -28,27 +28,27 @@ const cwd = path.join(__dirname, '..');
 describe('Compute tpu with topology', async () => {
   const nodePrefix = 'topology-node-name-2a2b3c';
   const nodeName = `${nodePrefix}${Math.floor(Math.random() * 1000 + 1)}`;
-  const zone = 'europe-west4-a';
+  const zone = 'us-central1-a';
   const tpuSoftwareVersion = 'tpu-vm-tf-2.17.0-pod-pjrt';
 
-  before(async () => {
-    // Cleanup resources
-    const nodes = await getStaleNodes(nodePrefix, zone);
-    await Promise.all(nodes.map(node => deleteNode(zone, node.nodeName)));
-  });
-
   after(async () => {
-    // Delete node
-    await deleteNode(zone, nodeName);
+    // Clean-up resources
+    const nodes = await getStaleNodes(nodePrefix);
+    await Promise.all(nodes.map(node => deleteNode(node.zone, node.nodeName)));
   });
 
-  it('should create a new tpu', () => {
-    const response = execSync(
-      `node ./vmCreateTopology.js ${nodeName} ${zone} ${tpuSoftwareVersion}`,
-      {
-        cwd,
-      }
+  it('should create a new tpu with topology', () => {
+    const acceleratorConfig = {type: 'V3', topology: '2x2'};
+
+    const response = JSON.parse(
+      execSync(
+        `node ./createTopologyVM.js ${nodeName} ${zone} ${tpuSoftwareVersion}`,
+        {
+          cwd,
+        }
+      )
     );
-    assert(response.includes(`TPU VM: ${nodeName} created.`));
+
+    assert.deepEqual(response.acceleratorConfig, acceleratorConfig);
   });
 });

tpu/test/util.js

@@ -18,30 +18,49 @@ const {TpuClient} = require('@google-cloud/tpu').v2;
 
 const tpuClient = new TpuClient();
 
+async function getTpuZones() {
+  const projectId = await tpuClient.getProjectId();
+  const parent = `projects/${projectId}/locations/-`; // List zones for the project
+
+  const [operations] = await tpuClient.listAcceleratorTypes({parent});
+  const zones = new Set();
+  operations.forEach(operation => {
+    zones.add(operation.name.split('/')[3]);
+  });
+
+  return Array.from(zones);
+}
+
 /**
  * Get nodes created more than one hour ago.
  */
-async function getStaleNodes(prefix, zone) {
+async function getStaleNodes(prefix) {
   const projectId = await tpuClient.getProjectId();
   const result = [];
   const currentDate = new Date();
-  currentDate.setHours(currentDate.getHours() - 3);
+  currentDate.setHours(currentDate.getHours() - 1);
 
-  const listNodesAsyncRequest = tpuClient.listNodesAsync({
-    parent: `projects/${projectId}/locations/${zone}`,
-  });
+  const zones = await getTpuZones();
 
-  for await (const tpuObject of listNodesAsyncRequest) {
-    const name = tpuObject.name.split('/').slice(-1)[0];
-    const data = new Date(tpuObject.createTime.nanos / 1000000);
-    if (data < currentDate && name.startsWith(prefix)) {
-      result.push({
-        nodeName: name,
-        timestamp: tpuObject.createTime,
-      });
+  for (const zone of zones) {
+    const [list] = await tpuClient.listNodes({
+      parent: `projects/${projectId}/locations/${zone}`,
+    });
+
+    for (const tpuObject of list) {
+      const name = tpuObject.name.split('/').slice(-1)[0];
+      const data = new Date(tpuObject.createTime.nanos / 1000000);
+      if (data < currentDate) {
+        if (data < currentDate && name.startsWith(prefix)) {
+          result.push({
+            zone,
+            nodeName: name,
+            timestamp: tpuObject.createTime,
+          });
+        }
+      }
     }
   }
-
   return result;
 }
 
@@ -52,7 +71,7 @@ async function deleteNode(zone, nodeName) {
     name: `projects/${projectId}/locations/${zone}/nodes/${nodeName}`,
   };
 
-  console.log('Deleting node: ', nodeName);
+  console.log('Deleting node:', nodeName);
 
   const [operation] = await tpuClient.deleteNode(request);