Skip to content

Commit 6de69a1

Browse files
gryczjJoanna Grycziennaekweinmeisterglasnt
authored
feat: tpu_queued_resources_create/delete_force/delete/get/list (#3904)
* feat: tpu_queued_resources_create * feat: tpu_queued_resources_get * refactor for create and delete * Fix samples and tests * feat: tpu_queued_resources_list * Use mocked TPUClient in tests --------- Co-authored-by: Joanna Grycz <[email protected]> Co-authored-by: Jennifer Davis <[email protected]> Co-authored-by: Karl Weinmeister <[email protected]> Co-authored-by: Katie McLaughlin <[email protected]> Co-authored-by: Maciej Strzelczyk <[email protected]>
1 parent 82a20f8 commit 6de69a1

File tree

8 files changed

+615
-1
lines changed

8 files changed

+615
-1
lines changed

‎tpu/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,4 @@
2121
"c8": "^10.0.0",
2222
"mocha": "^10.0.0"
2323
}
24-
}
24+
}

‎tpu/queuedResources/createQueuedResource.js‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
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(tpuClient) {
20+
// [START tpu_queued_resources_create]
21+
// Import the TPUClient
22+
// TODO(developer): Uncomment below line before running the sample.
23+
// const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
24+
const {Node, NetworkConfig, QueuedResource} =
25+
require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1;
26+
27+
// Instantiate a tpuClient
28+
// TODO(developer): Uncomment below line before running the sample.
29+
// tpuClient = new TpuClient();
30+
31+
/**
32+
* TODO(developer): Update these variables before running the sample.
33+
*/
34+
// Project ID or project number of the Google Cloud project, where you want to create queued resource.
35+
const projectId = await tpuClient.getProjectId();
36+
37+
// The name of the network you want the node to connect to. The network should be assigned to your project.
38+
const networkName = 'compute-tpu-network';
39+
40+
// The region of the subnetwork, that you want the node to connect to.
41+
const region = 'us-central1';
42+
43+
// The name for your queued resource.
44+
const queuedResourceName = 'queued-resource-1';
45+
46+
// The name for your node.
47+
const nodeName = 'node-name-1';
48+
49+
// The zone in which to create the node.
50+
// For more information about supported TPU types for specific zones,
51+
// see https://cloud.google.com/tpu/docs/regions-zones
52+
const zone = `${region}-f`;
53+
54+
// The accelerator type that specifies the version and size of the node you want to create.
55+
// For more information about supported accelerator types for each TPU version,
56+
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
57+
const tpuType = 'v2-8';
58+
59+
// Software version that specifies the version of the node runtime to install. For more information,
60+
// see https://cloud.google.com/tpu/docs/runtimes
61+
const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1';
62+
63+
async function callCreateQueuedResource() {
64+
// Create a node
65+
const node = new Node({
66+
name: nodeName,
67+
zone,
68+
acceleratorType: tpuType,
69+
runtimeVersion: tpuSoftwareVersion,
70+
// Define network
71+
networkConfig: new NetworkConfig({
72+
enableExternalIps: true,
73+
network: `projects/${projectId}/global/networks/${networkName}`,
74+
subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`,
75+
}),
76+
queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
77+
});
78+
79+
// Define parent for requests
80+
const parent = `projects/${projectId}/locations/${zone}`;
81+
82+
// Create queued resource
83+
const queuedResource = new QueuedResource({
84+
name: queuedResourceName,
85+
tpu: {
86+
nodeSpec: [
87+
{
88+
parent,
89+
node,
90+
nodeId: nodeName,
91+
},
92+
],
93+
},
94+
// TODO(developer): Uncomment next line if you want to specify reservation.
95+
// reservationName: 'reservation-name/ Before deleting the queued resource it is required to delete the TPU VM.'
96+
});
97+
98+
const request = {
99+
parent: `projects/${projectId}/locations/${zone}`,
100+
queuedResource,
101+
queuedResourceId: queuedResourceName,
102+
};
103+
104+
const [operation] = await tpuClient.createQueuedResource(request);
105+
106+
// Wait for the create operation to complete.
107+
const [response] = await operation.promise();
108+
109+
// You can wait until TPU Node is READY,
110+
// and check its status using callGetTpuVm() from `tpu_vm_get` sample.
111+
console.log(`Queued resource ${queuedResourceName} created.`);
112+
return response;
113+
}
114+
return await callCreateQueuedResource();
115+
// [END tpu_queued_resources_create]
116+
}
117+
118+
module.exports = main;
119+
120+
// TODO(developer): Uncomment below lines before running the sample.
121+
// main(...process.argv.slice(2)).catch(err => {
122+
// console.error(err);
123+
// process.exitCode = 1;
124+
// });

‎tpu/queuedResources/deleteQueuedResource.js‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
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(tpuClient) {
20+
// [START tpu_queued_resources_delete]
21+
// Import the TPUClient
22+
// TODO(developer): Uncomment below line before running the sample.
23+
// const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
24+
25+
// Instantiate a tpuClient
26+
// TODO(developer): Uncomment below line before running the sample.
27+
// tpuClient = new TpuClient();
28+
29+
/**
30+
* TODO(developer): Update these variables before running the sample.
31+
*/
32+
// Project ID or project number of the Google Cloud project, where you want to delete node.
33+
const projectId = await tpuClient.getProjectId();
34+
35+
// The name of queued resource.
36+
const queuedResourceName = 'queued-resource-1';
37+
38+
// The zone of your queued resource.
39+
const zone = 'us-central1-f';
40+
41+
async function callDeleteTpuVM(nodeName) {
42+
const request = {
43+
name: `projects/${projectId}/locations/${zone}/nodes/${nodeName}`,
44+
};
45+
46+
const [operation] = await tpuClient.deleteNode(request);
47+
48+
// Wait for the delete operation to complete.
49+
await operation.promise();
50+
51+
console.log(`Node: ${nodeName} deleted.`);
52+
}
53+
54+
async function callDeleteQueuedResource() {
55+
const request = {
56+
name: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
57+
};
58+
59+
// Retrive node name
60+
const [queuedResource] = await tpuClient.getQueuedResource(request);
61+
const nodeName = queuedResource.tpu.nodeSpec[0].nodeId;
62+
63+
// Before deleting the queued resource it is required to delete the TPU VM.
64+
await callDeleteTpuVM(nodeName);
65+
66+
const [operation] = await tpuClient.deleteQueuedResource(request);
67+
68+
// Wait for the delete operation to complete.
69+
const [response] = await operation.promise();
70+
71+
console.log(`Queued resource ${queuedResourceName} deleted.`);
72+
return response;
73+
}
74+
return await callDeleteQueuedResource();
75+
// [END tpu_queued_resources_delete]
76+
}
77+
78+
module.exports = main;
79+
80+
// TODO(developer): Uncomment below lines before running the sample.
81+
// main(...process.argv.slice(2)).catch(err => {
82+
// console.error(err);
83+
// process.exitCode = 1;
84+
// });

‎tpu/queuedResources/forceDeleteQueuedResource.js‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
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(tpuClient) {
20+
// [START tpu_queued_resources_delete_force]
21+
// Import the TPUClient
22+
// TODO(developer): Uncomment below line before running the sample.
23+
// const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
24+
25+
// Instantiate a tpuClient
26+
// TODO(developer): Uncomment below line before running the sample.
27+
// tpuClient = new TpuClient();
28+
29+
/**
30+
* TODO(developer): Update these variables before running the sample.
31+
*/
32+
// Project ID or project number of the Google Cloud project, where you want to delete node.
33+
const projectId = await tpuClient.getProjectId();
34+
35+
// The name of queued resource.
36+
const queuedResourceName = 'queued-resource-1';
37+
38+
// The zone of your queued resource.
39+
const zone = 'us-central1-f';
40+
41+
async function callForceDeleteQueuedResource() {
42+
const request = {
43+
name: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
44+
force: true,
45+
};
46+
47+
const [operation] = await tpuClient.deleteQueuedResource(request);
48+
49+
// Wait for the delete operation to complete.
50+
const [response] = await operation.promise();
51+
52+
console.log(`Queued resource ${queuedResourceName} deletion forced.`);
53+
return response;
54+
}
55+
return await callForceDeleteQueuedResource();
56+
// [END tpu_queued_resources_delete_force]
57+
}
58+
59+
module.exports = main;
60+
61+
// TODO(developer): Uncomment below lines before running the sample.
62+
// main(...process.argv.slice(2)).catch(err => {
63+
// console.error(err);
64+
// process.exitCode = 1;
65+
// });

‎tpu/queuedResources/getQueuedResource.js‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
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(tpuClient) {
20+
// [START tpu_queued_resources_get]
21+
// Import the TPUClient
22+
// TODO(developer): Uncomment below line before running the sample.
23+
// const {TpuClient} = require('@google-cloud/tpu').v2alpha1;
24+
25+
// Instantiate a tpuClient
26+
// TODO(developer): Uncomment below line before running the sample.
27+
// tpuClient = new TpuClient();
28+
29+
/**
30+
* TODO(developer): Update these variables before running the sample.
31+
*/
32+
// Project ID or project number of the Google Cloud project, where you want to retrive node.
33+
const projectId = await tpuClient.getProjectId();
34+
35+
// The name of queued resource.
36+
const queuedResourceName = 'queued-resource-1';
37+
38+
// The zone of your queued resource.
39+
const zone = 'us-central1-f';
40+
41+
async function callGetQueuedResource() {
42+
const request = {
43+
name: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`,
44+
};
45+
46+
const [response] = await tpuClient.getQueuedResource(request);
47+
48+
console.log(`Queued resource ${queuedResourceName} retrived.`);
49+
return response;
50+
}
51+
return await callGetQueuedResource();
52+
// [END tpu_queued_resources_get]
53+
}
54+
55+
module.exports = main;
56+
57+
// TODO(developer): Uncomment below lines before running the sample.
58+
// main(...process.argv.slice(2)).catch(err => {
59+
// console.error(err);
60+
// process.exitCode = 1;
61+
// });

0 commit comments

Comments
 (0)