Skip to content

Commit 964cafb

Browse files
author
Joanna Grycz
committed
feat: tpu_queued_resources_time_bound
1 parent 3949206 commit 964cafb

6 files changed

+234
-29
lines changed

‎tpu/queuedResources/createQueuedResourceNetwork.js

-3
Original file line numberDiff line numberDiff line change
@@ -112,9 +112,6 @@ async function main(
112112

113113
// You can wait until TPU Node is READY,
114114
// and check its status using getTpuVm() from `tpu_vm_get` sample.
115-
console.log(
116-
`Queued resource ${queuedResourceName} with specified network created.`
117-
);
118115
console.log(JSON.stringify(response));
119116
}
120117
await callCreateQueuedResourceNetwork();

‎tpu/queuedResources/createQueuedResourceStartupScript.js

-3
Original file line numberDiff line numberDiff line change
@@ -116,9 +116,6 @@ async function main(
116116

117117
// You can wait until TPU Node is READY,
118118
// and check its status using getTpuVm() from `tpu_vm_get` sample.
119-
console.log(
120-
`Queued resource ${queuedResourceName} with start-up script created.`
121-
);
122119
console.log(JSON.stringify(response));
123120
}
124121
await callCreateQueuedResourceStartupScript();

‎tpu/queuedResources/createQueuedResourceTimeBound.js

