feat: tpu_queued_resources_time_bound

GoogleCloudPlatform · Oct 25, 2024 · 83378d8 · 83378d8
1 parent da299ad
commit 83378d8
Showing 4 changed files with 308 additions and 119 deletions.
diff --git a/tpu/queuedResources/createQueuedResourceTimeBound.js b/tpu/queuedResources/createQueuedResourceTimeBound.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_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: {
+          // seconds format: hour * 3600s
+          seconds: 1 * 3600,
+        },
+      }),
+    });
+
+    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(`Time bound Queued resource ${queuedResourceName} created.`);
+    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
@@ -1,72 +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.
- */
+// /*
+//  * 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';
+// '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 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 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;
+// 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();
-  });
+//   before(async () => {
+//     const tpuClient = new TpuClient();
+//     projectId = await tpuClient.getProjectId();
+//   });
 
-  after(() => {
-    // Delete queued resource
-    execSync(
-      `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
-      {
-        cwd,
-      }
-    );
-  });
+//   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,
-    };
+//   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,
-      }
-    );
+//     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)));
-  });
-});
+//     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
@@ -1,64 +1,64 @@
-/*
- * 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.
- */
+// /*
+//  * 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';
+// 'use strict';
 
-const path = require('path');
-const assert = require('node:assert/strict');
-const {after, describe, it} = require('mocha');
-const cp = require('child_process');
+// 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 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';
+// 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,
-      }
-    );
-  });
+//   after(() => {
+//     // Delete queued resource
+//     execSync(
+//       `node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
+//       {
+//         cwd,
+//       }
+//     );
+//   });
 
-  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',
-    };
+//   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}`,
-      {
-        cwd,
-      }
-    );
+//     const response = execSync(
+//       `node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
+//       {
+//         cwd,
+//       }
+//     );
 
-    assert(
-      response.includes(
-        `Queued resource ${queuedResourceName} with start-up script created.`
-      )
-    );
-    assert(response.includes(JSON.stringify(metadata)));
-  });
-});
+//     assert(
+//       response.includes(
+//         `Queued resource ${queuedResourceName} with start-up script created.`
+//       )
+//     );
+//     assert(response.includes(JSON.stringify(metadata)));
+//   });
+// });
diff --git a/tpu/test/createQueuedResourceTimeBound.test.js 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-south1-b';
+  const tpuType = 'v5p-8';
+  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 = execSync(
+      `node ./queuedResources/createQueuedResourceTimeBound.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
+      {
+        cwd,
+      }
+    );
+    console.log(response);
+    assert(
+      response.includes(
+        `Time bound Queued resource ${queuedResourceName} created.`
+      )
+    );
+    // assert(response.includes(JSON.stringify(metadata)));
+  });
+});