feat: tpu_vm_create

Author: Joanna Grycz

compute/package.json

@@ -15,6 +15,7 @@
   },
   "dependencies": {
     "@google-cloud/compute": "^4.0.0",
+    "@google-cloud/tpu": "^3.5.0",
     "@sendgrid/mail": "^8.0.0",
     "nodemailer": "^6.0.0",
     "nodemailer-smtp-transport": "^2.7.4",

compute/test/tpu.test.js

@@ -0,0 +1,50 @@
+/*
+ * 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('chai');
+const {describe, xit} = 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;
+
+  //   before(async () => {
+  //     projectId = await tpuClient.getProjectId();
+  //   });
+
+  xit('should create a new tpu node', () => {
+    const response = execSync(
+      `node ./tpu/vmCreate.js ${nodeName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
+      {
+        cwd,
+      }
+    );
+
+    assert.include(response, `TPU VM: ${nodeName} created.`);
+  });
+});

compute/tpu/vmCreate.js

@@ -0,0 +1,82 @@
+/*
+ * 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]
+  // Import the TPU library
+  const tpuLib = require('@google-cloud/tpu');
+  const tpu = tpuLib.protos.google.cloud.tpu.v2;
+
+  // 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 create a node.
+  const projectId = await tpuClient.getProjectId();
+
+  // 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 = 'us-central1-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-32';
+
+  // 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';
+
+  async function callCreateTpuVM() {
+    // Create a node
+    const node = new tpu.Node({
+      name: nodeName,
+      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.
+      // 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}),
+      shieldedInstanceConfig: tpu.ShieldedInstanceConfig({
+        enableSecureBoot: true,
+      }),
+    });
+    const parent = `projects/${projectId}/locations/${zone}`;
+    const request = {parent, node, nodeId: nodeName};
+
+    const [response] = await tpuClient.createNode(request);
+
+    console.log(`TPU VM: ${nodeName} created.`);
+    console.log(JSON.stringify(response));
+  }
+
+  await callCreateTpuVM();
+  // [END tpu_vm_create]
+}
+
+main(...process.argv.slice(2)).catch(err => {
+  console.error(err);
+  process.exitCode = 1;
+});