Skip to content

Commit bdfab57

Browse files
author
Joanna Grycz
committed
feat: tpu_queued_resources_create_network
1 parent ff0c858 commit bdfab57

4 files changed

+209
-3
lines changed
Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
'use strict';
18+
19+
async function main(
20+
nodeName,
21+
queuedResourceName,
22+
zone,
23+
tpuType,
24+
tpuSoftwareVersion
25+
) {
26+
// [START tpu_queued_resources_create_network]
27+
// Import the TPU library
28+
const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
29+
const {Node, NetworkConfig, QueuedResource} =
30+
require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1;
31+
32+
// Instantiate a tpuClient
33+
const tpuClient = new TpuClient();
34+
35+
/**
36+
* TODO(developer): Update/uncomment these variables before running the sample.
37+
*/
38+
// Project ID or project number of the Google Cloud project, where you want to create queued resource.
39+
const projectId = await tpuClient.getProjectId();
40+
41+
// The name of the network you want the node to connect to. The network should be assigned to your project.
42+
const networkName = 'compute-tpu-network';
43+
44+
// The region of the network, that you want the node to connect to.
45+
const region = 'europe-west4';
46+
47+
// The name for your queued resource.
48+
// queuedResourceName = 'queued-resource-1';
49+
50+
// The name for your node.
51+
// nodeName = 'node-name-1';
52+
53+
// The zone in which to create the node.
54+
// For more information about supported TPU types for specific zones,
55+
// see https://cloud.google.com/tpu/docs/regions-zones
56+
// zone = 'europe-west4-a';
57+
58+
// The accelerator type that specifies the version and size of the node you want to create.
59+
// For more information about supported accelerator types for each TPU version,
60+
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
61+
// tpuType = 'v2-8';
62+
63+
// Software version that specifies the version of the node runtime to install. For more information,
64+
// see https://cloud.google.com/tpu/docs/runtimes
65+
// tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
66+
67+
async function callCreateQueuedResourceNetwork() {
68+
// Specify the network and subnetwork that you want to connect your TPU to.
69+
const networkConfig = new NetworkConfig({
70+
enableExternalIps: true,
71+
network: `projects/${projectId}/global/networks/${networkName}`,
72+
subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`,
73+
});
74+
75+
// Create a node
76+
const node = new Node({
77+
name: nodeName,
78+
zone,
79+
acceleratorType: tpuType,
80+
runtimeVersion: tpuSoftwareVersion,
81+
networkConfig,
82+
queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
83+
});
84+
85+
// Define parent for requests
86+
const parent = `projects/${projectId}/locations/${zone}`;
87+
88+
// Create queued resource
89+
const queuedResource = new QueuedResource({
90+
name: queuedResourceName,
91+
tpu: {
92+
nodeSpec: [
93+
{
94+
parent,
95+
node,
96+
nodeId: nodeName,
97+
},
98+
],
99+
},
100+
});
101+
102+
const request = {
103+
parent: `projects/${projectId}/locations/${zone}`,
104+
queuedResource,
105+
queuedResourceId: queuedResourceName,
106+
};
107+
108+
const [operation] = await tpuClient.createQueuedResource(request);
109+
110+
// Wait for the create operation to complete.
111+
const [response] = await operation.promise();
112+
113+
// You can wait until TPU Node is READY,
114+
// and check its status using getTpuVm() from `tpu_vm_get` sample.
115+
console.log(
116+
`Queued resource ${queuedResourceName} with specified network created.`
117+
);
118+
console.log(JSON.stringify(response));
119+
}
120+
await callCreateQueuedResourceNetwork();
121+
// [END tpu_queued_resources_create_network]
122+
}
123+
124+
main(...process.argv.slice(2)).catch(err => {
125+
console.error(err);
126+
process.exitCode = 1;
127+
});

tpu/queuedResources/createQueuedResourceStartupScript.js

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,8 @@ async function main(
8181
metadata: {
8282
// The script updates numpy to the latest version and logs the output to a file.
8383
'startup-script': `#!/bin/bash
84-
echo "Hello World" > /var/log/hello.log
85-
sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1`,
84+
echo "Hello World" > /var/log/hello.log
85+
sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1`,
8686
},
8787
});
8888

@@ -112,13 +112,14 @@ async function main(
112112
const [operation] = await tpuClient.createQueuedResource(request);
113113

114114
// Wait for the create operation to complete.
115-
await operation.promise();
115+
const [response] = await operation.promise();
116116

117117
// You can wait until TPU Node is READY,
118118
// and check its status using getTpuVm() from `tpu_vm_get` sample.
119119
console.log(
120120
`Queued resource ${queuedResourceName} with start-up script created.`
121121
);
122+
console.log(JSON.stringify(response));
122123
}
123124
await callCreateQueuedResourceStartupScript();
124125
// [END tpu_queued_resources_startup_script]
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
'use strict';
18+
19+
const path = require('path');
20+
const assert = require('node:assert/strict');
21+
const {after, before, describe, it} = require('mocha');
22+
const cp = require('child_process');
23+
const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
24+
25+
const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
26+
const cwd = path.join(__dirname, '..');
27+
28+
describe('TPU queued resource with specified network', async () => {
29+
const queuedResourceName = `queued-resource-with-network-${Math.floor(Math.random() * 1000 + 1)}`;
30+
const nodeName = `node-with-network-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`;
31+
const zone = 'us-south1-a';
32+
const tpuType = 'v5litepod-1';
33+
const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
34+
let projectId;
35+
36+
before(async () => {
37+
const tpuClient = new TpuClient();
38+
projectId = await tpuClient.getProjectId();
39+
});
40+
41+
after(() => {
42+
// Delete queued resource
43+
execSync(
44+
`node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
45+
{
46+
cwd,
47+
}
48+
);
49+
});
50+
51+
it('should create queued resource with specified network', () => {
52+
const networkConfig = {
53+
network: `projects/${projectId}/global/networks/compute-tpu-network`,
54+
subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`,
55+
enableExternalIps: true,
56+
};
57+
58+
const response = execSync(
59+
`node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
60+
{
61+
cwd,
62+
}
63+
);
64+
65+
assert(
66+
response.includes(
67+
`Queued resource ${queuedResourceName} with specified network created.`
68+
)
69+
);
70+
assert(response.includes(JSON.stringify(networkConfig)));
71+
});
72+
});

tpu/test/createQueuedResourceStartupScript.test.js

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ describe('TPU queued resource with start-up script', async () => {
4242
});
4343

4444
it('should create queued resource with start-up script', () => {
45+
const metadata = {
46+
'startup-script':
47+
'#!/bin/bash\n echo "Hello World" > /var/log/hello.log\n sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1',
48+
};
49+
4550
const response = execSync(
4651
`node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
4752
{
@@ -54,5 +59,6 @@ describe('TPU queued resource with start-up script', async () => {
5459
`Queued resource ${queuedResourceName} with start-up script created.`
5560
)
5661
);
62+
assert(response.includes(JSON.stringify(metadata)));
5763
});
5864
});

0 commit comments

Comments
 (0)