+155
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
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_time_bound]
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 callCreateQueuedResourceTimeBound() {
68+
// Create a node
69+
const node = new Node({
70+
name: nodeName,
71+
zone,
72+
acceleratorType: tpuType,
73+
runtimeVersion: tpuSoftwareVersion,
74+
// Define network
75+
networkConfig: new NetworkConfig({
76+
enableExternalIps: true,
77+
network: `projects/${projectId}/global/networks/${networkName}`,
78+
subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`,
79+
}),
80+
queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
81+
});
82+
83+
// Define parent for requests
84+
const parent = `projects/${projectId}/locations/${zone}`;
85+
86+
// Create queued resource
87+
const queuedResource = new QueuedResource({
88+
name: queuedResourceName,
89+
tpu: {
90+
nodeSpec: [
91+
{
92+
parent,
93+
node,
94+
nodeId: nodeName,
95+
},
96+
],
97+
},
98+
queueingPolicy: new QueuedResource.QueueingPolicy({
99+
// You can specify a duration after which a resource should be allocated.
100+
validAfterDuration: {
101+
// format: hour * 3600s
102+
seconds: 6 * 3600,
103+
},
104+
// You can specify how long a queued resource request remains valid.
105+
// validUntilDuration: {
106+
// // format: hour * 3600s
107+
// seconds: 6 * 3600,
108+
// },
109+
// You can specify a time after which a resource should be allocated.
110+
// validAfterTime: {
111+
// // format: new Date('YOUR_TIMESTAMP').getTime() / 1000
112+
// seconds: new Date('2024-10-25T11:45:00Z').getTime() / 1000,
113+
// },
114+
// You can specify a time before which the resource should be allocated.
115+
// validUntilTime: {
116+
// // format: new Date('YOUR_TIMESTAMP').getTime() / 1000
117+
// seconds: new Date('2024-10-25T11:45:00Z').getTime() / 1000,
118+
// },
119+
// You can specify an allocation interval. `startTime` specifies the beginning of the allocation interval
120+
// and `endTime` specifies the end of the allocation interval.
121+
// validInterval: {
122+
// // format: new Date('YOUR_TIMESTAMP').getTime() / 1000
123+
// startTime: {
124+
// seconds: new Date('2024-10-25T15:45:00Z').getTime() / 1000,
125+
// },
126+
// endTime: {
127+
// seconds: new Date('2024-10-26T11:45:00Z').getTime() / 1000,
128+
// },
129+
// },
130+
}),
131+
});
132+
133+
const request = {
134+
parent: `projects/${projectId}/locations/${zone}`,
135+
queuedResource,
136+
queuedResourceId: queuedResourceName,
137+
};
138+
139+
const [operation] = await tpuClient.createQueuedResource(request);
140+
141+
// Wait for the create operation to complete.
142+
const [response] = await operation.promise();
143+
144+
// You can wait until TPU Node is READY,
145+
// and check its status using getTpuVm() from `tpu_vm_get` sample.
146+
console.log(JSON.stringify(response));
147+
}
148+
await callCreateQueuedResourceTimeBound();
149+
// [END tpu_queued_resources_time_bound]
150+
}
151+
152+
main(...process.argv.slice(2)).catch(err => {
153+
console.error(err);
154+
process.exitCode = 1;
155+
});

‎tpu/test/createQueuedResourceNetwork.test.js

+11-11
Original file line numberDiff line numberDiff line change
@@ -48,25 +48,25 @@ describe('TPU queued resource with specified network', async () => {
4848
);
4949
});
5050

51-
it('should create queued resource with specified network', () => {
51+
it('should create queued resource', () => {
5252
const networkConfig = {
5353
network: `projects/${projectId}/global/networks/compute-tpu-network`,
5454
subnetwork: `projects/${projectId}/regions/europe-west4/subnetworks/compute-tpu-network`,
5555
enableExternalIps: true,
5656
};
5757

58-
const response = execSync(
59-
`node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
60-
{
61-
cwd,
62-
}
58+
const response = JSON.parse(
59+
execSync(
60+
`node ./queuedResources/createQueuedResourceNetwork.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
61+
{
62+
cwd,
63+
}
64+
)
6365
);
6466

65-
assert(
66-
response.includes(
67-
`Queued resource ${queuedResourceName} with specified network created.`
68-
)
67+
assert.deepEqual(
68+
response.tpu.nodeSpec[0].node.networkConfig,
69+
networkConfig
6970
);
70-
assert(response.includes(JSON.stringify(networkConfig)));
7171
});
7272
});

‎tpu/test/createQueuedResourceStartupScript.test.js

+9-12
Original file line numberDiff line numberDiff line change
@@ -41,24 +41,21 @@ describe('TPU queued resource with start-up script', async () => {
4141
);
4242
});
4343

44-
it('should create queued resource with start-up script', () => {
44+
it('should create queued resource', () => {
4545
const metadata = {
4646
'startup-script':
4747
'#!/bin/bash\n echo "Hello World" > /var/log/hello.log\n sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1',
4848
};
4949

50-
const response = execSync(
51-
`node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
52-
{
53-
cwd,
54-
}
55-
);
56-
57-
assert(
58-
response.includes(
59-
`Queued resource ${queuedResourceName} with start-up script created.`
50+
const response = JSON.parse(
51+
execSync(
52+
`node ./queuedResources/createQueuedResourceStartupScript.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
53+
{
54+
cwd,
55+
}
6056
)
6157
);
62-
assert(response.includes(JSON.stringify(metadata)));
58+
59+
assert.deepEqual(response.tpu.nodeSpec[0].node.metadata, metadata);
6360
});
6461
});

‎tpu/test/createQueuedResourceTimeBound.test.js

+59
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
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, describe, it} = require('mocha');
22+
const cp = require('child_process');
23+
24+
const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
25+
const cwd = path.join(__dirname, '..');
26+
27+
describe('TPU time bound queued resource', async () => {
28+
const queuedResourceName = `queued-resource-time-bound-${Math.floor(Math.random() * 1000 + 1)}`;
29+
const nodeName = `node-time-bound-2a2b3c${Math.floor(Math.random() * 1000 + 1)}`;
30+
const zone = 'us-west4-a';
31+
const tpuType = 'v5litepod-1';
32+
const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
33+
34+
after(() => {
35+
// Delete queued resource
36+
execSync(
37+
`node ./queuedResources/forceDeleteQueuedResource.js ${queuedResourceName} ${zone}`,
38+
{
39+
cwd,
40+
}
41+
);
42+
});
43+
44+
it('should create queued resource', () => {
45+
const response = JSON.parse(
46+
execSync(
47+
`node ./queuedResources/createQueuedResourceTimeBound.js ${nodeName} ${queuedResourceName} ${zone} ${tpuType} ${tpuSoftwareVersion}`,
48+
{
49+
cwd,
50+
}
51+
)
52+
);
53+
54+
assert.ok(response.queueingPolicy);
55+
assert.ok(response.queueingPolicy.validAfterTime);
56+
assert(typeof response.queueingPolicy.validAfterTime.seconds, 'string');
57+
assert(typeof response.queueingPolicy.validAfterTime.nano, 'number');
58+
});
59+
});

0 commit comments

Comments
 (0)